# Learning to Shard: RL for Co-optimizing the Parallelism Degrees and Per-operator Sharding Dimensions in Distributed LLM Inference

arXiv: https://arxiv.org/abs/2509.00217 · HTML: https://mlsanity.com/paper/2509.00217/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 19 references that are in the map.
- First posted on arXiv 29 Aug 2025.

## Builds on (most cited first, 13 of 13)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Seesaw: High-throughput LLM Inference via Model Re-sharding](https://mlsanity.com/paper/2503.06433/) (2025)
- [Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding](https://mlsanity.com/paper/2507.07120/) (2025)
- [TAPAS: Fast and Automatic Derivation of Tensor Parallel Strategies for Large Neural Networks](https://mlsanity.com/paper/2302.00247/) (2023)
- [UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming](https://mlsanity.com/paper/2307.16375/) (2023)
- [DFModel: Design Space Optimization of Large-Scale Systems Exploiting Dataflow Mappings](https://mlsanity.com/paper/2412.16432/) (2024)
- [COSMIC: Enabling Full-Stack Co-Design and Optimization of Distributed Machine Learning Systems](https://mlsanity.com/paper/2505.15020/) (2025)

## Abstract (arXiv)

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.