From eea0bdd40946ead6380aef968451d7c40ef5ebe0 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 08:13:42 +0000 Subject: [PATCH 1/3] Price the shared-expert return under context and expert parallelism The planner dropped a MoE layer's shared-expert return whenever CP or EP exceeded 1, although CP shards rows and shared experts are not expert- parallel. Qwen3.6-35B-A3B traces show the same gated pair per local token at CP1, CP2 and EP2, which left CP2 and EP2 checkpoint workspace 8 KB per token short. Keep TP, PP and ETP at 1. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 9 +++++++-- tests/unit/test_trainer_rank_shared_memory.py | 20 ++++++++++++++++++- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index d4718cab0..bc98864b0 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1355,12 +1355,17 @@ def _shared_expert_output_bytes_per_token(layer: torch.nn.Module) -> int: (type(getattr(config, name, None)) is not int or getattr(config, name) != 1) for name in ( "tensor_model_parallel_size", - "context_parallel_size", "pipeline_model_parallel_size", - "expert_model_parallel_size", "expert_tensor_parallel_size", ) ) + # CP shards rows and shared experts are not expert-parallel, so each + # local token returns the same shared output; Qwen3.6-35B-A3B traces + # show the same gated pair per token at CP1, CP2 and EP2. + or any( + (type(getattr(config, name, None)) is not int or getattr(config, name) < 1) + for name in ("context_parallel_size", "expert_model_parallel_size") + ) ): return 0 hidden = getattr(config, "hidden_size", None) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 79b73bbcb..765f300ab 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -201,7 +201,8 @@ def test_original_norm_installation_preserves_shared_return(layer, gated): x, "shared_experts_compute", lambda *a: None ), "topology bool": lambda x: setattr(x.config, "tensor_model_parallel_size", True), - "topology two": lambda x: setattr(x.config, "context_parallel_size", 2), + "topology two": lambda x: setattr(x.config, "pipeline_model_parallel_size", 2), + "topology zero": lambda x: setattr(x.config, "context_parallel_size", 0), "missing config": lambda x: delattr(x.shared_experts, "config"), } @@ -383,3 +384,20 @@ def test_invalid_pre_gate_cache_stays_inside_planning_status(layer, bad): ValueError, match="Invalid constructor checkpoint MoE coefficient" ): rank._plan_cost(plan) + + +@pytest.mark.parametrize( + "name", ["context_parallel_size", "expert_model_parallel_size"] +) +def test_shared_return_is_per_local_token_under_cp_and_ep(layer, name): + # CP shards rows and shared experts are not expert-parallel: each local + # token still returns one shared output. + shared_layer(layer) + setattr(layer.config, name, 2) + assert _shared_expert_output_bytes_per_token(layer) == 4096 + + +def test_cp_ranks_price_the_shared_return_beside_routed_rows(layer): + shared_layer(layer) + layer.config.context_parallel_size = 2 + assert _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=2)) == 192512 From 807ee7c3bc034625f17ef78bd795a2ed90e706d6 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 08:23:59 +0000 Subject: [PATCH 2/3] Test that the shared return escapes the EP routed-row allowance Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_shared_memory.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 765f300ab..007296224 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -3,7 +3,7 @@ from types import SimpleNamespace import pytest -from test_trainer_rank_moe_memory import _enclosing_moe, _rank +from test_trainer_rank_moe_memory import _enclosing_moe, _hybridep, _rank from test_trainer_rank_moe_memory import layer as layer from test_trainer_rank_pending_memory import full_requests, module, rank_with_moe import torch @@ -401,3 +401,19 @@ def test_cp_ranks_price_the_shared_return_beside_routed_rows(layer): shared_layer(layer) layer.config.context_parallel_size = 2 assert _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=2)) == 192512 + + +@pytest.mark.parametrize("checkpoint_grad,shared", [(False, 4096), (True, 8192)]) +def test_shared_return_escapes_the_ep_routed_allowance(layer, checkpoint_grad, shared): + shared_layer(layer) + local_experts = layer.token_dispatcher.num_local_experts + layer.config.context_parallel_size = layer.config.expert_model_parallel_size = 2 + _hybridep(layer, 2).token_dispatcher.num_local_experts = local_experts + # HybridEP's 1.5x allowance turns top-k 8 into 12 routed rows; the gated + # shared return (doubled for checkpoint backward) is per local token. + assert ( + _moe_output_bytes_per_token( + [layer], ParallelShape(tp=1, cp=2, ep=2), checkpoint_grad=checkpoint_grad + ) + == 12 * 11776 * 2 + shared + ) From 216bc1ef67b7677840f6c3fde5486d012cd6dac7 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 08:36:09 +0000 Subject: [PATCH 3/3] Give the EP2 test dispatcher the per-rank expert count Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_shared_memory.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_shared_memory.py b/tests/unit/test_trainer_rank_shared_memory.py index 007296224..36bc971d7 100644 --- a/tests/unit/test_trainer_rank_shared_memory.py +++ b/tests/unit/test_trainer_rank_shared_memory.py @@ -408,7 +408,8 @@ def test_shared_return_escapes_the_ep_routed_allowance(layer, checkpoint_grad, s shared_layer(layer) local_experts = layer.token_dispatcher.num_local_experts layer.config.context_parallel_size = layer.config.expert_model_parallel_size = 2 - _hybridep(layer, 2).token_dispatcher.num_local_experts = local_experts + # EP2 halves the experts each rank owns. + _hybridep(layer, 2).token_dispatcher.num_local_experts = local_experts // 2 # HybridEP's 1.5x allowance turns top-k 8 into 12 routed rows; the gated # shared return (doubled for checkpoint backward) is per local token. assert (