Class FeedForward

java.lang.Object

smile.llm.transformer.FeedForward

public class FeedForward
extends Object

Feedforward layer in Transformer. It has two linear transformations and an intermediate SiLU activation function.

Under tensor parallelism, w1/w3 are column-parallel and w2 is row-parallel (all-reduce after the down-projection).

Constructor Summary

Constructors

Constructor	Description
FeedForward(int dim, int intermediateSize)	Constructor with an explicit intermediate size, as provided by HuggingFace config.json.
FeedForward(int dim, int hiddenDim, int multipleOf, Double ffnDimMultiplier)	Constructor.
FeedForward(int dim, int globalIntermediate, TensorShardSpec shard, TensorParallelGroup tpGroup)	Tensor-parallel constructor.

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor x)	Feed forward.
MemorySegment	module()	Returns the PyTorch module handle.
void	replaceLinears(LinearOp w1, LinearOp w2, LinearOp w3)	Replaces FFN linears with quantized ops (already sharded/packed for TP).

Methods inherited from class Object

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

Constructor Details

FeedForward

public FeedForward(int dim,
 int intermediateSize)

Constructor with an explicit intermediate size, as provided by HuggingFace config.json.

Parameters:

dim - the dimension of input tensor.

intermediateSize - the FFN hidden dimension size, used directly without further computation.

FeedForward

public FeedForward(int dim,
 int globalIntermediate,
 TensorShardSpec shard,
 TensorParallelGroup tpGroup)

Tensor-parallel constructor.

Parameters:

dim - model dimension.

globalIntermediate - full (unsharded) intermediate size.

shard - local shard sizes (intermediateSize already divided).

tpGroup - TP group for row-parallel all-reduce; null when tp=1.

FeedForward

public FeedForward(int dim,
 int hiddenDim,
 int multipleOf,
 Double ffnDimMultiplier)

Constructor.

Parameters:

dim - the dimension of input tensor.

hiddenDim - the dimension of hidden layer. First, hiddenDim is set to two-thirds of the provided hiddenDim value. If ffnDimMultiplier is provided, hiddenDim is further multiplied by this value. The hiddenDim is then adjusted to ensure it is a multiple of multipleOf.

multipleOf - make SwiGLU hidden layer size multiple of large power of 2.

ffnDimMultiplier - the multiplier for the hidden dimension of the feedforward layers.

Method Details

replaceLinears

public void replaceLinears(LinearOp w1,
 LinearOp w2,
 LinearOp w3)

Replaces FFN linears with quantized ops (already sharded/packed for TP). Frees dense LinearLayer shells so empty GPU weights are released.

Parameters:

w1 - gate / up projection op.

w2 - down projection op.

w3 - up / parallel projection op.

module

public MemorySegment module()

Returns the PyTorch module handle.

Returns:

module handle.

forward

public Tensor forward(Tensor x)

Feed forward.

Parameters:

x - the input tensor.

Returns:

the output tensor.