CU11 — Training, outside the graph

from somatize.torch import Trainer, parameters

g = Graph.somatize(Encoder().on("cuda:0") >> Head().on("cuda:0"))
t = Trainer(g, objective=cross_entropy,
            optimizer=torch.optim.Adam(parameters(g), lr=1e-3))

t.fit(data, epochs=10)    # the sugar
t.step(batch)             # the primitive

Status: closed. 110 tests in Rust, 136 in Python, and zero new lines in core/ — the first use case that does not touch the core.

The question: does the training loop go inside the graph?

No, and there are two independent reasons.

The first is in the node contract. forward(input, ctx) → Done | Await describes one step: it executes once per run, it has a budget of 64 turns, and run() has no partial recovery. A training run lasts an afternoon, mutates its own state, emits metrics continuously and fails in ways one wants to recover from. The graph operates at the scale of a forward; a training run operates at the scale of an afternoon.

The original tried it: its node trait carries fn fit(&self, x, y). The bill shows in its own tests — soma-worker, soma-compiler, soma-runtime and soma-agent all implement an empty fit just to be able to exist. It is the same tax CU6 removed on the filter/step axis.

The second is that a graph describes a network, and searching is a family of networks. That graph is precisely the artifact CU13 serializes and sends; one carrying five configurations inside would be lying about the architecture.

The three levels

level	what it is	scale	what it spreads
the graph	a network	one `forward`	slices of a forward: waves, `Placement`, and `Remote` in its day
`Trainer`	one training run	an afternoon	nothing; it repeats forwards
a study	N training runs	an experiment	whole runs

And the rule that holds it up: no level knows the one above exists. The graph does not know it is being trained; the trainer does not know there are other trainers. Composition between levels is composition of functions, not of graphs.

Level 3 has no type, and that is on purpose

study = {lr: Trainer(net(), ..., lr).fit(data) for lr in (1e-4, 1e-2)}
best = min(study, key=lambda lr: study[lr].loss)

A graph earns its keep when there are dependencies to declare. N independent training runs have none: they are a list. Modelling a list as a graph is paying a DAG’s price without using it.

The alternative was even designed — the N configurations as branches of a |, with a node that picks the best — and rejected. And so was its clever variant, “one graph, one plan, N catalogs”, which falls over something concrete: Catalog is Clone, but it clones Arcs. The N replicas would share the node objects, i.e. the weights, and the five configurations would train the same model, giving results that look good. Each replica has to be built — and once you build it, you no longer have one graph with N plans, you have N graphs. There is a test.

Decisions taken

1. The Trainer receives the graph; never g.fit(...). That way the same graph is trained three ways without touching it, and it remains the artifact that travels.

2. It lives in somatize.torch. Loss, backward() and optimizer are torch; writing it neutrally would ask for a Backend with a single implementor. The core does not learn what training is, and that is the sign the separation holds.

3. step(batch) is the primitive; fit(data, epochs) is sugar. It is what avoids the god-trainer path: early stopping, odd schedules, federated rounds and PBT are a while the user writes over step, not a growing list of options and callbacks. A federated round is for _ in range(k): t.step(batch).

4. Parameters are collected by duck typing, and a graph without them fails when building the Trainer. It asks for .parameters() and skips whoever lacks it — a lemmatizer does not train and does not stop being a node for it. Putting it in the contract would be the original’s fit all over again. And since duck typing fails quietly — a graph without parameters trains nothing and shows a flat loss — the empty list blows up at construction, just like CU10’s postcondition.

5. They come without repeats, by identity. Two nodes can share a module — tied weights between embedding and output — and then the same Parameter comes out twice.

6. The optimizer is built by the caller. No optimizer="adam", which would end up being a name registry. The only thing checked is that the optimizer and the graph share some parameter: sharing none has no innocent reading. Covering only a part is legitimate and passes — freezing the encoder and training the head is exactly that.

7. The data is an iterable of (input, target). A DataLoader is one. Deliberately rejected: data as a source node, because a node produces one value per execution; to be a stream it would have to remember where it is, and then two executions of the same graph stop giving the same thing.

8. The loss is a callable, not a node. It is not part of the network: it is swapped without touching the model and at inference time it does not exist. As a node, the graph’s output would be a scalar and the graph would only be good for training.

What the GPU test found

The target does not cross the graph. The input does, and each node moves it to its device because that is what a placed node does; the target goes straight to the loss, so nobody moved it. With the last layer on cuda:0, the output comes from there, the target is still on the cpu and torch stops the training with “expected all tensors to be on the same device”.

The only one that sees both sides fixes it. And the target is moved, not the output: bringing the output to the cpu would drag the backward pass back over the wire at every step.

Questionnaire

Python (soma-python/tests/test_trainer.py)

What did NOT go in

Checkpoints and resumption · callbacks and early stopping, which are a while over step · metrics beyond the loss · schedulers · gradient accumulation · the study as a type · and exporting or loading a model’s state, which is CU15’s question and was not needed here: training locally extracts no state.

The question it reshaped

The state question stops blocking, and it stops being the core’s. It is no longer is a node’s state a Value? but what does a training run export? at level 2 — answered in CU15, with the case in front of us.

The three levels also split what the original had joined. Spreading one graph across hosts has dependencies halfway through the forward and needs Plan::Remote; spreading whole training runs — HPO, federated, data parallel — is “execute this thing over there”, level 3, and needs none of it. The original had them in one enum, ModelParallel beside DataParallel, Federated and PopulationBased.