nemo_automodel.components.speculative.serve_target

View as Markdown

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://&lt;host&gt;: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 pins transformers==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 native extract_hidden_states path. 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

NameDescription
_build_hf_targetLoad the target under HuggingFace and wrap it with aux-layer hooks.
_build_sglang_targetLoad the target under SGLang (imported here so hf runs need no SGLang).
_build_vllm_targetLoad the target under vLLM (imported here so hf runs need no vLLM).
_parse_args-
mainLoad the target model and run the blocking HTTP + NCCL server.

Data

logger

API

nemo_automodel.components.speculative.serve_target._build_hf_target(
args,
device: torch.device,
dtype: torch.dtype
) -> nemo_automodel.components.speculative.eagle.target.HFEagle3TargetModel

Load the target under HuggingFace and wrap it with aux-layer hooks.

nemo_automodel.components.speculative.serve_target._build_sglang_target(
args,
device: torch.device,
dtype: torch.dtype
)

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.

nemo_automodel.components.speculative.serve_target._build_vllm_target(
args,
device: torch.device,
dtype: torch.dtype
)

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.

nemo_automodel.components.speculative.serve_target._parse_args(
argv = None
) -> argparse.Namespace
nemo_automodel.components.speculative.serve_target.main(
argv = None
) -> None

Load the target model and run the blocking HTTP + NCCL server.

nemo_automodel.components.speculative.serve_target.logger = logging.getLogger(__name__)