Module tensor_parallel 

Tensor Parallelism

Implements tensor parallelism for splitting model weights and computations across multiple GPUs.

Weight Distribution

For a Linear layer with weight W of shape [out_features, in_features]:

• Column-parallel: Split along out_features, each GPU has W[:, start:end]
• Row-parallel: Split along in_features, each GPU has W[start:end, :]

Communication Patterns

• Column-parallel → Row-parallel: All-Reduce
• Row-parallel → Column-parallel: All-Gather

Structs

TensorParallelConfig

Tensor parallel configuration

TensorParallelGroup

Tensor parallel group for collective operations

TransformerParallelMapping

Tensor parallel layer mapping for transformer models

WeightShard

Weight sharding specification

Enums

LayerParallelType

Type of parallelism for a layer