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.
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.
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.
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.
Over the final 20 minutes:
| Group | Engines | Output tok/s per engine | vs default |
|---|---|---|---|
| Default head | 2 | 2,846.81 | — |
| Train, don’t push | 3 | 2,782.40 | −2.26% |
| Train and push | 3 | 3,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:
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.
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.
Over the last 10 minutes of each phase:
| Phase | Acceptance vs default | Output 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 phase | Hindi | Python | Japanese |
|---|---|---|---|
| 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.