[Bug]: cannot quantize models larger than available VRAM using v0.10.0 unless low_gpu_mem_usage = True

[davedgd]

Problem Description

With release V0.9.7, I was able to begin quantization of google/gemma-3-27b-it on a 48GB GPU (RTX 6000 Ada) despite the model not fitting entirely into VRAM. Now, the exact same call in v0.10.0 fails when the model is moving to CPU:

2026-02-12 15:00:32 INFO autoround.py L165: using MLLM mode for multimodal model.
Loading weights: 100%|████████████████| 1247/1247 [00:00<00:00, 2167.15it/s, Materializing param=model.vision_tower.vision_model.post_layernorm.weight]
2026-02-12 15:01:03 INFO base.py L486: using torch.bfloat16 for quantization tuning
2026-02-12 15:01:03 WARNING formats.py L154: some layers are skipped quantization (shape not divisible by 32).
2026-02-12 15:01:03 INFO base.py L1739: start to cache block inputs
Traceback (most recent call last):
  File "/home/daved/project/code/convert_autoround.py", line 43, in <module>
    ar.quantize_and_save(output_dir = quant_path, format = "auto_round")
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 949, in quantize_and_save
    model, _ = self.quantize()
               ^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 1740, in quantize
    all_inputs = self.try_cache_inter_data_gpucpu(all_first_block_names, self.nsamples, layer_names=layer_names)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 2209, in try_cache_inter_data_gpucpu
    new_max_memory[device] = max_memory[device] * 0.9
                             ~~~~~~~~~~^^^^^^^^
KeyError: 0

The error occurs here:

auto-round/auto_round/compressors/base.py

Line 2190 in 81c8ee4

new_max_memory[device] = max_memory[device] * 0.9

The cause is due to a mismatch between the list of devices here:

auto-round/auto_round/compressors/base.py

Line 2172 in 81c8ee4

devices = parse_available_devices(self.device_map)

And the max memory produced here:

auto-round/auto_round/compressors/base.py

Line 2174 in 81c8ee4

max_memory = get_max_memory()

The former includes device 0 for the GPU, whereas the latter only picks up the CPU. Hence, there is no 0 key on the later line:

['cuda:0'] # devices
{'cpu': 143219032064} # max_memory

By manually setting low_gpu_mem_usage to True, this KeyError can be avoided, but this is suboptimal since it leads to performance reductions. Note that correcting the max_memory to include the gpu ('0') led to additional (different) errors.

Reproduction Steps

1. Run this command:
   auto-round
   --model model
   --scheme "W4A16"
   --format "auto_round"
2. Where model is any model larger than available VRAM (e.g., google/gemma-3-27b-it for a 48GB GPU)
3. Run as shown
4. See error above

Environment Information

• OS: Ubuntu 24.04
• Python version: 3.12
• AutoRound version: v0.10.0
• Hardware: RTX 6000 Ada 48GB

Error Logs

Additional Context

No response

[lvliang-intel]

@davedgd, apologies for the delay — we were out for the Chinese New Year holiday.
I've submitted a fix in PR #1457: #1457.
I was able to reproduce the issue locally using mocks and have verified the fix. Could you please give it a try and let us know if it resolves things on your end? Happy to follow up if anything comes up.

[davedgd]

@lvliang-intel: No worries at all -- happy belated holiday!

I tried the fix a few minutes ago, and unfortunately it still leads to a (different) crash (please see below).

2026-02-21 13:09:29 INFO base.py L486: using torch.bfloat16 for quantization tuning
2026-02-21 13:09:30 WARNING formats.py L154: some layers are skipped quantization (shape not divisible by 32).
2026-02-21 13:09:30 INFO base.py L1720: start to cache block inputs
2026-02-21 13:09:30 WARNING base.py L2189: Device 0 is not available in accelerate's reported memory (max_memory keys: ['cpu']). This GPU may be out of memory or unavailable. Skipping it.
2026-02-21 13:09:30 INFO base.py L2250: switch to cpu to cache block inputs
cache block inputs: 100%|██████████████████████████████████████████████████████████████████████████████████| 128/128 [00:00<00:00, 135.37it/s]
2026-02-21 13:09:33 INFO base.py L1735: caching done
  0%|                                                                                                                  | 0/62 [00:00<?, ?it/s]Traceback (most recent call last):
  File "/home/daved/project/code/convert_autoround.py", line 46, in <module>
    ar.quantize_and_save(output_dir = quant_path, format = "auto_round")
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 949, in quantize_and_save
    model, _ = self.quantize()
               ^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 1759, in quantize
    self._quantize_blocks(
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 3122, in _quantize_blocks
    input_ids, input_others = self._preprocess_block_inputs(inputs)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 3079, in _preprocess_block_inputs
    input_ids = to_device(input_ids, self.cache_device)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/utils/model.py", line 1056, in to_device
    input_res.append(to_device(inp, device))
                     ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/utils/model.py", line 1046, in to_device
    return input.to(device)
           ^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: out of memory
Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

  0%|                                                                                                                  | 0/62 [00:00<?, ?it/s]

As shown, the warning from your patch does appear, and the GPU does seem to free up its memory when the 0% progress appears, but a CUDA OOM still occurs.

Here is the GPU utilization at initial load of model (which is larger than available VRAM). The VRAM fills and then eventually the warning (Device 0 is not available in accelerate's reported memory...) appears:

After the warning appears about skipping GPU, the VRAM clears:

However, the OOM noted above still occurs.

In the meantime, "low_gpu_mem_usage": True resolves all issues.

[lvliang-intel]

@davedgd, applied a fix for the secondary issue.
When OOM triggers a CPU fallback, self.cache_device isn't updated accordingly, causing _preprocess_block_inputs to re-trigger GPU allocation. Fix: set self.cache_device = torch.device("cpu") in the except torch.OutOfMemoryError branch and remove Accelerate hooks to avoid residual side effects. Give it a try and let us know if it resolves your issue. Thanks.

[davedgd]

Thanks again, @lvliang-intel. Unfortunately, I'm still getting an OOM with this code (see below). I did some further digging, and it's worth noting that adding PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to the call fixes this problem both with your patches present and in the current PyPI 0.10.1 release version.

Here's the latest error without expandable_segments but with the latest fix:

2026-02-22 10:39:09 INFO base.py L486: using torch.bfloat16 for quantization tuning
2026-02-22 10:39:10 WARNING formats.py L154: some layers are skipped quantization (shape not divisible by 32).
2026-02-22 10:39:10 INFO base.py L1720: start to cache block inputs
2026-02-22 10:39:10 WARNING base.py L2189: Device 0 is not available in accelerate's reported memory (max_memory keys: ['cpu']). This device may be out of memory or unavailable. Skipping it.
2026-02-22 10:39:10 INFO base.py L2250: switch to cpu to cache block inputs
cache block inputs: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 128/128 [00:01<00:00, 122.56it/s]
2026-02-22 10:39:13 INFO base.py L1735: caching done
Quantizing model.language_model.layers.0:   0%|                                                                                             | 0/62 [00:00<?, ?it/s]Traceback (most recent call last):
  File "/home/daved/project/code/convert_autoround.py", line 46, in <module>
    ar.quantize_and_save(output_dir = quant_path, format = "auto_round")
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 949, in quantize_and_save
    model, _ = self.quantize()
               ^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 1759, in quantize
    self._quantize_blocks(
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 3139, in _quantize_blocks
    q_input, input_ids = self._quantize_block(
                         ^^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/auto_round/compressors/base.py", line 2829, in _quantize_block
    block = block.to(device)
            ^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1381, in to
    return self._apply(convert)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/nn/modules/module.py", line 933, in _apply
    module._apply(fn)
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/nn/modules/module.py", line 933, in _apply
    module._apply(fn)
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/nn/modules/module.py", line 964, in _apply
    param_applied = fn(param)
                    ^^^^^^^^^
  File "/home/daved/miniforge3/envs/autoround/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1367, in convert
    return t.to(
           ^^^^^
torch.AcceleratorError: CUDA error: out of memory
Search for `cudaErrorMemoryAllocation' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Quantizing model.language_model.layers.0:   0%|                                                                                             | 0/62 [00:00<?, ?it/s]

As noted above, with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, it runs fine however. In the case of low_gpu_mem_usage, it's not necessary to set the expandable_segments flag.

[lvliang-intel]

@davedgd,

Thank you for the detailed report and for identifying the critical workaround with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True — that was the key clue that led us to the root cause.

The possible root cause of this issue: During initialization, _only_text_test() attempts to move the entire model to GPU to check whether it accepts text-only input. For gemma-3-27b-it (~54GB in bf16), this exceeds the 48GB VRAM. The failed model.to("cuda:0") causes PyTorch's caching allocator to issue large cudaMalloc calls to expand its memory pools, which then fail and roll back. The reserved-then-freed VA ranges fragment CUDA's virtual address space, so subsequent allocations fail with cudaErrorMemoryAllocation — not because physical memory is exhausted, but because cudaMalloc can‘t find a contiguous VA region to map the requested block. This is exactly why expandable_segments:True fixes it: it uses cuMemAddressReserve + cuMemMap to allow logical segments to grow on demand by stitching non-contiguous physical pages into contiguous virtual addresses, bypassing the need to find a contiguous VA range upfront.

A fix has been applied that should address this. Please give it a try and let us know how it goes. Thanks.

[davedgd]

@lvliang-intel: Awesome -- I can confirm all is working as expected with this version! It's no longer necessary to specify low_gpu_mem_usage with this latest patch -- nor is expandable_segments necessary.

Thank you so much for looking into this so diligently!