bench

Quick training-step benchmark

Execute run_accumulated_step from curry_train.benchmark over a small number of optimizer steps and report wall-clock, tokens/sec, peak memory, and final loss. The intent is smoke testing, not throughput tuning — see stage4-capacity-sweep and stage4-optuna-integration for proper benchmarking.

When to invoke

• User asks: "smoke test the runtime", "time a step", "is the model running at all", "how many tokens/sec", "verify the accumulation works".
• After scaffolding a new experiment via the new-experiment skill, to confirm the freshly scaffolded model produces a non-NaN loss in one step.

Inputs

• --model=<name> (optional): a registered model name. Defaults to whatever the project's configs/default.yaml selects.
• --steps=N (optional, default 5): number of optimizer steps to run.
• --accum=K (optional, default 1): gradient-accumulation microbatches per optimizer step. Should match GradientAccumulation.steps in the config.

Procedure

1. Resolve runtime and model. Use curry_train.models.create_runtime(<model>) and runtime.build_model() to get a ModelHandle.

2. Build the gradient-accumulation primitive.

   from curry_train.primitives import GradientAccumulation
   accum = GradientAccumulation(steps=K)
   

3. Loop N optimizer steps. Each step calls run_accumulated_step(runtime, handle, microbatches, accum, ...) with synthetic or fixture microbatches (the model package should expose a dummy_batch() helper for this purpose).

4. Measure. Record per-step:

   • wall time (time.perf_counter)
   • tokens processed (AccumulatedStepResult.tokens)
   • loss
   • peak GPU memory (torch.cuda.max_memory_allocated() if CUDA available)
   • grad-norm (OptimizerStepResult.grad_norm)

5. Report. Print a small markdown table. Add a single warning line if loss is NaN, grad-norm is 0, or step time variance > 50 % (suggests warmup or first-step compile cost not yet absorbed).

Output template

# bench — model=<name>, steps=N, accum=K

| step | loss   | grad-norm | tokens | step-time (s) | peak-mem (MiB) |
|------|--------|-----------|--------|---------------|----------------|
| 1    |   ?    |     ?     |   ?    |       ?       |       ?        |
| 2    |   ?    |     ?     |   ?    |       ?       |       ?        |
| ...

throughput: ? tokens/sec (mean of steps 2..N)
verdict: <ok | NaN-loss | grad-norm-zero | step-time-unstable>

Boundaries

• Do not use real datasets. Bench is a smoke test; use dummy_batch() from the model package or generate torch.randn of the expected shape.
• Do not run for thousands of steps. If the user wants a real throughput measurement, point them to stage4-capacity-sweep and Optuna-driven sweeps.
• Do not save a checkpoint. Bench is ephemeral.

Failure modes

• No registered model: list registered models from curry_train.models.list_models() and ask the user which one.
• Runtime build fails: surface the traceback verbatim, then suggest the diagnose skill (the user can ask to debug the failure in natural language).
• dummy_batch missing: tell the user to add it to their model package; refer them to stage1-scaffolder for the template.