`nemo_rl.models.megatron.draft.training`#

Module Contents#

Classes#

`DraftSpeculator`	Model-build operation selected by a configured draft speculator.
`Eagle3Speculator`	Current EAGLE-3 draft speculator.

Functions#

resolve_draft_speculator

Resolve an enabled draft configuration to its speculator.

Data#

_SPECULATOR_FACTORIES

API#

class nemo_rl.models.megatron.draft.training.DraftSpeculator#

Bases: typing.Protocol

Model-build operation selected by a configured draft speculator.

config: nemo_rl.models.policy.draft_config.Eagle3DraftConfig#: None

build_model( *, model_provider: megatron.bridge.models.model_provider.ModelProviderMixin, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, policy_model_chunk: megatron.core.transformer.MegatronModule, ) → megatron.core.transformer.MegatronModule | None#: Build the draft model for the policy chunk that owns it.

class nemo_rl.models.megatron.draft.training.Eagle3Speculator#

Current EAGLE-3 draft speculator.

config: nemo_rl.models.policy.draft_config.Eagle3DraftConfig#: None

build_model( *, model_provider: megatron.bridge.models.model_provider.ModelProviderMixin, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, policy_model_chunk: megatron.core.transformer.MegatronModule, ) → megatron.core.transformer.MegatronModule | None#: Build the EAGLE-3 draft model with the existing implementation.

nemo_rl.models.megatron.draft.training._SPECULATOR_FACTORIES: dict[str, type[nemo_rl.models.megatron.draft.training.Eagle3Speculator]]#: None

nemo_rl.models.megatron.draft.training.resolve_draft_speculator( config: nemo_rl.models.policy.draft_config.Eagle3DraftConfig | collections.abc.Mapping[str, Any] | None, ) → nemo_rl.models.megatron.draft.training.DraftSpeculator | None#