Skip to content

fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs#3281

Open
seonjinn wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
seonjinn:fix/fp8-param-optimizer-fp8-recipe
Open

fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs#3281
seonjinn wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
seonjinn:fix/fp8-param-optimizer-fp8-recipe

Conversation

@seonjinn

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes a first-train-step crash with policy.megatron_cfg.fp8_cfg.fp8_param=true (the MXFP8 param-storage path added in #2633):

TypeError: copy_(): argument 'other' (position 1) must be Tensor, not NoneType
  megatron/core/optimizer/distrib_optimizer.py, in _copy_main_params_to_param_buffer
  (called from megatron_policy_worker.train via _copy_main_params_to_param_buffer)

Root cause

_create_megatron_config sets fp8_recipe on the model config only and never on OptimizerConfig. Megatron's OptimizerConfig.__post_init__ derives:

use_precision_aware_optimizer_no_fp8_or_ds_fp8 = use_precision_aware_optimizer and (
    main_params_dtype != torch.float32
    or (fp8_recipe is None or fp8_recipe == "delayed")   # <-- None counts as "no fp8"
    or optimizer_cpu_offload
)

With use_precision_aware_optimizer: true (the default in the MoE recipes) and the unplumbed fp8_recipe=None, this flag mis-derives True: fp32 master params are held inside FusedAdam and shard_fp32_from_float16_groups is populated with None placeholders. The per-step master-param restage that #2633 added for reuse_grad_buf_for_mxfp8_param_ag then dereferences those Nones.

Megatron pretrain does not hit this because get_megatron_optimizer_config copies every args field - including args.fp8_recipe - into OptimizerConfig. CI does not hit it because the shipped fp8 example configs set use_precision_aware_optimizer: false.

Fix

Plumb fp8_cfg.fp8_recipe into optimizer_kwargs whenever fp8 is enabled, mirroring Megatron pretrain. With the recipe present, the flag derives False, fp32 master shards are materialized, and both the buffer restage and the mxfp8 quantize path see real tensors. No behavior change when fp8 is disabled.

Validation

  • Repro: Qwen3-30B-A3B GRPO, 2 nodes x 8 B200, fp8_cfg: {enabled: true, fp8: e4m3, fp8_recipe: mxfp8, fp8_param: true} - crashes at train step 1 with and without unrelated local patches (two independent arms).
  • With this fix, the optimizer builds correctly and step 1 completes on the same config.

Known remaining issue (out of scope for this PR)

After this fix, the same configuration fails at step 2 in the colocated GRPO flow: the refit offload cycle releases the shared grad/param buffer and the next step's master-param restage runs before the buffer is rematerialized (RuntimeError: setStorage ... storage of size 0 in the same function). This looks like a buffer-lifecycle contract issue between reuse_grad_buf_for_mxfp8_param_ag and the RL offload/refit cycle (related: NVIDIA/Megatron-LM#5083, NVIDIA/Megatron-LM#5236 / Megatron-Bridge#3980). I can file a separate issue with full repro stacks if useful - this PR fixes the config-derivation layer so that investigation can even begin.

Related: #1164 (the NaN caveat documented on fp8_param), #2633 (introduced the affected path).

@seonjinn
seonjinn requested a review from a team as a code owner July 18, 2026 23:44
@copy-pr-bot

copy-pr-bot Bot commented Jul 18, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant