Dynamo项目中TP4并行配置在disagg_router模式下的内存优化实践

背景介绍

在大型语言模型(LLM)推理服务部署中，Dynamo项目提供了一种高效的分布式推理框架。本文针对32B参数规模、32k最大序列长度的模型在16块A100 40GB GPU上的部署实践，重点分析了Tensor Parallelism 4(TP4)配置在disagg_router模式下出现的内存问题及解决方案。

问题现象

在采用disagg_router配置时，系统出现了CUDA内存不足的错误。错误日志显示，虽然GPU总容量为39.39GiB，但实际可用内存仅剩10.38MiB，导致模型加载失败。值得注意的是，相同的TP4配置在agg_router模式下能够正常工作。

技术分析

内存分配机制

在Dynamo的分布式推理架构中，disagg_router模式将工作负载分离到不同的工作节点：

1. VllmWorker负责解码阶段
2. PrefillWorker负责预填充阶段

每个工作节点都需要独立加载模型参数，因此必须正确配置各节点的并行策略。

关键配置项

原始配置中遗漏了PrefillWorker的tensor-parallel-size参数，导致系统默认使用TP1(单卡)模式加载模型。对于32B参数的模型，单卡加载显然会超出40GB显存容量。

解决方案

通过完善PrefillWorker的配置，显式指定TP4并行策略：

PrefillWorker:
  model: my_model
  kv-transfer-config: '{"kv_connector":"DynamoNixlConnector"}'
  block-size: 64
  max-model-len: 32768
  max-num-batched-tokens: 32768
  tensor-parallel-size: 4  # 关键配置项
  gpu-memory-utilization: 0.8
  ServiceArgs:
    workers: 1
    resources:
      gpu: 4

优化建议

1. 内存利用率调整：将gpu-memory-utilization从0.5提升到0.8，在保证稳定性的前提下提高资源利用率
2. 配置一致性检查：确保VllmWorker和PrefillWorker的并行配置匹配
3. 资源监控：部署前使用nvidia-smi工具监控各GPU内存使用情况
4. 渐进式加载：对于超大模型，考虑分阶段加载策略

经验总结

在Dynamo项目的分布式部署中，disagg_router模式提供了更灵活的工作负载分配能力，但也带来了配置复杂度的提升。工程师需要特别注意：

1. 所有涉及模型加载的工作节点都必须正确配置并行策略
2. TP参数需要与分配的GPU数量匹配
3. 不同组件的内存配置需要协调一致

通过这次实践，我们验证了32B参数模型在16块A100 GPU上使用TP4并行策略的可行性，为类似规模的模型部署提供了可靠参考。