feat(training): wire generalized tensor parallelism

[yaoyu-33]

Summary

• forward the dense and expert GTP rematerialization sizes derived from the existing weight-shard configuration into standard MCore process-group initialization
• configure GTP before model construction and classify all model/VPP chunks after distributed wrapping
• use the GTP-inclusive data-distribution groups for dataloaders, batch/sample accounting, training/evaluation loss reductions, and packed-sequence statistics
• reject decentralized process groups and collective averaging early when GTP is active, and report the Transformer Engine requirement actionably

This is Stage 1 of #5434. It intentionally keeps GTP implementation ownership in Megatron Core.

Validation

Pinned revisions:

• Bridge base: 064b77054eed990778d7ab4f092083449445b764
• Megatron-LM: 14346b65a2d0790e451919858f7771078105c5f0
• Transformer Engine main: c280e7cb192f57a182a6aab64d3d6bd40d933d92
• runtime: PyTorch 2.13.0a0+8145d630e8.nv26.6, Megatron Core 0.20.0+14346b65a, Transformer Engine 2.19.0.dev0

Targeted results:

• 105 passed across GTP, distributed initialization, setup, process-group wiring, and decentralized-group unit tests
• 2-GPU BF16 Bridge smoke: three finite steps with GTP=2, classified GTP parameters, full data-distribution size 2, replicate-DP size 1, consumed samples 2, 4, 6, and no skipped/NaN iterations
• GTP-off control: finite two-rank DP training with no GTP parameters
• uv run --active --no-sync pre-commit run --all-files

Reproducible Bridge/MLM scalar comparison

Both runs used a 2-layer GPT (hidden size 128, FFN size 256, 4 heads, sequence length 64, vocabulary size 128), TP1/PP1/CP1/GTP2, BF16, mock data, micro/global batch sizes 1/2, seed 1234, dropout 0, untied embeddings/output weights, no linear bias, attention softmax in FP32, weight decay 0.01, clip grad 1.0, and the same three-step cosine schedule (3e-4, 1.65e-4, 3e-5; minimum 3e-5; one warmup step). The optimizer was Adam with betas 0.9/0.999 and epsilon 1e-8, without distributed optimizer.

Bridge command:

uv run --active --no-sync python -m torch.distributed.run --nproc_per_node=2 --master_port=29555 \
  -m pytest -q -s \
  tests/functional_tests/test_groups/training/test_pretrain.py::TestPretrain::test_pretrain_with_generalized_tensor_parallelism

Megatron-LM command:

uv run --active --no-sync python -m torch.distributed.run --nproc_per_node=2 --master_port=29556 \
  3rdparty/Megatron-LM/pretrain_gpt.py \
  --mock-data --tokenizer-type NullTokenizer --vocab-size 128 \
  --num-layers 2 --hidden-size 128 --ffn-hidden-size 256 --num-attention-heads 4 \
  --seq-length 64 --max-position-embeddings 64 \
  --micro-batch-size 1 --global-batch-size 2 --train-iters 3 \
  --eval-interval 100 --eval-iters 0 \
  --lr 3e-4 --min-lr 3e-5 --lr-warmup-iters 1 --lr-decay-iters 3 --lr-decay-style cosine \
  --weight-decay 0.01 --weight-decay-incr-style constant --clip-grad 1.0 \
  --bf16 --disable-bias-linear --untie-embeddings-and-output-weights \
  --no-gradient-accumulation-fusion --attention-dropout 0 --hidden-dropout 0 \
  --attention-softmax-in-fp32 \
  --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 \
  --tensor-parallel-num-weight-shards 2 \
  --seed 1234 --log-interval 1 --dataloader-type single --num-workers 0

GTP is enabled from the same source setting in both frontends: Bridge sets tensor_model_parallel_size=1 and tensor_parallel_num_weight_shards=2 on its model config; MLM uses the corresponding two CLI flags. MCore derives gtp_weight_remat_size = tensor_parallel_num_weight_shards / tensor_model_parallel_size = 2; the derived field is not set directly.

Step	Samples	Bridge loss	MLM loss	Bridge grad	MLM grad	LR (both)
1	2	4.860444	4.860444	2.126	2.126	3.00e-4
2	4	4.827285	4.843864	2.480	2.480	1.65e-4
3	6	4.781815	4.781815	2.413	2.413	3.00e-5

The step-2 absolute loss difference is 0.016579, or 0.342268% relative to MLM. It reproduced across reruns. Matching dropout and FP32-softmax settings aligned step 1 and the displayed grad norms but did not remove step 2. Both frontends reduce token-summed loss and token counts over the GTP-inclusive data-distribution group. Disabling cross-entropy fusion did not remove the Bridge value, and a matched SGD control retained the same step-2 difference, so it is not explained solely by Adam state. A GTP-off run changes data topology and is not a clean isolation control.

These are rounded, reported scalar observations only. No per-rank inputs, logits, tensors, parameter deltas, per-parameter gradients, or optimizer slots were captured, so this does not establish bitwise/tensor parity or rule out transient state drift. The exact source of the step-2 discrepancy remains unresolved. The checked-in functional test therefore asserts the intended GTP topology, nonempty GTP parameter classification, progress, and finite results; it is deliberately not a numerical-parity assertion.

Deferred

• decentralized HyperCommGrid GTP support
• broader MoE/expert-GTP, PP/CP/VPP, checkpoint-topology, recipes/docs, and performance coverage
• upstream full-GPT limitations for bias-enabled linears and tied embedding/output weights; the correctness smoke follows upstream GTP coverage by disabling bias and using untied embeddings

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 008670b

[yaoyu-33]

/ok to test 40b50b4