如何在24GB的GPU上运行DeepSeek-R1-Distill-Qwen-32B

最新推荐文章于 2025-03-27 10:05:27 发布

Hi20240217

最新推荐文章于 2025-03-27 10:05:27 发布

阅读量1.6k

点赞数 22

分类专栏：代码片段学习文章标签： DeepSeek-R1 Qwen-32B LLM GPU 内存优化推理

本文链接：https://blog.csdn.net/m0_61864577/article/details/145620659

版权

如何在24GB的GPU上运行DeepSeek-R1-Distill-Qwen-32B

一、背景

随着深度学习的不断发展，大型语言模型（LLM，Large Language Model）在自然语言处理领域展现出了强大的能力。然而，伴随着模型参数规模的指数级增长，运行这些模型所需的计算资源也变得异常庞大，尤其是对显存（GPU内存）的需求。因此，如何在有限的GPU显存下有效地运行超大规模的LLM，成为了一个亟待解决的挑战。

本文验证在GPU显存受限的情况下，如何高效地运行超出GPU内存容量的LLM模型。通过对模型权重的量化和内存管理策略的优化，期望能够突破硬件瓶颈，为大型模型的部署和应用提供新的思路。

二、解决方案

下面的方案，主要包括权重量化、内存缓存机制以及自定义Linear的设计。具体方案如下：

权重的INT4块量化
- 量化策略：将模型的权重参数进行INT4（4位整数）块量化处理，量化的块大小设定为128。这种量化方式能够大幅度减少模型权重所占用的存储空间。
- 内存优势：经过INT4量化后的权重占用空间显著降低，使得所有权重可以加载到主机（HOST）内存中。这不仅缓解了GPU显存的压力，还为后续的高效读取奠定了基础。
减少磁盘I/O操作
- 全量加载：将所有量化后的INT4权重一次性加载到HOST内存中，避免了在模型运行过程中频繁进行磁盘读写操作。这种方式有效减少了磁盘I/O带来的时间开销和性能瓶颈。
设备内存缓存机制
- 缓存设计：在GPU设备内存中建立一个缓存机制，设定最大缓存条目数为N。N的取值与具体的GPU配置相关，目的是充分利用可用的设备内存，最大化其占用率，提升数据读取效率。
- 动态管理：缓存机制需要智能地管理内存的分配和释放，确保在不超过设备内存上限的情况下，高效地存取所需的数据。
权重预加载线程
- 职责分离：引入一个专门的权重预加载线程，负责将HOST内存中的INT4权重进行反量化处理（即将INT4还原为计算所需的格式），并将处理后的权重加载到GPU设备内存的缓存中。
- 效率优化：通过预加载线程的异步处理，提升了数据准备的效率，确保模型在需要数据时可以及时获取，最大程度减少等待时间。
自定义Linear模块
- 模块替换：将原有的nn.Linear层替换为自定义的Module。在模型构建和加载过程中，使用该自定义模块来承载线性计算任务。
- 运行机制：自定义的Module在前向传播（forward）过程中，从设备内存的缓存中获取所需的权重进行计算。计算完成后，立即释放权重占用的设备内存，以供后续的计算任务使用。
- 优势：这种动态加载和释放的机制，避免了在整个计算过程中权重长时间占用设备内存，极大地提高了内存的利用效率。

三、操作步骤

1.下载模型

# 模型介绍: https://www.modelscope.cn/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B

# 下载模型
apt install git-lfs -y
git clone https://www.modelscope.cn/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B.git

2.安装依赖

MAX_JOBS=4 pip install flash-attn==2.3.6
pip install torch-tb-profiler

3.量化

cat > extract_weights.py << EOF
import torch
import os
from tqdm import tqdm
from glob import glob
import torch
import sys
from safetensors.torch import safe_open, save_file

def quantize_tensor_int4(tensor):
    """
    将bfloat16的Tensor按照块大小128进行量化为int4，并返回每个块的scale。

    参数：
    tensor (torch.Tensor): bfloat16类型的输入Tensor。

    返回：
    int4_tensor (torch.Tensor): 量化后的uint8类型的Tensor，存储int4值，每个元素包含两个int4值。
    scales (torch.Tensor): 每个块对应的bfloat16类型的scale值。
    """
    # 确保输入Tensor为bfloat16类型
    tensor = tensor.to(torch.bfloat16)

    # 将Tensor展平为一维
    flat_tensor = tensor.flatten()
    N = flat_tensor.numel()
    block_size = 128
    num_blocks = (N + block_size - 1) // block_size  # 计算块的数量

    # 计算每个元素的块索引
    indices = torch.arange(N, device=flat_tensor.device)
    block_indices = indices // block_size  # shape: [N]

    # 计算每个块的x_max
    abs_tensor = flat_tensor.abs()
    zeros_needed = num_blocks * block_size - N
    # 对张量进行填充，使其长度为num_blocks * block_size
    if zeros_needed > 0:
        padded_abs_tensor = torch.cat([abs_tensor, torch.zeros(zeros_needed, device=abs_tensor.device, dtype=abs_tensor.dtype)])
    else:
        padded_abs_tensor = abs_tensor

    reshaped_abs_tensor = padded_abs_tensor.view(num_blocks, block_size)
    x_max = reshaped_abs_tensor.max(dim=1).values  # shape: [num_blocks]

    # 处理x_max为0的情况，避免除以0
    x_max_nonzero = x_max.clone()
    x_max_nonzero[x_max_nonzero == 0] = 1.0  # 防止除以0

    # 计算scale
    scales = x_max_nonzero / 7.0  # shape: [num_blocks]
    scales = scales.to(torch.bfloat16)

    # 量化
    scales_expanded = scales[block_indices]  # shape: [N]
    q = torch.round(flat_tensor / scales_expanded).clamp(-8, 7).to(torch.int8)

    # 将有符号int4转换为无符号表示
    q_unsigned = q & 0x0F  # 将范围[-8,7]映射到[0,15]

    # 如果元素数量是奇数，补充一个零
    if N % 2 != 0:
        q_unsigned = torch.cat([q_unsigned, torch.zeros(1, dtype=torch.int8, device=q.device)])

    # 打包两个int4到一个uint8
    q_pairs = q_unsigned.view(-1, 2)
    int4_tensor = (q_pairs[:, 0].to(torch.uint8) << 4) | q_pairs[:, 1].to(torch.uint8)

    return int4_tensor, scales

torch.set_default_device("cuda")
if len(sys.argv)!=3:
    print(f"{
     sys.argv[0]} input_model_dir output_dir")
else:
    input_model_dir=sys.argv[1]
    output_dir=sys.argv[2]
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    state_dicts = {
   }
    for file_path in tqdm(glob(os.path.join(input_model_dir, "*.safetensors"))):
        with safe_open(file_path, framework="pt", device="cuda") as f:
            for name in f.keys():
                param: torch.Tensor = f.get_tensor(name)
                #print(name,param.shape,param.dtype)
                if "norm" in name or "embed" in name:
                    state_dicts[name] = param
                else:
                    if "weight" in name:
                        int4_tensor, scales=quantize_tensor_int4(param)
                        state_dict={
   }
                        state_