Fix: Skip JIT warmup when fusion is disabled via arguments

megatron/training/initialize.py

@@ -14,6 +14,7 @@
 
 from megatron.core import mpu, tensor_parallel
 from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train
+from megatron.core.fusions.fused_bias_geglu import bias_geglu
 from megatron.core.fusions.fused_bias_gelu import bias_gelu
 from megatron.core.fusions.fused_bias_swiglu import bias_swiglu
 from megatron.core.parallel_state import create_group
@@ -526,7 +527,41 @@ def _warmup_jit_function(tp_size=None):
     else:
         dtype = torch.float32
 
-    # Warmup fused bias+gelu
+    # Check if TE activation function is used (in which case, no need to warmup custom fusions)
+    use_te_activation_func = getattr(args, 'use_te_activation_func', False)
+    gated_linear_unit = getattr(args, 'gated_linear_unit', False)
+
+    # Warmup bias_swiglu: swiglu activation (F.silu + GLU)
+    warmup_bias_swiglu = (
+        not use_te_activation_func
+        and args.swiglu
+        and args.bias_swiglu_fusion
+    )
+
+    warmup_bias_gelu = (
+        not use_te_activation_func
+        and not args.swiglu
+        and not getattr(args, 'quick_geglu', False)
+        and not gated_linear_unit
+        and args.bias_gelu_fusion
+    )
+
+    warmup_bias_geglu = (
+        not use_te_activation_func
+        and not args.swiglu
+        and not getattr(args, 'quick_geglu', False)
+        and gated_linear_unit
+        and args.bias_gelu_fusion
+    )
+
+    # NOTE: the torch.rand draws and the fused bias+dropout+add warmup below
+    # consume the default CUDA RNG stream before training starts. Downstream
+    # training numerics (and the golden values of the functional tests) encode
+    # exactly those consumptions, so they must stay UNCONDITIONAL with their
+    # historical shapes and order, independent of any argument; otherwise the
+    # RNG stream shifts for configs that differ from the golden baselines.
+    # Only the bias_swiglu/bias_gelu/bias_geglu calls are gated on arguments:
+    # they consume no RNG, so skipping them cannot perturb determinism.
     bias = torch.rand(
         args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda"
     )
@@ -539,18 +574,38 @@ def _warmup_jit_function(tp_size=None):
         dtype=dtype,
         device="cuda",
     )
+
     # Warmup JIT fusions with the input grad_enable state of both forward
-    # prop and recomputation
-    for bias_grad, input_grad in zip([True, True], [False, True]):
-        bias.requires_grad, input.requires_grad = bias_grad, input_grad
-        for _ in range(5):
-            if args.swiglu:
+    # prop and recomputation. The three branches are mutually exclusive.
+    if warmup_bias_swiglu:
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            bias.requires_grad, input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
                 output = bias_swiglu(input, bias)
-            else:
+        del output
+    elif warmup_bias_gelu:
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            bias.requires_grad, input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
                 output = bias_gelu(bias, input)
-    del bias, input, output
-
-    # Warmup fused bias+dropout+add
+        del output
+    elif warmup_bias_geglu:
+        # bias_geglu splits the last dim into two GLU halves, so its tensors
+        # are 2x ffn-sized. torch.cat consumes no RNG, so doubling the drawn
+        # tensors here keeps the RNG stream identical to the historical code.
+        geglu_bias = torch.cat([bias, bias], dim=-1)
+        geglu_input = torch.cat([input, input], dim=-1)
+        for bias_grad, input_grad in zip([True, True], [False, True]):
+            geglu_bias.requires_grad, geglu_input.requires_grad = bias_grad, input_grad
+            for _ in range(5):
+                output = bias_geglu(geglu_bias, geglu_input)
+        del output, geglu_bias, geglu_input
+    del bias, input
+
+    # Warmup fused bias+dropout+add. Intentionally NOT gated on
+    # args.bias_dropout_fusion: the fused train calls consume RNG for the
+    # dropout mask, and conditioning them would shift downstream determinism
+    # relative to the historical behavior encoded in the golden values.
     if args.sequence_parallel:
         # tp_size threaded by the caller (hetero MIMO language PGC); None -> mpu.
         seq_length = args.seq_length // (tp_size or mpu.get_tensor_model_parallel_world_size())
@@ -577,6 +632,7 @@ def _warmup_jit_function(tp_size=None):
         for _ in range(5):
             output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate)
     del bias, input, residual, output
+
     torch.cuda.empty_cache()
 
 

tests/unit_tests/test_initialize_jit_warmup.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""Tests that JIT warmup skips only RNG-free fused kernels when disabled.
+
+Uses mocks only (no CUDA required). Patches use
+``megatron.training.initialize.<name>`` because ``_warmup_jit_function`` closes
+over that module's globals (imported ``get_args``, ``torch``, fusion ops).
+
+Determinism contract under test: the ``torch.rand`` warmup draws and the fused
+bias+dropout+add warmup consume the default CUDA RNG stream before training
+starts, and downstream numerics (functional-test golden values) encode exactly
+those consumptions. They must therefore run unconditionally with their
+historical shapes and order regardless of arguments. Only the RNG-free
+bias_swiglu/bias_gelu/bias_geglu kernel calls may be skipped.
+"""
+
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from megatron.training.initialize import _warmup_jit_function
+
+
+def _args(**overrides):
+    base = dict(
+        bf16=False,
+        fp16=False,
+        use_te_activation_func=False,
+        gated_linear_unit=False,
+        quick_geglu=False,
+        swiglu=False,
+        bias_swiglu_fusion=True,
+        bias_gelu_fusion=True,
+        bias_dropout_fusion=True,
+        sequence_parallel=False,
+        ffn_hidden_size=64,
+        tensor_model_parallel_size=1,
+        seq_length=8,
+        context_parallel_size=1,
+        micro_batch_size=1,
+        hidden_size=32,
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+def _fake_rand(*size, **kwargs):
+    """Mirror ``torch.rand`` shapes used in ``_warmup_jit_function`` but on CPU."""
+    dtype = kwargs.get("dtype", torch.float32)
+    if len(size) == 1 and isinstance(size[0], tuple):
+        shape = size[0]
+    elif size:
+        shape = size if len(size) > 1 else (size[0],)
+    else:
+        shape = ()
+    return torch.zeros(shape, dtype=dtype, device="cpu")
+
+
+@mock.patch("megatron.training.initialize.torch.cuda.empty_cache")
+@mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+@mock.patch("megatron.training.initialize.torch.rand", side_effect=_fake_rand)
+@mock.patch("megatron.training.initialize.get_args")
+class TestJitWarmupSkippedWhenFusionDisabled:
+    @mock.patch("megatron.training.initialize.bias_swiglu")
+    def test_skips_bias_swiglu(self, mock_swiglu, mock_get_args, mock_rand, mock_dropout, _cache):
+        mock_get_args.return_value = _args(swiglu=True, bias_swiglu_fusion=False)
+        _warmup_jit_function()
+        mock_swiglu.assert_not_called()
+        # RNG-consuming draws and dropout warmup stay unconditional for
+        # golden-value parity; mock dropout call count: zip(...) × range(5).
+        assert mock_dropout.call_count == 10
+        assert mock_rand.call_count == 5
+
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    def test_skips_bias_gelu(self, mock_gelu, mock_get_args, mock_rand, mock_dropout, _cache):
+        mock_get_args.return_value = _args(
+            swiglu=False, bias_gelu_fusion=False, gated_linear_unit=False
+        )
+        _warmup_jit_function()
+        mock_gelu.assert_not_called()
+        assert mock_dropout.call_count == 10
+        assert mock_rand.call_count == 5
+
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    def test_skips_bias_geglu(self, mock_geglu, mock_get_args, mock_rand, mock_dropout, _cache):
+        mock_get_args.return_value = _args(
+            swiglu=False, bias_gelu_fusion=False, gated_linear_unit=True
+        )
+        _warmup_jit_function()
+        mock_geglu.assert_not_called()
+        assert mock_dropout.call_count == 10
+        assert mock_rand.call_count == 5
+
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    def test_skips_gelu_and_geglu_when_quick_geglu(
+        self, mock_gelu, mock_geglu, mock_get_args, mock_rand, mock_dropout, _cache
+    ):
+        # Even with bias_gelu_fusion True, quick_geglu skips gelu and geglu warmups.
+        mock_get_args.return_value = _args(
+            swiglu=False, quick_geglu=True, bias_gelu_fusion=True, gated_linear_unit=True
+        )
+        _warmup_jit_function()
+        mock_gelu.assert_not_called()
+        mock_geglu.assert_not_called()
+        # Draws and dropout warmup are still unconditional.
+        assert mock_dropout.call_count == 10
+        assert mock_rand.call_count == 5
+
+    @mock.patch("megatron.training.initialize.bias_geglu")
+    @mock.patch("megatron.training.initialize.bias_gelu")
+    @mock.patch("megatron.training.initialize.bias_swiglu")
+    def test_te_activation_skips_activation_warmups(
+        self, mock_swiglu, mock_gelu, mock_geglu, mock_get_args, mock_rand, mock_dropout, _cache
+    ):
+        # TE activation skips only the RNG-free activation kernels; the
+        # RNG-consuming dropout warmup still runs for golden-value parity.
+        mock_get_args.return_value = _args(
+            use_te_activation_func=True,
+            swiglu=True,
+            bias_swiglu_fusion=True,
+            gated_linear_unit=True,
+            bias_gelu_fusion=True,
+        )
+        _warmup_jit_function()
+        mock_swiglu.assert_not_called()
+        mock_gelu.assert_not_called()
+        mock_geglu.assert_not_called()
+        assert mock_dropout.call_count == 10
+        assert mock_rand.call_count == 5
+
+
+@mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+@mock.patch("megatron.training.initialize.bias_geglu")
+@mock.patch("megatron.training.initialize.bias_gelu")
+@mock.patch("megatron.training.initialize.torch.cuda.empty_cache")
+@mock.patch("megatron.training.initialize.torch.rand", side_effect=_fake_rand)
+@mock.patch("megatron.training.initialize.bias_swiglu")
+@mock.patch("megatron.training.initialize.get_args")
+def test_calls_bias_swiglu_when_fusion_enabled(
+    mock_get_args, mock_swiglu, mock_torch_rand, _empty_cache, mock_gelu, mock_geglu, mock_dropout
+):
+    mock_get_args.return_value = _args(swiglu=True, bias_swiglu_fusion=True)
+    mock_swiglu.return_value = torch.tensor(0.0)
+    _warmup_jit_function()
+    assert mock_swiglu.call_count == 10  # zip([True,True],[False,True]) × range(5)
+    mock_gelu.assert_not_called()
+    mock_geglu.assert_not_called()
+    assert mock_dropout.call_count == 10  # unconditional for RNG parity
+    assert mock_torch_rand.call_count == 5  # all draws unconditional
+
+
+@mock.patch("megatron.training.initialize.bias_geglu")
+@mock.patch("megatron.training.initialize.torch.cuda.empty_cache")
+@mock.patch("megatron.training.initialize.torch.rand", side_effect=_fake_rand)
+@mock.patch("megatron.training.initialize.bias_dropout_add_fused_train")
+@mock.patch("megatron.training.initialize.get_args")
+def test_dropout_warmup_and_geglu_cat(mock_get_args, mock_dropout, mock_rand, _cache, mock_geglu):
+    """Gated GEGLU warmup uses torch.cat (RNG-free) to get 2x-ffn tensors, and
+    the fused dropout warmup runs even when bias_dropout_fusion is False so
+    RNG consumption stays identical to the historical behavior."""
+    mock_get_args.return_value = _args(
+        swiglu=False, gated_linear_unit=True, bias_gelu_fusion=True, bias_dropout_fusion=False
+    )
+    _warmup_jit_function()
+    assert mock_geglu.call_count == 10
+    # bias_geglu receives 2x-ffn tensors built via torch.cat
+    bias_arg, input_arg = mock_geglu.call_args[0]
+    assert bias_arg.shape[-1] == 128  # 2 * (ffn_hidden_size // tp) = 2 * 64
+    assert input_arg.shape[-1] == 128
+    assert mock_dropout.call_count == 10  # unconditional for RNG parity
+    assert mock_rand.call_count == 5