From 4e5fa152544cf36b4a3c95595947ce40fe28c326 Mon Sep 17 00:00:00 2001 From: khazic Date: Fri, 5 Jun 2026 23:38:17 +0800 Subject: [PATCH 01/12] feat(speculative): add SGLang target backend for EAGLE-3 training Add a third Eagle3TargetBackend implementation that runs the frozen target through SGLang, alongside the co-located HF and remote backends. SGLang is the fastest serving path for mainstream architectures, so the remote target server can hold the target on dedicated GPUs while the draft trains elsewhere. The backend is split into two layers: - SGLangEagle3TargetModel (sglang_target.py) owns the supervision contract and assembles an Eagle3TargetBatch whose shift / aux-concatenation semantics are byte-for-byte identical to HFEagle3TargetModel, so a SGLang run is numerically equivalent to a co-located one. It depends only on a small runner protocol, so it is unit-testable on CPU without SGLang. - SGLangTargetRunner (sglang_runner.py) owns the SGLang-internal forward (ModelRunner + CaptureHiddenMode.FULL + a logits-processor wrap that returns all-position full-vocab logits plus the three concatenated aux hidden states). It is lazily imported and validated on the GPU server; CPU tests cover the surface that does not need SGLang. serve_target gains an --engine {hf,sglang} flag (engines share a builder signature via a dispatch map). SGLang is declared as an optional spec_sglang extra pinned to 0.5.9 and kept out of the main training image. Shared aux-layer-id default / validation helpers are extracted in target.py so all backends default identically (behavior-preserving refactor). Signed-off-by: khazic --- .../speculative/eagle/sglang_runner.py | 318 ++++++++++++++++++ .../speculative/eagle/sglang_target.py | 180 ++++++++++ .../components/speculative/eagle/target.py | 83 +++-- .../components/speculative/serve_target.py | 68 +++- pyproject.toml | 6 + .../speculative/test_eagle3_sglang.py | 292 ++++++++++++++++ 6 files changed, 902 insertions(+), 45 deletions(-) create mode 100644 nemo_automodel/components/speculative/eagle/sglang_runner.py create mode 100644 nemo_automodel/components/speculative/eagle/sglang_target.py create mode 100644 tests/unit_tests/speculative/test_eagle3_sglang.py diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py new file mode 100644 index 0000000000..5ff46f7552 --- /dev/null +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -0,0 +1,318 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SGLang ModelRunner forward for the EAGLE-3 target (server-side, GPU only). + +This module owns every SGLang-internal touch point so the rest of the +speculative stack stays SGLang-agnostic and importable without SGLang. It is +imported lazily (only from :meth:`SGLangEagle3TargetModel.from_pretrained`). + +Mechanism (mirrors SpecForge's SGLang backend, which is the only path that +returns the supervision tensors directly without a Mooncake transfer layer): + +1. Build a SGLang ``ModelRunner`` with ``enable_return_hidden_states=True``. +2. Wrap the model's ``LogitsProcessor`` so a single extend forward returns + *all-position* full-vocab logits (stock SGLang only keeps the last position) + alongside the three concatenated EAGLE-3 auxiliary hidden states. +3. Run one extend per request and stack the per-row results into batched + ``[batch, seq, *]`` tensors for :class:`SGLangEagle3TargetModel`. + +Unlike SpecForge (which embeds the target inside the training job and reuses the +trainer's TP process group), this runs in a *standalone* server process, so it +performs SGLang's own single-process distributed init rather than reusing an +external group, and it drops SpecForge's ``shard_returns`` / VLM paths that only +matter inside the training loop. + +SGLang's private ``LogitsProcessor`` helpers are version-coupled: the calls here +track ``sglang==0.5.9`` (SpecForge's pin). This forward path requires a GPU and +SGLang, so it is validated on the training server, not in CPU unit tests; the +CPU tests exercise the contract layer in +:mod:`nemo_automodel.components.speculative.eagle.sglang_target` against a fake +runner instead. +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional, Sequence + +import torch + +logger = logging.getLogger(__name__) + + +def _wrap_logits_processors_for_eagle3(model) -> None: # pragma: no cover - requires GPU + SGLang + """Replace every SGLang ``LogitsProcessor`` in ``model`` with an EAGLE-3 wrapper. + + The wrapper makes one extend forward return all-position full-vocab logits + plus the concatenated auxiliary hidden states, instead of stock SGLang's + last-position-only logits. Ported (simplified, no tensor-parallel sharding) + from SpecForge ``sglang_backend/utils.py`` for ``sglang==0.5.9``. + """ + from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + + class _LogitsProcessorForEagle3(torch.nn.Module): + def __init__(self, inner: LogitsProcessor): + super().__init__() + self.logits_processor = inner + + def forward( + self, + input_ids, + hidden_states, + lm_head, + logits_metadata, + aux_hidden_states: Optional[list] = None, + hidden_states_before_norm: Optional[torch.Tensor] = None, + ): + lp = self.logits_processor + # Force the extend forward down the decode branch so the patched + # path below is taken for every position, matching SpecForge. + logits_metadata.forward_mode = ForwardMode.DECODE + if isinstance(logits_metadata, ForwardBatch): + logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata) + + pruned_states, pruned_states_before_norm, aux_pruned_states, sample_indices, *_ = lp._get_pruned_states( + hidden_states, + hidden_states_before_norm, + aux_hidden_states, + logits_metadata, + ) + logits = self._compute_full_logits(lp, pruned_states, lm_head, logits_metadata) + aux = lp._get_hidden_states_to_store( + hidden_states, + hidden_states_before_norm, + aux_hidden_states, + pruned_states, + pruned_states_before_norm, + aux_pruned_states, + sample_indices, + logits_metadata, + ) + return _Eagle3LogitsOutput(logits=logits, aux_hidden_states=aux) + + @staticmethod + def _compute_full_logits(lp, hidden_states, lm_head, logits_metadata): + from sglang.srt.distributed import tensor_model_parallel_all_gather + from sglang.srt.utils.common import is_npu + + hidden_states, local_hidden_states = lp._gather_dp_attn_hidden_states(hidden_states, logits_metadata) + logits = lp._compute_lm_head(hidden_states, lm_head, None) + if lp.logit_scale is not None: + logits.mul_(lp.logit_scale) + if lp.do_tensor_parallel_all_gather: + if lp.use_attn_tp_group: + logits = lp._gather_attn_tp_logits(logits) + else: + logits = tensor_model_parallel_all_gather(logits) + logits = lp._scatter_dp_attn_logits(logits, local_hidden_states, logits_metadata) + logits = lp._copy_logits_to_buffer(logits, logits_metadata) + if lp.final_logit_softcapping: + if not is_npu(): + from sglang.srt.layers.logits_processor import fused_softcap + + fused_softcap(logits, lp.final_logit_softcapping) + else: + cap = lp.final_logit_softcapping + logits = cap * torch.tanh(logits / cap) + return logits + + for name, submodule in list(model.named_modules()): + if isinstance(submodule, LogitsProcessor): + setattr(model, name, _LogitsProcessorForEagle3(submodule)) + logger.info("wrapped %s with EAGLE-3 logits processor", name) + + +class _Eagle3LogitsOutput: # pragma: no cover - only built inside the GPU-only wrapper + """Carries the all-position logits + aux hidden states out of the wrapper.""" + + def __init__(self, logits: torch.Tensor, aux_hidden_states: torch.Tensor): + self.logits = logits + self.aux_hidden_states = aux_hidden_states + + +class SGLangTargetRunner: + """Standalone SGLang ModelRunner that returns EAGLE-3 supervision tensors. + + Built via :meth:`build`; consumed through the + :class:`~nemo_automodel.components.speculative.eagle.sglang_target.SGLangRunnerProtocol` + surface (``model`` / ``set_aux_layers`` / ``forward_eagle3`` / + ``input_embedding_weight``). + """ + + def __init__(self, model_runner): + self._model_runner = model_runner + + @property + def model(self): + """The loaded nn.Module (exposes ``.config`` and ``.parameters()``).""" + return self._model_runner.model + + @classmethod + def build( # pragma: no cover - requires GPU + SGLang + cls, + model_path: str, + *, + dtype: Optional[torch.dtype] = None, + tp_size: int = 1, + trust_remote_code: bool = False, + **sglang_kwargs, + ) -> "SGLangTargetRunner": + """Construct the SGLang ModelRunner for a standalone target server. + + ``sglang_kwargs`` are forwarded to ``ServerArgs`` (e.g. ``page_size``, + ``mem_fraction_static``, ``attention_backend``). The constructor mirrors + SpecForge's ``sglang==0.5.9`` usage and is GPU/SGLang-only. + """ + import torch.distributed as dist + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.distributed import init_distributed_environment, initialize_model_parallel + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.server_args import ServerArgs + + if not torch.cuda.is_available(): + raise RuntimeError("SGLangTargetRunner requires CUDA; run it on a GPU server, not the editing host.") + + server_args = ServerArgs( + model_path=model_path, + trust_remote_code=trust_remote_code, + dtype=dtype if dtype is not None else "auto", + enable_return_hidden_states=True, + disable_cuda_graph=True, # extend-only forward; CUDA graphs add no benefit here + disable_radix_cache=True, + tp_size=tp_size, + pp_size=1, + **sglang_kwargs, + ) + + gpu_id = torch.cuda.current_device() + # Standalone server: SGLang manages its own single-rank world rather than + # reusing a trainer process group (the SpecForge embedded case). + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", str(int(server_args.port or 8000) + 200)) + init_distributed_environment( + backend="nccl", + world_size=tp_size, + rank=0, + local_rank=gpu_id, + distributed_init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}", + ) + initialize_model_parallel(tensor_model_parallel_size=tp_size) + + model_config = ModelConfig.from_server_args(server_args) + model_runner = ModelRunner( + model_config=model_config, + mem_fraction_static=server_args.mem_fraction_static, + gpu_id=gpu_id, + tp_rank=0, + tp_size=tp_size, + pp_rank=0, + pp_size=1, + nccl_port=None, + server_args=server_args, + ) + _wrap_logits_processors_for_eagle3(model_runner.model) + return cls(model_runner) + + def set_aux_layers(self, aux_layer_ids: Sequence[int]) -> None: + """Tell the SGLang model which 3 decoder layers to capture.""" + self._model_runner.model.set_eagle3_layers_to_capture(list(aux_layer_ids)) + + def input_embedding_weight(self) -> torch.Tensor: + """Return the target input-embedding weight ``[vocab, hidden]``.""" + return self._model_runner.model.get_input_embeddings().weight + + @torch.no_grad() + def forward_eagle3( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run one extend per row and stack the per-position logits and aux states. + + Returns ``(logits[batch, seq, vocab], aux[batch, seq, 3 * hidden])``, + both unshifted; the contract layer applies the EAGLE-3 shift. Sequences + must share a length (training batches are padded), so the per-row + results stack cleanly. + """ + del attention_mask # each row is a full causal sequence in SGLang + logits_list, aux_list = self._extend(input_ids) + logits = torch.stack(logits_list, dim=0) + aux = torch.stack(aux_list, dim=0) + return logits, aux + + def _extend(self, input_ids: torch.Tensor) -> tuple[list, list]: # pragma: no cover - requires GPU + SGLang + from sglang.srt.managers.schedule_batch import Req, ScheduleBatch + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.radix_cache import RadixCache + from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch + from sglang.srt.sampling.sampling_params import SamplingParams + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + runner = self._model_runner + sampling_params = SamplingParams(temperature=0, max_new_tokens=1, top_k=1) + rows = torch.split(input_ids, 1, dim=0) + reqs, input_lens = [], [] + for idx, row in enumerate(rows): + token_ids = row.view(-1).tolist() + req = Req( + rid=str(idx), + origin_input_text="", + origin_input_ids=token_ids, + sampling_params=sampling_params, + ) + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + req.logprob_start_len = len(req.origin_input_ids) - 1 + reqs.append(req) + input_lens.append(len(token_ids)) + + cache_params = CacheInitParams( + disable=False, + req_to_token_pool=runner.req_to_token_pool, + token_to_kv_pool_allocator=runner.token_to_kv_pool_allocator, + page_size=runner.server_args.page_size, + ) + batch = ScheduleBatch.init_new( + reqs=reqs, + req_to_token_pool=runner.req_to_token_pool, + token_to_kv_pool_allocator=runner.token_to_kv_pool_allocator, + tree_cache=RadixCache(cache_params), + model_config=runner.model_config, + enable_overlap=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + ) + batch.prepare_for_extend() + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, runner) + forward_batch.capture_hidden_mode = CaptureHiddenMode.FULL + output = runner.forward(forward_batch).logits_output + + logits_list = list(torch.split(output.logits, input_lens, dim=0)) + aux_list = list(torch.split(output.aux_hidden_states, input_lens, dim=0)) + runner.req_to_token_pool.clear() + runner.token_to_kv_pool_allocator.clear() + return logits_list, aux_list + + def close(self) -> None: + """Release the SGLang model runner (best effort).""" + runner, self._model_runner = self._model_runner, None + if runner is None: + return + for name in ("token_to_kv_pool_allocator", "req_to_token_pool"): + pool = getattr(runner, name, None) + clear = getattr(pool, "clear", None) + if clear is not None: + clear() diff --git a/nemo_automodel/components/speculative/eagle/sglang_target.py b/nemo_automodel/components/speculative/eagle/sglang_target.py new file mode 100644 index 0000000000..dc9c34d75b --- /dev/null +++ b/nemo_automodel/components/speculative/eagle/sglang_target.py @@ -0,0 +1,180 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SGLang-backed EAGLE-3 target model. + +A third :class:`Eagle3TargetBackend` implementation: it runs the frozen target +through SGLang instead of HuggingFace, capturing the same EAGLE-3 supervision +(three auxiliary hidden states plus full-vocab logits) the co-located backend +produces. SGLang is the fastest serving path for mainstream architectures, so a +remote target server (:mod:`nemo_automodel.components.speculative.serve_target`) +can hold the target on dedicated GPUs while the draft trains elsewhere. + +The class is split into two layers so the contract is unit-testable without a +GPU or SGLang installed: + +- :class:`SGLangEagle3TargetModel` (this file) owns the *contract*: it assembles + an :class:`Eagle3TargetBatch` whose shift / aux-concatenation semantics are + identical to :class:`HFEagle3TargetModel`, so a SGLang run is numerically + equivalent to the co-located one. It depends only on a small runner surface + (``forward_eagle3`` / ``input_embedding_weight`` / ``set_aux_layers`` and a + ``model`` exposing ``.config`` + ``.parameters()``), which tests fake. +- :class:`~nemo_automodel.components.speculative.eagle.sglang_runner.SGLangTargetRunner` + owns the SGLang-internal forward and is lazily imported only by + :meth:`SGLangEagle3TargetModel.from_pretrained`, so importing this module + never pulls in SGLang. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING, Optional, Protocol, Sequence + +import torch + +from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend +from nemo_automodel.components.speculative.eagle.target import ( + Eagle3TargetBatch, + _shift_left_with_zero, + default_eagle3_aux_layer_ids, + validate_eagle3_aux_layer_ids, +) + +if TYPE_CHECKING: + import torch.nn as nn + + +class SGLangRunnerProtocol(Protocol): + """Minimal surface :class:`SGLangEagle3TargetModel` needs from a runner. + + Implemented for real by ``SGLangTargetRunner`` (GPU/SGLang) and faked in + unit tests, which is why the backend depends on this protocol rather than on + SGLang directly. + """ + + #: Loaded model handle exposing ``.config`` (with ``num_hidden_layers`` / + #: ``hidden_size`` / ``vocab_size``) and ``.parameters()`` for device + #: inference, mirroring what the server reads off ``HFEagle3TargetModel``. + model: "nn.Module" + + def set_aux_layers(self, aux_layer_ids: Sequence[int]) -> None: + """Tell the underlying model which 3 decoder layers to capture.""" + + def forward_eagle3( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the target once and return ``(logits, aux_hidden_states)``. + + ``logits`` is ``[batch, seq, vocab]`` (full vocab, unshifted) and + ``aux_hidden_states`` is ``[batch, seq, 3 * hidden]`` (the three capture + layers concatenated on the last dim, unshifted). + """ + + def input_embedding_weight(self) -> torch.Tensor: + """Return the target input-embedding weight ``[vocab, hidden]``.""" + + +class SGLangEagle3TargetModel(Eagle3TargetBackend): + """EAGLE-3 target backend that runs the frozen target through SGLang. + + Parameters + ---------- + runner: + A loaded runner implementing :class:`SGLangRunnerProtocol`. + aux_layer_ids: + The three decoder layers to capture (low / mid / high). When ``None`` + the shared EAGLE-3 default recipe is used, matching every other backend. + """ + + def __init__(self, runner: SGLangRunnerProtocol, aux_layer_ids: Optional[Sequence[int]] = None): + self._runner = runner + # Expose ``.model`` so the remote server can read the target's config + # and infer its device the same way it does for the co-located backend. + self.model = runner.model + num_layers = self.model.config.num_hidden_layers + if aux_layer_ids is None: + aux_layer_ids = default_eagle3_aux_layer_ids(num_layers) + self.aux_layer_ids = validate_eagle3_aux_layer_ids(aux_layer_ids, num_layers) + runner.set_aux_layers(self.aux_layer_ids) + + def get_input_embeddings(self) -> SimpleNamespace: + """Return an object exposing ``.weight`` (the target input embeddings). + + Matches the offline-cache / remote path: the draft's + ``copy_embeddings_from_target`` only reads ``.weight``. + """ + return SimpleNamespace(weight=self._runner.input_embedding_weight()) + + @torch.no_grad() + def generate_batch( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + loss_mask: torch.Tensor, + ) -> Eagle3TargetBatch: + """Run the SGLang target and capture aux hidden states plus logits. + + Produces an :class:`Eagle3TargetBatch` byte-for-byte compatible with + :meth:`HFEagle3TargetModel.generate_batch`: the logits, ``input_ids`` + and ``loss_mask`` are shifted left by one (next-token alignment) while + the aux hidden states are kept position-aligned. The draft-vocab + projection happens trainer-side / server-side from ``logits``, exactly + as in the co-located path. + + Note: sequences are assumed right-padded (loss_mask zeros the pad). With + causal attention, trailing pad tokens do not affect earlier positions, + so the captured supervision matches a masked HuggingFace forward. + """ + logits, aux_hidden_states = self._runner.forward_eagle3(input_ids, attention_mask) + return Eagle3TargetBatch( + aux_hidden_states=aux_hidden_states, + logits=_shift_left_with_zero(logits), + input_ids=_shift_left_with_zero(input_ids), + attention_mask=attention_mask, + loss_mask=_shift_left_with_zero(loss_mask), + ) + + def close(self) -> None: + """Release the SGLang runner (frees GPU memory / engine handles).""" + close = getattr(self._runner, "close", None) + if close is not None: + close() + + @classmethod + def from_pretrained( # pragma: no cover - requires GPU + SGLang + cls, + model_path: str, + *, + aux_layer_ids: Optional[Sequence[int]] = None, + dtype: Optional[torch.dtype] = None, + tp_size: int = 1, + trust_remote_code: bool = False, + **sglang_kwargs, + ) -> "SGLangEagle3TargetModel": + """Build a SGLang runner for ``model_path`` and wrap it as a target backend. + + SGLang is imported here (not at module load) so this module stays + importable in environments without SGLang; ``sglang_kwargs`` are passed + through to SGLang's ``ServerArgs`` for endpoint / parallelism tuning. + """ + from nemo_automodel.components.speculative.eagle.sglang_runner import SGLangTargetRunner + + runner = SGLangTargetRunner.build( + model_path, + dtype=dtype, + tp_size=tp_size, + trust_remote_code=trust_remote_code, + **sglang_kwargs, + ) + return cls(runner, aux_layer_ids=aux_layer_ids) diff --git a/nemo_automodel/components/speculative/eagle/target.py b/nemo_automodel/components/speculative/eagle/target.py index 862b92b144..d802980b51 100644 --- a/nemo_automodel/components/speculative/eagle/target.py +++ b/nemo_automodel/components/speculative/eagle/target.py @@ -38,6 +38,51 @@ def _shift_left_with_zero(tensor: torch.Tensor) -> torch.Tensor: return torch.cat((tensor[:, 1:], tail), dim=1) +def default_eagle3_aux_layer_ids(num_layers: int) -> list[int]: + """Return the EAGLE-3 default 3-layer (low / mid / high) aux capture recipe. + + The downstream draft model's ``fc`` projection is sized for exactly + ``num_aux_hidden_states`` layers (default 3) of concatenated target hidden + states. Silently deduplicating collisions on shallow targets would yield + fewer than 3 captured tensors and crash later inside the draft ``fc`` with a + confusing shape-mismatch error -- raise here instead so the caller picks 3 + distinct in-bounds ids that match the draft config. Shared by every target + backend (co-located, remote, SGLang) so they all default identically. + """ + candidates = [1, num_layers // 2 - 1, num_layers - 4] + if any(c < 0 or c >= num_layers for c in candidates) or len(set(candidates)) != 3: + raise ValueError( + f"Target model has num_hidden_layers={num_layers}, which is too shallow " + f"for the default EAGLE-3 aux recipe {candidates}. Pass aux_layer_ids " + f"explicitly (must be 3 distinct in-bounds layer indices, matching the " + f"draft model's num_aux_hidden_states)." + ) + return candidates + + +def validate_eagle3_aux_layer_ids(aux_layer_ids: Sequence[int], num_layers: int) -> list[int]: + """Validate an aux-layer selection against a target of ``num_layers`` depth. + + Shared by every target backend so an explicit ``aux_layer_ids`` is checked + identically whether the target runs co-located, remote, or under SGLang. + """ + aux_layer_ids = list(aux_layer_ids) + if len(aux_layer_ids) != 3: + raise ValueError( + f"EAGLE-3 expects exactly 3 aux_layer_ids, but got {len(aux_layer_ids)}: " + f"{aux_layer_ids}. This must match the draft model's num_aux_hidden_states." + ) + if len(set(aux_layer_ids)) != len(aux_layer_ids): + raise ValueError( + f"EAGLE-3 aux_layer_ids must be distinct, but got {aux_layer_ids}. " + "Duplicate ids would collapse the captured aux hidden states." + ) + for layer_id in aux_layer_ids: + if layer_id < 0 or layer_id >= num_layers: + raise ValueError(f"aux layer id {layer_id} is out of bounds for model with {num_layers} layers") + return aux_layer_ids + + @dataclass class Eagle3TargetBatch: """Target-model supervision for one draft-training batch. @@ -97,45 +142,11 @@ def __init__(self, model: nn.Module, aux_layer_ids: Sequence[int] | None = None) self.aux_layer_ids = self._validate_aux_layer_ids(candidate_ids) def _default_aux_layer_ids(self) -> list[int]: - # EAGLE-3 default 3-layer recipe (low / mid / high). - # - # The downstream draft model's ``fc`` projection is sized for - # exactly ``num_aux_hidden_states`` layers (default 3) of - # concatenated target hidden states. Silently deduplicating - # collisions on shallow targets would yield fewer than 3 - # captured tensors and crash later inside the draft ``fc`` with - # a confusing shape-mismatch error -- raise here instead so the - # caller picks 3 distinct in-bounds ids that match the draft - # config. - num_layers = self.model.config.num_hidden_layers - candidates = [1, num_layers // 2 - 1, num_layers - 4] - if any(c < 0 or c >= num_layers for c in candidates) or len(set(candidates)) != 3: - raise ValueError( - f"Target model has num_hidden_layers={num_layers}, which is too shallow " - f"for the default EAGLE-3 aux recipe {candidates}. Pass aux_layer_ids " - f"explicitly (must be 3 distinct in-bounds layer indices, matching the " - f"draft model's num_aux_hidden_states)." - ) - return candidates + return default_eagle3_aux_layer_ids(self.model.config.num_hidden_layers) def _validate_aux_layer_ids(self, aux_layer_ids: Sequence[int]) -> list[int]: """Validate aux-layer selection before any forward hooks are registered.""" - num_layers = self.model.config.num_hidden_layers - aux_layer_ids = list(aux_layer_ids) - if len(aux_layer_ids) != 3: - raise ValueError( - f"EAGLE-3 expects exactly 3 aux_layer_ids, but got {len(aux_layer_ids)}: " - f"{aux_layer_ids}. This must match the draft model's num_aux_hidden_states." - ) - if len(set(aux_layer_ids)) != len(aux_layer_ids): - raise ValueError( - f"EAGLE-3 aux_layer_ids must be distinct, but got {aux_layer_ids}. " - "Duplicate ids would collapse the captured aux hidden states." - ) - for layer_id in aux_layer_ids: - if layer_id < 0 or layer_id >= num_layers: - raise ValueError(f"aux layer id {layer_id} is out of bounds for model with {num_layers} layers") - return aux_layer_ids + return validate_eagle3_aux_layer_ids(aux_layer_ids, self.model.config.num_hidden_layers) def _get_transformer_layers(self) -> list[nn.Module]: """Return decoder layers as an ordered list indexable by integer. diff --git a/nemo_automodel/components/speculative/serve_target.py b/nemo_automodel/components/speculative/serve_target.py index 0a061f96f2..5b5bbbcb2b 100644 --- a/nemo_automodel/components/speculative/serve_target.py +++ b/nemo_automodel/components/speculative/serve_target.py @@ -33,6 +33,17 @@ Verify readiness with ``curl http://: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 two 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 pinned in a separate speculative-decoding + environment, so it is imported only when this engine is selected. + +Both engines emit the identical supervision contract, so the training client is +unchanged regardless of which one serves the target. """ from __future__ import annotations @@ -49,9 +60,52 @@ logger = logging.getLogger(__name__) +def _build_hf_target(args, device: torch.device, dtype: torch.dtype) -> HFEagle3TargetModel: + """Load the target under HuggingFace and wrap it with aux-layer hooks.""" + target_model = NeMoAutoModelForCausalLM.from_pretrained( + args.target, + torch_dtype=dtype, + trust_remote_code=args.trust_remote_code, + ) + target_model.to(device) + target_model.requires_grad_(False) + return HFEagle3TargetModel(target_model, aux_layer_ids=args.aux_layer_ids) + + +def _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. + """ + del device + from nemo_automodel.components.speculative.eagle.sglang_target import SGLangEagle3TargetModel + + return SGLangEagle3TargetModel.from_pretrained( + args.target, + aux_layer_ids=args.aux_layer_ids, + dtype=dtype, + tp_size=args.tp_size, + trust_remote_code=args.trust_remote_code, + ) + + def _parse_args(argv=None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Serve an EAGLE-3 target model for remote draft training.") parser.add_argument("--target", required=True, help="Target model name or path.") + parser.add_argument( + "--engine", + choices=("hf", "sglang"), + default="hf", + help="Inference engine for the frozen target (default: hf).", + ) + parser.add_argument( + "--tp-size", + type=int, + default=1, + help="Tensor-parallel size for the sglang engine (ignored for hf).", + ) parser.add_argument("--host", default="0.0.0.0", help="Bind address (0.0.0.0 for cross-machine).") parser.add_argument("--port", type=int, default=8001, help="HTTP control-plane port.") parser.add_argument( @@ -79,15 +133,11 @@ def main(argv=None) -> None: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 - logger.info("Loading target model %s on %s", args.target, device) - target_model = NeMoAutoModelForCausalLM.from_pretrained( - args.target, - torch_dtype=dtype, - trust_remote_code=args.trust_remote_code, - ) - target_model.to(device) - target_model.requires_grad_(False) - target_wrapper = HFEagle3TargetModel(target_model, aux_layer_ids=args.aux_layer_ids) + logger.info("Loading target model %s on %s via %s engine", args.target, device, args.engine) + # Engine builders share a ``(args, device, dtype)`` signature so adding an + # engine is one map entry; resolved at call time so tests can patch builders. + builders = {"hf": _build_hf_target, "sglang": _build_sglang_target} + target_wrapper = builders[args.engine](args, device, dtype) nccl_port = args.nccl_port if args.nccl_port is not None else args.port + 100 server_logic = TargetModelServer(target_wrapper, nccl_port=nccl_port, host=args.host) diff --git a/pyproject.toml b/pyproject.toml index 7218aebb95..ab6c0790d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,6 +160,12 @@ s3 = [ msc = [ "multi-storage-client>=0.13", ] +# Speculative-decoding (EAGLE-3) SGLang target backend. Kept out of the main +# training image and pinned in a separate dedicated SD environment so the +# SGLang version stays isolated from the rest of the stack. +spec_sglang = [ + "sglang==0.5.9", +] all = [ "nemo_automodel[cuda]", "nemo_automodel[delta-databricks]", diff --git a/tests/unit_tests/speculative/test_eagle3_sglang.py b/tests/unit_tests/speculative/test_eagle3_sglang.py new file mode 100644 index 0000000000..f07e588352 --- /dev/null +++ b/tests/unit_tests/speculative/test_eagle3_sglang.py @@ -0,0 +1,292 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for the SGLang EAGLE-3 target backend (contract layer). + +The SGLang forward itself needs a GPU + SGLang and is validated on the server; +here we fake the runner and verify the parts that must hold regardless of +engine: + +1. The supervision contract matches the co-located HF backend bit-for-bit when + the runner returns the same raw logits / aux states, so an SGLang run is + numerically equivalent to a co-located one. +2. aux-layer defaulting / validation is shared with the HF backend. +3. ``get_input_embeddings`` / ``set_aux_layers`` / ``close`` are wired through. +4. ``serve_target`` routes ``--engine`` to the right builder. +""" + +from __future__ import annotations + +import types + +import pytest +import torch +import torch.nn as nn +from transformers.modeling_outputs import CausalLMOutput + +from nemo_automodel.components.speculative.eagle.sglang_target import SGLangEagle3TargetModel +from nemo_automodel.components.speculative.eagle.target import ( + HFEagle3TargetModel, + default_eagle3_aux_layer_ids, + validate_eagle3_aux_layer_ids, +) + +_VOCAB = 32 +_HIDDEN = 16 +_LAYERS = 4 +_AUX = [0, 1, 3] + + +class _FakeCausalLM(nn.Module): + """Deterministic causal-LM stand-in shared by the HF and fake-SGLang paths.""" + + def __init__(self) -> None: + super().__init__() + self.config = type("Cfg", (), {"num_hidden_layers": _LAYERS, "hidden_size": _HIDDEN, "vocab_size": _VOCAB}) + self.embed_tokens = nn.Embedding(_VOCAB, _HIDDEN) + self.layers = nn.ModuleList([nn.Linear(_HIDDEN, _HIDDEN) for _ in range(_LAYERS)]) + self.lm_head = nn.Linear(_HIDDEN, _VOCAB, bias=False) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + def forward(self, input_ids, attention_mask=None, **kwargs): + h = self.embed_tokens(input_ids) + for layer in self.layers: + h = layer(h) + return CausalLMOutput(logits=self.lm_head(h)) + + +class _FakeSGLangRunner: + """Fake runner: returns the *unshifted* logits / aux a real runner would. + + It reuses ``_FakeCausalLM`` so the raw tensors are exactly what the HF + backend captures, letting the test assert engine-equivalence. + """ + + def __init__(self, model: _FakeCausalLM): + self.model = model + self.aux_layer_ids = None + self.closed = False + + def set_aux_layers(self, aux_layer_ids): + self.aux_layer_ids = list(aux_layer_ids) + + def input_embedding_weight(self): + return self.model.get_input_embeddings().weight + + def forward_eagle3(self, input_ids, attention_mask): + captured = {} + handles = [ + self.model.layers[i].register_forward_hook( + lambda _m, _i, out, i=i: captured.__setitem__(i, out) + ) + for i in self.aux_layer_ids + ] + try: + logits = self.model(input_ids=input_ids, attention_mask=attention_mask).logits + finally: + for h in handles: + h.remove() + aux = torch.cat([captured[i] for i in self.aux_layer_ids], dim=-1) + return logits, aux + + def close(self): + self.closed = True + + +def _inputs(): + torch.manual_seed(1) + input_ids = torch.randint(0, _VOCAB, (2, 5)) + attention_mask = torch.ones_like(input_ids) + loss_mask = torch.ones_like(input_ids) + return input_ids, attention_mask, loss_mask + + +def test_generate_batch_matches_colocated_hf(): + """SGLang supervision is bit-for-bit identical to the co-located HF path.""" + torch.manual_seed(0) + model = _FakeCausalLM().eval() + hf = HFEagle3TargetModel(model, aux_layer_ids=_AUX) + sgl = SGLangEagle3TargetModel(_FakeSGLangRunner(model), aux_layer_ids=_AUX) + + input_ids, attention_mask, loss_mask = _inputs() + hf_batch = hf.generate_batch(input_ids, attention_mask, loss_mask) + sgl_batch = sgl.generate_batch(input_ids, attention_mask, loss_mask) + + torch.testing.assert_close(sgl_batch.logits, hf_batch.logits) + torch.testing.assert_close(sgl_batch.aux_hidden_states, hf_batch.aux_hidden_states) + torch.testing.assert_close(sgl_batch.input_ids, hf_batch.input_ids) + torch.testing.assert_close(sgl_batch.loss_mask, hf_batch.loss_mask) + # The batch carries full logits (projection happens server-side), never the + # precomputed encoding. + assert sgl_batch.target_probs is None and sgl_batch.position_mask is None + + +def test_shift_semantics(): + """logits / input_ids / loss_mask shift left by one; aux stays aligned.""" + model = _FakeCausalLM().eval() + runner = _FakeSGLangRunner(model) + sgl = SGLangEagle3TargetModel(runner, aux_layer_ids=_AUX) + input_ids, attention_mask, loss_mask = _inputs() + raw_logits, raw_aux = runner.forward_eagle3(input_ids, attention_mask) + + batch = sgl.generate_batch(input_ids, attention_mask, loss_mask) + torch.testing.assert_close(batch.input_ids[:, :-1], input_ids[:, 1:]) + assert torch.all(batch.input_ids[:, -1] == 0) + torch.testing.assert_close(batch.logits[:, :-1], raw_logits[:, 1:]) + # aux is position-aligned (not shifted). + torch.testing.assert_close(batch.aux_hidden_states, raw_aux) + + +def test_default_aux_layer_ids_applied_and_forwarded(): + """When aux_layer_ids is None the shared default recipe is used and set on the runner.""" + + class _Deep(_FakeCausalLM): + def __init__(self): + super().__init__() + self.config = type("Cfg", (), {"num_hidden_layers": 32, "hidden_size": _HIDDEN, "vocab_size": _VOCAB}) + + runner = _FakeSGLangRunner(_Deep()) + sgl = SGLangEagle3TargetModel(runner, aux_layer_ids=None) + assert sgl.aux_layer_ids == default_eagle3_aux_layer_ids(32) == [1, 15, 28] + assert runner.aux_layer_ids == sgl.aux_layer_ids + + +def test_invalid_aux_layer_ids_rejected(): + runner = _FakeSGLangRunner(_FakeCausalLM()) + with pytest.raises(ValueError, match="exactly 3"): + SGLangEagle3TargetModel(runner, aux_layer_ids=[0, 1]) + with pytest.raises(ValueError, match="out of bounds"): + SGLangEagle3TargetModel(runner, aux_layer_ids=[0, 1, 99]) + + +def test_default_recipe_raises_on_shallow_target(): + # num_layers=4 -> [1, 1, 0] has a duplicate, so the default recipe must raise. + runner = _FakeSGLangRunner(_FakeCausalLM()) + with pytest.raises(ValueError, match="too shallow"): + SGLangEagle3TargetModel(runner, aux_layer_ids=None) + + +def test_input_embeddings_and_close(): + model = _FakeCausalLM().eval() + runner = _FakeSGLangRunner(model) + sgl = SGLangEagle3TargetModel(runner, aux_layer_ids=_AUX) + torch.testing.assert_close(sgl.get_input_embeddings().weight, model.get_input_embeddings().weight) + sgl.close() + assert runner.closed + + +def test_shared_validate_helper_matches_hf(): + """The extracted helpers back both backends, so HF behavior is preserved.""" + model = _FakeCausalLM().eval() + hf = HFEagle3TargetModel(model, aux_layer_ids=_AUX) + assert hf.aux_layer_ids == validate_eagle3_aux_layer_ids(_AUX, _LAYERS) == _AUX + + +def test_serve_target_engine_routing(monkeypatch): + from nemo_automodel.components.speculative import serve_target + + calls = {} + + def _record(engine): + def _builder(*_a, **_k): + calls["engine"] = engine + return _builder + + monkeypatch.setattr(serve_target, "_build_hf_target", _record("hf")) + monkeypatch.setattr(serve_target, "_build_sglang_target", _record("sglang")) + monkeypatch.setattr(serve_target, "TargetModelServer", lambda *a, **k: object()) + monkeypatch.setattr(serve_target, "serve", lambda *a, **k: None) + + serve_target.main(["--target", "x", "--engine", "sglang", "--tp-size", "2"]) + assert calls["engine"] == "sglang" + + serve_target.main(["--target", "x"]) + assert calls["engine"] == "hf" + + +def test_serve_target_arg_defaults(): + from nemo_automodel.components.speculative import serve_target + + args = serve_target._parse_args(["--target", "x"]) + assert args.engine == "hf" and args.tp_size == 1 + args = serve_target._parse_args(["--target", "x", "--engine", "sglang", "--tp-size", "4"]) + assert args.engine == "sglang" and args.tp_size == 4 + + +def test_build_sglang_target_delegates(monkeypatch): + """``_build_sglang_target`` forwards args to the backend's from_pretrained.""" + from nemo_automodel.components.speculative import serve_target + from nemo_automodel.components.speculative.eagle import sglang_target + + captured = {} + + def _fake_from_pretrained(model_path, **kwargs): + captured["model_path"] = model_path + captured.update(kwargs) + return "wrapper" + + monkeypatch.setattr(sglang_target.SGLangEagle3TargetModel, "from_pretrained", staticmethod(_fake_from_pretrained)) + args = serve_target._parse_args(["--target", "org/m", "--engine", "sglang", "--tp-size", "2"]) + result = serve_target._build_sglang_target(args, torch.device("cpu"), torch.float32) + assert result == "wrapper" + assert captured["model_path"] == "org/m" and captured["tp_size"] == 2 + + +# ── SGLangTargetRunner: surface that does not need SGLang (the forward itself +# needs a GPU and is covered on the server) ────────────────────────────── + + +class _FakeModelRunner: + """Stands in for SGLang's ModelRunner for the non-forward runner methods.""" + + def __init__(self): + self.model = _FakeCausalLM().eval() + self.model.set_eagle3_layers_to_capture = lambda ids: setattr(self, "captured_layers", list(ids)) + self.req_to_token_pool = types.SimpleNamespace(cleared=False) + self.req_to_token_pool.clear = lambda: setattr(self.req_to_token_pool, "cleared", True) + self.token_to_kv_pool_allocator = types.SimpleNamespace(cleared=False) + self.token_to_kv_pool_allocator.clear = lambda: setattr(self.token_to_kv_pool_allocator, "cleared", True) + + +def test_runner_surface_without_sglang(): + from nemo_automodel.components.speculative.eagle.sglang_runner import SGLangTargetRunner + + mr = _FakeModelRunner() + runner = SGLangTargetRunner(mr) + assert runner.model is mr.model + + runner.set_aux_layers([0, 1, 3]) + assert mr.captured_layers == [0, 1, 3] + torch.testing.assert_close(runner.input_embedding_weight(), mr.model.get_input_embeddings().weight) + + runner.close() + assert mr.req_to_token_pool.cleared and mr.token_to_kv_pool_allocator.cleared + runner.close() # idempotent after the runner is released + + +def test_runner_forward_stacks_per_row(monkeypatch): + """forward_eagle3 stacks the per-row extend outputs into batched tensors.""" + from nemo_automodel.components.speculative.eagle.sglang_runner import SGLangTargetRunner + + runner = SGLangTargetRunner(_FakeModelRunner()) + rows_logits = [torch.randn(5, _VOCAB), torch.randn(5, _VOCAB)] + rows_aux = [torch.randn(5, 3 * _HIDDEN), torch.randn(5, 3 * _HIDDEN)] + monkeypatch.setattr(runner, "_extend", lambda input_ids: (rows_logits, rows_aux)) + + logits, aux = runner.forward_eagle3(torch.zeros(2, 5, dtype=torch.long), torch.ones(2, 5)) + assert logits.shape == (2, 5, _VOCAB) and aux.shape == (2, 5, 3 * _HIDDEN) + torch.testing.assert_close(logits, torch.stack(rows_logits)) + torch.testing.assert_close(aux, torch.stack(rows_aux)) From 75ee08504c979d0cb23ca3052822e8c2a48aaca5 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 16:57:52 +0800 Subject: [PATCH 02/12] refactor(speculative): extract engine-agnostic TargetRunner contract Generalize the SGLang-specific runner protocol and target backend into an engine-agnostic seam so a second engine (vLLM) can plug in without touching the trainer, the remote server, or the supervision contract: - new target_runner.py: TargetRunner protocol + RunnerEagle3TargetModel, which owns the shift / aux-concatenation contract for any runner; - sglang_target.py: SGLangEagle3TargetModel now only adds SGLang construction on top of RunnerEagle3TargetModel (SGLangRunnerProtocol kept as a backwards-compatible alias of TargetRunner); - sglang_runner.py unchanged (still the GPU/SGLang forward). Behavior-preserving: the SGLang supervision is byte-for-byte identical, all existing CPU contract tests pass, plus a test locking that the backend is engine-agnostic. Signed-off-by: khazic --- .../components/speculative/eagle/__init__.py | 3 + .../speculative/eagle/sglang_runner.py | 4 +- .../speculative/eagle/sglang_target.py | 153 ++++------------- .../speculative/eagle/target_runner.py | 156 ++++++++++++++++++ .../speculative/test_eagle3_sglang.py | 22 +++ 5 files changed, 211 insertions(+), 127 deletions(-) create mode 100644 nemo_automodel/components/speculative/eagle/target_runner.py diff --git a/nemo_automodel/components/speculative/eagle/__init__.py b/nemo_automodel/components/speculative/eagle/__init__.py index 7288856a5f..7f0c21972c 100644 --- a/nemo_automodel/components/speculative/eagle/__init__.py +++ b/nemo_automodel/components/speculative/eagle/__init__.py @@ -38,6 +38,7 @@ resolve_eagle3_draft_spec, ) from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel +from nemo_automodel.components.speculative.eagle.target_runner import RunnerEagle3TargetModel, TargetRunner from nemo_automodel.components.speculative.eagle.target_v12 import HFEagleTargetModel __all__ = [ @@ -45,6 +46,8 @@ "Eagle3TrainerModule", "PEagleTrainerModule", "Eagle3TargetBackend", + "RunnerEagle3TargetModel", + "TargetRunner", "HFEagleTargetModel", "HFEagle3TargetModel", "LlamaEagleDraftModel", diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py index 5ff46f7552..c2e7c91ca4 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_runner.py +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -147,8 +147,8 @@ def __init__(self, logits: torch.Tensor, aux_hidden_states: torch.Tensor): class SGLangTargetRunner: """Standalone SGLang ModelRunner that returns EAGLE-3 supervision tensors. - Built via :meth:`build`; consumed through the - :class:`~nemo_automodel.components.speculative.eagle.sglang_target.SGLangRunnerProtocol` + Built via :meth:`build`; consumed through the engine-agnostic + :class:`~nemo_automodel.components.speculative.eagle.target_runner.TargetRunner` surface (``model`` / ``set_aux_layers`` / ``forward_eagle3`` / ``input_embedding_weight``). """ diff --git a/nemo_automodel/components/speculative/eagle/sglang_target.py b/nemo_automodel/components/speculative/eagle/sglang_target.py index dc9c34d75b..76dcbc67bc 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_target.py +++ b/nemo_automodel/components/speculative/eagle/sglang_target.py @@ -12,145 +12,48 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SGLang-backed EAGLE-3 target model. - -A third :class:`Eagle3TargetBackend` implementation: it runs the frozen target -through SGLang instead of HuggingFace, capturing the same EAGLE-3 supervision -(three auxiliary hidden states plus full-vocab logits) the co-located backend -produces. SGLang is the fastest serving path for mainstream architectures, so a -remote target server (:mod:`nemo_automodel.components.speculative.serve_target`) -can hold the target on dedicated GPUs while the draft trains elsewhere. - -The class is split into two layers so the contract is unit-testable without a -GPU or SGLang installed: - -- :class:`SGLangEagle3TargetModel` (this file) owns the *contract*: it assembles - an :class:`Eagle3TargetBatch` whose shift / aux-concatenation semantics are - identical to :class:`HFEagle3TargetModel`, so a SGLang run is numerically - equivalent to the co-located one. It depends only on a small runner surface - (``forward_eagle3`` / ``input_embedding_weight`` / ``set_aux_layers`` and a - ``model`` exposing ``.config`` + ``.parameters()``), which tests fake. -- :class:`~nemo_automodel.components.speculative.eagle.sglang_runner.SGLangTargetRunner` - owns the SGLang-internal forward and is lazily imported only by - :meth:`SGLangEagle3TargetModel.from_pretrained`, so importing this module - never pulls in SGLang. +"""SGLang adapter for the EAGLE-3 target backend. + +A thin engine adapter on top of the engine-agnostic contract in +:mod:`nemo_automodel.components.speculative.eagle.target_runner`: it builds a +SGLang runner and wraps it in :class:`RunnerEagle3TargetModel`. SGLang is the +fastest serving path for mainstream architectures, so a remote target server +(:mod:`nemo_automodel.components.speculative.serve_target`) can hold the target +on dedicated GPUs while the draft trains elsewhere. + +All supervision-contract logic (shift / aux-concatenation semantics, aux-layer +defaulting, embedding access) lives in ``target_runner`` and is shared with any +future engine adapter (e.g. vLLM). This module owns only SGLang construction, +and the SGLang-internal forward is isolated further in +:mod:`nemo_automodel.components.speculative.eagle.sglang_runner`, imported lazily +by :meth:`SGLangEagle3TargetModel.from_pretrained` so importing this module never +pulls in SGLang. """ from __future__ import annotations -from types import SimpleNamespace -from typing import TYPE_CHECKING, Optional, Protocol, Sequence +from typing import Optional, Sequence import torch -from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend -from nemo_automodel.components.speculative.eagle.target import ( - Eagle3TargetBatch, - _shift_left_with_zero, - default_eagle3_aux_layer_ids, - validate_eagle3_aux_layer_ids, +from nemo_automodel.components.speculative.eagle.target_runner import ( + RunnerEagle3TargetModel, + TargetRunner, ) -if TYPE_CHECKING: - import torch.nn as nn +#: Backwards-compatible alias. The runner surface is now engine-agnostic and +#: lives in ``target_runner`` as :class:`TargetRunner`; kept here so existing +#: imports of ``SGLangRunnerProtocol`` keep resolving. +SGLangRunnerProtocol = TargetRunner -class SGLangRunnerProtocol(Protocol): - """Minimal surface :class:`SGLangEagle3TargetModel` needs from a runner. +class SGLangEagle3TargetModel(RunnerEagle3TargetModel): + """EAGLE-3 target backend whose runner is SGLang. - Implemented for real by ``SGLangTargetRunner`` (GPU/SGLang) and faked in - unit tests, which is why the backend depends on this protocol rather than on - SGLang directly. + Adds only SGLang construction; the supervision contract is inherited from + :class:`RunnerEagle3TargetModel`. """ - #: Loaded model handle exposing ``.config`` (with ``num_hidden_layers`` / - #: ``hidden_size`` / ``vocab_size``) and ``.parameters()`` for device - #: inference, mirroring what the server reads off ``HFEagle3TargetModel``. - model: "nn.Module" - - def set_aux_layers(self, aux_layer_ids: Sequence[int]) -> None: - """Tell the underlying model which 3 decoder layers to capture.""" - - def forward_eagle3( - self, input_ids: torch.Tensor, attention_mask: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the target once and return ``(logits, aux_hidden_states)``. - - ``logits`` is ``[batch, seq, vocab]`` (full vocab, unshifted) and - ``aux_hidden_states`` is ``[batch, seq, 3 * hidden]`` (the three capture - layers concatenated on the last dim, unshifted). - """ - - def input_embedding_weight(self) -> torch.Tensor: - """Return the target input-embedding weight ``[vocab, hidden]``.""" - - -class SGLangEagle3TargetModel(Eagle3TargetBackend): - """EAGLE-3 target backend that runs the frozen target through SGLang. - - Parameters - ---------- - runner: - A loaded runner implementing :class:`SGLangRunnerProtocol`. - aux_layer_ids: - The three decoder layers to capture (low / mid / high). When ``None`` - the shared EAGLE-3 default recipe is used, matching every other backend. - """ - - def __init__(self, runner: SGLangRunnerProtocol, aux_layer_ids: Optional[Sequence[int]] = None): - self._runner = runner - # Expose ``.model`` so the remote server can read the target's config - # and infer its device the same way it does for the co-located backend. - self.model = runner.model - num_layers = self.model.config.num_hidden_layers - if aux_layer_ids is None: - aux_layer_ids = default_eagle3_aux_layer_ids(num_layers) - self.aux_layer_ids = validate_eagle3_aux_layer_ids(aux_layer_ids, num_layers) - runner.set_aux_layers(self.aux_layer_ids) - - def get_input_embeddings(self) -> SimpleNamespace: - """Return an object exposing ``.weight`` (the target input embeddings). - - Matches the offline-cache / remote path: the draft's - ``copy_embeddings_from_target`` only reads ``.weight``. - """ - return SimpleNamespace(weight=self._runner.input_embedding_weight()) - - @torch.no_grad() - def generate_batch( - self, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - loss_mask: torch.Tensor, - ) -> Eagle3TargetBatch: - """Run the SGLang target and capture aux hidden states plus logits. - - Produces an :class:`Eagle3TargetBatch` byte-for-byte compatible with - :meth:`HFEagle3TargetModel.generate_batch`: the logits, ``input_ids`` - and ``loss_mask`` are shifted left by one (next-token alignment) while - the aux hidden states are kept position-aligned. The draft-vocab - projection happens trainer-side / server-side from ``logits``, exactly - as in the co-located path. - - Note: sequences are assumed right-padded (loss_mask zeros the pad). With - causal attention, trailing pad tokens do not affect earlier positions, - so the captured supervision matches a masked HuggingFace forward. - """ - logits, aux_hidden_states = self._runner.forward_eagle3(input_ids, attention_mask) - return Eagle3TargetBatch( - aux_hidden_states=aux_hidden_states, - logits=_shift_left_with_zero(logits), - input_ids=_shift_left_with_zero(input_ids), - attention_mask=attention_mask, - loss_mask=_shift_left_with_zero(loss_mask), - ) - - def close(self) -> None: - """Release the SGLang runner (frees GPU memory / engine handles).""" - close = getattr(self._runner, "close", None) - if close is not None: - close() - @classmethod def from_pretrained( # pragma: no cover - requires GPU + SGLang cls, diff --git a/nemo_automodel/components/speculative/eagle/target_runner.py b/nemo_automodel/components/speculative/eagle/target_runner.py new file mode 100644 index 0000000000..e9072786a7 --- /dev/null +++ b/nemo_automodel/components/speculative/eagle/target_runner.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Engine-agnostic EAGLE-3 target backend built on a pluggable runner. + +The supervision contract (how a target's raw logits / aux hidden states become +an :class:`Eagle3TargetBatch`) is identical no matter which inference engine +runs the frozen target. This module owns that contract: + +- :class:`TargetRunner` is the narrow surface a concrete engine must implement + (``forward_eagle3`` / ``set_aux_layers`` / ``input_embedding_weight`` and a + ``model`` exposing ``.config`` + ``.parameters()``). It is the single seam + between the engine-coupled code and the rest of the stack: SGLang implements + it today (``sglang_runner.SGLangTargetRunner``), and a vLLM runner can + implement the same protocol and drop in without touching this file, the + trainer, or the remote server. +- :class:`RunnerEagle3TargetModel` is the :class:`Eagle3TargetBackend` that + assembles the supervision batch from any :class:`TargetRunner`. Its shift / + aux-concatenation semantics are identical to :class:`HFEagle3TargetModel`, so + a run through any runner is numerically equivalent to the co-located one. + +Keeping the contract here (and importing engines lazily in their own adapter +modules) means importing this module never pulls in SGLang or vLLM, so the +contract stays unit-testable on CPU against a fake runner. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING, Optional, Protocol, Sequence + +import torch + +from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend +from nemo_automodel.components.speculative.eagle.target import ( + Eagle3TargetBatch, + _shift_left_with_zero, + default_eagle3_aux_layer_ids, + validate_eagle3_aux_layer_ids, +) + +if TYPE_CHECKING: + import torch.nn as nn + + +class TargetRunner(Protocol): + """Minimal engine surface :class:`RunnerEagle3TargetModel` depends on. + + Implemented for real by an inference-engine runner (``SGLangTargetRunner`` + today, a vLLM runner later) and faked in unit tests, which is why the + backend depends on this protocol rather than on any engine directly. + """ + + #: Loaded model handle exposing ``.config`` (with ``num_hidden_layers`` / + #: ``hidden_size`` / ``vocab_size``) and ``.parameters()`` for device + #: inference, mirroring what the server reads off ``HFEagle3TargetModel``. + model: "nn.Module" + + def set_aux_layers(self, aux_layer_ids: Sequence[int]) -> None: + """Tell the underlying model which 3 decoder layers to capture.""" + + def forward_eagle3( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the target once and return ``(logits, aux_hidden_states)``. + + ``logits`` is ``[batch, seq, vocab]`` (full vocab, unshifted) and + ``aux_hidden_states`` is ``[batch, seq, 3 * hidden]`` (the three capture + layers concatenated on the last dim, unshifted). + """ + + def input_embedding_weight(self) -> torch.Tensor: + """Return the target input-embedding weight ``[vocab, hidden]``.""" + + +class RunnerEagle3TargetModel(Eagle3TargetBackend): + """EAGLE-3 target backend that runs the frozen target through a runner. + + Engine-agnostic: it owns only the supervision *contract* and delegates the + actual forward to a :class:`TargetRunner`. Concrete engines subclass this to + add their own ``from_pretrained`` (which builds the runner); everything else + is inherited. + + Parameters + ---------- + runner: + A loaded runner implementing :class:`TargetRunner`. + aux_layer_ids: + The three decoder layers to capture (low / mid / high). When ``None`` + the shared EAGLE-3 default recipe is used, matching every other backend. + """ + + def __init__(self, runner: TargetRunner, aux_layer_ids: Optional[Sequence[int]] = None): + self._runner = runner + # Expose ``.model`` so the remote server can read the target's config + # and infer its device the same way it does for the co-located backend. + self.model = runner.model + num_layers = self.model.config.num_hidden_layers + if aux_layer_ids is None: + aux_layer_ids = default_eagle3_aux_layer_ids(num_layers) + self.aux_layer_ids = validate_eagle3_aux_layer_ids(aux_layer_ids, num_layers) + runner.set_aux_layers(self.aux_layer_ids) + + def get_input_embeddings(self) -> SimpleNamespace: + """Return an object exposing ``.weight`` (the target input embeddings). + + Matches the offline-cache / remote path: the draft's + ``copy_embeddings_from_target`` only reads ``.weight``. + """ + return SimpleNamespace(weight=self._runner.input_embedding_weight()) + + @torch.no_grad() + def generate_batch( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + loss_mask: torch.Tensor, + ) -> Eagle3TargetBatch: + """Run the runner's target and capture aux hidden states plus logits. + + Produces an :class:`Eagle3TargetBatch` byte-for-byte compatible with + :meth:`HFEagle3TargetModel.generate_batch`: the logits, ``input_ids`` + and ``loss_mask`` are shifted left by one (next-token alignment) while + the aux hidden states are kept position-aligned. The draft-vocab + projection happens trainer-side / server-side from ``logits``, exactly + as in the co-located path. + + Note: sequences are assumed right-padded (loss_mask zeros the pad). With + causal attention, trailing pad tokens do not affect earlier positions, + so the captured supervision matches a masked HuggingFace forward. + """ + logits, aux_hidden_states = self._runner.forward_eagle3(input_ids, attention_mask) + return Eagle3TargetBatch( + aux_hidden_states=aux_hidden_states, + logits=_shift_left_with_zero(logits), + input_ids=_shift_left_with_zero(input_ids), + attention_mask=attention_mask, + loss_mask=_shift_left_with_zero(loss_mask), + ) + + def close(self) -> None: + """Release the runner (frees GPU memory / engine handles).""" + close = getattr(self._runner, "close", None) + if close is not None: + close() diff --git a/tests/unit_tests/speculative/test_eagle3_sglang.py b/tests/unit_tests/speculative/test_eagle3_sglang.py index f07e588352..fe3bae8dab 100644 --- a/tests/unit_tests/speculative/test_eagle3_sglang.py +++ b/tests/unit_tests/speculative/test_eagle3_sglang.py @@ -41,6 +41,7 @@ default_eagle3_aux_layer_ids, validate_eagle3_aux_layer_ids, ) +from nemo_automodel.components.speculative.eagle.target_runner import RunnerEagle3TargetModel _VOCAB = 32 _HIDDEN = 16 @@ -134,6 +135,27 @@ def test_generate_batch_matches_colocated_hf(): assert sgl_batch.target_probs is None and sgl_batch.position_mask is None +def test_engine_agnostic_backend_is_shared(): + """The contract layer is engine-agnostic: the same backend works on any runner. + + ``SGLangEagle3TargetModel`` only adds SGLang construction; the supervision + contract lives in ``RunnerEagle3TargetModel`` and a vLLM runner can reuse it + by implementing the same ``TargetRunner`` surface. + """ + assert issubclass(SGLangEagle3TargetModel, RunnerEagle3TargetModel) + + model = _FakeCausalLM().eval() + runner = _FakeSGLangRunner(model) + base = RunnerEagle3TargetModel(runner, aux_layer_ids=_AUX) + sgl = SGLangEagle3TargetModel(_FakeSGLangRunner(model), aux_layer_ids=_AUX) + + input_ids, attention_mask, loss_mask = _inputs() + base_batch = base.generate_batch(input_ids, attention_mask, loss_mask) + sgl_batch = sgl.generate_batch(input_ids, attention_mask, loss_mask) + torch.testing.assert_close(base_batch.logits, sgl_batch.logits) + torch.testing.assert_close(base_batch.aux_hidden_states, sgl_batch.aux_hidden_states) + + def test_shift_semantics(): """logits / input_ids / loss_mask shift left by one; aux stays aligned.""" model = _FakeCausalLM().eval() From 172071387332bc2c4866387f3417ab1a2ffd7fa9 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 18:19:04 +0800 Subject: [PATCH 03/12] fix(speculative): pass moe_ep_rank/moe_ep_size to SGLang ModelRunner sglang>=0.5.9 made moe_ep_rank/moe_ep_size required positional args on ModelRunner.__init__; the target runner is single-process with no expert parallelism, so pass (0, 1). Signed-off-by: khazic --- nemo_automodel/components/speculative/eagle/sglang_runner.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py index c2e7c91ca4..6a82f52e9a 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_runner.py +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -220,6 +220,11 @@ def build( # pragma: no cover - requires GPU + SGLang gpu_id=gpu_id, tp_rank=0, tp_size=tp_size, + # No expert parallelism for the target runner: a dense target has no + # experts, and a MoE target is run with plain TP here. sglang>=0.5.9 + # made these required positional args on ModelRunner. + moe_ep_rank=0, + moe_ep_size=1, pp_rank=0, pp_size=1, nccl_port=None, From 1175bddac52ed1abfbcc603c5d954707b54ba0f7 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 18:21:49 +0800 Subject: [PATCH 04/12] fix(speculative): don't double-init model parallel in SGLang runner ModelRunner.init_torch_distributed already calls initialize_model_parallel in sglang>=0.5.9, so calling it ourselves trips 'tensor model parallel group is already initialized'. Bring up only the world process group here and let ModelRunner build the TP group. Signed-off-by: khazic --- .../components/speculative/eagle/sglang_runner.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py index 6a82f52e9a..51cd8cb800 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_runner.py +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -179,7 +179,7 @@ def build( # pragma: no cover - requires GPU + SGLang """ import torch.distributed as dist from sglang.srt.configs.model_config import ModelConfig - from sglang.srt.distributed import init_distributed_environment, initialize_model_parallel + from sglang.srt.distributed import init_distributed_environment from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.server_args import ServerArgs @@ -199,8 +199,10 @@ def build( # pragma: no cover - requires GPU + SGLang ) gpu_id = torch.cuda.current_device() - # Standalone server: SGLang manages its own single-rank world rather than - # reusing a trainer process group (the SpecForge embedded case). + # Standalone server: bring up only the world process group here, then let + # ModelRunner.init_torch_distributed build the tensor-parallel group. sglang + # >=0.5.9 always calls initialize_model_parallel inside ModelRunner, so doing + # it here too trips "tensor model parallel group is already initialized". if not dist.is_initialized(): os.environ.setdefault("MASTER_ADDR", "127.0.0.1") os.environ.setdefault("MASTER_PORT", str(int(server_args.port or 8000) + 200)) @@ -211,7 +213,6 @@ def build( # pragma: no cover - requires GPU + SGLang local_rank=gpu_id, distributed_init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}", ) - initialize_model_parallel(tensor_model_parallel_size=tp_size) model_config = ModelConfig.from_server_args(server_args) model_runner = ModelRunner( From 85c0e98b59cef8e2d8e6b1c4b7dbfcb49f647d0f Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 18:44:29 +0800 Subject: [PATCH 05/12] fix(speculative): lazy-import HF AutoModel in serve_target hf builder The sglang engine never loads the HF AutoModel, so importing NeMoAutoModelForCausalLM at module top forced the sglang target server to pull in Automodel's full model stack. Move it into _build_hf_target, matching the lazy sglang import in _build_sglang_target, so the sglang server runs in a minimal sglang-only environment. Signed-off-by: khazic --- nemo_automodel/components/speculative/serve_target.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nemo_automodel/components/speculative/serve_target.py b/nemo_automodel/components/speculative/serve_target.py index 5b5bbbcb2b..bb39b02906 100644 --- a/nemo_automodel/components/speculative/serve_target.py +++ b/nemo_automodel/components/speculative/serve_target.py @@ -53,7 +53,6 @@ import torch -from nemo_automodel._transformers.auto_model import NeMoAutoModelForCausalLM from nemo_automodel.components.speculative.eagle.remote.server import TargetModelServer, serve from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel @@ -62,6 +61,11 @@ def _build_hf_target(args, device: torch.device, dtype: torch.dtype) -> HFEagle3TargetModel: """Load the target under HuggingFace and wrap it with aux-layer hooks.""" + # Imported here (not at module top) so the sglang engine, which never loads + # the HF AutoModel, can run in a minimal environment without Automodel's full + # model stack -- mirroring the lazy sglang import in ``_build_sglang_target``. + from nemo_automodel._transformers.auto_model import NeMoAutoModelForCausalLM + target_model = NeMoAutoModelForCausalLM.from_pretrained( args.target, torch_dtype=dtype, From 0834d70022e89bd4040649444a0011c71bc31cfb Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 19:07:44 +0800 Subject: [PATCH 06/12] fix(speculative): skip server NCCL init when the client can't join A client without sglang (the disaggregated case: sglang target server + sglang-free training client) cannot join the NCCL group, but _init_nccl still POSTed /init_nccl and let the server block on the rendezvous until its 120s timeout before both fell back to wire. Gate the request on a local nccl_transport_available() check so an sglang-free client goes straight to wire and never stalls the server. Also fix the serve_target test to patch the now lazily-imported NeMoAutoModelForCausalLM at its source. Signed-off-by: khazic --- .../speculative/eagle/remote/client.py | 13 +++++++++++- .../speculative/eagle/remote/transport.py | 13 ++++++++++++ .../test_eagle3_remote_coverage.py | 21 ++++++++++++++++++- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/nemo_automodel/components/speculative/eagle/remote/client.py b/nemo_automodel/components/speculative/eagle/remote/client.py index ae0438304e..b5ce179f8a 100644 --- a/nemo_automodel/components/speculative/eagle/remote/client.py +++ b/nemo_automodel/components/speculative/eagle/remote/client.py @@ -41,7 +41,7 @@ from nemo_automodel.components.speculative.eagle.backend import Eagle3TargetBackend from nemo_automodel.components.speculative.eagle.remote import protocol, wire -from nemo_automodel.components.speculative.eagle.remote.transport import NCCLTransport +from nemo_automodel.components.speculative.eagle.remote.transport import NCCLTransport, nccl_transport_available from nemo_automodel.components.speculative.eagle.target import Eagle3TargetBatch logger = logging.getLogger(__name__) @@ -103,6 +103,17 @@ def _init_nccl(self) -> bool: if self._nccl_attempted: return False self._nccl_attempted = True + + # Only ask the server to bring up NCCL if this process can actually + # join the group. A client without sglang (the common disaggregated + # case: sglang target server + sglang-free training client) cannot, + # and contacting the server anyway leaves it blocked on a rendezvous + # that never completes until it times out. Fall back to wire instead. + if not nccl_transport_available(): + logger.info("NCCL unavailable in this process (sglang not importable); using wire format") + self._nccl = None + return False + port = self._nccl_port() self._nccl = NCCLTransport(nccl_port=port, host=self._host(), is_server=False) diff --git a/nemo_automodel/components/speculative/eagle/remote/transport.py b/nemo_automodel/components/speculative/eagle/remote/transport.py index 8570939e8a..4e47457c0a 100644 --- a/nemo_automodel/components/speculative/eagle/remote/transport.py +++ b/nemo_automodel/components/speculative/eagle/remote/transport.py @@ -48,6 +48,19 @@ _HAS_SGLANG_PG, _init_custom_process_group = safe_import_from("sglang.srt.utils.common", "init_custom_process_group") + +def nccl_transport_available() -> bool: + """Whether GPU-direct NCCL transfer is usable in this process. + + NCCL transfer relies on sglang's ``init_custom_process_group``; without + sglang (e.g. a training client that intentionally keeps sglang out of its + env) NCCL cannot work and callers should use the wire fallback. Checking + this *before* asking the server to set up its NCCL side avoids leaving the + server blocked on a rendezvous the client can never complete. + """ + return _HAS_SGLANG_PG + + # dtypes NCCL P2P does not support; transmitted as raw uint8 views. _NCCL_UNSUPPORTED_DTYPES = {torch.int16, torch.int8, torch.bool} _ELEMENT_SIZE = {torch.int16: 2, torch.int8: 1, torch.bool: 1} diff --git a/tests/unit_tests/speculative/test_eagle3_remote_coverage.py b/tests/unit_tests/speculative/test_eagle3_remote_coverage.py index aa416feb79..10269f7a18 100644 --- a/tests/unit_tests/speculative/test_eagle3_remote_coverage.py +++ b/tests/unit_tests/speculative/test_eagle3_remote_coverage.py @@ -99,8 +99,11 @@ def test_parse_args_explicit(): ) def test_main_wires_server(monkeypatch, argv, expected_nccl_port): fake_model = mock.MagicMock() + # ``NeMoAutoModelForCausalLM`` is imported lazily inside ``_build_hf_target`` + # (so the sglang engine needs no HF stack), so patch it at its source module. monkeypatch.setattr( - serve_target.NeMoAutoModelForCausalLM, "from_pretrained", mock.MagicMock(return_value=fake_model) + "nemo_automodel._transformers.auto_model.NeMoAutoModelForCausalLM.from_pretrained", + mock.MagicMock(return_value=fake_model), ) monkeypatch.setattr(serve_target, "HFEagle3TargetModel", mock.MagicMock(return_value="wrapper")) captured = {} @@ -144,6 +147,22 @@ def test_nccl_port_nondigit_falls_back(monkeypatch): assert client._nccl_port() == 8100 # 8000 default + 100 +def test_init_nccl_skips_server_when_locally_unavailable(monkeypatch): + """No local NCCL support (e.g. an sglang-free client) must not ask the server + to init NCCL -- otherwise the server blocks on a rendezvous the client can + never complete. The client falls back to wire immediately instead.""" + from nemo_automodel.components.speculative.eagle.remote import client as client_mod + + monkeypatch.setattr(client_mod, "nccl_transport_available", lambda: False) + client = _ServerClient("http://h:8001", timeout=1, max_retries=0) + sent = [] + monkeypatch.setattr(client, "request", lambda endpoint, *_a, **_k: sent.append(endpoint) or b"{}") + + assert client._init_nccl() is False + assert protocol.EP_INIT_NCCL not in sent # server was never contacted + assert client._nccl is None + + # ── _ServerClient.request: retry / failure ─────────────────────────────── From b1f6e19e2c5a1fde784c79f5968807f697776dee Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 8 Jun 2026 19:31:57 +0800 Subject: [PATCH 07/12] chore(speculative): simplify SGLang target backend Drop the dead SGLangRunnerProtocol alias (the protocol and alias were both introduced on this branch, so there is no prior name to keep resolving), and cache the constant teacher-forcing SamplingParams on the runner instead of rebuilding it every extend. Signed-off-by: khazic --- .../components/speculative/eagle/sglang_runner.py | 5 ++++- .../components/speculative/eagle/sglang_target.py | 10 +--------- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py index 51cd8cb800..81072b6cfb 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_runner.py +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -155,6 +155,7 @@ class SGLangTargetRunner: def __init__(self, model_runner): self._model_runner = model_runner + self._sampling_params = None # constant teacher-forcing params, built once on first extend @property def model(self): @@ -268,7 +269,9 @@ def _extend(self, input_ids: torch.Tensor) -> tuple[list, list]: # pragma: no c from sglang.srt.speculative.spec_info import SpeculativeAlgorithm runner = self._model_runner - sampling_params = SamplingParams(temperature=0, max_new_tokens=1, top_k=1) + if self._sampling_params is None: + self._sampling_params = SamplingParams(temperature=0, max_new_tokens=1, top_k=1) + sampling_params = self._sampling_params rows = torch.split(input_ids, 1, dim=0) reqs, input_lens = [], [] for idx, row in enumerate(rows): diff --git a/nemo_automodel/components/speculative/eagle/sglang_target.py b/nemo_automodel/components/speculative/eagle/sglang_target.py index 76dcbc67bc..865d44ce29 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_target.py +++ b/nemo_automodel/components/speculative/eagle/sglang_target.py @@ -36,15 +36,7 @@ import torch -from nemo_automodel.components.speculative.eagle.target_runner import ( - RunnerEagle3TargetModel, - TargetRunner, -) - -#: Backwards-compatible alias. The runner surface is now engine-agnostic and -#: lives in ``target_runner`` as :class:`TargetRunner`; kept here so existing -#: imports of ``SGLangRunnerProtocol`` keep resolving. -SGLangRunnerProtocol = TargetRunner +from nemo_automodel.components.speculative.eagle.target_runner import RunnerEagle3TargetModel class SGLangEagle3TargetModel(RunnerEagle3TargetModel): From 7fe2dea42fefea7f7411085444f92fac505a2fa9 Mon Sep 17 00:00:00 2001 From: khazic Date: Wed, 10 Jun 2026 11:37:38 +0800 Subject: [PATCH 08/12] feat(speculative): co-located SGLang target backend for EAGLE-3 training Add target_model_backend: sglang to the EAGLE-3 recipe: the frozen target runs through SGLang's ModelRunner on the training GPU (single-process only; SGLang's parallel state must own every rank, so multi-GPU runs keep using serve_target --engine sglang + the remote backend). SGLang's memory pool defaults to half the GPU here so the draft trains in the remainder, tunable via recipe_args.sglang_args. Also fix the ServerArgs dtype handling (SGLang compares dtype against string literals, so torch.dtype objects silently missed every branch; addresses the review comment), guard runner construction against a process-group/tp_size mismatch with a clear error, and add a GPU smoke script that validates the SGLang forward against the HF backend on the server, including the pre-initialized process group of the co-located path. Signed-off-by: khazic --- .../eagle3/llama_eagle3_sglang.yaml | 59 +++++++ .../speculative/eagle/sglang_runner.py | 31 +++- nemo_automodel/recipes/llm/train_eagle3.py | 47 +++++- scripts/smoke_sglang_target.py | 149 ++++++++++++++++++ .../recipes/llm/test_eagle3_sglang_backend.py | 140 ++++++++++++++++ .../speculative/test_eagle3_sglang.py | 22 ++- 6 files changed, 442 insertions(+), 6 deletions(-) create mode 100644 examples/speculative/eagle3/llama_eagle3_sglang.yaml create mode 100644 scripts/smoke_sglang_target.py create mode 100644 tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py diff --git a/examples/speculative/eagle3/llama_eagle3_sglang.yaml b/examples/speculative/eagle3/llama_eagle3_sglang.yaml new file mode 100644 index 0000000000..06790037f4 --- /dev/null +++ b/examples/speculative/eagle3/llama_eagle3_sglang.yaml @@ -0,0 +1,59 @@ +recipe: TrainEagle3Recipe + +# Co-located EAGLE-3 training with the frozen target served through SGLang. +# +# Identical to the co-located MVP config (llama_eagle3_mvp.yaml) except the +# target forward runs through SGLang's ModelRunner instead of the HuggingFace +# eager forward, which is substantially faster for mainstream architectures. +# SGLang carves its weight + KV pool out of the training GPU up front +# (``sglang_args.mem_fraction_static``); the draft trains in the remainder. +# +# Single-process only: SGLang's parallel state must own every rank of the +# process group. For multi-GPU runs serve the target separately +# (``serve_target --engine sglang``) and use llama_eagle3_remote.yaml. + +dist_env: + backend: nccl + timeout_minutes: 30 + +recipe_args: + target_model_name_or_path: meta-llama/Llama-3.2-1B + + # --- sglang backend --- + # ``colocated`` (default) runs the target through HF on this GPU; ``sglang`` + # runs it through SGLang on this GPU; ``remote`` talks to serve_target. + target_model_backend: sglang + # Extra SGLang ServerArgs. ``mem_fraction_static`` is the fraction of GPU + # memory SGLang reserves for target weights + KV pool (default 0.5 here; + # raise it for big targets, lower it if draft training runs out of memory). + sglang_args: + mem_fraction_static: 0.5 + + train_data_path: /path/to/train.jsonl + val_data_path: null + train_split: null + val_split: null + output_dir: ./outputs/eagle3_llama_sglang + seq_length: 1024 + micro_batch_size: 1 + grad_accumulation_steps: 1 + num_workers: 0 + num_epochs: 1 + ttt_steps: 4 + draft_vocab_size: 8192 + freeze_embeddings: true + trust_remote_code: false + shuffle_seed: 42 + log_every_steps: 10 + max_grad_norm: 1.0 + +optimizer: + lr: 1.0e-4 + betas: [0.9, 0.95] + weight_decay: 0.0 + +checkpoint: + enabled: true + checkpoint_dir: ./outputs/eagle3_llama_sglang/checkpoints + model_save_format: safetensors + save_consolidated: true diff --git a/nemo_automodel/components/speculative/eagle/sglang_runner.py b/nemo_automodel/components/speculative/eagle/sglang_runner.py index 81072b6cfb..27859f0a18 100644 --- a/nemo_automodel/components/speculative/eagle/sglang_runner.py +++ b/nemo_automodel/components/speculative/eagle/sglang_runner.py @@ -52,6 +52,26 @@ logger = logging.getLogger(__name__) +_SGLANG_DTYPE_STRINGS = { + torch.float32: "float32", + torch.float16: "float16", + torch.bfloat16: "bfloat16", +} + + +def sglang_dtype_str(dtype: Optional[torch.dtype]) -> str: + """Map a torch dtype to the string form SGLang's ``ServerArgs.dtype`` expects. + + SGLang compares ``ServerArgs.dtype`` against string literals (``"auto"``, + ``"bfloat16"``, ...), so passing a raw ``torch.dtype`` silently misses every + branch. ``None`` means "let SGLang pick" (``"auto"``). + """ + if dtype is None: + return "auto" + if dtype not in _SGLANG_DTYPE_STRINGS: + raise ValueError(f"Unsupported SGLang target dtype {dtype}; expected one of {list(_SGLANG_DTYPE_STRINGS)}.") + return _SGLANG_DTYPE_STRINGS[dtype] + def _wrap_logits_processors_for_eagle3(model) -> None: # pragma: no cover - requires GPU + SGLang """Replace every SGLang ``LogitsProcessor`` in ``model`` with an EAGLE-3 wrapper. @@ -186,11 +206,20 @@ def build( # pragma: no cover - requires GPU + SGLang if not torch.cuda.is_available(): raise RuntimeError("SGLangTargetRunner requires CUDA; run it on a GPU server, not the editing host.") + # SGLang's global parallel state must own every rank of an existing + # process group (initialize_model_parallel raises unless world_size == + # tp_size * pp_size). Catch the mismatch here with a clear error instead + # of failing deep inside ModelRunner. + if dist.is_initialized() and dist.get_world_size() != tp_size: + raise RuntimeError( + f"SGLangTargetRunner.build inside an initialized process group requires " + f"world_size == tp_size, got world_size={dist.get_world_size()} and tp_size={tp_size}." + ) server_args = ServerArgs( model_path=model_path, trust_remote_code=trust_remote_code, - dtype=dtype if dtype is not None else "auto", + dtype=sglang_dtype_str(dtype), enable_return_hidden_states=True, disable_cuda_graph=True, # extend-only forward; CUDA graphs add no benefit here disable_radix_cache=True, diff --git a/nemo_automodel/recipes/llm/train_eagle3.py b/nemo_automodel/recipes/llm/train_eagle3.py index 06e24f6f05..42ef64ce7e 100644 --- a/nemo_automodel/recipes/llm/train_eagle3.py +++ b/nemo_automodel/recipes/llm/train_eagle3.py @@ -396,7 +396,10 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): """ # ``target_model_backend`` selects where the frozen target runs: # - ``colocated`` (default): load the target on this GPU and capture - # supervision in-process. + # supervision in-process via the HuggingFace forward. + # - ``sglang``: like ``colocated`` but the in-process forward runs + # through SGLang's ModelRunner, which is substantially faster than + # the HF eager forward for mainstream architectures. # - ``remote``: the target runs as a standalone server (see # ``serve_target``); this process only holds the draft and pulls # precomputed supervision over HTTP + NCCL. No target weights are @@ -404,10 +407,12 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): backend = recipe_cfg.get("target_model_backend", "colocated") if backend == "remote": self._setup_remote_target(recipe_cfg) + elif backend == "sglang": + self._setup_sglang_target(recipe_cfg, target_path) elif backend == "colocated": self._setup_colocated_target(recipe_cfg, target_path) else: - raise ValueError(f"Unknown target_model_backend={backend!r}; expected 'colocated' or 'remote'.") + raise ValueError(f"Unknown target_model_backend={backend!r}; expected 'colocated', 'sglang', or 'remote'.") self.train_dataloader = build_eagle3_dataloader( data_path=recipe_cfg.train_data_path, @@ -496,6 +501,44 @@ def _setup_colocated_target(self, recipe_cfg, target_path): aux_layer_ids=recipe_cfg.get("aux_layer_ids", None), ) + def _setup_sglang_target(self, recipe_cfg, target_path): + """Co-located SGLang target: serve the frozen target through SGLang on this GPU. + + Same supervision contract as ``colocated`` (full-vocab logits shipped to + the trainer, draft-vocab projection trainer-side), but the target forward + runs through SGLang's ModelRunner. SGLang carves its weight + KV pool out + of this GPU up front (``mem_fraction_static``), and the draft trains in + the remainder. + """ + if self.device.type != "cuda": + raise ValueError("target_model_backend='sglang' requires CUDA; use 'colocated' for CPU runs.") + if self.dist_env.world_size > 1: + # SGLang's global parallel state requires world_size == tp_size * pp_size, + # so per-rank tp=1 runners cannot share a multi-rank training process + # group. Multi-GPU runs split target and draft onto separate processes + # instead: ``serve_target --engine sglang`` + ``target_model_backend='remote'``. + raise ValueError( + "target_model_backend='sglang' supports single-process training only; " + "for multi-GPU runs serve the target separately (serve_target --engine sglang) " + "and set target_model_backend='remote'." + ) + from nemo_automodel.components.speculative.eagle.sglang_target import SGLangEagle3TargetModel + + sglang_args = recipe_cfg.get("sglang_args", None) or {} + sglang_kwargs = sglang_args.to_dict() if hasattr(sglang_args, "to_dict") else dict(sglang_args) + # SGLang's ServerArgs default (~0.88 of GPU memory) would starve the + # draft's optimizer states and activations; default to half the GPU and + # let ``recipe_args.sglang_args.mem_fraction_static`` override. + sglang_kwargs.setdefault("mem_fraction_static", 0.5) + self.target_model = None + self.target_wrapper = SGLangEagle3TargetModel.from_pretrained( + target_path, + aux_layer_ids=recipe_cfg.get("aux_layer_ids", None), + dtype=self.compute_dtype, + trust_remote_code=recipe_cfg.get("trust_remote_code", False), + **sglang_kwargs, + ) + def _setup_remote_target(self, recipe_cfg): """Connect to one or more remote target servers (no target loaded here).""" urls = recipe_cfg.get("remote_urls", None) diff --git a/scripts/smoke_sglang_target.py b/scripts/smoke_sglang_target.py new file mode 100644 index 0000000000..0d4525223e --- /dev/null +++ b/scripts/smoke_sglang_target.py @@ -0,0 +1,149 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU smoke test for the SGLang EAGLE-3 target backend (not part of CI). + +Run on the training server to validate that sglang 0.5.9's internal API really +lines up with the ported forward in ``sglang_runner.py`` (ModelRunner build, +``set_eagle3_layers_to_capture``, ``CaptureHiddenMode.FULL``, the private +``LogitsProcessor`` helpers). Two stages: + +1. SGLang only: build the target, run ``generate_batch`` on a tiny batch, and + assert shapes / dtype / finiteness of aux hidden states and logits. +2. ``--compare-hf``: also run the co-located HuggingFace backend on the SAME + inputs and report how closely they agree (next-token argmax match rate over + the loss positions + mean cosine similarity of the aux hidden states). The + contract claims numerical equivalence, so these should be high; small gaps + from kernel/dtype differences are expected, large gaps are a red flag. + +``--init-dist`` initializes ``torch.distributed`` (single-process NCCL group) +BEFORE building SGLang, mirroring ``target_model_backend: sglang`` inside the +training recipe, where torchrun owns the process group and SGLang must attach +to it instead of creating its own. Run the smoke once without and once with +this flag to validate both the ``serve_target`` and the co-located paths. + +Example (single GPU): + HF_HOME=/llm-align/liuchonghan/hf_cache CUDA_VISIBLE_DEVICES=0 \ + python scripts/smoke_sglang_target.py --target /llm-align/open_models/Qwen3/Qwen3-4B --compare-hf --init-dist +""" + +from __future__ import annotations + +import argparse +import os + +import torch + + +def _parse_args(): + p = argparse.ArgumentParser(description="SGLang EAGLE-3 target backend smoke test.") + p.add_argument("--target", default="/llm-align/open_models/Qwen3/Qwen3-4B", help="Target model path.") + p.add_argument("--batch", type=int, default=2) + p.add_argument("--seq", type=int, default=16) + p.add_argument("--tp-size", type=int, default=1) + p.add_argument("--trust-remote-code", action="store_true", default=True) + p.add_argument( + "--mem-fraction-static", + type=float, + default=0.4, + help="SGLang KV pool fraction (lower leaves room for the HF compare).", + ) + p.add_argument("--compare-hf", action="store_true", help="Also run the HF co-located backend and compare.") + p.add_argument( + "--init-dist", action="store_true", help="Pre-initialize torch.distributed like the training recipe does." + ) + return p.parse_args() + + +def main(): + """Build the SGLang target, validate the supervision batch, optionally compare against HF.""" + args = _parse_args() + assert torch.cuda.is_available(), "smoke test needs a GPU" + device = torch.device("cuda") + + if args.init_dist: + import torch.distributed as dist + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29511") + dist.init_process_group( + backend="nccl", world_size=1, rank=0, device_id=torch.device("cuda", torch.cuda.current_device()) + ) + print("[smoke] torch.distributed pre-initialized (world_size=1), as in the training recipe") + + from nemo_automodel.components.speculative.eagle.sglang_target import SGLangEagle3TargetModel + + print(f"[smoke] building SGLang target from {args.target} (tp_size={args.tp_size}) ...") + sgl = SGLangEagle3TargetModel.from_pretrained( + args.target, + dtype=torch.bfloat16, + tp_size=args.tp_size, + trust_remote_code=args.trust_remote_code, + mem_fraction_static=args.mem_fraction_static, + ) + cfg = sgl.model.config + hidden, vocab = cfg.hidden_size, cfg.vocab_size + print(f"[smoke] aux_layer_ids={sgl.aux_layer_ids} hidden_size={hidden} vocab_size={vocab}") + + torch.manual_seed(0) + input_ids = torch.randint(0, vocab, (args.batch, args.seq), device=device) + attention_mask = torch.ones(args.batch, args.seq, device=device, dtype=torch.long) + loss_mask = torch.ones(args.batch, args.seq, device=device, dtype=torch.long) + + batch = sgl.generate_batch(input_ids, attention_mask, loss_mask) + print(f"[smoke] aux {tuple(batch.aux_hidden_states.shape)} {batch.aux_hidden_states.dtype}") + print(f"[smoke] logits {tuple(batch.logits.shape)} {batch.logits.dtype}") + + assert batch.aux_hidden_states.shape == (args.batch, args.seq, 3 * hidden), "aux shape mismatch" + assert batch.logits.shape == (args.batch, args.seq, vocab), "logits shape mismatch" + assert batch.target_probs is None and batch.position_mask is None, "expected full-logits encoding" + assert torch.isfinite(batch.aux_hidden_states).all(), "aux has NaN/Inf" + assert torch.isfinite(batch.logits).all(), "logits has NaN/Inf" + print("[smoke] STAGE 1 OK: SGLang target produces well-formed supervision") + + if not args.compare_hf: + print("[smoke] done (pass --compare-hf for the HF equivalence check)") + return + + from nemo_automodel._transformers.auto_model import NeMoAutoModelForCausalLM + from nemo_automodel.components.speculative.eagle.target import HFEagle3TargetModel + + print("[smoke] building HF co-located target for comparison ...") + hf_model = NeMoAutoModelForCausalLM.from_pretrained( + args.target, torch_dtype=torch.bfloat16, trust_remote_code=args.trust_remote_code + ).to(device) + hf_model.requires_grad_(False) + hf = HFEagle3TargetModel(hf_model, aux_layer_ids=sgl.aux_layer_ids) + hf_batch = hf.generate_batch(input_ids, attention_mask, loss_mask) + + # Next-token argmax agreement over the supervised (shifted) positions. + valid = hf_batch.loss_mask.bool() + sgl_top = batch.logits.argmax(-1)[valid] + hf_top = hf_batch.logits.argmax(-1)[valid] + match = (sgl_top == hf_top).float().mean().item() + # Mean cosine similarity of the aux hidden states. + a = batch.aux_hidden_states.float().flatten(0, 1) + b = hf_batch.aux_hidden_states.float().flatten(0, 1) + cos = torch.nn.functional.cosine_similarity(a, b, dim=-1).mean().item() + + print(f"[smoke] logits argmax match rate (loss positions): {match:.4f}") + print(f"[smoke] aux hidden-state mean cosine similarity: {cos:.4f}") + if match >= 0.95 and cos >= 0.99: + print("[smoke] STAGE 2 OK: SGLang and HF backends agree") + else: + print("[smoke] STAGE 2 WARNING: agreement lower than expected; inspect aux-layer capture / shift") + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py new file mode 100644 index 0000000000..230bd8e5b1 --- /dev/null +++ b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for the co-located SGLang target backend in the EAGLE-3 recipe. + +The real SGLang construction needs a GPU and is validated on the server; here +``SGLangEagle3TargetModel.from_pretrained`` is mocked and the tests pin the +recipe-side wiring: backend dispatch, environment guards (CUDA-only, +single-process-only), and how ``recipe_args`` flow into the SGLang kwargs. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from nemo_automodel.recipes.llm.train_eagle3 import TrainEagle3Recipe + + +class _RecipeCfg(SimpleNamespace): + """recipe_args stand-in: attribute access plus the ConfigNode-style ``get``.""" + + def get(self, key, default=None): + return getattr(self, key, default) + + +class _ToDictNode: + """ConfigNode stand-in: a mapping exposed only through ``to_dict()``.""" + + def __init__(self, values): + self._values = dict(values) + + def to_dict(self): + return dict(self._values) + + +def _make_recipe(world_size: int = 1, device: str = "cuda") -> TrainEagle3Recipe: + recipe = TrainEagle3Recipe.__new__(TrainEagle3Recipe) + recipe.device = torch.device(device) + recipe.dist_env = SimpleNamespace(world_size=world_size) + recipe.compute_dtype = torch.bfloat16 + return recipe + + +_FROM_PRETRAINED = "nemo_automodel.components.speculative.eagle.sglang_target.SGLangEagle3TargetModel.from_pretrained" + + +def test_setup_sglang_target_requires_cuda(): + recipe = _make_recipe(device="cpu") + with pytest.raises(ValueError, match="requires CUDA"): + recipe._setup_sglang_target(_RecipeCfg(), "target/path") + + +def test_setup_sglang_target_rejects_multi_process(): + recipe = _make_recipe(world_size=2) + with pytest.raises(ValueError, match="single-process training only"): + recipe._setup_sglang_target(_RecipeCfg(), "target/path") + + +def test_setup_sglang_target_builds_wrapper_with_defaults(): + recipe = _make_recipe() + wrapper = object() + with patch(_FROM_PRETRAINED, MagicMock(return_value=wrapper)) as from_pretrained: + recipe._setup_sglang_target(_RecipeCfg(), "target/path") + + assert recipe.target_wrapper is wrapper + assert recipe.target_model is None + from_pretrained.assert_called_once_with( + "target/path", + aux_layer_ids=None, + dtype=torch.bfloat16, + trust_remote_code=False, + mem_fraction_static=0.5, + ) + + +@pytest.mark.parametrize("as_node", [False, True], ids=["plain-dict", "to_dict-node"]) +def test_setup_sglang_target_forwards_sglang_args(as_node): + """``recipe_args.sglang_args`` overrides the defaults and passes extras through.""" + recipe = _make_recipe() + args = {"mem_fraction_static": 0.35, "page_size": 32} + cfg = _RecipeCfg( + sglang_args=_ToDictNode(args) if as_node else args, + aux_layer_ids=[1, 2, 3], + trust_remote_code=True, + ) + with patch(_FROM_PRETRAINED, MagicMock(return_value=object())) as from_pretrained: + recipe._setup_sglang_target(cfg, "target/path") + + from_pretrained.assert_called_once_with( + "target/path", + aux_layer_ids=[1, 2, 3], + dtype=torch.bfloat16, + trust_remote_code=True, + mem_fraction_static=0.35, + page_size=32, + ) + + +class _DispatchReached(Exception): + """Sentinel: the dispatch reached the expected backend setup method.""" + + +@pytest.mark.parametrize( + "backend, method", + [ + ("sglang", "_setup_sglang_target"), + ("colocated", "_setup_colocated_target"), + ("remote", "_setup_remote_target"), + ], +) +def test_online_target_dispatches_backend(monkeypatch, backend, method): + recipe = _make_recipe() + + def _sentinel(self, *args, **kwargs): + raise _DispatchReached() + + monkeypatch.setattr(TrainEagle3Recipe, method, _sentinel) + with pytest.raises(_DispatchReached): + recipe._setup_online_target(_RecipeCfg(target_model_backend=backend), "target/path", None) + + +def test_online_target_rejects_unknown_backend(): + recipe = _make_recipe() + with pytest.raises(ValueError, match="expected 'colocated', 'sglang', or 'remote'"): + recipe._setup_online_target(_RecipeCfg(target_model_backend="bogus"), "target/path", None) diff --git a/tests/unit_tests/speculative/test_eagle3_sglang.py b/tests/unit_tests/speculative/test_eagle3_sglang.py index fe3bae8dab..f2035496a1 100644 --- a/tests/unit_tests/speculative/test_eagle3_sglang.py +++ b/tests/unit_tests/speculative/test_eagle3_sglang.py @@ -90,9 +90,7 @@ def input_embedding_weight(self): def forward_eagle3(self, input_ids, attention_mask): captured = {} handles = [ - self.model.layers[i].register_forward_hook( - lambda _m, _i, out, i=i: captured.__setitem__(i, out) - ) + self.model.layers[i].register_forward_hook(lambda _m, _i, out, i=i: captured.__setitem__(i, out)) for i in self.aux_layer_ids ] try: @@ -225,6 +223,7 @@ def test_serve_target_engine_routing(monkeypatch): def _record(engine): def _builder(*_a, **_k): calls["engine"] = engine + return _builder monkeypatch.setattr(serve_target, "_build_hf_target", _record("hf")) @@ -312,3 +311,20 @@ def test_runner_forward_stacks_per_row(monkeypatch): assert logits.shape == (2, 5, _VOCAB) and aux.shape == (2, 5, 3 * _HIDDEN) torch.testing.assert_close(logits, torch.stack(rows_logits)) torch.testing.assert_close(aux, torch.stack(rows_aux)) + + +def test_sglang_dtype_str_mappings(): + """ServerArgs.dtype is a string in SGLang; torch dtypes must map to it.""" + from nemo_automodel.components.speculative.eagle.sglang_runner import sglang_dtype_str + + assert sglang_dtype_str(None) == "auto" + assert sglang_dtype_str(torch.float32) == "float32" + assert sglang_dtype_str(torch.float16) == "float16" + assert sglang_dtype_str(torch.bfloat16) == "bfloat16" + + +def test_sglang_dtype_str_rejects_unsupported(): + from nemo_automodel.components.speculative.eagle.sglang_runner import sglang_dtype_str + + with pytest.raises(ValueError, match="Unsupported SGLang target dtype"): + sglang_dtype_str(torch.int8) From 02c4a63b4ac1c110907d2626d804b9307be62c87 Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 11 Jun 2026 11:23:34 +0800 Subject: [PATCH 09/12] fix(speculative): install sglang out-of-band, not as a pyproject extra sglang==0.5.9 hard-pins transformers==4.57.1, which conflicts with the project's transformers==5.8.1. Declaring it as the spec_sglang extra made uv's universal resolution unsatisfiable, failing every install-dependent CI job (lint, type-check, uv-lock, builds). The SGLang target backend is lazy-imported (safe_import), so the package needs no declared sglang dependency; install it in a separate dedicated SD venv/container as the serve_target docstring now documents. Signed-off-by: khazic --- nemo_automodel/components/speculative/serve_target.py | 8 ++++++-- pyproject.toml | 6 ------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/nemo_automodel/components/speculative/serve_target.py b/nemo_automodel/components/speculative/serve_target.py index bb39b02906..d6d33ee5ac 100644 --- a/nemo_automodel/components/speculative/serve_target.py +++ b/nemo_automodel/components/speculative/serve_target.py @@ -39,8 +39,12 @@ - ``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 pinned in a separate speculative-decoding - environment, so it is imported only when this engine is selected. + 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 Both engines emit the identical supervision contract, so the training client is unchanged regardless of which one serves the target. diff --git a/pyproject.toml b/pyproject.toml index 625dc351a1..c7c187d182 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -162,12 +162,6 @@ s3 = [ msc = [ "multi-storage-client>=0.13", ] -# Speculative-decoding (EAGLE-3) SGLang target backend. Kept out of the main -# training image and pinned in a separate dedicated SD environment so the -# SGLang version stays isolated from the rest of the stack. -spec_sglang = [ - "sglang==0.5.9", -] all = [ "nemo_automodel[cli]", "nemo_automodel[cuda]", From cebe206a4a3ea303893ec9592693f86fee54bd49 Mon Sep 17 00:00:00 2001 From: khazic Date: Fri, 26 Jun 2026 19:16:28 +0800 Subject: [PATCH 10/12] fix(speculative): reject sequence packing for non-colocated EAGLE-3 targets The packed_sequence_size guard only blocked the remote backend, but the SGLang runner processes each row as one full causal sequence with no per-document masking, so packing + sglang silently leaked supervision across document boundaries. Gate packing on backend != 'colocated' and hoist the backend-name validation ahead of the guard so a misspelled backend still reports the clearer 'unknown backend' error. Also fix the copyright year in the new test (2026 -> 2025) and add coverage for the packing guard. Signed-off-by: khazic --- nemo_automodel/recipes/llm/train_eagle3.py | 16 ++++++----- .../recipes/llm/test_eagle3_sglang_backend.py | 28 +++++++++++++++---- 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/nemo_automodel/recipes/llm/train_eagle3.py b/nemo_automodel/recipes/llm/train_eagle3.py index faf108a861..b4e44ecf38 100644 --- a/nemo_automodel/recipes/llm/train_eagle3.py +++ b/nemo_automodel/recipes/llm/train_eagle3.py @@ -366,22 +366,24 @@ def _setup_online_target(self, recipe_cfg, target_path, target_config): # precomputed supervision over HTTP + NCCL. No target weights are # loaded here, which frees the training GPU's memory. backend = recipe_cfg.get("target_model_backend", "colocated") - # Sequence packing is colocated-only (the remote server does not yet honor - # per-document masking). + if backend not in ("colocated", "sglang", "remote"): + raise ValueError(f"Unknown target_model_backend={backend!r}; expected 'colocated', 'sglang', or 'remote'.") + # Sequence packing is colocated-only: neither the remote server nor the + # SGLang runner honors per-document masking (SGLang treats each row as one + # full causal sequence), so a packed row would leak supervision across + # document boundaries. packed_sequence_size = recipe_cfg.get("packed_sequence_size", 0) - if packed_sequence_size > 0 and backend == "remote": + if packed_sequence_size > 0 and backend != "colocated": raise NotImplementedError( "packed_sequence_size > 0 is only supported with the colocated target backend; " - "the remote backend does not yet propagate per-document masking." + f"the {backend!r} backend does not propagate per-document masking." ) if backend == "remote": self._setup_remote_target(recipe_cfg) elif backend == "sglang": self._setup_sglang_target(recipe_cfg, target_path) - elif backend == "colocated": + else: # colocated self._setup_colocated_target(recipe_cfg, target_path) - else: - raise ValueError(f"Unknown target_model_backend={backend!r}; expected 'colocated', 'sglang', or 'remote'.") self.train_dataloader = build_eagle3_dataloader( data_path=recipe_cfg.train_data_path, diff --git a/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py index 230bd8e5b1..b6720d24c4 100644 --- a/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py +++ b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -115,6 +115,10 @@ class _DispatchReached(Exception): """Sentinel: the dispatch reached the expected backend setup method.""" +def _sentinel(self, *args, **kwargs): + raise _DispatchReached() + + @pytest.mark.parametrize( "backend, method", [ @@ -125,10 +129,6 @@ class _DispatchReached(Exception): ) def test_online_target_dispatches_backend(monkeypatch, backend, method): recipe = _make_recipe() - - def _sentinel(self, *args, **kwargs): - raise _DispatchReached() - monkeypatch.setattr(TrainEagle3Recipe, method, _sentinel) with pytest.raises(_DispatchReached): recipe._setup_online_target(_RecipeCfg(target_model_backend=backend), "target/path", None) @@ -138,3 +138,21 @@ def test_online_target_rejects_unknown_backend(): recipe = _make_recipe() with pytest.raises(ValueError, match="expected 'colocated', 'sglang', or 'remote'"): recipe._setup_online_target(_RecipeCfg(target_model_backend="bogus"), "target/path", None) + + +@pytest.mark.parametrize("backend", ["sglang", "remote"]) +def test_online_target_rejects_packing_on_non_colocated(backend): + """packed_sequence_size > 0 is colocated-only; SGLang/remote leak across docs.""" + recipe = _make_recipe() + cfg = _RecipeCfg(target_model_backend=backend, packed_sequence_size=4) + with pytest.raises(NotImplementedError, match="only supported with the colocated"): + recipe._setup_online_target(cfg, "target/path", None) + + +def test_online_target_allows_packing_on_colocated(monkeypatch): + """packed_sequence_size > 0 passes the guard for the colocated backend.""" + recipe = _make_recipe() + monkeypatch.setattr(TrainEagle3Recipe, "_setup_colocated_target", _sentinel) + cfg = _RecipeCfg(target_model_backend="colocated", packed_sequence_size=4) + with pytest.raises(_DispatchReached): + recipe._setup_online_target(cfg, "target/path", None) From 6af83445e04e3949f4f7853b196a41517366f364 Mon Sep 17 00:00:00 2001 From: khazic Date: Sat, 27 Jun 2026 23:44:23 +0800 Subject: [PATCH 11/12] test(speculative): set cfg on EAGLE-3 recipe stub for CP gate _setup_online_target now reads cfg.get("distributed.cp_size") for the context-parallelism gate (merged in from #2465). The test helper built the recipe via __new__ without a cfg, so the four tests that dispatch through _setup_online_target failed with AttributeError. Give the stub an empty _RecipeCfg so the gate defaults to cp_size=1 (no CP). Signed-off-by: khazic --- tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py index b6720d24c4..2f3074ae51 100644 --- a/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py +++ b/tests/unit_tests/recipes/llm/test_eagle3_sglang_backend.py @@ -53,6 +53,9 @@ def _make_recipe(world_size: int = 1, device: str = "cuda") -> TrainEagle3Recipe recipe.device = torch.device(device) recipe.dist_env = SimpleNamespace(world_size=world_size) recipe.compute_dtype = torch.bfloat16 + # ``_setup_online_target`` reads ``cfg.get("distributed.cp_size", ...)`` for the + # CP gate; an empty config stand-in defaults it to cp_size=1 (no CP). + recipe.cfg = _RecipeCfg() return recipe From b2b8e970de43945c9525169fbcb793f80feec4f0 Mon Sep 17 00:00:00 2001 From: khazic Date: Sat, 27 Jun 2026 23:58:24 +0800 Subject: [PATCH 12/12] docs(speculative): note SGLang/transformers version compatibility SGLang 0.5.9 pins transformers==4.57.1 while NeMoAutoModelForCausalLM needs transformers 5.x (AutoModelForMultimodalLM), so the smoke --compare-hf check cannot run as written in one environment. Document this in the EAGLE guide's environment-setup step, per maintainer request on #2449. Signed-off-by: khazic --- docs/guides/speculative/eagle.mdx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/guides/speculative/eagle.mdx b/docs/guides/speculative/eagle.mdx index 1b61519e99..bf7b919b4f 100644 --- a/docs/guides/speculative/eagle.mdx +++ b/docs/guides/speculative/eagle.mdx @@ -89,6 +89,18 @@ For SGLang serving (Step 5), install it in the same environment: uv pip install "sglang>=0.5.9" ``` + +**SGLang / transformers version compatibility.** SGLang `0.5.9` pins +`transformers==4.57.1`. The SGLang target backend (`target_model_backend: sglang`) +and `serve_sglang` run in that environment without issue. The `--compare-hf` +check in `smoke_sglang_target.py` additionally builds the HuggingFace target +through `NeMoAutoModelForCausalLM`, which imports `AutoModelForMultimodalLM` +(available only in `transformers` 5.x), so that single-process comparison cannot +run as written under `transformers==4.57.1`. To compare the two backends, use a +`transformers` build that satisfies both, or load the HuggingFace side with plain +`transformers.AutoModelForCausalLM` wrapped in `HFEagle3TargetModel`. + + --- ## Step 1 — Understand EAGLE Architecture