Colocated Megatron reference worker keeps ~17GB per GPU after offload

[rjpower]

With trainer.algorithm.use_kl_loss=true the reference model shares the policy GPUs (colocate_policy_ref). On the Megatron trainer the reference forward for the 67B-A2B snowball policy at the round-5 geometry (1024-token prompts, 8192-token generations, PP2 x EP8 x DP2) completes, but after ref_model.offload_to_cpu() and empty_cache the reference worker process still holds about 17GB on each GPU (nvidia-smi: "Process 14602 has 17.38 GiB memory in use" in the policy OOM report from /power/grug-megatron-snowball-r5-base-2). The policy update then fails with CUDA OOM at 61GB resident plus that residue.

offload_megatron_model_to_cpu moves named_parameters() of the non-DDP reference module to CPU, so the remainder is not parameters: the reference process's NCCL communicators (EP all-to-all, PP, DP), Transformer Engine and grouped-GEMM workspaces, and its CUDA context. Two directions:

• Tear down the reference worker's collective buffers between forwards (destroy and lazily recreate process groups, or size NCCL_BUFFSIZE for the reference role).
• Give the reference model its own GPUs on the Megatron path (colocate_policy_ref=false), at the cost of extra nodes.

Until then the Megatron benchmark configs run with use_kl_loss=false. The reference forward itself is fine: fwd_logprobs_values_reward with the reference model took 26s on Megatron against 32s on FSDP at the same geometry.