PithTrain: A Compact and Agent-Native MoE Training System

TL;DR

PithTrain is a CMU/MLC ~11K-line pure-Python MoE training framework, intentionally compact enough that an AI coding agent (Claude Code Opus 4.7) can hold the entire codebase in a single 200K–1M-token context window. It introduces agent-task efficiency (ATE) as a second optimization axis alongside training throughput — measuring the cost (in agent turns, GPU time, output tokens) of using coding agents to understand, operate, and extend a framework — and ships ATE-Bench (20 tasks across Q&A, Operate-and-Profile, New-Feature) for cross-framework comparison. On 5 representative MoE configurations (GPT-OSS-20B, Qwen3-30B-A3B, DeepSeek-V2-Lite) PithTrain matches or beats Megatron-LM throughput on 4 of 5; on ATE-Bench it cuts agent turns by up to 67% (Q&A) / 70% (Operate) and active GPU time by up to 64% (New-Feature) versus Megatron-LM and TorchTitan, with the same Claude Code agent. The headline framing: training frameworks should now be designed for agents to evolve, not just for humans to maintain.

Key claims

A pure-Python ~11K-line MoE training stack achieves training throughput parity with Megatron-LM (149K LoC) on 4 of 5 matched configurations across H100 / B200, BF16 / FP8, and PP+DP+CP+EP composed up to 4D — within 1.4% of Megatron-LM on the fifth [Table 4 §5.1]. The four named optimizations carrying this parity are DualPipeV pipeline schedule with compute–communication overlap (built on DeepSeek-V3’s DualPipe scaffold), torch.compile(fullgraph=True) on all non-MoE transformer computation, wgrad delay, fused SwiGLU forward/backward, and EP-dispatch deduplication [§3.2].
Four design principles are framed as primary constraints for agent-native ML frameworks: code compactness (codebase fits in one agent context window), Python-native components (single-language navigation, readable tracebacks, no rebuild cycle), no implicit indirection (direct calls, each model in a self-contained file under models/, no plugin registries or string-keyed resolution), and shipped task-specific agent skills [§3.1, Table 2]. Production-framework comparison: Megatron-LM scores ✗/✗/✗/◐, DeepSpeed ✗/✗/✗/✗, TorchTitan ◐/✓/◐/◐, PithTrain ✓/✓/✓/✓ [Table 2 footnote 2 — LoC at pinned commits 3bec9aa / 44c51e3 / d84e83d / v0.1.2].
“Agent-task efficiency” is introduced as a metric orthogonal to throughput, defined as the cost (session duration, active GPU time, agent turns, per-turn context, output tokens) of using a coding agent to understand, operate, and extend a framework [§1, §4]. ATE-Bench inverts standard agent benchmarks (SWE-bench / HumanEval / MLE-bench): instead of varying the agent on a fixed codebase, it varies the framework on a fixed agent and task suite [§4, Fig. 5].
ATE-Bench covers 20 tasks in three categories with monotonically deepening agent involvement: Q&A (12 read-only tasks like “how is the device mesh built?”), Operate-and-Profile (4 tasks like “capture Nsight Systems profile, identify expensive CUDA kernels”), and New-Feature (4 tasks: port Differential Transformer, Dynamic MoE, MoBA, MoE++ end-to-end against published reference implementations) [§4]. Cross-framework correctness was 100% across all attempts on Claude Code Opus 4.7 at xhigh effort [§5.2].
On Q&A, PithTrain uses up to 67% fewer agent turns than Megatron-LM to reach correct answers; per-turn context shrinks from ~46K to ~33K tokens [Table 5 §5.2]. The mechanism is mechanical: smaller codebase + no implicit indirection shrinks the agent’s search space.
On Operate-and-Profile, PithTrain cuts agent turns by up to 70% vs Megatron-LM and 57% vs TorchTitan, and output tokens by up to 78% and 65% respectively [Table 6 §5.2]. The agent auto-invokes in-repo skills (capture-nsys-profile) without being told.
On New-Feature (porting a new architecture against a published reference), PithTrain cuts active GPU time by up to 44% vs Megatron-LM and 64% vs TorchTitan, primarily because the test–debug cycle converges in fewer training runs [Table 7 §5.2]. Two named failure modes on Megatron-LM are causally attributed to violated principles: a hidden CLI-argument registry collides with auto-derived flags (implicit indirection), and TransformerEngine grouped-GEMM segfaults emit opaque C++ errors that drive speculative configuration toggles (not Python-native).
An agent-skill ablation on the wgrad-delay validation workflow shows skills cut agent turns by 70% (validate-correctness) and 52% (capture-nsys-profile) at constant active GPU time, with similar drops in session duration, per-turn context, and output tokens [Table 8 §5.3]. Active GPU time stays near parity because the skill prescribes the same training runs; the savings are all in the agent’s reasoning overhead.
MoBA-integration case study: the agent’s output tokens decompose into Editing / Exploring / Other categories; Megatron-LM spends 10.2K tokens on Exploring vs PithTrain’s 2.2K — a 4.6× reduction explained by codebase size + no implicit indirection lowering both Exploring tokens and per-turn context [Fig. 6 §5.4]. TorchTitan’s elevated cost in two of three runs is caused by initial OOM forcing repeated debug-edit cycles, isolating runtime memory headroom as a factor independent of codebase structure.
Shipped agent skills are designed around three properties — specific scope, explicit prerequisites, verifiable PASS/FAIL success — and ship as self-contained folders with a top-level SKILL.md, optional helper scripts, and optional supplementary markdown [§3.3, Fig. 4]. The four shipped skills are add-new-model, add-memory-prints, capture-nsys-profile, validate-correctness [Table 3].

Method

PithTrain is organized as three layers (application / engine / operator) totaling ~11K lines of Python [Fig. 3]. The application layer drives training loops; the engine layer composes pipeline parallelism (DualPipeV from DeepSeek-V3), data parallelism via FSDP, context parallelism (ring-attention style), and expert parallelism into a 4D-parallel scheduler with FP8 training and DCP checkpointing on NVIDIA Hopper and Blackwell GPUs; the operator layer provides fused Triton kernels (EP token scatter, FP8 quantization, fused SwiGLU forward/backward). Each transformer layer is decomposed into five stages at EP communication boundaries, with EP all-to-alls on a separate stream so the forward of one micro-batch overlaps with the backward of another. torch.compile(fullgraph=True) is applied everywhere except MoE forward/backward (excluded because per-expert input shapes are data-dependent under EP).

The agent-native design constraints are not novel individually (compact code, Python-native, no indirection, shipped playbooks); the paper’s argument is that treating them jointly as primary design constraints, and then measuring what they buy you on a fixed-agent benchmark, is what’s new. Each MoE model lives in a self-contained file under models/ rather than being composed via a plugin registry or runtime spec, which trades cross-model code reuse for local readability.

Evaluation uses Claude Code (Opus 4.7) at “xhigh” effort as the fixed agent across Megatron-LM (commit 3bec9aa), TorchTitan (commit d84e83d), and PithTrain (v0.1.2 / 23db182). Each ATE-Bench task runs three times; medians are reported across five effort metrics (session duration, active GPU time, agent turns, per-turn context, output tokens). Training correctness is validated against Megatron-LM on both pretraining loss curves and downstream accuracy (Appendix A). The skill ablation strips skills from both the working tree and the git history so the agent cannot recover the procedure from either source.

Results

Training throughput parity on 5 MoE configurations on H100 / B200: PithTrain matches or beats Megatron-LM on 4 of 5 (e.g. GPT-OSS-20B PP2/DP1/CP1/EP4 BF16 on 18-B200: PithTrain 140.9K vs Megatron-LM 129.5K tokens/sec/GPU), within 1.4% on the fifth (28-H100 Qwen3-30B-A3B PP2/EP8 BF16: PithTrain 124.9K vs Megatron-LM 126.7K) [Table 4]. TorchTitan OOMs or doesn’t support some configurations; DeepSpeed excluded because PP+EP for MoE is unsupported.
Q&A: across 12 questions, agent turns drop from up to 54 (Megatron-LM) to 9–21 on PithTrain. Per-turn context drops from ~46K to ~33K tokens on the hardest questions [Table 5].
Operate-and-Profile “Getting Started” task: session duration 40.5 → 6.6 min (Megatron-LM → PithTrain), active GPU time 5.4 → 3.1 min, agent turns 88 → 26, output tokens 26.9K → 5.8K [Table 6]. “Train and Evaluate” task: agent turns 163 (Megatron-LM) / 212 (TorchTitan) → 92 (PithTrain).
New-Feature: porting MoBA — session duration 61.6 / 105.1 / 38.7 min (Megatron-LM / TorchTitan / PithTrain), active GPU time 49.5 / 77.9 / 27.7, agent turns 134 / 91 / 57, output tokens 53.8K / 111.8K / 32.4K [Table 7]. DynMoE shows similar ratios (199 / 197 / 76 agent turns).
Skill ablation on validate-correctness: 114 → 34 agent turns (70% drop), 30.2K → 11.3K output tokens (63% drop), at near-identical active GPU time (20.8 → 22.5 min) [Table 8].
MoBA case-study editing-token decomposition: 13.1K (Megatron-LM) / 22.2K (TorchTitan) / 4.7K (PithTrain) on Editing alone — TorchTitan’s elevated cost is OOM-driven (runtime memory failures, not codebase structure) [Fig. 6].

Why it’s interesting

The paper positions itself at the intersection of two concepts the wiki already tracks: it’s another entry in the Distributed training parallelism catalog (FSDP / PP / CP / EP composed for MoE on H100/B200, building on the same DualPipe scaffold that the page already cites via SpectraX — True MPMD Pipeline Parallelism for JAX and Fireworks’s frontier stack), but the novel argument is on the AI-for-AI Research axis: it inverts ASI-ARCH’s “have agents discover architectures” framing to ask “have agents maintain the training framework”, and proposes a measurable framework-design objective (ATE) for that goal. Contrasts cleanly with Scaling and Optimizing Frontier Model Training (Fireworks AI) which optimizes for human-engineer throughput at frontier (1T MoE / 1M context / GB200) — PithTrain accepts narrower scope (Hopper/Blackwell only, fewer features) to win on a different axis. Complements Scalable Training of Mixture-of-Experts Models with Megatron Core (Megatron-Core, the production baseline PithTrain measures against) by reframing what “production-quality” should mean now that AI coding agents are doing the maintenance work. The 4.6× drop in Exploring tokens on the MoBA case study, isolated by category, is the cleanest filed evidence that codebase structure (not just LoC) is doing measurable work on agent cost.