Package smile.llm.parallel

package smile.llm.parallel
Tensor / pipeline / data-parallel primitives for LLM inference.

Phase 1 provides in-process tensor parallelism (TensorParallelGroup, column/row-parallel linears). Pipeline axes on ParallelConfig / ParallelState are reserved for multi-node PP.

  • Class
    Description
    ColumnParallelLinear
    Megatron-style column-parallel linear: shards out_features across TP ranks.
    DeviceMesh
    Maps logical mesh coordinates to CUDA devices (phase 1) or process ranks (phase 2 multi-node).
    ParallelConfig
    Process-mesh configuration for tensor / pipeline / data parallelism.
    ParallelState
    Per-rank view of a ParallelConfig mesh.
    RowParallelLinear
    Megatron-style row-parallel linear: shards in_features across TP ranks.
    TensorParallelGroup
    Single-process tensor-parallel group: one ParallelState and CUDA device per TP rank, plus NCCL (or peer-copy fallback) collectives.
    TensorShardSpec
    Local shard sizes for a dense decoder under tensor parallelism.
    WeightSharding
    Slices full (unsharded) HuggingFace / Meta weights into the local TP shard.