bridge.models.muse_spark_open.modeling_muse_spark_open.blocks

Missing Muse Spark Open operations composed with existing MCore building blocks.

Module Contents

Classes

InputScaledRMSNorm	Apply learned input scaling before native, parameter-free RMSNorm.
MuseSparkResidualGate	Post-normalize and mix a branch with its residual via learned gates.
MuseSparkSelfAttention	Reuse GQA and gating, inserting the missing per-head output norm.
MuseSparkRouter	Reuse sigmoid top-k routing with the model’s pre-sigmoid logit scale.
MuseSparkExperts	Normalize each expert’s output before applying its routing weight.
MuseSparkMoE	Reuse latent projections and dispatch with a pre-expert RMSNorm.

Functions

get_muse_spark_hybrid_stack_spec

Compose independent attention and MoE layers for a repeated *E stack.

API

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.InputScaledRMSNorm(

hidden_size: int,

*,

eps: float,

config: Any,

sequence_parallel: bool,

)

Bases: torch.nn.Module

Apply learned input scaling before native, parameter-free RMSNorm.

TE’s affine RMSNorm scales after normalization and cannot express this operation. Native RMSNorm supplies the normalization without an extra weight.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor

Normalize scaled activations using the existing framework kernel.

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkResidualGate(config: Any, *, input_scaled_norm: bool)

Bases: torch.nn.Module

Post-normalize and mix a branch with its residual via learned gates.

Implements the callable interface expected by TransformerLayer’s BDA slot, keeping the standard layer’s recompute, offload and checkpoint machinery.

Initialization

forward(training: bool, fused: bool) → Callable

Return the residual operation for MCore’s BDA call convention.

_apply_branch(

output_with_bias: tuple[torch.Tensor, torch.Tensor | None],

residual: torch.Tensor,

dropout: float,

) → torch.Tensor

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkSelfAttention(

config: Any,

*args: Any,

**kwargs: Any,

)

Bases: megatron.bridge.models.muse_glimmer.modeling_muse_glimmer.MuseSelfAttention

Reuse GQA and gating, inserting the missing per-head output norm.

Initialization

_apply_output_gate(

hidden_states: torch.Tensor,

gate: torch.Tensor,

) → torch.Tensor

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkRouter(

config: Any,

*args: Any,

**kwargs: Any,

)

Bases: megatron.core.transformer.moe.router.TopKRouter

Reuse sigmoid top-k routing with the model’s pre-sigmoid logit scale.

Initialization

_apply(

fn: Callable,

recurse: bool = True,

) → bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkRouter

Preserve the FP32 correction buffer across device and model-dtype moves.

gating(input: torch.Tensor) → torch.Tensor

Scale router logits before MCore performs expert selection.

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkExperts(

num_local_experts: int,

config: Any,

*args: Any,

**kwargs: Any,

)

Bases: megatron.core.transformer.moe.experts.TEGroupedMLP

Normalize each expert’s output before applying its routing weight.

Initialization

forward(

permuted_local_hidden_states: torch.Tensor,

tokens_per_expert: torch.Tensor,

permuted_probs: torch.Tensor,

) → tuple[torch.Tensor, None]

Keep grouped GEMMs while delaying routing weights until after normalization.

sharded_state_dict(

prefix: str = '',

sharded_offsets: tuple = (),

metadata: dict | None = None,

) → dict

Save the shared norm with ordinary TP replication, independently of ETP.

class bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.MuseSparkMoE(

config: Any,

*args: Any,

**kwargs: Any,

)

Bases: megatron.core.transformer.moe.moe_layer.MoELayer

Reuse latent projections and dispatch with a pre-expert RMSNorm.

Initialization

preprocess(

hidden_states: torch.Tensor,

probs: torch.Tensor,

routing_map: torch.Tensor,

padding_mask: torch.Tensor | None = None,

) → tuple[torch.Tensor, torch.Tensor]

Normalize the shared latent representation before dispatching tokens.

As in MoELayer.preprocess, only the flex dispatcher takes the batch-first padding mask.

bridge.models.muse_spark_open.modeling_muse_spark_open.blocks.get_muse_spark_hybrid_stack_spec(

config: Any,

) → megatron.core.transformer.spec_utils.ModuleSpec

Compose independent attention and MoE layers for a repeated *E stack.