diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 8d5a11a..91bc886 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 6aacd39..fc84461 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -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) diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index 393cb0a..2c175b2 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -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() diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 647cd88..f85ca11 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -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) diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index cc023b4..1f770fb 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -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): diff --git a/extensions_built_in/diffusion_models/flux2/src/pipeline.py b/extensions_built_in/diffusion_models/flux2/src/pipeline.py index aec4dd0..fd7e619 100644 --- a/extensions_built_in/diffusion_models/flux2/src/pipeline.py +++ b/extensions_built_in/diffusion_models/flux2/src/pipeline.py @@ -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( diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index a67ba7f..ef053ff 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -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) diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 7e947d0..d2b6944 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -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() diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index 1a39802..9694420 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -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() diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 4c4be72..3abfe91 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -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() diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index fa43a58..b6e1463 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -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 = ( diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index fcc2ce3..0cc7c86 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -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() diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index a1f45b1..ef6ebb9 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -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 =====") diff --git a/toolkit/dequantize.py b/toolkit/dequantize.py index 54c8ec7..8587833 100644 --- a/toolkit/dequantize.py +++ b/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) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 910c426..19faf7f 100644 --- a/toolkit/memory_management/manager.py +++ b/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): diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index d559f18..402981c 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/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"] diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 773e26b..94560be 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/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 diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 35641e2..fb981b6 100644 --- a/toolkit/models/v2/_mixin.py +++ b/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): diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 73fde4c..f22c4ab 100644 --- a/toolkit/models/wan21/wan21.py +++ b/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) diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index ec0645a..6fc5b76 100644 --- a/toolkit/util/convrot_quant.py +++ b/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 diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index 0b55cbe..a4b2ad7 100644 --- a/toolkit/util/quantize.py +++ b/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)