fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs#3281
Open
seonjinn wants to merge 1 commit into
Open
fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs#3281seonjinn wants to merge 1 commit into
seonjinn wants to merge 1 commit into
Conversation
Signed-off-by: sna <sna@nvidia.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
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.
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):Root cause
_create_megatron_configsetsfp8_recipeon the model config only and never onOptimizerConfig. Megatron'sOptimizerConfig.__post_init__derives:With
use_precision_aware_optimizer: true(the default in the MoE recipes) and the unplumbedfp8_recipe=None, this flag mis-derivesTrue: fp32 master params are held inside FusedAdam andshard_fp32_from_float16_groupsis populated withNoneplaceholders. The per-step master-param restage that #2633 added forreuse_grad_buf_for_mxfp8_param_agthen dereferences thoseNones.Megatron pretrain does not hit this because
get_megatron_optimizer_configcopies everyargsfield - includingargs.fp8_recipe- intoOptimizerConfig. CI does not hit it because the shipped fp8 example configs setuse_precision_aware_optimizer: false.Fix
Plumb
fp8_cfg.fp8_recipeintooptimizer_kwargswhenever fp8 is enabled, mirroring Megatron pretrain. With the recipe present, the flag derivesFalse, 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
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).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 0in the same function). This looks like a buffer-lifecycle contract issue betweenreuse_grad_buf_for_mxfp8_param_agand 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).