Skip to content

Conditional JIT warmup #6677

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

Closed
maanug-nv wants to merge 6 commits into from
Closed

Conditional JIT warmup #6677

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7c8e4ef
Skip JIT warmup when fusion is disabled via arguments
Jan 6, 2026
8e5c2be
Fix: Add use_te_activation_func and geglu checks for JIT warmup
Feb 5, 2026
66c893d
Remove some comments
Feb 10, 2026
61986f8
test: add CPU-only unit tests for JIT warmup fusion gating
Apr 8, 2026
e9b6a54
style: apply black formatting to test_initialize_jit_warmup.py
asolergi-nv Apr 27, 2026
c38f6cb
update relevant ghci mr golden values
maanug-nv Aug 19, 2026
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
160 changes: 115 additions & 45 deletions megatron/training/initialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -526,57 +527,126 @@ def _warmup_jit_function(tp_size=None):
else:
dtype = torch.float32

# Warmup fused bias+gelu
bias = torch.rand(
args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device="cuda"
# 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
)
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_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
)
# 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:

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)
else:
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
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
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
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
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
for _ in range(5):
output = bias_dropout_add_fused_train([input, bias], residual, dropout_rate)
del bias, input, residual, output
torch.cuda.empty_cache()


Expand Down
Loading
Loading