[TRTLLM-11614][feat] Fixing multigpu tests

examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml

@@ -16,3 +16,23 @@ model_kwargs:
   #   num_hidden_layers: 6
   # vision_config:
   #   depth: 2
+transforms:
+  detect_sharding:
+    sharding_dims: ['tp','ep', 'bmm']
+    # use only manual config for TP sharding
+    sharding_source: ['manual']
+    manual_config:
+      tp_plan:
+        # GDN layer
+        "in_proj_qkv": "delta"
+        # attention layer
+        "q_proj": "colwise"
+        "k_proj": "colwise"
+        "v_proj": "colwise"
+        "o_proj": "rowwise"
+        # replicating shared experts (keep them commented out)
+        # "shared_expert_gate_proj": "colwise"
+        # "shared_expert_up_proj": "colwise"
+        # "shared_expert_down_proj": "rowwise"
+        # gating layer should be replicated as well
+        # "gate": "gather"

examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml

@@ -16,3 +16,22 @@ model_kwargs:
 transforms:
   export_to_gm:
     num_moe_experts_for_export: 2
+  detect_sharding:
+    sharding_dims: ['tp','ep', 'bmm']
+    # use only manual config for TP sharding
+    sharding_source: ['manual']
+    manual_config:
+      tp_plan:
+        # GDN layer
+        "in_proj_qkv": "delta"
+        # attention layer
+        "q_proj": "colwise"
+        "k_proj": "colwise"
+        "v_proj": "colwise"
+        "o_proj": "rowwise"
+        # replicating shared experts (keep them commented out)
+        # "shared_expert_gate_proj": "colwise"
+        # "shared_expert_up_proj": "colwise"
+        # "shared_expert_down_proj": "rowwise"
+        # gating layer should be replicated as well
+        # "gate": "gather"

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
 
@@ -85,16 +86,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()
 
@@ -129,32 +130,23 @@ def _set_distributed_env_vars(local_rank: int, world_size: int, port: int) -> No
     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(("127.0.0.1", 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 +155,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 +181,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 +243,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,22 +302,33 @@ 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:
             kwargs["input_queue"] = input_queues[rank]
         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 +345,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):

tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py

@@ -239,7 +239,9 @@ def validate_config(self, sources: Union[ShardingSource, List[ShardingSource]] =
             supported_modes = {
                 "colwise",  # row split and no collective
                 "rowwise",  # column split and all-reduce
+                "mla",  # mla layer
                 "mamba",  # mamba SSM layer
+                "delta",  # gated delta net layer
                 "gather",  # simple shard (row + all_gather)
                 # TODO: remaining values are not supported yet.
                 # They require hybrid EP+TP and/or SP support.
@@ -2663,6 +2665,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 +2728,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 +2794,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)

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):