File size: 2,823 Bytes
d3bcf72 | 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 | #!/usr/bin/env python3
"""Block-level group offloading for a fused MiniMax-H3 transformer.
Fused blocks read their packed weight groups directly, so they stream as whole blocks
(``offload_type="block_level"``). Two parts of the model sit outside ``transformer_blocks``:
* the token refiner (two transformer blocks over the text rows, about 1.5 GB) depends only on
the prompt embeddings. It is computed once per prompt from its own synchronously offloaded
blocks and never occupies the GPU during the denoising forwards;
* the remaining top-level modules (embedders, output norm and projections) form the group that
Diffusers copies to the GPU and back synchronously around every forward. The forward needs
them on the GPU anyway, so they stay there for the whole denoising loop.
"""
from __future__ import annotations
import weakref
import torch
from diffusers.hooks import apply_group_offloading
class _LastInputCache:
"""``compute(tensor)`` for the last input tensor object, while its in-place version holds."""
def __init__(self, compute):
self.compute = compute
self.key = None
self.value = None
def __call__(self, tensor):
if (
self.key is not None
and self.key[0]() is tensor
and self.key[1] == tensor._version
):
return self.value
self.value = self.compute(tensor)
self.key = (weakref.ref(tensor), tensor._version)
return self.value
def enable_fused_group_offload(transformer, **group_offload_kwargs) -> None:
refiner = transformer.token_refiner
del transformer.token_refiner
transformer.enable_group_offload(**group_offload_kwargs)
apply_group_offloading(
refiner,
onload_device=group_offload_kwargs["onload_device"],
offload_device=group_offload_kwargs["offload_device"],
offload_type="block_level",
num_blocks_per_group=1,
)
transformer.token_refiner = refiner
# The denoiser passes the same prompt-embedding tensor at every step; the context embedder
# then returns the same tensor object, which the refiner cache recognizes.
transformer.context_embedder.forward = _LastInputCache(
transformer.context_embedder.forward
)
refiner.forward = _LastInputCache(refiner.forward)
def keep_unmatched_group_on_cuda_while_denoising(loop_wrapper_cls, transformer) -> None:
group = transformer._diffusers_hook.get_hook("group_offloading").group
loop_call = loop_wrapper_cls.__call__
def call(self, components, state):
group.offload_ = lambda: None
try:
return loop_call(self, components, state)
finally:
del group.offload_
group.offload_()
torch.cuda.empty_cache()
loop_wrapper_cls.__call__ = call
|