nemo_rl.models.policy.draft_config

Module Contents

Classes

DraftOptimizerConfig	Optional optimizer schedule for draft-model parameters.
Eagle3DraftConfig	Configuration for EAGLE-3 draft-model co-training with the policy.

Functions

coerce_draft_config	Accept either a validated model or a raw mapping at API boundaries.
draft_refit_enabled	Return whether generation must accept refitted draft weights.

API

class nemo_rl.models.policy.draft_config.DraftOptimizerConfig

Bases: pydantic.BaseModel

Optional optimizer schedule for draft-model parameters.

lr: Annotated[float, Field(gt=0)]

None

min_lr: Annotated[float, Field(ge=0)] | None

None

weight_decay: Annotated[float, Field(ge=0)] | None

None

validate_lr_range() → Self

Require the draft minimum learning rate to fit its schedule.

class nemo_rl.models.policy.draft_config.Eagle3DraftConfig

Bases: pydantic.BaseModel

Configuration for EAGLE-3 draft-model co-training with the policy.

speculator_type: Literal[eagle3]

‘eagle3’

enabled: bool

False

model_name: str | None

None

loss_weight: float

0.1

num_layers: int | None

None

aux_layer_indices: list[int] | None

None

optimizer: nemo_rl.models.policy.draft_config.DraftOptimizerConfig | None

None

_reject_near_miss_extra_keys() → nemo_rl.models.policy.draft_config.Eagle3DraftConfig

nemo_rl.models.policy.draft_config.coerce_draft_config(

config: Eagle3DraftConfig | Mapping[str, Any] | None,

) → nemo_rl.models.policy.draft_config.Eagle3DraftConfig | None

Accept either a validated model or a raw mapping at API boundaries.

MasterConfig validation normally produces the model, but PolicyConfig is a TypedDict, so callers that assemble one by hand still pass a plain dict.

nemo_rl.models.policy.draft_config.draft_refit_enabled(

config: Eagle3DraftConfig | Mapping[str, Any] | None,

) → bool

Return whether generation must accept refitted draft weights.