Remove relayouting linear weights logic from hf-adapters load-common.py

ARCHITECTURE.md

@@ -560,8 +560,8 @@ SigLIP-specific Spyre adaptations:
   docs/siglip_vision_spyre_findings.md), so the patch embedding + learned
   position add run on CPU and the result is moved to Spyre. CPU copies of the
   conv weight/bias and position table are snapshotted at prepare time so the
-  closure survives the blanket device move (`_embedding_param_ids` can't exclude
-  a 4-D conv weight).
+  closure survives the blanket device move via `load_model_to_spyre`
+  (4-D conv weights are not left on CPU by the move).
 
 **Combined two-tower adapter — Granite** (`hf_granite_vision_mm.py`): runs the
 Spyre SigLIP tower, projects/packs its features, splices them into the
@@ -636,7 +636,7 @@ Multimodal-specific Spyre adaptations (beyond those shared with Granite VLM):
   explicitly stores Mistral decoder blocks in `model._spyre_text_blocks` to avoid
   collision with the vision tower's compiled blocks stored by
   `hf_pixtral_vision.prepare_for_spyre` in `model._spyre_compiled_blocks`.
-- **`multi_modal_projector` pinned to CPU** after `_move_to_spyre_with_layout`
+- **`multi_modal_projector` pinned to CPU** after the Spyre device move
   (same pattern as Granite's `layerwise_projectors` pin).
 - **`Mistral3PatchMerger`** (`nn.functional.unfold` + `merging_layer`) runs on
   CPU inside the projector — `unfold` doesn't lower on Spyre.

hf_adapters/_dspark_common.py

@@ -261,8 +261,8 @@ def snapshot_cpu_embeddings(model):
     """Keep CPU copies of the gather-only embedding weights.
 
     ``nn.Embedding`` is a gather with no Spyre kernel (CPU fallback), so the
-    noise-block and markov ``w1`` lookups must run on CPU. But
-    ``_move_to_spyre_with_layout`` (which runs AFTER ``prepare_for_spyre``) moves
+    noise-block and markov ``w1`` lookups must run on CPU. But the device move
+    via ``load_model_to_spyre`` (which runs AFTER ``prepare_for_spyre``) moves
     every parameter — including these embeddings — onto the Spyre device, which
     would collide CPU ids with a Spyre weight. Snapshot the weights on CPU here so
     the lookups are device-independent of where the move leaves the module; the
@@ -315,7 +315,7 @@ def snapshot_cpu_fc(model):
     tensor, off the autoregressive loop, so we run it on CPU (fp32, exact) and
     return the result on device — the same host boundary spyre_draft.py used, and
     analogous to the RoPE / embedding CPU fallbacks. Snapshot the weights here so
-    the projection is independent of where ``_move_to_spyre_with_layout`` leaves
+    the projection is independent of where the Spyre device move leaves
     ``fc``/``hidden_norm``.
     """
     fc = getattr(model, "fc", None)

hf_adapters/hf_common.py

@@ -1089,45 +1089,6 @@ def allocate_kv_caches(model, batch_size, max_cache_len, dtype, device=None):
 # ---------------------------------------------------------------------------
 
 
-def _patch_torch_empty():
-    """Workaround for torch_spyre spyre_empty() not accepting size= kwarg.
-
-    Upstream fix: https://github.com/torch-spyre/torch-spyre/issues/1729
-    """
-    _orig = torch.empty
-
-    def _patched(*args, size=None, **kwargs):
-        if size is not None:
-            return _orig(size, **kwargs)
-        return _orig(*args, **kwargs)
-
-    if getattr(torch.empty, "_hf_adapters_patched", False):
-        return
-    torch.empty = _patched
-    torch.empty._hf_adapters_patched = True
-
-
-def _embedding_param_ids(model):
-    """Data-pointers of weights that must keep the default (column-major) layout.
-
-    Gather-only embedding weights (used via ``nn.Embedding``, not matmul) must not
-    receive a row-major SpyreTensorLayout. Returns the set of ``data_ptr()``
-    values for all such weights.
-
-    Found by walking ``named_modules`` for ``nn.Embedding`` rather than matching
-    known attribute names. The name-matching version missed ModernBERT's
-    ``embeddings.tok_embeddings`` (and would miss the next new spelling), which
-    silently sent a [180000, 384] table down the matmul-weight path.
-    """
-    return {
-        module.weight.data_ptr()
-        for module in model.modules()
-        if isinstance(module, nn.Embedding)
-        and module.weight is not None
-        and module.weight.dim() == 2
-    }
-
-
 def untie_embedding_and_lm_head(model):
     """If the token-embedding weight and the LM head weight share storage, clone
     the LM head's weight so each can take a different Spyre layout.
@@ -1162,8 +1123,12 @@ def get_model_dtype(model: nn.Module) -> torch.dtype:
 
 
 def _move_to_spyre_with_layout(model, dtype):
-    """Move all parameters and buffers to Spyre with row-major layout for 2D
-    matmul weights, except embedding weights which keep the default layout.
+    """Prepare RoPE then transfer the model to Spyre via torch-spyre.
+
+    Layout selection (``dim_order=[1,0]`` for ``nn.Linear`` weights) is owned by
+    ``torch_spyre.model_utils.load_model_to_spyre``. This wrapper only handles
+    HF-specific RoPE prep and the CPU-test early return when ``DEVICE`` is not
+    Spyre.
     """
     # Propagate dtype to the precomputed RoPE module(s) so the freq cache
     # matches the chosen weight dtype (avoids fp16/bf16 mismatch in
@@ -1180,47 +1145,9 @@ def _move_to_spyre_with_layout(model, dtype):
         model.to(dtype=dtype)
         return
 
-    # Prime torch-spyre autoload before importing torch_spyre._C or calling
-    # torch.empty(..., device_layout=...). Calls with the spyre-only
-    # device_layout kwarg fail kwarg validation before dispatch.
-    torch.empty(1, device=DEVICE)
-
-    from torch_spyre._C import SpyreTensorLayout  # type: ignore[import-not-found]
-
-    skip_layout_ptrs = _embedding_param_ids(model)
-
-    def _alloc_on_spyre(t: torch.Tensor) -> torch.Tensor:
-        # The row-major [1, 0] dim_order describes a 2-D permutation, so it only
-        # applies to 2-D matmul weights. 1-D tensors (norms, biases) and any
-        # higher-rank weight (e.g. the 3-D/4-D Conv2d and position-embedding
-        # tables in a multimodal checkpoint's vision/audio towers) keep the
-        # default layout — forcing [1, 0] on them raises "Incompatible host_size
-        # and dim_order". Embedding tables are gather-only and also skipped.
-        if t.dim() == 2 and t.data_ptr() not in skip_layout_ptrs:
-            stl = SpyreTensorLayout(t.shape, t.stride(), dtype, [1, 0])
-        else:
-            stl = None
-        new: torch.Tensor = torch.empty(  # type: ignore[call-overload]
-            t.shape,
-            device=torch.device(DEVICE),
-            device_layout=stl,
-            dtype=dtype,
-        )
-        new.copy_(t.to(dtype))
-        return new
-
-    for name, param in list(model.named_parameters()):
-        new = _alloc_on_spyre(param.data)
-        module_path, _, attr = name.rpartition(".")
-        owner = model.get_submodule(module_path) if module_path else model
-        setattr(owner, attr, nn.Parameter(new, requires_grad=False))
+    from torch_spyre.model_utils import load_model_to_spyre
 
-    for name, buf in list(model.named_buffers()):
-        new = _alloc_on_spyre(buf)
-        module_path, _, attr = name.rpartition(".")
-        owner = model.get_submodule(module_path) if module_path else model
-        persistent = attr not in owner._non_persistent_buffers_set
-        owner.register_buffer(attr, new, persistent=persistent)
+    load_model_to_spyre(model, dtype=dtype)
 
 
 def load_model_common(model_path, module, dtype=torch.float16, auto_model_cls=None):
@@ -1255,8 +1182,6 @@ def load_model_common(model_path, module, dtype=torch.float16, auto_model_cls=No
 def move_model_to_spyre(model, module, dtype: torch.dtype) -> None:
     untie_embedding_and_lm_head(model)
     module.prepare_for_spyre(model)
-    # print("Moving model to Spyre ...")
-    _patch_torch_empty()
     _move_to_spyre_with_layout(model, dtype)
     for submod_name in getattr(model, "_spyre_cpu_submodules", []):
         model.get_submodule(submod_name).to("cpu")

hf_adapters/hf_gpt2.py

@@ -28,11 +28,11 @@
   - a GELU MLP (``c_fc`` → ``gelu_new`` → ``c_proj``);
   - the backbone at ``model.transformer`` (not ``model.model``).
 
-``prepare_for_spyre`` rewrites every ``Conv1D`` into an ``nn.Linear`` (so the
-Spyre row-major matmul layout applies) and splits the fused ``c_attn`` into
-separate q/k/v linears, then compiles one block per layer via the shared
-``make_decoder_block`` (the non-RoPE causal-decoder factory, also used by the
-OPT/BLOOM/MPT family).
+``prepare_for_spyre`` rewrites every ``Conv1D`` into an ``nn.Linear`` (so
+``torch_spyre.model_utils.load_model_to_spyre`` can apply optimal Linear
+layout) and splits the fused ``c_attn`` into separate q/k/v linears, then
+compiles one block per layer via the shared ``make_decoder_block`` (the
+non-RoPE causal-decoder factory, also used by the OPT/BLOOM/MPT family).
 
 Usage::
 
@@ -60,9 +60,9 @@ def _conv1d_to_linear(conv):
 
     ``Conv1D`` computes ``y = x @ W + b`` with ``W`` of shape ``[in, out]`` and
     ``b`` of shape ``[out]``. ``nn.Linear`` computes ``y = x @ W.T + b`` with
-    ``W`` of shape ``[out, in]``, so the weight is the transpose. Spyre's
-    row-major matmul layout (see ``_move_to_spyre_with_layout``) targets 2-D
-    ``nn.Linear`` weights, so we convert before moving to device.
+    ``W`` of shape ``[out, in]``, so the weight is the transpose.
+    ``load_model_to_spyre`` applies optimal layout only to ``nn.Linear``
+    weights, so we convert before moving to device.
     """
     in_features, out_features = conv.weight.shape
     linear = nn.Linear(in_features, out_features, bias=conv.bias is not None)

hf_adapters/hf_granite_vision_mm.py

@@ -106,7 +106,7 @@ def prepare_for_spyre(model):
 def _embed_text(model, input_ids):
     """Token embeddings * embedding_multiplier (Granite scales its embeddings).
 
-    The gather runs on ``embed_tokens``' device — after ``_move_to_spyre_with_layout``
+    The gather runs on ``embed_tokens``' device — after the Spyre device move
     the table lives on Spyre, so ``input_ids`` is moved to match (mirrors the
     decode-step ``embed_ids``). Returns embeddings on the embedding's device.
     """
@@ -147,7 +147,7 @@ def _deepstack_features(model, pixel_values, image_sizes):
 
     # The deepstack/spatial projectors (Blip2 QFormers) and the image_newline
     # parameter are stock CPU modules: _project_and_pack / pack_image_features run
-    # them on CPU (vision features are moved to CPU first). _move_to_spyre_with_layout
+    # them on CPU (vision features are moved to CPU first). load_model_to_spyre
     # blanket-moves every param to Spyre, so re-pin these to CPU before use — the
     # same CPU-fallback contract as the patch-embed conv (idempotent; .to(cpu) on an
     # already-CPU module is a no-op).

hf_adapters/hf_mistral3_vision_mm.py

@@ -116,10 +116,10 @@ def prepare_for_spyre(model):
     hf_pixtral_vision.prepare_for_spyre(model)
 
     # --- Text decoder ---
-    # Re-pin the multi_modal_projector to CPU: _move_to_spyre_with_layout
-    # will blanket-move every param; the projector must run on CPU because
-    # it processes CPU vision features (same pattern as granite_vision_mm's
-    # layerwise_projectors pin).
+    # Re-pin the multi_modal_projector to CPU: the device move via
+    # load_model_to_spyre blanket-moves every param; the projector must run
+    # on CPU because it processes CPU vision features (same pattern as
+    # granite_vision_mm's layerwise_projectors pin).
     if hasattr(model, "model") and hasattr(model.model, "multi_modal_projector"):
         model.model.multi_modal_projector.to("cpu")
 

hf_adapters/hf_pixtral_vision.py

@@ -376,7 +376,7 @@ def prepare_for_spyre(model):
     patch_rmsnorm(PixtralRMSNorm)
 
     # Snapshot the CPU patch-embed closure (Conv2d + ln_pre; keeps working on
-    # CPU after _move_to_spyre_with_layout relocates the tower's params).
+    # CPU after the device move relocates the tower's params).
     model._spyre_pixtral_patch_embed = _make_patch_embed_fn(tower)
 
     # Snapshot the 2D RoPE inv_freq table on CPU — used in prefill_vision_tower

hf_adapters/hf_siglip_vision.py

@@ -144,12 +144,11 @@ def _make_patch_embed(inner):
     flags; the patch path is small and fixed, so we inline it.
 
     The Conv2d weight/bias and position table are captured as CPU copies here, at
-    prepare time, so this closure keeps running on CPU after
-    ``_move_to_spyre_with_layout`` relocates the module's own params to Spyre
-    (``nn.Conv2d`` is not assumed to lower on Spyre — see
-    docs/siglip_vision_spyre_findings.md). ``_embedding_param_ids`` cannot exclude
-    these because the conv weight is 4-D / the position table is reached via a
-    tower-specific path, so we snapshot rather than skip-the-move.
+    prepare time, so this closure keeps running on CPU after the device move via
+    ``_move_to_spyre_with_layout`` / ``load_model_to_spyre`` relocates the
+    module's own params to Spyre (``nn.Conv2d`` is not assumed to lower on Spyre
+    — see docs/siglip_vision_spyre_findings.md). We snapshot rather than relying
+    on the move to leave these on CPU.
     """
     emb = inner.embeddings
     weight = emb.patch_embedding.weight.detach().cpu()