mstar.model.components.norm

mstar.model.components.norm#

RMSNorm and AdaRMSNorm.

RMSNorm supports both the standard Llama-style normalization (normed * weight) and Gemma’s variant (normed * (1 + weight)) through the gemma_mode flag. Standard mode dispatches to a fused accelerator kernel when available; Gemma mode uses an fp32 computation that matches Hugging Face Gemma exactly.

AdaRMSNorm adds adaRMS conditioning (scale / shift / gate from a condition vector). Used by pi05’s action expert flow-matching path; the output is consumed by GatedDecoderLayer rather than a plain residual add.

Classes

AdaRMSNorm(hidden_size, cond_dim[, eps])

RMSNorm with adaRMS conditioning.

RMSNorm(hidden_size[, eps, gemma_mode])

RMSNorm with optional Gemma-style (1 + weight) scaling.

RMSNormGated(hidden_size[, eps])

RMSNorm scaled by a SiLU gate, for the delta-net family.
class mstar.model.components.norm.AdaRMSNorm(hidden_size, cond_dim, eps=1e-6)[source]#

Bases: Module

RMSNorm with adaRMS conditioning.

A per-norm nn.Linear(cond_dim, hidden_size*3) maps a shared condition vector to (scale, shift, gate). The normalization is rmsnorm(x) * (1 + scale) + shift and the gate is returned for the enclosing decoder layer to apply at the residual.

The dense.weight and dense.bias are zero-initialized so the norm starts as the identity (matches HF Gemma / lerobot openpi).

Parameters:
forward(x, cond)[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:
Return type:

tuple[Tensor, Tensor]

class mstar.model.components.norm.RMSNorm(hidden_size, eps=1e-6, gemma_mode=False)[source]#

Bases: Module

RMSNorm with optional Gemma-style (1 + weight) scaling.

Parameters:

  • hidden_size (int) – feature dimension to normalize over.

  • eps (float) – variance epsilon.

  • gemma_mode (bool) – if True, use (1 + weight) and a fp32 manual implementation (matches HF Gemma exactly; the loaded checkpoint weight is centered around zero, not one). If False, use weight and dispatch to the device-appropriate implementation.

extra_repr()[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

forward(hidden_states)[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:

hidden_states (Tensor)

Return type:

Tensor

forward_residual(hidden_states, residual, comm_group=None)[source]#

r = hidden_states + residual; returns (norm(r), r).

Keeps add and norm together so inductor fuses them, rather than folding the add into the producing GEMM as an addmm (an extra copy of residual). Under TP hidden_states is a row-parallel partial sum that comm_group reduces in the same fused kernel.

Parameters:
Return type:

tuple[Tensor, Tensor]

class mstar.model.components.norm.RMSNormGated(hidden_size, eps=1e-6)[source]#

Bases: Module

RMSNorm scaled by a SiLU gate, for the delta-net family.

Normalizes first, then gates (order matters). The weight is [head_v_dim], so this runs per head over [tokens, num_v_heads, head_v_dim]. Kept apart from RMSNorm because its casts follow HF’s for parity, not FlashInfer’s.

Parameters:
extra_repr()[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

forward(hidden_states, gate)[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:
Return type:

Tensor