`nemo_rl.models.megatron.draft.optimizer`#

Module Contents#

`DraftOptimizerConfigOwner`	Draft config fields required to opt in to optimizer overrides.
`DraftOptimizerConfigOverrideProvider`	Add scheduler overrides only for parameters owned by a draft model.

`_is_draft_parameter`
`build_draft_optimizer_override_provider`	Build an override provider only when a draft schedule is configured.

class nemo_rl.models.megatron.draft.optimizer.DraftOptimizerConfigOwner#

Bases: typing.Protocol

Draft config fields required to opt in to optimizer overrides.

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

nemo_rl.models.megatron.draft.optimizer._is_draft_parameter(parameter: torch.nn.Parameter) → bool#

class nemo_rl.models.megatron.draft.optimizer.DraftOptimizerConfigOverrideProvider#

Bases: megatron.bridge.training.config.OptimizerConfigOverrideProvider

Add scheduler overrides only for parameters owned by a draft model.

draft_optimizer: nemo_rl.models.policy.draft_config.DraftOptimizerConfig#: None

build_config_overrides( context: megatron.bridge.training.config.OptimizerConfigOverrideProviderContext, ) → dict[megatron.core.optimizer.ParamKey, megatron.core.optimizer_param_scheduler.ParamGroupOverride] | None#: Combine standard Megatron groups with the opt-in draft schedule.

nemo_rl.models.megatron.draft.optimizer.build_draft_optimizer_override_provider( draft_config: nemo_rl.models.megatron.draft.optimizer.DraftOptimizerConfigOwner | None, ) → nemo_rl.models.megatron.draft.optimizer.DraftOptimizerConfigOverrideProvider | None#: Build an override provider only when a draft schedule is configured.