Skip to content

Hybrid MoE checkpoint support

The specialized Core adapter supports config-driven KDA/full-attention layouts, latent MoE, per-head Q/K gains and scalable softmax. The 12.5B profile has 16 blocks (14 KDA and two full-attention), hidden width 1024, latent width 512, 512 experts and top-16 routing. These are checkpoint properties, not constants in the adapter.

Use the sources in the runtime lock. The adapter lives in MILES; OLMo-core supplies model construction and conversion, and olmo-sglang supplies the serving implementation. See serving modes for compatible arithmetic and limits.

Conversion checks

scripts/miles/validate_moe_conversion.py accepts explicit --native and --hf paths and performs two exhaustive checks:

  1. Load the saved native architecture and model weights, including flattened optimizer-backed master weights where applicable. Stream the HF mapping and compare every tensor to the reference export.
  2. Build the Core adapter from the HF configuration, import its tensors, then stream them back and compare every tensor again.

Both passes reject missing, duplicate, extra, wrong-shaped and nonfinite weights. Values must match exactly after the recorded export dtype conversion. The native pass may trim only embedding/LM-head vocabulary padding, using the saved sizes. Reports retain tensor/config hashes, dtype casts and packing metadata. This check loads no optimizer moments and does not establish generation parity.

Tiny runtime tests cover multiple widths and expert counts, distinct non-unit attention gains/scales and the real distributed-checkpoint reader. For a full checkpoint, use a CPU-only unallocated job with the required WEKA mounts: omit context.minRuntime, try Saturn first, and use Jupiter after checking scheduler events and canceling an unschedulable attempt. Follow the launch procedure and commit source changes before submission.

Attention and publication contracts

Conversion preserves the two-dimensional Q/K gain tensors and learned self_attn.ssmax_scale, with strict tensor coverage and native dense-MLP packing detection. Scalable softmax multiplies Q by log(position + 1) * ssmax_scale in the reference dtype and rounding order; the usual 1/sqrt(head_dim) factor still applies. Positions must be local to each request, including chunked prefill and decode. Publication includes these parameters and copies into existing storage so captured graphs retain valid pointers.

Current limits

The pinned Olmo KDA kernel supports SGLang's extra_buffer cache path using tracked-state snapshots written directly from FP32 accumulators. Casting BF16 intermediate states to FP32 does not satisfy this contract. Rebuild the runtime image when updating the adapter pin; older images reject tracked prefill batches. The named olmo35-b300-v1 preset enables caching for its supported geometry; olmo35-h100-v1 disables it. Named presets apply whole: an explicit inference.radix_cache value conflicts, even when equal. To change cache mode, select inference.profile = "generic" and configure serving settings explicitly. Generic and colocated setups default to disabled caching; recurrent state storage is still required for active requests.

Synchronous barrier publication with EP4 training and TP1 fused-rounding serving is supported for this model geometry. Fully async refresh supports generation across weight updates, including training on mixed-version responses; qualify the intended hardware, placement and workload using the throughput considerations. Short checks do not establish long-run learning equivalence or save/resume support.

Matching greedy tokens does not imply matching probabilities against an HF reference. Use Core's actual scoring policy for rollout comparisons and retain both token and probability checks.

Before extending the profile, test cached generation, mixed request lengths, chunked prefill, graph replay with changing positions and publication of changed gains/scales. Larger tensor parallelism and different model geometries need their own checks. Dense Olmo 3 uses a separate FSDP adapter; results do not transfer between the two paths.

See the support matrix for the other supported paths.