[Bug]: imbalanced memory usages across cards

[wenhuach21]

Problem Description

CUDA_VISIBLE_DEVICES=2,3,4,7 python3 -m auto_round ../GLM-5.3-Flash-BF16/ --device_map auto --batch_size 4 --gradient_accumulate_steps 2 --disable_torch_compile
2026-08-27 17:35:24 INFO main.py L294: torch.compile is explicitly disabled with --disable_torch_compile.
2026-08-27 17:35:24 INFO main.py L299: start to quantize ../GLM-5.3-Flash-BF16
Loading weights: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1609/1609 [02:30<00:00, 10.68it/s]
[transformers] The following generation flags are not valid and may be ignored: ['top_p']. Set TRANSFORMERS_VERBOSITY=info for more details.
2026-08-27 17:38:16 INFO base.py L2041: format is not set, using default auto_round format.
2026-08-27 17:38:16 INFO replace_modules.py L128: Experts (before replacement) [model.language_model.layers.3.mlp.experts] (Glm5NextTextExperts):
Glm5NextTextExperts()
[transformers] loss_type=None was set in the config but it is unrecognized. Using the default loss: ForCausalLMLoss.
2026-08-27 17:38:16 INFO device.py L1446: Before applying custom replacements 'peak_ram': 780.18GB
2026-08-27 17:39:48 INFO moe_experts_interface.py L681: [MoE Prep] Unfused 42 MOE experts modules
2026-08-27 17:39:48 INFO device.py L1446: After applying custom replacements 'peak_ram': 780.23GB
2026-08-27 17:39:48 INFO replace_modules.py L101: Prepared 42 MOE modules for quantization
2026-08-27 17:39:48 INFO replace_modules.py L128: Experts (after replacement/skip) [model.language_model.layers.3.mlp.experts] (Glm5NextTextExperts):
Glm5NextTextExperts(
(0-287): 288 x _ExpertContainer(
(down_proj): Linear(in_features=2048, out_features=4096, bias=False)
(gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
2026-08-27 17:39:48 INFO replace_modules.py L101: Prepared 42 MOE modules for quantization
2026-08-27 17:39:48 INFO replace_modules.py L128: Experts (after replacement/skip) [model.language_model.layers.3.mlp.experts] (Glm5NextTextExperts):
Glm5NextTextExperts(
(0-287): 288 x _ExpertContainer(
(down_proj): Linear(in_features=2048, out_features=4096, bias=False)
(gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
(up_proj): Linear(in_features=4096, out_features=2048, bias=False)
)
)
2026-08-27 17:39:48 INFO base.py L964: Using predefined ignore_layers: indexer, layers.0.mlp, weights_proj
2026-08-27 17:39:48 INFO layer_config_resolver.py L365: Ignored layers: indexer, layers.0.mlp, model.language_model.layers.0.mlp.down_proj, model.language_model.layers.0.mlp.gate_proj, model.language_model.layers.0.mlp.up_proj, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.weights_proj, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.wk, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.wq_b, weights_proj
2026-08-27 17:39:50 INFO layer_config_resolver.py L365: Ignored layers: indexer, layers.0.mlp, model.language_model.layers.0.mlp.down_proj, model.language_model.layers.0.mlp.gate_proj, model.language_model.layers.0.mlp.up_proj, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.weights_proj, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.wk, model.language_model.layers.[3,7,11,15,19,23,27,31,35,39,43].self_attn.indexer.wq_b, weights_proj
2026-08-27 17:39:53 INFO base.py L1184: 'enable_torch_compile' is disabled. Enabling it can reduce tuning cost by about 20%.
2026-08-27 17:39:53 INFO base.py L1288: torch.compile is disabled, as the user disabled it
2026-08-27 17:39:57 INFO orchestrator.py L570: start to cache block inputs
2026-08-27 17:39:57 INFO mllm.py L86: Using MLLM template: glm5_next
2026-08-27 17:39:57 INFO calib_dataset.py L1113: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
2026-08-27 17:40:30 INFO device.py L1448: 'peak_ram': 788.91GB
2026-08-27 17:40:30 INFO orchestrator.py L602: caching done
Quantizing model.language_model.layers.0: 0%| | 0/45 [00:03<?, ?it/s]2026-08-27 17:45:38 INFO quantizer.py L545: Unquantized layers: ['mlp.gate_proj', 'mlp.up_proj', 'mlp.down_proj']
quantized 9/12 layers in the block, loss iter 0: 3.632e-07 -> iter 195: 1.290e-07
2026-08-27 17:45:42 INFO device.py L1448: 'peak_ram': 788.91GB, 'peak_vram': {'0': 128.04GB, '1': 1.16GB, '2': 1.25GB, '3': 1.22GB}
Quantizing model.language_model.layers.1: 2%|████▎ | 1/45 [05:13<3:49:33, 313.04s/it]quantized 12/12 layers in the block, loss iter 0: 2.632e-07 -> iter 182: 9.384e-08
2026-08-27 17:50:53 INFO device.py L1448: 'peak_ram': 788.91GB, 'peak_vram': {'0': 134.08GB, '1': 2.35GB, '2': 2.36GB, '3': 2.04GB}
Quantizing model.language_model.layers.2: 4%|████████▌ | 2/45 [10:23<3:43:26, 311.78s/it]quantized 12/12 layers in the block, loss iter 0: 3.801e-06 -> iter 182: 1.298e-06
2026-08-27 17:56:04 INFO device.py L1448: 'peak_ram': 788.91GB, 'peak_vram': {'0': 134.08GB, '1': 2.35GB, '2': 2.36GB, '3': 2.04GB}
Quantizing model.language_model.layers.3: 7%|████████████▊ | 3/45 [15:34<3:37:57, 311.38s/it]/software/wenhuach/envs/wenhuach/lib/python3.12/site-packages/torch/autograd/graph.py:979: UserWarning: Memory Efficient attention defaults to a non-deterministic algorithm. To explicitly enable determinism call torch.use_deterministic_algorithms(True, warn_only=False). (Triggered internally at /__w/pytorch/pytorch/aten/src/ATen/native/transformers/cuda/attention_backward.cu:833.)
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
2026-08-27 18:03:32 INFO quantizer.py L545: Unquantized layers: ['self_attn.indexer.wq_b', 'self_attn.indexer.wk', 'self_attn.indexer.weights_proj']
quantized 872/875 layers in the block, loss iter 0: 1.985e-05 -> iter 144: 6.055e-06
2026-08-27 18:03:48 INFO device.py L1448: 'peak_ram': 788.91GB, 'peak_vram': {'0': 134.08GB, '1': 48.85GB, '2': 46.25GB, '3': 46.65GB}
Quantizing model.language_model.layers.4: 9%|█████████████████ | 4/45 [24:04<4:26:08, 389.48s/it]quantized 876/876 layers in the block, loss iter 0: 8.160e-06 -> iter 116: 4.181e-06
2026-08-27 18:17:29 INFO device.py L1448: 'peak_ram': 788.91GB, 'peak_vram': {'0': 169.61GB, '1': 50.4GB, '2': 47.6GB, '3': 47.67GB}
Quantizing model.language_model.layers.5: 11%|█████████████████████▎ | 5/45 [38:34<6:15:12, 562.82s/it][W827 18:19:44.282443104 CUDACachingAllocator.cpp:3933] memory allocation failed with OOM on device 0 while trying to allocate 17179869184 bytes (free: 16141582336, total: 191503007744).
[W827 18:19:44.427305568 CUDACachingAllocator.cpp:3933] memory allocation failed with OOM on device 0 while trying to allocate 17179869184 bytes (free: 16787505152, total: 191503007744).
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/home/wenhuach/auto-round/auto_round/main.py", line 20, in
run()
File "/home/wenhuach/auto-round/auto_round/cli/main.py", line 513, in run
start(argv=command_argv)
File "/home/wenhuach/auto-round/auto_round/cli/main.py", line 244, in start
tune(args)
File "/home/wenhuach/auto-round/auto_round/cli/main.py", line 408, in tune
model, folders = autoround.quantize_and_save( # pylint: disable=no-member
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/wenhuach/auto-round/auto_round/compressors/base.py", line 2089, in quantize_and_save
self.output_dir = output_dir
^^^^^^^^^^^^^^^
File "/home/wenhuach/auto-round/auto_round/compressors/orchestrator.py", line 400, in quantize
return self._quantize_data_driven()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/wenhuach/auto-round/auto_round/compressors/orchestrator.py", line 725, in _quantize_data_driven
self._quantize_blocks(

Reproduction Steps

~

Environment Information

No response

Error Logs

Additional Context

No response