优化大模型训练通信，让多卡并行计算更快更省带宽
SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization
- 只传标量统计量，减少跨卡通信开销
- 相比传统方法，计算延迟降低81.4%~93.9%
- 适合大规模语言模型训练，提升推理效率
查询-键归一化（QK-Norm）能提升现代大语言模型的训练稳定性和质量。但在张量并行（TP）下，逐层QK-Norm因依赖完整隐藏向量而引入额外跨卡通信。我们提出SwiftQK，一种多卡RMSNorm内核，仅交换标量归一化统计量，并将剩余点对点归约操作与独立元素计算重叠，实现无死锁的持续运行。在近期大语言模型上的评估显示，SwiftQK相比基于全向量All-Gather的标准TP QK-Norm，将延迟降低了81.4%~93.9%。在端到端服务中，相较于基于All-Gather的基线，平均减少了29.5%的TPOT；相较于优化的标量聚合实现，仍降低14.3%。
原文摘要 · Abstract (English)
Query-Key Normalization (QK-Norm) improves the training stability and quality of modern Large Language Models (LLMs). However, under Tensor Parallelism (TP), layerwise QK-Norm introduces additional cross-GPU communication because the normalization factor depends on the full hidden vector. We present SwiftQK, a multi-GPU RMSNorm kernel that exchanges only scalar normalization statistics and overlaps the remaining Peer-to-Peer reduction with independent element-wise computation in a deadlock-safe persistent kernel. Evaluations on recent LLMs show that SwiftQK reduces QK-Norm latency by 81.4--93.9% relative to the standard TP QK-Norm using full-vector All-Gather. In end-to-end serving, SwiftQK reduces TPOT on average by 29.5% over the All-Gather-based baseline and by 14.3% over an optimized scalar-aggregation implementation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。