Conversation
Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/ok to test b64a8e3 |
Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/ok to test de7b76c |
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
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.
Summary
Align NeMo RL selective MXFP8 generation with the construction-time TE recipe design in NVIDIA/Megatron-LM#7300. No batch-invariance changes.
policy.megatron_cfg.te_precision_config_filefor training and Megatron generation. Remove this PR's obsolete inference parameter regex filters and refit keyword arguments.first_last_layers_bf16,num_layers_at_start_in_bf16, andnum_layers_at_end_in_bf16to the provider. These previously were not applied by NeMo RL setup.MCore handles parameter selection before checkpoint loading and automatically inherits the enclosing MXFP8 init context where appropriate. Users do not need a checkpoint helper or an explicit
inherit_model_init_contextflag. The recipe keeps shared experts, attention, and MTP in BF16.Dependency
Requires NVIDIA/Megatron-LM#7300 and a NeMo RL environment containing that MCore revision. Keep this draft until the MCore change and dependency update land. The current pinned MCore is not sufficient for the new selective-storage behavior.
Validation
Using this PR's NeMo RL checkout with MCore
8faece7a713e53e6a85b54cff42c8d5083897d61, cached Torch 2.11.0+cu130 / TE 2.18.0+27486e03 environment on GB200:nemo_rl/weight_sync/xferdtensor.pyagainst the cached NCCL M-to-N dependency. This PR does not modify that file; this is not a full clean-environment CI pass.End-to-end smoke test passed (exit 0): Qwen/Qwen3-30B-A3B on 4 GB200 GPUs, TP1/EP4, colocated Megatron generation with native in-place refit, vLLM-selected grouped GEMM, NVLS, global MXFP8 parameter storage, routed-experts-only TE recipe, first 2 / last 4 layers BF16, and eager execution. Both short GRPO steps completed rollout, policy logprobs, backward, and optimizer step.
Reported generation KL errors were 0.0039 / 0.0011. The smoke run had zero reward/loss; this is execution coverage, not a convergence, performance, or zero-logprob-parity claim. Batch invariance and R3 router replay were disabled. Nano and non-colocated M-to-N end-to-end runs were not repeated for this PR.