diff --git a/tests/unit_tests/test_minimal_async_ep_kernels.py b/tests/unit_tests/test_minimal_async_ep_kernels.py index c8e986b23c..26b33f00b3 100644 --- a/tests/unit_tests/test_minimal_async_ep_kernels.py +++ b/tests/unit_tests/test_minimal_async_ep_kernels.py @@ -19,7 +19,7 @@ from torch.fx.experimental.proxy_tensor import make_fx from torchtitan.distributed.minimal_async_ep.kernels import ( - _copy_rows_to_peer_ptrs_kernel, + _copy_rows_to_peer_ptrs_persistent_kernel, copy_full_counts_to_peers_kernel, copy_rows_to_peers_kernel, expand_topk_grad_kernel, @@ -284,6 +284,33 @@ def test_receive_capacity_rejects_global_routing_overflow(): ) +def test_overlap_copy_grid_propagates_to_minimal_async_ep_dispatchers(): + from torchtitan.experiments.graph_trainer.deepseek_v3.config_registry import ( + graph_trainer_deepseek_v3_671b_bf16_minimal_async_ep, + ) + from torchtitan.models.common.token_dispatcher import MinimalAsyncEPTokenDispatcher + + config = graph_trainer_deepseek_v3_671b_bf16_minimal_async_ep() + config.compile.memory_policy = "full" + config.model_spec.model.update_from_config(config=config) + dispatchers = [ + layer.moe.routed_experts.token_dispatcher + for layer in config.model_spec.model.layers + if layer.moe is not None + ] + assert dispatchers + assert all(dispatcher.num_row_copy_ctas is None for dispatcher in dispatchers) + + config.compile.ep_overlap.enabled = True + config.model_spec.model.update_from_config(config=config) + + assert all( + isinstance(dispatcher, MinimalAsyncEPTokenDispatcher.Config) + and dispatcher.num_row_copy_ctas == 50 + for dispatcher in dispatchers + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA required") class TestMinimalAsyncEPKernels(unittest.TestCase): @unittest.skipUnless( @@ -354,6 +381,7 @@ def run_exchange(x, scores, expert_ids): dtype=torch.float32, device=device, receive_capacity_factor=2.0, + num_row_copy_ctas=17, ) rank = dist.get_rank() x = ( @@ -404,8 +432,32 @@ def wrapped(*args, **kwargs): buffer_state = minimal_async_ep_api._buffer_state assert buffer_state is not None - comm_stream = buffer_state.comm_stream.cuda_stream - self.assertEqual({stream for _, stream in launch_streams}, {comm_stream}) + prep_stream = buffer_state.prep_stream.cuda_stream + copy_stream = buffer_state.copy_stream.cuda_stream + self.assertNotEqual(prep_stream, copy_stream) + self.assertLess( + buffer_state.copy_stream.priority, + buffer_state.prep_stream.priority, + ) + self.assertEqual( + buffer_state.num_row_copy_ctas, + 17, + ) + prep_functions = { + "copy_full_counts_to_peers_kernel", + "fill_dispatch_metadata_kernel", + "fill_combine_metadata_kernel", + "invert_flat_indices_kernel", + } + for name, stream in launch_streams: + if name in prep_functions: + self.assertEqual(stream, prep_stream) + elif name == "copy_rows_to_peers_kernel": + self.assertEqual(stream, copy_stream) + wait_streams = { + stream for name, stream in launch_streams if name == "_wait_ready" + } + self.assertEqual(wait_streams, {prep_stream, copy_stream}) self.assertEqual( {name for name, _ in launch_streams}, set(launch_functions) ) @@ -784,6 +836,26 @@ def test_copy_kernels_handle_strides_and_active_row_masks(self): assert_equal(row_dsts[0][4], src[0]) assert_equal(row_dsts[1][5], torch.zeros(2, device="cuda")) + bounded_src = torch.arange(128, device="cuda", dtype=torch.float32).view(64, 2) + bounded_dst = torch.zeros_like(bounded_src) + bounded_dst_ptrs = torch.tensor( + [bounded_dst.data_ptr()], device="cuda", dtype=torch.int64 + ) + copy_rows_to_peers_kernel( + bounded_src, + [bounded_dst], + torch.zeros(64, device="cuda", dtype=torch.int64), + torch.arange(64, device="cuda", dtype=torch.int64), + ep_size=1, + num_rows=64, + num_cols=2, + dst_ptrs=bounded_dst_ptrs, + block_m=1, + num_ctas=2, + ) + torch.cuda.synchronize() + assert_equal(bounded_dst, bounded_src) + def test_copy_rows_uses_int64_for_source_stride_arithmetic(self): row = 1_048_576 stride = 2048 @@ -814,7 +886,7 @@ def test_copy_rows_uses_int64_for_source_stride_arithmetic(self): src_storage[base_offset + high_offset] = 93 torch.cuda.synchronize() - _copy_rows_to_peer_ptrs_kernel[(metadata_numel, 1)]( + _copy_rows_to_peer_ptrs_persistent_kernel[(20,)]( src, dst_ptrs, dst_ranks, @@ -832,6 +904,7 @@ def test_copy_rows_uses_int64_for_source_stride_arithmetic(self): SRC_ROW_DIVISOR=1, BLOCK_M=1, BLOCK_N=1, + NUM_COL_TILES=1, ) torch.cuda.synchronize() diff --git a/torchtitan/distributed/minimal_async_ep/api.py b/torchtitan/distributed/minimal_async_ep/api.py index 918b9de96c..2181ba238b 100644 --- a/torchtitan/distributed/minimal_async_ep/api.py +++ b/torchtitan/distributed/minimal_async_ep/api.py @@ -92,7 +92,9 @@ class _MinimalAsyncEPBufferState: counts_recv_handle: Any counts_recv_peer_buffers: list[torch.Tensor] counts_recv_peer_ptrs: torch.Tensor - comm_stream: torch.cuda.Stream + prep_stream: torch.cuda.Stream + copy_stream: torch.cuda.Stream + num_row_copy_ctas: int | None check_receive_capacity: bool hidden_recv_buffer_index: int = 0 pending_events: dict[tuple[str, int], deque[_PendingEvent]] = field( @@ -184,6 +186,12 @@ def maybe_update_minimal_async_ep_config(model_config: Any, config: Any) -> None "MinimalAsyncEP receive capacity factor must be finite and at least " "1.0, or None." ) + overlap_enabled = bool(getattr(overlap_config, "enabled", False)) + num_row_copy_ctas = ( + getattr(overlap_config, "minimal_async_ep_num_copy_ctas", None) + if overlap_enabled + else None + ) for token_dispatcher_cfg in dispatcher_cfgs: token_dispatcher_cfg.hidden_dim = model_config.dim @@ -193,6 +201,7 @@ def maybe_update_minimal_async_ep_config(model_config: Any, config: Any) -> None token_dispatcher_cfg.dtype = TORCH_DTYPE_MAP[ config.training.mixed_precision_param ] + token_dispatcher_cfg.num_row_copy_ctas = num_row_copy_ctas token_dispatcher_cfg.force_load_balance = config.debug.moe_force_load_balance token_dispatcher_cfg.receive_capacity_factor = receive_capacity_factor token_dispatcher_cfg.receive_capacity = _buffer_rows( @@ -285,6 +294,7 @@ def init_buffer( device: torch.device, force_load_balance: bool = False, receive_capacity_factor: float | None = None, + num_row_copy_ctas: int | None = None, ) -> None: """Initialize the process-local MinimalAsyncEP symmetric-memory buffer.""" global _buffer_state @@ -294,6 +304,11 @@ def init_buffer( num_experts = ep_size * num_local_experts assert _buffer_state is None + if num_row_copy_ctas is not None and num_row_copy_ctas < 1: + raise ValueError( + f"num_row_copy_ctas must be positive or None, got {num_row_copy_ctas}." + ) + dispatch_rows, combine_rows = _buffer_rows( tokens_per_rank, top_k, @@ -316,7 +331,7 @@ def init_buffer( logger.info( "Initializing MinimalAsyncEP buffer: hidden_dim=%d, tokens_per_rank=%d, " "top_k=%d, num_local_experts=%d, ep_size=%d, capacity_mode=%s, " - "buffer_rows=%d", + "buffer_rows=%d, num_row_copy_ctas=%s", hidden_dim, tokens_per_rank, top_k, @@ -324,6 +339,7 @@ def init_buffer( ep_size, capacity_mode, dispatch_rows, + num_row_copy_ctas if num_row_copy_ctas is not None else "unbounded", ) backend = symm_mem.get_backend(device) if backend != "CUDA": @@ -386,7 +402,9 @@ def init_buffer( counts_recv_handle=counts_recv_handle, counts_recv_peer_buffers=counts_recv_peer_buffers, counts_recv_peer_ptrs=counts_recv_peer_ptrs, - comm_stream=torch.cuda.Stream(device=device), + prep_stream=torch.cuda.Stream(device=device), + copy_stream=torch.cuda.Stream(device=device, priority=-1), + num_row_copy_ctas=num_row_copy_ctas, check_receive_capacity=( receive_capacity_factor is not None and not force_load_balance ), @@ -442,7 +460,7 @@ def _copy_rows_to_peers_async_cuda( num_valid_rows: torch.Tensor | None = None, retained: tuple[torch.Tensor, ...] = (), ) -> torch.Tensor: - """Launch a row copy through symmetric memory on the EP comm stream.""" + """Launch a row copy through symmetric memory on the EP copy stream.""" assert _buffer_state is not None buffer_rows = _buffer_state.hidden_recv_buffers[0].shape[0] @@ -468,9 +486,9 @@ def _copy_rows_to_peers_async_cuda( ) launch_stream = torch.cuda.current_stream(x.device) - comm_stream = _buffer_state.comm_stream - comm_stream.wait_stream(launch_stream) - with torch.cuda.stream(comm_stream): + copy_stream = _buffer_state.copy_stream + copy_stream.wait_stream(launch_stream) + with torch.cuda.stream(copy_stream): copy_rows_to_peers_kernel( x, hidden_recv_peer_buffers, @@ -485,9 +503,10 @@ def _copy_rows_to_peers_async_cuda( src_row_divisor=src_row_divisor, dst_ptrs=hidden_recv_peer_ptrs, num_valid_rows=num_valid_rows, + num_ctas=_buffer_state.num_row_copy_ctas, ) _wait_ready(hidden_recv_handle, _HIDDEN_READY_CHANNEL) - _record_pending_event(exchange, hidden_recv_view, comm_stream, retained) + _record_pending_event(exchange, hidden_recv_view, copy_stream, retained) return hidden_recv_view @@ -764,9 +783,9 @@ def dispatch_op( ) launch_stream = torch.cuda.current_stream(dispatch_input.device) - comm_stream = _buffer_state.comm_stream - comm_stream.wait_stream(launch_stream) - with torch.cuda.stream(comm_stream): + prep_stream = _buffer_state.prep_stream + prep_stream.wait_stream(launch_stream) + with torch.cuda.stream(prep_stream): torch.argsort( T_row_to_expert_N, stable=True, diff --git a/torchtitan/distributed/minimal_async_ep/kernels.py b/torchtitan/distributed/minimal_async_ep/kernels.py index 086b31b3df..fcbeeb9f47 100644 --- a/torchtitan/distributed/minimal_async_ep/kernels.py +++ b/torchtitan/distributed/minimal_async_ep/kernels.py @@ -351,6 +351,64 @@ def _copy_rows_to_peer_ptrs_kernel( tl.store(dst_ptr, values, mask=mask & dst_rank_mask[:, None]) +# Preserve the 2-D launch above for unbounded copies. Flattening its grid changes +# peer issue order and causes severe rank imbalance at EP64. +@triton.jit +def _copy_rows_to_peer_ptrs_persistent_kernel( + src, + dst_ptrs: tl.pointer_type(tl.int64), + dst_ranks: tl.pointer_type(tl.int64), + dst_rows: tl.pointer_type(tl.int64), + num_valid_rows: tl.pointer_type(tl.int64), + src_rows: tl.pointer_type(tl.int64), + NUM_ROWS: tl.constexpr, + NUM_COLS: tl.constexpr, + SRC_ROW_STRIDE: tl.constexpr, + SRC_COL_STRIDE: tl.constexpr, + DST_ROW_STRIDE: tl.constexpr, + DST_DTYPE: tl.constexpr, + HAS_NUM_VALID_ROWS: tl.constexpr, + HAS_SRC_ROWS: tl.constexpr, + SRC_ROW_DIVISOR: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_COL_TILES: tl.constexpr, +) -> None: + row_limit = NUM_ROWS + if HAS_NUM_VALID_ROWS: + row_limit = tl.load(num_valid_rows) + num_tiles = tl.cdiv(row_limit, BLOCK_M) * NUM_COL_TILES + tile = tl.program_id(0) + while tile < num_tiles: + row_tile = tile // NUM_COL_TILES + col_tile = tile % NUM_COL_TILES + row = (row_tile * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64) + col = col_tile * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = row < row_limit + col_mask = col < NUM_COLS + mask = row_mask[:, None] & col_mask[None, :] + src_row = row + if HAS_SRC_ROWS: + src_row = tl.load(src_rows + row, mask=row_mask, other=0).to(tl.int64) + if SRC_ROW_DIVISOR != 1: + src_row = src_row // SRC_ROW_DIVISOR + dst_rank = tl.load(dst_ranks + row, mask=row_mask, other=-1) + dst_row = tl.load(dst_rows + row, mask=row_mask, other=0).to(tl.int64) + dst_rank_mask = row_mask & (dst_rank >= 0) + values = tl.load( + src + src_row[:, None] * SRC_ROW_STRIDE + col[None, :] * SRC_COL_STRIDE, + mask=mask, + ) + dst_base = tl.load(dst_ptrs + dst_rank, mask=dst_rank_mask, other=0) + dst_ptr = ( + dst_base.to(tl.pointer_type(DST_DTYPE))[:, None] + + dst_row[:, None] * DST_ROW_STRIDE + + col[None, :] + ) + tl.store(dst_ptr, values, mask=mask & dst_rank_mask[:, None]) + tile += tl.num_programs(0) + + def copy_full_counts_to_peers_kernel( counts: torch.Tensor, dsts: list[torch.Tensor], @@ -397,22 +455,29 @@ def copy_rows_to_peers_kernel( src_rows: torch.Tensor | None = None, src_row_divisor: int = 1, num_valid_rows: torch.Tensor | None = None, + num_ctas: int | None = None, ) -> None: if len(dsts) != ep_size: raise ValueError(f"expected {ep_size} destination buffers, got {len(dsts)}.") + if num_ctas is not None and num_ctas <= 0: + raise ValueError(f"num_ctas must be positive, got {num_ctas}.") + block_n = min(_MAX_BLOCK_N, triton.next_power_of_2(num_cols)) - grid = (triton.cdiv(num_rows, block_m), triton.cdiv(num_cols, block_n)) + num_row_tiles = triton.cdiv(num_rows, block_m) + num_col_tiles = triton.cdiv(num_cols, block_n) dst_dtype = _HIDDEN_ROW_DTYPES.get(src.dtype) if dst_dtype is None: raise ValueError(f"Unsupported MinimalAsyncEP row-copy dtype: {src.dtype}.") - _copy_rows_to_peer_ptrs_kernel[grid]( + common_args = ( src, dst_ptrs, dst_ranks, dst_rows, num_valid_rows if num_valid_rows is not None else dst_rows[:1], src_rows if src_rows is not None else dst_rows, + ) + common_kwargs = dict( NUM_ROWS=num_rows, NUM_COLS=num_cols, SRC_ROW_STRIDE=src.stride(0), @@ -426,6 +491,18 @@ def copy_rows_to_peers_kernel( BLOCK_N=block_n, num_warps=num_warps, ) + if num_ctas is None: + _copy_rows_to_peer_ptrs_kernel[(num_row_tiles, num_col_tiles)]( + *common_args, + **common_kwargs, + ) + else: + num_tiles = num_row_tiles * num_col_tiles + _copy_rows_to_peer_ptrs_persistent_kernel[(min(num_tiles, num_ctas),)]( + *common_args, + NUM_COL_TILES=num_col_tiles, + **common_kwargs, + ) def fill_dispatch_metadata_kernel( diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 5fa1b5bcfe..ecec578ea7 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -61,6 +61,13 @@ class EpOverlapConfig: the configured capacity instead of dropping tokens. """ + minimal_async_ep_num_copy_ctas: int | None = 50 + """Persistent row-copy grid size used by MinimalAsyncEP during overlap. + + ``None`` leaves the copy kernel unbounded. This setting has no effect on + other EP backends or when EP overlap is disabled. + """ + disable_early_grad_accumulation: bool = False """Disable chunked parameter-gradient accumulation before communication. @@ -225,6 +232,13 @@ def validate_ep_overlap_config( "must be finite and at least 1.0, or None" ) + num_copy_ctas = ep_overlap_config.minimal_async_ep_num_copy_ctas + if num_copy_ctas is not None and num_copy_ctas < 1: + raise ValueError( + "--compile.ep_overlap.minimal_async_ep_num_copy_ctas must be " + "positive or None" + ) + return chunk_dim, chunk_strategy, module_fqn diff --git a/torchtitan/experiments/graph_trainer/tests/test_passes.py b/torchtitan/experiments/graph_trainer/tests/test_passes.py index a06d3da67a..05617d4f53 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_passes.py @@ -39,6 +39,7 @@ from torchtitan.experiments.graph_trainer.configs import ( EpOverlapConfig, GraphTrainerCompileConfig, + validate_ep_overlap_config, ) from torchtitan.experiments.graph_trainer.cudagraph import ( insert_kernel_annotations_pass, @@ -3618,6 +3619,12 @@ def test_ep_overlap_pass_pipeline_order(self): disabled_names.index("joint_transformer_block_bucketing_reordering_pass"), ) + def test_ep_overlap_rejects_invalid_minimal_async_ep_copy_grid(self): + with self.assertRaisesRegex(ValueError, "must be positive or None"): + validate_ep_overlap_config( + EpOverlapConfig(enabled=True, minimal_async_ep_num_copy_ctas=0) + ) + def test_graph_ep_chunking_rejects_tensor_parallel(self): cases = ( ("seq", "layers.*.moe"), diff --git a/torchtitan/models/common/token_dispatcher.py b/torchtitan/models/common/token_dispatcher.py index d041ea4e5a..1b63842ba1 100644 --- a/torchtitan/models/common/token_dispatcher.py +++ b/torchtitan/models/common/token_dispatcher.py @@ -1114,6 +1114,7 @@ class MinimalAsyncEPTokenDispatcher(LocalTokenDispatcher): tokens_per_rank: int | None dtype: torch.dtype | None buffer_device: torch.device + num_row_copy_ctas: int | None force_load_balance: bool receive_capacity_factor: float | None receive_capacity: int | None @@ -1124,6 +1125,7 @@ class Config(LocalTokenDispatcher.Config): tokens_per_rank: int | None = None dtype: torch.dtype | None = None device: torch.device | None = None + num_row_copy_ctas: int | None = None force_load_balance: bool = False receive_capacity_factor: float | None = None receive_capacity: int | None = None @@ -1135,6 +1137,7 @@ def __init__(self, config: Config): self.hidden_dim = config.hidden_dim self.tokens_per_rank = config.tokens_per_rank self.dtype = config.dtype + self.num_row_copy_ctas = config.num_row_copy_ctas self.force_load_balance = config.force_load_balance self.receive_capacity_factor = config.receive_capacity_factor self.receive_capacity = config.receive_capacity @@ -1157,6 +1160,7 @@ def __init__(self, config: Config): int, torch.dtype, torch.device, + int | None, bool, float | None, int, @@ -1221,6 +1225,7 @@ def init_buffer(self) -> None: self.top_k, self.dtype, self.buffer_device, + self.num_row_copy_ctas, self.force_load_balance, self.receive_capacity_factor, self.receive_capacity, @@ -1243,6 +1248,7 @@ def init_buffer(self) -> None: device=self.buffer_device, force_load_balance=self.force_load_balance, receive_capacity_factor=self.receive_capacity_factor, + num_row_copy_ctas=self.num_row_copy_ctas, ) MinimalAsyncEPTokenDispatcher._global_buffer_key = buffer_key