From cc7079a9d7a1d69b95c3bfb76667f0b1140d1c68 Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Fri, 14 Aug 2026 14:35:43 -0700 Subject: [PATCH 1/3] Update [ghstack-poisoned] --- .../integration_test_8gpu_features.yaml | 4 ++-- torchtitan/components/optimizer.py | 17 ++++++++++++----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_features.yaml b/.github/workflows/integration_test_8gpu_features.yaml index 6aa16888d5..59d368dceb 100644 --- a/.github/workflows/integration_test_8gpu_features.yaml +++ b/.github/workflows/integration_test_8gpu_features.yaml @@ -144,8 +144,8 @@ jobs: # whose CPU synchronization is incompatible with CUDA graph capture. python3 scripts/loss_compare.py . . \ --baseline-module=qwen3 --baseline-config=qwen3_moe_debug \ - --baseline-options="--parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 4 --training.disable_cuda_graphs" \ - --test-options="--parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 4 --training.disable_cuda_graphs" \ + --baseline-options="--parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 4 --parallelism.spmd_backend spmd_types --training.disable_cuda_graphs" \ + --test-options="--parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 4 --parallelism.spmd_backend spmd_types --training.disable_cuda_graphs" \ --job-dump-folder="${RUNNER_TEMP}/artifacts-to-be-uploaded/moe_loss_comparison" \ --import-result="${MOE_LOSS_FILE}" --assert-equal --steps=100 rm -rf $RUNNER_TEMP/artifacts-to-be-uploaded/* diff --git a/torchtitan/components/optimizer.py b/torchtitan/components/optimizer.py index fe1979a135..d31a8f1f04 100644 --- a/torchtitan/components/optimizer.py +++ b/torchtitan/components/optimizer.py @@ -26,6 +26,7 @@ from torchtitan.config import Configurable from torchtitan.distributed import ParallelDims from torchtitan.distributed.flex_shard import build_distributed_muon +from torchtitan.distributed.spmd_types import spmd_dense_mesh from torchtitan.tools.logging import logger __all__ = [ @@ -480,16 +481,22 @@ def _update_expert_bias( placements=[Replicate()] * dtensor_mesh.ndim ) else: - # non-full_dtensor: DTensor mesh only has TP/EP (if enabled). - # full_tensor() reduces on TP/EP, then all-reduce on loss_mesh - # covers DP/CP separately. + # non-full_dtensor: reduce TP in-band, then DP/CP on loss_mesh. is_dtensor = isinstance( tokens_per_expert_E_by_layer, torch.distributed.tensor.DTensor ) if is_dtensor: dtensor_mesh = tokens_per_expert_E_by_layer.device_mesh - tokens_per_expert_E_by_layer = ( - tokens_per_expert_E_by_layer.full_tensor() + tokens_per_expert_E_by_layer = tokens_per_expert_E_by_layer.to_local() + tp_group = dtensor_mesh.get_group("tp") + elif parallel_dims.tp > 1: + tp_group = spmd_dense_mesh()["tp"].get_group() + else: + tp_group = None + if tp_group is not None: + torch.distributed.all_reduce( + tokens_per_expert_E_by_layer, + group=tp_group, ) if loss_mesh is not None: torch.distributed.all_reduce( From f912c6c5f0239f33eb003ea4b5a23b7db3a7297b Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Fri, 14 Aug 2026 14:35:43 -0700 Subject: [PATCH 2/3] Update (base update) [ghstack-poisoned] From e1fb02b9e337c6a21fbd2ea644e727d8ce130ece Mon Sep 17 00:00:00 2001 From: Pian Pawakapan Date: Fri, 14 Aug 2026 16:40:08 -0700 Subject: [PATCH 3/3] Update [ghstack-poisoned] --- tests/unit_tests/test_optimizer_param_groups.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/test_optimizer_param_groups.py b/tests/unit_tests/test_optimizer_param_groups.py index 58d3e7d0df..43bcbf0700 100644 --- a/tests/unit_tests/test_optimizer_param_groups.py +++ b/tests/unit_tests/test_optimizer_param_groups.py @@ -76,6 +76,7 @@ def __init__(self, load_balance_coeffs=(0.1, 0.2)): class FakeParallelDims: spmd_backend = "none" + tp = 1 def get_optional_mesh(self, name): return None