One-sentence Summary

Researchers from Carnegie Mellon University, Xlue, and NVIDIA propose PithTrain, a compact and agent-native Mixture-of-Experts training system that introduces agent-task efficiency (ATE) and ATE-Bench, enabling coding agents to understand and extend the framework with up to 62% fewer agent turns and 64% less active GPU time while matching the throughput of production frameworks.

Key Contributions

• PithTrain, a roughly 11K-line Python-native MoE training framework designed agent-native from the start, matches the training throughput of production frameworks on NVIDIA H100 and B200 GPUs.
• Four design principles for agent-native ML training frameworks guide PithTrain: code compactness, Python-native components, no implicit indirection, and agent skills.
• ATE-Bench varies the training framework while holding the agent and tasks fixed, and an empirical study shows up to 62% fewer Agent Turns and 64% less Active GPU Time on PithTrain, including a skills ablation and case study.

Introduction

The authors address the challenge of evolving distributed training frameworks for Mixture-of-Experts (MoE) language models, where production stacks rely on layered Python designs with heavy compiled extensions, plugin systems, and registry-based indirection. While these designs deliver peak throughput, they make it costly for AI coding agents to locate, trace, and verify changes when adapting the framework to new architectures or optimizations. Existing evaluations focus solely on training throughput, ignoring the engineering effort required to extend the system. To tackle this, the authors introduce PithTrain, an end-to-end MoE training framework built from the start to be agent-native. It follows four principles: code compactness, Python-native components, direct calls that avoid implicit indirection, and bundled agent skills. They also define agent-task efficiency (ATE) and create ATE-Bench, a benchmark that holds the coding agent fixed while varying the framework to measure how design choices affect agent cost. PithTrain matches the training throughput of production frameworks while enabling coding agents to complete the same tasks with significantly fewer turns and less GPU time.

Dataset

The authors introduce ATE-Bench, a curated benchmark designed to evaluate agent-task efficiency across training frameworks. It holds the agent and tasks fixed while varying the framework, so performance differences can be attributed to framework design. The benchmark consists of 20 tasks organized into three categories, each reflecting a common pattern in framework work.

Dataset composition and sources

• The benchmark is a hand-crafted suite of tasks, not a scraped or automatically generated dataset.
• Tasks are designed to be valid across multiple frameworks (e.g., PithTrain and production frameworks) and avoid framework-specific quirks.
• The suite covers three recurring developer activities: understanding a framework, operating it as a profiling tool, and extending it with new functionality.

Key details for each subset

• Q&A (12 tasks): Read-only questions whose answers are properties of the code, not runtime measurements. Example: “How is the device mesh built?” These are filtered to work identically across all evaluated frameworks.
• Operate and Profile (4 tasks): Tasks that require running, instrumenting, and profiling the framework as a tool. Example: capture an Nsight Systems profile and identify the most expensive CUDA kernels.
• New Feature (4 tasks): End-to-end porting of a new model architecture against a published reference implementation. Example: implementing Mixture of Block Attention (MoBA). These involve substantial code modification and test-driven iteration.

How the paper uses the data

• The benchmark is used to evaluate agent-task efficiency by running a fixed agent on each task across different frameworks.
• The authors report five independent effort metrics: session duration, active GPU time, agent turns, per-turn context size, and output tokens. No single scalar metric is used.
• Agent involvement deepens across categories: Q&A is read-only, Operate and Profile adds execution and minor instrumentation, and New Feature requires significant modification and iteration.

Processing details

• Each task comes with a detailed description and per-category correctness checks (defined in Appendix B of the paper).
• Q&A questions are explicitly filtered to exclude framework-specific behaviors, ensuring cross-framework validity.
• No cropping or metadata construction beyond task definitions and evaluation criteria is described; the benchmark is used as-is for framework comparison.

Method

The authors introduce PithTrain, a system designed around four agent-native principles to facilitate coding agent interaction. First, they prioritize a compact codebase, ensuring the entire framework can fit within a single context window of frontier coding agents, contrasting with production frameworks that exceed 160K lines. Second, the framework is Python-native, allowing agents to navigate the code in a single language and eliminating compiled-extension rebuild cycles. Third, they enforce a no implicit indirection rule. Production frameworks often resolve submodules through runtime specs or plugin registries, which obscures local readability. As shown in the figure below, PithTrain instantiates layers directly in self-contained files rather than relying on a shared layer skeleton with implicit indirection.

Fourth, the system incorporates task-specific agent skills, which are procedural playbooks encoding verified procedures for recurring tasks.

Regarding the system architecture, the codebase is organized into three layers: application, engine, and operator. Refer to the framework diagram to see the component breakdown and line counts. To maintain compactness, PithTrain narrows its scope to exactly the necessary components for distributed Mixture of Experts (MoE) training, supporting pipeline parallelism, data parallelism via FSDP, context parallelism, expert parallelism, FP8 training, and DCP checkpointing.

To achieve competitive training throughput, the authors implement several standard MoE optimizations. They utilize a DualPipeV pipeline schedule with compute and communication overlap, decomposing each transformer layer into five stages at expert parallelism boundaries and running all-to-all communications on a separate stream. They apply torch.compile with fullgraph mode to all transformer computation except the MoE forward and backward passes to reject graph breaks at compile time. Additional optimizations include wgrad delay, fused SwiGLU kernels, expert parallelism dispatch deduplication, an FP8 weight cache across micro-batches, and fused Triton kernels for token scatter and quantization.

Finally, the agent skills are designed around three core properties: specific scope, explicit prerequisites, and verifiable success. The figure below illustrates these properties using the validate-correctness skill. The description and triggers define the scope, the prerequisites section enumerates environment and data assumptions to catch missing state early, and the procedure concludes with a script call that returns a reproducible PASS or FAIL verdict.

Experiment

The evaluation compares PithTrain against Megatron-LM and TorchTitan on training throughput and agent-task efficiency using ATE-Bench, a benchmark that fixes the agent and tasks to isolate framework design effects across Q&A, profiling, and new-feature integration. PithTrain matches or exceeds production frameworks in training throughput while dramatically reducing agent effort: its compact, Python-native codebase with no implicit indirection yields fewer agent turns, lower context usage, and less GPU time, as failures surface locally with readable tracebacks rather than opaque C++ errors or hidden registries. An ablation study confirms that in-repo skills further cut agent reasoning overhead by providing fixed plans for recurring tasks, and a case study on integrating MoBA shows that editing dominates across frameworks but PithTrain's lower exploration cost and single-file fixes contrast with the multi-file debugging and memory-pressure reruns in other systems.

PithTrain is the only framework that satisfies all four agent-native design principles: a compact Python codebase with no implicit indirection and built-in agent skills. Production frameworks like Megatron-LM and DeepSpeed have large, multi-language codebases with implicit indirection, while TorchTitan is Python-native but still relies on indirection and a larger codebase. These structural differences directly impact agent effort, with compactness and local readability reducing exploration and keeping fixes local. PithTrain's codebase (11K lines) is over three times smaller than TorchTitan's (38K) and more than an order of magnitude smaller than Megatron-LM's (149K) and DeepSpeed's (167K). Only PithTrain avoids implicit indirection, trading cross-model reuse for local readability that helps agents identify what code runs at a call site. PithTrain and TorchTitan are Python-native, eliminating the compiled-extension rebuild cycle and surfacing readable tracebacks, unlike Megatron-LM and DeepSpeed. Agent skills, present only in PithTrain, reduce agent turns by up to 70% on recurring training-system tasks by providing a fixed procedure. In a MoBA integration case study, PithTrain's compactness and lack of indirection kept fixes local to the agent's edits, whereas Megatron-LM required multi-file fixes.

PithTrain keeps the codebase compact at roughly 11K lines, with the entire stack written in Python except for a Python DSL used for custom kernels. A flat file structure and absence of plugin registries eliminate implicit indirection, while each MoE model resides in a self-contained file to favor local readability. Agent skills are placed as specialized files under the project root, making recurring tasks straightforward to automate. The full stack is Python-native, relying on a Python DSL only for custom kernels, which avoids mixed-language complexity. A flat code organization with no plugin registries or runtime specs removes implicit indirection, so models are self-contained and easy to navigate. Agent skills live as specialized files under the project root, directly encoding procedures for tasks like correctness validation and profiling.

PithTrain provides a set of skills that encode procedures for recurring training-framework tasks, each with a specific scope, explicit prerequisites, and a verifiable success check. When agents use these skills for validation and profiling tasks, their reasoning overhead drops substantially while GPU time stays nearly unchanged. In a MoBA integration case study, PithTrain's compact codebase and absence of implicit indirection lead to less exploration and lower per-turn context compared to other frameworks. Agent turns fall by 70% and 52% for the two tasks when skills are enabled, indicating the agent follows a fixed plan instead of iteratively deriving one. Active GPU time remains nearly identical with and without skills, since the GPU work is determined by the fixed training runs, not the agent. In the MoBA integration, PithTrain's editing output tokens are the lowest among the three frameworks, and its per-turn context window stays consistently lower than Megatron-LM's. Failures on PithTrain are localized to a single file, while Megatron-LM failures span multiple files, reflecting PithTrain's compactness and lack of implicit indirection.

Across five MoE training configurations on NVIDIA H100 and B200 GPUs, PithTrain matches or exceeds Megatron-LM throughput in four cases and trails by only 1.4% in the fifth. TorchTitan either runs out of memory or delivers lower throughput than both other frameworks. These results confirm that a compact, Python-native codebase can achieve competitive training throughput. PithTrain outperforms Megatron-LM by up to 27% on B200 nodes with FP8 precision. TorchTitan encounters out-of-memory errors on two configurations and shows the lowest throughput on all others.

On codebase Q&A tasks, PithTrain consistently reduces agent effort compared to Megatron-LM and TorchTitan. Agent turns drop by up to 67% relative to Megatron-LM, and per-turn context and output tokens are also lower, reflecting a compact codebase with less implicit indirection. PithTrain required fewer agent turns than Megatron-LM on all but one question, with the largest reduction reaching 67% (e.g., Q2: 54 turns on Megatron-LM vs 18 on PithTrain). Per-turn context was consistently lower on PithTrain, often by thousands of tokens, indicating a smaller search space during Q&A tasks.

The evaluation compares PithTrain, a compact and Python-native training framework with built-in agent skills, against production frameworks Megatron-LM, DeepSpeed, and TorchTitan across codebase Q&A, a MoBA integration case study, and throughput benchmarks on H100 and B200 GPUs. PithTrain's small codebase, flat structure, and avoidance of implicit indirection substantially reduce agent exploration, keep fixes local, and lower per-turn context, while its agent skills further decrease reasoning overhead on recurring tasks. Despite its minimal design, PithTrain matches or exceeds Megatron-LM throughput in most configurations, confirming that a locally readable, agent-native codebase can achieve competitive training performance with far less agent effort.