Class TensorParallelGroup

java.lang.Object
smile.llm.parallel.TensorParallelGroup
All Implemented Interfaces:
AutoCloseable
public final class TensorParallelGroup extends Object implements AutoCloseable
Single-process tensor-parallel group: one ParallelState and CUDA device per TP rank, plus NCCL (or peer-copy fallback) collectives.

When NCCL is available, each rank calls allReduceSumInPlace(int, Tensor) on its own tensor without a host barrier gather. The communicator is retained for future pipeline-parallel send/recv on the same process group.

  • Constructor Details

    • TensorParallelGroup

      public TensorParallelGroup(ParallelConfig config)
      Creates a TP group for the given mesh configuration.
      Parameters:
      config - process-mesh configuration (ppSize must be 1 in phase 1).

  • Method Details

    • config

      public ParallelConfig config()
      Returns the process-mesh configuration.
      Returns:
      parallel config.

    • mesh

      public DeviceMesh mesh()
      Returns the device mesh for this group.
      Returns:
      device mesh.

    • tpSize

      public int tpSize()
      Returns the tensor-parallel size.
      Returns:
      TP size.

    • hasNccl

      public boolean hasNccl()
      Returns whether NCCL backs collectives for this group.
      Returns:
      true when an NCCL communicator was created.

    • state

      public ParallelState state(int tpRank)
      Returns the parallel state for a TP rank.
      Parameters:
      tpRank - tensor-parallel rank.
      Returns:
      per-rank state.

    • device

      public Device device(int tpRank)
      Returns the CUDA device for a TP rank.
      Parameters:
      tpRank - tensor-parallel rank.
      Returns:
      device for that rank.

    • allReduceSumInPlace

      public void allReduceSumInPlace(int tpRank, Tensor local)
      In-place sum all-reduce. Every TP rank's worker thread must call this with its local tensor (same shape/dtype). No-op when tpSize == 1.
      Parameters:
      tpRank - calling rank.
      local - local tensor to reduce in place.

    • broadcast

      public void broadcast(Tensor[] locals, int root)
      Broadcasts locals[root] to every other slot. Caller must supply the full array (one tensor per rank) from a single coordinating thread, or use per-rank buffers that are already filled.
      Parameters:
      locals - one tensor per TP rank.
      root - source rank whose tensor is copied to the others.

    • close

      public void close()
      Specified by:
      close in interface AutoCloseable

    • isClosed

      public boolean isClosed()
      Returns whether close() has been called.
      Returns:
      true if closed.