[Bug] Encountering torch.OutOfMemoryError During int8 Quantization

[SolomonLeon]

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.

Describe the bug

I'm attempting to quantize the DeepSeek-R1-Distill-Qwen-32B model on my server equipped with 4x V100 GPUs using the following command:

CUDA_VISIBLE_DEVICES=0,1,2,3 lmdeploy lite smooth_quant ./DeepSeek-R1-Distill-Qwen-32B --work-dir ./DeepSeek-R1-Distill-Qwen-32B-int8 --quant-dtype int8

However, I keep encountering a torch.OutOfMemoryError even though CUDA_VISIBLE_DEVICES is properly set. Could you please advise on how to successfully complete the quantization process?

Reproduction

CUDA_VISIBLE_DEVICES=0,1,2,3 lmdeploy lite smooth_quant ./DeepSeek-R1-Distill-Qwen-32B --work-dir ./DeepSeek-R1-Distill-Qwen-32B-int8 --quant-dtype int8

Environment

sys.platform: linux                                                               
Python: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0]         
CUDA available: True                                                              
MUSA available: False
numpy_random_seed: 2147483648            
GPU 0,1,2,3: Tesla V100-SXM2-16GB
CUDA_HOME: /usr                          
NVCC: Cuda compilation tools, release 12.0, V12.0.140                                                                                                               
GCC: gcc (Ubuntu 12.3.0-17ubuntu1) 12.3.0                                                                                                                           
PyTorch: 2.5.1+cu124                                                                                                                                                
PyTorch compiling details: PyTorch built with:                                                                                                                      
  - GCC 9.3                                                                       
  - C++ Version: 201703                                                           
  - Intel(R) oneAPI Math Kernel Library Version 2024.2-Product Build 20240605 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v3.5.3 (Git Hash 66f0cb9eb66affd2da3bf5f8d897376f04aae6af)
  - Intel(R) oneAPI Math Kernel Library Version 2024.2-Product Build 20240605 for Intel(R) 64 architecture applications                            09:07:15 [0/1798]
  - Intel(R) MKL-DNN v3.5.3 (Git Hash 66f0cb9eb66affd2da3bf5f8d897376f04aae6af)   
  - OpenMP 201511 (a.k.a. OpenMP 4.5)    
  - LAPACK is enabled (usually provided by MKL)                                   
  - NNPACK is enabled 
  - CPU capability usage: AVX2           
  - CUDA Runtime 12.4                                                             
  - NVCC architecture flags: -gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,co
de=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90                                                
  - CuDNN 90.1                           
  - Magma 2.6.1                                                                   
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=12.4, CUDNN_VERSION=9.1.0, CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= -D_GLI
BCXX_USE_CXX11_ABI=0 -fabi-version=11 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DLIBKINETO_NOXPUPTI=ON -DUSE_FBGE
MM -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation
 -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-strict-overflow -Wno-strict-aliasing
 -Wno-stringop-overflow -Wsuggest-override -Wno-psabi -Wno-error=old-style-cast -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-var
iable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, TORCH_V
ERSION=2.5.1, USE_CUDA=ON, USE_CUDNN=ON, USE_CUSPARSELT=1, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_GLOO=ON, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, U
SE_NCCL=1, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, USE_ROCM_KERNEL_ASSERT=OFF,                                                                                  
                                         
TorchVision: 0.20.1+cu124          
LMDeploy: 0.7.0.post3+     
transformers: 4.45.2                     
gradio: 5.1.0                            
fastapi: 0.115.2                         
pydantic: 2.9.2                          
triton: 3.1.0                            
NVIDIA Topology:                         
        GPU0    GPU1    GPU2    GPU3    CPU Affinity    NUMA Affinity   GPU NUMA ID                                                                                 
GPU0     X      NV1     NV1     NV2     0-15    0               N/A
GPU1    NV1      X      NV2     NV1     0-15    0               N/A               
GPU2    NV1     NV2      X      NV1     0-15    0               N/A
GPU3    NV2     NV1     NV1      X      0-15    0               N/A               
                                         
Legend:                                  
                                         
  X    = Self                            
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)                                                              
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node                                                        
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)                                                                               
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)                                                                      
  PIX  = Connection traversing at most a single PCIe bridge                       
  NV#  = Connection traversing a bonded set of # NVLinks

Error traceback

Traceback (most recent call last):                                                
  File "/home/username/miniconda3/envs/lmdeploy/bin/lmdeploy", line 8, in <module>
    sys.exit(run())                                                               
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/cli/entrypoint.py", line 39, in run
    args.run(args)                                                                
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/cli/lite.py", line 131, in smooth_quant
    smooth_quant(**kwargs)                                                        
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/lite/apis/smooth_quant.py", line 45, in smooth_quant
    vl_model, model, tokenizer, work_dir = calibrate(model,                       
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/lite/apis/calibrate.py", line 312, in calibrate
    calib_ctx.calibrate(all_data)                                                 
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/lmdeploy/lite/quantization/calibration.py", line 224, in calibrate
    _ = model(data.to(self.device))                                               
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)                                       
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)                                          
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 995, in forward
    hidden_states = self.norm(hidden_states)           
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)                                       
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1844, in _call_impl
    return inner()                                                                
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1790, in inner
    result = forward_call(*args, **kwargs)                                        
  File "/home/username/miniconda3/envs/lmdeploy/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 132, in forward
    variance = hidden_states.pow(2).mean(-1, keepdim=True)                        
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 5.00 GiB. GPU 0 has a total capacity of 15.77 GiB of which 3.93 GiB is free. Including non-PyTorch memory, this process has 11.84 GiB memory in use. Of the allocated memory 11.46 GiB is allocated by PyTorch, and 12.85 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[AllentDan]

We currently only support one device to do quantization. As for OOM case, may refer to https://lmdeploy.readthedocs.io/en/latest/quantization/w4a16.html#faqs