fix(vlm): use FP32 master weights for Nemotron Omni - #3444
Merged
HuiyingLi merged 3 commits intoAug 15, 2026
Merged
Conversation
Contributor
Author
|
/ok to test 48328c1 |
Contributor
Author
|
/claude review |
HuiyingLi
changed the title
HuiyingLi
marked this pull request as ready for review
August 14, 2026 21:29
HuiyingLi
changed the title
Contributor
Author
|
/claude review |
Contributor
Author
|
/ok to test afa8737 |
Contributor
|
LGTM |
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Remove the projector-freezing feature and the freeze_embeddings config change, keeping only the FP32 master-weights optimizer and activation checkpointing changes. - restore freeze_config to the original tower-wide freezing (freeze_embeddings reinstated, freeze_vision_projector/freeze_audio_projector removed) - revert apply_parameter_freezing to its main-branch behavior - drop the projector-freezing unit tests and freeze_config recipe assertions Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
HuiyingLi
force-pushed
the
huiyingl/fix/nemotron-omni-fp32-master
branch
from
August 14, 2026 22:18
37cf6d0 to
3ea2b46
Compare
akoumpa
approved these changes
Aug 15, 2026
HuiyingLi
enabled auto-merge (squash)
August 15, 2026 00:47
Contributor
Author
|
/ok to test 3ea2b46 |
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
2 participants
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 changed
6.0e-6Why
At
6.0e-6, directly updating BF16-resident parameters causes many optimizer steps to round back to the same BF16 value. Megatron Bridge avoids this by accumulating updates in FP32 main parameters. Using FusedAdam FP32 master weights gives AutoModel the same update precision while retaining BF16 model storage.Impact
The shipped Nemotron Omni CORD-v2 recipe can train at
6.0e-6on one node with 8 H100 GPUs. The change is scoped to that single recipe plus its test — no framework code and no other recipe is touched.Validation
pytest tests/unit_tests/recipes/test_nemotron_omni_cord_recipe.py— 1 passed20-step single-node 8xH100 run with FSDP2, EP8, activation checkpointing, and FP32 master weights: loss
0.7625 -> 0.0811, peak memory approximately 59.6 GiB/GPUre-run on the current head (
3ea2b46), 20 steps + validation on one node with 8xH100, clean exit — wandb: https://wandb.ai/Nemo-automodel/omni-fp32-master/runs/hvhxolfvValidation loss tracks training loss. Three signals confirm the new optimizer path is the one actually running: the
Detected torch.optim.Adam/AdamW with trainable bf16 model parameterswarning thatmainemits every run is absent;lrlogs as6.00e-06rather thanmain's1e-4; and Transformer Engine accepts bothmaster_weightsandmaster_weight_dtype.