You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
MegatronPolicyWorker.offload_after_refit() leaves params on GPU after a colocated refit, instead of offloading it. Root cause is in #2530:
feat: PPO with MCore #2530 changed offload_before_refit() to unconditionally onload params to GPU (move_model("cuda", move_params=True, ...)) — needed because PPO's new finish_training() offloads params before refit, and streaming needs them back.
But offload_after_refit() reuses offload_before_refit() as its final step (to drop optimizer/grads/caches). So it re-onloads the params it just moved to CPU.
Impact / how it surfaced
Caught by test_vllm_megatron_weight_update_memory (assert current_allocated <= 0.1 → observed 1884.59). The test is mcore + deselected as expensive, so it's latent on main and only runs on gates that include expensive tests. DTensor is unaffected (no finish_training; its offload_before_refit is pure-offload and lands at 0 GB).
Notes:
The onload is only needed for megatron + PPO (after finish_training()); it's a no-op in GRPO (params already resident). Having offload_before_refit onload params seems a bit counterintuitive. Can we instead explicitly re-onload only in this case (PPO?)
MegatronPolicyWorker.offload_after_refit()leaves params on GPU after a colocated refit, instead of offloading it. Root cause is in #2530:move_model("cuda", move_params=True, ...)) — needed because PPO's newfinish_training()offloads params before refit, and streaming needs them back.offload_after_refit()reusesoffload_before_refit()as its final step (to drop optimizer/grads/caches). So it re-onloads the params it just moved to CPU.Impact / how it surfaced
Caught by
test_vllm_megatron_weight_update_memory(assert current_allocated <= 0.1 → observed 1884.59). The test is mcore + deselected as expensive, so it's latent on main and only runs on gates that include expensive tests. DTensor is unaffected (no finish_training; its offload_before_refit is pure-offload and lands at 0 GB).Notes:
The onload is only needed for megatron + PPO (after finish_training()); it's a no-op in GRPO (params already resident). Having
offload_before_refitonload params seems a bit counterintuitive. Can we instead explicitly re-onload only in this case (PPO?)