From 78df0fb0e2b14f0f95c1c2d544af770ad339543a Mon Sep 17 00:00:00 2001 From: factnn <166481866+factnn@users.noreply.github.com> Date: Sat, 23 May 2026 21:26:12 +0800 Subject: [PATCH 1/2] Fix build_mlp() missing name kwarg in custom-pg test transformer_layer.py now passes name= to submodules.mlp() (added in #4358), but the test helper build_mlp() did not accept it, causing TypeError in all test_transformer_block_custom_pgs parametrizations. Add name=None parameter and suppress it with the other unused args. Fixes #4934 --- .../transformer/test_transformer_block_custom_pgs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py b/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py index b9a080b41bc..2d25ff1ce70 100644 --- a/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py +++ b/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py @@ -167,9 +167,9 @@ def _gpt_te_layer_spec_with_hetro_pgs( ): def build_mlp( - config: TransformerConfig, pg_collection: ProcessGroupCollection, is_mtp_layer: bool + config: TransformerConfig, pg_collection: ProcessGroupCollection, is_mtp_layer: bool, name=None ): - del pg_collection, is_mtp_layer + del pg_collection, is_mtp_layer, name return MLP( config, submodules=MLPSubmodules( From 6dd6c5f09a63f43a8611946422a9f4a9f5c627f6 Mon Sep 17 00:00:00 2001 From: factnn <166481866+factnn@users.noreply.github.com> Date: Wed, 27 May 2026 14:21:37 +0800 Subject: [PATCH 2/2] Revert flaky_in_dev marker for custom-pg test Now that the root cause (missing name kwarg) is fixed, remove the emporary flaky_in_dev quarantine added in 741ae4e9. --- .../unit_tests/transformer/test_transformer_block_custom_pgs.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py b/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py index 2d25ff1ce70..473b6155260 100644 --- a/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py +++ b/tests/unit_tests/transformer/test_transformer_block_custom_pgs.py @@ -365,7 +365,6 @@ def test_params_and_grads_match_transformer_block(self, world_size, tp_size, cp_ default_param.main_grad is not None and custom_param.main_grad is not None ), f"Gradient is None for parameter '{param_name}' at index {i}" - @pytest.mark.flaky_in_dev @pytest.mark.skipif( version.parse(torch.__version__) < version.parse('2.3.0'), reason="Device mesh feature requires PyTorch 2.3 or later",