bridge.training.post_training.distillation#

Module Contents#

Classes#

ModelOptDistillConfig

Configuration settings for Model Optimizer distillation.

Functions#

create_kd_loss_function

Masked next-token loss, wrapped with the KD loss when the model is distilling.

loss_func_kd

Loss function (with KD Loss support).

_mask_loss

API#

class bridge.training.post_training.distillation.ModelOptDistillConfig#

Bases: modelopt.torch.distill.plugins.megatron.DistillationConfig

Configuration settings for Model Optimizer distillation.

bridge.training.post_training.distillation.create_kd_loss_function(
loss_mask: torch.Tensor,
model: megatron.core.transformer.MegatronModule,
check_for_nan_in_loss: bool,
check_for_spiky_loss: bool,
) → Callable#

Masked next-token loss, wrapped with the KD loss when the model is distilling.

Forward steps share this so a model family only has to supply its own batch handling to become distillable; a step that builds its loss directly would silently train without KD.

bridge.training.post_training.distillation.loss_func_kd(
output_tensor: torch.Tensor,
loss_mask: torch.Tensor,
original_loss_fn: Callable,
model: megatron.core.transformer.MegatronModule,
)#

Loss function (with KD Loss support).

Parameters:

  • output_tensor (Tensor) – The tensor with the losses

  • loss_mask (Tensor) – Used to mask out some portions of the loss

  • original_loss_fn (Callable) – The original loss function

  • model (GPTModel) – The model (can be wrapped)

bridge.training.post_training.distillation._mask_loss(output_tensor: torch.Tensor, loss_mask: torch.Tensor)#