PithTrain: A Compact and Agent-Native MoE Training System

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Lai, Ruihang, Kang, Hao, Tang, Haozhan, Parthasarathy, Akaash R., Yu, Zichun, Shao, Junru, Mowry, Todd C., Xiong, Chenyan, Chen, Tianqi
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913174030123008
author Lai, Ruihang
Kang, Hao
Tang, Haozhan
Parthasarathy, Akaash R.
Yu, Zichun
Shao, Junru
Mowry, Todd C.
Xiong, Chenyan
Chen, Tianqi
author_facet Lai, Ruihang
Kang, Hao
Tang, Haozhan
Parthasarathy, Akaash R.
Yu, Zichun
Shao, Junru
Mowry, Todd C.
Xiong, Chenyan
Chen, Tianqi
contents Mixture-of-Experts (MoE) has become the dominant architecture for frontier language models. To meet this demand, production frameworks have built optimized MoE training stacks over years of engineering effort. Yet evolving these stacks for new architectures and system optimizations remains expensive. With the rise of AI coding agents, they could automate parts of training-framework development and accelerate this evolution. But applying them to these existing frameworks carries hidden costs, invisible to today's throughput-only evaluations. We name this missing dimension agent-task efficiency (ATE): the cost of using coding agents to understand, operate, and extend a framework. Grounded in four agent-native design principles, we build PithTrain, a compact, agent-native MoE training framework. We further introduce ATE-Bench, covering real-world training-framework tasks. Our evaluation shows PithTrain matches the throughput of production frameworks, and on ATE-Bench, PithTrain enables higher agent-task efficiency, with up to 62% fewer Agent Turns and 64% less Active GPU Time.
format Preprint
id arxiv_https___arxiv_org_abs_2605_31463
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle PithTrain: A Compact and Agent-Native MoE Training System
Lai, Ruihang
Kang, Hao
Tang, Haozhan
Parthasarathy, Akaash R.
Yu, Zichun
Shao, Junru
Mowry, Todd C.
Xiong, Chenyan
Chen, Tianqi
Machine Learning
Artificial Intelligence
Computation and Language
Distributed, Parallel, and Cluster Computing
Mixture-of-Experts (MoE) has become the dominant architecture for frontier language models. To meet this demand, production frameworks have built optimized MoE training stacks over years of engineering effort. Yet evolving these stacks for new architectures and system optimizations remains expensive. With the rise of AI coding agents, they could automate parts of training-framework development and accelerate this evolution. But applying them to these existing frameworks carries hidden costs, invisible to today's throughput-only evaluations. We name this missing dimension agent-task efficiency (ATE): the cost of using coding agents to understand, operate, and extend a framework. Grounded in four agent-native design principles, we build PithTrain, a compact, agent-native MoE training framework. We further introduce ATE-Bench, covering real-world training-framework tasks. Our evaluation shows PithTrain matches the throughput of production frameworks, and on ATE-Bench, PithTrain enables higher agent-task efficiency, with up to 62% fewer Agent Turns and 64% less Active GPU Time.
title PithTrain: A Compact and Agent-Native MoE Training System
topic Machine Learning
Artificial Intelligence
Computation and Language
Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2605.31463