diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 374d7810..95dec62e 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -237,12 +237,23 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: _apply_nemotron_h_generate_patch(model) model_device = _get_model_device(model, device) + gen_ids = torch.from_numpy(input_ids).to(model_device) + max_new = case.generation_params.get("max_new_tokens", 20) with torch.no_grad(): - gen_output = model.generate( - torch.from_numpy(input_ids).to(model_device), - max_new_tokens=case.generation_params.get("max_new_tokens", 20), - do_sample=False, - ) + try: + gen_output = model.generate(gen_ids, max_new_tokens=max_new, do_sample=False) + except ValueError as e: + # All-attention GraniteMoeHybrid variants (e.g. granite-4.0-1b) + # trip transformers' hybrid Mamba/attention generation cache, + # which assumes at least one linear-attention (Mamba) layer: + # "`has_previous_state` can only be called on LinearAttention + # layers". Greedy output is cache-independent, so fall back to + # the (slower) cache-free path. + if "has_previous_state" not in str(e): + raise + gen_output = model.generate( + gen_ids, max_new_tokens=max_new, do_sample=False, use_cache=False + ) generated_ids = gen_output[0, seq_len:].cpu().numpy() save_golden_ref( diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 2ecbc655..475b69b9 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -2226,10 +2226,39 @@ class GraniteMoeHybridConfig(BambaConfig): @classmethod def from_transformers(cls, config, parent_config=None) -> GraniteMoeHybridConfig: - # Reuse BambaConfig.from_transformers for mamba fields + layer_types conversion - # (converts HF "mamba"→"mamba2" and "attention"→"full_attention") + # Reuse BambaConfig.from_transformers for mamba fields, MoE/RoPE/multiplier + # extraction, then rebuild layer_types from GraniteMoeHybrid's own naming. bamba = BambaConfig.from_transformers(config, parent_config) bamba_fields = _shallow_fields(bamba) + + # GraniteMoeHybrid names layers "full_attention" / "linear_attention" + # (linear_attention == Mamba2/SSD), unlike Bamba's "attention" / "mamba". + raw_layer_types = ( + getattr(config, "layer_types", None) + or getattr(config, "layers_block_type", None) + or [] + ) + _attn = {"full_attention", "attention"} + _mamba = {"linear_attention", "mamba", "mamba2"} + layer_types: list[str] = [] + for ltype in raw_layer_types: + if ltype in _attn: + layer_types.append("full_attention") + elif ltype in _mamba: + layer_types.append("mamba2") + else: + raise ValueError(f"Unknown GraniteMoeHybrid layer type: {ltype!r}") + if layer_types: + bamba_fields["layer_types"] = layer_types + + # Respect position_embedding_type: GraniteMoeHybrid checkpoints ship + # default ``rope_parameters`` even for the NoPE variant + # (granite-4.0-tiny-preview: position_embedding_type='nope'). Only apply + # RoPE when explicitly requested; otherwise disable it so + # ``initialize_rope`` returns None and attention runs NoPE. + if getattr(config, "position_embedding_type", "rope") != "rope": + bamba_fields["rope_type"] = None + return cls( **bamba_fields, shared_intermediate_size=getattr(config, "shared_intermediate_size", 1024), diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index edcaa5e0..095e3211 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -334,7 +334,18 @@ def torch_forward( ) kwargs["past_key_values"] = cache - outputs = model(**kwargs) + try: + outputs = model(**kwargs) + except ValueError as e: + # All-attention hybrid models (e.g. GraniteMoeHybrid granite-4.0-1b) + # trip transformers' recurrent-mask builder, which assumes the hybrid + # cache contains a linear-attention (Mamba) layer: "`has_previous_state` + # can only be called on LinearAttention layers". A single-pass forward's + # logits don't depend on caching, so retry without a cache. + if "has_previous_state" not in str(e) or "past_key_values" in kwargs: + raise + kwargs["use_cache"] = False + outputs = model(**kwargs) logits = outputs.logits.cpu().numpy() # Extract KV cache if available (Mamba models don't have it) diff --git a/src/mobius/models/granitemoehybrid.py b/src/mobius/models/granitemoehybrid.py index 1a244930..c278b858 100644 --- a/src/mobius/models/granitemoehybrid.py +++ b/src/mobius/models/granitemoehybrid.py @@ -1,22 +1,26 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""GraniteMoeHybrid: Mamba2/SSD + Attention hybrid with MoE FFN on all layers. +"""GraniteMoeHybrid: Mamba2/SSD + Attention hybrid with optional MoE FFN. -Every layer has both a routed MoE block (``block_sparse_moe``) and a dense -shared MLP (``shared_mlp``). The layer type ("mamba2" or "full_attention") -controls whether the attention sub-block is a Mamba2/SSD or standard GQA. -Attention layers use NoPE (no rotary position embeddings). +Each layer has a dense shared MLP (``shared_mlp``) and, when the config has +routed experts (``num_local_experts > 0``), a routed MoE block +(``block_sparse_moe``); variants with no experts (e.g. granite-4.0-1b) run +only the shared MLP. The layer type ("mamba2" or "full_attention") controls +whether the token-mixing sub-block is a Mamba2/SSD or standard GQA. Attention +layers use RoPE when ``position_embedding_type == 'rope'`` (granite-4.0-1b) +and NoPE otherwise (granite-4.0-tiny-preview). Granite scaling multipliers +(embedding/attention/residual/logits) are applied throughout. Forward pass per layer:: residual = x x = input_layernorm(x) - x = mamba(x) OR self_attn(x, position_embeddings=None) + x = mamba(x) OR self_attn(x, position_embeddings) x = residual + x * residual_multiplier residual = x x = post_attention_layernorm(x) - x = block_sparse_moe(x) + shared_mlp(x) + x = shared_mlp(x) [+ block_sparse_moe(x) if experts] x = residual + x * residual_multiplier HuggingFace reference: ``GraniteMoeHybridForCausalLM``. @@ -39,6 +43,7 @@ TopKGate, create_attention_bias, get_activation, + initialize_rope, ) if TYPE_CHECKING: @@ -183,6 +188,24 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: # --------------------------------------------------------------------------- +def _feedforward( + op: OpBuilder, + hidden_states: ir.Value, + block_sparse_moe: nn.Module | None, + shared_mlp: nn.Module, +) -> ir.Value: + """Combined routed-MoE + shared-MLP feedforward. + + When routed experts exist, the output is ``moe(x) + shared_mlp(x)``; for + variants with ``num_local_experts == 0`` (e.g. granite-4.0-1b) only the + dense shared MLP runs. Mirrors HF ``GraniteMoeHybridDecoderLayer``. + """ + shared_out = shared_mlp(op, hidden_states) + if block_sparse_moe is None: + return shared_out + return op.Add(block_sparse_moe(op, hidden_states), shared_out) + + class _GraniteMoeHybridMambaDecoderLayer(nn.Module): """GraniteMoeHybrid Mamba2 layer. @@ -212,8 +235,12 @@ def __init__(self, config: GraniteMoeHybridConfig): ) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # Routed MoE with fused 3D expert weights - self.block_sparse_moe = _FusedMoEBlock(config) + # Routed MoE with fused 3D expert weights. GraniteMoeHybrid variants with + # ``num_local_experts == 0`` (e.g. granite-4.0-1b) have no routed experts; + # only the dense shared MLP runs. Mirrors HF ``block_sparse_moe = MoE(...) + # if num_local_experts > 0 else None``. + self._has_experts = bool(config.num_local_experts) + self.block_sparse_moe = _FusedMoEBlock(config) if self._has_experts else None # Dense shared MLP with fused gate+up weight self.shared_mlp = _FusedSharedMLP(config) @@ -225,7 +252,7 @@ def forward( op: OpBuilder, hidden_states: ir.Value, attention_bias: ir.Value, - position_embeddings: tuple, + position_embeddings: tuple | None, past_key_value: tuple | None, ): """Forward pass. Returns (hidden_states, (conv_state, ssm_state)). @@ -250,11 +277,7 @@ def forward( # MoE + shared-MLP path with pre-norm residual = hidden_states hidden_states = self.post_attention_layernorm(op, hidden_states) - # Both routed MoE and shared MLP run on every layer; outputs are summed - hidden_states = op.Add( - self.block_sparse_moe(op, hidden_states), - self.shared_mlp(op, hidden_states), - ) + hidden_states = _feedforward(op, hidden_states, self.block_sparse_moe, self.shared_mlp) rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states) hidden_states = op.Add(residual, op.Mul(hidden_states, rm)) @@ -274,12 +297,16 @@ class _GraniteMoeHybridAttentionDecoderLayer(nn.Module): def __init__(self, config: GraniteMoeHybridConfig): super().__init__() self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # GraniteMoeHybrid attention uses NoPE (no position embeddings) - self.self_attn = Attention(config) + # Attention scale is Granite's ``attention_multiplier`` (a fixed value, + # not the default 1/sqrt(head_dim)); RoPE is applied only when the text + # model supplies ``position_embeddings`` (position_embedding_type='rope'). + self.self_attn = Attention(config, scale=config.attention_multiplier) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # Routed MoE with fused 3D expert weights - self.block_sparse_moe = _FusedMoEBlock(config) + # Routed MoE with fused 3D expert weights. Absent when there are no + # routed experts (num_local_experts == 0), e.g. granite-4.0-1b. + self._has_experts = bool(config.num_local_experts) + self.block_sparse_moe = _FusedMoEBlock(config) if self._has_experts else None # Dense shared MLP with fused gate+up weight self.shared_mlp = _FusedSharedMLP(config) @@ -291,20 +318,22 @@ def forward( op: OpBuilder, hidden_states: ir.Value, attention_bias: ir.Value, - position_embeddings: tuple, + position_embeddings: tuple | None, past_key_value: tuple | None, ): - """Forward pass. Returns (hidden_states, (key, value)).""" - del position_embeddings # GraniteMoeHybrid uses NoPE: no rotary embeddings + """Forward pass. Returns (hidden_states, (key, value)). - # GQA attention path (no RoPE) + ``position_embeddings`` is ``None`` for NoPE variants and a + ``(cos, sin)`` tuple when ``position_embedding_type == 'rope'``. + """ + # GQA attention path (RoPE applied iff position_embeddings is not None) residual = hidden_states hidden_states = self.input_layernorm(op, hidden_states) attn_out, present_kv = self.self_attn( op, hidden_states=hidden_states, attention_bias=attention_bias, - position_embeddings=None, # NoPE: skip rotary embedding application + position_embeddings=position_embeddings, past_key_value=past_key_value, ) rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), attn_out) @@ -313,10 +342,7 @@ def forward( # MoE + shared-MLP path with pre-norm residual = hidden_states hidden_states = self.post_attention_layernorm(op, hidden_states) - hidden_states = op.Add( - self.block_sparse_moe(op, hidden_states), - self.shared_mlp(op, hidden_states), - ) + hidden_states = _feedforward(op, hidden_states, self.block_sparse_moe, self.shared_mlp) rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states) hidden_states = op.Add(residual, op.Mul(hidden_states, rm)) @@ -332,7 +358,9 @@ class _GraniteMoeHybridTextModel(nn.Module): """GraniteMoeHybrid text backbone: embedding -> N x (Mamba2|Attention) layers -> norm. Layer type ("mamba2" or "full_attention") is read from ``config.layer_types``. - No rotary embeddings are used (NoPE for attention layers). + Attention layers use RoPE when ``config`` declares rotary parameters + (``position_embedding_type == 'rope'``, e.g. granite-4.0-1b) and NoPE + otherwise (e.g. granite-4.0-tiny-preview). """ def __init__(self, config: GraniteMoeHybridConfig): @@ -352,7 +380,10 @@ def __init__(self, config: GraniteMoeHybridConfig): self.layers.append(_GraniteMoeHybridAttentionDecoderLayer(config)) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - # No rotary_emb: GraniteMoeHybrid uses NoPE (no positional encodings) + # RoPE when the config declares rotary params; None => NoPE (attention + # layers then receive position_embeddings=None). + self.rotary_emb = initialize_rope(config) + self.embedding_multiplier = config.embedding_multiplier def forward( self, @@ -362,10 +393,15 @@ def forward( position_ids: ir.Value, past_key_values: list | None = None, ): - del position_ids # unused: NoPE architecture has no positional embeddings - # (batch, seq, hidden) hidden_states = self.embed_tokens(op, input_ids) + # Granite scales embeddings by embedding_multiplier after lookup. + hidden_states = op.Mul(hidden_states, self.embedding_multiplier) + + # (cos, sin) tuple for RoPE, or None for NoPE variants. + position_embeddings = ( + self.rotary_emb(op, position_ids) if self.rotary_emb is not None else None + ) attention_bias = create_attention_bias( op, input_ids=input_ids, @@ -380,7 +416,7 @@ def forward( op, hidden_states=hidden_states, attention_bias=attention_bias, - position_embeddings=None, # NoPE: no RoPE + position_embeddings=position_embeddings, past_key_value=past_kv, ) present_key_values.append(present_kv) @@ -413,6 +449,7 @@ def __init__(self, config: GraniteMoeHybridConfig): self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight + self.logits_scaling = config.logits_scaling def forward( self, @@ -430,6 +467,8 @@ def forward( past_key_values=past_key_values, ) logits = self.lm_head(op, hidden_states) + # Granite divides final logits by logits_scaling. + logits = op.Div(logits, self.logits_scaling) return logits, present_key_values def preprocess_weights( diff --git a/testdata/cases/causal-lm/granite-4_0-1b.yaml b/testdata/cases/causal-lm/granite-4_0-1b.yaml new file mode 100644 index 00000000..a46fc75f --- /dev/null +++ b/testdata/cases/causal-lm/granite-4_0-1b.yaml @@ -0,0 +1,23 @@ +model_id: "ibm-granite/granite-4.0-1b" +model_type: "granitemoehybrid" +revision: "main" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +notes: >- + GraniteMoeHybrid dense variant: all full_attention layers (no Mamba), RoPE + (position_embedding_type='rope'), and no routed experts (num_local_experts=0) + so only the shared MLP runs. Exercises the Granite scaling multipliers + (embedding/attention/residual/logits). Unlike granite-4.0-tiny-preview this + runs natively in transformers (no trust_remote_code, no mamba-ssm), so a + golden can be generated in CI. diff --git a/testdata/golden/causal-lm/granite-4_0-1b.json b/testdata/golden/causal-lm/granite-4_0-1b.json new file mode 100644 index 00000000..e7c97fd2 --- /dev/null +++ b/testdata/golden/causal-lm/granite-4_0-1b.json @@ -0,0 +1,41 @@ +{ + "top1_id": 330, + "top2_id": 4815, + "top10_ids": [ + 330, + 4815, + 720, + 57277, + 578, + 763, + 3146, + 2355, + 1035, + 362 + ], + "top10_logits": [ + "0x1.2c68300000000p+4", + "0x1.2894c40000000p+4", + "0x1.2880740000000p+4", + "0x1.215f900000000p+4", + "0x1.1de6be0000000p+4", + "0x1.1bb2a80000000p+4", + "0x1.1aa89c0000000p+4", + "0x1.17b8080000000p+4", + "0x1.12ce680000000p+4", + "0x1.0c2f880000000p+4" + ], + "logits_summary": [ + "0x1.2c68300000000p+4", + "-0x1.80fc560000000p+3", + "-0x1.01df734ac7e0ap-2", + "0x1.f3583dce9b92bp+1" + ], + "input_ids": [ + 8586, + 374, + 856, + 33894, + 25 + ] +} diff --git a/testdata/golden/causal-lm/granite-4_0-1b_generation.json b/testdata/golden/causal-lm/granite-4_0-1b_generation.json new file mode 100644 index 00000000..8564b1b6 --- /dev/null +++ b/testdata/golden/causal-lm/granite-4_0-1b_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "ibm-granite/granite-4.0-1b", + "prompt": "Here is my poem:", + "generated_tokens": [ + 330, + 791, + 8219, + 596, + 35969, + 685, + 1, + 555, + 35266, + 98337, + 1473, + 791, + 8219, + 596, + 35969, + 685, + 198, + 32, + 21411, + 37466 + ], + "generated_text": " \"The Sun's Radiance\" by Emily Dickinson:\n\nThe Sun's Radiance\nA golden orb" +}