From 845a0e3151010c7468faad92548f3591c9717e80 Mon Sep 17 00:00:00 2001 From: yanliang Date: Tue, 6 Jan 2026 14:39:37 +0800 Subject: [PATCH 1/6] Skip JIT warmup when fusion is disabled via arguments --- megatron/training/initialize.py | 96 +++++++++++++++++---------------- 1 file changed, 49 insertions(+), 47 deletions(-) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index e88222fe7fe..76363089e04 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -472,55 +472,57 @@ def _warmup_jit_function(): dtype = torch.float32 # Warmup fused bias+gelu - bias = torch.rand( - args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" - ) - input = torch.rand( - ( - args.seq_length // args.context_parallel_size, - args.micro_batch_size, - args.ffn_hidden_size // args.tensor_model_parallel_size, - ), - 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: - output = bias_swiglu(input, bias) - else: - output = bias_gelu(bias, input) - del bias, input, output + if (args.swiglu and args.bias_swiglu_fusion) or (not args.swiglu and args.bias_gelu_fusion): + bias = torch.rand( + args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" + ) + input = torch.rand( + ( + args.seq_length // args.context_parallel_size, + args.micro_batch_size, + args.ffn_hidden_size // args.tensor_model_parallel_size, + ), + 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: + output = bias_swiglu(input, bias) + else: + output = bias_gelu(bias, input) + del bias, input, output # Warmup fused bias+dropout+add - if args.sequence_parallel: - seq_length = args.seq_length // mpu.get_tensor_model_parallel_world_size() - else: - seq_length = args.seq_length - input = torch.rand( - (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), - dtype=dtype, - device="cuda", - ) - residual = torch.rand( - (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), - dtype=dtype, - device="cuda", - ) - bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual) - dropout_rate = 0.1 - # Warmup JIT fusions with the input grad_enable state of both forward - # prop and recomputation - for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): - input.requires_grad = input_grad - bias.requires_grad = bias_grad - residual.requires_grad = residual_grad - for _ in range(5): - output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate) - del bias, input, residual, output + if args.bias_dropout_fusion: + if args.sequence_parallel: + seq_length = args.seq_length // mpu.get_tensor_model_parallel_world_size() + else: + seq_length = args.seq_length + input = torch.rand( + (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), + dtype=dtype, + device="cuda", + ) + residual = torch.rand( + (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), + dtype=dtype, + device="cuda", + ) + bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual) + dropout_rate = 0.1 + # Warmup JIT fusions with the input grad_enable state of both forward + # prop and recomputation + for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): + input.requires_grad = input_grad + bias.requires_grad = bias_grad + residual.requires_grad = residual_grad + for _ in range(5): + output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate) + del bias, input, residual, output torch.cuda.empty_cache() From 7396dc414b5e30f7d4ae06687e10c99e86a37b92 Mon Sep 17 00:00:00 2001 From: yanliang Date: Thu, 5 Feb 2026 11:23:39 +0800 Subject: [PATCH 2/6] Fix: Add use_te_activation_func and geglu checks for JIT warmup --- megatron/training/initialize.py | 100 ++++++++++++++++++++++++++++++-- 1 file changed, 94 insertions(+), 6 deletions(-) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 76363089e04..90acbf73cb2 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -13,6 +13,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 @@ -471,8 +472,75 @@ def _warmup_jit_function(): else: dtype = torch.float32 - # Warmup fused bias+gelu - if (args.swiglu and args.bias_swiglu_fusion) or (not args.swiglu and args.bias_gelu_fusion): + # 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) + + # Determine which activation fusion to warmup based on args + # Reference: megatron/core/transformer/mlp.py and megatron/core/transformer/moe/shared_experts.py + # + # In MLP forward (when bias_activation_fusion=True and use_te_activation_func=False): + # - activation_func=F.gelu + gated_linear_unit=True -> bias_geglu_impl + # - activation_func=F.gelu + gated_linear_unit=False -> bias_gelu_impl + # - activation_func=F.silu + gated_linear_unit=True -> bias_swiglu_impl + # + # Args mapping: + # - args.swiglu=True -> gated_linear_unit=True, activation_func=F.silu + # - args.quick_geglu=True -> gated_linear_unit=True, activation_func=quick_gelu (no fusion warmup needed) + # - default (neither set) -> gated_linear_unit=False, activation_func=F.gelu + + # gated_linear_unit can be set via YAML config or other means + 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: non-gated gelu activation (F.gelu without GLU) + 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: gated gelu activation (F.gelu + GLU) + # This is triggered when gated_linear_unit=True with gelu activation + 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 + ) + + # Warmup fused bias+swiglu + if warmup_bias_swiglu: + bias = torch.rand( + args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" + ) + input = torch.rand( + ( + args.seq_length // args.context_parallel_size, + args.micro_batch_size, + args.ffn_hidden_size // args.tensor_model_parallel_size, + ), + 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): + output = bias_swiglu(input, bias) + del bias, input, output + + # Warmup fused bias+gelu (non-gated) + if warmup_bias_gelu: bias = torch.rand( args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" ) @@ -490,10 +558,30 @@ def _warmup_jit_function(): 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: - output = bias_swiglu(input, bias) - else: - output = bias_gelu(bias, input) + output = bias_gelu(bias, input) + del bias, input, output + + # Warmup fused bias+geglu (gated gelu) + if warmup_bias_geglu: + # For geglu, input size is 2x ffn_hidden_size (will be split into two halves) + bias = torch.rand( + (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2, dtype=dtype, device="cuda" + ) + input = torch.rand( + ( + args.seq_length // args.context_parallel_size, + args.micro_batch_size, + (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2, + ), + 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): + output = bias_geglu(bias, input) del bias, input, output # Warmup fused bias+dropout+add From 0fd720f0283a2f9330861b3057b2a502d804111b Mon Sep 17 00:00:00 2001 From: yanliang Date: Tue, 10 Feb 2026 09:45:06 +0800 Subject: [PATCH 3/6] Remove some comments --- megatron/training/initialize.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index b0739d808d8..50e337b4b5d 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -480,21 +480,6 @@ def _warmup_jit_function(): # 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) - - # Determine which activation fusion to warmup based on args - # Reference: megatron/core/transformer/mlp.py and megatron/core/transformer/moe/shared_experts.py - # - # In MLP forward (when bias_activation_fusion=True and use_te_activation_func=False): - # - activation_func=F.gelu + gated_linear_unit=True -> bias_geglu_impl - # - activation_func=F.gelu + gated_linear_unit=False -> bias_gelu_impl - # - activation_func=F.silu + gated_linear_unit=True -> bias_swiglu_impl - # - # Args mapping: - # - args.swiglu=True -> gated_linear_unit=True, activation_func=F.silu - # - args.quick_geglu=True -> gated_linear_unit=True, activation_func=quick_gelu (no fusion warmup needed) - # - default (neither set) -> gated_linear_unit=False, activation_func=F.gelu - - # gated_linear_unit can be set via YAML config or other means gated_linear_unit = getattr(args, 'gated_linear_unit', False) # Warmup bias_swiglu: swiglu activation (F.silu + GLU) @@ -504,7 +489,6 @@ def _warmup_jit_function(): and args.bias_swiglu_fusion ) - # Warmup bias_gelu: non-gated gelu activation (F.gelu without GLU) warmup_bias_gelu = ( not use_te_activation_func and not args.swiglu @@ -513,8 +497,6 @@ def _warmup_jit_function(): and args.bias_gelu_fusion ) - # Warmup bias_geglu: gated gelu activation (F.gelu + GLU) - # This is triggered when gated_linear_unit=True with gelu activation warmup_bias_geglu = ( not use_te_activation_func and not args.swiglu @@ -523,7 +505,6 @@ def _warmup_jit_function(): and args.bias_gelu_fusion ) - # Warmup fused bias+swiglu if warmup_bias_swiglu: bias = torch.rand( args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" @@ -545,7 +526,6 @@ def _warmup_jit_function(): output = bias_swiglu(input, bias) del bias, input, output - # Warmup fused bias+gelu (non-gated) if warmup_bias_gelu: bias = torch.rand( args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" From 47a6aae71ac8566bb6264b7d649e91f69b0a56ba Mon Sep 17 00:00:00 2001 From: liyanliang Date: Wed, 8 Apr 2026 09:42:16 +0800 Subject: [PATCH 4/6] test: add CPU-only unit tests for JIT warmup fusion gating Cover skip paths when bias_* fusion flags are off, quick_geglu, and use_te_activation_func, using mocks so tests run without CUDA. --- .../unit_tests/test_initialize_jit_warmup.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 tests/unit_tests/test_initialize_jit_warmup.py diff --git a/tests/unit_tests/test_initialize_jit_warmup.py b/tests/unit_tests/test_initialize_jit_warmup.py new file mode 100644 index 00000000000..6b6b461e1e0 --- /dev/null +++ b/tests/unit_tests/test_initialize_jit_warmup.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +"""Tests that JIT warmup skips fused kernels when disabled via arguments. + +Uses mocks only (no CUDA required). Patches use +``megatron.training.initialize.`` because ``_warmup_jit_function`` closes +over that module's globals (imported ``get_args``, ``torch``, fusion ops). +""" + +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=False, + 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.get_args") +class TestJitWarmupSkippedWhenFusionDisabled: + @mock.patch("megatron.training.initialize.bias_swiglu") + def test_skips_bias_swiglu(self, mock_swiglu, mock_get_args, _empty_cache): + mock_get_args.return_value = _args(swiglu=True, bias_swiglu_fusion=False) + _warmup_jit_function() + mock_swiglu.assert_not_called() + + @mock.patch("megatron.training.initialize.bias_gelu") + def test_skips_bias_gelu(self, mock_gelu, mock_get_args, _empty_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() + + @mock.patch("megatron.training.initialize.bias_geglu") + def test_skips_bias_geglu(self, mock_geglu, mock_get_args, _empty_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() + + @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, _empty_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() + + @mock.patch("megatron.training.initialize.bias_dropout_add_fused_train") + def test_skips_bias_dropout_fusion(self, mock_dropout, mock_get_args, _empty_cache): + # With swiglu=True, gelu/geglu warmups are off; disable swiglu fusion so no CUDA path runs. + mock_get_args.return_value = _args( + swiglu=True, bias_swiglu_fusion=False, bias_dropout_fusion=False + ) + _warmup_jit_function() + mock_dropout.assert_not_called() + + @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.bias_swiglu") + def test_te_activation_skips_activation_warmups( + self, + mock_swiglu, + mock_gelu, + mock_geglu, + mock_dropout, + mock_get_args, + _empty_cache, + ): + # bias_dropout_fusion is independent of use_te_activation_func; keep False + # so this test stays CPU-only without mocking the dropout path + mpu. + 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, + bias_dropout_fusion=False, + ) + _warmup_jit_function() + mock_swiglu.assert_not_called() + mock_gelu.assert_not_called() + mock_geglu.assert_not_called() + mock_dropout.assert_not_called() + + +@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() + mock_dropout.assert_not_called() + assert mock_torch_rand.call_count == 2 From c61e8a5c2c8a3f48fd6651e0e74714839f835fdc Mon Sep 17 00:00:00 2001 From: Antoni-Joan Solergibert Date: Mon, 27 Apr 2026 19:29:12 +0200 Subject: [PATCH 5/6] style: apply black formatting to test_initialize_jit_warmup.py Repo black config uses --skip-magic-trailing-comma with line_length=100, which collapses argument lists that fit on one line. Fixes CI linting. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../unit_tests/test_initialize_jit_warmup.py | 21 +++---------------- 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/tests/unit_tests/test_initialize_jit_warmup.py b/tests/unit_tests/test_initialize_jit_warmup.py index 6b6b461e1e0..826f59950ee 100644 --- a/tests/unit_tests/test_initialize_jit_warmup.py +++ b/tests/unit_tests/test_initialize_jit_warmup.py @@ -82,10 +82,7 @@ def test_skips_gelu_and_geglu_when_quick_geglu( ): # 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, + swiglu=False, quick_geglu=True, bias_gelu_fusion=True, gated_linear_unit=True ) _warmup_jit_function() mock_gelu.assert_not_called() @@ -105,13 +102,7 @@ def test_skips_bias_dropout_fusion(self, mock_dropout, mock_get_args, _empty_cac @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_dropout, - mock_get_args, - _empty_cache, + self, mock_swiglu, mock_gelu, mock_geglu, mock_dropout, mock_get_args, _empty_cache ): # bias_dropout_fusion is independent of use_te_activation_func; keep False # so this test stays CPU-only without mocking the dropout path + mpu. @@ -138,13 +129,7 @@ def test_te_activation_skips_activation_warmups( @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, 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) From ed28045bbcbd689af346630b778b55a3e1f2b420 Mon Sep 17 00:00:00 2001 From: kisseternity <15059072+kisseternity@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:26:00 +0800 Subject: [PATCH 6/6] Fix: skip only RNG-free fusion warmups; keep RNG draws unconditional The JIT warmup runs after _set_random_seed (pretrain() -> set_jit_fusion_options()), and historically always executed the same torch.rand warmup draws and fused bias+dropout+add warmup for every config. Those consumptions are baked into determinism-sensitive output, e.g. the golden values of the functional tests, so gating them on arguments shifts the default CUDA RNG stream and breaks golden-value comparisons. This was seen on the two --disable-bias-linear GPT jobs in the merge queue (validate_args forces bias_gelu_fusion=False there), and restoring the RNG state instead would have shifted every baseline. Make warmup RNG consumption identical to before for every config: - keep all torch.rand draws unconditional with historical shapes/order; - gate only the RNG-free bias_swiglu/bias_gelu/bias_geglu kernel calls on the fusion arguments; - for GEGLU configs warm up bias_geglu (the historical code warmed bias_gelu), doubling the drawn tensors via torch.cat so no extra RNG is consumed; - run the fused bias+dropout+add warmup unconditionally. Update the unit tests to pin this determinism contract. Signed-off-by: kisseternity <15059072+kisseternity@users.noreply.github.com> --- megatron/training/initialize.py | 150 ++++++++---------- .../unit_tests/test_initialize_jit_warmup.py | 80 +++++++--- 2 files changed, 125 insertions(+), 105 deletions(-) diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 52b1ae64712..4f79cc78e40 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -516,99 +516,85 @@ def _warmup_jit_function(tp_size=None): 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" + ) + input = torch.rand( + ( + args.seq_length // args.context_parallel_size, + args.micro_batch_size, + args.ffn_hidden_size // args.tensor_model_parallel_size, + ), + dtype=dtype, + device="cuda", + ) + + # Warmup JIT fusions with the input grad_enable state of both forward + # prop and recomputation. The three branches are mutually exclusive. if warmup_bias_swiglu: - bias = torch.rand( - args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" - ) - input = torch.rand( - ( - args.seq_length // args.context_parallel_size, - args.micro_batch_size, - args.ffn_hidden_size // args.tensor_model_parallel_size, - ), - 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): output = bias_swiglu(input, bias) - del bias, input, output - - if warmup_bias_gelu: - bias = torch.rand( - args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda" - ) - input = torch.rand( - ( - args.seq_length // args.context_parallel_size, - args.micro_batch_size, - args.ffn_hidden_size // args.tensor_model_parallel_size, - ), - dtype=dtype, - device="cuda", - ) - # Warmup JIT fusions with the input grad_enable state of both forward - # prop and recomputation + 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+geglu (gated gelu) - if warmup_bias_geglu: - # For geglu, input size is 2x ffn_hidden_size (will be split into two halves) - bias = torch.rand( - (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2, dtype=dtype, device="cuda" - ) - input = torch.rand( - ( - args.seq_length // args.context_parallel_size, - args.micro_batch_size, - (args.ffn_hidden_size // args.tensor_model_parallel_size) * 2, - ), - dtype=dtype, - device="cuda", - ) - # Warmup JIT fusions with the input grad_enable state of both forward - # prop and recomputation + 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]): - bias.requires_grad, input.requires_grad = bias_grad, input_grad - for _ in range(5): - output = bias_geglu(bias, input) - del bias, input, output - - # Warmup fused bias+dropout+add - if args.bias_dropout_fusion: - 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()) - else: - seq_length = args.seq_length - input = torch.rand( - (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), - dtype=dtype, - device="cuda", - ) - residual = torch.rand( - (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), - dtype=dtype, - device="cuda", - ) - bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual) - dropout_rate = 0.1 - # Warmup JIT fusions with the input grad_enable state of both forward - # prop and recomputation - for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): - input.requires_grad = input_grad - bias.requires_grad = bias_grad - residual.requires_grad = residual_grad + geglu_bias.requires_grad, geglu_input.requires_grad = bias_grad, input_grad for _ in range(5): - output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate) - del bias, input, residual, output + 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()) + else: + seq_length = args.seq_length + input = torch.rand( + (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), + dtype=dtype, + device="cuda", + ) + residual = torch.rand( + (seq_length // args.context_parallel_size, args.micro_batch_size, args.hidden_size), + dtype=dtype, + device="cuda", + ) + bias = torch.rand((args.hidden_size), dtype=dtype, device="cuda").expand_as(residual) + dropout_rate = 0.1 + # Warmup JIT fusions with the input grad_enable state of both forward + # prop and recomputation + for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]): + input.requires_grad = input_grad + bias.requires_grad = bias_grad + residual.requires_grad = residual_grad + for _ in range(5): + output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate) + del bias, input, residual, output + torch.cuda.empty_cache() diff --git a/tests/unit_tests/test_initialize_jit_warmup.py b/tests/unit_tests/test_initialize_jit_warmup.py index 826f59950ee..345a00480b6 100644 --- a/tests/unit_tests/test_initialize_jit_warmup.py +++ b/tests/unit_tests/test_initialize_jit_warmup.py @@ -1,10 +1,17 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -"""Tests that JIT warmup skips fused kernels when disabled via arguments. +"""Tests that JIT warmup skips only RNG-free fused kernels when disabled. Uses mocks only (no CUDA required). Patches use ``megatron.training.initialize.`` 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 @@ -25,7 +32,7 @@ def _args(**overrides): swiglu=False, bias_swiglu_fusion=True, bias_gelu_fusion=True, - bias_dropout_fusion=False, + bias_dropout_fusion=True, sequence_parallel=False, ffn_hidden_size=64, tensor_model_parallel_size=1, @@ -51,34 +58,44 @@ def _fake_rand(*size, **kwargs): @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, _empty_cache): + 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, _empty_cache): + 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, _empty_cache): + 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, _empty_cache + 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( @@ -87,38 +104,31 @@ def test_skips_gelu_and_geglu_when_quick_geglu( _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_dropout_add_fused_train") - def test_skips_bias_dropout_fusion(self, mock_dropout, mock_get_args, _empty_cache): - # With swiglu=True, gelu/geglu warmups are off; disable swiglu fusion so no CUDA path runs. - mock_get_args.return_value = _args( - swiglu=True, bias_swiglu_fusion=False, bias_dropout_fusion=False - ) - _warmup_jit_function() - mock_dropout.assert_not_called() - - @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.bias_swiglu") def test_te_activation_skips_activation_warmups( - self, mock_swiglu, mock_gelu, mock_geglu, mock_dropout, mock_get_args, _empty_cache + self, mock_swiglu, mock_gelu, mock_geglu, mock_get_args, mock_rand, mock_dropout, _cache ): - # bias_dropout_fusion is independent of use_te_activation_func; keep False - # so this test stays CPU-only without mocking the dropout path + mpu. + # 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, - bias_dropout_fusion=False, ) _warmup_jit_function() mock_swiglu.assert_not_called() mock_gelu.assert_not_called() mock_geglu.assert_not_called() - mock_dropout.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") @@ -137,5 +147,29 @@ def test_calls_bias_swiglu_when_fusion_enabled( assert mock_swiglu.call_count == 10 # zip([True,True],[False,True]) × range(5) mock_gelu.assert_not_called() mock_geglu.assert_not_called() - mock_dropout.assert_not_called() - assert mock_torch_rand.call_count == 2 + 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