huggingface / huggingface/diffusers

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

Offen
#14,433 2 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug group-offloading
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

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.

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Start in diffusers/hooks/group_offloading.py, especially _init_cpu_param_dict(), _to_cpu(), _pinned_memory_tensors(), and _onload_from_memory(); trace how low_cpu_mem_usage and use_stream select these paths. Reproduce with the supplied minimal Linear stack and torchao Int8WeightOnlyConfig(version=2). Done means the failing combination completes without attempting to pin CUDA qdata while the two working combinations and group offloading behavior remain intact.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
ai, performance
Issue-Typ
Bug
Schwierigkeit
3/5
Geschätzter Aufwand
1-2 Tage
Aktivitätsstatus
Ruhig
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
65/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.