Skip to content

fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs #3281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
seonjinn wants to merge 1 commit into
base: main
Choose a base branch
from
+8 −0
Open

fix: plumb fp8_recipe into OptimizerConfig for fp8_param runs #3281

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions nemo_rl/models/megatron/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1044,6 +1044,14 @@ def _create_megatron_config(
"overlap_param_gather": overlap_param_gather,
"reuse_grad_buf_for_mxfp8_param_ag": reuse_grad_buf_for_mxfp8_param_ag,
}
# OptimizerConfig.__post_init__ treats fp8_recipe=None as "no fp8 params" and
# lets the precision-aware optimizer keep fp32 masters inside FusedAdam,
# leaving None placeholders in shard_fp32_from_float16_groups; with
# reuse_grad_buf_for_mxfp8_param_ag the shared param buffer must be refilled
# from those masters each step, so the recipe has to be plumbed to the
# optimizer just like Megatron pretrain's get_megatron_optimizer_config does.
if fp8_cfg is not None and fp8_cfg.get("enabled", False):
optimizer_kwargs["fp8_recipe"] = fp8_cfg.get("fp8_recipe")

# Fused linear logprobs run the decoder but read output_layer.weight directly
# instead of calling output_layer.forward(). Megatron's distributed-optimizer
Expand Down
Loading