[Feat]: VRAM grows monotonically during training (CPU offloading + torch.compile) on Windows — progressive slowdown, 12GB card

[wykawake]

Describe your use-case.

Environment

• OS: Windows 10 (64-bit)
• GPU: NVIDIA RTX 4070, 12 GB VRAM
• PyTorch: 2.12.0+cu130 (CUDA build 13.0)
• OneTrainer: commit 23df3832e8f213d64c337e39f9365ba46e99fed0 (master, 2026-08-19)
• Model: Krea2 (DiT), LoRA fine-tune

Training configuration (relevant parts)

• Optimizer: AdamW8bit, cosine schedule (LR 1e-4, warmup 200)
• Weight precision/quantization: INT_W8A8
• Offloading: offload_fraction ≈ 0.7 ; activation_offloading=off
• torch.compile : enabled (top-level "compile": true`)
• batch_size = 1 , resolution cap 1024
• Dataset: ~56 images

Observed behavior

During a single continuous training run:

• Dedicated VRAM grows monotonically: from ~9 GB at the start to ~11.7 GB within the first few hundred steps, and continues rising.
• Step time degrades in lockstep: from ~5 s/it to ~20 s/it as VRAM fills.
• The growth is linear with step count (in one measured run, ~+200 MB every 25 steps).

Profiling evidence (torch.cuda memory snapshots)

Comparing active allocations at step 100 vs step 400:

• The leaked tensors are ~8 MB each, and their count grows linearly with steps (~0.5–1 new block per step; ~150 additional blocks over 300 steps).
• These ~8 MB blocks have no Python call stack — they are allocated inside compiled / checkpoint internals, not ordinary Python code.
• Captured call stacks at a later stage point into:
  • eval_frame → runtime_wrappers._backward_impl (torch.compile backward path)
  • checkpoint.py:unpack_hook → load_tensors (checkpoint saved-tensor offload hook)

Code-path context

When CPU offloading is active (offload_fraction > 0 or activation_offloading), OneTrainer uses reentrant checkpointing internally (see modules/util/LayerOffloadConductor.py:728 and modules/util/checkpointing_util.py:160). There is no user-facing toggle for reentrant checkpointing — it is an internal consequence of enabling offload. The growth is observed with this offloading + torch.compile combination enabled.

Additional observations

• The same Krea2 model in a different tool with torch.compile also enabled shows VRAM pinned at ~9.3 GB, so the leak appears specific to OneTrainer's offload + compile + checkpoint combination rather than the model or compile alone.
• During investigation, disabling torch.compile and switching the checkpoint to use_reentrant=False were each tried separately; neither stopped the VRAM growth (speed was slightly faster with reentrant off, but the leak continued).

Request

Please help investigate the VRAM growth. The captured call stacks implicate the torch.compile backward path and the checkpoint saved-tensor offload hook under Windows + PyTorch 2.12. Related context in this area: #1306 (offloading race / memory corruption), #1649 (offloading rewrite, currently open). I can provide the full memory-snapshot stack traces, or re-run with any instrumentation you suggest.

What would you like to see as a solution?

Help fix the VRAM increase

Have you considered alternatives? List them here.

No response