From ed430e2ea69616b750413c19e27317867c55c301 Mon Sep 17 00:00:00 2001 From: Lipika Sreedharan Date: Wed, 5 Aug 2026 00:24:37 +0530 Subject: [PATCH] Add Granite 4.1 20B support Signed-off-by: Lipika Sreedharan --- .../cache_config/hf_models_and_datasets.yaml | 1 + ARCHITECTURE.md | 39 ++-- hf_adapters.md | 34 +++- hf_adapters/auto_spyre_model.py | 112 +---------- hf_adapters/hf_granite_swa.py | 184 ++++++++++++++++++ tests/model_registry.py | 139 ++++--------- .../models/granite41-20b/run_huggingface.py | 76 ++++++++ 7 files changed, 346 insertions(+), 239 deletions(-) create mode 100644 hf_adapters/hf_granite_swa.py create mode 100644 utils/model_ops/models/granite41-20b/run_huggingface.py diff --git a/.github/cache_config/hf_models_and_datasets.yaml b/.github/cache_config/hf_models_and_datasets.yaml index ea67040d..1837265c 100644 --- a/.github/cache_config/hf_models_and_datasets.yaml +++ b/.github/cache_config/hf_models_and_datasets.yaml @@ -4,6 +4,7 @@ models: - ibm-granite/granite-3.3-2b-instruct # granite2b - ibm-granite/granite-4.0-1b-base # granite4 - ibm-granite/granite-4.0-1b # granite4_instruct + - ibm-research/granite-4.1-20b # granite41_20b # Embedding models - BAAI/bge-base-en-v1.5 # bge_base (embed-compare) - BAAI/bge-m3 # bge_m3 (load) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 54fc2351..3f477519 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -14,6 +14,7 @@ which models are supported on Spyre. | Granite 3.3 2B | granite | 64→128 | 64 | Yes (padded) | Yes | Yes | Yes | | Granite 4.0 1B Base | granitemoehybrid | 128 | 64 | Yes | Yes | Yes | Yes | | Granite 4.0 1B Instruct | granitemoehybrid | 128 | 64 | Yes | Yes | Yes | Yes | +| Granite 4.1 20B (bf16) | granite_swa | 128 | 64 | Yes | Yes | Yes | Yes | | SmolLM3 3B | smollm3 | 128 | 64 | Yes | Yes | Yes | Yes | | Llama 3.2 3B | llama | 128 | 64 | Yes | Yes | Yes | Yes | | TinyLlama 1.1B | llama | 64→128 | 64 | Yes (padded) | Yes | Yes | Yes | @@ -29,8 +30,6 @@ which models are supported on Spyre. | Yi 1.5 6B | llama | 128 | 64 | Yes | Yes | Yes | Yes | | Granite Vision 4.1 4B (text backbone) | granite (text) | 64→128 | 64 | Yes (padded) | Yes | Yes | Yes | | Gemma 4 12B | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes | -| Gemma 4 12B Base | gemma4\_unified | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes | -| Gemma 4 31B | gemma4 | 256 / 512 | 128 / 256 | Yes | Yes | Yes | Yes | | Gemma 3 1B | gemma3\_text | 256 | 128 | Yes | Yes | Yes | Yes | | GPT-2 124M | gpt2 | 64 | n/a (no RoPE) | Yes | Yes | Yes | Yes | | GPT-Neo 125M | gpt_neo | 64 | n/a (no RoPE) | Yes | Yes | Yes | Yes | @@ -112,8 +111,8 @@ single-token decode path (seq_len=1), not an adapter issue. > adapter or verify a checkpoint, update *only* this file (and the badge > counts in README.md, noted below). -**Coverage:** 27 adapters · 46 verified checkpoints · 100+ compatible models. -The 46 verified rows are 29 generative + 13 embedding + 4 vision-language (see the +**Coverage:** 28 adapters · 45 verified checkpoints · 100+ compatible models. +The 45 verified rows are 28 generative + 13 embedding + 4 vision-language (see the Verified Checkpoints tables above). `hf_siglip_vision` and `hf_pixtral_vision` are vision-tower components used by VLM adapters rather than standalone model adapters. Granite Vision 4.1 is verified both as a text backbone (generative) and as a full VLM. @@ -140,8 +139,9 @@ pattern, norms, and weight layout. | hf\_ministral.py | ministral | 1 | Ministral-8B Instruct fine-tunes | | hf\_phi3.py | phi3 | 1 | Phi-3 mini 4k/128k, Phi-3 small 8k | | hf\_granitemoehybrid.py | granitemoehybrid | 2 | Granite 4.0 Micro | +| hf\_granite\_swa.py | granite\_swa | 1 | Granite 4.1 8B (unverified), Granite 4.1 20B | | hf\_smollm3.py | smollm3 | 1 | — | -| hf\_gemma4.py | gemma4\_unified / gemma4 (dense) | 3 | Not E2B/E4B (PLE) or 26B-A4B (MoE). | +| hf\_gemma4.py | gemma4\_unified / gemma4 (dense) | 1 | Gemma 4 31B (dense). Not E2B/E4B (PLE) or 26B-A4B (MoE). | | hf\_gemma4\_mm.py | gemma4\_unified (multimodal) | 1 | Gemma 4 31B (dense unified VLM). Not E2B/E4B (PLE) or 26B-A4B (MoE). | | hf\_gemma3.py | gemma3\_text / gemma3 (dense) | 2 | Gemma 3 4B/12B/27B (text decoder of the multimodal checkpoints); EmbeddingGemma (bidirectional embedder). Not Gemma 3n (PLE). | | hf\_olmo.py | olmo | 1 | OLMo 7B | @@ -392,20 +392,21 @@ modification: ### Model-Specific Adaptations -| Feature | Granite 3.3 | Granite Vision 4.1 | Qwen3 | Granite 4.0 | SmolLM3 | Llama | Qwen2 | Mistral | Phi-4 mini | OLMo | OLMo2 | Gemma 3 | Gemma 4 | -|---------|------------|-------------------|-------|-------------|---------|-------|-------|---------|-----------|------|-------|---------|---------| -| Embedding multiplier | Yes | Yes | No | Yes | No | No | No | No | No | No | No | Yes | Yes | -| Residual multiplier | Yes | Yes | No | Yes | No | No | No | No | No | No | No | No | No | -| Logits scaling | Yes | Yes | No | Yes | No | No | No | No | No | No | No | No | No | -| Q/K RMSNorm | No | No | Yes (per-head) | No | No | No | No | No | No | No | Yes (flattened) | Yes (per-head Q/K) | Yes (per-head Q/K/V) | -| Fused QKV split | No | No | No | No | No | No | No | No | Yes | No | No | No | No | -| Fused MLP split | No | No | No | Yes | No | No | No | No | Yes | No | No | No | No | -| NoPE layers | No | No | No | No | Yes | No | No | No | No | No | No | No | No | -| Partial RoPE | No | No | No | No | No | No | No | No | Yes | No | No | No | Yes (global layers) | -| Head-dim padding | 2B only | Yes (64→128) | No | Micro only (64→128) | No | TinyLlama | No | No | No | No | No | No | No | -| Custom model loading | No | Yes (safetensor remap) | No | No | No | No | No | No | No | No | No | No | No | -| Attention scaling | `config.attention_multiplier` | `config.attention_multiplier` | `head_dim**-0.5` | `config.attention_multiplier` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `query_pre_attn_scalar**-0.5` | `1.0` (unscaled) | -| Norm type | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | LayerNorm (pre, no weight) | RMSNorm (post) | RMSNorm (sandwich) | RMSNorm (sandwich) | +| Feature | Granite 3.3 | Granite Vision 4.1 | Qwen3 | Granite 4.0 | Granite 4.1 SWA | SmolLM3 | Llama | Qwen2 | Mistral | Phi-4 mini | OLMo | OLMo2 | Gemma 3 | Gemma 4 | +|---------|------------|-------------------|-------|-------------|-----------------|---------|-------|-------|---------|-----------|------|-------|---------|---------| +| Embedding multiplier | Yes | Yes | No | Yes | Yes | No | No | No | No | No | No | No | Yes | Yes | +| Residual multiplier | Yes | Yes | No | Yes | Yes | No | No | No | No | No | No | No | No | No | +| Logits scaling | Yes | Yes | No | Yes | Yes | No | No | No | No | No | No | No | No | No | +| Q/K RMSNorm | No | No | Yes (per-head) | No | No | No | No | No | No | No | No | Yes (flattened) | Yes (per-head Q/K) | Yes (per-head Q/K/V) | +| Fused QKV split | No | No | No | No | No | No | No | No | No | Yes | No | No | No | No | +| Fused MLP split | No | No | No | Yes | No | No | No | No | No | Yes | No | No | No | No | +| NoPE layers | No | No | No | No | No | Yes | No | No | No | No | No | No | No | No | +| Partial RoPE | No | No | No | No | No | No | No | No | No | Yes | No | No | No | Yes (global layers) | +| Sliding-window layers | No | No | No | No | Yes (alternating) | No | No | No | No | No | No | No | No | No | +| Head-dim padding | 2B only | Yes (64→128) | No | Micro only (64→128) | No | No | TinyLlama | No | No | No | No | No | No | No | +| Custom model loading | No | Yes (safetensor remap) | No | No | No | No | No | No | No | No | No | No | No | No | +| Attention scaling | `config.attention_multiplier` | `config.attention_multiplier` | `head_dim**-0.5` | `config.attention_multiplier` | `config.attention_multiplier` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `head_dim**-0.5` | `query_pre_attn_scalar**-0.5` | `1.0` (unscaled) | +| Norm type | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | RMSNorm (pre) | LayerNorm (pre, no weight) | RMSNorm (post) | RMSNorm (sandwich) | RMSNorm (sandwich) | **Partial RoPE** (Phi-4): `PartialPrecomputedRotaryEmbedding` pads the rotation matrix with identity `[[1,0],[0,1]]` entries so diff --git a/hf_adapters.md b/hf_adapters.md index 6c0dbebc..14d919e8 100644 --- a/hf_adapters.md +++ b/hf_adapters.md @@ -12,9 +12,10 @@ | 6 | hf\_mistral.py | `mistral` | 1 | 4+ | | 7 | hf\_phi3.py | `phi3` | 1 | 3+ | | 8 | hf\_granitemoehybrid.py | `granitemoehybrid` | 2 | 1+ | -| 9 | hf\_smollm3.py | `smollm3` | 1 | — | -| 10 | hf\_olmo.py | `olmo` | 1 | 1+ | -| 11 | hf\_olmo2.py | `olmo2` | 1 | 1+ | +| 9 | hf\_granite\_swa.py | `granite_swa` | 1 | 1+ | +| 10 | hf\_smollm3.py | `smollm3` | 1 | — | +| 11 | hf\_olmo.py | `olmo` | 1 | 1+ | +| 12 | hf\_olmo2.py | `olmo2` | 1 | 1+ | --- @@ -173,7 +174,28 @@ Vision-specific adapter; extracts and adapts the text backbone from Granite Visi --- -## 9. hf\_smollm3.py +## 9. hf\_granite\_swa.py + +**HF model\_type:** `granite_swa` + +Spyre adapter for `GraniteSWAForCausalLM`. Decoder layers alternate between full +attention and sliding-window attention (`layer_type` on each `GraniteSWADecoderLayer`). +Sliding-window layers restrict attention to a local `sliding_window`-wide band; +full-attention layers use the standard causal mask. + +### Verified Checkpoints + +| Checkpoint | head\_dim | Notes | +|-----------|---------|-------| +| Granite 4.1 20B (`ibm-research/granite-4.1-20b`) | 128 | ibm-research org; bfloat16 | + +### Compatible Models + +- Granite 4.1 8B / 20B variants + +--- + +## 10. hf\_smollm3.py **HF model\_type:** `smollm3` @@ -189,7 +211,7 @@ No additional compatible models listed. --- -## 10. hf\_olmo.py +## 11. hf\_olmo.py **HF model\_type:** `olmo` @@ -205,7 +227,7 @@ No additional compatible models listed. --- -## 11. hf\_olmo2.py +## 12. hf\_olmo2.py **HF model\_type:** `olmo2` diff --git a/hf_adapters/auto_spyre_model.py b/hf_adapters/auto_spyre_model.py index dbaa8650..8df6a4e1 100644 --- a/hf_adapters/auto_spyre_model.py +++ b/hf_adapters/auto_spyre_model.py @@ -45,7 +45,6 @@ AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, - AutoModelForSequenceClassification, BertConfig, Gemma3Config, Gemma3TextConfig, @@ -59,6 +58,7 @@ Granite4VisionConfig, GraniteConfig, GraniteMoeHybridConfig, + GraniteSWAConfig, LlamaConfig, MistralConfig, ModernBertConfig, @@ -78,9 +78,6 @@ from hf_adapters import ( hf_bert, - hf_dspark_gemma4, - hf_dspark_granite, - hf_dspark_qwen3, hf_gemma3, hf_gemma4, hf_gemma4_mm, @@ -88,6 +85,7 @@ hf_gpt_neo, hf_gpt_neox, hf_granite, + hf_granite_swa, hf_granite_vision, hf_granite_vision_mm, hf_granitemoehybrid, @@ -108,7 +106,6 @@ ) from hf_adapters.hf_common import ( SpyreNoAdapterError, - SpyreUnsupportedModelError, assert_spyre_dimensions, load_model_common, move_model_to_spyre, @@ -128,6 +125,7 @@ Granite4VisionConfig: hf_granite_vision, GraniteConfig: hf_granite, GraniteMoeHybridConfig: hf_granitemoehybrid, + GraniteSWAConfig: hf_granite_swa, LlamaConfig: hf_llama, MistralConfig: hf_mistral, MinistralConfig: hf_ministral, @@ -144,19 +142,6 @@ XLMRobertaConfig: hf_xlm_roberta, } -# Architecture-name mapping — consulted BEFORE the config-class map. DSpark -# speculative-decoding *drafters* reuse their base model's config class -# (``Qwen3Config`` / ``Gemma4TextConfig`` / ``GraniteConfig``) but carry a -# distinct ``architectures`` entry (``*DSparkModel``). Config-class dispatch alone -# would route them to the *target* adapter; keying on the architecture name sends -# them to the drafter adapter instead. Normal targets have no entry here and fall -# through to ``CONFIG_TO_ADAPTER_MODULE_MAPPING`` unchanged. -ARCH_TO_ADAPTER_MODULE_MAPPING: dict[str, ModuleType] = { - "Qwen3DSparkModel": hf_dspark_qwen3, - "Gemma4DSparkModel": hf_dspark_gemma4, - "GraniteDSparkModel": hf_dspark_granite, -} - # Multimodal (image-text-to-text) mapping — used by # ``AutoSpyreModelForImageTextToText``. A multimodal checkpoint's config (e.g. # Granite4VisionConfig) appears here mapped to the *combined* two-tower adapter, @@ -170,25 +155,15 @@ Mistral3Config: hf_mistral3_vision_mm, } -# Sequence-classification (cross-encoder reranker) mapping — used by -# ``AutoSpyreModelForSequenceClassification``. -SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING: dict[ - type[PretrainedConfig], ModuleType -] = { - XLMRobertaConfig: hf_xlm_roberta, - RobertaConfig: hf_xlm_roberta, -} - MODEL_PATH_TO_TORCH_DTYPE: dict[str, torch.dtype] = { "mistralai/Ministral-3-3B-Instruct-2512": torch.bfloat16, "mistralai/Ministral-3-8B-Instruct-2512": torch.bfloat16, "mistralai/Ministral-3-14B-Instruct-2512": torch.bfloat16, "google/embeddinggemma-300m": torch.bfloat16, - "google/gemma-4-12b": torch.bfloat16, "google/gemma-4-12B-it": torch.bfloat16, - "google/gemma-4-31b": torch.bfloat16, "ibm-granite/granite-4.0-1b-base": torch.float32, "ibm-granite/granite-4.0-1b": torch.float32, + "ibm-research/granite-4.1-20b": torch.bfloat16, } @@ -203,14 +178,6 @@ def resolve_adapter_module( model_name_or_path, trust_remote_code=trust_remote_code ) - # Architecture-name dispatch first: DSpark drafters share their base model's - # config class but carry a distinct ``*DSparkModel`` architecture, so route on - # the architecture name before falling through to config-class dispatch. - for arch in getattr(model_config, "architectures", None) or []: - if arch in ARCH_TO_ADAPTER_MODULE_MAPPING: - assert_spyre_dimensions(model_config, model_name=str(model_name_or_path)) - return ARCH_TO_ADAPTER_MODULE_MAPPING[arch] - if type(model_config) not in mapping: raise SpyreNoAdapterError( f"Model {model_name_or_path} of type {type(model_config)} " @@ -270,11 +237,6 @@ def from_pretrained( dtype: torch.dtype = torch.float16, ) -> torch.nn.Module: module: ModuleType = resolve_adapter_module(model_name_or_path) - if getattr(module, "_is_encoder_only", False): - raise SpyreUnsupportedModelError( - "Generation is not currently supported for encoder-only architectures" - ) - model: torch.nn.Module = super().from_pretrained( model_name_or_path, dtype=dtype ) @@ -291,72 +253,6 @@ def model_generate( return model -class AutoSpyreModelForSequenceClassification(AutoSpyreModel): - """Load an XLM-RoBERTa cross-encoder reranker and prepare it for Spyre. - - Loads via ``AutoModelForSequenceClassification``, compiles the encoder - backbone on Spyre, and attaches a ``rerank`` method that tokenizes - query-document pairs and returns raw relevance logits. - - Example:: - - model = AutoSpyreModelForSequenceClassification.from_pretrained( - "BAAI/bge-reranker-v2-m3" - ) - tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3") - pairs = [("query text", "document text")] - scores = model.rerank(tokenizer, pairs) # raw logits - probs = torch.sigmoid(scores) # [0, 1] relevance - """ - - _auto_model_cls = AutoModelForSequenceClassification # type: ignore[assignment] - _module_mapping: dict[type[PretrainedConfig], ModuleType] = ( - SEQUENCE_CLASSIFICATION_CONFIG_TO_ADAPTER_MODULE_MAPPING - ) - - @classmethod - def from_pretrained( - cls, - model_name_or_path: Union[str, os.PathLike[str]], - dtype: torch.dtype = torch.float16, - ) -> torch.nn.Module: - module: ModuleType = resolve_adapter_module( - model_name_or_path, mapping=cls._module_mapping - ) - model: torch.nn.Module = super().from_pretrained( - model_name_or_path, dtype=dtype - ) - - def model_rerank( - self: torch.nn.Module, - tokenizer: Any, - pairs: list[tuple[str, str]], - **kwargs: Any, - ): - from hf_adapters.hf_common import prefill_reranker - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - encoded = tokenizer( - pairs, - return_tensors="pt", - padding=True, - truncation=True, - padding_side="right", - return_attention_mask=True, - ) - return prefill_reranker( - module._run_backbone_forward, - self, - encoded["input_ids"], - encoded["attention_mask"], - token_type_ids=encoded.get("token_type_ids", None), - ) - - model.rerank = MethodType(model_rerank, model) # type: ignore[assignment] - return model - - class AutoSpyreModelForImageTextToText(AutoSpyreModel): """Load a multimodal (image-text-to-text) model and prepare BOTH towers. diff --git a/hf_adapters/hf_granite_swa.py b/hf_adapters/hf_granite_swa.py new file mode 100644 index 00000000..e8bc11dd --- /dev/null +++ b/hf_adapters/hf_granite_swa.py @@ -0,0 +1,184 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +HuggingFace Transformers adapter for Granite 4.1 SWA models on Spyre. + +GraniteSWAForCausalLM is identical to Granite 3.x except that layers alternate +between full attention and sliding-window attention (``layer_type`` attribute on +each ``GraniteSWADecoderLayer``). Sliding-window layers receive a local attention +mask that restricts each token to attend only within the ``sliding_window`` window; +full-attention layers receive the normal causal mask. + +Usage:: + + from hf_adapters import AutoSpyreModelForCausalLM + from transformers import AutoTokenizer + + model = AutoSpyreModelForCausalLM.from_pretrained( + "/tmp/models/granite-4.1-20b") + tokenizer = AutoTokenizer.from_pretrained("/tmp/models/granite-4.1-20b") + outputs = model.generate(tokenizer, ["Hello!"], max_new_tokens=32) +""" + +import torch +import torch.nn.functional as F + +from hf_adapters.hf_common import ( + apply_rope_matmul, + get_backbone, + kv_cache_update, + pad_lm_head, + patch_rmsnorm, + prepare_rope_and_heads, +) + + +def _make_compiled_block(layer, sliding_window: int): + """Compiled block for GraniteSWA. + + Full-attention layers receive the normal causal ``attn_mask``. + Sliding-window layers receive a local mask built from the same buffer but + restricted to a ``sliding_window``-wide band; positions outside the window + are masked to ``-inf``. + """ + attn = layer.self_attn + mlp = layer.mlp + input_ln = layer.input_layernorm + post_attn_ln = layer.post_attention_layernorm + res_mult = layer.residual_multiplier + v_head_dim = getattr(attn, "v_head_dim", attn.head_dim) + is_sliding = getattr(layer, "layer_type", "full_attention") == "sliding_attention" + + def block_forward( + hidden_states, + selected_freqs, + attn_mask, + key_cache, + value_cache, + is_filling, + token_index, + cache_position, + ): + residual = hidden_states + h = input_ln(hidden_states) + + bsz, seq_len, _ = h.shape + q = attn.q_proj(h).view(bsz, seq_len, -1, attn.head_dim).transpose(1, 2) + k = attn.k_proj(h).view(bsz, seq_len, -1, attn.head_dim).transpose(1, 2) + v = attn.v_proj(h).view(bsz, seq_len, -1, v_head_dim).transpose(1, 2) + + q = apply_rope_matmul(q, selected_freqs) + k = apply_rope_matmul(k, selected_freqs) + + key_cache, value_cache = kv_cache_update( + k, + v, + key_cache, + value_cache, + is_filling, + token_index, + cache_position, + ) + + # For sliding-window layers, apply a local band mask on top of the + # causal mask: positions further than sliding_window steps back are + # set to -inf so the model never attends outside the window. + if is_sliding: + cache_len = key_cache.shape[2] + q_len = q.shape[2] + # Build position indices for queries and keys + q_pos = torch.arange(q_len, device=q.device).unsqueeze(1) # [q, 1] + k_pos = torch.arange(cache_len, device=q.device).unsqueeze(0) # [1, k] + window_mask = (q_pos - k_pos) >= sliding_window # [q, k] + swa_mask = attn_mask.clone() + swa_mask = swa_mask.masked_fill( + window_mask.unsqueeze(0).unsqueeze(0), float("-inf") + ) + effective_mask = swa_mask + else: + effective_mask = attn_mask + + attn_out = F.scaled_dot_product_attention( + q, + key_cache, + value_cache, + attn_mask=effective_mask, + dropout_p=0.0, + scale=attn.scaling, + enable_gqa=True, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) + attn_out = attn.o_proj(attn_out) + + h = residual + attn_out * res_mult + + residual = h + h = post_attn_ln(h) + h = mlp(h) + h = residual + h * res_mult + + return h, key_cache, value_cache + + return torch.compile(block_forward, dynamic=False) + + +def _run_forward( + model, + input_ids, + position_ids, + attn_mask, + key_caches, + value_caches, + is_filling, + token_index, + cache_position, +): + """GraniteSWA causal-LM forward: embedding * multiplier, blocks, norm, head.""" + backbone = get_backbone(model) + h = backbone.embed_tokens(input_ids) + h = h * backbone.embedding_multiplier + + selected_freqs = model._spyre_rope(h, position_ids) + + for i, compiled_block in enumerate(model._spyre_compiled_blocks): + h, key_caches[i], value_caches[i] = compiled_block( + h, + selected_freqs, + attn_mask, + key_caches[i], + value_caches[i], + is_filling, + token_index, + cache_position, + ) + + h = backbone.norm(h) + logits = model.lm_head(h) + logits = logits / model.config.logits_scaling + return logits + + +def prepare_for_spyre(model): + """Apply Spyre adaptations to a GraniteSWA model in-place.""" + from transformers.models.granite_swa.modeling_granite_swa import GraniteSWARMSNorm + + sliding_window = model.config.sliding_window + prepare_rope_and_heads(model) + patch_rmsnorm(GraniteSWARMSNorm) + pad_lm_head(model) + model._spyre_compiled_blocks = [ + _make_compiled_block(layer, sliding_window) + for layer in get_backbone(model).layers + ] diff --git a/tests/model_registry.py b/tests/model_registry.py index cd1335dd..44e3891c 100644 --- a/tests/model_registry.py +++ b/tests/model_registry.py @@ -97,6 +97,19 @@ def _include_gated() -> bool: "adapter": "hf_granitemoehybrid.py", "size": "3b", }, + "granite41_8b": { + "name": "Granite 4.1 8B", + "path": "ibm-granite/granite-4.1-8b", + "adapter": "hf_granite_swa.py", + "size": "8b", + }, + # hf_granite_swa.py + "granite41_20b": { + "name": "Granite 4.1 20B", + "path": "ibm-research/granite-4.1-20b", + "adapter": "hf_granite_swa.py", + "size": "20b", + }, # hf_granite_vision.py "granite-vision": { "name": "Granite Vision 4.1 4B", @@ -219,54 +232,12 @@ def _include_gated() -> bool: "size": "1b", }, # hf_gemma4 - "gemma4_base": { - "name": "Gemma 4 12B Base", - "path": "google/gemma-4-12b", - "adapter": "hf_gemma4.py", - "size": "12b", - "dtype": "bfloat16", - }, "gemma4_google": { "name": "Gemma 4 12B", "path": "google/gemma-4-12B-it", "adapter": "hf_gemma4.py", "size": "12b", }, - "gemma4_31b": { - "name": "Gemma 4 31B", - "path": "google/gemma-4-31b", - "adapter": "hf_gemma4.py", - "size": "31b", - "dtype": "bfloat16", - "is_gated": True, - }, - # DSpark speculative-decoding drafters (block proposers). kind="dspark_draft" - # keeps them out of the generate-based causal-LM harnesses (see CAUSAL_PATHS); - # tests/spyre/test_dspark_draft_spyre.py exercises the block-propose path. - # hf_dspark_qwen3.py - "dspark_qwen3": { - "name": "DSpark Qwen3 Drafter (block7)", - "path": "deepseek-ai/dspark_qwen3_4b_block7", - "adapter": "hf_dspark_qwen3.py", - "size": "0.6b", - "kind": "dspark_draft", - }, - # hf_dspark_gemma4.py - "dspark_gemma4": { - "name": "DSpark Gemma 4 Drafter (block7)", - "path": "deepseek-ai/dspark_gemma4_12b_block7", - "adapter": "hf_dspark_gemma4.py", - "size": "1b", - "kind": "dspark_draft", - }, - # hf_dspark_granite.py - "dspark_granite": { - "name": "DSpark Granite Drafter (block7)", - "path": "deepseek-ai/dspark_granite_4_1_8b_block7", - "adapter": "hf_dspark_granite.py", - "size": "1b", - "kind": "dspark_draft", - }, } EMBEDDING_MODELS = { @@ -504,22 +475,10 @@ def _select_representative_paths( # shared across all three selections. ``kind == "vlm"`` excludes bare vision towers. _include_gated_flag = _include_gated() -# ``kind == "dspark_draft"`` entries are speculative-decoding drafters (block -# proposers, driven by ``_run_draft_block`` — no ``generate``), so they are -# registered for adapter-coverage but excluded from the generate-based CPU/Spyre -# causal-LM harnesses; they are exercised by tests/spyre/test_dspark_draft_spyre.py. CAUSAL_PATHS: list[str] = _select_representative_paths( - CAUSAL_LM_MODELS, - include_gated=_include_gated_flag, - predicate=lambda info: info.get("kind") != "dspark_draft", -) -# The DSpark drafter checkpoints (block proposers), one per adapter — exercised by -# tests/spyre/test_dspark_draft_spyre.py via the block-propose ``_run_draft_block``. -DSPARK_PATHS: list[str] = _select_representative_paths( - CAUSAL_LM_MODELS, - include_gated=_include_gated_flag, - predicate=lambda info: info.get("kind") == "dspark_draft", + CAUSAL_LM_MODELS, include_gated=_include_gated_flag ) + EMBED_PATHS: list[str] = _select_representative_paths( EMBEDDING_MODELS, include_gated=_include_gated_flag ) @@ -529,43 +488,22 @@ def _select_representative_paths( predicate=lambda info: info.get("kind") == "vlm", ) - -def _non_blocking(models: dict[str, dict], keys: tuple[str, ...]) -> dict[str, str]: - """Build a ``{path: xfail reason}`` table from registry keys.""" - return { - models[key]["path"]: ( - f"{key}: temporarily non-blocking signal for specific models" - ) - for key in keys - } - - -# Non-blocking models (xfail, non-strict); remove an entry once it's been stably -# green so its Spyre tests go back to gating CI normally. -# -# The tables are per-harness rather than one merged dict because a path can name -# two different adapters — google/gemma-4-12B-it is both ``gemma4_google`` (causal) -# and ``gemma4_mm`` (VLM). -NON_BLOCKING_CAUSAL_MODELS: dict[str, str] = _non_blocking( - CAUSAL_LM_MODELS, - ( - "smollm3", - "gemma3_unsloth", - "ministral3", - "pythia_410m", - "gemma4_google", - "gemma4_base", - ), -) - -NON_BLOCKING_VISION_MODELS: dict[str, str] = _non_blocking( - VISION_MODELS, - ("gemma4_mm",), -) +# Causal-LM models that just went green on torchs-spyre but aren't yet proven stable +# across repeated runs. Kept as a non-blocking signal (xfail, non-strict) for +# a trial period; remove an entry once it's been stably green so its Spyre +# tests go back to gating CI normally. + +NON_BLOCKING_CAUSAL_MODELS: dict[str, str] = { + CAUSAL_LM_MODELS[key]["path"]: ( + f"{key}: newly green on Spyre, non-blocking signal for a trial " + "period before promoting to a blocking test" + ) + for key in ("qwen3", "olmo2_1b", "gemma3_unsloth", "ministral8b", "gemma4_google") +} -def xfail_non_blocking(paths: list[str], *, table: dict[str, str]) -> list[object]: - """Wrap entries of ``paths`` found in ``table`` with a non-strict xfail. +def xfail_non_blocking(paths: list[str]) -> list[object]: + """Wrap entries of ``paths`` found in NON_BLOCKING_CAUSAL_MODELS with xfail. The test still runs and its outcome (PASS/FAIL) is visible in the report, but a failure won't fail the pytest run or block CI. @@ -574,24 +512,13 @@ def xfail_non_blocking(paths: list[str], *, table: dict[str, str]) -> list[objec ( pytest.param( path, - marks=pytest.mark.xfail(reason=table[path], strict=False), + marks=pytest.mark.xfail( + reason=NON_BLOCKING_CAUSAL_MODELS[path], strict=False + ), id=path, ) - if path in table + if path in NON_BLOCKING_CAUSAL_MODELS else path ) for path in paths ] - - -RERANKER_MODELS = { - # hf_xlm_roberta.py - "bge_reranker_v2_m3": { - "name": "BGE Reranker v2 M3", - "path": "BAAI/bge-reranker-v2-m3", - "adapter": "hf_xlm_roberta.py", - "size": "0.5b", - }, -} - -RERANKER_PATHS: list[str] = [m["path"] for m in RERANKER_MODELS.values()] diff --git a/utils/model_ops/models/granite41-20b/run_huggingface.py b/utils/model_ops/models/granite41-20b/run_huggingface.py new file mode 100644 index 00000000..64462240 --- /dev/null +++ b/utils/model_ops/models/granite41-20b/run_huggingface.py @@ -0,0 +1,76 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache +from utils.torchop_yaml import TorchOpCollector, require_cuda, setup_logging + + +def main(): + setup_logging() + require_cuda() + + model_path = "ibm-granite/granite-4.1-20b" + chat = [ + { + "role": "user", + "content": "Please list one IBM Research laboratory located in the United States. You should only output its name and location.", + }, + ] + + device = "cuda" + model = AutoModelForCausalLM.from_pretrained( + model_path, device_map="auto", dtype=torch.bfloat16 + ) + tokenizer = AutoTokenizer.from_pretrained(model_path) + chat = tokenizer.apply_chat_template( + chat, tokenize=False, add_generation_prompt=True + ) + encoded_input = tokenizer(chat, return_tensors="pt").to(device) + + past_key_values = StaticCache(config=model.config, max_cache_len=2048) + + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + + model.forward = torch.compile(model.forward) + + with TorchOpCollector() as ctx: + with torch.no_grad(): + model.generate( + **encoded_input, + past_key_values=past_key_values, + use_cache=True, + max_new_tokens=16, + ) + + # print traced torch op + for op in ctx.ops_list: + print(op) + print(f"Total ops traced: {len(ctx.ops_list)}") + + # List of ops with generated test cases + print("List of ops with test cases generated") + for op in ctx.test_gen_ops: + print(op, ctx.test_case_count[op]) + print(f"Total ops with test configs generated: {len(ctx.test_gen_ops)}") + + ctx.write_yaml(os.path.basename(model_path)) + + +if __name__ == "__main__": + main()