Generalized Tensor Parallelism (GTP)

[fanshiqing]

Description

Core-idea: add Generalized Tensor Parallelism (GTP), which is a flexible fine-grained sharding/just-in time materialization of both activations and parameters with efficient computation-communication overlap.

Mission: improve LLM pretraining efficiency through generalized tensor parallelism, enabling high performance, memory efficiency, ease of use, and strong scalability.

GTP introduction

GTP_introduction.md

How Mcore interacts with TE

• TE ships self-contained: Protocol + dispatcher no-op.
• GTP lives 100% in Mcore as one protocol implementer.

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:

• distributed_weight.py (NEW): the DistributedWeight protocol + dispatchers — TE's GTP-agnostic extension point.
• module/linear.py: rewired fwd/bwd to is_distributed_weight + dispatchers;
• module/layernorm_linear.py: same protocol rewiring as linear.
• module/grouped_linear.py: same, for the grouped N-weight case.
• ops/fused/grouped_mlp.py: protocol rewiring; FC2 gather split; FC1 NVFP4 dgrad materialized before the precision dispatch.
• distributed.py: added general external_coalescing flag; reduce_scatter output=; async NVFP4 handle.
• tests/pytorch/test_distributed_weight.py (NEW): 7 GPU-free protocol/dispatcher tests.

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 introduces Generalized Tensor Parallelism (GTP) support into TransformerEngine by adding a DistributedWeight protocol (distributed_weight.py) and wiring its dispatchers into Linear, LayerNormLinear, GroupedLinear, ops.GroupedLinear, and the fused GroupedMLP CuTe kernel path. TE ships the protocol as a no-op extension point; the actual GTP implementation lives in Megatron-LM.

• New protocol file (distributed_weight.py): defines DistributedWeight as a runtime_checkable Protocol, enforces torch.Tensor subclass at runtime, and ships list-normalizing dispatchers; covered by 7 GPU-free unit tests.
• Module rewiring (linear.py, layernorm_linear.py, grouped_linear.py): forward materializes and backward re-gathers sharded weights instead of saving gathered workspaces; finalize_weight_grads issues the async reduce-scatter.
• Ops rewiring (ops/basic/grouped_linear.py, ops/fused/grouped_mlp.py): same pattern extended to the fusible ops layer; FC1 dgrad materialization in fused backward is now correctly placed before the NVFP4 dispatch; distributed.py gains output=/external_coalescing parameters and fixes a null-handle crash in _NVFP4AllGatherAsyncHandle.wait().

Confidence Score: 3/5

The fused CuTeGEMM path embeds GTP-specific attribute access outside the protocol, and the wgrad finalization relies on an in-place side-effect contract that is not enforced or clearly documented.

The fused CuTeGEMM path hard-codes GTP-internal attributes outside the protocol contract (is_routed_expert, weight_list), and the finalize_weight_grads in-place side-effect is load-bearing but undocumented — the test FakeDistributedWeight does not exercise the production assumption, leaving the in-place contract untested.

transformer_engine/pytorch/ops/fused/grouped_mlp.py (GTP-specific assertions at lines 921-922) and transformer_engine/pytorch/ops/basic/grouped_linear.py (finalize_weight_grads return-value contract in both backward methods).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed_weight.py	New file defining the GTP-agnostic DistributedWeight protocol and dispatchers; clean design with runtime type enforcement and list normalization.
transformer_engine/pytorch/distributed.py	Adds output= to reduce_scatter_along_first_dim, output_tensor/external_coalescing to NVFP4 and MXFP8 gather paths, and fixes the null async_handle crash in _NVFP4AllGatherAsyncHandle.wait().
transformer_engine/pytorch/module/linear.py	Wires DistributedWeight dispatchers into forward/backward; fuse_wgrad_accumulation dummy-wgrad branch can still fire for GTP params but is harmless when the in-place main_grad contract is followed.
transformer_engine/pytorch/module/layernorm_linear.py	Adds DistributedWeight dispatchers for forward materialization, backward re-gather, wgrad finalization; symmetric with linear.py.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds dispatch helpers; finalize_weight_grads return value discarded in both backward paths, relying on undocumented in-place side-effect contract.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Adds distributed-weight support to fused CuTeGEMM path; FC1 dgrad materialization correctly placed before NVFP4 branch; lines 921-922 embed GTP-specific attributes outside the DistributedWeight protocol.

_{Reviews (26): Last reviewed commit: "Unify distributed-weight wgrad finalize ..." | Re-trigger Greptile}

[fanshiqing]

/te-ci L1 pytorch

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[timmoon10]

/te-ci pytorch L1

[timmoon10]

LGTM, pending CI

[timmoon10]

/te-ci pytorch L1

[fanshiqing]

@timmoon10 Hi Tim, the online building (core/Jax/Pytorch) failed but from the log it seems it's nccl_ep related, I have also saw the same build failure from other opening MRs, is that one know issue and will it block the merge process?

nccl_ep.cc:23:10: fatal error: nccl_device.h: No such file or directory
     23 | #include <nccl_device.h>
        |          ^~~~~~~~~~~~~~~
  compilation terminated.
  make: *** [Makefile:156: /__w/TransformerEngine/TransformerEngine/3rdparty/nccl/build/obj/nccl_ep/nccl_ep.o] Error 1
  make: *** Waiting for unfinished jobs....
  make: Leaving directory '/__w/TransformerEngine/TransformerEngine/3rdparty/nccl/contrib/nccl_ep'
  [NCCL EP] Building libnccl_ep.a (gencode='-gencode=arch=compute_90,code=sm_90')
  Traceback (most recent call last):
    File "/usr/lib/python3/dist-packages/pip/_vendor/pep517/in_process/_in_process.py", line 363, in <module>
      main()
    File "/usr/lib/python3/dist-packages/pip/_vendor/pep517/in_process/_in_process.py", line 345, in main
      json_out['return_val'] = hook(**hook_input['kwargs'])
    File "/usr/lib/python3/dist-packages/pip/_vendor/pep517/in_process/_in_process.py", line 164, in prepare_metadata_for_build_wheel
      return hook(metadata_directory, config_settings)
    File "/usr/lib/python3/dist-packages/setuptools/build_meta.py", line 174, in prepare_metadata_for_build_wheel
      self.run_setup()
    File "/usr/lib/python3/dist-packages/setuptools/build_meta.py", line 267, in run_setup
      super(_BuildMetaLegacyBackend,
    File "/usr/lib/python3/dist-packages/setuptools/build_meta.py", line 158, in run_setup
      exec(compile(code, __file__, 'exec'), locals())
    File "setup.py", line 365, in <module>
      ext_modules = [setup_common_extension()]
    File "setup.py", line 94, in setup_common_extension
      nccl_home = build_nccl_ep_submodule()
    File "setup.py", line 275, in build_nccl_ep_submodule
      subprocess.check_call(
    File "/usr/lib/python3.10/subprocess.py", line 369, in check_call
      raise CalledProcessError(retcode, cmd)
  subprocess.CalledProcessError: Command '['make', '-j', '4', '-C', 'contrib/nccl_ep', 'lib']' returned non-zero exit status 2.
  error: subprocess-exited-with-error
  
  × Preparing metadata (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> See above for output.