diff --git a/hf_adapters/_dspark_common.py b/hf_adapters/_dspark_common.py index 8b27c150..b660e2b0 100644 --- a/hf_adapters/_dspark_common.py +++ b/hf_adapters/_dspark_common.py @@ -399,7 +399,7 @@ def install_spyre_markov(model): correction runs on host with only the ``w2`` matmul on-device — no device mix. No-op when the drafter has no markov head. """ - from deepspec.utils.sampling import sample_tokens + from deepspec.utils.sampling import sample_tokens # type: ignore[import-not-found] mh = getattr(model, "markov_head", None) if mh is None or not hasattr(mh, "markov_w2"): diff --git a/hf_adapters/hf_dspark_gemma4.py b/hf_adapters/hf_dspark_gemma4.py index 21eb39aa..90779edb 100644 --- a/hf_adapters/hf_dspark_gemma4.py +++ b/hf_adapters/hf_dspark_gemma4.py @@ -46,12 +46,10 @@ _pad_markov_w2, install_spyre_compute_logits, install_spyre_markov, + run_draft_block, snapshot_cpu_embeddings, snapshot_cpu_fc, ) -from hf_adapters._dspark_common import ( # noqa: F401 (re-exported as the public forward) - run_draft_block as _run_draft_block, -) from hf_adapters.hf_common import ( InvFreqShim, PrecomputedRotaryEmbedding, @@ -60,6 +58,7 @@ ) from hf_adapters.hf_gemma4 import _patch_gemma4_rmsnorm +_run_draft_block = run_draft_block # reuse the common runner CTX_PAD = 56 @@ -132,7 +131,9 @@ def block_forward(hidden_states, target_hidden_states, selected_freqs, attn_mask def prepare_for_spyre(model): """Apply Spyre adaptations to the Gemma4 DSpark drafter in-place.""" - from deepspec.modeling.dspark.gemma4.modeling import Gemma4RMSNorm + from deepspec.modeling.dspark.gemma4.modeling import ( # type: ignore[import-not-found] + Gemma4RMSNorm, + ) block_size = int(model.block_size) kv_pad = ((CTX_PAD + block_size + 31) // 32) * 32 diff --git a/hf_adapters/hf_dspark_granite.py b/hf_adapters/hf_dspark_granite.py index 57f1bc50..07162c18 100644 --- a/hf_adapters/hf_dspark_granite.py +++ b/hf_adapters/hf_dspark_granite.py @@ -34,13 +34,9 @@ Usage: see ``hf_dspark_qwen3``. """ -# The block-propose forward is shared across families — reuse the common runner -# under the adapter's public ``_run_draft_block`` name (see hf_dspark_qwen3). -from hf_adapters._dspark_common import prepare_dspark_common -from hf_adapters._dspark_common import ( # noqa: F401 (re-exported as the public forward) - run_draft_block as _run_draft_block, -) +from hf_adapters._dspark_common import prepare_dspark_common, run_draft_block +_run_draft_block = run_draft_block # reuse the common runner CTX_PAD = 56 diff --git a/hf_adapters/hf_dspark_qwen3.py b/hf_adapters/hf_dspark_qwen3.py index ac66b8b3..384d7ebb 100644 --- a/hf_adapters/hf_dspark_qwen3.py +++ b/hf_adapters/hf_dspark_qwen3.py @@ -37,14 +37,9 @@ model, draft_input_ids, target_hidden_states, selected_freqs, ctx_valid_len) """ -# The block-propose forward is identical across the drafter families — reuse the -# shared runner directly (exposed under the adapter's public ``_run_draft_block`` -# name that ``resolve_adapter_module``/tests look up). -from hf_adapters._dspark_common import prepare_dspark_common -from hf_adapters._dspark_common import ( # noqa: F401 (re-exported as the public forward) - run_draft_block as _run_draft_block, -) +from hf_adapters._dspark_common import prepare_dspark_common, run_draft_block +_run_draft_block = run_draft_block # reuse the common runner # Fixed context / kv widths (stick-aligned). CTX_PAD=56 keeps kv_pad = # round32(CTX_PAD + block_size=7) = 64 — the proven-compilable attention width. CTX_PAD = 56 @@ -52,7 +47,9 @@ def prepare_for_spyre(model): """Apply Spyre adaptations to the Qwen3 DSpark drafter in-place.""" - from deepspec.modeling.dspark.qwen3.modeling import Qwen3RMSNorm + from deepspec.modeling.dspark.qwen3.modeling import ( # type: ignore[import-not-found] + Qwen3RMSNorm, + ) block_size = int(model.block_size) kv_pad = ((CTX_PAD + block_size + 31) // 32) * 32 diff --git a/hf_adapters/hf_granite.py b/hf_adapters/hf_granite.py index c4a96fb1..d2593809 100644 --- a/hf_adapters/hf_granite.py +++ b/hf_adapters/hf_granite.py @@ -32,6 +32,7 @@ pad_lm_head, patch_rmsnorm, prepare_rope_and_heads, + text_config, ) @@ -93,8 +94,7 @@ def _run_forward( cache_position, ) logits = model.lm_head(h) - logits = logits / model.config.logits_scaling - return logits + return logits / text_config(model.config).logits_scaling def prepare_for_spyre(model): diff --git a/hf_adapters/hf_granite_vision.py b/hf_adapters/hf_granite_vision.py index 3297e1f9..f2893c67 100644 --- a/hf_adapters/hf_granite_vision.py +++ b/hf_adapters/hf_granite_vision.py @@ -41,15 +41,17 @@ import torch +from hf_adapters import hf_granite from hf_adapters.hf_common import ( get_backbone, make_standard_gqa_block, pad_lm_head, patch_rmsnorm, prepare_rope_and_heads, - text_config, ) -from hf_adapters.hf_granite import _run_backbone_forward + +_run_backbone_forward = hf_granite._run_backbone_forward +_run_forward = hf_granite._run_forward def load_hf_model(model_path, dtype=torch.float16): @@ -96,34 +98,3 @@ def prepare_for_spyre(model): model._spyre_compiled_blocks = [ make_standard_gqa_block(layer, True) for layer in get_backbone(model).layers ] - - -def _run_forward( - model, - input_ids, - position_ids, - attn_mask, - key_caches, - value_caches, - is_filling, - token_index, - cache_position, -): - """Granite Vision text causal-LM forward: backbone + head / scaling. - - Identical to ``hf_granite._run_forward`` except ``logits_scaling`` lives on - the nested ``text_config`` rather than the top-level VLM config. - """ - h = _run_backbone_forward( - model, - input_ids, - position_ids, - attn_mask, - key_caches, - value_caches, - is_filling, - token_index, - cache_position, - ) - logits = model.lm_head(h) - return logits / text_config(model.config).logits_scaling diff --git a/hf_adapters/hf_granitemoehybrid.py b/hf_adapters/hf_granitemoehybrid.py index 6c0414d5..422a8fe2 100644 --- a/hf_adapters/hf_granitemoehybrid.py +++ b/hf_adapters/hf_granitemoehybrid.py @@ -44,6 +44,7 @@ import torch.nn as nn import torch.nn.functional as F +from hf_adapters import hf_granite from hf_adapters.hf_common import ( apply_rope_matmul, get_backbone, @@ -54,6 +55,9 @@ split_fused_linear, ) +_run_backbone_forward = hf_granite._run_backbone_forward +_run_forward = hf_granite._run_forward + def _make_compiled_block(layer, res_mult, gate_proj, up_proj): """Compiled block for Granite 4.0 dense: split MLP, multipliers.""" @@ -119,68 +123,6 @@ def block_forward( return torch.compile(block_forward, dynamic=False) -def _run_backbone_forward( - model, - input_ids, - position_ids, - attn_mask, - key_caches, - value_caches, - is_filling, - token_index, - cache_position, -): - """Granite 4.0 backbone: embedding * multiplier, blocks, norm.""" - backbone = get_backbone(model) - h = backbone.embed_tokens(input_ids) - h = h * model.config.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) - return h - - -def _run_forward( - model, - input_ids, - position_ids, - attn_mask, - key_caches, - value_caches, - is_filling, - token_index, - cache_position, -): - """Granite 4.0 causal-LM forward: backbone + head / scaling.""" - h = _run_backbone_forward( - model, - input_ids, - position_ids, - attn_mask, - key_caches, - value_caches, - is_filling, - token_index, - cache_position, - ) - logits = model.lm_head(h) - logits = logits / model.config.logits_scaling - return logits - - def prepare_for_spyre(model): """Apply Spyre adaptations to Granite 4.0 dense model in-place.""" from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( diff --git a/hf_adapters/hf_mistral3.py b/hf_adapters/hf_mistral3.py index d3e722cc..25c6efe6 100644 --- a/hf_adapters/hf_mistral3.py +++ b/hf_adapters/hf_mistral3.py @@ -49,14 +49,14 @@ outputs = model.generate(tokenizer, ["Hello!"], max_new_tokens=32) """ +from hf_adapters import hf_mistral from hf_adapters.hf_common import ( get_backbone, prepare_standard_gqa, ) -from hf_adapters.hf_mistral import ( - _run_backbone_forward, # noqa: F401 re-exported as adapter module API - _run_forward, # noqa: F401 re-exported as adapter module API -) + +_run_backbone_forward = hf_mistral._run_backbone_forward +_run_forward = hf_mistral._run_forward def load_hf_model(model_path, dtype): diff --git a/hf_adapters/hf_xlm_roberta.py b/hf_adapters/hf_xlm_roberta.py index 4da0536d..9b3272b2 100644 --- a/hf_adapters/hf_xlm_roberta.py +++ b/hf_adapters/hf_xlm_roberta.py @@ -15,7 +15,7 @@ """ HuggingFace Transformers adapter for XLM-RoBERTa encoder-only models on Spyre. -Supports models with XLMRobertaConfig (e.g. BAAI/bge-m3, +Supports models with RoBERTa/XLMRobertaConfig (e.g. BAAI/bge-m3, intfloat/multilingual-e5-large, sentence-transformers/paraphrase-multilingual-*). Structurally identical to BERT (same attention/FFN/post-LN module names) so the @@ -27,17 +27,12 @@ ``padding_idx``. ``prefill_encoder`` synthesizes 0-based position ids that are correct for BERT but wrong for XLM-R, so this adapter overrides ``_run_backbone_forward`` to recompute position ids from ``input_ids`` itself. -- Embedding sum order differs (``word + position`` then ``+ token_type`` vs - BERT's three-way add). Mathematically equivalent — same final tensor. The - token-type add goes through ``add_token_type_embedding``, which broadcasts the - single row directly when ``type_vocab_size == 1`` (as in bge-m3 / - bge-reranker-v2-m3) because a one-row gather does not lower on Spyre. """ from hf_adapters.hf_bert import _make_compiled_encoder_block from hf_adapters.hf_common import ( BLOCK_SIZE, - add_token_type_embedding, + encoder_backbone_forward, fairseq_position_ids, get_backbone, pad_attention_heads_simple, @@ -45,29 +40,13 @@ def _run_backbone_forward(model, input_ids, attn_mask, position_ids, token_type_ids): - """Encoder backbone forward with XLM-R position ids. - - Modified version of ``encoder_backbone_forward``. To maintain the signature - of the original function, we pass in ``position_ids`` as an argument, but - compute the XLM-R-style positions from ``input_ids``. Otherwise - follows ``encoder_backbone_forward``: word + position + token_type embed, - LayerNorm, then the compiled encoder blocks with the Spyre layout-fixup - clones around each block. - """ - backbone = get_backbone(model) - emb = backbone.embeddings - + """Encoder backbone forward with RoBERTa/XLM-R position ids.""" + emb = get_backbone(model).embeddings pos_ids = fairseq_position_ids(input_ids, emb.padding_idx) - h = emb.word_embeddings(input_ids) + emb.position_embeddings(pos_ids) - h = add_token_type_embedding(h, emb, token_type_ids) - h = emb.LayerNorm(h) - h = h.clone() if h.device.type == "spyre" else h - for compiled_block in model._spyre_compiled_blocks: - h = compiled_block(h, attn_mask) - if h.device.type == "spyre": - h = h.clone() - return h + return encoder_backbone_forward( + model, input_ids, attn_mask, pos_ids, token_type_ids + ) _is_encoder_only = True diff --git a/tests/spyre/test_vlm_e2e_spyre.py b/tests/spyre/test_vlm_e2e_spyre.py index ddcef588..c6e4cc44 100644 --- a/tests/spyre/test_vlm_e2e_spyre.py +++ b/tests/spyre/test_vlm_e2e_spyre.py @@ -146,7 +146,6 @@ def _adapter_teacher_forced_steps( model_d_type = get_model_dtype(model) backbone = adapter.get_backbone(model) - # emb_mult = backbone.embedding_multiplier # Falls back to 1.0 for models (Mistral) that don't scale embeddings emb_mult = getattr(backbone, "embedding_multiplier", 1.0)