fix: forward use_te_activation_func flag in non-MoE GPT layer spec

[saakshigupta2002]

What does this PR do ?

Fixes the --use-te-activation-func CLI flag being silently ignored for non-MoE GPT models by forwarding the parameter through the _get_transformer_layer_spec() code path.

Fixes: #2770

Problem

The --use-te-activation-func flag is correctly parsed from the command line and stored in TransformerConfig.use_te_activation_func, but it is never forwarded to get_gpt_layer_with_transformer_engine_spec() when building layer specs for non-MoE GPT models.

In gpt_builders.py, the function _get_transformer_layer_spec() calls get_gpt_layer_with_transformer_engine_spec() without passing use_te_activation_func, causing it to silently default to False. This means TransformerEngine activation functions are never used in non-MoE GPT models regardless of the CLI flag.

Root Cause

The call chain is:

1. --use-te-activation-func is parsed into args and transferred to TransformerConfig
2. core_transformer_config_from_args(args) creates a TransformerConfig with use_te_activation_func=True
3. _get_transformer_layer_spec(use_te, config) calls get_gpt_layer_with_transformer_engine_spec() without use_te_activation_func
4. get_gpt_layer_with_transformer_engine_spec(..., use_te_activation_func=False) defaults to False
5. get_mlp_module_spec_for_backend(..., use_te_activation_func=False) selects PyTorch activation functions instead of TE's fused implementations

Notably, the MoE code path (get_gpt_decoder_layer_specs in gpt_layer_specs.py) and the experimental attention variant code path (experimental_attention_variant_module_specs.py) both correctly forward use_te_activation_func=config.use_te_activation_func. Only the non-MoE path in gpt_builders.py has this omission.

Fix

Added use_te_activation_func=config.use_te_activation_func to the get_gpt_layer_with_transformer_engine_spec() call in _get_transformer_layer_spec(), consistent with how other code paths already forward this parameter.

Changed Files

• gpt_builders.py: Added one line to forward use_te_activation_func from the config to the layer spec builder function.
• tests/unit_tests/models/test_gpt_model.py: Added regression test test_get_transformer_layer_spec_forwards_use_te_activation_func that uses mock patching to verify the parameter is correctly forwarded from config to the downstream spec function.

Diff

 def _get_transformer_layer_spec(use_te, config):
     args = get_args()
     if use_te:
         return get_gpt_layer_with_transformer_engine_spec(
             args.num_experts,
             args.moe_grouped_gemm,
             args.qk_layernorm,
             args.multi_latent_attention,
             args.experimental_attention_variant,
             moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
             qk_l2_norm=args.qk_l2_norm,
             use_kitchen=config.use_kitchen,
+            use_te_activation_func=config.use_te_activation_func,
             use_kitchen_attention=config.use_kitchen_attention,
             kitchen_attention_backend=config.kitchen_attention_backend,
         )

How to Reproduce the Original Bug

As described in #2770:

1. Run a non-MoE GPT training job with --use-te-activation-func and --transformer-impl transformer_engine
2. Run the same job without --use-te-activation-func
3. Observe identical behavior — the flag has no effect

After this fix, setting --use-te-activation-func correctly enables TransformerEngine activation functions (e.g., TE's fused GELU/SiLU) in the MLP layer spec for non-MoE GPT models.

Testing

Unit Test Added

A regression test test_get_transformer_layer_spec_forwards_use_te_activation_func was added to tests/unit_tests/models/test_gpt_model.py. The test:

• Mocks get_args() and get_gpt_layer_with_transformer_engine_spec() to isolate the forwarding behavior
• Creates a config with use_te_activation_func=True
• Calls _get_transformer_layer_spec(use_te=True, config=config)
• Asserts that get_gpt_layer_with_transformer_engine_spec was called with use_te_activation_func=True

This test does not require CUDA and directly validates the fix for issue #2770.

Existing Coverage

The existing test_gpt_with_te_activation_func in the same file validates the downstream behavior of get_gpt_layer_with_transformer_engine_spec(use_te_activation_func=True) end-to-end with actual model construction and forward pass.

Verification

• The fix adds exactly one line forwarding an existing, already-validated parameter
• The config object passed to _get_transformer_layer_spec() already contains use_te_activation_func (set from CLI args via core_transformer_config_from_args)
• The target function get_gpt_layer_with_transformer_engine_spec() already accepts use_te_activation_func as a keyword argument (gpt_layer_specs.py:181)
• Other code paths (MoE at gpt_layer_specs.py:536,548 and experimental attention at experimental_attention_variant_module_specs.py:398,448) forward this parameter in the same way
• No new imports, no new functions, no architectural changes
• Black (24.4.2) and isort (5.13.2) formatting verified on changed files

Contribution process

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• 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.

[gautham-kollu]

/ok to test 509ef8d

[yaox12]

/ok to test f3b227e

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22453236145