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. In standard mode the FlashInfer fused kernel is used; Gemma mode falls back to a fp32 manual computation that matches HF 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.

class mstar.model.components.norm.AdaRMSNorm(hidden_size, cond_dim, eps=1e-6)

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:

• hidden_size (int)

• cond_dim (int)

• eps (float)

forward(x, cond)

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)

• cond (Tensor)

Return type:

tuple[Tensor, Tensor]

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

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 FlashInfer’s fused kernel.

extra_repr()

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)

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