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

    • TensorParallelGroup

      public TensorParallelGroup(ParallelConfig localConfig, int worldSize, int rankOffset, byte[] uniqueId)
      A group whose ranks span several processes (and nodes): this process owns the world ranks rankOffset .. rankOffset + localConfig.tpSize() - 1, one per device of localConfig, and joins the NCCL communicator identified by uniqueId (the same 128 bytes in every process, from Native.ncclUniqueId()). Construction returns once every rank of the world has joined, so all processes must construct their groups concurrently. allReduceSumInPlace(int, Tensor) takes WORLD ranks.
      Parameters:
      localConfig - the devices of this process (one rank each).
      worldSize - ranks of the whole communicator.
      rankOffset - world rank of this process's first rank.
      uniqueId - the shared NCCL unique id.

  • Method Details

    • worldSize

      public int worldSize()
      Returns the ranks of the whole communicator (equals tpSize() unless it spans processes).
      Returns:
      ranks of the whole communicator (equals tpSize() unless it spans processes).

    • rankOffset

      public int rankOffset()
      Returns the world rank of this process's first local rank.
      Returns:
      world rank of this process's first local rank.

    • 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.