AIInfra之模型視訊記憶體管理分析

阿里妹導讀
本文圍繞某線上客戶部署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日誌中已經顯示如下對單卡的視訊記憶體的佔用資訊。
INFO03-2019:40:50 worker.py:267] Memory profiling takes 16.56 secondsINFO03-2019: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 GiBINFO03-2019: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 Cacheis4x.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_paramsprompt=[    "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].textprint(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 inrange(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);constsize_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 poolget_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());returnalloc_found_block(        params, orig_size, std::move(context), split_remainder);  }
free
voidfree(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 reportingauto 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_ttry_merge_blocks(Block* dst, Block* src, BlockPool& pool){if (!src || src->allocated || src->event_count > 0 ||        !src->stream_uses.empty() || dst->mapped != src->mapped) {return0;    }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;      }    }constsize_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 轉化為講解影片,提高了製作效率,使創作者能專注於內容創新。    
點選閱讀原文檢視詳情。

相關文章