mstar.model.components.distributed.mlp

mstar.model.components.distributed.mlp#

TP-aware SwiGLU MLPs (parallel counterparts of mstar.model.components.GatedMLP).

ParallelGatedMLP fuses the gate/up projections into a single MergedColumnParallelLinear (sharded along the intermediate dim) with a RowParallelLinear down projection that all-reduces the partial sums. The checkpoint stores gate_proj.weight and up_proj.weight separately; the model’s weight loader calls self.gate_up_proj.weight.weight_loader(loaded_weight, loaded_shard_id=0) for gate and loaded_shard_id=1 for up.

ParallelGatedMLPUnfused keeps gate/up as separate ColumnParallelLinear projections so state_dict() keys match a checkpoint’s one-to-one — for loaders that stream weights by name with no stacked-parameter rules.

Classes

ParallelGatedMLP(hidden_size, intermediate_size)

SwiGLU-style gated MLP partitioned across TP ranks.

ParallelGatedMLPUnfused(hidden_size, ...[, ...])

SwiGLU MLP with separate (unfused) gate/up column shards.
class mstar.model.components.distributed.mlp.ParallelGatedMLP(hidden_size, intermediate_size, comm_group=None, activation='silu', bias=False)[source]#

Bases: Module

SwiGLU-style gated MLP partitioned across TP ranks.

Parameters:

  • comm_group (CommGroup | None) – TP comm group for this MLP’s parallel linears.

  • hidden_size (int) – model hidden dim (full, not per-partition).

  • intermediate_size (int) – SwiGLU intermediate dim (full).

  • activation (str | Callable) – HF activation name (silu, gelu, gelu_tanh).

  • 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

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

Bases: Module

SwiGLU MLP with separate (unfused) gate/up column shards.

Same math and sharding as ParallelGatedMLP, with gate_proj and up_proj kept as separate ColumnParallelLinear projections so the parameter names match a checkpoint’s gate_proj.weight / up_proj.weight one-to-one (plain name-matched loading, no stacked-parameter loader rules). A trivial comm group (world size 1) makes the projections plain linears.

Parameters:

  • hidden_size (int)

  • intermediate_size (int)

  • comm_group (CommGroup | None)

  • activation (str | Callable)

  • bias (bool)

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