[RFC] Disaggregated streaming backend for speculative-decoding draft training

[RFC] Disaggregated streaming backend for speculative-decoding draft training

Ownership

This feature is assigned to @kashif (Kashif Rasul).

Summary

Introduce a streaming, producer/consumer feature pipeline for draft-model
(EAGLE / DFlash / DSpark) training so that target-model feature generation and
draft-model training run as independently-paced pools connected by a
lightweight reference queue plus a pluggable feature store, instead of the
current synchronous request/response remote backend.

The first producer backend keeps using AutoModel's own target-model forward
(the existing Eagle3TargetBackend path); no external inference engine is
required for the initial landing. This is a design proposal only. Implementation
is split into an independently-reviewable PR chain (see "Phased plan").

Motivation

Today the only train/inference disaggregation path is the EAGLE-3 remote backend
(nemo_automodel/components/speculative/eagle/remote/). It cleanly separates the
target GPUs from the draft GPUs over HTTP (control) + NCCL (tensors), with
prefetch and multiple remote URLs. But the trainer and the target server are
still bound in a tight request/response relationship:

• The trainer decides when to request a batch and blocks on the matching reply;
  each server serializes to one in-flight generate (_generate_lock,
  remote/server.py), and the client keeps one in-flight request per server to
  preserve NCCL recv ordering (train_eagle3.py prefetch loop, around line 775).
• There is no global queue, no producer/consumer backpressure, and no durable,
  resumable feature references. If a target server errors mid-step, it directly
  stalls the training step.
• Target and draft throughput must roughly match. Scaling the target side means
  changing remote_urls, and the trainer prefetch depth is capped to the number
  of remote servers (_resolve_prefetch_depth, train_eagle3.py).
• The remote backend exists for EAGLE-3 only. DFlash and DSpark are
  colocated-only (train_dflash.py, train_dspark.py load the target
  in-process), so their large-target story is worse.

For very large MoE targets (100B+ class, tens of GPUs, multi-node, long runs),
we want target inference and draft training to overlap continuously, scale
independently, absorb producer jitter, and tolerate a single producer failing
without killing the run. That requires turning the current "run one remote
forward on demand" into an asynchronous feature-production system.

Non-goal for small setups: 8B / 30B single-node runs where target and draft
throughput are close are well served by the existing colocated or remote
backend. This RFC targets the large-target, production-scale regime and must not
regress or complicate the small-setup path.

Current state (anchors this builds on)

• Producer interface: Eagle3TargetBackend(ABC) in
  components/speculative/eagle/backend.py with generate_batch(...),
  get_input_embeddings(), optional set_vocab_mapping(...), supports_async,
  generate_batch_async(...).
• Trainer data contract: Eagle3TargetBatch dataclass in
  components/speculative/eagle/target.py, carrying aux_hidden_states,
  input_ids, attention_mask, loss_mask, and exactly one supervision
  encoding (logits, or target_probs + position_mask), plus optional packing
  metadata. to_trainer_inputs() dispatches on the encoding.
• Backend selection: recipe_args.target_model_backend in
  {colocated, sglang, vllm, remote} (train_eagle3.py), with remote keys
  remote_urls / remote_url, remote_timeout, remote_max_retries,
  target_prefetch_depth.
• Existing transport primitives we can reuse: NCCL tensor transport
  (remote/transport.py), compact wire codec (remote/wire.py), HTTP control
  protocol (remote/protocol.py).
• Per-algorithm batch schemas already exist: Eagle3TargetBatch (EAGLE-3),
  DFlashTargetBatch (dflash/target.py), DSparkTargetBatch
  (dspark/target.py).

Proposed design

Add a small, algorithm-agnostic data plane that sits between any target-feature
producer and the draft trainer. Control-plane records carry no tensors; tensors
live in a store and are referenced by key.

1. SampleRef (tensor-free reference)

A frozen dataclass describing one produced sample without holding its tensors:

• sample_id, run_id
• store_uri (which store + generation, e.g. mem://..., file://...)
• feature_keys: dict[name -> store_key]
• feature_specs: dict[name -> FeatureSpec(shape, dtype)]
• algorithm (eagle3 / dflash / dspark), schema_version
• num_tokens, estimated_bytes
• versions for correctness gating: target_model_version, draft_weight_version
  (relevant once train-with-decode / weight resync is added)

An assert_no_tensors guard enforces the no-tensor invariant on every
control-plane hop. feature_specs lets the consumer preallocate a receive
buffer from the ref alone (mirrors how remote/protocol.py ships dtype+shape
metadata so the client preallocates NCCL recv buffers today).

2. FeatureStore (pluggable tensor transport)

Abstract API: put(sample_id, tensors) -> keys, get(ref, device) -> (tensors, handle), release(handle) (consume-once), gc(), health() -> capacity ints.

Backends land incrementally:

• LocalFeatureStore: in-process dict, resident-byte cap, MemoryError
  backstop. Enough to build and unit-test the whole pipeline colocated with no
  network.
• SharedDirFeatureStore: a POSIX shared mount (torch.save / safetensors).
  A natural fit for clusters where a single filesystem is mounted across all
  nodes, so a producer node writes and a trainer node reads with no explicit
  staging.
• NcclFeatureStore (optional, later): reuse remote/transport.py for
  zero-copy GPU-to-GPU, keeping the RDMA/Mooncake-style fast path as a drop-in
  backend behind the same API.

Consume-once semantics: clone-on-fetch by default, release(handle) immediately
after materialize so prefetch cannot race a free; gc() retries failed frees.

3. SampleRefQueue + backpressure

A metadata-only queue between producers and consumers with lease / ack / fail
and visibility-timeout reclaim, so an unacked ref from a crashed consumer is
redelivered.

Backpressure uses a high/low watermark hysteresis band read from
FeatureStore.health() (ints only, no tensors):

• Producers pause when resident bytes cross high_watermark_bytes.
• Producers resume only when they fall back below low_watermark_bytes.

This bounds store residency, prevents a fast target from OOMing the store, and
prevents a slow target from silently starving the trainer (separate
producer-starved vs consumer-starved counters make which side is the bottleneck
observable).

4. Consumer: FeatureDataLoader

Trainer-side loader that leases a ref, calls store.get(ref, device),
materializes tensors, releases the handle, applies the per-algorithm collate,
and yields the existing *TargetBatch (Eagle3TargetBatch first). This slots in
where the trainer currently calls to_trainer_inputs(), so the trainer module
is unchanged. Empty / short loss-mask samples are neutralized, not dropped, to
keep FSDP DP ranks in lockstep (a known EAGLE data-pipeline requirement).

Consumer-side DP resharding: partition by a stable sample_id hash so a stream
produced under one DP width is consumable under another. Partitioning is a
consumer decision, not pinned by the producer, which is what lets the two pools
scale independently.

5. Producer: AutoModel target forward (first backend)

The first FeatureProducer wraps the existing target forward
(HFEagle3TargetModel.generate_batch and, where available, the remote
generate_batch_async), draws prompts, runs the forward, and puts the
resulting tensors into the store plus a SampleRef onto the queue, honoring
backpressure. No external inference engine dependency is introduced. The
existing sglang / vllm co-located backends remain available and can later be
adapted as additional producers behind the same interface without touching the
trainer.

Data flow

prompts -> FeatureProducer (AutoModel target forward)
              -> FeatureStore.put(tensors) -> keys
              -> SampleRefQueue.put(SampleRef)   [tensor-free]
SampleRefQueue -> FeatureDataLoader.lease(ref)
              -> FeatureStore.get(ref) -> tensors -> release(handle)
              -> *TargetBatch -> Eagle3TrainerModule.forward
backpressure: producer pauses on store.health() >= high_watermark,
              resumes below low_watermark

Feature schema per algorithm

The store/queue are algorithm-agnostic; each algorithm registers a schema so the
consumer can validate a ref before materializing.

Algorithm	Required feature keys	Supervision	Notes
EAGLE-3	aux_hidden_states (3 aux layers concat, H*3), input_ids, attention_mask, loss_mask	logits [B,S,V] OR target_probs + position_mask (draft-vocab)	exactly one supervision encoding; draft-vocab projection via set_vocab_mapping
DFlash	hidden_states [B,S,len(layers)*H], input_ids, attention_mask, loss_mask	optional logits [B,S,V] (only with capture_logits=True, needed by JetSpec forward-KL)	hard-label CE needs no logits
DSpark	target_hidden_states [B,S,len(layers)*H], target_last_hidden_states [B,S,H], input_ids, loss_mask	none	needs both intermediate and final hidden states; no attention_mask

schema_version on each ref gates consumer/producer compatibility.

Phased plan (independently reviewable PRs)

Each PR is self-contained, unit-tested toward full coverage of new code, and
does not regress the existing colocated / remote paths.

• PR 1: Data-plane contracts and local store. SampleRef, FeatureSpec,
  FeatureStore ABC + LocalFeatureStore, SampleRefQueue, assert_no_tensors.
  No trainer wiring yet; pure library + contract tests. Behavior-neutral.
• PR 2: EAGLE-3 producer/consumer over the local store. FeatureProducer
  wrapping the AutoModel target forward, FeatureDataLoader yielding
  Eagle3TargetBatch, selected by a new target_model_backend: streaming value
  with the loader colocated in-process. Closes the minimal end-to-end loop
  (ref queue + get-by-key) with numerical parity against the colocated path.
• PR 3: Asynchronous production + backpressure. Producer runs ahead of the
  trainer with the high/low watermark hysteresis; overlap target forward with
  draft training. Adds SharedDirFeatureStore so producer and trainer can be
  separate processes/nodes on the shared filesystem.
• PR 4: Independent scaling + fault isolation. Multiple producer replicas, a
  controller managing them, consumer-side DP resharding, visibility-timeout
  redelivery, and optional NcclFeatureStore reusing remote/transport.py for
  the fast tensor path. World sizes of the two pools fully decoupled.

Generalizing the producer/consumer path to DFlash and DSpark (giving them their
first disaggregated backend) is a natural follow-up once the EAGLE-3 path is
proven, and is out of scope for the initial chain.

Scope and non-goals

• Not required for, and must not complicate, small single-node runs. Colocated
  stays the default; streaming is opt-in via target_model_backend.
• First producer is the AutoModel target forward only. External inference
  engines (SGLang / vLLM / TRT-LLM) as producers are explicitly out of scope for
  the initial chain, though the interface is designed not to preclude them.
• No new heavy dependency (Ray, Mooncake) is required to land PR 1-3. The local
  and shared-dir stores cover build-out and testing; an RDMA/GPUDirect backend
  is an optional later addition behind the same FeatureStore API.

Open questions

1. Should target_model_backend: streaming be a new backend value, or a
   streaming: true modifier layered on top of the existing
   {colocated, remote} backends? (Leaning: new value, to keep selection flat.)
2. Store residency policy for packed sequences: cap by sample count, by resident
   bytes, or both? (Leaning: both, bytes as the hard backstop.)
3. Do we want durable ref metadata (resume a trainer mid-run from committed,
   unacked refs) in the initial chain, or defer to the multi-replica PR?
4. Minimum viable controller for PR 4: a plain process supervisor, or reuse the
   existing HTTP control protocol in remote/protocol.py for producer health /
   registration?

Start with the existing contracts in components/speculative/eagle/backend.py and components/speculative/eagle/target.py, then review remote/transport.py, remote/wire.py, and remote/protocol.py. The phased plan identifies PR 1's SampleRef, FeatureSpec, FeatureStore, LocalFeatureStore, SampleRefQueue, and contract tests as the first boundary; done means the tensor-free contracts and local store are unit-tested without trainer wiring.