nemo_automodel.components.speculative.eagle.msd_curriculum

Two-stage data curriculum for Multimodal Speculative Decoding training.

Module Contents

Classes

Name	Description
`MSDCurriculum`	Stateless-by-step two-stage curriculum with checkpointable progress.
`MSDCurriculumConfig`	Configuration for the two-stage MSD data curriculum.
`MSDCurriculumPhase`	High-level phase of the two-stage MSD training schedule.
`MSDDataSource`	Data source selected for one MSD optimizer step.

API

class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum(
    config: nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig,
    global_step: int = 0
)

Stateless-by-step two-stage curriculum with checkpointable progress.

Call :meth:source_for_step from a training loop that owns separate text and multimodal dataloaders, then draw the selected dataloader’s next batch. :meth:next_source is a convenience for single-process loops and stores the next global step in its checkpoint state.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.load_state_dict(
    state_dict: dict[str, typing.Any]
) -> None

Restore curriculum progress from a checkpoint state dictionary.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.multimodal_probability(
    global_step: int
) -> float

Return the multimodal batch probability at global_step.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.next_source() -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource

Return the current source selection and advance one optimizer step.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.phase_for_step(
    global_step: int
) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase

Return the curriculum phase active at global_step.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.source_for_step(
    global_step: int
) -> nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource

Deterministically choose the data source for global_step.

The random draw is seeded from the global step rather than mutable RNG state. Resumed and distributed runs therefore make the same sampling decision for a given optimizer step.

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculum.state_dict() -> dict[str, typing.Any]

Return progress required to resume :meth:next_source.

class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig(
    text_only_steps: int,
    multimodal_ramp_steps: int,
    final_multimodal_ratio: float = 1.0,
    seed: int = 42
)

Dataclass

Configuration for the two-stage MSD data curriculum.

The first stage consumes text-only instruction data. In the second stage, the probability of selecting a multimodal batch rises linearly to final_multimodal_ratio. A final ratio of one makes the post-ramp phase entirely multimodal, matching the curriculum described in the MSD paper.

Parameters:

text_only_steps

int

Number of initial optimizer steps using only text batches.

multimodal_ramp_steps

int

Number of stage-two steps over which the multimodal sampling probability rises to its final value. Set to zero to switch immediately after the text-only stage.

final_multimodal_ratio

floatDefaults to 1.0

Multimodal batch probability after the ramp.

seed

intDefaults to 42

Seed used for stateless, step-indexed source selection. All distributed ranks therefore select the same source for a global optimizer step without communicating RNG state.

final_multimodal_ratio

float = 1.0

multimodal_ramp_steps

int

seed

int = 42

text_only_steps

int

nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumConfig.__post_init__() -> None

class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDCurriculumPhase

Bases: enum.Enum

High-level phase of the two-stage MSD training schedule.

MULTIMODAL

= 'multimodal'

MULTIMODAL_RAMP

= 'multimodal_ramp'

TEXT_ONLY

= 'text_only'

class nemo_automodel.components.speculative.eagle.msd_curriculum.MSDDataSource

Bases: enum.Enum

Data source selected for one MSD optimizer step.

MULTIMODAL

= 'multimodal'

TEXT

= 'text'