diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/common.py b/tensorrt_llm/_torch/auto_deploy/distributed/common.py index 9f1dc85f1063..aba29b21dce2 100644 --- a/tensorrt_llm/_torch/auto_deploy/distributed/common.py +++ b/tensorrt_llm/_torch/auto_deploy/distributed/common.py @@ -2,6 +2,7 @@ import atexit import os +import socket import sys from typing import Callable, List, Optional, Tuple @@ -15,6 +16,8 @@ # TODO: check to what extend we can reuse _torch/distributed.py +_MASTER_ADDR = "127.0.0.1" + class _DistGroup: """Global instance to set/get the default process group for distributed ops.""" @@ -85,16 +88,16 @@ def initialize_or_skip( rank: int = 0, world_size: int = 1, port: Optional[int] = None, - shared_port: Optional["mp.Value"] = None, - port_ready_barrier: Optional["mp.Barrier"] = None, + port_recv_conn=None, + port_send_conns=None, ) -> Tuple[int, int]: if not dist.is_initialized(): return initialize( rank=rank, world_size=world_size, port=port, - shared_port=shared_port, - port_ready_barrier=port_ready_barrier, + port_recv_conn=port_recv_conn, + port_send_conns=port_send_conns, ) return get_rank(), get_world_size() @@ -124,37 +127,28 @@ def _set_distributed_env_vars(local_rank: int, world_size: int, port: int) -> No """Set environment variables required by NCCL's env:// init method.""" os.environ["RANK"] = str(local_rank) os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_ADDR"] = _MASTER_ADDR os.environ["MASTER_PORT"] = str(port) os.environ["LOCAL_RANK"] = str(local_rank) -def _try_init_process_group(local_rank: int, world_size: int, port: int) -> bool: - """Attempt to initialize process group. Returns True on success, False on EADDRINUSE.""" - _set_distributed_env_vars(local_rank, world_size, port) - +def _is_port_available(port: int) -> bool: + """Lightweight check: try to bind to the port and release immediately.""" try: - dist.init_process_group( - "nccl", - world_size=world_size, - rank=local_rank, - device_id=torch.device(local_rank), - ) - return True - except Exception as e: - # Check if this is a port-in-use error (only rank 0 binds, so only rank 0 can get this) - if "EADDRINUSE" in str(e) or "address already in use" in str(e).lower(): - ad_logger.warning(f"Port {port} already in use, will retry with new port") - return False - raise + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind((_MASTER_ADDR, port)) + return True + except OSError: + return False def initialize( rank: int = 0, world_size: int = 1, port: Optional[int] = None, - shared_port: Optional["mp.Value"] = None, - port_ready_barrier: Optional["mp.Barrier"] = None, + port_recv_conn=None, + port_send_conns=None, max_retries: int = 5, ) -> Tuple[int, int]: """Initialize distributed process group. @@ -163,8 +157,8 @@ def initialize( rank: Process rank (ignored for OMPI/torchelastic). world_size: Total number of processes (ignored for OMPI/torchelastic). port: Initial port to try. If None, a free port will be selected. - shared_port: Optional mp.Value for rank 0 to share the final port with other ranks. - port_ready_barrier: Optional mp.Barrier to synchronize port selection. + port_recv_conn: Pipe read-end for this rank to receive the resolved port from rank 0. + port_send_conns: List of Pipe write-ends for rank 0 to broadcast the port. max_retries: Maximum number of port retry attempts for rank 0. """ if is_ompi(): @@ -189,65 +183,57 @@ def initialize( # Necessary to assign a device to each rank. torch.cuda.set_device(local_rank) - # If we have shared port synchronization (multiprocess spawn mode) - if shared_port is not None and port_ready_barrier is not None: - if local_rank == 0: - # Rank 0: try ports until one works, then share with other ranks - init_success = False - init_error = None - try: - for attempt in range(max_retries): - ad_logger.info( - f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=} (attempt {attempt + 1})" - ) - if _try_init_process_group(local_rank, world_size, port): - # Success! Share the working port with other ranks - shared_port.value = port - init_success = True - break - else: - # Port was taken, try a new one - port = get_free_port() - else: - # All retries exhausted - init_error = RuntimeError( - f"Failed to find available port after {max_retries} attempts" - ) - except Exception as e: - # Catch any unexpected error so we can still signal other ranks - init_error = e - finally: - # ALWAYS signal other ranks, even on error, to prevent deadlock - if not init_success: - shared_port.value = -1 - port_ready_barrier.wait() - - if init_error is not None: - raise init_error - else: - # Other ranks: wait for rank 0 to find a working port - port_ready_barrier.wait() - port = shared_port.value - if port == -1: - raise RuntimeError("Rank 0 failed to initialize, cannot proceed") - ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}") - _set_distributed_env_vars(local_rank, world_size, port) - dist.init_process_group( - "nccl", - world_size=world_size, - rank=local_rank, - device_id=torch.device(local_rank), - ) - else: - # Original path: no retry mechanism (OMPI, torchelastic, or single process) - ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}") - _set_distributed_env_vars(local_rank, world_size, port) - dist.init_process_group( - "nccl", - world_size=world_size, - rank=local_rank, - device_id=torch.device(local_rank), - ) + # Pipe-based port synchronization (multiprocess spawn mode). + # Port resolution MUST happen before init_process_group because that call + # is collective — all ranks must participate simultaneously. Rank 0 does a + # lightweight socket-bind check (not a full init) to find a usable port, + # broadcasts it through the pipes, and then ALL ranks call + # init_process_group together. + if port_send_conns is not None: + assert local_rank == 0, "port_send_conns should only be provided to rank 0" + port_ok = False + init_error = None + try: + for attempt in range(max_retries): + ad_logger.info( + f"Checking port for: {lib=}, {local_rank=}, {world_size=}, {port=} (attempt {attempt + 1})" + ) + if _is_port_available(port): + port_ok = True + break + ad_logger.warning(f"Port {port} already in use, will retry with new port") + port = get_free_port() + else: + init_error = RuntimeError( + f"Failed to find available port after {max_retries} attempts" + ) + except Exception as e: + init_error = e + finally: + # ALWAYS broadcast to other ranks, even on error, to prevent deadlock + final_port = port if port_ok else -1 + for conn in port_send_conns: + conn.send(final_port) + conn.close() + + if init_error is not None: + raise init_error + elif port_recv_conn is not None: + # Non-rank-0: wait for rank 0 to broadcast the working port + port = port_recv_conn.recv() + port_recv_conn.close() + if port == -1: + raise RuntimeError("Rank 0 failed to initialize, cannot proceed") + + # All paths converge here: initialize the process group collectively + ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}") + _set_distributed_env_vars(local_rank, world_size, port) + dist.init_process_group( + "nccl", + world_size=world_size, + rank=local_rank, + device_id=torch.device(local_rank), + ) # Register cleanup function to be called at exit atexit.register(cleanup) @@ -259,11 +245,11 @@ def initialize( def init_and_run_process( - job, rank, size, port, shared_port=None, port_ready_barrier=None, **kwargs + job, rank, size, port, port_recv_conn=None, port_send_conns=None, **kwargs ): try: initialize_or_skip( - rank, size, port, shared_port=shared_port, port_ready_barrier=port_ready_barrier + rank, size, port, port_recv_conn=port_recv_conn, port_send_conns=port_send_conns ) job(rank, size, **kwargs) except Exception as e: @@ -318,11 +304,17 @@ def _start_multiprocess_job( ctx = mp.get_context("spawn") processes: List[mp.Process] = [] - # Create shared state for port synchronization with retry mechanism: - # - shared_port: rank 0 writes the final working port here - # - port_ready_barrier: all ranks wait here until rank 0 has bound successfully - shared_port = ctx.Value("i", port) # 'i' = signed int - port_ready_barrier = ctx.Barrier(size) + # Create pipes for port synchronization with retry mechanism. + # Pipes use sockets (not POSIX named semaphores like Value/Barrier), + # so they work reliably on all filesystems including network mounts. + # Rank 0 finds a working port and sends it through each pipe; + # other ranks block on recv() until rank 0 is ready. + port_recv_conns = [] + port_send_conns = [] + for _ in range(size - 1): + recv_conn, send_conn = ctx.Pipe(duplex=False) + port_recv_conns.append(recv_conn) + port_send_conns.append(send_conn) for rank in range(size): if input_queues: @@ -330,10 +322,15 @@ def _start_multiprocess_job( if output_queue: kwargs["output_queue"] = output_queue if rank == 0 else None + rank_kwargs = { + **kwargs, + "port_recv_conn": port_recv_conns[rank - 1] if rank > 0 else None, + "port_send_conns": port_send_conns if rank == 0 else None, + } p = ctx.Process( target=init_and_run_process, args=(job, rank, size, port), - kwargs={**kwargs, "shared_port": shared_port, "port_ready_barrier": port_ready_barrier}, + kwargs=rank_kwargs, daemon=True, ) p.start() @@ -350,9 +347,11 @@ def _join_multiprocess_job(processes): for p in processes: p.join() - # Ensure that all processes have exited successfully - if isinstance(p, mp.Process): - assert not p.exitcode, f"Process {p.pid} exited with code {p.exitcode}" + # Ensure that all processes have exited successfully. + # Check exitcode via hasattr rather than isinstance(p, mp.Process), because + # spawn-context processes (SpawnProcess) don't inherit from mp.Process. + if hasattr(p, "exitcode"): + assert p.exitcode == 0, f"Process {p.pid} exited with code {p.exitcode}" def spawn_multiprocess_job(job: Callable[[int, int], None], size: Optional[int] = None): diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py index 3c9398a946e3..3acbdeaf74ee 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py @@ -189,6 +189,7 @@ def _init_mapping(self): moe_tp_size=self.dist_mapping.get("moe_tp", 1), moe_ep_size=self.dist_mapping.get("moe_ep", self.world_size), moe_cluster_size=self.dist_mapping.get("moe_cluster", 1), + enable_attention_dp=self.enable_attention_dp, ) except ValueError as e: ad_logger.warning(f"Invalid parallel grid config: {e}") @@ -1358,9 +1359,9 @@ def _shard_parameter_node( # Shard weight using the unified function (also updates the parameter) weight_nodes = extract_weight_nodes(node) - # Parametrized nodes must have at least one weight (for debugging) - assert len(weight_nodes.weights) > 0, ( - f"Node {node.name} has no weights - weight mapping may be incorrect" + # Parametrized nodes must have at least one weight or bias (for debugging) + assert len(weight_nodes.weights) > 0 or len(weight_nodes.biases) > 0, ( + f"Node {node.name} has no weights or biases - weight mapping may be incorrect" ) for weight_node in weight_nodes.weights: @@ -1725,7 +1726,7 @@ def _insert_sharded_mxfp4_mlp_ep( def _process_simple_shard( nodes_linear: Dict[Node, List[Node]], transform_container: ShardingTransformContainer, - layer_type: LayerType = LayerType.MLP, + layer_type: LayerType = LayerType.UNKNOWN, ) -> int: # for every linear node: # --> row_split (dim 0 of weight) + all_gather (dim -1 of output) @@ -2663,6 +2664,8 @@ def detect_sharding_from_config( num_row_col_shards = 0 num_attention_shards = 0 num_ssm_shards = 0 + num_mla_shards = 0 + num_delta_shards = 0 linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op)) # use layer_subgraphs to determine the layer_type @@ -2724,7 +2727,14 @@ def detect_sharding_from_config( if _process_ssm_sharding(layer_subgraph, transform_container) > 0: num_ssm_shards += 1 num_row_col_shards += 1 - + elif config == "delta": + if _process_delta_sharding(layer_subgraph, transform_container) > 0: + num_delta_shards += 1 + num_row_col_shards += 1 + elif config == "mla": + if _process_mla_sharding(layer_subgraph, transform_container) > 0: + num_mla_shards += 1 + num_row_col_shards += 1 elif "sequence" in config: # TODO: Sequence parallelism is not supported yet. ad_logger.warning("Sequence parallelism is not supported yet. Skipping.") @@ -2783,7 +2793,8 @@ def detect_sharding_from_config( num_shards = num_simple_shards + num_row_col_shards ad_logger.info( f"Applied {num_shards} TP shards from config. Simple: {num_simple_shards}, " - f"row-col: {num_row_col_shards} (including: ssm: {num_ssm_shards}, attention: {num_attention_shards})" + f"row-col: {num_row_col_shards} (attention: {num_attention_shards}, " + f"mla: {num_mla_shards}, delta: {num_delta_shards}, ssm: {num_ssm_shards})" ) num_matches = len(transform_container.weight_sharding_transforms) diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index ddd38e4e9629..b57249b0300f 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py @@ -45,11 +45,11 @@ class LayerType(Enum): class LayerSubgraph(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) + layer_type: LayerType opening_nodes: List[Node] - subgraph_nodes: List[Node] terminating_node: Union[Node, None] - layer_type: LayerType min_local_shape: int = 1 + subgraph_nodes: List[Node] class WeightNode(BaseModel): @@ -413,11 +413,18 @@ def find_get_attr_node(weight_node: Node) -> Node: def get_weight_node(node: Node) -> Node: - """Get the primary weight node for a compute node.""" + """Get the primary weight node for a compute node. + + When the node itself is a bias get_attr node (i.e. extract_weight_nodes + puts it into .biases rather than .weights), return the bias node so that + num_users_of_weight_node gives the correct user count instead of 0. + """ weight_nodes = extract_weight_nodes(node) - if len(weight_nodes.weights) == 0: - raise ValueError(f"Node {node.name} has no weight") - return weight_nodes.weights[0].node + if len(weight_nodes.weights) > 0: + return weight_nodes.weights[0].node + if len(weight_nodes.biases) > 0: + return weight_nodes.biases[0].node + raise ValueError(f"Node {node.name} has no weight or bias") def num_users_of_weight_node(node: Node) -> int: diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py index 38a6ec03ed6c..863cc56625d9 100644 --- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py +++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py @@ -18,12 +18,13 @@ ) from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer from tensorrt_llm._torch.auto_deploy.utils._graph import lint, recompile +from tensorrt_llm._torch.auto_deploy.utils.mapping_utils import deserialize_mapping from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op from tensorrt_llm.functional import AllReduceStrategy def _run_ep_shard_job( - num_experts: int, rank: int, world_size: int, enable_attention_dp: bool + num_experts: int, enable_attention_dp: bool, rank: int, world_size: int ) -> None: device = "cuda" hidden_size = 32 @@ -68,7 +69,13 @@ def transform_check(gm): moe_nodes = [n for n in gm.graph.nodes if is_op(n, torch.ops.auto_deploy.torch_moe)] assert len(moe_nodes) == 1, f"Expected 1 MoE node, got {len(moe_nodes)}" moe_node = moe_nodes[0] - return moe_node.kwargs["enable_alltoall"] == (world_size > 1) + print(f"\nMoE node {moe_node.name}, kwargs: {moe_node.kwargs}\n\n") + assert "mapping_config" in moe_node.kwargs, ( + f"Mapping config not found in MoE node {moe_node.name}" + ) + # deserialize the mapping config string to dict + mapping_config = deserialize_mapping(moe_node.kwargs["mapping_config"]) + return mapping_config.enable_attention_dp else: def transform_check(gm): @@ -76,6 +83,9 @@ def transform_check(gm): is_op(n, torch.ops.auto_deploy.torch_dist_all_reduce) for n in gm.graph.nodes ) == (world_size > 1) + # NOTE: we need to skip output assert when enable_attention_dp is True because + # a single rank does not contain the full output (lack of final collective), + # so the output will be inherently incomplete w.r.t to the single node. run_test_transformed_gm( model, x, @@ -83,6 +93,7 @@ def transform_check(gm): check_transformed_graph=transform_check, _get_expected_num_params=partial(_get_expected_num_params, rank, world_size), test_load_hook=False, + skip_output_assert=enable_attention_dp, ) diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py index f42d4856800c..e2d649daf410 100644 --- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py +++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py @@ -26,43 +26,6 @@ from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_linear_op, is_op, is_weight_node from tensorrt_llm.functional import AllReduceStrategy -base_model_tp_plan = { - "q_proj": "colwise", - "k_proj": "colwise", - "v_proj": "colwise", - "o_proj": "rowwise", - "gate_proj": "colwise", - "up_proj": "colwise", - "down_proj": "rowwise", - "linear1": "colwise", - "linear2": "rowwise", - "linear": "gather", - # Mamba2 specific projections - "in_proj": "mamba", - "out_proj": "rowwise", - # MLA specific projections - "q_a_proj": "gather", - "q_b_proj": "colwise", - "kv_a_proj_with_mqa": "gather", - "kv_b_proj": "colwise", - # "input_layernorm.weight": "sequence_parallel", - # "post_attention_layernorm.weight": "sequence_parallel", - # "norm.weight": "sequence_parallel", - # "shared_expert.gate_proj": "local_colwise", - # "shared_expert.up_proj": "local_colwise", - # "shared_expert.down_proj": "local_rowwise", - # "experts.gate_up_proj": "local_packed_rowwise", - # "experts.down_proj": "local_colwise", - # "experts": "local", - "feed_forward": "gather", - "self": "gather", - "weight": "gather", -} - -predefined_config = { - "tp_plan": base_model_tp_plan, -} - class GQA_Block(nn.Module): def __init__( @@ -417,6 +380,7 @@ def _run_sharding_execution_job( rank: int, world_size: int, ) -> None: + # global predefined_config # init model and input batch_size = 4 sequence_len = 8 @@ -433,8 +397,27 @@ def _run_sharding_execution_job( hidden_size=num_features, num_key_value_heads=num_key_value_heads, ).to(device="cuda", dtype=torch.float16) + predefined_config = { + "tp_plan": { + "q_proj": "colwise", + "k_proj": "colwise", + "v_proj": "colwise", + "o_proj": "rowwise", + } + } elif model_cls == FP8MLP: model = model_cls(num_features, num_features, bias=bias).to("cuda") + predefined_config = { + "tp_plan": { + # gated MLP + "gate_proj": "colwise", + "up_proj": "colwise", + "down_proj": "rowwise", + # GPT-style MLP + "linear1": "colwise", + "linear2": "rowwise", + } + } elif model_cls == NemotronHMamba2Mixer: # Create config for Mamba2 based on Nemotron models # Scaled down from typical values: hidden_size=5120, ssm_state_size=128 @@ -460,6 +443,7 @@ def _run_sharding_execution_job( num_hidden_layers=1, ) model = model_cls(mamba_config, layer_idx=0).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"in_proj": "mamba"}} elif model_cls == MLA_Block: # Use actual DeepSeek-V3/R1 production values # From HuggingFace config (HunYuanPretrainedConfig defaults): @@ -483,6 +467,7 @@ def _run_sharding_execution_job( v_head_dim=v_head_dim, bias=bias, ).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"q_a_proj": "mla"}} elif model_cls in (GDN_Block, GDN_Block_Unfused): num_k_heads_gdn = 4 num_v_heads_gdn = 8 @@ -497,10 +482,28 @@ def _run_sharding_execution_job( head_v_dim=head_v_dim, bias=bias, ).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"in_proj_qkv": "delta"}} + elif model_cls == nn.Linear: + model = model_cls(num_features, num_features, bias=bias).to( + device="cuda", dtype=torch.float16 + ) + # update the tp_plan in predefined_config to force simple sharding of the single linear layer + predefined_config = {"tp_plan": {"*": "gather"}} else: model = model_cls(num_features, num_features, bias=bias).to( device="cuda", dtype=torch.float16 ) + predefined_config = { + "tp_plan": { + # gated MLP + "gate_proj": "colwise", + "up_proj": "colwise", + "down_proj": "rowwise", + # GPT-style MLP + "linear1": "colwise", + "linear2": "rowwise", + } + } x = torch.randn(batch_size, sequence_len, num_features, device="cuda", dtype=torch.float16) # Scale down GDN_Block weights if needed to avoid numerical instability @@ -541,10 +544,15 @@ def _get_expected_num_params(num_p_og: int) -> int: else: num_params = num_p_og // world_size + num_update if model_cls == MLA_Block: - # since q_a_proj is simple-sharded and followed by q_a_layernorm, the layernorm params - # are NOT sharded - they have to be replicated. To account for this, we need to add the - # number of parameters of the layernorm (weight and bias)to the number of parameters of the model. - num_params += 2 * kv_lora_rank * (world_size - 1) // world_size + # Both q_a_layernorm and kv_a_layernorm are replicated because their preceding + # projections (q_a_proj, kv_a_proj_with_mqa) use simple-shard/all_gather. + # Each LayerNorm has weight + bias of size kv_lora_rank → 4 * kv_lora_rank + # replicated params total. The base formula (num_p_og // world_size) treated + # these as sharded, so we add back the under-counted fraction. + replicated_ln_params = ( + 4 * kv_lora_rank + ) # q_a_layernorm + kv_a_layernorm, weight+bias each + num_params += replicated_ln_params * (world_size - 1) // world_size return num_params def verify_local_weight_sizes(gm) -> bool: @@ -563,13 +571,18 @@ def verify_local_weight_sizes(gm) -> bool: op_expected = getattr(torch.ops.auto_deploy, dist_op_expected) gm = torch_export_to_gm(model, args=(x,), clone=True) + + sharding_source = "heuristic" if from_config else "manual" gm_transformed = InferenceOptimizer( None, { "detect_sharding": { "stage": "sharding", "sharding_scope": ["tp"], - "tp_sharding_source": ["heuristic"], + "sharding_source": [sharding_source], + "manual_config": predefined_config, + # needed for a solitary Linear node that is not part of any layer subgraph + "shard_all_unprocessed": True, }, "sharding_transform_executor": { "stage": "sharding", @@ -617,6 +630,14 @@ def _run_pattern_detection_job( hidden_size=num_features, num_key_value_heads=num_key_value_heads, ).to(device="cuda", dtype=torch.float16) + predefined_config = { + "tp_plan": { + "q_proj": "colwise", + "k_proj": "colwise", + "v_proj": "colwise", + "o_proj": "rowwise", + } + } elif model_cls == NemotronHMamba2Mixer: # Create config for Mamba2 mamba_config = SimpleNamespace( @@ -641,6 +662,7 @@ def _run_pattern_detection_job( num_hidden_layers=1, ) model = model_cls(mamba_config, layer_idx=0).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"in_proj": "mamba"}} elif model_cls == MLA_Block: # Create simplified MLA based on DeepSeek-V3 architecture qk_nope_head_dim = 2 @@ -658,6 +680,7 @@ def _run_pattern_detection_job( v_head_dim=v_head_dim, bias=bias, ).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"q_a_proj": "mla"}} elif model_cls in (GDN_Block, GDN_Block_Unfused): num_k_heads_gdn = 4 num_v_heads_gdn = 8 @@ -671,10 +694,26 @@ def _run_pattern_detection_job( head_v_dim=head_v_dim, bias=bias, ).to(device="cuda", dtype=torch.float16) + predefined_config = {"tp_plan": {"in_proj_qkv": "delta"}} + elif model_cls == nn.Linear: + model = model_cls(num_features, num_features, bias=bias).to( + device="cuda", dtype=torch.float16 + ) + # update the tp_plan in predefined_config to force simple sharding of the single linear layer + predefined_config = {"tp_plan": {"*": "gather"}} else: model = model_cls(num_features, num_features, bias=bias).to( device="cuda", dtype=torch.float16 ) + predefined_config = { + "tp_plan": { + "gate_proj": "colwise", + "up_proj": "colwise", + "down_proj": "rowwise", + "linear1": "colwise", + "linear2": "rowwise", + } + } x = torch.randn(batch_size, sequence_len, num_features, device="cuda", dtype=torch.float16) # Test pattern detection - create expected transformations for validation @@ -746,7 +785,7 @@ def _run_pattern_detection_job( config=config, dist_op="all_gather", min_local_shape=1, - layer_type=LayerType.MLP, + layer_type=LayerType.UNKNOWN, ) ) elif model_cls == FP8MLP: @@ -931,6 +970,7 @@ def _run_pattern_detection_job( ) ) + sharding_source = "heuristic" if from_config else "manual" # get detected transformations optimizer = InferenceOptimizer( None, @@ -938,8 +978,9 @@ def _run_pattern_detection_job( "detect_sharding": { "stage": "sharding", "sharding_scope": ["tp"], - "tp_sharding_source": ["heuristic"], + "sharding_source": [sharding_source], "shard_all_unprocessed": True, + "manual_config": predefined_config, }, }, )