Plan: evaluate TrainCheck-assisted agents on training implementation and distributed compatibility #165

Goal

Evaluate whether TrainCheck helps coding agents implement meaningful changes to training systems while preserving the contracts with the rest of the training stack. Include actual small-scale distributed training from the first pilot. An optimizer-formula-only task is insufficient.

The primary question is: Does access to TrainCheck help an agent deliver a correct implementation across the required configurations, identify and repair regressions earlier, or reduce the total cost of doing so? Improved model quality or throughput counts only when the implementation satisfies the task's correctness and compatibility requirements.

This issue plans a small research pilot before investing in a full benchmark. It does not claim that TrainCheck already supports all of the proposed distributed observations.

Why change the task selection?

The previous reduced RSI pilot (#164) evaluated loss design and optional optimizer equations through a restricted tensor interpreter. The trainer, model, parameter application and distributed integration were outside the agent's editable surface. Across two pairs, the mean held-out accuracy difference was +0.652 percentage points, with opposite pair directions; neither treatment agent reported changing its design or selection because of the diagnostics. These results do not establish benefit or harm for implementation-level training tasks.

That pilot also used API-only tracing, without model-variable coverage. It cannot establish that parameter updates, optimizer-state relationships or distributed contracts were checked.

RSI itself is broader than hyperparameter tuning. Useful starting points include its public OPSD sample, which allows actual trainer/collator/loss changes, and MoE router and hybrid architecture tasks with module/state/export contracts. Reuse appropriate code and realistic requirements, but label reduced or adapted tasks separately from official RSI evaluations.

Candidate task families

Agents edit real Python/PyTorch components and their integration, with documented public requirements. Start by qualifying these three families; choose two for the first agent comparison based on reproducibility, integration depth and resource cost, before looking at agent outcomes.

Use small real models and short runs that still exercise these interactions. Include meaningful quality or memory/throughput objectives where appropriate, so agents must complete a useful feature rather than merely satisfy a collection of bug-specific assertions. Optimizer implementations may appear within these tasks, but are not the entire suite.

Two distinct evaluation tracks

1. Diagnostic calibration and bug repair. Start from independently verified healthy implementations and introduce realistic, separately documented faults: e.g. incorrect loss scaling across ranks, stale parameter references after wrapping/conversion, missing state at resume, or skipped synchronization. Use minimized historical regressions where reproducible. Establish that the fault causes the claimed failure using an independent oracle before measuring TrainCheck. Include healthy and legitimately changed controls. Keep calibration faults separate from held-out repair variants.
2. Natural implementation. Give fresh agents a correct baseline and a feature request. Do not inject bugs into this track. Compare agents with and without TrainCheck while they implement the feature; record errors they introduce and whether diagnostics cause a useful revision. This is the primary evidence for the proposed use case.

Report the tracks separately. Faults caught in calibration do not prove that the tool improves agent decisions. Do not select the task suite solely from faults TrainCheck can already detect; publish unsupported contracts and the coverage denominator.

Execution matrix and cost controls

• Development: local/labmach/Ring18, isolated environments; CPU single-process and two-process Gloo checks for logic that they actually exercise.
• Required pilot: a real 2-GPU, single-node configuration, with DDP and a separately qualified FSDP configuration for applicable tasks. Confirm available hardware, framework versions and allocations first. Use the UCSD NAIRR allocation when needed. CPU multiprocess success is not evidence of CUDA/NCCL, AMP or GPU sharding compatibility.
• Optional extension: selected 4-GPU cases or a second parallelization scheme after the two-GPU pipeline works. Avoid a full Cartesian product. Multi-node execution is not required for the first pilot.
• Keep a one-device reference where useful, but grade on the declared distributed configurations. Fix global sample/token accounting, seeds and numerical tolerances for each comparison. Do not assert every FSDP state should be equal on every rank; validate ownership and reconstructed logical state instead.
• Pilot sizing proposal: 2 qualified tasks x 3 independent matched agent pairs = 12 agent sessions. This is early descriptive evidence, not a powered estimate. Set per-session wall/token limits, candidate-run counts and total CPU/GPU-hour caps after calibration and freeze them before agent trials. Include tracing, inference, checking, failed runs and queue time separately in cost reports.
• Run heavy work only on the designated lab resources or allocated cluster compute nodes. Use rank-specific artifacts, bounded timeouts and cleanup of all owned worker processes on failure.

TrainCheck preparation and validity requirements

• Build a capability matrix for API events, model parameters, gradients, optimizer state, tensor representations and rank/step identity. Distinguish supported observations from missing instrumentation or missing relation/checker support. Do not assume the existing model tracker observes optimizer state or supports FSDP/DTensor.
• Verify plain versus instrumented behavior on healthy runs, including outputs, state, RNG, exceptions and checkpoint compatibility. Collect traces per rank with stable logical parameter identifiers. A cross-rank rule needs an explicit aggregation/checking path; independent per-rank checks are not a substitute. Hangs, missing rank traces and tracer crashes are explicit failures, not zero violations.
• Learn reference invariants only from independently validated executions. Include valid configurations/implementations where needed. Declare which contracts survive each intended modification: a new algorithm may legitimately change values, shapes, dtypes and call sequences. Do not certify an unverified candidate by learning new rules from that same candidate.
• Freeze the reference set, applicability policy, diagnostics and grading before agent trials. Do not remove inconvenient alarms after seeing outcomes. Record false alarms, unresolved alarms, active/dormant rules, checks attempted/completed, skips, errors and unsupported contracts. Coverage changes require a new study revision.
• Fix or explicitly account for the prior pilot's synthetic-identity/selective-metadata issues before using affected rules as evidence. Preserve full training context in short probes (including schedule horizon); a diagnostic probe should not silently change the task.

Paired agent protocol and metrics

Both arms receive identical task requirements, source, public tests, ordinary logs/profilers and development feedback, with the same model/settings and matched budgets. Treatment additionally receives the TrainCheck workflow and diagnostic access. Count its compute against the declared budget and also report it separately. Freeze all artifacts and final selections before independent grading.

Publicly document the required compatibility configurations; withhold test instances and compositions, not surprise requirements. Keep reference implementations, held-out faults and grading code outside agent-readable/editable paths. Reset each run and isolate agents from other runs and operator diagnoses.

Primary outcome: feature completion with all required contract checks passing, including the distributed cases. Also report per-contract outcomes, crash/hang versus silent regression, time/compute to a valid solution, residual defects, task quality/performance, and total diagnostic cost. Verify correctness through independent reference/differential/metamorphic checks; "zero TrainCheck violations" is never the correctness oracle.

Record the evidence chain when claiming a diagnostic helped: observation -> agent hypothesis/change -> independent test improvement. Report cases with no defect, no applicable rule, no tool use, or ignored/misleading diagnostics rather than imputing a benefit. Separate tool coverage, diagnostic quality and agent response when interpreting negative results.

Milestones and decision gates

• M0: Qualify tasks and resources. Pin sources/licenses, two initial task specifications, editable boundaries, supported parallelization versions, hardware availability and resource caps. Reproduce the healthy implementations and independent oracles. Record rejected candidates and reasons.
• M1: Establish diagnostic capability. Validate instrumentation transparency on actual two-GPU runs, map coverage to contracts, and run healthy/fault controls. Missing distributed/tensor/state support becomes a focused engineering issue, not a successful check. Only run a TrainCheck-benefit comparison where some task-relevant diagnostic capability is demonstrated; retain the broader coverage audit.
• M2: Freeze the agent experiment. Checkpoint task packets, diagnostic reference artifacts, grading, seeds, budgets and the 12-session assignment. Finish qualification of both selected tasks before starting the paired comparison.
• M3: Run and independently grade. Execute isolated paired agents, archive traces and changes, verify final artifacts on the declared distributed matrix, and audit causal-use claims.
• M4: Decide whether to scale. Report complete outcomes and costs, including failures. Expand only if we observe actionable diagnostic evidence and a credible mechanism for agent benefit. A single successful example validates a mechanism, not a general success-rate improvement; an inconclusive pilot should lead to a specific gap analysis rather than a larger uninformative run.

Checkpoint implementation milestones with focused signed commits and PRs, including baseline/candidate revisions, reproducible commands, correctness results and measured costs. Keep large traces, datasets and checkpoints outside Git and provide artifact manifests. This issue itself is the planning deliverable; no new training jobs have been launched for this proposal.