feat(megatron): add manual CPU offload as an alternative to torch_mem…

[yeqinghe609]

…ory_saver

Colocate training offloads the train actor to CPU on every sleep/wake. The only path was torch_memory_saver (TMS) VMM pause()/resume(), which is semantically blind (backs up the whole runtime pool) and unsafe on backends without expandable_segments (e.g. Kunlunxin P800). This adds an explicit, application-level selective offload and lets it be selected per job.

Behavior:

• Default stays TMS.
• New --manual-offload switches to selective offload: only live train state (param flat buffers + optimizer master/Adam state) is copied to CPU; grads and recomputable buffers are just released, then empty_cache() returns the freed pages to the driver.

Use TMS (default) when

torch_memory_saver hook LD_PRELOAD'ed (the launcher does this automatically for offload_train + megatron).
The allocator supports it (notably expandable_segments), so pause()/resume() can unmap/remap physical pages cleanly.
You want the simplest, well-tested path and don't mind that TMS backs up the entire runtime memory pool (it can't tell live optimizer state from dead activation/cache blocks).

Use --manual-offload when

You want selective offload: only the live train state (param flat buffers + optimizer master/Adam state) is copied to CPU; gradients and recomputable buffers are just released and empty_cache() returns the freed pages to the driver. This avoids blindly copying the ~70% dead data that whole-pool TMS backup moves, cutting sleep+wake time significantly on large/MoE/long-sequence models.
You need a hardware-neutral path that relies only on plain PyTorch storage ops (no .so / VMM), consistent across NVIDIA/Others.

Changes:

• weight_update/train_offload.py (new): MegatronTrainStateOffloader facade delegating to _TmsOffloadStrategy / _ManualOffloadStrategy behind an _OffloadStrategy Protocol. Owns TMS margin init and strategy selection (TMS when LD_PRELOAD'ed, else manual; --manual-offload forces manual).
• backends/megatron/actor.py: sleep()/wake_up()/update_weights() delegate to the offloader; drop the inline _torch_memory_saver_enabled block and the torch_memory_saver / nullcontext imports.
• weight_update/common.py: _maybe_get_cpu_backup() reads from the offloader's _relax_cpu_offload_data when a param's GPU storage was freed, and only takes the TMS get_cpu_backup() branch when the hook is actually LD_PRELOAD'ed.
• distributed/ray/actor_group.py: skip LD_PRELOAD + TMS_INIT env vars under --manual-offload so TMS is not the active mechanism.
• utils/arguments.py: add --manual-offload.
• scripts/training/text/run-qwen35-8xklx.sh: pass --manual-offload for P800.

No behavior change for existing TMS/CUDA users (default path untouched).

Testing

Qwen3.6-35B

performance

• wake_up_time: 10s->2s
• sleep_time: 20s->5s

Memory

3G vs 9G(tms offload vs manual offload)
tms offload

manual offload

• pre-commit run --all-files passes
• Tests pass (pytest tests/)
• New tests added (if applicable)
• Documentation updated (if applicable)

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• [ ✅] New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation update
• Refactoring (no functional changes)
• Performance improvement
• CI/CD or build changes

Screenshots / Logs