用强化学习同时优化大模型分布式推理的并行度与算子分片，提升性能3.5倍。
Learning to Shard: RL for Co-optimizing the Parallelism Degrees and Per-operator Sharding Dimensions in Distributed LLM Inference
- 基于注意力机制的策略网络，从历史高性能配置中学习搜索最优分片方案。
- 在1.6万亿参数MoE模型上，比元启发式方法快3.5倍，比Megatron快1.06倍。
- 适合大规模分布式推理场景，尤其适用于硬件异构、模型扩增的生产环境。
分布式大模型推理需要在数百至数千个NPU间协调并行策略以满足生产级服务等级协议（SLO）。现有系统如Megatron-LM依赖静态启发式方法，分别配置粗粒度并行度和细粒度算子分片维度，导致模型规模扩大和硬件拓扑多样化时性能严重受限。我们提出Learn to Shard，据我们所知首个基于强化学习的方法，用于联合优化分布式大模型推理中的粗粒度并行度与细粒度算子分片维度。该方法采用基于注意力的策略网络，从精英历史中学习高绩效策略，高效探索庞大的组合搜索空间。在配备H100集群的MoE模型（最大1.6万亿参数）上评估，该方法相比元启发式基线最高提升3.5倍吞吐量，相较Megatron启发式提升1.06倍。
原文摘要 · Abstract (English)
Distributed LLM inference requires careful coordination of parallelization strategies across hundreds to thousands of NPUs to meet production SLOs. Current systems like Megatron-LM rely on static heuristics that separately configure parallelism degrees and per-operator sharding dimensions, leaving significant performance on the table as models scale and hardware topologies diversify. We introduce Learn to Shard, to our knowledge, the first RL-based approach to co-optimize both coarse-grained parallelism degrees and fine-grained per-operator sharding dimensions for distributed LLM inference. Our method employs an attention-based policy over an elite history that learns from high-performing strategies to efficiently navigate the vast combinatorial search space. Evaluated on H100 clusters with MoE models up to 1.6T parameters, Learn to Shard achieves up to 3.5x throughput improvement over metaheuristic baselines and 1.06x over Megatron heuristics.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。