Performance fixes and bug fixes

toolkit/dequantize.py

@@ -46,6 +46,10 @@ def hacked_state_dict(self, *args, **kwargs):
 
 # hacks the state dict so we can dequantize before saving
 def patch_dequantization_on_save(model):
+    # idempotent: re-patching would capture the hacked state_dict as
+    # orig_state_dict and recurse on save
+    if getattr(model, "orig_state_dict", None) is not None:
+        return
     model.orig_state_dict = model.state_dict
     model.state_dict = partial(hacked_state_dict, model)
   

toolkit/memory_management/manager.py

@@ -6,6 +6,7 @@ from .manager_modules import (
     LinearLayerMemoryManager,
     ConvLayerMemoryManager,
     OstrisLinearLayerMemoryManager,
+    EmbeddingLayerMemoryManager,
     _DEVICE_STATE,
 )
 import random
@@ -92,12 +93,59 @@ class MemoryManager:
         module._mm_to = module.to
         module.to = module._memory_manager.memory_managed_to
 
-        # add ignore modules to unmanaged list
+        # a fully offloaded module's parameters all live on cpu, which makes
+        # ModelMixin.device (and pipelines deriving their execution device
+        # from it) report "cpu" and plant latents/timesteps there. Report the
+        # compute device instead via an in-place subclass (same class-swap
+        # pattern as OstrisLinear/adopt_component); detach() restores it.
+        try:
+            module._mm_orig_class = module.__class__
+            managed_cls = type(
+                module.__class__.__name__,
+                (module.__class__,),
+                {
+                    "device": property(
+                        lambda self: self._memory_manager.process_device,
+                        lambda self, value: self.__dict__.__setitem__(
+                            "_mm_device_shadow", value
+                        ),
+                    )
+                },
+            )
+            # transformers keys per-class registries (e.g. the hidden-states
+            # capture specs) by str(model.__class__); make the subclass
+            # stringify identically so those lookups still hit
+            managed_cls.__module__ = module.__class__.__module__
+            managed_cls.__qualname__ = module.__class__.__qualname__
+            module.__class__ = managed_cls
+        except TypeError:
+            # exotic class layouts (__slots__ etc.): keep the original class
+            module._mm_orig_class = None
+
+        # add ignore modules to unmanaged list; they must stay RESIDENT on the
+        # compute device (fp32 tables, pad tokens) — a model attached while
+        # parked on cpu would otherwise feed cpu tensors into gpu math
         for im in ignore_modules:
             module._memory_manager.unmanaged_modules.append(im)
+            try:
+                if isinstance(im, torch.Tensor):
+                    im.data = im.data.to(device)
+                elif isinstance(im, torch.nn.Module):
+                    im.to(device)
+            except Exception:
+                pass
 
         # count ignore modules as processed
         modules_processed = [x for x in ignore_modules]
+
+        # weights tied to an embedding (lm_head <-> embed_tokens) must not be
+        # managed: pinning the linear's weight to cpu strands the (unmanaged)
+        # embedding that shares the same tensor
+        embedding_weight_ptrs = {
+            m.weight.data_ptr()
+            for m in module.modules()
+            if isinstance(m, torch.nn.Embedding)
+        }
         # attach to all modules
         for name, sub_module in module.named_modules():
             for child_name, child_module in sub_module.named_modules():
@@ -110,6 +158,12 @@ class MemoryManager:
                         # randomly skip some modules
                         if random.random() > offload_percent:
                             skip = True
+                    if (
+                        not getattr(child_module, "is_ostris_quantized", False)
+                        and isinstance(getattr(child_module, "weight", None), torch.Tensor)
+                        and child_module.weight.data_ptr() in embedding_weight_ptrs
+                    ):
+                        skip = True
                     if skip:
                         module._memory_manager.unmanaged_modules.append(child_module)
                     else:
@@ -158,12 +212,40 @@ class MemoryManager:
                                 )
                             modules_processed.append(ara)
                     modules_processed.append(child_module)
+                elif (
+                    isinstance(child_module, torch.nn.Embedding)
+                    and child_module not in modules_processed
+                    and child_module.weight.numel()
+                    * child_module.weight.element_size()
+                    > 64 * 1024 * 1024
+                ):
+                    # (the cpu gather is autograd-transparent, so a trainable
+                    # embedding still gets grads — they just land on cpu)
+                    # large frozen vocab table: cpu-resident, rows gathered on
+                    # cpu (only the looked-up tokens cross the bus)
+                    EmbeddingLayerMemoryManager.attach(
+                        child_module, module._memory_manager
+                    )
+                    # the cpu move replaced weight.data, so a tied lm_head no
+                    # longer matches the (original-ptr) tied guard and gets
+                    # managed normally — correct: the shared weight now lives
+                    # on cpu, so the linear must bounce it
+                    modules_processed.append(child_module)
                 elif child_module.__class__.__name__ in UNMANAGED_MODULES or any(
                     inc in child_module.__class__.__name__
                     for inc in UNMANAGED_MODULES_INCLUDES
                 ):
-                    # unmanaged
-                    module._memory_manager.unmanaged_modules.append(child_module)
+                    # unmanaged — but never re-list a module the nested walk
+                    # already managed (a managed Embedding landing here would
+                    # get hauled back to the gpu by memory_managed_to), and
+                    # don't append duplicates
+                    if (
+                        child_module not in modules_processed
+                        and not hasattr(child_module, "_layer_memory_manager")
+                        and child_module
+                        not in module._memory_manager.unmanaged_modules
+                    ):
+                        module._memory_manager.unmanaged_modules.append(child_module)
                 else:
                     continue
 
@@ -179,6 +261,10 @@ class MemoryManager:
         if not hasattr(module, "_memory_manager"):
             return
 
+        if getattr(module, "_mm_orig_class", None) is not None:
+            module.__class__ = module._mm_orig_class
+            del module._mm_orig_class
+
         for unmanaged in module._memory_manager.unmanaged_modules:
             try:
                 if isinstance(unmanaged, torch.Tensor):

toolkit/memory_management/manager_modules.py

@@ -312,6 +312,12 @@ class _BouncingLinearFn(torch.autograd.Function):
             ctx.device = torch.device("cpu")
             return out.to(x.device)
 
+        if x.device != device:
+            # a pipeline that derives its execution device from the module's
+            # (fully offloaded -> "cpu") parameters can feed cpu activations;
+            # compute happens on the staged device and the output stays there
+            x = x.to(device, non_blocking=True)
+
         state = _get_device_state(device)
         # the guard makes current_stream() (used by the staging helpers' event
         # waits/records) resolve to the process device; without it they hit
@@ -461,6 +467,11 @@ class _BouncingConv2dFn(torch.autograd.Function):
             ctx.meta = ("cpu", stride, padding, dilation, groups, target_dtype)
             return out.to(x.device)
 
+        if x.device != device:
+            # cpu activations from a fully-offloaded pipeline: see
+            # _BouncingLinearFn.forward
+            x = x.to(device, non_blocking=True)
+
         state = _get_device_state(device)
         # device guard: see _BouncingLinearFn.forward
         with torch.cuda.device(device):
@@ -631,6 +642,42 @@ class BaseLayerMemoryManager:
             param._is_memory_managed = True
 
 
+class EmbeddingLayerMemoryManager(BaseLayerMemoryManager):
+    """Offloads a (large, frozen) nn.Embedding by keeping the weight on cpu
+    and doing the row gather THERE: instead of staging a multi-GB vocab table
+    to the gpu, only the looked-up rows (tokens x dim, ~KBs) cross the bus.
+    Lookups happen once per prompt, so the cpu gather is free."""
+
+    def __init__(self, module: nn.Module, manager: "MemoryManager"):
+        super().__init__(module, manager)
+
+        # cpu-resident weight; no pinning — the weight never crosses the bus
+        module.weight.data = module.weight.data.to("cpu")
+
+        self._original_forward = module.forward
+
+        def _mm_forward(input_ids, *args, **kwargs):
+            if args or kwargs:
+                return self._original_forward(input_ids, *args, **kwargs)
+            out_device = (
+                input_ids.device
+                if input_ids.device.type == "cuda"
+                else self.manager.process_device
+            )
+            out = F.embedding(
+                input_ids.to("cpu"),
+                self.module.weight,
+                self.module.padding_idx,
+                self.module.max_norm,
+                self.module.norm_type,
+                self.module.scale_grad_by_freq,
+                self.module.sparse,
+            )
+            return out.to(out_device, non_blocking=True)
+
+        module.forward = _mm_forward
+
+
 class LinearLayerMemoryManager(BaseLayerMemoryManager):
     def __init__(
         self,
@@ -739,6 +786,11 @@ class OstrisLinearLayerMemoryManager(BaseLayerMemoryManager):
                 # already resident on device
                 return self._original_forward(x)
 
+            if x.device != device:
+                # cpu activations from a fully-offloaded pipeline: see
+                # _BouncingLinearFn.forward
+                x = x.to(device, non_blocking=True)
+
             state = _get_device_state(device)
             d = state["depth"]
             idx = state["forward_clk"]

toolkit/models/v2/PLANNING.md

@@ -470,11 +470,75 @@ metrics_baseline.json). Profiling findings and fixes:
 - posix_fadvise(WILLNEED) readahead on single-file loads (mixin._readahead)
   — ~10-15% on cold reads, free when warm, no-op off POSIX.
 
+Round 2 (same day):
+- ARA path (`attach_ara_and_quantize`) now quantizes each hijacked linear on
+  the gpu (was: wherever it lived, historically the cpu) with the same
+  keep_on_device policy; the uint8 extras pass follows suit. Verified with
+  hidream + its uint3 ARA end to end.
+- `quantize_model` delegates its non-ARA branch to `quantize_module`
+  (keep_on_device derived from model_config); dead `_quantize_model_blocks`
+  removed; `patch_dequantization_on_save` made idempotent (double-patching
+  would have recursed on save).
+- Local-first hub loads: `_local_first` wraps from_pretrained / load_config /
+  AutoTokenizer / AutoProcessor — fully-cached repos load with
+  local_files_only (13x on VAE metadata, 0.53s -> 0.04s; a failed local
+  attempt is instant, so uncached repos fall through to the online path at
+  zero cost). Tokenizer-load slowness was DIAGNOSED as these per-process hub
+  etag HEADs, not the use_fast=False sites.
+
+Round 3 — vram spike elimination (2026-08-28):
+- keep_on_device's extras pass had moved the whole unquantized remainder to
+  the gpu in bf16 before quantizing (+7-10GB transient over final residency
+  on big-extras models). Now `quantize(keep_on_quantize_device=True)`
+  quantizes each extra layer on the gpu and leaves it — transient is one
+  layer; non-quantized leftovers move only at final placement. Same for the
+  ARA uint8 pass.
+- convrot8 `quantize_` chunked: takes the weight at stored dtype
+  (wants_fp32_weight=False; ConvRotIntN keeps True) and rotates+quantizes in
+  256MB row-chunks — a 152k-vocab TE projection's transient fell 7.6GB ->
+  1.4GB. Sub-threshold layers keep the one-shot path, verified bit-identical.
+- qwen_image measured: standard load peak 35.3 -> 27.5GB (below the 28.0
+  baseline; process peak now set by generation); low_vram load quantizes the
+  whole 20B stack never holding >2.4GB gpu.
+
 Ideas not yet done: pinned-buffer staged h2d (~+60% warm-upload bandwidth,
 9.5 vs 5.8GB/s measured), overlapping h2d with quantize kernels (small — the
 kernels are ~2% of the pass), wan22 low_vram generate choreography (165s of
-unload/reload per sample), extending keep_on_device to the legacy
-quantize_model/trainer path.
+unload/reload per sample), holders' legacy "just to make sure" .to(device)
+lines that move quantized TEs to gpu even under low_vram (qwen's Preparing
+Model stage parks 7.6GB — policy question, not a quantize spike).
+
+## Offload hardening + final certification (2026-08-28)
+
+The harness gained a per-arch 100%-offload probe (attach MemoryManager at
+offload_percent=1.0 to the DiT + TEs + connectors/vision towers, then a
+2-step sample — 2 not 1: diffusers' shift_terminal stretch divides by zero
+at a single step). It flushed out and fixed, in the manager:
+- fully-offloaded modules now REPORT the compute device (class swap with
+  stringification-stable __module__/__qualname__ — transformers keys its
+  hidden-states capture registry by str(class)); bouncing ops also pull cpu
+  activations to the staged device
+- weights tied to embeddings (lm_head <-> embed_tokens) are never pinned out
+  from under the unmanaged embedding
+- ignore_modules are moved to the compute device at attach
+- the nested attach walk no longer double-lists managed modules as unmanaged
+  (a managed Embedding was getting hauled back to gpu by pipeline.to)
+- NEW: bouncing embeddings — vocab tables >64MB stay cpu-resident with the
+  row gather done ON cpu (only looked-up rows cross the bus); tied lm_heads
+  fall through to normal bouncing
+- wan22's dual wrapper skips its whole-14B .to() swaps when memory-managed,
+  and offload attaches per sub-transformer (pipelines hold those directly)
+
+Measured offload-phase peaks: wan22_14b 14.7 -> 1.5GB, ltx2.5 14.0 ->
+11.8GB (remainder: video activations + resident VAEs/vocoder by design),
+most image archs 1-4GB vs 10-53GB resident. flux2's Mistral warnings
+silenced behavior-preservingly (tie_word_embeddings=False, explicit
+fix_mistral_regex=False, processor_kwargs — tokenization verified
+bit-identical). ConvRot skip notices gated to layers >=1M weights.
+
+FINAL CERTIFICATION: full sweep, 33/34 PASS, 0 FAIL, 1 SKIP (mageflow
+upstream 404) — load + convrot8 quantize + generate + 100%-offload forward
+for every arch on the complete optimized stack.
 
 ## Testing
 

toolkit/models/v2/_mixin.py

@@ -139,13 +139,29 @@ class OstrisModelMixin:
 
     @classmethod
     def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
-        return cls.from_pretrained(
-            path, subfolder=subfolder, torch_dtype=dtype, **kwargs
+        return cls._local_first(
+            cls.from_pretrained, path, subfolder=subfolder, torch_dtype=dtype, **kwargs
         )
 
     @classmethod
     def aitk_load_config(cls, path, subfolder=None):
-        return cls.load_config(path, subfolder=subfolder)
+        return cls._local_first(cls.load_config, path, subfolder=subfolder)
+
+    @staticmethod
+    def _local_first(loader, *args, **kwargs):
+        """Serve fully-cached repos without touching the hub. Online
+        from_pretrained does an etag HEAD round-trip per file (seconds of pure
+        network latency on every warm load); a local_files_only attempt is
+        instant when the cache can't satisfy it, so failing over to the
+        online path costs nothing."""
+        import os as _os
+
+        if _os.environ.get("HF_HUB_OFFLINE") == "1":
+            return loader(*args, **kwargs)
+        try:
+            return loader(*args, local_files_only=True, **kwargs)
+        except Exception:
+            return loader(*args, **kwargs)
 
     @classmethod
     def aitk_from_config(cls, config):
@@ -247,6 +263,17 @@ class OstrisModelMixin:
             exclude = list(exclude_quant_modules or []) + list(
                 self.get_quantization_exclude_modules() or []
             )
+            q_device = quantize_device if quantize_device is not None else device
+            # final home is the quantize gpu (no offloading, no cpu
+            # parking): quantized weights stay put — one bus crossing
+            # instead of three, and no model-sized host-ram copy
+            keep_on_device = (
+                not offload
+                and device is not None
+                and q_device is not None
+                and torch.device(device) == torch.device(q_device)
+                and torch.device(device).type != "cpu"
+            )
             if ara_path is not None:
                 if base_model is None:
                     raise ValueError(
@@ -254,20 +281,16 @@ class OstrisModelMixin:
                         "network wiring"
                     )
                 status_fn("Quantizing with accuracy recovery adapter")
-                attach_ara_and_quantize(base_model, self, ara_path, exclude=exclude)
+                attach_ara_and_quantize(
+                    base_model,
+                    self,
+                    ara_path,
+                    exclude=exclude,
+                    device=q_device,
+                    keep_on_device=keep_on_device,
+                )
             else:
                 status_fn(f"Quantizing ({qtype})")
-                q_device = quantize_device if quantize_device is not None else device
-                # final home is the quantize gpu (no offloading, no cpu
-                # parking): quantized weights stay put — one bus crossing
-                # instead of three, and no model-sized host-ram copy
-                keep_on_device = (
-                    not offload
-                    and device is not None
-                    and q_device is not None
-                    and torch.device(device) == torch.device(q_device)
-                    and torch.device(device).type != "cpu"
-                )
                 quantize_module(
                     self,
                     qtype,
@@ -623,8 +646,8 @@ class OstrisModelMixin:
             os.path.join(source, subfolder)
         ):
             subfolder = None
-        return AutoTokenizer.from_pretrained(
-            source, subfolder=subfolder or "", **kwargs
+        return cls._local_first(
+            AutoTokenizer.from_pretrained, source, subfolder=subfolder or "", **kwargs
         )
 
     @classmethod
@@ -645,8 +668,8 @@ class OstrisModelMixin:
             os.path.join(source, subfolder)
         ):
             subfolder = None
-        return AutoProcessor.from_pretrained(
-            source, subfolder=subfolder or "", **kwargs
+        return cls._local_first(
+            AutoProcessor.from_pretrained, source, subfolder=subfolder or "", **kwargs
         )
 
     # ------------------------------------------------------------------
@@ -778,15 +801,17 @@ class OstrisTransformersMixin(OstrisModelMixin):
 
     @classmethod
     def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
-        return cls.from_pretrained(
-            path, subfolder=subfolder or "", torch_dtype=dtype, **kwargs
+        return cls._local_first(
+            cls.from_pretrained, path, subfolder=subfolder or "", torch_dtype=dtype, **kwargs
         )
 
     @classmethod
     def aitk_load_config(cls, path, subfolder=None):
         from transformers import AutoConfig
 
-        return AutoConfig.from_pretrained(path, subfolder=subfolder or "")
+        return cls._local_first(
+            AutoConfig.from_pretrained, path, subfolder=subfolder or ""
+        )
 
     @classmethod
     def aitk_from_config(cls, config):

toolkit/models/wan21/wan21.py

@@ -433,7 +433,10 @@ class Wan21(BaseModel):
         pipe.transformer = pipe.transformer.to(self.device_torch)
 
         flush()
-        text_encoder.to(self.device_torch)
+        # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
+        # to the gpu on demand
+        if not self.model_config.low_vram:
+            text_encoder.to(self.device_torch)
         text_encoder.requires_grad_(False)
         text_encoder.eval()
         pipe.transformer = pipe.transformer.to(self.device_torch)

toolkit/util/convrot_quant.py

@@ -662,7 +662,7 @@ class ConvRotQuantizer(OstrisQuantizer):
         d = module.in_features
         rot = self._rot_for(d)
         if d % BLOCK != 0 or module.out_features % BLOCK != 0 or rot < 16:
-            if d not in _skip_warned:
+            if d not in _skip_warned and d * module.out_features >= 1 << 20:
                 _skip_warned.add(d)
                 print_acc(
                     f"ConvRot: skipping linears with in_features={d} "
@@ -1312,7 +1312,7 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
     def can_quantize(self, module: torch.nn.Linear) -> bool:
         d = module.in_features
         if d % BLOCK != 0 or module.out_features % 8 != 0 or self._rot_for(d) < 16:
-            if d not in _skip_warned:
+            if d not in _skip_warned and d * module.out_features >= 1 << 20:
                 _skip_warned.add(d)
                 print_acc(
                     f"ConvRot: skipping linears with in_features={d} "
@@ -1324,7 +1324,7 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
             # (torch._int_mm raises CUBLAS_STATUS_NOT_SUPPORTED), e.g.
             # omnigen2's 64 -> 2520 x_embedder
             key = (d, module.out_features)
-            if key not in _skip_warned:
+            if key not in _skip_warned and d * module.out_features >= 1 << 20:
                 _skip_warned.add(key)
                 print_acc(
                     f"ConvRot: skipping linears with in_features={d}, "
@@ -1334,9 +1334,32 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
             return False
         return True
 
+    # takes the weight in its stored dtype and casts per row-chunk below: a
+    # full-size fp32 copy of a huge layer (e.g. a TE's vocab projection) was
+    # a multi-GB vram transient on top of the full-size rotated fp32 tensor
+    wants_fp32_weight = False
+
+    # fp32 working-set target per chunk; layers under this keep the one-shot
+    # path (identical math and bit-identical output to the historical code)
+    _QUANT_CHUNK_BYTES = 256 * 1024 * 1024
+
     def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
         rot = self._rot_for(module.in_features)
-        q, scales = quantize_int8_rows(rotate(weight_fp32, rot))
+        rows, in_f = weight_fp32.shape
+        chunk_rows = max(1, self._QUANT_CHUNK_BYTES // (in_f * 4))
+        if rows <= chunk_rows:
+            q, scales = quantize_int8_rows(rotate(weight_fp32.float(), rot))
+        else:
+            # row-chunked: rotation and per-row scaling are row-independent,
+            # so the fp32 transient stays ~3 chunks instead of ~4x the layer
+            q = torch.empty(rows, in_f, dtype=torch.int8, device=weight_fp32.device)
+            scales = torch.empty(rows, dtype=torch.float32, device=weight_fp32.device)
+            for i in range(0, rows, chunk_rows):
+                qc, sc = quantize_int8_rows(
+                    rotate(weight_fp32[i : i + chunk_rows].float(), rot)
+                )
+                q[i : i + chunk_rows] = qc
+                scales[i : i + chunk_rows] = sc
         module.register_buffer("cr8_qdata", q, persistent=False)
         # fp32 scales stored as a uint8 byte view (see convrot4: module.to(dtype=...)
         # would otherwise cast them)
@@ -2270,6 +2293,10 @@ class ConvRotIntNQuantizer(ConvRotInt8Quantizer):
     weight grid and bitpacked storage. One instance per bit-width, shareable
     across modules."""
 
+    # this class's own quantize_ still expects the fp32 copy (the parent's
+    # chunked-cast path does not apply to it)
+    wants_fp32_weight = True
+
     def __init__(self, bits: int, rot_size: int = 256):
         super().__init__(rot_size=rot_size)
         self.bits = bits

toolkit/util/quantize.py

@@ -133,6 +133,7 @@ def quantize(
     include: Optional[Union[str, List[str]]] = None,
     exclude: Optional[Union[str, List[str]]] = None,
     quantize_device: Optional[torch.device] = None,
+    keep_on_quantize_device: bool = False,
 ):
     """Quantize the specified model submodules
 
@@ -163,6 +164,11 @@ def quantize(
             If provided, each module is moved to this device to quantize, then moved
             back to the device its weights were on initially. Lets a CPU-resident
             model (low vram) quantize layer-by-layer on the GPU.
+        keep_on_quantize_device (`bool`):
+            With quantize_device set: leave each layer on the quantize device after
+            quantizing instead of moving it back. The gpu transient is then one
+            bf16 layer at a time — never the whole unquantized remainder at once —
+            for models whose final home is that device.
     """
     if include is not None:
         include = [include] if isinstance(include, str) else include
@@ -237,7 +243,7 @@ def quantize(
                         optimizer=optimizer,
                     )
             finally:
-                if orig_device is not None:
+                if orig_device is not None and not keep_on_quantize_device:
                     # quanto replaces the module in its parent, so re-fetch by name
                     model.get_submodule(name).to(orig_device)
         except Exception as e:
@@ -361,10 +367,21 @@ def quantize_module(
 
     status_fn(" - quantizing extras")
     if keep_on_device:
-        # blocks already live on the gpu; bring the remainder over and run the
-        # extras pass there (fast kernels instead of an all-cores cpu burst)
+        # blocks already live on the gpu; quantize each remaining layer there
+        # one at a time and leave it — the transient is one bf16 layer, never
+        # the whole unquantized remainder at once (that transient bump could
+        # exceed the model's final footprint by GBs on big-extras models)
+        quantize(
+            module,
+            weights=quantization_type,
+            exclude=exclude,
+            quantize_device=device,
+            keep_on_quantize_device=True,
+            **quantize_kwargs,
+        )
+        # non-quantized leftovers (norms, embeddings, buffers) follow — final
+        # residency, not a transient
         module.to(device)
-        quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs)
     else:
         # cpu-resident model: quantize each extra layer with a gpu round-trip
         quantize(
@@ -395,15 +412,35 @@ def quantize_model(
     # sensitive modules to keep in full precision (fnmatch patterns)
     exclude_modules = base_model.get_quantization_exclude_modules() or []
 
-    if base_model.model_config.accuracy_recovery_adapter is not None:
+    mc = base_model.model_config
+    device = base_model.device_torch
+    keep_on_device = (
+        not mc.low_vram
+        and not (mc.layer_offloading and mc.layer_offloading_transformer_percent > 0)
+        and torch.device(device).type != "cpu"
+    )
+
+    if mc.accuracy_recovery_adapter is not None:
         attach_ara_and_quantize(
             base_model,
             model_to_quantize,
-            ara_path=base_model.model_config.accuracy_recovery_adapter,
+            ara_path=mc.accuracy_recovery_adapter,
             exclude=exclude_modules,
+            device=device,
+            keep_on_device=keep_on_device,
         )
     else:
-        _quantize_model_blocks(base_model, model_to_quantize, exclude_modules)
+        quantize_module(
+            model_to_quantize,
+            mc.qtype,
+            device=device,
+            dtype=base_model.torch_dtype,
+            block_names=base_model.get_transformer_block_names(),
+            exclude=exclude_modules,
+            quantize_kwargs=mc.quantize_kwargs,
+            status_fn=base_model.print_and_status_update,
+            keep_on_device=keep_on_device,
+        )
 
 
 @torch.no_grad()
@@ -412,11 +449,18 @@ def attach_ara_and_quantize(
     model_to_quantize: torch.nn.Module,
     ara_path: str,
     exclude: Optional[List[str]] = None,
+    device=None,
+    keep_on_device: bool = False,
 ):
     """Load an accuracy recovery adapter as a live network on the module and
     quantize around it (adapter-hijacked linears at the configured qtype,
     everything else uint8). The network lands on
-    base_model.accuracy_recovery_adapter."""
+    base_model.accuracy_recovery_adapter.
+
+    ``device``: quantize each hijacked linear there (gpu kernels) instead of
+    wherever it happens to live (historically the cpu — slow). With
+    ``keep_on_device`` the quantized linears stay there (final home is that
+    gpu); otherwise each returns to the device it came from."""
     from toolkit.config_modules import NetworkConfig
     from toolkit.lora_special import LoRASpecialNetwork
 
@@ -537,12 +581,19 @@ def attach_ara_and_quantize(
     base_model.accuracy_recovery_adapter = network
 
     # quantize it
+    keep_on_device = keep_on_device and device is not None
     lora_exclude_modules = []
     quantization_type = get_qtype(base_model.model_config.qtype)
     for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"):
         # the lora has already hijacked the original module
         orig_module = lora_module.org_module[0]
-        orig_module.to(base_model.torch_dtype)
+        orig_device = None
+        if device is not None:
+            param = next(orig_module.parameters(), None)
+            orig_device = param.device if param is not None else None
+            orig_module.to(device, dtype=base_model.torch_dtype)
+        else:
+            orig_module.to(base_model.torch_dtype)
         # make the params not require gradients
         for param in orig_module.parameters():
             param.requires_grad = False
@@ -550,78 +601,23 @@ def attach_ara_and_quantize(
         freeze(orig_module)
         module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '')
         lora_exclude_modules.append(module_name)
-        if base_model.model_config.low_vram:
-            # move it back to cpu
+        if not keep_on_device and orig_device is not None:
+            orig_module.to(orig_device)
+        elif base_model.model_config.low_vram and device is None:
+            # legacy behavior when no quantize device was given
             orig_module.to("cpu")
-    pass
     # quantize additional layers
     print_acc(" - quantizing additional layers")
     quantization_type = get_qtype('uint8')
     quantize(
         model_to_quantize,
         weights=quantization_type,
-        exclude=lora_exclude_modules + exclude_modules
+        exclude=lora_exclude_modules + exclude_modules,
+        quantize_device=device,
+        keep_on_quantize_device=keep_on_device,
     )
+    if keep_on_device:
+        # non-quantized leftovers follow — final residency, not a transient
+        model_to_quantize.to(device)
 
 
-def _quantize_model_blocks(base_model, model_to_quantize, exclude_modules):
-    if True:
-        # quantize model the original way without an accuracy recovery adapter
-        # move and quantize only certain pieces at a time.
-        quantization_type = get_qtype(base_model.model_config.qtype)
-        quantize_kwargs = base_model.model_config.quantize_kwargs or {}
-        # all_blocks = list(model_to_quantize.transformer_blocks)
-        all_blocks: List[torch.nn.Module] = []
-        transformer_block_names = base_model.get_transformer_block_names() or []
-        for name in transformer_block_names:
-            # name may be a dotted path for models that nest their blocks
-            # (e.g. hidream_o1's "model.language_model.layers").
-            block_list = model_to_quantize
-            for part in name.split('.'):
-                block_list = getattr(block_list, part, None)
-                if block_list is None:
-                    break
-            if block_list is not None:
-                all_blocks += list(block_list)
-        base_model.print_and_status_update(
-            f" - quantizing {len(all_blocks)} transformer blocks"
-        )
-        already_quantized = 0
-        for block in tqdm(all_blocks):
-            if not _has_quantizable_linear(block, quantization_type, exclude_modules):
-                # pre-quantized checkpoint with a matching qtype: nothing in
-                # this block would change — skip the device round-trip and the
-                # dtype cast entirely so the load stays byte-identical
-                already_quantized += 1
-                continue
-            block.to(base_model.device_torch, dtype=base_model.torch_dtype, non_blocking=True)
-            # exclude patterns with a leading wildcard (e.g. "*adaln_proj*")
-            # also apply inside blocks, where names are block-relative
-            quantize(
-                block,
-                weights=quantization_type,
-                exclude=exclude_modules,
-                **quantize_kwargs,
-            )
-            freeze(block)
-            # NOT non_blocking: an async D2H allocates the cpu destination in pinned
-            # memory, which the caching host allocator keeps forever (with power-of-2
-            # bucket rounding on top) — that silently retained a model-sized chunk of
-            # host ram after the weights moved back to the gpu for training
-            block.to("cpu")
-        if already_quantized:
-            base_model.print_and_status_update(
-                f" - {already_quantized} blocks already quantized with a matching qtype; left untouched"
-            )
-
-        # todo, on extras find a universal way to quantize them on device and move them back to their original
-        # device without having to move the transformer blocks to the device first
-        base_model.print_and_status_update(" - quantizing extras")
-        # model_to_quantize.to(base_model.device_torch, dtype=base_model.torch_dtype)
-        quantize(
-            model_to_quantize,
-            weights=quantization_type,
-            exclude=exclude_modules,
-            **quantize_kwargs,
-        )
-        freeze(model_to_quantize)