NVIDIA / NVIDIA/Megatron-LM

🐛 CI failure: build_mlp() got unexpected kwarg 'name' in test_transformer_block_custom_pgs

Open
#4,934 1 comment 0 reactions 1 assignee Claimed by @santhnm2 View on GitHub
bug
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Describe the bug**

CI test `tests/unit_tests/transformer/test_transformer_block_custom_pgs.py::TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block` is failing across all 5 parametrizations with a `TypeError` from a recently-added `name=` kwarg:

```
megatron/core/transformer/transformer_layer.py:420: TypeError
E TypeError: _gpt_te_layer_spec_with_hetro_pgs..build_mlp() got an unexpected keyword argument 'name' when instantiating HeterogenousTransformerLayer
```

Failing nodes:

- `TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block[8-1-1-8]`
- `TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block[8-8-1-1]`
- `TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block[8-2-1-4]`
- `TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block[8-4-1-2]`
- `TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block[8-2-2-2]`

Tag @NVIDIA/mcore-oncall to get oncall's attention to this issue.

**Root cause (likely)**

Commit `4c63602607` ([#4358](https://github.com/NVIDIA/Megatron-LM/pull/4358), 2026-05-22, @xrennvidia) added a `name=(name + ".mlp") if name is not None else None,` kwarg to `submodules.mlp(...)` at `megatron/core/transformer/transformer_layer.py:424`. The test helper `_gpt_te_layer_spec_with_hetro_pgs..build_mlp` defined at `tests/unit_tests/transformer/test_transformer_block_custom_pgs.py:165` does not accept `name`, so any invocation through this spec now crashes.

**Failing run**

| Field | Value |
|-------|-------|
| PR | [#4931: test: enable NVTE_CUTEDSL_FUSED_GROUPED_MLP via pytest fixture](https://github.com/NVIDIA/Megatron-LM/pull/4931) (surfaced here; the PR itself does not touch `transformer_layer.py` or the custom-pg test) |
| Run | [26290097610](https://github.com/NVIDIA/Megatron-LM/actions/runs/26290097610) |
| Job | [`tests/unit_tests/transformer/**/*.py - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/26290097610/job/77389398611) |

**Error (verbatim, abridged)**

```
megatron/core/transformer/transformer_layer.py:420: TypeError
E TypeError: _gpt_te_layer_spec_with_hetro_pgs..build_mlp() got an unexpected keyword argument 'name' when instantiating HeterogenousTransformerLayer
```

**Steps/Code to reproduce bug**

Re-run the failing CI job linked above, or locally inside the dev container:

```bash
pytest tests/unit_tests/transformer/test_transformer_block_custom_pgs.py::TestTransformerBlockWithProcessGroups::test_fwd_bwd_pass_non_uniform_transformer_block
```

**Additional context**

Triaged automatically via `/create-issue`. Fix is either to extend `build_mlp` in the test helper to accept and ignore (or forward) `name`, or to make the call site at `transformer_layer.py:424` only pass `name` when the inner spec supports it.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.