diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 54fc2351..8b52c713 100644 --- a/ARCHITECTURE.md +++ b/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. diff --git a/hf_adapters/_dspark_common.py b/hf_adapters/_dspark_common.py index 8b27c150..b99c70d9 100644 --- a/hf_adapters/_dspark_common.py +++ b/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) diff --git a/hf_adapters/hf_common.py b/hf_adapters/hf_common.py index 07c4c938..9d0aefc1 100644 --- a/hf_adapters/hf_common.py +++ b/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") diff --git a/hf_adapters/hf_gpt2.py b/hf_adapters/hf_gpt2.py index fdc36260..8df73710 100644 --- a/hf_adapters/hf_gpt2.py +++ b/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) diff --git a/hf_adapters/hf_granite_vision_mm.py b/hf_adapters/hf_granite_vision_mm.py index 761ba928..69f5244d 100644 --- a/hf_adapters/hf_granite_vision_mm.py +++ b/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). diff --git a/hf_adapters/hf_mistral3_vision_mm.py b/hf_adapters/hf_mistral3_vision_mm.py index af36400d..e0b49498 100644 --- a/hf_adapters/hf_mistral3_vision_mm.py +++ b/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") diff --git a/hf_adapters/hf_pixtral_vision.py b/hf_adapters/hf_pixtral_vision.py index 5fc3071d..464f77f7 100644 --- a/hf_adapters/hf_pixtral_vision.py +++ b/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 diff --git a/hf_adapters/hf_siglip_vision.py b/hf_adapters/hf_siglip_vision.py index 1c6f221e..3b8f409d 100644 --- a/hf_adapters/hf_siglip_vision.py +++ b/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()