Fix: Skip JIT warmup when fusion is disabled via arguments

[kisseternity]

What does this PR do ?

This PR adds conditional checks to _warmup_jit_function in initialize.py to ensure JIT fusions (Bias+GeLU/SwiGLU and Bias+Dropout+Add) are only warmed up when they are enabled via arguments (e.g., checks args.bias_gelu_fusion, args.bias_dropout_fusion). This prevents unnecessary warmup execution when users explicitly disable these fusions using flags like --no-bias-gelu-fusion.

Pre-checks

• [done ] I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• [done ] I have added relevant unit tests
• [done ] I have added relevant functional tests
• [done ] I have added proper typing to my code Typing guidelines
• [done ] I have added relevant documentation
• [done ] I have run the autoformatter.sh on my PR

[copy-pr-bot]

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.

[asolergi-nv]

Looks nice, but we should check exhaustively the conditions that trigger the torch.compiled functions. Could you also add geglu? Thanks!

[kisseternity]

Looks nice, but we should check exhaustively the conditions that trigger the torch.compiled functions. Could you also add geglu? Thanks!

Thanks for the feedback! I've updated the code to add exhaustive condition checks based on the MLP & shared experts implementation. Please let me know if any additional conditions should be considered.

[asolergi-nv]

Just remove some comments & we are ready!

[kisseternity]

Just remove some comments & we are ready!

Addressed, thanks!

[asolergi-nv]

/ok to test 201e816

[Phlip79]

/ok to test b0d1dd1

[dimapihtar]

/ok to test d3af9b3

[Phlip79]

/claude fix

[svcnvidia-nemo-ci]

❌ Claude fix stopped because a workflow step failed. Inspect the run.

[Phlip79]

@kisseternity please fix failing tests

[guihong-nv]

@kisseternity May you please fix failing tests?

[ilml]

/ok to test d3af9b3

[ilml]

waiting on @kisseternity

[Phlip79]

/claude fix

[svcnvidia-nemo-ci]

❌ Claude fix stopped because a workflow step failed. Inspect the run.

[janEbert]

I'm assuming that the contribution may have changed seeding, since certain torch.rand calls aren't being executed. Therefore, the golden values may also need to be updated.

[kisseternity]

I'm assuming that the contribution may have changed seeding, since certain torch.rand calls aren't being executed. Therefore, the golden values may also need to be updated.

@cspades @janEbert — figured out what's going on and pushed a fix: 8b94191.

Root cause. The PR didn't change the seed itself, but it changed how many
torch.rand(..., device="cuda") calls _warmup_jit_function makes. Every
warmup draw consumes the default CUDA RNG stream, and warmup executes
after seeding: pretrain() → initialize_megatron()
(training/training.py:1084) → _set_random_seed →
model_parallel_cuda_manual_seed() (which calls torch.cuda.manual_seed() on
the default generator), then set_jit_fusion_options() (training.py:1110) →
_warmup_jit_function(). Downstream code also relies on that default stream,
e.g. dropout in non-tensor-parallel regions (see the docstring of
model_parallel_cuda_manual_seed in core/tensor_parallel/random.py).

Previously the warmup executed 5 unconditional draws. With the new conditional
checks, configs that end up with a fusion arg disabled skip some draws —
explicit --no-bias-gelu-fusion / --no-bias-dropout-fusion / --quick-geglu,
and also --disable-bias-linear, because validate_args then forces
args.bias_gelu_fusion = False ("Disable bias gelu fusion if we are disabling
bias altogether", arguments.py:1410). The new GEGLU warmup additionally draws
2×-sized tensors for --gated-linear-unit configs. So the RNG stream position
at training start shifts — exactly as @janEbert suspected.

This is now concrete for the two failing jobs: both ran
gpt3_mcore_te_tp1_pp4_resume_torch_dist(_persistent)_disable_bias_linear,
which set --disable-bias-linear: true → bias_gelu_fusion = False, so the 2
activation-warmup draws that previously ran unconditionally are skipped (the 3
bias-dropout draws still run) — precisely the shift behind the golden-value
mismatch. It also explains why only these two jobs failed: the fusion flags
default to True (store_false), so configs without --no-*-fusion /
--disable-bias-linear / --quick-geglu / --gated-linear-unit still draw
byte-identically to before.

Fix. Instead of regenerating golden values, the new commit makes the warmup
side-effect free: snapshot the CUDA RNG state on entry
(torch.cuda.get_rng_state_all()) and restore it before returning
(torch.cuda.set_rng_state_all()), guarded by torch.cuda.is_available() so
CPU-only runs are unaffected. Skipping arbitrary warmup blocks can no longer
perturb training determinism, so the existing golden values should pass
unchanged. +9 lines in megatron/training/initialize.py only.

(The 2× sizing in the GEGLU warmup itself is intentional and kept: GEGLU
splits the last dim into two halves (torch.chunk(y, 2, -1)), so the previous
code was warming up with wrong shapes. With the RNG snapshot/restore in place
it no longer affects determinism either way.)

Could you re-trigger the merge queue / functional tests when you get a chance?
If the two disable-bias-linear gpt jobs go green with the golden values
untouched, that closes this out.

[janEbert]

/ok to test 8b94191

[maanug-nv]

@kisseternity
I guess we're still waiting on functional tests to confirm, but won't this last commit break a lot of functional tests?
As you mentioned, the warmup previously executed unconditional draws. The golden values for all functional tests are depending on that. In other words, the numerics of the tests expect those unconditional draws in jit warmup.
Before 8b94191, some tests would skip those draws (again as you mentioned), affecting the numerics for those tests. After that commit, all tests are skipping those draws, correct? So, the numerics for all tests will change. After that commit.
The result of 8b94191 is that the new numerics are independent of RNG draws in jit warmup, which is a good design choice. however, if my understanding is correct and it requires updating golden values for many many more tests, i'm not sure if that's worthwhile. i dont think we see changes to the jit warmup that often.

Let me know if there's any flaw in that understanding. I guess the functional tests run on this commit will confirm/refute.

[copy-pr-bot]

/ok to test 845a0e3

@maanug-nv, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[maanug-nv]

/ok to test ed28045

[Phlip79]

/claude fix

[svcnvidia-nemo-ci]

🛠️ Claude fix commit c7671d8126b4 (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Applied black's exact one-hunk reformat to pr-head／tests／unit＿tests／test＿initialize＿jit＿warmup.py: collapsed the 3-line signature of test＿dropout＿warmup＿and＿geglu＿cat onto one 98-char line (limit 100). Whitespace-only; 1 file, +1／-3. Verified via py＿compile, a no-line-over-100 check, and byte-for-byte match with black's diff in the CI log. black is not installed and there is no network, so the formatter was not re-run and the unit test was not executed.

Files changed by Claude

• tests/unit_tests/test_initialize_jit_warmup.py

Why
Run 32280254211 had one real failure: the linting job, where ’black --check --diff’ reported "would reformat tests／unit＿tests／test＿initialize＿jit＿warmup.py" and exited 1, printing the full proposed diff. The other logged error is the Nemo＿CICD＿Test aggregator re-reporting that same job. No unit-test failures and no unmerged paths, so no conflict resolution was needed. A re-regression: c61e8a5 formatted the file, later commit ed28045 reintroduced it.

DCO
@kisseternity, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test c7671d8

[svcnvidia-nemo-ci]

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.