KTransformers项目在双路服务器上的权重镜像与NUMA优化实践

背景与问题现象

在大型语言模型推理场景中，KTransformers项目为多路服务器提供了高效的推理支持。近期有用户反馈，在双路Intel E5和AMD EPYC 9004平台上运行相同模型时，出现了截然不同的内存行为模式：

1. Intel E5平台：正常触发了权重镜像机制，模型权重被复制到两个NUMA节点的内存中，总内存占用约740GB，QPI互联流量较低，推理性能良好
2. AMD EPYC 9004平台：权重仅加载到缓存中，未实现跨NUMA节点的权重镜像，导致双路性能与单路相比无明显提升

技术原理分析

NUMA架构与权重镜像

现代多路服务器采用NUMA(Non-Uniform Memory Access)架构，每个CPU插槽及其直连内存构成一个NUMA节点。KTransformers通过以下机制优化多路性能：

1. 权重镜像：将模型权重复制到各NUMA节点的本地内存，避免跨节点访问带来的延迟惩罚
2. NUMA感知调度：确保计算任务优先访问本地内存，减少QPI/UPI互联流量

平台差异解析

造成两种平台表现差异的技术原因包括：

1. NUMA库依赖：AMD平台缺少libnuma-dev库导致NUMA功能未正确启用
2. EPYC架构特性：AMD EPYC处理器支持NPS(Numa Per Socket)配置模式，可能需要特殊优化
3. 内存控制器差异：Intel和AMD的内存控制器架构不同，影响跨节点访问效率

解决方案与实践

环境配置要点

1. 基础依赖安装：

   sudo apt-get install libnuma-dev
   

2. 编译选项设置：

   export USE_NUMA=1
   make dev_install
   

3. 运行监控：

   numactl --hardware  # 验证NUMA节点识别
   numastat -m  # 监控内存分布
   

AMD平台特殊优化

对于EPYC平台，建议尝试以下配置：

1. 测试不使用USE_NUMA=1的情况
2. 在BIOS中调整NPS模式为NPS1或NPS4
3. 检查/sys/devices/system/node/下的NUMA节点信息

性能调优建议

1. 内存分配策略：

   • 优先使用numactl --localalloc确保内存本地化
   • 对于大模型，考虑--interleave=all交错分配

2. 线程绑定：

   numactl --cpunodebind=0 --membind=0 python script.py
   

3. 监控指标：

   • 使用perf stat监控QPI/UPI流量
   • 通过numastat观察跨节点访问情况

典型问题排查

1. 权重未镜像：

   • 检查NUMA库安装
   • 验证环境变量设置
   • 检查BIOS中的NUMA设置

2. 性能不达预期：

   • 确认线程绑定策略
   • 检查内存带宽利用率
   • 监控CPU缓存命中率

3. 显存管理：

   • 使用nvidia-smi监控显存占用
   • 考虑使用CUDA MPS提高GPU利用率

总结

KTransformers在多路服务器上的性能优化需要综合考虑硬件架构、NUMA配置和软件环境。通过正确的NUMA配置和权重镜像机制，可以显著提升大型语言模型在多路服务器上的推理效率。不同平台可能需要针对性的优化策略，建议用户根据具体硬件特性进行调优。

对于AMD EPYC平台，除了标准的NUMA优化外，还需要关注其特有的NPS配置和内存控制器特性，以获得最佳性能表现。未来随着KTransformers项目的持续发展，预期将提供更完善的多路服务器支持能力。