nemo_rl.models.megatron.draft.step_state

Module Contents

Classes

DraftStepPayload	Detached raw EAGLE statistics carried through loss metrics.
DraftStepState	Accumulate one-bin EAGLE loss state across split microbatches.

Data

DRAFT_STEP_PAYLOAD_KEY
DRAFT_LOSS_METRIC_KEY

API

nemo_rl.models.megatron.draft.step_state.DRAFT_STEP_PAYLOAD_KEY

‘_draft_step_payload’

nemo_rl.models.megatron.draft.step_state.DRAFT_LOSS_METRIC_KEY

‘draft_loss’

class nemo_rl.models.megatron.draft.step_state.DraftStepPayload

Detached raw EAGLE statistics carried through loss metrics.

stats: nemo_rl.algorithms.loss.draft.DraftLossStats

None

class nemo_rl.models.megatron.draft.step_state.DraftStepState

Accumulate one-bin EAGLE loss state across split microbatches.

_local_numerators: torch.Tensor | None

None

_local_counts: torch.Tensor | None

None

_weights: torch.Tensor | None

None

_global_counts: torch.Tensor | None

None

static metric_payload(

stats: nemo_rl.algorithms.loss.draft.DraftLossStats,

) → nemo_rl.models.megatron.draft.step_state.DraftStepPayload

Build a non-differentiable payload safe for metric transport.

property active: bool

property local_numerators: torch.Tensor

property local_counts: torch.Tensor

accumulate(

payload: nemo_rl.models.megatron.draft.step_state.DraftStepPayload,

) → None

Add one detached microbatch payload to this optimizer step.

counts_for_reduction(reference: torch.Tensor) → torch.Tensor

Return local counts on the existing policy-count reduction tensor.

set_global_counts(counts: torch.Tensor) → None

Record the DP-reduced one-bin denominator.

_normalization_scale() → float

normalize_metric(value: Any) → Any

Normalize a raw local draft numerator by the global draft count.

correct_main_grads(

parameters: Iterable[Any],

*,

policy_normalization_count: torch.Tensor,

) → None

Correct draft-tagged main grads after policy normalization.

active is a job-wide predicate: the payload reaches every pipeline rank through broadcast_loss_metrics_from_last_stage, while the draft module is attached only to the post-process chunk. Ranks that hold no draft-tagged parameter are therefore expected and must no-op here.