File size: 5,114 Bytes
a5e8975 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Device-only offload workaround for the pinned bitsandbytes 0.50.2 stack.
Parent Module.to() recursively invokes _apply(), not each child's to().
Linear8bitLt.to() moves auxiliary tensors, but this path is bypassed by a
parent move. Keep CB aliases and SCB/outlier state colocated with the weight.
This is instance-local; checkpoint tensors and installed packages are unchanged.
"""
from types import MethodType
import inspect
from weakref import WeakSet
import torch
def check_sampling_support(config):
"""Fail before loading weights when an old runtime would ignore a saved grid."""
if config.get('sample_sigmas') is not None:
from diffusers import QwenImage21Pipeline
if 'sample_sigmas' not in inspect.signature(QwenImage21Pipeline.__init__).parameters:
raise RuntimeError('This checkpoint requires Diffusers sampling-grid support (PR #14950)')
def release_prefix_cache_before_decode(pipe):
"""Release finished denoising prefix tensors before the full-resolution VAE.
The pinned upstream pipeline holds its local KV containers through decode.
Their tensors are no longer read after the final transformer call. Clearing
them at decode entry preserves the denoising computation and latent input.
Weak references avoid retaining caches on latent-only calls or failures.
"""
if getattr(pipe, '_prefix_cache_release', False):
return
caches = WeakSet()
telemetry = {'calls': 0, 'tensor_bytes': 0}
def capture(_module, _args, kwargs):
cache = kwargs.get('kv_cache')
if cache is not None:
caches.add(cache)
handle = pipe.transformer.register_forward_pre_hook(capture, with_kwargs=True)
original = pipe.vae.decode
def decode(self, *args, **kwargs):
released = 0
for cache in list(caches):
for layer in cache.layer_caches:
released += sum(t.numel() * t.element_size() for t in (layer.k, layer.v)
if isinstance(t, torch.Tensor))
layer.k = None
layer.v = None
caches.clear()
telemetry.update(calls=telemetry['calls'] + 1, tensor_bytes=released)
return original(*args, **kwargs)
pipe.vae.decode = MethodType(decode, pipe.vae)
pipe._prefix_cache_release_hook = handle
pipe._prefix_cache_release = True
pipe._prefix_cache_telemetry = telemetry
def patch_int8_device_moves(model):
from bitsandbytes.nn import Linear8bitLt
for layer in model.modules():
if not isinstance(layer, Linear8bitLt) or getattr(layer, '_int8_move_patched', False):
continue
if layer.state.has_fp16_weights:
raise ValueError('Offload helper only supports inference INT8 weights')
original = layer._apply
def apply(self, fn, recurse=True, _original=original):
result = _original(fn, recurse=recurse)
device = self.weight.device
# Parameter auxiliaries are used before init_8bit_state(); state
# auxiliaries are used afterwards. Never duplicate the CB storage.
if self.weight.CB is not None:
self.weight.CB = self.weight.data
if self.weight.SCB is not None:
self.weight.SCB = self.weight.SCB.to(device)
for name, value in vars(self.state).items():
if isinstance(value, torch.Tensor):
setattr(self.state, name, self.weight.data if name == 'CB' else value.to(device))
return result
layer._apply = MethodType(apply, layer)
layer._int8_move_patched = True
def enable_int8_cpu_offload(pipe):
"""Install auxiliary-tensor movement before standard model CPU offload."""
for name in ('text_encoder', 'transformer'):
patch_int8_device_moves(getattr(pipe, name))
pipe.enable_model_cpu_offload()
def load_int8_pipeline(model, local_files_only=False):
"""Load/offload components sequentially to avoid a combined CUDA load peak."""
from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel
from transformers import Qwen3VLForConditionalGeneration
check_sampling_support(QwenImage21Pipeline.load_config(model, local_files_only=local_files_only))
components = {}
for name, cls in [('transformer', QwenImage21Transformer2DModel),
('text_encoder', Qwen3VLForConditionalGeneration)]:
component = cls.from_pretrained(model, subfolder=name, dtype=torch.bfloat16,
local_files_only=local_files_only)
if not getattr(component, 'is_loaded_in_8bit', False):
raise ValueError('Expected saved bitsandbytes INT8 component: '+name)
patch_int8_device_moves(component)
component.to('cpu')
torch.cuda.empty_cache()
components[name] = component
pipe = QwenImage21Pipeline.from_pretrained(model, dtype=torch.bfloat16,
local_files_only=local_files_only, **components)
enable_int8_cpu_offload(pipe)
release_prefix_cache_before_decode(pipe)
return pipe
|