Markdown

PithTrain: agent-native MoE training and agent-task efficiency¶

Scope: PithTrain, a compact (roughly 11K line) Python MoE training framework from CMU, Xlue, and NVIDIA, designed so an AI coding agent can navigate, operate, and extend it cheaply, and the agent-task efficiency (ATE) metric plus ATE-Bench benchmark the same paper introduces to measure that cost. Covers the four agent-native design principles, the DualPipeV pipeline schedule, the in-repo SKILL.md agent skills, the training-throughput results against Megatron-LM, DeepSpeed, and TorchTitan, and the ATE-Bench agent-effort results. Verified against arXiv 2605.31463v1 ("PithTrain: A Compact and Agent-Native MoE Training System", 29 May 2026) and the companion github.com/mlc-ai/pith-train repository. Distinct from MoE expert backends and grouped-GEMM kernels, the kernel layer any of these frameworks dispatch into, and from NeMo AutoModel, a throughput-first MoE fine-tuning accelerator built on Transformers v5; this page is about training-framework design as a cost surface for coding agents, not GPU kernel throughput.

The numpy block under How to develop with it runs on stock python3 (3.12) plus numpy 2.4.6, is asserted, and includes two adversarial cases pulled from the paper's own per-attempt data (Table 12). Every training-throughput and agent-effort number elsewhere on this page is the paper's own reported figure (Tables 1 and 4 through 12, Appendices A and B); none of PithTrain's training runs were reproduced here, since there is no GPU in this environment. The install and bash commands under How to use it are reference templates copied from the upstream README as of 2026-07-14; the repo is at tag v0.1.3 and moving, so pin a commit and re-verify before running.

What it is¶

PithTrain is an end-to-end Mixture-of-Experts (MoE) training framework built around a single design question the paper poses directly: "can we redesign an MoE training framework that optimizes for agent-task efficiency?"² It defines agent-task efficiency (ATE) as "the cost of using coding agents to understand, operate, and extend a framework, measured along dimensions such as session duration and output tokens,"² treats that cost as a dimension orthogonal to training throughput, and builds PithTrain to score well on both. The framework supports pipeline parallelism (PP), FSDP-based data parallelism (DP), context parallelism (CP), expert parallelism (EP), FP8 training, and PyTorch Distributed Checkpoint (DCP) checkpointing, on NVIDIA Hopper and Blackwell GPUs; it also trains dense (non-MoE) models, skipping the MoE-specific machinery such as EP for those.⁶

The codebase is organized into three layers, application, engine, and operator, and the paper reports it at roughly 11K total lines of Python across the pinned v0.1.2 tag,³ against 149K for Megatron-LM and 167K for DeepSpeed at their own pinned commits, and 38K for TorchTitan.³ The companion repository (github.com/mlc-ai/pith-train, Apache-2.0, confirmed live and public as of 2026-07-14) matches this: pithtrain/ holds config.py, dualpipe/, layers/, models/, modules/, operators/, and tasks/, with no plugin-registry or spec-resolution machinery in sight.¹⁹

PithTrain is built on four agent-native design principles, contrasted against Megatron-LM, DeepSpeed, and TorchTitan in the paper's own Table 1:³

Principle	Megatron-LM	DeepSpeed	TorchTitan	PithTrain
Compact codebase	No (149K LoC)	No (167K LoC)	Partial (38K LoC)	Yes (11K LoC)
Python-native	No	No	Yes	Yes
No implicit indirection	No	No	Partial	Yes
Agent skills	Partial	No	Partial	Yes

The paper's own footnote pins the exact snapshots this table measures: "LoC measured at the pinned commits 3bec9aa (Megatron-LM), 44c51e3 (DeepSpeed), d84e83d (TorchTitan), and PithTrain at v0.1.2."³ The companion benchmark, ATE-Bench, is described as the inverse of SWE-bench and similar agent-capability suites: "We invert this with ATE-Bench... holding the agent fixed so that differences in agent cost isolate framework design."² Where SWE-bench varies the agent against a fixed codebase to score capability, ATE-Bench varies the framework (Megatron-LM, TorchTitan, PithTrain) against one fixed agent, Claude Code running Opus 4.7 at "xhigh" effort, and one fixed task suite, to score how much the framework's own design costs that agent.⁴

Why use it¶

Throughput parity at a fraction of the code. PithTrain "matches or exceeds Megatron-LM on 4 of 5 configurations, and stays within 1.4% of Megatron-LM on the fifth"⁸ (details in Architecture and How to run it in production), while running roughly 13x fewer lines than Megatron-LM and 15x fewer than DeepSpeed.³
Lower measured agent cost on real tasks. Across the three ATE-Bench categories the paper reports, against Megatron-LM: up to 67% fewer Agent Turns on Q&A,⁹ up to 70% fewer Agent Turns and 78% fewer Output Tokens on Operate-and-Profile,¹⁰ and up to 44% less Active GPU Time on New-Feature integration (64% against TorchTitan).¹¹
Skills measurably help, isolated from codebase structure. Toggling two of PithTrain's own skills on and off on an otherwise identical commit dropped Agent Turns 70% and 52% on the two ablated tasks, with Active GPU Time held near constant, since "each task runs a fixed set of training runs pinned by the workflow, so the GPU work is determined by the task rather than the agent."¹² That isolates the skill's effect from the compact-codebase effect, which the paper does not do for its top-line framework comparison.
A principled, checkable torch.compile boundary. PithTrain compiles with fullgraph=True, the strict mode that "rejects graph breaks at compile time rather than silently degrading speedup," and excludes only the MoE forward and backward, "because per-expert input shapes are data-dependent under expert parallelism."⁶ That is the identical constraint MoE expert backends and grouped-GEMM kernels documents for HuggingFace Transformers' eager backend, "not compatible with full graph compilation" for the same data-dependent-shape reason.

Read the headline percentages carefully. The abstract states PithTrain achieves higher agent-task efficiency "with up to 62% fewer Agent Turns and 64% less Active GPU Time,"¹ and the introduction repeats this, scoped to "the hardest new-feature tasks."² Section 5.2's own New-Feature paragraph, however, states an Active GPU Time percentage (44% against Megatron-LM, 64% against TorchTitan) but never restates an Agent-Turns percentage for this category at all.¹¹ The 64% figure is traceable directly to prose. The 62% figure is not; it has to be reconstructed from Table 7's per-task medians, and doing that (see the executed block below) shows it is the Differential Transformer task's Agent-Turns reduction against Megatron-LM (125 to 47 turns, 62.4%), a different task, metric pairing than the 64% figure (the MoBA task's Active GPU Time reduction against TorchTitan, 77.9 to 27.7 minutes, 64.4%). The abstract compresses two different tasks' two different maxima, against two different baseline frameworks, into one sentence that reads as a single comparison. Neither number is wrong, but they do not describe the same run.

When to use it (and when not)¶

Use it for research or prototyping-scale MoE training where you want a codebase a coding agent, or a human, can hold in one context pass and extend quickly, on Hopper or Blackwell hardware, for one of the demonstrated architectures (the Qwen3 MoE family, DeepSeek-V2/V3-style models, GPT-OSS) or a model you are willing to port yourself with the add-new-model skill, which is itself one of ATE-Bench's own benchmarked tasks.
Do not expect Megatron-LM's or DeepSpeed's breadth yet. Throughput parity is demonstrated on five configurations across three model families, on up to 4 nodes (32 H100s) or a single 8x B200 node.⁸ Nothing in the paper shows PithTrain at the multi-thousand-GPU scale TorchTitan explicitly targets,¹⁷ and the project is young: four release tags (v0.1.0 through v0.1.3 as of 2026-07-14) and one arXiv preprint, not years of production hardening.
Do not read Table 4 as a PithTrain-versus-DeepSpeed result. DeepSpeed is absent from every throughput row because, at the paper's pinned commit, "the Megatron-DeepSpeed integration script states 'Currently we don't support PP for MoE,'"⁸ so DeepSpeed was never run on any of the five configurations. That is a documented DeepSpeed gap, not a PithTrain win.
Weigh TorchTitan's own gaps in the same table. TorchTitan is marked unsupported ("-") for GPT-OSS-20B with pipeline parallelism and OOMs on two of the four remaining rows,⁸ so three of five throughput comparisons against TorchTitan are incomplete by TorchTitan's own limits, not head-to-head losses.
Treat the ATE numbers as evidence for one fixed agent. Every ATE-Bench result is Claude Code running Opus 4.7 at xhigh effort, "completed every task across all attempts and frameworks with no failure."⁹ The paper does not test whether the same framework-design effect holds for a different agent, a different model family, or a lower-effort setting.
Note where an independent, same-vendor model judges the agent. New-Feature correctness rules are graded by "an independent claude-opus-4-7 session at xhigh effort,"¹⁴ the same model family as the benchmarked agent. The paper mitigates this with a mandatory human spot-check of every verdict ("a human reader opens the diff at the cited file and confirms the line evidence for both PASS judgements and any FAIL or partial verdicts"¹⁴), which is a real safeguard, but the judge-and-agent affinity is worth knowing before treating "every attempt satisfies all three rules" as fully independent verification.

Architecture¶

Three layers, application, engine, and operator, with two structural choices load-bearing throughout: a flat, one-file-per-model layout instead of a shared layer skeleton, and a pipeline schedule (DualPipeV) that owns its own compute/communication overlap rather than inheriting it from a generic scaffold.

flowchart TB
  subgraph PRINC["Four agent-native design principles"]
    P1["Compact codebase: constraint on growth, not a snapshot"]
    P2["Python-native: no compiled-extension rebuild cycle"]
    P3["No implicit indirection: flat, direct calls"]
    P4["Task-specific in-repo agent skills (SKILL.md)"]
  end
  subgraph STACK["PithTrain, three layers, roughly 11K LoC"]
    APP["Application: pretrain / SFT training loop"]
    subgraph ENGINE["Engine"]
      MODEL["Model: Qwen / DeepSeek / GPT-OSS, one file each"]
      PIPE["Pipeline Engine: DualPipeV, 5-stage EP-boundary overlap"]
      BLOCKS["Building Blocks: FP8 linear, ring attention, expert dispatch+dedup, deferred wgrad"]
      DIST["Distributed Training: PP x FSDP x CP x EP mesh"]
      INFRA["Training Infrastructure: torch.compile(fullgraph=True), optimizer, DCP checkpointing"]
    end
    OP["Operator: PyTorch/NCCL basic ops, DeepGEMM/FlashAttention, Triton DSL kernels"]
  end
  APP --> MODEL --> PIPE --> OP
  DIST --> PIPE
  INFRA --> APP
  PRINC -.shape.-> STACK
  STACK --> THR["Training throughput: matches or exceeds Megatron-LM on 4 of 5 configs"]
  STACK --> ATE["ATE-Bench: fixed Claude Opus 4.7 agent, Q&A / Operate-and-Profile / New-Feature"]

The most consequential architectural choice is the one Figure 2 of the paper illustrates directly: in Megatron-LM, TransformerLayer resolves self.mlp from a ModuleSpec built in a separate file (gpt_layer_specs.py), so reading transformer_layer.py alone does not tell you what actually runs. In PithTrain's Qwen3DecoderLayer, self.mlp = Qwen3MoE(cfg) if use_moe else Qwen3MLP(cfg) is defined in the same file the class lives in.⁵ "The implicit-indirection pattern resolves submodules through a runtime spec, supporting model variation from a shared layer skeleton; PithTrain instantiates layers directly, favoring local readability."⁵ That trade is explicit: Megatron-LM's registry buys cross-model code reuse, PithTrain's flat structure buys local readability, and the paper picks the latter deliberately, at the cost of some duplication across model files.

The pipeline schedule, DualPipeV, "builds on DualPipe from DeepSeek-V3,"⁶ the same schedule comms-compute overlap documents with its (PP/2-1)(F&B+B-3W) bubble formula and even-pipeline-stage requirement. PithTrain's contribution on top of the open-source DualPipe scaffold is the overlap itself: "DeepSeek's open-source version provides a minimal pipeline-orchestration scaffold, on top of which PithTrain implements the actual compute-communication overlap. Each transformer layer is decomposed into five stages at EP communication boundaries. EP all-to-alls run on a separate communication stream, and the schedule overlaps the forward of one micro-batch with the backward of another."⁶ The "deferred wgrad" building block in Figure 3 is the same weight-gradient deferral (splitting backward into B and W) that produces the -2W and -3W terms in the Zero Bubble and DualPipe bubble formulas; PithTrain cites the Zero Bubble Pipeline Parallelism paper directly for it.¹⁶

How to use it¶

Installation and the basic pretraining loop, copied from the upstream README (reference template, unexecuted; pin a commit and verify against the current repo):¹⁸

# Reference template. Requires NVIDIA Hopper (SM90) or Blackwell (SM100), CUDA >= 13.0,
# Python >= 3.12. Uses uv for dependency management.
git clone https://github.com/mlc-ai/pith-train.git && cd pith-train
uv venv                       # skip if you already have a virtual environment
uv pip install .              # for users; `uv sync` instead for a development checkout

# 1. Tokenize the DCLM pretraining corpus into mmap-friendly packed sequences.
bash examples/tokenize_corpus/launch.sh dclm-qwen3

# 2. Edit examples/pretrain_lm/qwen3-30b-a3b/script.py and config.json for parallelism,
#    batch size, learning rate. The launch script auto-detects GPUs and supports both
#    single-node and multi-node (SLURM) setups; training resumes from the latest
#    checkpoint automatically.
bash examples/pretrain_lm/launch.sh qwen3-30b-a3b

# 3. Export a training checkpoint to standard HuggingFace format (also imports HF
#    checkpoints for continued pretraining).
bash examples/convert_checkpoint/launch.sh qwen3-30b-a3b

The ATE-Bench Operate-and-Profile and New-Feature tasks all share one fixed configuration rather than the pretraining example above: base model DeepSeek-V2-Lite, mesh PP=4, EP=2, DP=1, sequence length 2048, global batch size 1024, precision BF16, chosen so the whole suite "fits within a single node with 8 NVIDIA H100 GPUs."¹⁴

How to develop with it¶

The four design principles are development-time constraints as much as architectural claims. "We treat compactness as a constraint on growth: PithTrain may grow over time, but additions should respect the four principles,"⁵ which means a contribution that reintroduces a plugin registry, or reaches for a compiled extension where a Python-native op will do, works against the paper's own stated intent even if it improves some other metric.

Agent skills follow a fixed contract. "Each skill is a self-contained folder with a top-level SKILL.md playbook, optionally additional markdown documents, and optionally helper scripts,"⁷ designed around three properties: "specific scope, explicit prerequisites, and verifiable success."⁷ The paper's own worked example, validate-correctness, ends in a script call rather than the agent's self-assessment: "worktree the base branch, run N steps on each branch, run python compare.py base.log feat.log, exit 0 = PASS or 1 = FAIL, clean up worktree."⁷ Table 3 names four representative skills: add-new-model (port a new MoE architecture end-to-end: model, FSDP, tests), add-memory-prints (instrument training for memory profiling), capture-nsys-profile (capture an Nsight Systems profile of a short PithTrain run), and validate-correctness (compare per-step loss curves between branches).⁷ The live repository, checked directly against .agents/skills/ on 2026-07-14, carries eight: those four plus analyze-nsys-profile, estimate-memory, launch-with-slurm, and setup-benchmark-inputs.¹⁹ The repo has grown past what the paper's own table shows, which is itself a small instance of the pattern this KB tracks on every fast-moving upstream: verify against the current source, not the paper's snapshot.

The torch.compile(fullgraph=True) boundary is a discipline, not a knob. Because it is strict mode, adding a new op that introduces a graph break inside the compiled region fails loudly at compile time. The one carve-out, the MoE forward and backward, exists for the identical reason MoE expert backends and grouped-GEMM kernels documents on the HuggingFace side: which experts a batch of tokens actually hits is a runtime, data-dependent property under expert parallelism, so the shape the compiler would need to specialize on is not known ahead of time.

Validating ATE-Bench's own arithmetic¶

ATE-Bench reports "the median of three independent attempts" per task and metric,⁴ and the abstract's headline percentages are maxima taken across that median table. Nothing about that arithmetic can be checked from the paper's prose alone, since Section 5.2 states only the final medians (Table 7), not the raw attempts. Appendix C's Table 12, however, publishes every individual attempt, which makes the aggregation independently checkable: recompute the median from the three raw attempts and diff it against the paper's own reported Table 7 value. The block below does that for every New-Feature task, framework, and the two headline metrics (Agent Turns, Active GPU Time), transcribed directly from the PDF, then reconstructs the "up to 44%/64%/62%" claims from the raw data and stress-tests the aggregation with two cases the paper's own numbers actually contain: a 3x outlier attempt, and a task where PithTrain used more effort than a baseline.

# ate_model.py, validated: system python3 + numpy 2.4.6.
# Reproduces ATE-Bench's median-of-three aggregation (arXiv 2605.31463 Sec. 4/5.2) from the
# paper's OWN per-attempt data (Table 12) and cross-checks it byte-exact against the paper's
# OWN reported medians (Table 7). Then derives the "up to 62% fewer Agent Turns / 64% less
# Active GPU Time" abstract headline from the New-Feature category and stress-tests the
# aggregation with two adversarial cases present in the paper's own numbers: a 3x outlier
# attempt, and a task/baseline/metric triple where PithTrain is WORSE, not better.
import numpy as np

# Table 12 per-attempt data (New-Feature tasks), Agent Turns and Active GPU Time (minutes).
TURNS = {
    "diff":   {"megatron": (62, 125, 141),  "torchtitan": (61, 53, 58),   "pithtrain": (47, 35, 83)},
    "dynmoe": {"megatron": (200, 199, 186), "torchtitan": (253, 126, 197), "pithtrain": (76, 73, 118)},
    "moba":   {"megatron": (134, 135, 91),  "torchtitan": (61, 104, 91),  "pithtrain": (57, 55, 66)},
    "moepp":  {"megatron": (829, 145, 140), "torchtitan": (85, 87, 125),  "pithtrain": (141, 85, 90)},
}
GPU = {
    "diff":   {"megatron": (33.7, 44.9, 33.6), "torchtitan": (40.3, 49.2, 40.1), "pithtrain": (27.6, 33.2, 27.3)},
    "dynmoe": {"megatron": (45.4, 49.1, 62.1), "torchtitan": (97.8, 94.4, 90.3), "pithtrain": (41.9, 42.3, 41.3)},
    "moba":   {"megatron": (49.1, 49.5, 49.9), "torchtitan": (39.4, 107.6, 77.9), "pithtrain": (27.7, 32.0, 26.9)},
    "moepp":  {"megatron": (78.0, 58.7, 53.6), "torchtitan": (72.1, 51.9, 45.2), "pithtrain": (36.7, 39.9, 40.1)},
}

# Table 7's OWN reported medians, for the cross-check.
TABLE7_TURNS = {
    "diff": {"megatron": 125, "torchtitan": 58, "pithtrain": 47},
    "dynmoe": {"megatron": 199, "torchtitan": 197, "pithtrain": 76},
    "moba": {"megatron": 134, "torchtitan": 91, "pithtrain": 57},
    "moepp": {"megatron": 145, "torchtitan": 87, "pithtrain": 90},
}
TABLE7_GPU = {
    "diff": {"megatron": 33.7, "torchtitan": 40.3, "pithtrain": 27.6},
    "dynmoe": {"megatron": 49.1, "torchtitan": 94.4, "pithtrain": 41.9},
    "moba": {"megatron": 49.5, "torchtitan": 77.9, "pithtrain": 27.7},
    "moepp": {"megatron": 58.7, "torchtitan": 51.9, "pithtrain": 39.9},
}


def median3(triple):
    """'median of three independent attempts'. n=3, so the median is just the middle
    order statistic; sorted()[1] and np.median must agree."""
    m = sorted(triple)[1]
    assert np.isclose(m, float(np.median(np.asarray(triple, dtype=float))))
    return m


def pct_reduction(baseline_triple, pith_triple):
    """(baseline - pith) / baseline, on medians. None (not inf, not silently 0) when the
    baseline median is 0: the formula is undefined there, not "100% better"."""
    b, p = median3(baseline_triple), median3(pith_triple)
    if b == 0:
        return None
    return (b - p) / b


# --- Boundary check on median3 itself: order invariance and the degenerate identical triple ---
assert median3((5, 5, 5)) == 5
assert median3((3, 1, 2)) == median3((1, 2, 3)) == median3((2, 3, 1)) == 2

# --- Cross-check: recompute every New-Feature median from Table 12 and diff against Table 7 ---
mismatches = []
for task in TURNS:
    for fw in ("megatron", "torchtitan", "pithtrain"):
        got_t = median3(TURNS[task][fw])
        got_g = median3(GPU[task][fw])
        if got_t != TABLE7_TURNS[task][fw]:
            mismatches.append((task, fw, "turns", got_t, TABLE7_TURNS[task][fw]))
        if not np.isclose(got_g, TABLE7_GPU[task][fw]):
            mismatches.append((task, fw, "gpu", got_g, TABLE7_GPU[task][fw]))
assert not mismatches, mismatches  # every recomputed median must equal the paper's own Table 7

# --- Reproduce the "up to X%" headline claims from raw per-attempt data ---
turns_vs_meg = {t: pct_reduction(TURNS[t]["megatron"], TURNS[t]["pithtrain"]) for t in TURNS}
turns_vs_tt = {t: pct_reduction(TURNS[t]["torchtitan"], TURNS[t]["pithtrain"]) for t in TURNS}
gpu_vs_meg = {t: pct_reduction(GPU[t]["megatron"], GPU[t]["pithtrain"]) for t in GPU}
gpu_vs_tt = {t: pct_reduction(GPU[t]["torchtitan"], GPU[t]["pithtrain"]) for t in GPU}

max_gpu_meg_task = max(gpu_vs_meg, key=gpu_vs_meg.get)
max_gpu_tt_task = max(gpu_vs_tt, key=gpu_vs_tt.get)
max_turns_meg_task = max(turns_vs_meg, key=turns_vs_meg.get)

# The two explicitly-stated body-text numbers (Sec. 5.2, "New Feature" paragraph):
# "up to 44% lower than Megatron-LM's and 64% lower than TorchTitan's" (Active GPU Time).
assert max_gpu_meg_task == "moba" and round(gpu_vs_meg["moba"] * 100) == 44
assert max_gpu_tt_task == "moba" and round(gpu_vs_tt["moba"] * 100) == 64
# The abstract's "up to 62% fewer Agent Turns" is never stated as a New-Feature Agent-Turns
# sentence anywhere in the body text (Sec. 5.2 only gives GPU-time percentages for this
# category); it is only reconstructible from Table 7/12, and the true maximum is here:
assert max_turns_meg_task == "diff" and round(turns_vs_meg["diff"] * 100) == 62

# --- Adversarial 1: a task/baseline/metric triple where PithTrain is WORSE, not better ---
regression_task = min(turns_vs_tt, key=turns_vs_tt.get)
assert turns_vs_tt[regression_task] < 0, "expected a real regression in this dataset"

# --- Adversarial 2: the median suppresses a 3x outlier attempt; the mean would not ---
meg_moepp_turns = TURNS["moepp"]["megatron"]      # (829, 145, 140): attempt 1 is a runaway
pith_moepp_turns = TURNS["moepp"]["pithtrain"]
median_reduction = pct_reduction(meg_moepp_turns, pith_moepp_turns)
mean_reduction = (np.mean(meg_moepp_turns) - np.mean(pith_moepp_turns)) / np.mean(meg_moepp_turns)
assert mean_reduction > 1.5 * median_reduction, "outlier should inflate the mean-based claim"

# --- Boundary: a zero-median baseline must not silently become a fake 100% or divide-by-zero ---
assert pct_reduction((0, 0, 0), (1, 1, 1)) is None

print("median3 cross-check: 24/24 Turns+GPU medians recomputed from Table 12 match Table 7 exactly")
print()
print("Active GPU Time reduction vs Megatron-LM (New Feature):",
      {k: f"{v*100:.1f}%" for k, v in gpu_vs_meg.items()}, "-> max", max_gpu_meg_task, f"{gpu_vs_meg[max_gpu_meg_task]*100:.1f}%")
print("Active GPU Time reduction vs TorchTitan   (New Feature):",
      {k: f"{v*100:.1f}%" for k, v in gpu_vs_tt.items()}, "-> max", max_gpu_tt_task, f"{gpu_vs_tt[max_gpu_tt_task]*100:.1f}%")
print("Agent Turns    reduction vs Megatron-LM (New Feature):",
      {k: f"{v*100:.1f}%" for k, v in turns_vs_meg.items()}, "-> max", max_turns_meg_task, f"{turns_vs_meg[max_turns_meg_task]*100:.1f}%")
print("Agent Turns    reduction vs TorchTitan   (New Feature):",
      {k: f"{v*100:.1f}%" for k, v in turns_vs_tt.items()})
print()
print(f"ADVERSARIAL regression: {regression_task} Agent Turns vs TorchTitan = "
      f"{turns_vs_tt[regression_task]*100:+.1f}% (PithTrain used MORE turns; 'up to X%' is a max, not a floor)")
print(f"ADVERSARIAL outlier: MoE++/Megatron-LM Turns {meg_moepp_turns} -> "
      f"median-based reduction {median_reduction*100:.1f}%, mean-based reduction {mean_reduction*100:.1f}% "
      f"(mean inflates the claim {mean_reduction/median_reduction:.2f}x off one runaway attempt)")
print("boundary: zero-median baseline ->", pct_reduction((0, 0, 0), (1, 1, 1)))

Executed output:

median3 cross-check: 24/24 Turns+GPU medians recomputed from Table 12 match Table 7 exactly

Active GPU Time reduction vs Megatron-LM (New Feature): {'diff': '18.1%', 'dynmoe': '14.7%', 'moba': '44.0%', 'moepp': '32.0%'} -> max moba 44.0%
Active GPU Time reduction vs TorchTitan   (New Feature): {'diff': '31.5%', 'dynmoe': '55.6%', 'moba': '64.4%', 'moepp': '23.1%'} -> max moba 64.4%
Agent Turns    reduction vs Megatron-LM (New Feature): {'diff': '62.4%', 'dynmoe': '61.8%', 'moba': '57.5%', 'moepp': '37.9%'} -> max diff 62.4%
Agent Turns    reduction vs TorchTitan   (New Feature): {'diff': '19.0%', 'dynmoe': '61.4%', 'moba': '37.4%', 'moepp': '-3.4%'}

ADVERSARIAL regression: moepp Agent Turns vs TorchTitan = -3.4% (PithTrain used MORE turns; 'up to X%' is a max, not a floor)
ADVERSARIAL outlier: MoE++/Megatron-LM Turns (829, 145, 140) -> median-based reduction 37.9%, mean-based reduction 71.6% (mean inflates the claim 1.89x off one runaway attempt)
boundary: zero-median baseline -> None

Two findings fall out of the raw data that the paper's own prose does not spell out. First, on the MoE++ task, one of Megatron-LM's three attempts spirals across every metric at once (829 Agent Turns, 516.5K per-turn context, 326.4K output tokens, 1.5x to nearly 6x its own other two attempts depending on the metric),¹⁵ and the median-of-three convention suppresses it cleanly: the reported 145-turn median matches the two well-behaved attempts, not the outlier. Reporting the mean instead would have inflated PithTrain's advantage on this exact comparison from a real 37.9% to an artifact-driven 71.6%, nearly double, off a single anomalous session. Second, "up to X%" is a maximum, not a per-task guarantee: on MoE++ specifically, PithTrain used more Agent Turns than TorchTitan's median (90 against 87, a 3.4% regression). The paper's own framing is honest about this ("up to"), but a reader skimming only the abstract's two headline numbers would not know that a real, verified regression sits in the same table.

How to maintain it¶

Re-run validate-correctness on every change. It is the paper's own recommended workflow for exactly this: compare per-step loss curves between the base branch and a feature branch, PASS or FAIL by script, not by the agent's self-assessment.⁷
Pin an exact commit, not a tag. The paper itself measures PithTrain at two different points: 23db182 for the Table 4 training-throughput comparisons,⁸ and the v0.1.2 tag for the Table 1 line-of-code count.³ Checking the four published release tags directly against the GitHub API on 2026-07-14 shows none of them, v0.1.0 through v0.1.3, matches commit 23db182; the throughput benchmark's pinned commit is an untagged snapshot distinct from any tagged release.¹⁹ That is ordinary practice for a fast-moving research repo measured over time, not a defect, but it means "PithTrain" in this paper does not refer to one single fixed snapshot throughout, and a reader who wants to reproduce a specific number should pin the specific commit that number is attributed to, not the nearest release tag.
Watch the growth-versus-compactness tension mechanically. "Compactness as a constraint on growth"⁵ is a code-review norm, not something check_pages.py-style tooling can gate on here; the repo already carries four skills beyond what Table 3 shows (see How to develop with it), which is the kind of drift worth tracking explicitly as the project moves past its v0.1.x line.

How to run it in production¶

Correctness is validated against Megatron-LM at one configuration, not swept across the whole matrix: Qwen3-30B-A3B, 4x8-H100, PP4, DP1, CP1, EP8, DCLM dataset, sequence length 2048, global batch 1024, 4096 steps, Adam, cosine decay from 2e-4 to 1e-5.¹³ "The trajectories align across the full run, with Megatron-LM showing two transient spikes that recover within a few steps,"¹³ and six downstream benchmarks (OpenBookQA and WinoGrande 0-shot; ARC-Challenge, ARC-Easy, HellaSwag, and PIQA 10-shot) "achieve comparable accuracy within statistical noise across all six benchmarks at every checkpoint."¹³
Multi-node launch is SLURM-based, and checkpoints are DCP-backed and explicitly "reshardable across different parallelism," with the convert_checkpoint tool round-tripping to and from standard HuggingFace safetensors for handoff to serving.¹⁸ See weight loading in inference engines for what that HF-format checkpoint has to satisfy on the serving side.
Demonstrated scale tops out at 4 nodes. The paper's own throughput table covers a single 8x B200 node and up to 4x8 H100 nodes (32 GPUs);⁸ nothing here demonstrates the framework at the scale distributed training recipes or pipeline parallelism discuss for frontier training runs. Treat multi-thousand-GPU deployment as unverified until you have measured it yourself.
Do not compute a single MFU number across mixed precision. The paper deliberately omits Model FLOPs Utilization from its own throughput table "because Tensor Core peak FLOPS differs between BF16 and FP8, making the metric ambiguous for mixed-precision steps."⁸ Use tokens/sec, the paper's own metric, as the comparable figure across the BF16 and FP8 rows in Table 4, and treat that as a caution against building an MFU-based SLO or MFU regression alert that silently mixes precisions.

Failure modes¶

Trusting the abstract's compressed headline as one comparison. "Up to 62% fewer Agent Turns and 64% less Active GPU Time" describes two different New-Feature tasks against two different baseline frameworks, not one run; see Why use it and the executed block above.
Reading Table 4 as PithTrain beating DeepSpeed. DeepSpeed was never run on any of the five throughput configurations; it is absent because it lacks PP+EP support for MoE at the pinned commit, a DeepSpeed limitation the paper documents rather than a measured PithTrain advantage.
Assuming TorchTitan's OOM and unsupported cells are PithTrain wins. Three of five rows involve a TorchTitan limitation (no pipeline parallelism for GPT-OSS-20B, OOM on two Qwen3-30B-A3B configurations) rather than a measured throughput loss.
Skipping the judge-affinity caveat on New-Feature grading. The rule-based correctness judge is the same model family (Opus 4.7 xhigh) as the benchmarked agent; the paper's mandatory human spot-check of every verdict is the mitigation, and it is worth confirming that check happened before trusting "all 36 attempts satisfy all three rules" at face value.
Pinning v0.1.2 and expecting the Table 4 throughput numbers. Those numbers were measured at commit 23db182, which independent verification against the repo's release tags shows is not any of the four tagged releases; pin the throughput-specific commit if you need to reproduce Table 4, not the LoC-table tag.
Losing the torch.compile(fullgraph=True) boundary on a new feature. Adding code to the compiled region that introduces a data-dependent shape outside the deliberately-excluded MoE forward/backward will hard-fail at compile time (the intended behavior), not silently degrade; do not "fix" this by disabling fullgraph globally, which trades away the strict-mode guarantee for the whole framework.
Generalizing the ATE-Bench results past one agent. Every Agent Turns, Active GPU Time, and Output Token number on this page describes Claude Code running Opus 4.7 at xhigh effort. A different agent, harness, or effort setting is untested here.

References¶

PithTrain: A Compact and Agent-Native MoE Training System, arXiv:2605.31463v1 (29 May 2026) — https://arxiv.org/abs/2605.31463 · PDF fetched and read in full for this page: https://arxiv.org/pdf/2605.31463
PithTrain repository (Apache-2.0, confirmed live 2026-07-14) — https://github.com/mlc-ai/pith-train
PithTrain announcement blog — https://blog.mlc.ai/2026/06/01/pithtrain-compact-agent-native-moe-training-system
DeepSeek-AI, DualPipe (the pipeline-orchestration scaffold DualPipeV builds on) — https://github.com/deepseek-ai/DualPipe
Qi et al., Zero Bubble (Almost) Pipeline Parallelism, ICLR 2024 (the wgrad-delay technique PithTrain's "Deferred Wgrad" building block implements) — https://arxiv.org/abs/2401.10241
Qiu et al., Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models, arXiv:2501.11873 (the "Global-Level" auxiliary loss used in the Appendix A pretraining hyperparameters)
Anthropic, Equipping agents for the real world with Agent Skills — https://www.anthropic.com/engineering/equipping-agents-for-the-real-world-with-agent-skills
Jimenez et al., SWE-bench: Can Language Models Resolve Real-World GitHub Issues?, ICLR 2024 — https://arxiv.org/abs/2310.06770
Liang et al., TorchTitan: One-stop PyTorch Native Solution for Production Ready LLM Pretraining, ICLR 2025 — https://arxiv.org/abs/2410.06511
Lu et al., MoBA: Mixture of Block Attention for Long-Context LLMs, arXiv:2502.13189 — https://arxiv.org/abs/2502.13189
Guo et al., Dynamic Mixture of Experts: An Auto-Tuning Approach for Efficient Transformer Models, arXiv:2405.14297 — https://arxiv.org/abs/2405.14297
Ye et al., Differential Transformer, arXiv:2410.05258 — https://arxiv.org/abs/2410.05258
Li et al., DataComp-LM: In Search of the Next Generation of Training Sets for Language Models, arXiv:2406.11794 (the DCLM pretraining corpus used in both the correctness validation and ATE-Bench) — https://arxiv.org/abs/2406.11794

arXiv 2605.31463v1, Abstract: "PithTrain enables higher agent-task efficiency, with up to 62% fewer Agent Turns and 64% less Active GPU Time." ↩
Sec. 1, Introduction: "We define agent-task efficiency (ATE) as the cost of using coding agents to understand, operate, and extend a framework, measured along dimensions such as session duration and output tokens"; "Plugin systems, registry-based indirection, and heavy compiled extensions raise the cost of locating relevant code, tracing what runs at a given call site, and verifying a change is complete"; "We invert this with ATE-Bench, a comprehensive benchmark that varies the framework on real-world training-framework tasks, holding the agent fixed so that differences in agent cost isolate framework design"; "a coding agent completes the same training-framework tasks on PithTrain with up to 62% fewer Agent Turns and 64% less Active GPU Time than on production frameworks for the hardest new-feature tasks, with similar reductions on other metrics." ↩↩↩↩
Sec. 3.1, Table 1 and its footnote 3: LoC by framework (Megatron-LM 149K, DeepSpeed 167K, TorchTitan 38K, PithTrain 11K) with satisfied/partial/not-satisfied marks across the four principles; "LoC measured at the pinned commits 3bec9aa (Megatron-LM), 44c51e3 (DeepSpeed), d84e83d (TorchTitan), and PithTrain at v0.1.2." ↩↩↩↩↩↩
Sec. 4, ATE-Bench: task categories (Q&A 12 tasks, Operate and Profile 4 tasks, New Feature 4 tasks); "Each metric reports the median of three independent attempts"; the fixed agent is "Claude Code (Opus 4.7 at xhigh effort)" (Sec. 5.2). ↩↩
Sec. 3.1: "We treat compactness as a constraint on growth: PithTrain may grow over time, but additions should respect the four principles"; "Production frameworks compose many model variants from a shared layer skeleton via implicit indirection... This pattern enables code reuse across models, while making what runs at a given call site harder to identify by local reading"; Figure 2 caption: "The implicit-indirection pattern resolves submodules through a runtime spec, supporting model variation from a shared layer skeleton; PithTrain instantiates layers directly, favoring local readability." ↩↩↩↩
Sec. 3.2: "PithTrain supports standard MoE training with pipeline parallelism (PP), data parallelism (DP) via FSDP, context parallelism (CP), expert parallelism (EP), FP8 training, and DCP checkpointing, on NVIDIA Hopper and Blackwell GPUs"; footnote 4: "PithTrain also supports dense (non-MoE) models, with MoE-specific machinery (e.g., expert parallelism) skipped"; "PithTrain's pipeline scheduler builds on DualPipe from DeepSeek-V3... Each transformer layer is decomposed into five stages at EP communication boundaries. EP all-to-alls run on a separate communication stream, and the schedule overlaps the forward of one micro-batch with the backward of another"; "PithTrain applies torch.compile(fullgraph=True) to all transformer computation except the MoE forward and backward... We exclude the MoE forward and backward because per-expert input shapes are data-dependent under expert parallelism." ↩↩↩↩
Sec. 3.1 and 3.3, Table 3, Figure 4: "Megatron-LM ships six skills for CI/PR hygiene; TorchTitan ships one developer-workflow skill plus four editor rules orienting the agent to the codebase; DeepSpeed ships two markdown files with generic rules. None of these target recurring training-framework tasks"; "Each skill is a self-contained folder with a top-level SKILL.md playbook, optionally additional markdown documents, and optionally helper scripts"; "Each skill in PithTrain is designed around three properties: specific scope, explicit prerequisites, and verifiable success"; the validate-correctness procedure: "1. worktree the base branch, 2. run N steps on each branch, 3. run python compare.py base.log feat.log -> exit 0=PASS, 1=FAIL, 4. clean up worktree." ↩↩↩↩↩
Sec. 5.1, Table 4 and footnotes 5-7: "PithTrain matches or exceeds Megatron-LM on 4 of 5 configurations, and stays within 1.4% of Megatron-LM on the fifth"; footnote 5, "TorchTitan does support GPT-OSS-20B training, but not with pipeline parallelism"; footnote 7, "Verified at commit 44c51e3: the Megatron-DeepSpeed integration script states 'Currently we don't support PP for MoE.'"; "We omit Model FLOPs Utilization (MFU) because Tensor Core peak FLOPS differs between BF16 and FP8, making the metric ambiguous for mixed-precision steps." Table 4 values (tokens/sec): GPT-OSS-20B 1x8-B200 PP2/DP1/CP1/EP4 8192 BF16, Megatron-LM 129.5K, TorchTitan unsupported, PithTrain 140.9K; Qwen3-30B-A3B 1x8-B200 PP2/DP1/CP1/EP4 8192 FP8, Megatron-LM 106.2K, TorchTitan OOM, PithTrain 134.5K; Qwen3-30B-A3B 2x8-H100 PP2/DP1/CP1/EP8 2048 BF16, Megatron-LM 126.7K, TorchTitan 90.5K, PithTrain 124.9K; Qwen3-30B-A3B 4x8-H100 PP4/DP1/CP1/EP8 4096 BF16, Megatron-LM 264.1K, TorchTitan OOM, PithTrain 280.0K; DeepSeek-V2-Lite 1x8-H100 PP2/DP1/CP1/EP4 2048 BF16, Megatron-LM 107.3K, TorchTitan 74.1K, PithTrain 114.6K. ↩↩↩↩↩↩↩
Sec. 5.2, "Q&A" paragraph: "Across the 12 questions, the agent uses up to 67% fewer Agent Turns to reach the final answer on PithTrain than on Megatron-LM"; "Opus 4.7 completed every task across all attempts and frameworks with no failure." ↩↩
Sec. 5.2, "Operate and Profile" paragraph: "PithTrain's Agent Turns are up to 70% lower than Megatron-LM's and 57% lower than TorchTitan's, and its Output Tokens are up to 78% and 65% lower respectively... the agent invokes in-repo skills on its own when applicable; for example, the Report Heavy Kernels task triggers capture-nsys-profile." ↩
Sec. 5.2, "New Feature" paragraph: "PithTrain's Active GPU Time is up to 44% lower than Megatron-LM's and 64% lower than TorchTitan's, primarily because PithTrain converges in fewer training runs. Two patterns inflate Megatron-LM's reruns: a hidden argument registry causes the agent's manually-added CLI flags to collide with auto-derived ones (implicit indirection), and C++ paths like TransformerEngine's grouped-GEMM emit opaque segfaults that drive speculative configuration toggles (not Python-native). TorchTitan's reruns are dominated by memory-pressure debugging." ↩↩
Sec. 5.3: "Agent turns drop the most (70% and 52% respectively)"; "Active GPU Time stays near parity across both tasks: each task runs a fixed set of training runs pinned by the workflow, so the GPU work is determined by the task rather than the agent." Table 8 (minutes / count): validate-correctness skills off, Session 26.0, GPU 20.8, Turns 114, Context 96.3K, Tokens 30.2K; skills on, Session 22.9, GPU 22.5, Turns 34, Context 43.5K, Tokens 11.3K. capture-nsys-profile skills off, Session 9.4, GPU 5.6, Turns 75, Context 62.1K, Tokens 25.3K; skills on, Session 6.6, GPU 5.5, Turns 36, Context 40.3K, Tokens 14.5K. Note validate-correctness's own Active GPU Time rises slightly with skills on (20.8 to 22.5 minutes), consistent with the paper's fixed-workflow explanation but not itself an improvement on that one metric. ↩
Appendix A: training config (Qwen3-30B-A3B, 4x8-H100, PP4/DP1/CP1/EP8, DCLM, seq 2048, BF16, global batch 1024, 4096 steps, 128 warmup steps, Adam, max LR 2e-4, min LR 1e-5, cosine decay, Global-Level aux loss, coefficient 1e-3); "the trajectories align across the full run, with Megatron-LM showing two transient spikes that recover within a few steps"; downstream evaluation on OpenBookQA and WinoGrande (0-shot), ARC-Challenge, ARC-Easy, HellaSwag, and PIQA (10-shot), "comparable accuracy within statistical noise across all six benchmarks at every checkpoint." ↩↩↩
Appendix B: fixed configuration for Operate-and-Profile and New-Feature tasks, "the base model is DeepSeek-V2-Lite (its training fits within a single node with 8 NVIDIA H100 GPUs), the parallelism mesh is PP=4, EP=2, DP=1, sequence length 2048, global batch size 1024, precision BF16"; New-Feature grading: "Each attempt is judged by an independent claude-opus-4-7 session at xhigh effort, given the three rules and the git diff produced by the agent against the main branch of the framework; the judge returns PASS or FAIL per rule with a one-sentence justification quoting a specific code reference... We additionally inspect every verdict and its cited justification by hand: a human reader opens the diff at the cited file and confirms the line evidence for both PASS judgements and any FAIL or partial verdicts." ↩↩↩
Appendix C, Table 12: per-attempt New-Feature data. MoE++, Megatron-LM: attempt 1, Session 144.6 min, Active GPU 78.0 min, Agent Turns 829, Per-Turn Context 516.5K, Output Tokens 326.4K; attempt 2, 88.5 / 58.7 / 145 / 170.8K / 117.0K; attempt 3, 75.5 / 53.6 / 140 / 188.7K / 86.9K. ↩
Sec. 3.2, "Other optimizations": "PithTrain also implements wgrad delay... fused SwiGLU forward and backward kernels for throughput and reduced activation memory; EP dispatch deduplication for lower all-to-all communication volume; an FP8 weight cache across micro-batches to avoid re-quantization; and fused Triton kernels for EP token scatter and FP8 quantization," citing Qi et al.'s Zero Bubble (Almost) Pipeline Parallelism for the wgrad-delay technique. ↩
Sec. 2, Related Work: "TorchTitan more recently scaled PyTorch-native training to multi-thousand-GPU clusters, sharing PithTrain's Python-native goal but without making agent-native design a primary axis or matching production-framework throughput on our configurations." ↩
mlc-ai/pith-train README, fetched 2026-07-14: installation (uv venv, uv pip install . for users, uv sync for developers; NVIDIA Hopper SM90 or Blackwell SM100, CUDA >= 13.0, Python >= 3.12); getting-started steps (tokenize_corpus, edit script.py/config.json, pretrain_lm launch script, convert_checkpoint); "The launch script auto-detects GPUs and supports both single-node and multi-node (SLURM) setups. Training resumes from the latest checkpoint automatically, and checkpoints are reshardable across different parallelism." ↩↩
Independent verification against the live repository, 2026-07-14 (not from the paper): github.com/mlc-ai/pith-train is public, Apache-2.0 licensed, at release tags v0.1.0 through v0.1.3 (commits e610429, cc8c3e9, 2a97744, 42e35c1 respectively, none matching the paper's throughput-benchmark commit 23db182); pithtrain/ contains config.py, dualpipe/, layers/, models/, modules/, operators/, tasks/; .agents/skills/ (symlinked from .claude/skills/) contains eight skill directories: add-memory-prints, add-new-model, analyze-nsys-profile, capture-nsys-profile, estimate-memory, launch-with-slurm, setup-benchmark-inputs, validate-correctness. ↩↩↩