ATE-Bench: ML Agent Efficiency Benchmark

• ATE-Bench is a benchmark measuring agent-task efficiency by quantifying the operational cost of understanding, operating, and extending ML training frameworks.
• It reports detailed metrics—including session duration, active GPU time, agent turns, and token counts—across a suite of real-world tasks.
• The benchmark highlights trade-offs in framework design by isolating cost dimensions without aggregating them into a single composite score.

ATE-Bench is a benchmark for agent-task efficiency (ATE) in ML training frameworks, introduced alongside PithTrain as a way to measure the operational cost incurred by a coding agent when it must understand, operate, and extend a framework rather than merely train models at high throughput (Lai et al., 29 May 2026). It is motivated by the observation that conventional framework evaluations emphasize training throughput such as tokens/sec or MFU, while leaving unmeasured the effort required to navigate large codebases, resolve indirection, debug failures, and execute edit–run–profile loops. ATE-Bench therefore holds the agent and the tasks fixed, varies only the framework, and reports a vector of effort metrics—session duration, active GPU time, agent turns, per-turn context tokens, and output tokens—across a suite of real-world training-framework tasks (Lai et al., 29 May 2026).

1. Definition and conceptual scope

ATE-Bench operationalizes agent-task efficiency as the cost incurred by a coding agent to complete real training-system tasks on a given framework (Lai et al., 29 May 2026). The benchmark is explicitly framed against production frameworks that have grown large and complex, often relying on layered plugin architectures, registry-based indirection, and C++/CUDA extensions. According to the benchmark description, these design choices may support human experts and peak throughput, but they also create hidden costs for agents: more code to inspect, harder call-site resolution, non-Python error surfaces, and rebuild cycles after changes (Lai et al., 29 May 2026).

The benchmark’s central methodological choice is to treat ATE as a multi-dimensional resource-and-time measurement problem rather than as a correctness-only evaluation. The measured costs are: agent time as session duration, GPU time as active GPU minutes consumed during agent-initiated training or profiling runs, and agent operations and usage as tool-action counts and token consumption (Lai et al., 29 May 2026). This makes ATE-Bench distinct from evaluations that ask whether an agent can solve a task at all; in the reported experiments, all tasks across all attempts and frameworks succeeded, so the differentiating signal lies in effort rather than binary success (Lai et al., 29 May 2026).

A common misconception is to treat ATE-Bench as a throughput benchmark. It is not. Throughput comparisons appear in the surrounding PithTrain study, but they are described as orthogonal to ATE-Bench. Another misconception is to assume that ATE-Bench produces a single scalar ranking. The benchmark explicitly reports each effort dimension independently and does not aggregate across tasks into a composite index (Lai et al., 29 May 2026).

2. Benchmark structure and task taxonomy

ATE-Bench covers 20 tasks in three categories: 12 codebase Q&A tasks, 4 operate-and-profile tasks, and 4 new-feature tasks (Lai et al., 29 May 2026). The suite was deliberately chosen to be valid across all compared frameworks.

Category	Number of tasks	Core objective
Codebase Q&A	12	Answer questions whose answers are properties of the code
Operate and Profile	4	Set up, run, instrument, and profile the framework
New Feature	4	Integrate new research architectures end-to-end

The Codebase Q&A tasks require read-only exploration using Read, Grep, Glob, and Bash, with Edit/Write disabled (Lai et al., 29 May 2026). Each answer must cite exact file paths and line numbers, and success is determined by two independent human graders, with disagreements resolved by a third reader. The examples span distributed initialization, configuration propagation, data sharding, seed management, attention kernel dispatch, RoPE, SwiGLU, normalization placement, sequence/context parallelism, FSDP/DDP wrapping, global gradient clipping, and checkpointing (Lai et al., 29 May 2026). All 108 attempts—computed as 12 tasks × 3 frameworks × 3 attempts—were judged correct (Lai et al., 29 May 2026).

The Operate and Profile tasks treat the framework as an executable toolchain. The agent is allowed the full toolset, including edit and shell operations, and may invoke in-repo skills if present (Lai et al., 29 May 2026). The four tasks are: Getting Started, Train and Evaluate, Collect Routing Trace, and Report Heavy Kernels. These tasks use a fixed training configuration: DeepSeek-V2-Lite base model; 8×H100 (1 node), PP=4, EP=2, DP=1, seq len 2048, global batch 1024, BF16. Public HuggingFace checkpoints are used to ensure steady, load-balanced MoE routing (Lai et al., 29 May 2026).

The New Feature tasks require end-to-end integration of published research architectures against reference implementations, thereby exercising the edit–run–debug loop (Lai et al., 29 May 2026). The four target architectures are Differential Transformer (Diff), Dynamic Mixture of Experts (DynMoE), Mixture of Block Attention (MoBA), and MoE++. Each task imposes three architecture-specific rule checks, and success further requires that cross-entropy decreases over 64 steps and remains finite. All 36 attempts—4 tasks × 3 frameworks × 3 attempts—satisfied both correctness axes (Lai et al., 29 May 2026).

This structure suggests that ATE-Bench is intended to approximate recurrent engineering work rather than synthetic puzzle-solving. Its task mix spans static comprehension, operational orchestration, performance instrumentation, and architectural extension.

3. Metrics, formalization, and aggregation

ATE-Bench evaluates each framework $f$ on a suite of tasks $t$, with each task attempted three times using a fixed agent and harness, indexed by $s \in \{1,2,3\}$ (Lai et al., 29 May 2026). The benchmark defines the following per-attempt metrics:

• Session duration (minutes):

$D_{t,f,s} = t_{\mathrm{end}} - t_{\mathrm{start}}$

• Active GPU time (minutes):

$G_{t,f,s}$, the GPU-active minutes accrued by framework workloads launched during the episode

• Agent turns (count):

$A_{t,f,s}$, the number of tool actions executed by the agent

• Per-turn context (tokens):

$C_{t,f,s}$, the tokens consumed as input context per agent action, aggregated by the harness into a single per-attempt value

• Output tokens (tokens):

$O_{t,f,s}$, the total assistant tokens generated over the episode; thinking tokens not exposed by the API are excluded (Lai et al., 29 May 2026)

Task success is represented by the binary indicator

$$S_{t,f,s} = 1 \text{ if the task’s correctness checks pass; } 0 \text{ otherwise.}$$

In the reported experiments, all tasks across all attempts and frameworks succeeded, so success rate is 1.0 (Lai et al., 29 May 2026).

Aggregation is by median across the three attempts for each task/framework pair:

• $\tilde{D}_{t,f} = \mathrm{median}_s \; D_{t,f,s}$
• $t$0
• $t$1
• $t$2
• $t$3 (Lai et al., 29 May 2026)

The benchmark description is explicit that there is no single-scalar “ATE score” or composite index, and that no confidence intervals or statistical tests are reported (Lai et al., 29 May 2026). This is methodologically significant: ATE-Bench does not assume that wall-clock time, GPU minutes, tool actions, and token expenditure are reducible to a common unit without loss of meaning.

A plausible implication is that the benchmark is designed to surface trade-offs that would be obscured by composite scoring. For example, a framework could reduce GPU time while increasing agent turns, or reduce output tokens while leaving session duration unchanged.

4. Agent environment, protocol, and validation logic

ATE-Bench uses Claude Code (Opus 4.7) at xhigh effort setting as the fixed agent (Lai et al., 29 May 2026). The standard coding-agent primitives are Read, Grep, Glob, Bash, and, for modification tasks, Edit/Write/NotebookEdit. For Q&A tasks, only the read-only tools are enabled (Lai et al., 29 May 2026). An episode is one end-to-end attempt of a task, and a turn is a single agent tool action such as a read, grep, edit, or shell command (Lai et al., 29 May 2026).

The benchmark harness logs turns, input “Per-Turn Context” tokens, and output tokens (Lai et al., 29 May 2026). For PithTrain specifically, a set of in-repo, task-specific skills—such as validate-correctness and capture-nsys-profile—is available, and the agent may discover and invoke these skills on its own during an episode (Lai et al., 29 May 2026).

The framework comparison uses the following pinned targets: PithTrain v0.1.2, Megatron-LM commit 3bec9aa, and TorchTitan commit d84e83d (Lai et al., 29 May 2026). DeepSpeed is excluded from the operate-and-profile and new-feature suite because, at pinned commit 44c51e3, it does not support PP combined with EP for MoE training in the required configurations (Lai et al., 29 May 2026). The operate/profile/new-feature tasks run on 1×8 H100 node (8 GPUs) (Lai et al., 29 May 2026).

Reproducibility controls include pinned framework commits, fixed training configurations per task category, read-only scripts such as train.sh and evaluate.sh in the harness, pre-staged datasets and model checkpoints, and validation via programmatic checks and/or human adjudication (Lai et al., 29 May 2026). For example:

• Getting Started succeeds when the framework reaches step 5 with finite loss, as verified by harness log parsing.
• Train and Evaluate succeeds when the framework trains 25 steps from random init, exports to HuggingFace, evaluates zero-shot HellaSwag via vLLM + lm-evaluation-harness, and yields a finite score.
• Collect Routing Trace succeeds when four step-XXXXXX.npz files contain expert IDs and gating weights with correct shapes and values, and the weights are non-negative and sum to 1 over selected top-k.
• Report Heavy Kernels succeeds when a CSV with schema (kernel_name,total_time_ms,instances,mean_time_ms) is produced and validated against the raw .nsys-rep in Nsight Systems (Lai et al., 29 May 2026).

This protocol makes ATE-Bench closer to a controlled systems evaluation than to an open-ended agent benchmark. The fixed agent, fixed hardware, fixed scripts, and pinned commits are intended to isolate the effect of framework design on agent cost.

5. Reported results and diagnosed sources of cost

The reported results are medians across three attempts and indicate that PithTrain reduces agent effort substantially relative to the production frameworks on the benchmark suite (Lai et al., 29 May 2026).

For Codebase Q&A, the paper reports up to 67% fewer Agent Turns versus Megatron-LM (Lai et al., 29 May 2026). In the Q1 Device Mesh example, Agent Turns are 15 for PithTrain, 33 for Megatron-LM, and 18 for TorchTitan; Per-Turn Context is 33.4K vs. 45.8K vs. 44.3K; and Output Tokens are 4.1K vs. 9.0K vs. 7.1K (Lai et al., 29 May 2026). The accompanying interpretation is that compactness and the absence of implicit indirection shrink the search space and reduce per-turn context (Lai et al., 29 May 2026).

For Operate and Profile, the benchmark reports reductions in both agent effort and GPU time (Lai et al., 29 May 2026). In Getting Started, Agent Turns are 26 (PithTrain) vs. 88 (Megatron) vs. 54 (Titan); Active GPU time is 3.1 vs. 5.4 vs. 5.2 minutes; Session Duration is 6.6 vs. 40.5 vs. 11.4 minutes; and Output Tokens are 5.8K vs. 26.9K vs. 15.8K (Lai et al., 29 May 2026). In Train and Evaluate, Agent Turns are 92 vs. 163 vs. 212, Active GPU is 22.7 vs. 36.0 vs. 36.3 minutes, and Output Tokens are 34.2K vs. 52.9K vs. 97.8K (Lai et al., 29 May 2026). In Collect Routing Trace, Agent Turns are 58 vs. 112 vs. 103, with Active GPU 2.8 vs. 5.5 vs. 10.4 minutes (Lai et al., 29 May 2026). In Report Heavy Kernels, Agent Turns are 42 vs. 60 vs. 40, Active GPU is 3.6 vs. 12.1 vs. 6.7 minutes, and Output Tokens are 16.0K vs. 23.9K vs. 22.5K (Lai et al., 29 May 2026). The aggregate deltas described in the text are up to 70% fewer Agent Turns than Megatron-LM and 57% fewer than TorchTitan, with Output Tokens up to 78% and 65% lower, respectively (Lai et al., 29 May 2026).

For New Feature tasks, the paper emphasizes reduced reruns and lower Active GPU Time (Lai et al., 29 May 2026). In Differential Transformer, Agent Turns are 47 (PithTrain) vs. 125 (Megatron) vs. 58 (Titan) and Active GPU is 27.6 vs. 33.7 vs. 40.3 minutes (Lai et al., 29 May 2026). In DynMoE, the corresponding values are 76 vs. 199 vs. 197 and 41.9 vs. 49.1 vs. 94.4 minutes (Lai et al., 29 May 2026). For MoBA, they are 57 vs. 134 vs. 91 and 27.7 vs. 49.5 vs. 77.9 minutes (Lai et al., 29 May 2026). For MoE++, they are 90 vs. 145 vs. 87 and 39.9 vs. 58.7 vs. 51.9 minutes (Lai et al., 29 May 2026). The paper summarizes this as up to 62% fewer Agent Turns and 64% less Active GPU Time on the hardest new-feature tasks relative to production frameworks (Lai et al., 29 May 2026).

The benchmark also reports diagnosed sources of cost in the production frameworks. For Megatron-LM, these include implicit-indirection collisions between manual and auto-derived CLI flags, C++/TE grouped-GEMM segfaults leading to speculative toggling, and opaque errors away from the edit site (Lai et al., 29 May 2026). For TorchTitan, repeated OOM debugging drove additional edit–run cycles (Lai et al., 29 May 2026). By contrast, PithTrain failures are described as surfacing as Python tracebacks within the just-edited file, keeping fixes localized (Lai et al., 29 May 2026).

6. Interpretation, relation to PithTrain, and limits of generalization

ATE-Bench is presented as a mechanism for quantifying how PithTrain’s four agent-native design principles affect agent efficiency: compact codebase (~11K LoC), Python-native stack, no implicit indirection, and agent skills (Lai et al., 29 May 2026). The benchmark’s ablations support that relation. When two in-repo skills are turned off while keeping the code and agent fixed, the benchmark records substantial increases in agent-side overhead. For validate-correctness, Agent Turns rise from 34 to 114 and Output Tokens from 11.3K to 30.2K, while Active GPU remains near parity (22.5 min vs. 20.8 min) because the GPU work is fixed by the workflow (Lai et al., 29 May 2026). For capture-nsys-profile, Agent Turns rise from 36 to 75 and Output Tokens from 14.5K to 25.3K, again with near-parity GPU time (5.5 min vs. 5.6 min) (Lai et al., 29 May 2026). The reported interpretation is that skills encode procedural knowledge and verifiable endpoints, allowing the agent to execute a fixed plan instead of iteratively deriving one (Lai et al., 29 May 2026).

A case study on integrating MoBA further attributes the observed reductions to framework structure. Editing dominates across frameworks, but the median editing tokens are much lower on PithTrain: 4.7K versus 13.1K for Megatron-LM and 22.2K for TorchTitan (Lai et al., 29 May 2026). Exploration tokens and per-turn context are much higher on Megatron-LM due to code size and implicit indirection, while TorchTitan’s spikes align with OOM debugging (Lai et al., 29 May 2026).

The benchmark also positions itself against existing agent evaluations. Standard agent benchmarks such as SWE-bench, HumanEval, MLE-bench, FlashInfer-Bench, and KernelBench hold codebases fixed and vary the agent, using correctness as the primary metric (Lai et al., 29 May 2026). ATE-Bench inverts that axis by holding the agent and tasks fixed, varying the framework, and measuring cost and time. The stated novelty claim is that it is the first benchmark to measure developer/agent operational cost for ML training frameworks across understanding, operating, and extending tasks rather than only throughput or correctness (Lai et al., 29 May 2026).

Its limitations are explicitly stated. The tasks were selected to be valid across all three compared frameworks and therefore exclude scenarios where production frameworks’ implicit indirection could be beneficial, such as propagating a shared change across many models (Lai et al., 29 May 2026). Results depend on a single fixed agent—Claude Code Opus 4.7—and may vary with agent capability, prompting, or tools (Lai et al., 29 May 2026). The suite centers on MoE training frameworks, so its representativeness for non-MoE or non-training systems remains to be established (Lai et al., 29 May 2026). Finally, the benchmark reports separate effort dimensions and medians across three attempts, with no composite ATE metric, confidence intervals, or hypothesis tests (Lai et al., 29 May 2026).

A nomenclatural ambiguity should be distinguished carefully. In separate work on assistive technology, “@Bench (Assistive Technology Benchmark)” has been described as equivalent in spirit to an Assistive Technology Evaluation benchmark (ATE-Bench) because it evaluates AT-relevant capabilities across tasks selected directly by people with visual impairments, but the authors do not use the specific term “ATE-Bench” for that benchmark (Jiang et al., 2024). In the ML systems literature, by contrast, ATE-Bench specifically names the benchmark for agent-task efficiency introduced with PithTrain (Lai et al., 29 May 2026).