nemo_rl.models.automodel.draft.runtime

Per-worker draft co-training runtimes: loss computation and capture lifecycle.

Module Contents

Classes

_DraftRuntimeBase	Shared per-worker draft co-training state and distributed plumbing.
DSparkRuntime	Per-worker DSpark co-training state: draft model, capture, and loss.
Eagle3Runtime	Per-worker EAGLE3 co-training state: draft model, capture, TTT loss.

Functions

anchor_sampling_seed	Deterministic seed for the draft’s anchor sampling.
draft_capture_ctx	Keep the hidden-capture hooks active for the duration of a train call.
next_token_position_mask	Shift a per-TOKEN mask onto the POSITIONS whose logits predict it.
_shift_left_with_zero	Shift a sequence tensor left by one position along dim 1, zero-tailed.
draft_refit_export_name	Strip the eagle3 draft’s model. prefix so refit keys match what the vLLM drafter expects.
build_draft_runtime	Build the per-worker draft runtime, attach hidden capture to the policy.

API

nemo_rl.models.automodel.draft.runtime.anchor_sampling_seed(

dp_rank: int,

global_batch_index: int,

microbatch_index: int,

) → int

Deterministic seed for the draft’s anchor sampling.

A pure function of the DP rank and the training progress counters: every TP/CP peer of the same DP slice derives the identical seed (the draft is replicated across them, so they must sample identical anchors or the replicas drift apart), while DP ranks, global batches, and microbatches all draw differently. Independent of each rank’s ambient RNG state.

class nemo_rl.models.automodel.draft.runtime._DraftRuntimeBase(

draft_model: torch.nn.Module,

options: Any,

loss_weight: float,

dp_group: Optional[torch.distributed.ProcessGroup],

tp_group: Optional[torch.distributed.ProcessGroup] = None,

cp_group: Optional[torch.distributed.ProcessGroup] = None,

)

Shared per-worker draft co-training state and distributed plumbing.

Subclasses supply the loss (compute_loss) and the layer-id source for the hidden capture (_capture_layer_ids); everything here is common to every drafter family: the option/group state, the microbatch-slot normalization, and the teacher-logits stash.

Initialization

_algo_label

‘Draft’

property _cp_size: int

begin_global_batch(num_microbatch_slots: int) → None

Record this global batch’s microbatch-slot count (identical on all ranks).

The training loop sums per-microbatch losses across gradient accumulation, while each draft microbatch loss is a mean over that microbatch’s (DP-all-reduced) supervised tokens. Dividing every microbatch term by the slot count turns the sum into an average of per-slot global means, keeping the draft gradient scale independent of gbs/mbs — the same effective scale as the policy loss, which normalizes each microbatch by the whole-global-batch token denominator instead.

abstractmethod _capture_layer_ids() → list[int]

attach_capture(policy_model: torch.nn.Module) → None

stash_teacher_logits(logits: torch.Tensor, need_clone: bool) → None

Record the policy’s raw logits before any in-place temperature scaling.

With temperature 1.0 no scaling happens and a detached view is exact and free; otherwise the in-place div would corrupt the view, so clone first.

Under TP the logits arrive as a vocab-sharded DTensor; it is stashed as-is and gathered to the full vocab lazily in compute_loss so the full-size tensor does not coexist with the policy-loss peak.

_require_capture_and_teacher() → torch.Tensor

Common compute_loss preamble: guards plus the stashed teacher.

class nemo_rl.models.automodel.draft.runtime.DSparkRuntime(

draft_model: nemo_rl.models.automodel.draft.draft_qwen3.Qwen3DSparkModel,

dspark_options: dict[str, Any],

loss_weight: float,

dp_group: Optional[torch.distributed.ProcessGroup],

tp_group: Optional[torch.distributed.ProcessGroup] = None,

cp_group: Optional[torch.distributed.ProcessGroup] = None,

)

Bases: nemo_rl.models.automodel.draft.runtime._DraftRuntimeBase

Per-worker DSpark co-training state: draft model, capture, and loss.

Initialization

_algo_label

‘DSpark’

_capture_layer_ids() → list[int]

compute_loss(

data_dict: Any,

) → tuple[torch.Tensor, dict[str, float]]

Run the draft forward on the captured hiddens and compute the DSpark loss.

static _terms_to_metrics(

terms: dict[str, torch.Tensor],

microbatch_slots: int,

) → dict[str, float]

Flatten the loss terms into shape-stable scalar metrics.

loss/ce_loss/l1_loss/confidence_loss use DP-global denominators and are divided by microbatch_slots to match the backward value (see compute_loss). Ratios (tau, accept_rate) are emitted as raw *_num/*_den pairs, NOT pre-divided: the training-loop metric aggregation (grpo.py’s per-key reduction) sums every metric not on its small mean-reduction allowlist across all microbatches and DP ranks, so a pre-divided per-microbatch ratio would be summed into a meaningless value (see nemo_rl.algorithms.utils.finalize_draft_ratio_metrics, which the training loop calls after that reduction to turn the summed num/den pairs back into the correct token-weighted global ratio).

nemo_rl.models.automodel.draft.runtime.draft_capture_ctx(runtime: Any)

Keep the hidden-capture hooks active for the duration of a train call.

Works for any draft runtime exposing capture (dspark/dflash/eagle3).

nemo_rl.models.automodel.draft.runtime.next_token_position_mask(token_mask: torch.Tensor) → torch.Tensor

Shift a per-TOKEN mask onto the POSITIONS whose logits predict it.

Rollout token_mask[t] marks token t as supervised (a response token), but logits at position t predict token t+1 (the policy loss paths apply token_mask[:, 1:] for the same reason). The eagle3 TTT forward gates loss at logit/teacher positions, so its mask must be mask[t] = token_mask[t + 1] with a zero tail: the final position has no next-token label, and the last-prompt-token position (whose label is the FIRST response token — exactly where drafting starts) is supervised.

nemo_rl.models.automodel.draft.runtime._shift_left_with_zero(tensor: torch.Tensor) → torch.Tensor

Shift a sequence tensor left by one position along dim 1, zero-tailed.

shifted[:, t] = tensor[:, t + 1]; the last position is zeroed. Matches Automodel’s Eagle3Target._shift_left_with_zero convention (speculative/eagle/target.py).

class nemo_rl.models.automodel.draft.runtime.Eagle3Runtime(

draft_model: nemo_rl.models.automodel.draft.eagle3_llama.Eagle3DraftModel,

eagle3_options: nemo_rl.models.policy.Eagle3DraftConfig,

loss_weight: float,

dp_group: Optional[torch.distributed.ProcessGroup],

tp_group: Optional[torch.distributed.ProcessGroup] = None,

cp_group: Optional[torch.distributed.ProcessGroup] = None,

)

Bases: nemo_rl.models.automodel.draft.runtime._DraftRuntimeBase

Per-worker EAGLE3 co-training state: draft model, capture, TTT loss.

Exposes the same runtime protocol as DSparkRuntime so the worker and loss wrapper treat all draft algos uniformly.

Initialization

_algo_label

‘EAGLE3’

_capture_layer_ids() → list[int]

compute_loss(

data_dict: Any,

) → tuple[torch.Tensor, dict[str, float]]

Run the TTT draft forward on captured hiddens and compute the loss.

nemo_rl.models.automodel.draft.runtime.draft_refit_export_name(name: str) → str

Strip the eagle3 draft’s model. prefix so refit keys match what the vLLM drafter expects.

nemo_rl.models.automodel.draft.runtime.build_draft_runtime(

draft_model: torch.nn.Module,

draft_config: Any,

policy_model: torch.nn.Module,

policy_config: dict[str, Any],

dp_group: Optional[torch.distributed.ProcessGroup],

tp_group: Optional[torch.distributed.ProcessGroup] = None,

cp_group: Optional[torch.distributed.ProcessGroup] = None,

) → nemo_rl.models.automodel.draft.runtime._DraftRuntimeBase

Build the per-worker draft runtime, attach hidden capture to the policy.

Also validates generation.vllm_kwargs.speculative_config.num_speculative_tokens against the drafter’s actual proposal count (eagle3’s ttt_steps, or the block drafter’s block_size minus the unsupervised anchor slot for dflash) – this needs the built drafter’s own config (block_size / sample_from_anchor), so config validation alone cannot do it.