Class RowParallelLinear

java.lang.Object

smile.llm.parallel.RowParallelLinear

public final class RowParallelLinear
extends Object

Megatron-style row-parallel linear: shards in_features across TP ranks. After the local matmul the caller must TensorParallelGroup.allReduceSumInPlace(int, Tensor) so outputs are replicated.

Supports dense LinearLayer or quantized LinearOp. Quantized weights must be sharded then packed before wrapping.

Constructor Summary

Constructors

Constructor	Description
RowParallelLinear(int globalInFeatures, int outFeatures, boolean bias, int tpSize, int tpRank)	Creates a dense row-parallel linear layer.
RowParallelLinear(LinearOp local, int globalInFeatures, int tpSize, int tpRank)	Wraps an already-sharded (and packed, if quantized) local linear op.

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor input)	Local matmul; caller must all-reduce for replicated outputs.
LinearLayer	linear()	Returns the dense linear layer when present.
LinearOp	linearOp()	Returns the underlying linear op.
int	localInFeatures()	Returns local input features (globalInFeatures / tpSize).
MemorySegment	module()	Returns the dense module handle.
int	tpRank()	Returns this TP rank.
int	tpSize()	Returns the TP world size.

Methods inherited from class Object

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

Constructor Details

RowParallelLinear

public RowParallelLinear(int globalInFeatures,
 int outFeatures,
 boolean bias,
 int tpSize,
 int tpRank)

Creates a dense row-parallel linear layer.

Parameters:

globalInFeatures - global input features (must be divisible by tpSize).

outFeatures - output features.

bias - whether to include bias.

tpSize - tensor-parallel world size.

tpRank - this rank.

RowParallelLinear

public RowParallelLinear(LinearOp local,
 int globalInFeatures,
 int tpSize,
 int tpRank)

Wraps an already-sharded (and packed, if quantized) local linear op.

Parameters:

local - local linear op.

globalInFeatures - global input features.

tpSize - tensor-parallel world size.

tpRank - this rank.

Method Details

linearOp

public LinearOp linearOp()

Returns the underlying linear op.

Returns:

linear op.

linear

public LinearLayer linear()

Returns the dense linear layer when present.

Returns:

dense LinearLayer.

module

public MemorySegment module()

Returns the dense module handle.

Returns:

module handle.

localInFeatures

public int localInFeatures()

Returns local input features (globalInFeatures / tpSize).

Returns:

local input features.

tpRank

public int tpRank()

Returns this TP rank.

Returns:

TP rank.

tpSize

public int tpSize()

Returns the TP world size.

Returns:

TP world size.

forward

public Tensor forward(Tensor input)

Local matmul; caller must all-reduce for replicated outputs.

Parameters:

input - activation tensor.

Returns:

local output.