Class WeightSharding

java.lang.Object

smile.llm.parallel.WeightSharding

public final class WeightSharding
extends Object

Slices full (unsharded) HuggingFace / Meta weights into the local TP shard.

Linear weights use PyTorch layout [out_features, in_features]. Column-parallel layers take a contiguous out-feature band; row-parallel layers take a contiguous in-feature band.

Method Summary

Modifier and Type	Method	Description
static Tensor	columnParallel(Tensor weight, int tpSize, int tpRank)	Column-parallel slice along dim 0 (output features).
static Tensor	columnParallelBias(Tensor bias, int tpSize, int tpRank)	Column-parallel bias [globalOut].
static Tensor	columnParallelHeads(Tensor weight, int numHeads, int headWidth, int tpSize, int tpRank)	Head-wise column shard for fused QKV / gate projections packed as [numHeads * headDim * pack, dim] where pack is 1 or 2.
static Tensor	rowParallel(Tensor weight, int tpSize, int tpRank)	Row-parallel slice along dim 1 (input features).
static Tensor	rowParallelHeads(Tensor weight, int numHeads, int headDim, int tpSize, int tpRank)	Head-wise row shard for output projections [dim, numHeads * headDim].

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Method Details

columnParallel

public static Tensor columnParallel(Tensor weight,
 int tpSize,
 int tpRank)

Column-parallel slice along dim 0 (output features).

Parameters:

weight - full weight [globalOut, in].

tpSize - TP size.

tpRank - this rank.

Returns:

a view/narrow of the local out-feature rows (caller should clone if needed).

columnParallelBias

public static Tensor columnParallelBias(Tensor bias,
 int tpSize,
 int tpRank)

Column-parallel bias [globalOut].

Parameters:

bias - full bias vector.

tpSize - TP size.

tpRank - this rank.

Returns:

local bias shard.

rowParallel

public static Tensor rowParallel(Tensor weight,
 int tpSize,
 int tpRank)

Row-parallel slice along dim 1 (input features).

Parameters:

weight - full weight [out, globalIn].

tpSize - TP size.

tpRank - this rank.

Returns:

local in-feature columns.

columnParallelHeads

public static Tensor columnParallelHeads(Tensor weight,
 int numHeads,
 int headWidth,
 int tpSize,
 int tpRank)

Head-wise column shard for fused QKV / gate projections packed as [numHeads * headDim * pack, dim] where pack is 1 or 2.

Parameters:

weight - full projection weight.

numHeads - global head count.

headWidth - channels per head (headDim or headDim * pack).

tpSize - TP size.

tpRank - this rank.

Returns:

local head rows.

rowParallelHeads

public static Tensor rowParallelHeads(Tensor weight,
 int numHeads,
 int headDim,
 int tpSize,
 int tpRank)

Head-wise row shard for output projections [dim, numHeads * headDim].

Parameters:

weight - full output-projection weight.

numHeads - global head count.

headDim - per-head dimension.

tpSize - TP size.

tpRank - this rank.

Returns:

local head columns.