Skip to content

Fix build_mlp() missing name kwarg in custom-pg test #4954

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

Draft
factnn wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Fix build_mlp() missing name kwarg in custom-pg test #4954

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
78df0fb
Fix build_mlp() missing name kwarg in custom-pg test
factnn May 23, 2026
6dd6c5f
Revert flaky_in_dev marker for custom-pg test
factnn May 27, 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
5 changes: 2 additions & 3 deletions tests/unit_tests/transformer/test_transformer_block_custom_pgs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down