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
AutoModel LLM KD with independently weighted main/MTP objectives. |
|
AutoModel VLM KD, including a teacher-first PP pass. |
Functions
Allow DCP to save and restore disjoint pipeline-stage state dictionaries. |
- class KnowledgeDistillationRecipeForNextTokenPrediction#
Bases:
_WeightedObjectiveMixin,KnowledgeDistillationRecipeForNextTokenPredictionAutoModel LLM KD with independently weighted main/MTP objectives.
- run_train_validation_loop()#
- setup()#
- class KnowledgeDistillationRecipeForVLM#
Bases:
_WeightedObjectiveMixin,KnowledgeDistillationRecipeForVLMAutoModel 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