1 Introduction↩︎

Modern AI systems are increasingly powered by Mixture-of-Experts (MoE) language models such as DeepSeek-V3, Qwen3, Kimi-K2, and GPT-OSS [1]–[4], whose training depends on systems that scale across distributed clusters. Production frameworks [5]–[8] have built end-to-end MoE training stacks over years of engineering effort, pairing layered Python designs with extra compiled extensions to deliver broad model coverage, peak throughput, and multi-platform support needed for diverse training workloads. However, evolving these stacks for new model architectures and system optimizations demands in-depth expertise and substantial engineering effort.

AI coding agents [9]–[12] have begun to automate parts of this work and could in principle accelerate training-system development. Most current practice applies agents to existing frameworks. But the design choices that helped human engineers cost agents differently. 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. Today’s throughput-only evaluations leave this cost unmeasured, and existing frameworks were not designed with it in mind.

This raises a design question: can we redesign an MoE training framework that optimizes for agent-task efficiency? 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. We answer this question with PithTrain, an end-to-end MoE training framework designed agent-native from the start. PithTrain is built on four design principles. First, we favor code compactness over coverage: PithTrain focuses on a compact MoE training stack rather than the broad model and feature coverage of production frameworks, while remaining straightforward for agents to extend with new features. Second, we use minimal, Python-native components covering the key layers of MoE training: operators, training engine, and applications. Third, we use direct calls and avoid implicit indirection in module composition, so that what runs at a given call site can be identified by static reading. Finally, we ship agent skills [13] for recurring training-framework tasks.

Beyond building the framework, we systematically evaluate how framework design affects agent-task efficiency on real training-system tasks. Existing AI-coding benchmarks [14]–[16] vary the agent on a fixed codebase to score agent capability, focusing on general software-engineering tasks such as issue resolution. 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. 1 summarizes PithTrain’s overall design.

This paper makes three contributions:

PithTrain delivers dual efficiency: strong training throughput together with high agent-task efficiency. It matches the throughput of production frameworks across a range of MoE models and training settings on NVIDIA H100 and B200 GPUs. On ATE-Bench, 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. A qualitative case study illustrates how framework design shapes agent behavior.

2 Related Work↩︎

Mature production frameworks have driven large-scale MoE training. Megatron-LM [5], [6] established the pipeline-parallel Transformer recipe most subsequent frameworks build on, and DeepSpeed [7], [8] introduced ZeRO sharded-optimizer techniques. Both rely on layered designs with plugin systems, registry-based indirection, and compiled extension chains, delivering broad model coverage and peak throughput on production hardware. TorchTitan [17] 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. PithTrain takes a different point on this tradeoff curve with agent-native design as a first-class goal, achieving higher agent-task efficiency along Agent Turns, Output Tokens, etc., while matching the training throughput of production frameworks.

AI coding agents combine a tool-using language model with a small, stable vocabulary of primitives (read, edit, shell, search). Recent work designs better agent flows and harnesses for software engineering on top of these primitives [18]–[20]. PithTrain takes a complementary direction: rather than improving the agent flow, we design the software framework so that existing flows work better, lowering agent cost without changing the agent.

Existing agent benchmarks score capability on fixed codebases and tasks: SWE-bench [14] on GitHub-issue resolution, MLE-bench [21] on Kaggle-style ML engineering, and HumanEval [15] on function-level code generation. Closer to ML systems engineering, FlashInfer-Bench [22] and KernelBench [23] target inference operators and GPU kernels respectively. The agent is the variable; aggregate task correctness is the metric. ATE-Bench inverts this: holding the agent and task fixed, we vary the training framework so that differences in agent cost and task outcome isolate framework design. This axis is complementary to existing capability benchmarks.

3 The PithTrain System↩︎

3.1 Agent-Native Design Principles↩︎

This subsection introduces the four agent-native design principles that guide PithTrain’s system design, and contrasts how production training frameworks align with each principle today. 1 summarizes this comparison; production frameworks adopt different subsets of the four principles, reflecting different design priorities.

Compact codebase. Production frameworks such as Megatron-LM and DeepSpeed offer broad coverage of models, training features, and hardware platforms, accumulated over years of engineering effort, with core codebases exceeding 160K lines. A larger codebase inflates the cost of locating relevant code, tracking cross-file dependencies, and verifying a change is complete. A compact codebase reduces this cost; with frontier coding agents operating at context windows of 200K to 1M tokens, a sufficiently compact codebase can also fit in a single context pass. We treat compactness as a constraint on growth: PithTrain may grow over time, but additions should respect the four principles.

Python-native codebase. Python is the dominant language in modern ML. A pure-Python framework lets an agent navigate the full framework in a single language, surfaces readable Python tracebacks instead of opaque native errors, and eliminates the compiled-extension rebuild cycle. Megatron-LM composes its core transformer layers from out-of-tree TransformerEngine [24] modules, and DeepSpeed bundles extensive in-tree extensions. These deliver peak kernel performance and vendor-tuned numerics, but push an agent across language boundaries and force a rebuild on change.

No implicit indirection. Production frameworks compose many model variants from a shared layer skeleton via implicit indirection (a stored callable, plugin registry, or string-keyed resolution). This pattern enables code reuse across models, while making what runs at a given call site harder to identify by local reading. 2 shows an instance of model construction in a production framework: TransformerLayer resolves its submodules from a runtime spec in a separate file. A flat code structure trades cross-model reuse for local readability, reducing the effort an agent spends building an end-to-end understanding.

Task-specific agent skills. An agent skill [13] is a procedural playbook a coding agent loads on demand. Skills encode procedural knowledge an agent cannot recover from static reading alone, so it runs a verified procedure. Agent skills are a recent practice that existing training frameworks have not yet adopted at scale: 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.

3.2 System Architecture and Optimizations↩︎

The codebase is organized in three layers (application, engine, operator), as shown in 3. To realize a compact codebase, we identified the necessary components for a distributed MoE training framework, and PithTrain covers exactly those. Where production frameworks deliver broad out-of-the-box coverage of models, features, and hardware, PithTrain narrows scope to keep the codebase compact and reachable in a single context window. Table [tab:principle-map] maps each principle to where it is realized: we adopt a flat code structure with no plugin registries or runtime specs, with each MoE model living in a self-contained file under models/ rather than being composed through a shared layer skeleton. This favors local readability over cross-model code reuse.

PithTrain supports standard MoE training³ with pipeline parallelism (PP), data parallelism (DP) via FSDP [25], context parallelism (CP) [26], expert parallelism (EP) [27], FP8 training [28], and DCP checkpointing [29], on NVIDIA Hopper and Blackwell GPUs. Despite its compact, Python-native codebase, PithTrain aims for training throughput competitive with production frameworks by adopting standard MoE optimizations; these techniques are not novel, but they are central to PithTrain’s training throughput and worth calling out.

3.3 Agent Skills↩︎

A skill encodes the procedure for one recurring training-framework task. PithTrain ships a suite of skills covering several common ones (Table 4). Each skill is a self-contained folder with a top-level SKILL.md playbook, optionally additional markdown documents, and optionally helper scripts. Some skills are pure markdown, like add-new-model; others bundle scripts that offload deterministic work.

Each skill in PithTrain is designed around three properties: specific scope, explicit prerequisites, and verifiable success. [fig:skill-anatomy] illustrates these on validate-correctness. The description and triggers together encode the skill’s specific scope. The prerequisites section enumerates environment, data, and configuration assumptions, so missing state is caught before the skill begins to run. The procedure ends in a script call that returns a reproducible PASS/FAIL verdict rather than the agent’s self-assessment. Skills designed around these properties are not technically hard to author, and we expect other training frameworks to ship comparable coverage as the practice matures.

4 ATE-Bench: A Benchmark for Agent-Task Efficiency↩︎

Evaluating the agent-task efficiency of a training framework requires varying the framework while holding the agent and task fixed. This is the inverse of standard benchmarks for AI coding agents [14]–[16], which hold the codebase and task fixed and vary the agent to score capability. We introduce ATE-Bench, a benchmark with a fixed agent and curated task suite run across frameworks, so that differences in agent performance are attributable to framework design. The suite spans the kinds of work researchers typically perform on training frameworks, organized around three recurring patterns: understanding the framework without modifying it, operating it as a tool for instrumentation and profiling, and extending it with new functionality. Tasks are distributed across three categories:

5 illustrates the agent loop for each category, with agent involvement deepening from Q&A (read-only) to Operate and Profile (running, minor instrumentation) to New Feature (substantial modification, test-driven iteration). Full task descriptions and per-category correctness checks are in Appendix 8. Using ATE-Bench, we evaluate PithTrain and production frameworks, reporting five effort metrics: session duration, active GPU time, agent turns, per-turn context, and output tokens. Without a single-scalar metric for agent-task efficiency, we report each dimension independently.

Q&A questions are chosen to be valid across all three frameworks, excluding framework-specific behaviors. ATE-Bench does not cover tasks like cross-model propagation of a shared change, where production frameworks’ implicit indirection may lower agent effort; we leave these as future work.

5 Evaluation↩︎

We evaluate PithTrain on both axes of dual efficiency: training efficiency and agent-task efficiency. We validate training correctness against Megatron-LM on both pretraining loss curves and downstream accuracy in Appendix 7. This section is organized to answer the following questions:

5.1 Training Efficiency↩︎

3pt

We compare PithTrain (23db182), Megatron-LM [6] (3bec9aa) and TorchTitan [17] (d84e83d) on three representative MoE models (GPT-OSS-20B [3], Qwen3-30B-A3B [2], and DeepSeek-V2-Lite [33]) under matched parallelism (PP, DP, CP, EP), sequence length, and precision. Configurations span single-node and multi-node deployments on NVIDIA H100 and B200 GPUs. For Megatron-LM, we follow NVIDIA’s documented best practices⁴. DeepSpeed is excluded as it does not currently support PP combined with EP for MoE training⁵, so it cannot run any of the configurations in our suite.

To ensure the MoE router exhibits steady-state load-balanced routing across frameworks and thus comparable throughput, we load public model checkpoints rather than random initializations. Each experiment runs 25 steps, and we report the median step time over the last 10. We omit Model FLOPs Utilization (MFU) [34] because Tensor Core peak FLOPS differs between BF16 and FP8, making the metric ambiguous for mixed-precision steps. Training hyperparameters follow Appendix 7.

As 2 shows, PithTrain matches or exceeds Megatron-LM on 4 of 5 configurations, and stays within 1.4% of Megatron-LM on the fifth. This parity comes from optimizations such as DualPipeV’s compute–communication overlap and torch.compile(fullgraph=True). These results demonstrate that a compact, Python-native codebase can achieve competitive training throughput.

5.2 Agent-Task Efficiency↩︎

In this section, we evaluate agent-task efficiency across frameworks. We run ATE-Bench(§4) on Megatron-LM, TorchTitan, and PithTrain with Claude Code (Opus 4.7 at xhigh effort⁶) as the fixed agent. Each task runs three times and we report medians; hardware configuration, task descriptions, correctness criteria, and per-attempt values are described in Appendix 8 and 9. Opus 4.7 completed every task across all attempts and frameworks with no failure.

Q&A. Answering a question requires locating where a behavior lives in the codebase. We omit Session Duration (all tasks finish in under three minutes) and Active GPU Time (no training runs). All 12 questions are answered correctly across attempts and frameworks (grading details in Appendix 8.1.2). Across the 12 questions, the agent uses up to 67% fewer Agent Turns to reach the final answer on PithTrain than on Megatron-LM. A compact codebase and the absence of implicit indirection shrink the search space, lowering Per-Turn Context (3) accordingly.

Operate and Profile. Across all tasks (4), 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. PithTrain’s compact codebase explains these reductions. In addition, the agent invokes in-repo skills (§3.3) on its own when applicable; for example, the Report Heavy Kernels task triggers capture-nsys-profile.

3pt

New Feature. New-feature tasks exercise the test–debug cycle: edit, run training, read crash, edit again. Across all tasks ([tab:agent-efforts]), 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. On PithTrain, failures surface inside the file the agent just wrote with a readable Python traceback, and fixes stay in the same file. §5.4 provides a detailed case study on MoBA.

3pt

5.3 Ablation Study on Agent Skills↩︎

In this section, we isolate the effect of agent skills via ablation. They are a self-contained set of files shipped in the repository, so we can toggle them on and off against an otherwise fixed codebase. We pick two of PithTrain’s skills, validate-correctness and capture-nsys-profile, which mirror the natural follow-up after a system optimization: validate that training correctness is preserved, then capture an Nsight Systems profile to examine whether the optimization is effective.

We run this ablation on the wgrad delay [30] commit in PithTrain, repeating each task three times with skills and three times without. The codebase, agent, and harness are otherwise identical, and we report the same five effort metrics as in §5.2. When skills are disabled, we strip them from both the working tree and the git history, so the agent cannot recover the procedure from either. All twelve runs completed successfully, reporting the correct verdict or generating a valid Nsight Systems profile.

5 reports the results. 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. The four agent-side metrics, which capture the agent’s reasoning overhead in setup, launch, and interpretation, all drop substantially with skills enabled. Agent turns drop the most (70% and 52% respectively), suggesting that with the procedure encoded in the skill, the agent acts on a fixed plan rather than iteratively deriving one through repeated tool calls. These results demonstrate that task-specific in-repo skills, comprising the markdown playbook and any helper scripts, substantially reduce agent effort on the recurring training-system tasks they target.

3pt

5.4 Case Study: Integrating MoBA↩︎

To examine where per-framework differences in agent cost come from, we conduct a case study on integrating MoBA, decomposing the agent’s output tokens by action category⁷ (6 (a)) and tracing its per-turn context window (6 (b)). Editing dominates across all three frameworks (PithTrain 4.7K, Megatron-LM 13.1K, TorchTitan 22.2K). Megatron-LM also spends substantially more on Exploring (10.2K vs.3.8K for TorchTitan and 2.2K for PithTrain), and its per-turn context sits well above PithTrain’s. The agent reads the codebase to locate edits and interpret tracebacks, so a compact codebase with no implicit indirection lowers both Exploring and per-turn context. TorchTitan’s elevated Editing and ~2\(\times\) context spike in two of three runs have a different cause: the agent’s initial implementation runs out of memory (OOM), forcing repeated debug-edit cycles. This points to runtime properties like memory headroom as a factor independent of codebase structure.

Beyond TorchTitan’s memory failures, we examine the other failures the agent encountered. Across the three PithTrain runs, two complete without any failure; the third hits a tensor-stride mismatch in the agent’s custom attention kernel, fixed in the same file as the traceback. Megatron-LM’s three runs hit two distinct failures: two runs fail with a duplicate command-line flag registration that conflicts with one defined in framework code, and one run fails with a BF16 overflow in the agent’s code. Each fix on Megatron-LM spans multiple files. This contrast reflects PithTrain’s compactness and absence of implicit indirection, which keep each fix local to the agent’s edit.

6 Conclusion↩︎

We presented PithTrain, a compact, agent-native MoE training framework built on four design principles. PithTrain matches the throughput of production frameworks across a range of models, and on ATE-Bench, a coding agent achieves higher agent-task efficiency on PithTrain than on production frameworks. We hope PithTrain serves as a starting point for future agent-native training framework.

7 Training Correctness↩︎

This appendix validates the training correctness of PithTrain against Megatron-LM at matched configuration. We report two complementary measurements: pretraining loss trajectories (§7.1) and downstream task accuracy (§7.2).

7.1 Pretraining Loss↩︎

We pretrain Qwen3-30B-A3B with Megatron-LM and PithTrain under identical configuration. 7 reports the full training configuration alongside the two loss curves; the trajectories align across the full run, with Megatron-LM showing two transient spikes that recover within a few steps.

7.2 Downstream Accuracy↩︎

We evaluate downstream task accuracy across six standard benchmarks: OpenBookQA [35] and WinoGrande [36] in the 0-shot setting, and ARC-Challenge, ARC-Easy [37], HellaSwag [38], and PIQA [39] in the 10-shot setting. 8 plots accuracy for each task; Megatron-LM and PithTrain achieve comparable accuracy within statistical noise across all six benchmarks at every checkpoint.

8 Task Suite↩︎

This appendix expands the per-category task descriptions and the correctness checks used to validate each attempt. All operate-and-profile and new-feature tasks share a fixed training configuration: the base model is DeepSeek-V2-Lite [33] (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.

8.1 Q&A Tasks↩︎

Each Q&A task asks the agent to locate where a specific behavior lives in the framework codebase. The agent receives a single prompt consisting of the universal instruction below followed verbatim by the full query text for that task (Q1–Q12 in §8.1.1), and has read-only access to Read, Grep, Glob, and Bash; tools that modify the working tree (Edit, Write, NotebookEdit) are explicitly disabled. Correctness is validated by two independent human graders (§8.1.2).

Universal instruction (prepended to every query). Your task is to explore this training framework codebase and answer the question that follows. Every claim must be verified against the code on disk before you state it: use your CLI tools (Read, Grep, Glob, Bash) to locate the exact symbol, file, and line, then cite the path/to/file.py:LINE you actually read. Wrap your final response in <final_answer> tags, trace the execution path step by step, and give one citation per claim. If a feature is genuinely absent from this codebase, say so explicitly and cite negative evidence (the grep command and pattern that returned zero matches, or the file you inspected that does not contain it). A correct “absent, verified” answer is preferred over a fabricated one.

8.1.1 Task Descriptions↩︎

Q1: Process Groups / Device Mesh. Trace the sequence of function calls from the main training entry script down to the initialization of the PyTorch distributed process groups or DeviceMesh. Detail the exact file paths, function names, and line numbers where the world size and ranks for the parallel groups present in this codebase (any of DP, TP, PP, EP, CP that this repo actually supports) are assigned.

Q2: Configuration Propagation. Locate the exact file and line numbers where the user-provided configuration for hidden_size or dim is parsed. Trace how this specific variable propagates down to the instantiation of the first Transformer block.

Q3: Data Loading & Sharding. Trace the initialization of the dataset and dataloader. Identify the file, class, and line number where the global dataset is sliced or partitioned among the Data Parallel ranks to ensure each rank receives a unique, non-overlapping subset of data.

Q4: Distributed Seed Management. Find where the random seeds are set for data loading versus model initialization. Provide the exact file paths and line numbers, and trace how seeds are differentiated across different parallel groups (or, if they are deliberately shared, where and why).

Q5: Attention Kernel Dispatch. Locate the core attention module. Trace the logic that selects between an optimised kernel (FlashAttention, PyTorch SDPA flash backend, ring attention, etc.) and any fallback math implementation. Identify the exact file, class, line number, and configuration flag/condition controlling this dispatch. If the dispatch is delegated to the PyTorch scaled_dot_product_attention, say so and cite the call site.

Q6: RoPE Implementation. Find where the Rotary Positional Embedding (RoPE) is implemented. Locate the function that precomputes the frequency tensor (sine/cosine cache). Provide the file path, class/function name, and exact line numbers of the tensor operations where the rotary transform is applied to the Query and Key tensors.

Q7: SwiGLU / MLP Block. Locate the Feed-Forward Network (FFN) implementation. Identify the file and line number where the SwiGLU activation (or equivalent gated linear unit) is mathematically applied. Trace how the up-projection tensor is chunked or split before the activation function.

Q8: Normalization Placement. Find the implementation of the main Transformer layer. Provide the file and exact line numbers where the normalization (RMSNorm or LayerNorm) is applied before the attention module. Detail where the epsilon (eps) term is added for numerical stability.

Q9: Context / Sequence Parallelism. Find where sequence parallelism (or context parallelism) is handled. Identify the exact file and line numbers where the sequence dimension of the input tensor is sharded, gathered, or scattered across ranks during the forward and backward passes. If the codebase does not support CP/SP, say so and cite the negative evidence.

Q10: FSDP / DDP Wrapping. Locate the exact file and line numbers where the core model is wrapped with Fully Sharded Data Parallel (FSDP1 or FSDP2) or Distributed Data Parallel (DDP) wrappers, and identify the device-mesh axes used for sharding/replication.

Q11: Global Gradient Clipping. Trace the logic for global gradient norm clipping. Find the file, function, and exact line number where the global norm of the gradients is computed (including any reduction across distributed ranks) before the optimizer step.

Q12: Distributed Checkpoint Serialization. Locate the model checkpoint saving logic. Detail whether the system uses a rank-0 gather approach or distributed sharded saving (e.g., PyTorch DCP). Provide the exact file path and line numbers where the disk serialization occurs.

8.1.2 Correctness Checks↩︎

Each Q&A answer the agent returns cites the file paths, function names, and line numbers in the framework codebase that implement the behavior the question asks about. Two human graders independently trace these citations: they open each cited code location and verify that the code there does what the agent claims in the answer. An attempt is recorded as satisfied when both graders confirm every citation against the code on disk; disagreements are resolved by a third reader who looks only at the cited evidence. All 108 attempts (12 questions \(\times\) 3 frameworks \(\times\) 3 attempts) were judged satisfied by both graders.

8.2 Operate-and-profile Tasks↩︎

Each operate-and-profile task asks the agent to drive a real training-system workflow end-to-end, so success requires the agent to set up the environment, run the framework as intended, and produce the expected artifact. The four tasks span the workflows a researcher typically performs before any code change: getting the framework running, executing a full train-and-evaluate pipeline, instrumenting the model to capture behavior, and profiling the system to surface bottlenecks.

8.2.1 Task Descriptions↩︎

Getting Started. Set up the Python environment for the framework and run the provided 5-step smoke training script. Success means the script reaches step 5 with a finite loss. The agent must install all dependencies the script needs so that running bash train.sh as-is succeeds; train.sh itself documents best practices for training MoE models and is read-only. Pre-tokenized DCLM and the converted base-model checkpoint are pre-staged.

Train and Evaluate. Drive the full setup-train-export-evaluate pipeline for the base model. The agent trains from random initialization for 25 steps, exports the resulting checkpoint to HuggingFace format, and runs lm-evaluation-harness HellaSwag (zero-shot) on it via vLLM. The task tests pipeline correctness, not model quality: the HellaSwag score is expected to be near-random because 25 steps from random init produces a barely-trained model, and the agent must produce some score from a working pipeline. Initialization must be random; everything outside the fixed mesh and step count (LR, optimizer, scheduler, data preprocessing) is left to the agent.

Collect Routing Trace. Instrument training to dump the per-token MoE routing trace for the first 8M training tokens. The routing decision in each MoE layer, the top-\(k\) expert IDs and their gating weights, must be captured for every token in the global batch. Training resumes from the released HuggingFace checkpoint so the router is in its trained, load-balanced regime; routing decisions are model-intrinsic and valid from step 1, so no warmup is required. The output schema is one step-<step_id:08d>.npz file per step under workspace/<framework>/routing-traces/, each containing the expert-ID and gating-weight arrays.

Report Heavy Kernels. Profile a 7-step training run with Nsight Systems and identify the top 3 most expensive CUDA kernels by total GPU time, aggregated across all 8 ranks. Training resumes from the released HuggingFace checkpoint; the agent profiles only step 7 because steps 1–6 are warmup for cudagraph capture, NCCL handshake, and allocator priming. The output is a single CSV at workspace/<framework>/heavy-kernels/top-kernels.csv with header kernel_name,total_time_ms,instances,mean_time_ms and exactly three rows sorted by total_time_ms descending, plus the raw profile.nsys-rep so the result is reproducible.

8.2.2 Correctness Checks↩︎

Each operate-and-profile task verifies the artifact the agent produces, not the path the agent took to produce it. We use a mix of programmatic checks baked into the task harness and human inspection where the artifact is non-textual.

Getting Started. The task succeeds when running train.sh reaches step 5 with a finite loss. After the agent finishes installing dependencies, the harness re-runs the (read-only) training script and parses the resulting log; the artifact is accepted only if step 5 prints a finite loss value.

Train and Evaluate. The harness ships a read-only evaluate.sh script that consumes the pipeline output produced by the agent: the agent trains for 25 steps from random init, exports the resulting checkpoint to HuggingFace format, and evaluate.sh loads the export into vLLM and runs lm-evaluation-harness HellaSwag (zero-shot). An attempt is satisfied when evaluate.sh runs to completion and writes a finite HellaSwag score; the score is not required to clear any quality threshold, since 25 steps from random initialization is expected to yield near-random accuracy.

Collect Routing Trace. The harness verifies that four step-<step_id:08d>.npz files are present at workspace/<framework>/routing-traces/ and that each carries expert-ID and gating-weight arrays of the correct shape for every MoE layer over every token in the global batch. A programmatic check additionally confirms that the expert-ID values are within the valid range [0, num_experts) for the MoE configuration of the model, and that the gating weights are non-negative and sum to 1 over the top-\(k\) selected experts for every token.

Report Heavy Kernels. The agent submits both top-kernels.csv and the raw profile.nsys-rep. The harness checks the CSV programmatically against the prescribed schema (header, exactly three rows sorted by total time descending). The kernel names themselves are validated by a human reader who opens profile.nsys-rep in the Nsight Systems profiler GUI, reads the CUDA GPU Kernel Summary in the Stats System view, and confirms that the top three kernel names in the CSV produced by the agent match the top three entries shown by the profiler.

All 36 attempts (4 tasks \(\times\) 3 frameworks \(\times\) 3 attempts) produced an artifact accepted by both the harness and the human reader.

8.3 New-feature Tasks↩︎

Each new-feature task mimics the workflow of a researcher or ML engineer integrating a recently published modeling architecture into a training framework. The agent is given the same materials such an engineer would normally consult: the arXiv paper and the reference implementation. From these inputs, the agent must produce a training script that integrates the new feature into the base model and runs on DCLM [40] under the fixed configuration above.

The four tasks are chosen for both coverage and a fair starting point. Two revise the attention mechanism and two require changes at the FFN (MoE) layer, so the suite spans the architectural subsystems a training framework must accommodate. None of the four architectures had been integrated into any of the three frameworks we compare at the time of this study, so all frameworks start the task from the same point and the comparison reflects only how the design of each framework affects integration effort.

Correctness is validated on two axes. First, the cross-entropy loss must decrease across the 64-step run and remain finite (no explosion, no NaN), evidence that the modified training pipeline produces a learnable model. Second, the changes made by the agent must satisfy three task-specific rules, each validating one required component of the new feature (§8.3.2); the rules target the architectural elements that distinguish each new feature from the baseline framework, ensuring the agent has implemented the intended mechanism rather than producing a passing-loss reading from an unchanged model. An attempt is recorded as satisfied when both axes hold.

8.3.1 Task Descriptions↩︎

Diff [41]. Differential attention replaces a single softmax-attention map with the difference of two parallel softmax maps, \(\mathrm{Attn}(Q,K,V) = (\mathrm{softmax}(Q_1 K_1^\top) - \lambda\, \mathrm{softmax}(Q_2 K_2^\top))\,V\), where \(\lambda\) is a learned per-head scalar. The construction cancels common-mode attention noise, sharpening which tokens receive mass. Integrating it involves splitting the per-head \(Q\)/\(K\) projections in half and introducing \(\lambda\) as a trainable parameter with the published initialisation schedule.

DynMoE [42]. Standard MoE fixes the number of activated experts \(k\) per token. DynMoE replaces top-\(k\) selection with a per-expert sigmoid gate and a learned threshold, so the count of active experts varies per token. Integrating it involves replacing the top-\(k\) routing kernel of the framework and reformulating the load-balancing auxiliary loss for a variable-\(k\) regime.

MoBA [32]. MoBA partitions the key/value sequence into fixed-size blocks and routes each query to the top-\(k\) blocks selected by a learned gate, yielding sub-quadratic attention cost in the sequence length. Integrating it involves inserting block-level routing between the \(Q\)/\(K\) projection and the attention computation, composing with the existing attention backend in the framework, and preserving causal masking inside each selected block.

MoE++ [43]. MoE++ augments a standard MoE expert pool with three zero-computation expert types (zero, copy, and constant), allowing easy tokens to be routed past the feed-forward layer entirely. Integrating it involves extending the expert pool definition with the zero-computation variants, widening the router output to cover them, and adjusting the load-balancing loss to prevent the zero experts from absorbing all easy tokens.

8.3.2 Rule-based Correctness Checks↩︎

Each new-feature task decomposes into three required components, the architectural elements that distinguish the new feature from the baseline framework. The rules below name each component and describe how it should be implemented; they were fixed before inspecting any attempt. 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, and an attempt is recorded as satisfied when all three rules pass. 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. Across all 36 attempts (4 tasks \(\times\) 3 frameworks \(\times\) 3 attempts), every attempt satisfies all three rules of its task.

Diff.

• The attention forward path produces two separate softmax outputs that are combined as \(\mathrm{attn}_1 - \lambda \cdot \mathrm{attn}_2\) (literal subtraction with a learnable coefficient).

• A learnable parameter (\(\lambda\), e.g.named lambda or lambda_init) is registered as an nn.Parameter with shape compatible with per-head broadcasting.

• The \(Q\)/\(K\) projections are split into two halves, either via an output dimension of \(2 \cdot n_\mathrm{heads} \cdot d_\mathrm{head}\) followed by a split, or via two separate projection layers.

DynMoE.

• The router uses per-expert sigmoid gates (or sigmoid-gated activation), not pure softmax with top-\(k\) selection.

• A learnable threshold parameter for expert activation is registered as an nn.Parameter (often named tau, threshold, or gate_threshold).

• Active-expert selection is not a fixed top_k = N; the count depends on which experts pass the threshold (boolean mask or variable-length selection).

MoBA.

• The \(K\)/\(V\) sequence is partitioned into fixed-size blocks (a reshape or view into shape […,\(n_\mathrm{blocks}\),\(B\), …] or equivalent).

• Each query computes a per-block score (typically \(\mathrm{query} \cdot \mathrm{pooled\_block\_key}\)) followed by top-\(k\) block selection.

• Final attention is computed only over the selected blocks, with causal masking preserved at block boundaries (each query attends only to blocks at positions \(\leq\) its own).

MoE++.

• The expert pool includes at least one zero-computation expert type (zero, copy, or constant), identifiable by class name, string literal, or a specialised expert factory.

• The router emits logits over a pool that contains both regular FFN experts and zero-computation experts, so the router learns to dispatch tokens to the no-op variants. The zero-computation experts may be added on top of the existing pool (router output dim grows to \(n_{\mathrm{FFN}} + n_{\mathrm{zero}}\)) or carved out of an unchanged total expert count.

• The load-balancing auxiliary loss handles zero experts explicitly: they are excluded from the balance term, weighted differently, or balanced under a separate penalty.

9 Per-Task Results↩︎

This appendix reports per-attempt agent effort across the three task categories: Tables 6 and 7 cover the 12 Q&A tasks (split 6/6 across two pages), 8 covers the four operate-and-profile tasks, and 9 covers the four new-feature tasks. Each row is one independent attempt; per-task medians appear in the corresponding main-text tables in §5.2. Session Duration and Active GPU Time are reported in minutes. Lower is better on every metric.

4pt

4pt

4pt

4pt

References↩︎

---

1. Equal contribution↩︎

2. Equal advising↩︎

3. PithTrain also supports dense (non-MoE) models, with MoE-specific machinery (e.g., expert parallelism) skipped.↩︎

4. https://docs.nvidia.com/nemo/megatron-bridge/nightly/training/moe-optimization.html↩︎

5. Verified at commit 44c51e3: the Megatron-DeepSpeed integration script states “Currently we don’t support PP for MoE.”↩︎

6. https://platform.claude.com/docs/en/build-with-claude/effort↩︎

7. Thinking tokens are excluded as they are not surfaced by the API.↩︎