Skip to content

Megatron offload_after_refit leaves params on GPU (re-onload via offload_before_refit) #2775

Description

@yfw

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?)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

MemorybugSomething isn't working

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions