Module Ocannl.Train

Sets the tensor's value as materialized (device-resident, inspectable on demand via the context), and returns the tensor's forward code with a label-derived comment.

A scalar non-differentiable accumulator for grad_update's ?accum_loss: zero-initialized at allocation and materialized. Read it with Context.get_values (which awaits the device) and reset it with Context.set_values ctx t.value [| 0. |] — e.g. once per epoch.

The subset of loss.params that loss's backprop actually trains: the parameters whose gradient the backprop code writes (Asgns.collect_written). t.params deliberately answers a broader question — which parameter leaves the forward graph reads, hence what init_params must initialize and Persistence must save — so a parameter detached behind Operation.stop_gradient (a frozen backbone) stays in loss.params while its gradient is neither zeroed nor written. Stepping it anyway would apply weight decay to weights the user froze (gh-ocannl-673); the optimizer-side helpers (sgd_update, grad_l2_norm, clip_by_global_norm, grad_checksum, zero_params_grads) therefore derive this set instead of trusting loss.params. Empty when loss is not differentiable.

Replaces the zeroing of the given gradient nodes with Noop inside a zero_grads computation (each per-tensor zeroing is a Fetch of zeros — see Tensor.fetch_zeros). Used by the gradient-accumulation variant of grad_update, which must keep zeroing the intermediate gradients every micro-step (their backprop contributions are plain =+ accumulations relying on a same-routine reset) while the parameter gradients accumulate across micro-steps.

The match is by Fetch constructor and tnode identity, so a change to how zeroing is emitted (a different constructor, or a parameter gradient dropping out of the tree) would silently keep zeroing parameter gradients — corrupting the accumulation. A gradient that the backprop writes is accumulated into, so it is zeroed in the tree: we raise when such a grads member had nothing removed rather than let the drift through. Pass only the gradients backprop reaches — a parameter detached from the loss (behind Operation.stop_gradient, say) stays in loss.params while its gradient is neither zeroed nor accumulated.

Returns the tensor's forward, zeroing gradients, and backprop code wrapped with label-derived comments. Sets the tensor's value as materialized. If setup_for_parallel is true (false by default), sets the trained parameters' gradients as "non-local" (on-device). When accum_loss is given (see loss_accumulator), the update also accumulates the loss value into it (accum_loss =+ loss): training loops can then read the loss sum once per epoch instead of once per step — on GPU backends a per-step Context.get_values awaits the whole device, serializing the stream, while steps that only accumulate on device queue up and overlap with host-side scheduling. When loss_scale is given (see Mixed_prec.Loss_scaler), the backprop is seeded with the scale's value instead of 1 (loss.grad =: loss_scale), so all gradients come out multiplied by the scale — unscale them before the optimizer update (the grad_unscale argument of sgd_update).

When accum_steps is given (gh-ocannl-465), the returned code is a micro-step of gradient accumulation, llm.c-style: parameter gradients are NOT zeroed here (they are materialized so they persist across runs, and each micro-step's backprop =+-accumulates into them — run zero_params_grads at the start of each accumulation cycle instead), while intermediate gradients are still zeroed every micro-step; and the backprop seed is pre-scaled by 1/accum_steps (folded with loss_scale if both are given), so after accum_steps runs the parameter gradients hold the mean of the micro-batch gradients — matching a single batch accum_steps times larger under a mean-reduced loss. Run the optimizer step once per cycle, after the last micro-step.

Code zeroing the trained parameters' gradients (trainable_params, or ?params), as a standalone computation: the gradient-accumulation counterpart of grad_update ~accum_steps (which deliberately does not zero them). Compile it as its own routine and run it at the start of each accumulation cycle — before the first micro-step, llm.c-style ("we're about to += accumulate into them").

A scalar checksum over the trained parameters' gradients (trainable_params, or ?params) of loss: returns the flag tensor and the code that resets it to 0 and accumulates the sum of every gradient cell into it. The sum is non-finite if and only if some gradient cell is non-finite (a finite sum cannot arise from non-finite cells: same-sign infinities stay infinite, opposite-sign infinities and NaNs produce NaN; a spurious overflow of large finite gradients only triggers a benign extra backoff). Sequence it after grad_update in the same routine, read the flag with Context.get_values and gate the optimizer step on Float.is_finite — the dynamic loss scaling recipe (Mixed_prec.step) does exactly this.

See: https://github.com/tinygrad/tinygrad/blob/master/tinygrad/nn/optim.py

When grad_unscale is given (the reciprocal of grad_update's loss_scale), the gradient is first multiplied in place by it, so the optimizer math below — including the momentum buffer — sees unscaled gradients, and so does any later reader of p.grad (e.g. gradient clipping).

When grad_scale is given (a broadcastable scalar, e.g. field-grad_clipping.grad_scale of clip_by_global_norm, gh-ocannl-465), the gradient is multiplied by it as read into the update — the gradient buffer itself is left untouched (llm.c folds its clipping scale into the optimizer kernel the same way): later readers, the next accumulation cycle, and logged gradient norms all see the unclipped values, and no extra per-parameter sweep is emitted. The scale applies to the gradient only, before weight decay's p term joins the delta.

When update_gate is given (a broadcastable scalar holding 1 to apply the step and 0 to skip it, computed on device — see Mixed_prec.gated_scaled_update, gh-ocannl-492 task 5), every optimizer-state mutation is gated by Where selection: on a skipped step the parameter and the momentum buffer keep their previous values exactly. Selection, not multiplication — the skipped steps are the ones whose gradients hold inf/nan, and 0 * inf is nan.

Maps sgd_one over the parameters loss trains (trainable_params, or ?params): a parameter frozen behind Operation.stop_gradient takes no step — in particular no weight decay (gh-ocannl-673).

All and only bindings with associated ranges are iterated, with the binding's initial value lost. Bindings without ranges remain at their initial values, as do symbolic extents (gh-490): an extent is a size set once by the user, not an index to iterate.

f need not wait for the device: Context.run reads the bindings at the dispatch, so the next iteration may rebind them immediately whatever the backend schedules asynchronously.

Host-side learning-rate schedules: pure functions from the step number to a float, fed to the device via scheduled_learning_rate (or any host-written scalar). All schedules start with a linear warmup over warmup_steps steps (base_lr * (step+1) / warmup_steps; no warmup when 0) and decay toward base_lr *. final_frac at total_steps. Steps beyond total_steps clamp to the final value.

A device-resident, broadcastable scalar the host can overwrite with Context.set_values. Data-backed on purpose: a 1-element term_init would come out as a Constant fetch, re-fetched by every step's forward code, silently undoing Context.set_values. The bcast_if_1 axis basis (as in Tensor.number) lets the scalar broadcast into tensors of any shape.

A learning-rate scalar driven by a host-side schedule: returns the tensor (pass it as sgd_update's ~learning_rate) and a setter overwriting it with the schedule's value at step — call the setter once per step, before running the optimizer routine. The overwrite is a tiny host-to-device transfer, not a recompilation.

A scalar holding the global L2 norm over the gradients of the parameters loss trains (trainable_params, or ?params): returns the norm tensor (materialized and observable — read it with Context.get_values for logging or Outlier_detector feeding) and the code that computes it: per-parameter sum-of-squares einsum reductions (deterministic by construction — OCANNL emits no atomics) followed by a square root. Sequence it after grad_update in the same routine or a later one. Like grad_checksum, it must be called only after the model and the loss are fully constructed.

When grad_unscale is given (the reciprocal of a Mixed_prec.Loss_scaler's scale), the norm is multiplied by it after the square root, so the result is the true-magnitude norm even when backprop was seeded with a loss scale (the buffers themselves still hold scaled gradients at this point — sgd_one unscales them in place later).

Global-norm gradient clipping, llm.c-style (llmc/global_norm.cuh feeding grad_scale into the AdamW launch): the returned scale leaves gradient buffers untouched and multiplies the gradients as the optimizer reads them (sgd_update ~grad_scale).

Behavior at the edges of the finite range, deliberately (matching llm.c's grad_scale = grad_clip / grad_norm, which shares both properties): a nan norm fails the ordered comparison and selects scale 1 — clipping does not gate non-finite gradients, combine with grad_checksum or Mixed_prec.gated_scaled_update for inf/nan defense. An infinite norm — including the overflow of squaring a finite f32 gradient component above ~1.8e19 (the accumulator is f32 and narrow-storage gradients widen at load, so storage precision does not lower that threshold) — makes the ratio 0, suppressing the gradient term entirely: the step degrades to weight decay alone, strictly more conservative than rescaling an explosion of that magnitude to max_norm, and self-recovering since the buffers are untouched.

A sliding-window z-score outlier detector for host-observed scalars (per llm.c's llmc/outlier_detector.h): feed it the per-step loss and/or grad_l2_norm values and skip the optimizer step when the returned z-score exceeds a threshold. Pure host-side state; use one detector per monitored quantity.

Materializes every non-literal embedded tensor node of t (so its value is inspectable on demand via the context). Replaces the old every_non_literal_on_host now that there is no hosted memory mode (gh-ocannl-333).

Which placement arm tune_placements ships (gh-ocannl-638).

Measured_winner is the default and the only setting a normal run should use: both arms are timed and the faster one ships. The two forcing settings are measurement-only. They ship a chosen arm whatever the timings said, which is how a measurement gets executed values out of the artifact it profiles rather than out of whichever artifact happened to win — the gap benchmarks/report-gh612-hip.md had to state in its verdict, where three of four cells shipped arm B while every ratio was computed on arm A's never-executed routines.

Forcing changes what ships, not what is measured: both arms are still searched, so the A-vs-B comparison a report quotes stays available and the positional ?report contract is untouched. A forced arm that failed has no fallback — its failure propagates rather than the other arm shipping in its place, since the caller asked for that artifact and there is none.

Parses the tune_ship_arm spelling of a placement_arm. source names what is being parsed, for the error message.

Placement A/B autotuning: Autotune.tune on comp under the graph's current (default) placements — virtual intermediates plus the compiler's promotions — and again with every embedded node of loss materialized, keeping the measured winner (the arms' best_ms are min-of-N timings on the same device, so directly comparable). Under the default ship_arm = Measured_winner the result is by construction at least as fast as the better of the default and materialize-all placements, whichever the search would find; this generalizes the old "materialize everything before tuning" recipe instead of replacing one fixed placement policy with another. That guarantee is exactly what the other ship_arm settings give up (gh-ocannl-638): a forced arm ships whether or not it was the faster one, which is the point — they exist so a measurement can execute the artifact it profiles — so every "keeps the winner" and "at least as fast" statement here is about Measured_winner alone. Respecting the two-level memory-mode split (docs/proposals/context-scoped-memory-modes.md) — tnode-level memory_mode is declared, semantics-bearing intent, while placement decisions are context-level and functional — the B arm does not touch intent: it tunes from Context.decide_materialized siblings of ctx (and of timing_ctx), so the arms are hermetic and tune_placements leaves no trace on the graph or on the caller's contexts beyond the returned winner. See test/operations/materialize_after_compile.ml. report, when given, observes both arms' reports in order — arm A first, then arm B — so a consumer holding both reports can attribute every per-arm fact to the arm that produced it. What the reports do not determine is which arm SHIPPED: read on_ship for that. The two came apart in stages — a winning flip refinement ships a placement vector that is neither arm (gh-ocannl-555), and ship_arm (gh-ocannl-638) overrides the time comparison outright — so "the smaller best_ms shipped" is no longer a rule a consumer can apply, and applying it is exactly the misattribution on_ship exists to prevent. The reports are measurements of two searches; on_ship is the identity of the returned artifact. That separation is what makes "a Schedule.Tensorize was crowned in an arm that did not ship" reportable (gh-ocannl-546): best_tensorized on the other arm's report, with mma_best_ms against best_ms for the margin. The same conclusion is logged here under config autotune_log. Other arguments are forwarded to Autotune.tune; the same caveats apply (notably timing_ctx and non-idempotent routines — both arms share timing_ctx's device for their searches). name included (gh-ocannl-669): it names both arms' compiles and the flip refinement's decision-surface lowerings, and is what lets a comp carrying no Ir.Assignments.t.Block_comment — one Context.compile would name at the call site — be tuned here at all.

An arm that fails is a losing arm, not a failed run (gh-ocannl-550): a search that terminates on a fatal failure ranks at infinity, the other arm's completed winner ships and stays cached, and the failed arm's own report — Autotune.Search_died, carrying the failure — still reaches report in position, so the failure is recorded rather than downgraded to "that arm merely lost". A failed arm's best_ms is deliberately not shippable and not compared: Autotune.tune raised, so no routine was compiled from ctx for it. Only when every arm fails does tune_placements propagate, with the first failure's original backtrace — with two exceptions that are not the arm's to absorb and propagate at once: process-level failures (Out_of_memory, Sys.Break) and compiler invariant violations (Assert_failure, Stack_overflow), the same classes Ir.Schedule_outcome.classify_raw refuses to classify. A report callback's own exception likewise propagates rather than counting as an arm failure, and so does a failure that poisoned the lineage the arms share (Context.poisoned_failure): every timing run in the sibling would then refuse to execute, so the second arm can only burn a search proving it. Consumers attributing arms by arrival order should read Autotune.terminal_failure — or match the report's outcome — before best_ms, as benchmarks/runners/ocannl/bench_harness.ml does.

The in-position guarantee is Autotune.tune's reporting contract, and inherits its one carve-out: an argument-precondition violation (an incompatible timing_ctx) is detected before the call reaches any phase, so it reports nothing and propagates. Both arms are given the same contexts, so both raise it; there is no surviving arm to misattribute.

The arms differ in which candidates exist, not only in how they rank: a tensorized candidate is seeded only when the matmul site's operand and destination storage precisions resolve to a tile the backend advertises (Autotune.mma_tile_for_precisions), and placement decides which nodes the site reads. Under the mixed-precision recipe on a uniform-format backend (Metal's simdgroup matrices) that makes arm A tensorization-free: the reduced-precision cast twins are virtual there, so the site reads f32 masters into a reduced-precision destination — a mixed triple no tile matches — while materialize-all turns the twins into real reduced-precision nodes and the seeds fire. Materializing just the twins (Mixed_prec.Twin_materialized) reaches the same seeds at arm A's cost; see benchmarks/report-gh546-metal.md.

gh-555: the A/B is the coarse level of the hierarchical inlining search — inlining decided first, tiling/scheduling within each arm by the nested Autotune.tune. inline_flips (config tune_inline_flips, default 0) adds a greedy per-node refinement level: the default-policy arm's compile reports its searchable decision dimensions (Ir.Low_level.flip_candidates), and the candidates are tried one at a time from arm A's context — Materialize via Context.decide_materialized (walking toward arm B one node at a time), Inline via Context.decide_inline — each accepted flip becoming the base for the next, and the refined result shipping only if it beats the A/B winner. Every measured flip costs a full search like an arm, so the budget is explicit and defaults to zero. Flip searches report through flip_report, not report: the positional arm-A-then-arm-B contract of report is preserved regardless of the budget.

gh-514, the tuned placement-space search: the chain walks the surface in Autotune.placement_surface's ranking — family-unlocking (enablement) Materialize flips before cost, per the gh-558 lesson; config tune_flip_ordering=cost restores the legacy recompute-cost order as the evaluation baseline — weighed, under the default tune_flip_ordering=profitable (gh-ocannl-579), against what the two arms just MEASURED about the family that prior points at (Autotune.family_profit_of_reports over both arm reports: arm B is the all-materialized specialization the enablement set is derived from, so its best tensorized time against its best time prices the promotion for free). A family measured to lose here by more than tune_flip_profit_margin voids the prior and the surface ranks by cost — the prior models expressibility, and on gh-514's metal/f16 cell promoting a hopeless family took budget slots 1-2 and pushed the winning cheap flip out of a budget-5 chain. And, under config autotune_bound_pruning, fathoms a Materialize flip pre-search when the roofline floor of the chain's partial placement vector extended by it already meets the best measured time (admissible: the floor lower-bounds every completion, so the flip cannot win). Fathomed flips do not consume the budget, which counts measured flips.

gh-ocannl-638, ship_arm (config tune_ship_arm, default Measured_winner): ship a chosen placement_arm instead of the measured winner. It exists for measurement — a profile of arm A's kernels is evidence about arm A's routine, and until that routine is the one that ships, nothing ever executes it against a reference, so a value-changing regression inside it leaves every structural and timing figure plausible. Deliberately loud: a non-default setting announces itself on stderr regardless of autotune_log, both when it is resolved and at the decision it changes, because it is the optimizer's shipping path. Forcing also skips the flip refinement (which walks away from the chosen arm one node at a time, so its result is neither arm), and a forced arm that failed propagates its failure rather than falling back to the other arm.

on_ship is called exactly once with "A", "B" or "flip" on the path that returns a routine, and not at all when nothing ships. It is what a consumer should attribute the returned artifact by: deriving the shipped arm from the reports' best_ms is only valid while nothing can override the comparison, which ship_arm now can, and it never described a flip-refined result at all.

gh-ocannl-498 rematerialization: the configured device-memory budget for a compiled routine, or None when there is none (the default — under which the planning pass never runs and compilation is bit-for-bit what it was). The setting is a byte count with an optional K/M/G suffix (powers of 1024), the word minimize, or 0 / off / false / none.

gh-ocannl-498: plan comp's inlining decision vector against a device-memory budget and return the context to compile it from. budget overrides the config key memory_budget; with neither, this is the identity on ctx and returns None — the default-off path does not lower, score or decide anything. See Memory_budget.fit for what the planner does and what it requires (config buffer_aliasing).

Dumps comp as a .cd file in the build directory, for the ?output_cd_file argument of to_routine and run_once. caller names the calling function in the error raised when the global setting output_debug_files_in_build_directory is false.

The gh-ocannl-498 rematerialization seam shared by to_routine and run_once: plan comp's inlining decision vector against the budget, compile from the planned context, and only then let budget_report observe the plan.

The report fires AFTER the compile so it observes the plan that SHIPPED — a compile or link failure must not have announced one. It also keeps a callback from reaching the compile it is reporting on: the config gates the scoring depends on (buffer_aliasing) are re-read at each compile, so a callback that flipped one would make the routine use a layout other than the one just scored.

Compiles comp and returns the post-compile context together with the routine. budget, max_candidates and budget_report are the gh-ocannl-498 rematerialization seam, forwarded to fit_memory_budget: with no budget and no memory_budget config key nothing is planned and the compile is exactly what it was; budget_report, if given, observes the plan that shipped.

The post-compile context is returned rather than discarded (gh-ocannl-772), matching Context.compile and run_once: chain it into the next compile instead of reaching into routine.context for the same value. A caller that only wants the routine can snd this.

init_params initializes the parameters of t, via running their forward code or copying from the host as appropriate. If reinit_all is true, all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

run_once is a wrapper around init_params that additionally runs code of f t and returns the context. If skip_init is true (false by default), no initialization is performmed. If reinit_all is true (false by default), all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

If output_cd_file is true, the global setting output_debug_files_in_build_directory must be true, and the update code is output to a file before shape inference potentially crashes at init_params.

forward_once is a wrapper around run_once that runs the forward code of t.

update_once is a wrapper around run_once that runs the gradient update code of t: both forward and backprop.

printf is a wrapper around Tensor.print that assumes ~force:true, and by default sets ~with_code:false, ~with_grad:true, and ~style:`Default. It takes an explicit context and retrieves values on demand (gh-ocannl-333). If the tensor's value is not already materialized in ctx, it is recomputed via the for_print copy trick so real values are still shown.

printf_tree is a wrapper around Tensor.print_tree that assumes ~force:true, and by default sets ~with_value:true, ~with_grad:true, and ~depth:9. It takes an explicit context and retrieves values on demand (recomputing via for_print if not already materialized).