[Bug]: --dataset opencode-instruct costs 70G ram before tuning

[wenhuach21]

Problem Description

(autoround) wenhuach@mlp-dgx-01:~/auto-round-main$ CUDA_VISIBLE_DEVICES=0 python3 -m auto_round /models/Qwen3-8B --dataset opencode-instruct
2026-09-21 10:40:04 INFO main.py L332: start to quantize /models/Qwen3-8B
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 399/399 [00:00<00:00, 4335.36it/s]
2026-09-21 10:40:05 INFO base.py L2143: format is not set, using default auto_round format.
2026-09-21 10:40:06 INFO base.py L1366: torch.compile is enabled
2026-09-21 10:40:07 INFO orchestrator.py L594: start to cache block inputs
2026-09-21 10:40:07 INFO calib_dataset.py L1213: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
Resolving data files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:00<00:00, 12037.38it/s]
Resolving data files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:00<00:00, 17215.17it/s]
'timed out' thrown while requesting GET https://huggingface.co/datasets/nvidia/OpenCodeInstruct/resolve/8f3ba5bafe4d6e8db46082cf7ae6741bc370604d/data/train-00007-of-00050.parquet
Retrying in 1s [Retry 1/5].
Map: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:01<00:00, 5274.66 examples/s]
Filter: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2155/2155 [00:01<00:00, 1945.36 examples/s]
Casting the dataset: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2155/2155 [01:43<00:00, 20.86 examples/s]
Resolving data files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:00<00:00, 12225.44it/s]
Resolving data files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:00<00:00, 17171.47it/s]
Map: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:01<00:00, 5361.83 examples/s]
Filter: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2155/2155 [00:01<00:00, 1938.98 examples/s]
Casting the dataset: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2155/2155 [00:11<00:00, 192.94 examples/s]
2026-09-21 10:44:40 INFO device.py L1560: 'peak_ram': 71.53GB
2026-09-21 10:44:40 INFO orchestrator.py L626: caching done
Quantizing model.layers.0: 0%| | 0/36 [00:01<?, ?it/s]quantized 7/7 layers in the block, loss iter 0: 7.124e-09 -> iter 192: 2.650e-09
2026-09-21 10:45:06 INFO device.py L1560: 'peak_ram': 83.22GB, 'peak_vram': 11.08GB
Quantizing model.layers.1: 3%|█████▋ | 1/36 [00:26<15:40, 26.87s/it]quantized 7/7 layers in the block, loss iter 0: 1.388e-08 -> iter 189: 3.604e-09

Reproduction Steps

~

Environment Information

No response

Error Logs

Additional Context

No response

[changwangss]

The high peak RAM comes from the shared calibration preprocessing path, so this is not specific to OpenCodeInstruct.

The original order is filter → select → cast. Filtering creates row indices, and reading those rows produces Arrow slices that still reference the underlying column buffer. When converting input_ids
from int32 to int64, the affected casting path converts that full buffer for each slice, causing large temporary allocations.

I reproduced the same behavior with synthetic inputs through the loaders for pile-10k, new-title-chinese, MBPP, AudioCaps, local JSON, and the ModelScope FineWeb-Edu path.

The fix moves casting before filtering and selection for regular datasets. Streaming datasets retain their existing order: lazy filtering, bounded materialization, then casting.

[changwangss]

For the OpenCodeInstruct example in this issue, the log shows 2,155 samples entering Casting the dataset. With the default sequence length of 2,048, their int64 token IDs should occupy only:

2,155 × 2,048 × 8 bytes ≈ 33.7 MiB

However, filter() creates row indices. Each resulting Arrow row slice still references the full token buffer. During the int32-to-int64 cast, that entire buffer is converted separately for each row.

With the default cast batch size of 1,000:

One batch: 33.7 MiB × 1,000 ≈ 32.9 GiB
Two overlapping batches: ≈ 65.8 GiB

Together with other process allocations, this can explain the reported 71.53 GB peak and the slow casting step.

The fix casts regular datasets before filtering and selecting rows. It is applied in the shared preprocessing path because other loaders can trigger the same behavior.

In a reduced reproduction using the same 2,155 samples, with sequence length 32, observed additional Arrow memory dropped from 1,053 MiB to 2.40 MiB, with identical output.