github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Performance fixes and bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-28 22:21:09 -06:00
parent 5ddc5f8ca7
commit 7195abc32b
21 changed files with 511 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -206,11 +206,13 @@ class ChromaModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

8
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -192,11 +192,13 @@ class ChromaRadianceModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

6
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
View File

@@ -133,8 +133,10 @@ class ErnieImageModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/f_light/f_light.py
View File

@@ -115,8 +115,10 @@ class FLiteModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

13
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

@@ -98,12 +98,21 @@ class Flux2Model(BaseModel):
self.print_and_status_update("Loading Mistral")
# load + quantize + offload + placement, all driven by model_config
# tie_word_embeddings=False: the checkpoint carries both embed_tokens
# and lm_head with different values; the config's tie claim is wrong
text_encoder = Mistral3TextEncoder.load(
MISTRAL_PATH, subfolder="", **self.component_load_kwargs("te")
MISTRAL_PATH,
subfolder="",
tie_word_embeddings=False,
**self.component_load_kwargs("te"),
)
flush()
tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH)
# fix_mistral_regex=False: keep the exact tokenization flux2 has always
# used (True would change the pre-tokenizer and shift conditioning)
tokenizer = AutoProcessor.from_pretrained(
MISTRAL_PATH, fix_mistral_regex=False
)
return text_encoder, tokenizer
def load_model(self):

10
extensions_built_in/diffusion_models/flux2/src/pipeline.py
View File

@@ -109,15 +109,19 @@ class Flux2Pipeline(DiffusionPipeline):
# Process all messages at once
# with image processing a too short max length can throw an error in here.
try:
# tokenization kwargs ride in processor_kwargs (same values end up
# in the same place; loose **kwargs just warn on new transformers)
inputs = self.tokenizer.apply_chat_template(
messages_batch,
add_generation_prompt=False,
tokenize=True,
return_dict=True,
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=max_sequence_length,
processor_kwargs={
"padding": "max_length",
"truncation": True,
"max_length": max_sequence_length,
},
)
except ValueError as e:
print(

8
extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
View File

@@ -137,11 +137,13 @@ class FluxKontextModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

6
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -490,8 +490,10 @@ class LTX2Model(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py
View File

@@ -144,8 +144,10 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -161,8 +161,10 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

20
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@@ -95,6 +95,13 @@ class DualWanTransformer3DModel(torch.nn.Module):
def dtype(self) -> torch.dtype:
return self.torch_dtype
def get_offload_ignore_modules(self):
# both DiTs' fp32 modulation tables stay resident on the compute device
return (
self.transformer_1.get_offload_ignore_modules()
+ self.transformer_2.get_offload_ignore_modules()
)
@property
def config(self):
return self.transformer_1.config
@@ -128,17 +135,26 @@ class DualWanTransformer3DModel(torch.nn.Module):
else:
t_name = "transformer_2"
# memory-managed transformers own their placement (layers bounce
# from cpu per forward); whole-model swaps would haul the full
# 14B up and defeat the offloading
managed = (
hasattr(self, "_memory_manager")
or hasattr(self.transformer_1, "_memory_manager")
or hasattr(self.transformer_2, "_memory_manager")
)
# check if we are changing the active transformer, if so, we need to swap the one in
# vram if low_vram is enabled
# todo swap the loras as well
if t_name != self._active_transformer_name:
if self.low_vram:
if self.low_vram and not managed:
getattr(self, self._active_transformer_name).to("cpu")
getattr(self, t_name).to(self.device_torch)
torch.cuda.empty_cache()
self._active_transformer_name = t_name
if self.transformer.device != hidden_states.device:
if not managed and self.transformer.device != hidden_states.device:
if self.low_vram:
# move other transformer to cpu
other_tname = (

6
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -149,8 +149,10 @@ class ZetaChromaModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].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.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

77
testing/test_model_loading.py
View File

@@ -338,6 +338,71 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict:
gen_profile.disable()
gen_profile.dump_stats(os.path.join(out_dir, "gen.prof"))
# ---- 100% offload forward: attach full layer offloading to the DiT and
# text encoder(s), then run a 1-step sample so the offload path (TE
# encode + denoise + decode with every layer staged from cpu) is proven
# for every arch ----
import torch
from toolkit.memory_management import MemoryManager
prof.stage("offload 100%: attach")
offload_targets = []
model_module = getattr(sd, "model", None)
if model_module is not None:
subs = [
getattr(model_module, name, None)
for name in ("transformer_1", "transformer_2")
]
if all(s is not None for s in subs):
# dual-DiT wrappers (wan22 14b): pipelines hold and .to() the
# sub-transformers directly, so each needs its own manager
offload_targets.extend(subs)
else:
offload_targets.append(model_module)
tes = getattr(sd, "text_encoder", None)
for te in tes if isinstance(tes, list) else [tes]:
offload_targets.append(te)
# auxiliary conditioning stacks where archs have them (ltx connectors,
# i2v vision towers)
offload_targets.append(getattr(sd, "image_encoder", None))
offload_targets.append(getattr(getattr(sd, "pipeline", None), "connectors", None))
attached = 0
for m in offload_targets:
if not isinstance(m, torch.nn.Module):
continue
if type(m).__name__.startswith("Fake"):
continue
if next(m.parameters(), None) is None:
continue
get_ignore = getattr(m, "get_offload_ignore_modules", None)
ignore = get_ignore() if callable(get_ignore) else None
MemoryManager.attach(
m,
torch.device(device),
offload_percent=1.0,
ignore_modules=list(ignore or []),
)
attached += 1
prof.stage("offload 100%: generate")
offload_profile = cProfile.Profile()
offload_profile.enable()
t_off0 = time.perf_counter()
# 2 steps, not 1: diffusers' shift_terminal stretch divides by
# one_minus_z[-1], which is zero for a single step — NaN timesteps and a
# NaN image would mask real numerical breakage in the offload path
gen_offload = GenerateImageConfig(
prompt="a photo of a cat sitting on a wooden table",
output_folder=out_dir,
output_ext="png",
**{**sample_kwargs, "num_inference_steps": 2},
)
sd.generate_images([gen_offload], **gen_kwargs)
offload_seconds = time.perf_counter() - t_off0
offload_profile.disable()
offload_profile.dump_stats(os.path.join(out_dir, "offload.prof"))
stages = prof.finish()
produced = [
@@ -350,8 +415,6 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict:
if not produced:
raise RuntimeError(f"no output file produced in {out_dir}")
import torch
result = {
"arch": arch,
"status": "PASS",
@@ -360,12 +423,16 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict:
"qtype_te": model_config.qtype_te if model_kwargs.get("quantize_te") else None,
"load_seconds": round(load_seconds, 3),
"gen_seconds": round(gen_seconds, 3),
"offload_seconds": round(offload_seconds, 3),
"offload_modules_attached": attached,
"stages": stages,
"profile": {
"load_top": profile_top(load_profile),
"gen_top": profile_top(gen_profile, limit=15),
"offload_top": profile_top(offload_profile, limit=15),
"load_prof": os.path.join(out_dir, "load.prof"),
"gen_prof": os.path.join(out_dir, "gen.prof"),
"offload_prof": os.path.join(out_dir, "offload.prof"),
},
"env": {
"torch_num_threads": torch.get_num_threads(),
@@ -577,7 +644,8 @@ def main():
total = len(MODEL_TESTS)
sweep_t0 = time.perf_counter()
for i, arch in enumerate(MODEL_TESTS, start=1):
print(f"\n===== {arch} ({i}/{total}) =====")
# flush so tail -f shows headers/progress live when stdout is a file
print(f"\n===== {arch} ({i}/{total}) =====", flush=True)
result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json")
cmd = [
sys.executable,
@@ -610,7 +678,8 @@ def main():
print(
f">>> progress: {i}/{total} ({i * 100 // total}%) — "
f"{tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip — "
f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m"
f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m",
flush=True,
)
print("\n===== summary =====")

4
toolkit/dequantize.py
View File

@@ -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)

92
toolkit/memory_management/manager.py
View File

@@ -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):

52
toolkit/memory_management/manager_modules.py
View File

@@ -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"]

68
toolkit/models/v2/PLANNING.md
View File

@@ -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

69
toolkit/models/v2/_mixin.py
View File

@@ -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):

5
toolkit/models/wan21/wan21.py
View File

@@ -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)

35
toolkit/util/convrot_quant.py
View File

@@ -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

144
toolkit/util/quantize.py
View File

@@ -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)