cover_image

AI Infra之模型显存管理分析

望雲、十雨、言澈 阿里云开发者
2025年04月07日 00:30
图片

阿里妹导读


本文围绕某线上客户部署DeepSeek-R1满血版模型时进行多次压测后,发现显存占用一直上升,从未下降的现象,记录了排查过程。


背景

某线上客户部署DeepSeek-R1满血版模型,进行多次压测后,发现显存占用一直上升,从未下降,于是发出了灵魂三问:

1. 增加的显存被谁消耗?
2. 显存增加是否会达到一个阈值后就不再增加?
3. 如果会一直增加,有什么办法可以释放显存?

在分析问题之前,必须要了解客户的环境,客户的环境如下:

1. 模型文件:DeepSeek-R1 671B
2. GPU型号: 2*8 GPU RDMA(ACS分布式部署)
3. 推理框架:vllm 0.7.2
4. 启动参数:vllm serve /data/DeepSeek-R1/ --port 8000 --trust-remote-code --served-model-name ds --max-model-len 8192 --gpu-memory-utilization 0.95 --tensor-parallel-size 16 --enforce-eager --api-key token_xxxx --max-num-seqs 128

排查分析

显存分配是否合理

面对客户疑问,抱着“显存真的一直上升吗?”的怀疑尝试复现客户的情况。利用vllm启动模型时候,vllm日志中已经显示如下对单卡的显存的占用信息。

INFO 03-20 19:40:50 worker.py:267] Memory profiling takes 16.56 secondsINFO 03-20 19:40:50 worker.py:267] the current vLLM instance can use total_gpu_memory (9x.00 GiB) x gpu_memory_utilization (0.95) = 9x.00 GiBINFO 03-20 19:40:50 worker.py:267] model weights take 42.59GiB; non_torch_memory takes 2.30GiB; PyTorch activation peak memory takes 1.42GiB; the rest of the memory reserved for KV Cache is 4x.00 GiB.

在ACS中,分布式部署的大模型,每个Pod是8张GPU卡。所以Pod维度是9x.00 GiB*8=7xxGiB

模型占用/Pod:42.59GiB*8=340.72GiB

non_torch_memory/Pod:2.30GiB*8=18.4GiB

PyTorch activation peak memory/每Pod:1.42Gib*8=11.36GiB

KV Cache/每Pod 4x.00GiB*8=3xxGiB

而Prometheus监控显示显存Pod占用7xxGiB的显存,这个是和日志匹配上的。

图片

尝试复现抓住现场

第一次请求

以下面的语句对模型进行query请求。

  time curl http://172.17.98.24:8080/v1/chat/completions   -H "Content-Type: application/json"   -d '{    "model""ds",    "messages": [      {        "role""system"        "content""你是个友善的AI助手。"      },      {        "role""user",        "content""介绍一下深度学习好吗。"      }    ],    "max_tokens"2048  }'

Prometheus监控显示显存使用上升2GiB, 并且持续2h未下降。

图片

会不会是KV Cache不足

在请求过程中KV cache使用率在0.1左右,所以原则上不存在KV显存不足,占用系统层面显存。

图片

通过配置VLLM指标抓取,也从侧面佐证了KV cache占用很低,不存在KV Cache不足。同时通过Ray的监控也证明了在Query时候,显存使用率增长,但不会下降。

图片
图片

监控指标是否异常

因为指标是来自Prometheus,所以可能是Prometheus的指标问题,或Prometheus的PromQL语句展现的是累计值。

分析Prometheus的指标的Expore如下,采集的是DCGM_FI_DEV_FB_USED,并且该值和nvidia-smi命令中Memory-Usage的已使用值相对应。

sum(DCGM_FI_DEV_FB_USED{PodName=~"", NamespaceName=~""})by(NamespaceName, PodName)

指标名称

指标类型

单位

说明

DCGM_FI_DEV_FB_USED

Gauge

MiB

表示帧缓存已使用数。

该值与nvidia-smi命令中Memory-Usage的已使用值对应。

到pod对比前后的nvidia-smi情况,可以看到每张GPU卡的显存都会提高,并且一直不会被释放。

图片

显存增长是否有上限

对推理服务进行了100个请求,10并发,INPUT 3500,OUTPUT 1250的压测,发现压测后Pod显存占用增长了将近60GiB,压测结束后依然不会下降。

图片

是NCCL还是CUDA

到目前为止,问题可以明确在推理框架以下了,但是不确定是NCCL还是CUDA引发的异常显存占用,所以需要对模型服务抓取nsight进行分析。容器环境抓取比较费劲,需要挂载nas等存储以便储存临时的nsight文件。

#抓取vllm nsightnohup nsys launch  -t cuda,nvtx,osrt,cudnn,cublas,opengl,cudla-verbose,cusolver-verbose   --cudabacktrace=all --cuda-memory-usage=true --cuda-um-cpu-page-faults=true --cuda-um-gpu-page-faults=true --session-new vllm-capture1 vllm serve /model/DeepSeek-R1/ --port 8080 --trust-remote-code --served-model-name ds --enable-reasoning  --reasoning-parser deepseek_r1 --max-num-batched-tokens 8192  --max-model-len  8192  --enable-prefix-caching  --gpu-memory-utilization 0.90 --tensor-parallel-size 16 --enforce-eager > /tmp/vllm-capture.log 2>&1 &

看nsight文件,收到query请求后,由pytorch调用cudamalloc申请显存,共申请246M x 8 (Phase2: 显存增加2GB)与复现现象一致。

NCCL初始化配置在模型加载时已经完成,可以排除NCCL导致。

图片
图片
图片

设置环境变量,GPU设置为阻塞同步方式,方便观察GPU信息。由于deepseek模型参数量大,加载缓慢,更换Qwen 1.5B模型,使用单机环境方便,现象依然存在,并且增加prompt长度,显存申请数量增加,与压侧显存不断增长现象相对应。

#阻塞同步export CUDA_LAUNCH_BLOCKING=1
#PDBfrom vllm import LLM,sampling_params
prompt=[        "Hello, my name is",    "The president of the United States is",    "The capital of France is",    "The future of AI is"]model_name = "Qwen/Qwen2.5-1.5B"
llm=LLM(model= model_name,trust_remote_code=True,gpu_memory_utilization=0.95,max_model_len=8192)ans=llm.generate(prompts=prompt)for output in ans:    prompt = output.prompt    generated_text = output.outputs[0].text    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

通过pdb+nvidia-smi配合,观察到显存增加处为with ctx_factory()

用关系如下:

llm.generate()LLMEngine.step()model_executor.execute_model()driver_worker.execute_model()model_runner.execute_model()

  /usr/local/lib/python3.10/bdb.py(598)run()-> exec(cmd, globals, locals)  <string>(1)<module>()  /mnt/workspace/gx/main.py(11)<module>()-> ans=llm.generate(prompts=prompt)  /usr/local/lib/python3.10/site-packages/vllm/utils.py(838)inner()-> return fn(*args, **kwargs)  /usr/local/lib/python3.10/site-packages/vllm/entrypoints/llm.py(316)generate()-> outputs = self._run_engine(use_tqdm=use_tqdm)  /usr/local/lib/python3.10/site-packages/vllm/entrypoints/llm.py(569)_run_engine()-> step_outputs = self.llm_engine.step()  /usr/local/lib/python3.10/site-packages/vllm/engine/llm_engine.py(911)step()-> output = self.model_executor.execute_model(  /usr/local/lib/python3.10/site-packages/vllm/executor/gpu_executor.py(110)execute_model()-> output = self.driver_worker.execute_model(execute_model_req)  /usr/local/lib/python3.10/site-packages/vllm/worker/worker_base.py(272)execute_model()-> output = self.model_runner.execute_model(> /usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py(114)decorate_context()->[SamplerOutput..._metrics=None)]-> with ctx_factory():(Pdb

ctx_factory()是PyTorch 的上下文管理器,常用来管理GPU资源,性能监控,异常捕获等场景。所以基本可得出此现象与pytorch缓存机制有关。

进一步佐证定位Pytorch

在/root/miniconda3/envs/niqi-vllm/lib/python3.10/site-packages/vllm/worker/model_runner.py中的def execute_model下插入:

  for device_id in range(8): #8表示几张卡            print(torch.cuda.memory_summary(device_id))
图片

压测前后,可以看到Nvidia 0和6卡增加了1346MiB,对应的是 GPU reserved memory中 “from large pool”这种增加,其他的GPU显存占用在测试前后并无区别,进一步证明显存的增加唯一源是pytorch的缓存机制。

图片

Pytorch显存cache管理剖析

主要数据结构

图片

Block:

  • 分配 / 管理内存块的基本单位,(stream_id, size, ptr) 三元组可以特异性定位一个 Block,即 Block 维护一个 ptr 指向大小为 size 的内存块,隶属于 stream_id 的 CUDA Stream。
  • 所有地址连续的 Block(不论是否为空闲,只要是由 Allocator::malloc 得来的)都被组织在一个双向链表里,便于在释放某一个 Block 时快速检查前后是否存在相邻碎片,若存在可以直接将这三个 Block 合成为一个。

BlockPool:

  • 内存池,用 std::set 存储 Block 的指针,按照 (cuda_stream_id -> block size -> addr) 的优先级从小到大排序。所有保存在 BlockPool 中的 Block 都是空闲的。

  • DeviceCachingAllocator 中维护两种 BlockPool (large_blocks, small_blocks),分别存放较小的块和较大的块(为了分别加速小需求和大需求),简单地将 <= 1MB 的 Block 归类为小块,> 1MB 的为大块。

直观理解 Block、BlockPool 见下图:

图片

BlockPool 示意图,左边两个 500MB 的块 size 相同,因此上面的比下面的地址更

Block 在 Allocator 内有两种组织方式,一种是显式地组织在 BlockPool(红黑树)中,按照大小排列;另一种是具有连续地址的 Block 隐式地组织在一个双向链表里(通过结构体内的 prev, next 指针),可以以 O(1) 时间查找前后 Block 是否空闲,便于在释放当前 Block 时合并碎片。

申请block过程

图片

根据 size 决定实际上的 alloc_size:

  • 为小于 1MB 的 size 分配 2MB;
  • 为 1MB ~ 10MB 的 size 分配 20MB;
  • 为 >= 10MB 的 size 分配 { size 向上取整至 2MB 倍数 } MB。
图片

碎片管理

通过环境变量会指定一个阈值 max_split_size_mb,实际上从变量名可以看出,指定的是最大的可以被 “split” 的 Block 的大小。

alloc_block 函数通过 cuda 新申请的 Block(这是 Allocator 唯一一个创建新 Block 的途径)大小是计算出的 alloc_size 而不是用户通过 malloc 请求的大小 size。因此,如果 malloc 成功返回了一个 Block,Allocator 会通过函数 should_split 检查:

  • 由于过量分配内存,Block 内部产生的大小为 alloc_size - size 的碎片大小称为 remaining;
  • 如果这个 Block 属于 small_blocks 且 remaining >= 512 Bytes;或者 remaining >= 1MB 且该 Block 大小没有超过上述的阈值,则这个 Block 需要被 split。
split 的操作很简单,当前的 Block 会被拆分成两个 Block,第一个大小正好为请求分配的 size,第二个则大小为 remaining,被挂到当前 Block 的 next 指针上。这样一来,这两个 Block 的地址自然而然成为连续的了。随着程序运行,较大的 Block(只要仍小于阈值 max_split_size_mb)会不断被分成小的 Block。值得注意的是,由于新 Block 的产生途径只有一条,即通过 alloc_block 函数经由 cudaMalloc 申请,无法保证新 Block 与其他 Block 地址连续,因此所有被维护在双向链表内的有连续地址空间的 Block 都是由一个最初申请来的 Block 拆分而来的。

一段连续空间内部(由双向链表组织的 Block 们)如下图所示:

图片

释放block

当 Block 被释放时,会检查其 prev、next 指针是否为空(以此判断是否被使用)。若没有在被使用,则会使用 try_merge_blocks 合并相邻的 Block。由于每次释放 Block 都会检查,因此不会出现两个相邻的空闲块,于是只需检查相邻的块是否空闲即可。这一检查过程见 free_block 函数。又因为只有在释放某个 Block 时才有可能使多个空闲块地址连续,所以只需要在释放 Block 时整理碎片即可。

释放cache

CUDACachingAllocator 申请并维护的缓存,不会自动释放,直到显式调用torch.cuda.memory.empty_cache()实际会走到CUDACachingAllocator:emptyCache,最终通过release_blocks释放large/small 两个blocks。

图片

结论及建议

1. 显存占用:
推理过程中增加的显存是由于pytorch的显存cache机制预留了一定数量显存。nvidia-smi是reserved_memory和torch context显存之和。在推理过程中虽然不会主动释放,但是随着block的增多,也会有大量新的请求使用已经存在的空闲block,增长趋势会愈发变慢最终达到稳态。
2. 避免显存碎片化:
当出现显存剩余量足够,却无法分配连续的显存,可以通过环境变量CUDA_PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb调小,来减少显存碎片化的影响。
3. 尝试其他优化策略:除了调整max_split_size_mb外,还可以考虑其他优化策略来减少显存碎片化,如使用显存清理工具(如torch.cuda.empty_cache())或调整模型和数据加载策略。

本文撰写过程中感谢@复东@倪祺给予的帮助。

附录

模型显存计算逻辑

总显存计算

推理所需显存=模型参数部分+激活参数部分+KV Cache部分

  • 模型参数部分=模型参数量 × 精度系数

  • 量化精度与字节数关系
    • FP16/BF16:每个参数占用 2字节
    • INT8:每个参数占用 1字节
    • INT4:每个参数占用 0.5字节

  • 激活参数部分=激活参数量 × 精度系数

  • KV Cache部分=并发数 × (输入Token数+输出Token数) × 2 × 层数 × hidden_size × Sizeof(精度系数)
  • 关键参数说明
  • 层数:模型层数(如DeepSeek-R1 671B可能包含80层)
  • 隐藏层维度:模型每层的神经元数量(如1024或更大)
  • 序列长度:上下文长度(例如1K、4K、8K等)
  • 量化精度:与模型参数一致,如FP16/BF16对应2字节
以FP8精度的满血版DeepSeek-R1 671B为例,假设batch size=30,isl=2048,out=2048,num_layers=61,hidden_size=7168,激活参数量37B

总的显存容量评估如下:=671×1GB+37x1G+30×(2048+2048)×2×61×7168×1Bytes=671 GB + 100.08GB=808.08GB

首token及每Token延时评估

大模型推理其实不只看GPU的显存大小,其他参数如GPU的算力、GPU的显存带宽以及模型参数量大小都是影响实际使用的重要因素:

  • GPU显存大小:决定了能够加载并运行多大量&精度的模型;
  • GPU算力大小:与首token延时密切相关,算力越高首Token越快;
  • GPU带宽大小:与每token延时有关,带宽越大token并发越高;
  • 模型参数&精度通常模型参数越大,准确率越高,逻辑能力越强;

1、首token延时评估公式如下:

首token延时=模型参数×并发数×输入Token数×sizeof(精度) /(GPU峰值算力*90%)

2、每token延时计算方法

每token延时=模型权重内存大小÷(显存带宽*利用率)+KV Cahce内存大小÷(显存带宽*利用率)+卡间通信时延

卡间通信延时:hidden_size ×并发数× sizeof(精度) × 2 × layers ÷卡间带宽

  • 单卡推理:显存带宽利用率 按经验值60%, 
  • 多卡并行:显存带宽利用率 按经验值40%

上下文长度对显存的影响

上下文长度直接影响KV缓存的大小。以FP16精度为例:

  • 每Token的KV缓存需求 ≈ 2 × 层数 × 隐藏层维度 × 2字节
  • 总KV缓存需求 ≈ 每Token需求 × 序列长度 × 批大小
  • 示例
    对于671B模型(层数80,隐藏维度8192),1K上下文长度的KV缓存需求为:
    2 × 80 × 8192 × 2字节 × 1024 ≈ 2.6GB
    若扩展到8K上下文,则需求升至 21GB

优化策略与显存压缩

量化技术
使用INT8或INT4量化可显著降低显存需求。例如,671B模型在INT4下显存占用仅需 335.5GB(参数)+ 约5.3GB(KV缓存)≈ 340.8GB

大模型量化:降低占用显存,加快推理速度但损失精度

大模型的量化精度INT8与INT4:INT8适用于大多数需要推理加速的场景;INT4导致较大精度损失,特定场景仍有价值

异构计算
通过将稀疏MoE矩阵卸载到CPU内存,仅保留稠密部分在GPU显存中(如KTransformers项目),可将671B模型的显存需求从200GB+压缩至24GB。

算子优化
使用Marlin算子加速量化计算,结合CUDA Graph减少显存碎片,提升利用率

总结公式

  • 总显存 ≈ (参数量 × 量化字节数) + (2 × 层数 × 序列长度 × 隐藏维度 × 量化字节数× 批大小) + 框架开销
  • 关键变量:参数量、量化精度、上下文长度、批大小、模型架构参数(层数、隐藏维度)。

Pytorch显存管理源码

malloc
  Block* malloc(      c10::DeviceIndex device,      size_t orig_size,      cudaStream_t stream) {    // done outside the lock because we don't know what locks the recorder needs    // to have...    auto context = maybeGatherContext(RecordContext::STATE);
    std::unique_lock<std::recursive_mutex> lock(mutex);
    if (C10_LIKELY(captures_underway.empty())) {      // Processes end-of-life events for outstanding allocations used on      // multiple streams (checks if their GPU-side uses are complete and      // recycles their memory if so)      //      // Q. Why skip process_events if a capture might be underway?      // A. process_events involves cudaEventQueries, illegal during CUDA graph      //    capture.      //    Dumb simple solution: defer reclaiming these allocations until after      //    capture. Cross-stream memory use is uncommon, so the deferral's      //    effect on memory use during capture should be small.      process_events(context);    }    size_t size = round_size(orig_size);    auto& pool = get_pool(size, stream);    const size_t alloc_size = get_allocation_size(size);    AllocParams params(device, size, stream, &pool, alloc_size, stats);    params.stat_types = get_stat_types_for_pool(pool);
    // First, try to get a block from the existing pool.    bool block_found =        // Search pool        get_free_block(params)        // Trigger callbacks and retry search        || (trigger_free_memory_callbacks(params) && get_free_block(params));
    // Can't reuse an existing block; try to get a new one.    if (!block_found) {      // Do garbage collection if the flag is set.      if (C10_UNLIKELY(              set_fraction &&              CUDAAllocatorConfig::garbage_collection_threshold() > 0.0)) {        garbage_collect_cached_blocks(context);      }      // Attempt allocate      // WARNING: alloc_block may release the allocator lock when calling      // cudaMalloc. So far this function has not modified allocator state, but      // keep in mind that any observed allocator state may change across calls      // to alloc_block since it may release the lock.      block_found = alloc_block(params, false, context, lock)          // Free enough available cached blocks to satisfy alloc and retry          // alloc.          || (release_available_cached_blocks(params, context) &&              alloc_block(params, false, context, lock))          // Free all non-split cached blocks and retry alloc.          || (C10_LIKELY(captures_underway.empty()) &&              release_cached_blocks(context) &&              alloc_block(params, true, context, lock));    }
    if (!block_found) {      //OOM流程    }
    bool split_remainder = should_split(params.block, params.size());    return alloc_found_block(        params, orig_size, std::move(context), split_remainder);  }
free
void free(Block* block) {    std::shared_ptr<GatheredContext> context =    maybeGatherContext(RecordContext::ALL);    std::lock_guard<std::recursive_mutex> lock(mutex);
    block->allocated = false;
    // following logic might modifying underlaying Block, causing the size    // changed. We store ahead for reporting    auto orig_block_ptr = block->ptr;    auto orig_block_size = block->size;
    StatTypes stat_types = get_stat_types_for_pool(*block->pool);    for_each_selected_stat_type(stat_types, [&](size_t stat_type) {        stats.allocation[stat_type].decrease(1);        stats.allocated_bytes[stat_type].decrease(block->size);    });    auto allocated_bytes_gauge =    STATIC_GAUGE(pytorch.CUDACachingAllocator.allocated_bytes);    allocated_bytes_gauge.record(        stats.allocated_bytes[static_cast<int64_t>(StatType::AGGREGATE)]        .current);
    record_trace(        TraceEntry::FREE_REQUESTED,        int64_t(block->ptr),        block->requested_size,        block->stream,        block->device,        context ? context : block->context_when_allocated);
    if (block->size >= CUDAAllocatorConfig::max_split_size())        stats.oversize_allocations.decrease(1);
    if (!block->stream_uses.empty()) {        if (C10_UNLIKELY(!captures_underway.empty())) {            // It's forbidden to cudaEventQuery an event recorded during CUDA graph            // capture. We conservatively defer recording end-of-life events until            // the next call to process_events() (which won't happen until no            // captures are underway)            needs_events_deferred_until_no_capture.push_back(block);        } else {            insert_events(block);        }    } else {        free_block(block, context);    }
    c10::reportMemoryUsageToProfiler(        orig_block_ptr,        -static_cast<int64_t>(orig_block_size),        stats.allocated_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,        stats.reserved_bytes[static_cast<size_t>(StatType::AGGREGATE)].current,        c10::Device(c10::DeviceType::CUDA, block->device));}
try_merge_blocks
  size_t try_merge_blocks(Block* dst, Block* src, BlockPool& pool) {    if (!src || src->allocated || src->event_count > 0 ||        !src->stream_uses.empty() || dst->mapped != src->mapped) {      return 0;    }
    AT_ASSERT(dst->is_split() && src->is_split());
    if (dst->prev == src) { // [src dst]      dst->ptr = src->ptr;      dst->prev = src->prev;      if (dst->prev) {        dst->prev->next = dst;      }      dst->context_when_segment_allocated =          std::move(src->context_when_segment_allocated);    } else { // [dest src]      dst->next = src->next;      if (dst->next) {        dst->next->prev = dst;      }    }    const size_t subsumed_size = src->size;    dst->size += subsumed_size;    // NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)    auto erased =        src->mapped ? pool.blocks.erase(src) : pool.unmapped.erase(src);    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(erased == 1);    delete src;
    return subsumed_size;  }

智能理解 PPT 内容,快速生成讲解视频


在制作线上课程、自媒体内容或者活动宣传视频时,用户通常需要撰写解说词、录制音频和剪辑视频,制作流程繁琐且周期较长。本方案利用大模型的理解和生成能力自动将 PPT 转化为讲解视频,提高了制作效率,使创作者能专注于内容创新。    


点击阅读原文查看详情。

继续滑动看下一个
阿里云开发者
向上滑动看下一个