global_kd_recipe

Small Puzzletron extensions to NeMo AutoModel’s native KD recipes.

AutoModel owns model construction, FSDP2/TP/EP/CP/SP/PP, dataloading, checkpointing, and the train loop. This module only adds independent weighted main/MTP CE and KD terms, TVD selection, and the VLM PP teacher pass missing from the upstream VLM KD recipe.

Classes

`KnowledgeDistillationRecipeForNextTokenPrediction`	AutoModel LLM KD with independently weighted main/MTP objectives.
`KnowledgeDistillationRecipeForVLM`	AutoModel VLM KD, including a teacher-first PP pass.

Functions

install_pp_checkpoint_state_dict_support

Allow DCP to save and restore disjoint pipeline-stage state dictionaries.

class KnowledgeDistillationRecipeForNextTokenPrediction

Bases: _WeightedObjectiveMixin, KnowledgeDistillationRecipeForNextTokenPrediction

AutoModel LLM KD with independently weighted main/MTP objectives.

run_train_validation_loop()

setup()

class KnowledgeDistillationRecipeForVLM

Bases: _WeightedObjectiveMixin, KnowledgeDistillationRecipeForVLM

AutoModel VLM KD, including a teacher-first PP pass.

setup()

install_pp_checkpoint_state_dict_support()

Allow DCP to save and restore disjoint pipeline-stage state dictionaries.

PyTorch DCP’s strict model-state verification assumes every distributed rank describes the same module tree. Pipeline ranks intentionally own disjoint parameter FQNs, so strict verification can reduce a valid local state dict to an empty mapping. NeMo AutoModel’s checkpoint wrappers already describe multi-stage PP as non-strict on load; apply the same rule consistently to model and optimizer save/load calls for this PP training process.

Return type:: None