huggingface / huggingface/diffusers

Title: Group offloading with use_stream=True fails on torchao int8 version=2 (pin_memory on CUDA qdata)

Abierto
#14,433 2 comentarios 0 reacciones 0 asignados Ver en GitHub
bug group-offloading
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

Title: Group offloading with use_stream=True fails on torchao int8 version=2 (pin_memory on CUDA qdata)

> Context: found while building an app that runs MiniMax-H3 through diffusers. The
> investigation was done with an AI coding agent. The reproduction below was executed on my
> own hardware and the output is pasted verbatim, and I am accountable for its accuracy.

### Describe the bug

`apply_group_offloading(...)` fails deterministically on the first forward when
`use_stream=True` and `low_cpu_mem_usage=True` are combined and the module's weights are
torchao `Int8Tensor` produced by `Int8WeightOnlyConfig(version=2)`. Both flags are
documented options, and the pairing is the natural one for a memory-constrained setup
(streamed prefetch + reduced host RAM):

```
RuntimeError: cannot pin 'torch.cuda.CharTensor' only dense CPU tensors can be pinned
```

This is a different failure from the one fixed in #14112. That PR was explicitly scoped to
the v1 `AffineQuantizedTensor` path and states:

> the int8 compile path that uses TorchAO `version=2` [...] support `is_pinned()` and
> `pin_memory()` on current main

`Int8Tensor` does implement `pin_memory()` — that part holds. The problem is what its
implementation pins: it forwards to the inner `qdata`, and at that point `qdata` is on
CUDA, so the pin raises. So v2 was left on the unconditional-pinning path on the
assumption that pinning works, and it does not.

### Reproduction

Self-contained, no pretrained weights (a 256-dim, 4-block Linear stack is enough):

```python
import torch
from diffusers.hooks.group_offloading import apply_group_offloading
from torchao.quantization import Int8WeightOnlyConfig, quantize_

def build_stack(num_blocks=4, dim=256):
class Block(torch.nn.Module):
def __init__(self):
super().__init__()
self.lin1 = torch.nn.Linear(dim, dim * 2, bias=False, dtype=torch.bfloat16)
self.lin2 = torch.nn.Linear(dim * 2, dim, bias=False, dtype=torch.bfloat16)

def forward(self, x):
return self.lin2(torch.nn.functional.gelu(self.lin1(x)))

class Stack(torch.nn.Module):
_supports_group_offloading = True

def __init__(self):
super().__init__()
self.blocks = torch.nn.ModuleList([Block() for _ in range(num_blocks)])

def forward(self, x):
for b in self.blocks:
x = b(x)
return x

stack = Stack()
for block in stack.blocks:
quantize_(block, Int8WeightOnlyConfig(version=2))
return stack

def run(use_stream, low_cpu_mem_usage):
stack = build_stack()
apply_group_offloading(
stack,
onload_device=torch.device("cuda"),
offload_device=torch.device("cpu"),
offload_type="block_level",
num_blocks_per_group=1,
use_stream=use_stream,
low_cpu_mem_usage=low_cpu_mem_usage,
)
x = torch.randn(2, 256, dtype=torch.bfloat16, device="cuda")
for _ in range(3):
with torch.no_grad():
stack(x)
return "OK"

for label, stream, low_cpu in [
("use_stream=True, low_cpu_mem_usage=True ", True, True),
("use_stream=False, low_cpu_mem_usage=True ", False, True),
("use_stream=True, low_cpu_mem_usage=False", True, False),
]:
try:
print(f"{label} -> {run(stream, low_cpu)}", flush=True)
except Exception as e:
print(f"{label} -> {type(e).__name__}: {e}", flush=True)
```

Output:

```
use_stream=True, low_cpu_mem_usage=True -> RuntimeError: cannot pin 'torch.cuda.CharTensor' only dense CPU tensors can be pinned
use_stream=False, low_cpu_mem_usage=True -> OK
use_stream=True, low_cpu_mem_usage=False -> OK
```

### Traceback (tail; paths shortened)

```
File "diffusers/hooks/group_offloading.py", line 289, in _onload_from_memory
with self._pinned_memory_tensors() as pinned_memory:
File "diffusers/hooks/group_offloading.py", line 205, in _pinned_memory_tensors
param: tensor.pin_memory() if not tensor.is_pinned() else tensor
File "torchao/quantization/quantize_/workflows/int8/int8_tensor.py", line 427, in _
pinned_qdata = args[0].qdata.pin_memory()
RuntimeError: cannot pin 'torch.cuda.CharTensor' only dense CPU tensors can be pinned
```

### Analysis

Two places disagree about who pins, and the disagreement is only observable for tensor
subclasses whose payload lives in a member:

- `_init_cpu_param_dict()` runs once at enable time and, via `_to_cpu()`, **skips** pinning
when `low_cpu_mem_usage=True` (`return t if low_cpu_mem_usage else t.pin_memory()`).
- `_pinned_memory_tensors()` is entered from `_onload_from_memory()` on **every step** when
`use_stream=True`, and pins unconditionally.

For a plain tensor the second path is merely redundant work. For `Int8Tensor` it reaches
`qdata.pin_memory()` on data that is not on the host, and raises.

The two working combinations in the output above are consistent with this: `use_stream=False`
never enters `_pinned_memory_tensors()` at all, and `low_cpu_mem_usage=False` pins once at
enable time, while the module is still entirely on CPU, so the per-step path finds
`is_pinned() == True` and skips.

### Suggested fix

Any of these would close the gap:

1. Make `_pinned_memory_tensors()` honour `low_cpu_mem_usage` (skip when the init path
skipped), mirroring the two functions' intent.
2. Wrap the pin in a try/except and fall back to the unpinned CPU copy — this is what
#14112's description says it does for tensors whose pinning ops are unavailable, but
the current code has no such fallback for this path.
3. At minimum, reject the combination up front instead of failing mid-denoise, since the
failure surfaces far from its cause.

Worth noting for (1): `use_stream=True, low_cpu_mem_usage=False` is not just "the
workaround" — in my measurements on a 33B transformer it is **4–5x faster to onload**
(0.04–0.07s vs 0.1–0.26s per block), because pinned memory can't be paged out. The cost is
~14–16GB of page-locked host RAM and ~22s spent pinning at enable time. Users who hit this
crash and retreat to `use_stream=False` lose that speedup.

### System Info

- diffusers: `main` (verified `hooks/group_offloading.py` is byte-identical to the commit
I tested, f37ab93)
- torch 2.9.0+cu128, torchao 0.17.0, transformers 5.14.1
- GPU: NVIDIA RTX PRO 6000 Blackwell (sm_120), Ubuntu 24.04, Python 3.12
- I have only tested torchao 0.17.0; I have not checked whether a newer torchao changes
`Int8Tensor.pin_memory()`'s behaviour here.

Originally hit while running MiniMax-H3 (33B) with the transformer int8-quantized and
group-offloaded; the reproduction above is the reduced case. Full context and the probe
scripts: https://github.com/animede/Diffusers_minimax-h3

### Who can help?

@sayakpaul @a-r-r-o-w — the issue template has no line for group offloading /
quantization, so I picked based on recent commits to `hooks/group_offloading.py`.
Apologies if that routes to the wrong person.

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

Comienza en diffusers/hooks/group_offloading.py, especialmente en _init_cpu_param_dict(), _to_cpu(), _pinned_memory_tensors() y _onload_from_memory(); sigue cómo low_cpu_mem_usage y use_stream seleccionan estas rutas. Reproduce el problema con la pila Linear mínima proporcionada y torchao Int8WeightOnlyConfig(version=2). Se considera terminado cuando la combinación que falla se completa sin intentar fijar qdata de CUDA, mientras las dos combinaciones que funcionan y el comportamiento de group offloading permanecen intactos.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
ai, performance
Tipo de issue
Error
Dificultad
3/5
Tiempo estimado
1-2 días
Estado de actividad
Tranquilo
Claridad
Bastante claro
Aptitud para principiantes
65/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.