nemo_automodel.components.speculative.serve_target
nemo_automodel.components.speculative.serve_target
Launch a remote EAGLE-3 target server (train-inference disaggregation).
Loads the frozen target model on this process’s GPU and serves draft-vocab
supervision (aux hidden states, target_probs, position_mask) to a
training client over HTTP (control plane) + NCCL (data plane).
Typical usage (single-GPU server)::
CUDA_VISIBLE_DEVICES=0 python -m nemo_automodel.components.speculative.serve_target
—target meta-llama/Llama-3.1-8B-Instruct
—host 0.0.0.0 —port 8001
Then point training at it::
recipe_args.target_model_backend: remote recipe_args.remote_urls: [“http://<server-host>:8001”] recipe_args.target_prefetch_depth: 1
Verify readiness with curl http://<host>:8001/health. NCCL GPU-direct
transfer requires sglang installed in the server’s environment; without it the
server transparently falls back to the binary wire format.
The frozen target runs under one of three inference engines (--engine):
-
hf(default): HuggingFace forward with aux-layer hooks (:class:HFEagle3TargetModel); works for any AutoModel target. -
sglang: SGLang forward (:class:SGLangEagle3TargetModel); faster for mainstream architectures. SGLang is not a NeMo-AutoModel dependency — it pinstransformers==4.57.1, which conflicts with the training stack — so it is imported only when this engine is selected. Install it yourself in a separate, dedicated speculative-decoding venv/container on the server::uv pip install sglang==0.5.9
-
vllm: vLLM forward (:class:VLLMEagle3TargetModel) via vLLM’s nativeextract_hidden_statespath. Like sglang, vLLM is not a NeMo-AutoModel dependency and is imported only when this engine is selected. Install it yourself in a separate, dedicated speculative-decoding venv/container on the server::uv pip install vllm==0.23.0
All engines emit the identical supervision contract, so the training client is unchanged regardless of which one serves the target.
Module Contents
Functions
Data
API
Load the target under HuggingFace and wrap it with aux-layer hooks.
Load the target under SGLang (imported here so hf runs need no SGLang).
SGLang places the model itself (torch.cuda.current_device()), so
device is unused; the signature matches :func:_build_hf_target so both
are interchangeable in the engine dispatch.
Load the target under vLLM (imported here so hf runs need no vLLM).
vLLM places the model itself (torch.cuda.current_device()), so device
is unused; the signature matches :func:_build_hf_target so both are
interchangeable in the engine dispatch.
Load the target model and run the blocking HTTP + NCCL server.