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

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):

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
   Fix TorchAO v1 group offloading with use_stream=True #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.