laion/snowball-64k-eagle3-draft-base

EAGLE-3 base draft for GrugMoe, the starting point for your own adaptation

A one-layer draft trained from scratch on the target's own outputs to general SFT-corpus prompts, with no agent-workload specialisation. It requires a compatible GrugMoe target and the pinned EAGLE-3 serving fork. It is not a standalone language model.

This is the checkpoint the R2E-Gym adapted draft (https://huggingface.co/laion/snowball-64k-eagle3-draft-r2egym) was warm-started from. Publish-time intent: if your workload is not R2E-Gym, start from this one and adapt it on your own target's rollouts with the toolkit at https://github.com/lukedhlee/grugmoe-eagle3; adapting from the R2E-Gym draft instead carries that workload's specialisation with it. Served as-is it already gives a measured 1.50× on held-out R2E-Gym agent turns.

This package omits the target embedding. The serving engine supplies the active target's embedding. The draft weights, output head and vocabulary maps are unchanged from the trained checkpoint. Optimizer state, training logs, target weights and training examples are not included.

Compatibility and use

The measured target was the internal 64k-context Snowball S3 Nemotron-Terminal step-1888 checkpoint: 26 layers, hidden size 2560, vocabulary 128256. This draft has one Llama-style layer, a 32,000-token draft vocabulary, a 2,048-token attention window, and auxiliary hidden-state slots [2,13,23].

Use the vLLM fork at 00d81ae11404ab68afef4142cdc883f7ba77e9a6 (branch lukedhlee/grugmoe-eagle3) from https://github.com/lukedhlee/vllm. Pass the downloaded draft directory to your target server:

{"method":"eagle3","model":"/path/to/downloaded/draft","num_speculative_tokens":3}

For SkyRL, this dictionary belongs in generator.engine_init_kwargs.speculative_config. Every inference worker needs the patched serving environment and access to the draft directory. The config deliberately has no saved target filesystem path. Select the target explicitly in your serving command.

The embedding-free export has passed the exporter's tensor-content checks and the pinned serving config loader. It has not been separately served on a GPU; the adapted draft's export, produced by the same exporter from a checkpoint of identical architecture, was served on one GH200 node on 2026-09-11 with zero errors. Do not read the measurements below as measurements of this export. Another target checkpoint needs its own acceptance and speed benchmark. The public 262k Grug model has not been tested, and equal tensor dimensions do not establish tokenizer or positional-encoding compatibility.

Training and measurements

Trained from scratch on 12,099 of the target's own responses to SFT-corpus prompts (about 159 million input tokens and 43 million supervised tokens per epoch), with loss only on tokens the target sampled, for two epochs in 25 minutes on one GH200 node. R2E-Gym was held out of training entirely. Training used prefixes up to 32k tokens and three training-time unroll steps.

On four GH200 GPUs, with 16 concurrent streams and three speculative tokens, on held-out 64k R2E-Gym agent turns:

• 139/140 decode tokens per second per stream in two same-node repeats, versus 93 without a draft, about 1.50×.
• Median turn latency 11.5 s down to 6 to 7 s.
• Accepted tokens per verification step, bonus token included: 2.55 on prompts over 40k, 2.15 under 16k; 2.19/2.20 median in a later same-node screen.
• Zero request errors; sampled log-probabilities and top-1 mass matched the no-draft server.

After warm-starting this draft on 2,177 of the target's R2E-Gym rollouts, the adapted draft served 15 % and 12 % faster than this one in two same-node repeats, which is the gain adaptation on your own workload can be expected to fight for.

Serving was tested up to 64k context. Context beyond 64k and other target families are unvalidated.

Adapt to your workload

Generate your exact target's own outputs on your prompts, extract its hidden states, and build a task-disjoint dataset. Keep the included d2t-32000.npy and t2d-32000.npy when warm-starting this draft. The toolkit's toolkit.py train --draft <this directory> --d2t ... --t2d ... prepares a separate working copy that names your target stub and restores its embedding before training. Never silently replace the vocabulary mapping of learned draft weights. Select the result by served speed on a held-out replay, not by validation loss alone.

The frozen target embedding is intentionally absent; this package still contains the draft output head and all learned draft tensors.