feat: enable knowledge distillation

There are many different forms of model training which exist. One popular form of training is knowledge distillation, where a student model learns the output distributions from a teacher model.  This commit introduces support for knowledge distillation in the training library.

This commit also exposes the `weight_decay` hyperparameter which is often used to help deep learning models generalize.

Lastly, this commit changes the useage from `torch.distributed` to just `dist`, as it is a common module used throughout the codebase.

Signed-off-by: Oleg S <[email protected]>

Author: RobotSail

src/instructlab/training/config.py

@@ -121,6 +121,16 @@ class DeepSpeedOptions(BaseModel):
     save_samples: int | None = None
 
 
+class DistillationConfig(BaseModel):
+    """
+    Config to use when performing knowledge distillation during training.
+    """
+
+    temperature: float = Field(1.0, gt=0.0)
+    alpha: float = Field(1.0, le=1.0, ge=0.0)
+    teacher_path: str
+
+
 # public API
 class ShardingStrategies(Enum):
     FULL_SHARD = "FULL_SHARD"
@@ -179,6 +189,11 @@ class TrainingArgs(BaseModel):
     is_padding_free: bool = False  # TODO: deprecate
     checkpoint_at_epoch: bool = True
     accelerate_full_state_at_epoch: bool = True
+    weight_decay: float = Field(0.0, ge=0.0)
+
+    # settings for knowledge distillation
+    distillation_options: Optional[DistillationConfig] = None
+    use_distillation: bool = False
 
     mock_data: Optional[bool] = False
     mock_data_len: int = 0