SWARMS: Speculator Weight Adaptation via Remote Model Synchronization


A speculative decoder is only as good as its match to the traffic it serves. A draft head trained on general chat data is a poor guesser on Hindi, or on one customer’s code base, and every rejected draft token is wasted work. The usual fix is to collect traffic, train a new draft model offline and redeploy it. This post is about doing that continuously instead: every serving engine trains its own draft head on the requests it is serving right now, and a small central service merges what all the engines learn.

It works because speculative decoding already produces the training signal for free. We take that signal off the GPU, train on the host CPUs, stream the new weights back into a spare buffer, and federate the engines with DiLoCo. The GPU never runs a backward pass. We call the system SWARMS: Speculator Weight Adaptation via Remote Model Synchronization.

This post is a small-scale test of the idea with a small model, Qwen3.5-4B and its built-in MTP draft head, on a single 8× L40S node. To show the training engine improving the speculator as traffic changes, we serve public Hindi, Python and Japanese prompts. The results:

  • Throughput (90 min, Hindi): engines that serve the weights they are learning overtake the default speculator after about 7 minutes and end +13.0% output tok/s per worker ahead. Running the whole training pipeline without serving its weights costs 2.3%.
  • Elasticity (90 min, Hindi): two engines are hard-killed, a fresh one joins and the aggregator is killed and restarted. The surviving engines see zero request errors and the fleet ends +13.6% in throughput.
  • Distribution shift (150 min, Hindi → Python → Japanese): the fleet follows Hindi and Japanese, but never gets back above the default on Python, and it forgets each old language as it learns the new one. That last part is an open problem.

1. The free training signal

A speculative decoding step has two halves. A small draft model, here the MTP (multi-token prediction) head that ships with Qwen3.5, proposes the next three tokens. The large target model then runs one forward pass over those proposals and checks them: it keeps the longest prefix it agrees with and replaces the first token it disagrees with by its own choice.

To check the drafts, the target computes its own next token at every drafted position. So after every step we already know the target’s hidden states and the tokens the target actually chose. That is exactly a labelled training example for the draft head: given the target’s hidden state at position t and the next token, predict the target’s next three tokens.

One decode step with three drafted tokens. The target model's verify pass over four positions produces its own next token at each one; two drafts are accepted and the third is replaced. The hidden states and the target's chosen tokens are left over as a training example for the draft head, with no extra target computation.

The loss is plain cross-entropy at each of the three draft depths against the tokens the target produced, over the full vocabulary. It needs no labels, no extra forward passes of the 4B target, and the data is on-policy: it comes from exactly the traffic the engine is serving. With greedy decoding, the tokens the target emits are the tokens it would emit with no speculation at all, so the draft head can only change how fast the output arrives, not what it is.

The trainable part is small: the MTP module is one decoder layer plus its input projection, 120.6M parameters. The embedding and LM head it shares with the target stay frozen, so training cannot change the target’s outputs.

2. Keeping training off the serving path

A serving GPU is busy, so the design rule was that serving should never wait for training. Each engine gets a CPU trainer process on 8 host cores, and data and weights move between them asynchronously.

A worker is one serving engine on one GPU plus CPU cores on the same host. The GPU holds the frozen target, two banks of draft-head weights (one serving, one spare being filled) and capture buffers. Captured hidden states and tokens go to pinned host memory on their own CUDA stream. The CPU trainer holds FP32 master weights and AdamW state and writes new BF16 weights into the spare GPU bank in 4 MiB chunks. Every 100 local steps the worker exchanges a 482 MB weight delta with the DiLoCo aggregator.

Capture. On every 6th verify step, the engine copies the hidden states, tokens and recent draft-head K/V for up to 32 requests into preallocated buffers, then to pinned host memory on a separate CUDA stream (about 33 MB per capture). The input slots are double-buffered, so a slot is only reused after its copy has finished.

Train on the CPU. The trainer keeps FP32 master weights and AdamW state. One local step (forward and backward over the three draft depths, then AdamW at LR 2.47e-6) takes about 1.5 s on 8 cores, and the trainer gets through roughly 45 training positions per second. The engine serves about 2,800 tokens per second, so the head learns from under 2% of them, and that rate, not the amount of traffic, is the real limit on how fast it learns.

Swap, don’t copy. The GPU holds two copies of the draft head, each with its own captured CUDA graphs. One serves while the trainer writes new BF16 weights into the other in 4 MiB chunks. When the copy completes, the scheduler sees it with a non-blocking event query and flips the bank pointer at the next step boundary. No forward pass ever sees half-written weights, and nothing on the serving stream waits for a copy.

Each local step’s weights go live this way, about every 1.5 s. The second bank and its graphs cost about 1.5 GB of GPU memory.

3. Local learning plus a central trainer

Each engine on its own only sees its own share of the traffic. To let every engine benefit from what all of them learn, we treat each one as a DiLoCo worker:

  • Each worker takes H = 100 local AdamW steps from the current global weights.
  • It then uploads its delta, the global weights minus its local weights, to an aggregator. The aggregator is a CPU service and could sit anywhere.
  • Once a quorum of workers has reported, the aggregator averages their deltas and treats the average as a gradient for an outer Nesterov step (learning rate 0.7, momentum 0.9), then sends the new global weights back.
  • Each worker resets its local weights to the new global but keeps its AdamW moments.

Sequence diagram. In the local loop, about every 1.5 s: the GPU sends captured hidden states and tokens to the CPU trainer, which runs forward and backward (about 1.1 s) and an AdamW step (178 ms), then writes new BF16 weights into the spare bank; the GPU swaps banks at the next step boundary. Every 100 local steps: the trainer uploads its delta (482 MB), keeps training while the aggregator waits for a quorum, averages the deltas and takes an outer Nesterov step, then sends 482 MB of new global weights to every worker, which resets to them but keeps its AdamW moments.

Nothing here is synchronous with serving. A worker keeps training and serving while its upload and the outer step are in flight; the local progress it made in that window is dropped when it resets to the new global. A round took about 3.4 minutes in these runs.

The exchange is one weight-sized payload each way per round: 482 MB up and 482 MB down in FP32, with no compression yet: about 12.5 GB each way per worker over a 90-minute run of 26 rounds. We also tried averaging after every update instead (plain federated averaging). In a 3-worker, 60-minute comparison it reached a similar Hindi gain (+9.4% against +9.1% for DiLoCo) and forgot less Python, but moved 5.5× more data (280 GB against 51 GB). DiLoCo’s appeal is that the aggregator can be far away.

4. Results

All three experiments start from the stock Qwen3.5-4B MTP head, with fresh optimizer state and the settings above. Each engine draws prompts from its own disjoint shard of the traffic pool. Throughput is output tokens per second per engine at 192 concurrent requests. Acceptance length is output tokens per target verify step, including the target’s own token, so the default head’s 2.23 on Hindi means each verify step yields 2.23 tokens on average.

For held-out results we replay saved weights on prompts that never enter training (900 Hindi, 200 Python and 900 Japanese) and report the change in acceptance length against the untrained head.

Experiment 1: what it costs and what it buys

All 8 GPUs serve Hindi traffic for 90 minutes. Two engines serve the default head with training switched off. Three run the full training and federation pipeline but never serve the result (train, don’t push): their GPU still receives every weight update into the spare bank, but never swaps it in. Three serve their trained weights (push). All six trainers are one DiLoCo fleet.

Output tokens per second per worker over 90 minutes, five-minute rolling. The default speculator stays near 2,850 tok/s. Engines that train but don't push run about 2.3% below it. Engines that push their trained weights overtake the default at 6.4 minutes and reach about 3,220 tok/s. Red ticks along the bottom mark the 26 new global weights.

Over the final 20 minutes:

GroupEnginesOutput tok/s per enginevs default
Default head22,846.81—
Train, don’t push32,782.40−2.26%
Train and push33,217.55+13.02%

The 2.3% is the price of the pipeline: the capture copies, the spare-bank fills and the CPU cores the trainer takes from the host. The +13.0% is the net gain after paying it. The pushing engines overtake the default 6.4 minutes in (7.0 minutes after correcting each engine by its own pre-run calibration) and stay ahead for the rest of the run.

The live curve mixes local and global weights, so we also replayed every one of the 26 global versions on held-out prompts:

Held-out acceptance of each of the 26 global versions against the untrained speculator. Hindi rises steadily to +12.4%. Python, which never appears in the training traffic, falls slowly to −2.2%. Error bars are ±1.96 paired standard errors.

Hindi acceptance rises from 2.23 to 2.51 (+12.4%). It is still rising at 90 minutes. Python, which the fleet never sees, slips by 2.2%. The replay runs with deterministic FP32 numerics, and in it the target’s output was byte-identical to the untrained run on all 34,100 trained-weight comparisons, confirming that training changes only speed.

Experiment 2: workers come and go

A fleet that is also a serving fleet has to cope with engines disappearing. Here three engines train and a fourth serves the default head as a control, on Hindi, for 90 minutes, on four of the GPUs while Experiment 3 ran on the other four. The aggregator closes a round once max(2, ceil(0.6 × live workers)) workers have reported; workers hold a 120-second lease renewed by heartbeats. During the run:

  • 20 min: one worker is hard-killed. Its lease expires and the other two carry on.
  • 42 min: a fresh worker joins with a new identity, an unused shard and no optimizer state. It downloads the latest global, serves it, and starts training.
  • 60 min: the aggregator is killed and restarted from its persisted state. It is back at the same version in 3.4 s.
  • 75 min: a second worker is hard-killed mid-round.

Top: live acceptance length, five-minute rolling. The mean over live trainers keeps rising from 2.26 to 2.56 through every event, while the frozen default stays near 2.21. The fresh worker joins slightly below the others and catches up within minutes. Bottom: held-out acceptance of each global version: Hindi reaches +9.9% and Python falls to −3.7%.

The surviving engines recorded zero request errors, were never paused, and no global version was skipped. The curve keeps rising straight through every event. Over the last 10 minutes, each live trainer served +13.6% more output tokens per second than the control. The held-out Hindi gain of the final global is +9.9%, with Python down 3.7%.

Joining is not free for the new engine: its own decode had two gaps of about 5 seconds around its ready time, which is cold start. The engines that were already running saw no multi-second gaps.

This was the second attempt. The first one stopped at 60.8 minutes, just after the aggregator restart. The aggregator had come back fine, but while its DNS name was briefly unresolvable, the workers’ clients reported a transient retry status, and our experiment monitor read that as a fatal trainer failure and tore the fleet down. We made the terminal failure an explicit state, gave the client exponential backoff on transient network errors (0.25 s up to 4 s, while serving and local training carry on), tested a real kill-and-restart, then reran the whole experiment from scratch.

Experiment 3: following a shift

Three trainers and a control serve Hindi for 45 minutes, then Python for 45, then Japanese for 60. All engines switch together, and nothing is reset at the boundaries. Since the default head’s acceptance itself jumps between languages, the useful measure is the trainers’ acceptance divided by the control’s on the same traffic.

Top: live acceptance of the trainers divided by the frozen default on the same traffic. It climbs to about 1.13 on Hindi, drops to about 0.99 when traffic switches to Python and stays there, then dips to 0.98 on Japanese and crosses 1 after 8.9 minutes, ending near 1.07. Bottom: held-out acceptance every 5 minutes. Hindi gains 8.5% during the Hindi phase and gives it back during Python and Japanese; Python ends at −8.3%; Japanese ends at +6.1%.

Over the last 10 minutes of each phase:

PhaseAcceptance vs defaultOutput tok/s vs default
Hindi+13.17%+9.95%
Python−0.82%−2.10%
Japanese+6.66%+3.64%

On Japanese, the fleet starts below the default (it arrives carrying Hindi and Python weights), crosses it 8.9 minutes in, and keeps climbing. On Python it never gets back above the default within the 45 minutes. The default head is already strong on code, with 3.2 tokens per step, which leaves little to gain, and the fleet starts the phase as a Hindi specialist.

The held-out replays show the cost. Relative to the untrained head:

After phaseHindiPythonJapanese
Hindi (45 min)+8.5%−1.6%−3.0%
Python (90 min)+0.5%−0.3%−4.6%
Japanese (150 min)−2.2%−8.3%+6.1%

Each language’s gain lasts only while it is the traffic. By the end, Hindi is below where it started and Python is 8.3% worse. With a 120M-parameter draft head and a learning signal that only ever reflects the current traffic, the fleet adapts by overwriting. For a fleet serving one slowly changing workload that may be fine. For mixed or quickly changing traffic it is the main open problem: the obvious next steps are replaying a small sample of past traffic, keeping per-domain heads, or anchoring to the original weights, and none of them is tested yet.