Conversation
Async GRPO previously rejected colocated inference outright. Allow it when the generation backend is MegatronGeneration, where the engine shares the training model's GPUs and process groups instead of refitting weights across worlds. The colocated engine is initialized once, before trajectory collection starts, and stays resident from then on: it is paused for the training step and resumed after, but never torn down at a validation boundary, so its KV cache survives across steps. Weight sync degenerates to offloading the training buffers and bumping the weight version, since generation reads the training weights in place. On the generation side, a colocated MegatronGeneration now brings the engine and its HTTP server up at construction and publishes dp_openai_server_base_urls, so nemo_gym can build clients against the same workers. Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
inference_optimized layers are dual-mode: they run the fast decode kernels under InferenceMode and fall back to the trainable Transformer Engine path otherwise. A colocated run can therefore build one model that both trains and generates quickly, with no second model and no weight sync. Gate the existing training-worker guard behind megatron_cfg.allow_inference_optimized_training so this stays opt-in, and register InferenceColumnParallelLinear with the Megatron->HF exporter, whose registry does not know the type used by shared_experts.linear_fc1 in the dual-mode model. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
A colocated engine and the trainer share one NCCL communicator, so decode collectives and training collectives must never interleave on it. Each rank has to drain its own GPU before the barrier that hands the communicator over, in both directions. Two things were wrong. torch.cuda.synchronize() drains the calling thread's current device, and the Ray caller thread has not necessarily called set_device, so it could drain the wrong GPU entirely; target the LOCAL_RANK device explicitly instead. And on wake the barrier ran after the resume signal, which let a rank replay decode graphs while another was still finishing training work; drain and barrier before signalling resume. With graph-captured decode the mismatch is not obvious right away -- the communicator desyncs only after several suspend/resume cycles. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
Distributed checkpoint save runs collectives over the whole world on the same GPUs the colocated engine is decoding on, and those two collective streams race: the run deadlocks partway through the save. Suspend the engine and the collector for the duration of the save and resume only once every checkpoint collective has completed. This means the engine must stay down from the end of the training step through validation when a checkpoint follows, rather than being woken as usual. Decide whether a checkpoint is due once, up front, and reuse that decision -- the timeout check latches, so calling it twice would report the save at one site and skip it at the other. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
Async GRPO's step cadence is uneven -- steps stall waiting for the replay buffer to fill -- so consumed samples plotted against step number misstates throughput and is not comparable between two runs with different rollout behaviour. Log train/consumed_samples against elapsed wall-clock time instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
The collector caps in-flight prompt groups at num_prompts_per_step * max_trajectory_age_steps. With native in-flight Megatron generation each group fans out one actor call per generated sample, so a large enough configured cap saturates the Ray asyncio actor's 1,000-call concurrency limit. Lifecycle calls such as pause and resume_after_refit then queue behind rollout work and the run appears to hang. Bound the cap by actor capacity as well, reserving headroom for control calls, and report when the configured value is lowered. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
…eration mcore_generation_config is merged into megatron_cfg only when MegatronGeneration stands up its own inference Policy. A colocated run reuses the training model as-is, so a generation-side parallelism setting there is silently dropped and the run proceeds with the training layout -- with no indication that the configured one had no effect. Raise on a weight-sharding mismatch (TP, PP, EP, ETP), naming both values, and warn on context_parallel_size, which shards activations rather than weights but still cannot take effect because the engine inherits the training process groups. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
Describe when to use colocated Megatron generation, how to enable it including the dual-mode inference_optimized model, why generation parallelism must match training, and how the engine's lifecycle interacts with training steps, validation, and checkpoint saves. Also correct the claim that async GRPO requires non-colocated inference, which now holds only for the vLLM backend. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
Nemotron-3-Nano-30B-A3B on 2n8g with Megatron generation colocated with training, exercising the async colocated path end to end. Raises the nightly compute ceiling to 3320 GPU hours to accommodate it. Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Signed-off-by: Anil Thomas <anlthms@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
these are the commits to be upstreamed to nemo-rl for async colo with megatron backend. there are a couple of commits that are not must-fixes: