Bugfix spec check for DP+FSDP configurations

[dinodeep]

Description

Please include a brief summary of the changes, relevant motivation and context.
In the collective GEMM code, there is a small bug where the code misses the general case that (lhs|rhs)_non_cspecs can have elements that are tuples when a dimension in sharded on multiple resource axes (e.g. (data, fsdp) on the batch dimension).

This leads to an incorrect check which results in a sharding in which a tensor is sharded on the same device axis twice.

  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/opt/maxtext/src/maxtext/trainers/pre_train/train.py", line 988, in <module>
  File "/usr/local/lib/python3.12/dist-packages/absl/app.py", line 410, in run
  File "/usr/local/lib/python3.12/dist-packages/absl/app.py", line 355, in _run_main
  File "/opt/maxtext/src/maxtext/trainers/pre_train/train.py", line 984, in main
  File "/opt/maxtext/src/maxtext/trainers/pre_train/train.py", line 974, in train_func
  File "/opt/maxtext/src/maxtext/trainers/pre_train/train.py", line 944, in run
  File "/opt/maxtext/src/maxtext/trainers/pre_train/train.py", line 817, in train_loop
  File "/opt/jax/jax/_src/stages.py", line 593, in compile
  File "/opt/jax/jax/_src/interpreters/pxla.py", line 1260, in compile
  File "/opt/jax/jax/_src/interpreters/pxla.py", line 1757, in from_hlo
  File "/opt/jax/jax/_src/interpreters/pxla.py", line 1537, in _cached_compilation
  File "/opt/jax/jax/_src/compiler.py", line 464, in compile_or_get_cached
  File "/opt/jax/jax/_src/compiler.py", line 729, in _compile_and_write_cache
  File "/opt/jax/jax/_src/profiler.py", line 420, in wrapper
  File "/opt/jax/jax/_src/compiler.py", line 350, in backend_compile_and_load
  File "/opt/jax/jax/_src/custom_partitioning.py", line 160, in _custom_partitioning_partition
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/jax/cpp_extensions/gemm.py", line 1139, in partition
  File "/opt/jax/jax/_src/named_sharding.py", line 515, in check_pspec
  File "/opt/jax/jax/_src/named_sharding.py", line 544, in _check_unique_resources
DuplicateSpecError: A single NamedSharding spec specification can map every mesh axis to at most one positional dimension, but P(('data', 'fsdp'), 'tensor_sequence', 'fsdp') has duplicate entries for `fsdp`

The change in this PR resolves that problem by flattening the tuple so that dimensions that are sharded on multiple device mesh axes are considered correctly.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Fixes sharding spec check so that tuple elements are checked appropriately.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR extends collective GEMM sharding handling and distributed coverage.

• Adds an eight-device DP+FSDP+TP test configuration.
• Uses combined DP/FSDP axes in distributed dense input sharding.
• Flattens tuple-valued non-contracting specs when filtering GEMM output sharding.

Confidence Score: 3/5

The PR is not yet safe to merge because partially overlapping tuple-valued operand specs can still produce an invalid output PartitionSpec.

The revised checks recognize exact tuple equality but do not detect a mesh axis shared by different tuples, so both specs can survive and assign that axis to multiple output dimensions.

Files Needing Attention: transformer_engine/jax/cpp_extensions/gemm.py

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/gemm.py	Updates both collective GEMM partition branches to account for tuple-valued non-contracting sharding specs.
tests/jax/distributed_test_base.py	Adds an eight-device mesh configuration spanning DP, FSDP, and tensor/sequence parallel axes.
tests/jax/test_distributed_dense.py	Combines DP and FSDP resources into the distributed dense input PartitionSpec.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Operand PartitionSpecs] --> B[Extract non-contracting specs]
  B --> C[Flatten opposite operand axes]
  C --> D[Filter overlapping specs]
  D --> E[Concatenate output PartitionSpec]

Loading

_{Reviews (2): Last reviewed commit: "Add in check for if spec is a tuple" | Re-trigger Greptile}