mstar.model.components.mlp

mstar.model.components.mlp#

MLP and SwiGLU-style GatedMLP.

Used for transformer FFN blocks and for small projection MLPs (timestep embedders, Talker resize projections, etc.).

GatedMLP’s gate + up projections start as separate nn.Linear``s that match the HF checkpoint layout. After load, calling ``consolidate_gate_up_weight() concatenates them into a single gate_up_proj_weight buffer (one fused GEMM instead of two) and nulls out the originals. The forward branches on whether consolidation has happened.

Classes

FusedGatedMLP(hidden_size, intermediate_size)

SwiGLU-style gated MLP with the gate + up projections fused into a single FusedColumnLinear (one GEMM instead of two): down(act(gate) * up). Unlike GatedMLP (separate Linears + post-load consolidate_gate_up_weight), this fuses from construction and loads the separate gate_proj / up_proj checkpoint tensors straight into the fused parameter via the loader's stacked-param rules (gate is shard 0, up is shard 1). Use for models that fuse but don't need TP. :param hidden_size: input/output feature dim. :param intermediate_size: gate/up output and down input feature dim. :param activation: activation applied to the gate path. Either an HF string (silu / gelu / gelu_tanh) or a callable. :param bias: whether the linears have a bias term.

GatedMLP(hidden_size, intermediate_size[, ...])

SwiGLU-style gated MLP: down(act(gate(x)) * up(x)).

MLP(input_size, intermediate_size[, ...])

Plain two-layer MLP: out(act(in(x))).
class mstar.model.components.mlp.FusedGatedMLP(hidden_size, intermediate_size, activation='silu', bias=False)[source]#

Bases: Module

SwiGLU-style gated MLP with the gate + up projections fused into a single FusedColumnLinear (one GEMM instead of two): down(act(gate) * up). Unlike GatedMLP (separate Linears + post-load consolidate_gate_up_weight), this fuses from construction and loads the separate gate_proj / up_proj checkpoint tensors straight into the fused parameter via the loader’s stacked-param rules (gate is shard 0, up is shard 1). Use for models that fuse but don’t need TP. :param hidden_size: input/output feature dim. :param intermediate_size: gate/up output and down input feature dim. :param activation: activation applied to the gate path. Either an HF

string (silu / gelu / gelu_tanh) or a callable.

Parameters:

  • bias (bool) – whether the linears have a bias term.

  • hidden_size (int)

  • intermediate_size (int)

  • activation (str | Callable)

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.components.mlp.GatedMLP(hidden_size, intermediate_size, activation='silu', bias=False)[source]#

Bases: Module

SwiGLU-style gated MLP: down(act(gate(x)) * up(x)).

Parameters:

  • hidden_size (int) – input/output feature dim.

  • intermediate_size (int) – gate/up output and down input feature dim.

  • activation (str | Callable) – activation applied to the gate path. Either an HF string (silu / gelu / gelu_tanh) or a callable.

  • bias (bool) – whether the linears have a bias term.

consolidate_gate_up_weight()[source]#

Fuse gate_proj and up_proj weights into a single gate_up_proj_weight buffer and null out the originals. Idempotent; safe to call multiple times.

Return type:

None

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.components.mlp.MLP(input_size, intermediate_size, output_size=None, activation='silu', bias=True)[source]#

Bases: Module

Plain two-layer MLP: out(act(in(x))).

Used for small projection MLPs that aren’t gated (Talker resize projection, bagel timestep embedder MLP, etc.). For SwiGLU-style transformer FFNs, use GatedMLP.

Parameters:

  • input_size (int) – input feature dim.

  • intermediate_size (int) – hidden feature dim.

  • output_size (int | None) – output feature dim. Defaults to input_size if None.

  • activation (str | Callable) – activation between the two linears.

  • bias (bool) – whether the linears have a bias term.

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor