[micahtyong]

Summary

Spindle serves two kinds of training requests. forward_backward runs the model forward and backward and accumulates gradients. forward (Tinker's forward_only=True) runs the model forward only, to score tokens without training. Both return the same payload: a log probability for every token and every candidate label in the batch.

A request travels through three hops: the SDK sends it to the frontend, the frontend hands it to the engine, and the engine calls the trainer. The result comes back along the same path. In #42, forward_backward got a binary transport for this trip, with results encoded as protobuf. forward never did. Its request went to the trainer as JSON, its result came back as JSON through the engine and the frontend, and only then did the frontend re-encode it as protobuf for the SDK. For a batch of 128 examples with 5 labels per position that result is about a million floats. As JSON, where every float is written out as text and parsed twice, that is about 20 MB per call. As protobuf it is 4 MB.

This change gives forward the same binary transport and protobuf results that forward_backward has. Two smaller things come with it:

• The engine batches queued forwards the way it already batches forward_backward. The SDK splits a large forward into chunks; those chunks now run as one trainer call instead of several. Forwards never share a call with forward_backward, because the two need different GPU work.
• If a shared forward call fails, the engine reruns each request on its own. One client's bad input then fails only that client's request.

E2E results

Same setup as #152: a 4B LoRA classifier on 4×H100; each step runs two forward calls, one forward_backward, then optim_step, one after another; 128 examples of about 1,631 tokens with 5 candidate labels per position; 16 steady-state steps. Step-0 log probabilities are bit-identical before and after.

Before WandB / After WandB

Median per step	Before	After
Time outside the trainer	8.95 s	5.67 s
One forward: encode the result and send it trainer → engine	0.63 s	0.11 s
One forward: engine → frontend	0.97 s	0.03 s

🤖 Generated with Claude Code