Module tensor_parallel 

Tensor parallelism: column-parallel and row-parallel linear layers — lazy-only.

Tensor parallelism (TP) splits weight matrices across devices so each device computes a partial result, then collective communication merges the outputs.

Sharding strategies

Strategy	Weight split	Communication
ColumnParallel	Split output dim	None (outputs are disjoint shards)
RowParallel	Split input dim	AllReduce after matmul

In a standard MLP (Y = XA):

• Column-parallel: A is split column-wise → each rank gets A_i (columns [i*cols/N..(i+1)*cols/N]) → Y_i = X @ A_i → results are concatenated.
• Row-parallel: A is split row-wise → each rank gets A_i (rows [i*rows/N..(i+1)*rows/N]) → Y_i = X_i @ A_i → results are summed via all-reduce.

Usage pattern

use fuel_parallel::tensor_parallel::{TensorParallelConfig, ShardDim};

let config = TensorParallelConfig::new(4)?; // 4-way TP
assert_eq!(config.world_size(), 4);

// Compute shard boundaries for a [4096, 4096] weight on rank 1
let (start, end) = config.shard_range(1, 4096, ShardDim::Column);
assert_eq!(start, 1024);
assert_eq!(end, 2048);

Structs

ColumnParallel

Column-parallel linear layer.

LayerParallelPlan

Describes how a model layer should be parallelized.

Linear

Minimal inlined Linear layer, lazy.

RowParallel

Row-parallel linear layer.

TensorParallelConfig

Configuration for tensor parallelism.

TensorShard

Metadata for one weight shard.

Enums

ShardDim

Which dimension to shard.