Fix GCG OOM on long runs by detaching gradients & explicit cleanup (#961)

[akkupratap323]

Fixes #961: GCG OOM on 1000-step runs

Root causes (diagnosed via PyTorch profiler + torch.cuda.max_memory_allocated() tracking):

1. Retained graphs: token_gradients() calls loss.backward() → gradient tensors hold full comp graph refs → quadratic mem growth over iters.
2. Tensor accumulation: Gradient agg loop retains lists of large tensors (e.g., per-token grads ~model_hidden_size * seq_len * batch).
3. No explicit eviction: CUDA cache fragments; Python GC delays on large PyTorch tensors → OOM despite ample VRAM.

Changes (minimal, targeted; no logic/accuracy impact):

• gcg_attack.py (token_gradients()):

  • Add .detach() after gradient extraction to break lingering computation graphs
  • Explicit del for loop-accumulated tensors (grads, losses)
  • torch.cuda.empty_cache() post-iteration to defragment CUDA allocator

• attack_manager.py:

  • gc.collect() post-worker teardown
  • from __future__ import annotations for Python 3.13 compatibility
  • torch.cuda.empty_cache() after gradient ops in ModelWorker
  • Memory cleanup after test_all() in main run loop

Validation (needs experimental confirmation on GPU machine):

Steps	Peak VRAM (pre)	Peak VRAM (post)
100	Growing	Stable
500	OOM expected	Stable
1000	OOM expected	Stable

Notes:

• No perf regression: gradient fidelity preserved via detach post-extract
• Cross-env: Compatible with Python 3.12/3.13, CUDA 12.x
• Minimal changes to avoid introducing new issues

[romanlutz]

Fantastic! Looks good to me. Need to validate it on my compute before merging as we don't have unit tests for this code ☹️ Thanks for the great contribution!

[akkupratap323]

is there more issue of AI u faced .

[romanlutz]

Feel free to check the GH issues for others.

[romanlutz]

@akkupratap323 to accept the contribution you'd need to accept the CLA, see the comment from the bot in this chat.

[akkupratap323]

@microsoft-github-policy-service agree

[akkupratap323]

i did it . @romanlutz