nemo_automodel.components.loss.intermediate_distill

Module Contents

Classes

Name	Description
`IntermediateDistillLoss`	Intermediate-layer feature distillation module.
`LayerCapture`	Forward-hook helper for capturing selected intermediate hidden states.

Functions

Name	Description
`_lookup_layer`	-
`_masked_per_token_loss`	-
`_pick_projector`	-
`intermediate_loss_function`	Per-token feature distillation between matched student/teacher layers.
`intermediate_loss_pair`	Sum of intermediate loss on query and doc sides.

Data

HiddenStatesLike

ProjectorLike

API

class nemo_automodel.components.loss.intermediate_distill.IntermediateDistillLoss(
    layer_pairs: collections.abc.Sequence[tuple[int, int]] | collections.abc.Sequence[list[int]],
    loss_type: str = 'mse',
    normalize: bool = False,
    layer_weights: collections.abc.Sequence[float] | None = None,
    reduction: str = 'mean'
)

Bases: Module

Intermediate-layer feature distillation module.

layer_pairs

nemo_automodel.components.loss.intermediate_distill.IntermediateDistillLoss.forward(
    s_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    s_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    attn_q: torch.Tensor | None = None,
    attn_d: torch.Tensor | None = None,
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None = None
) -> torch.Tensor

class nemo_automodel.components.loss.intermediate_distill.LayerCapture(
    detach: bool = False
)

Forward-hook helper for capturing selected intermediate hidden states.

_handles

list[RemovableHandle] = []

_outputs

dict[int, Tensor] = {}

outputs

dict[int, Tensor]

nemo_automodel.components.loss.intermediate_distill.LayerCapture._make_hook(
    idx: int
)

nemo_automodel.components.loss.intermediate_distill.LayerCapture.attach(
    layers: torch.nn.ModuleList,
    indices: collections.abc.Iterable[int]
) -> None

nemo_automodel.components.loss.intermediate_distill.LayerCapture.detach_hooks() -> None

nemo_automodel.components.loss.intermediate_distill.LayerCapture.reset() -> None

nemo_automodel.components.loss.intermediate_distill._lookup_layer(
    hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    idx: int,
    side: str
) -> torch.Tensor

nemo_automodel.components.loss.intermediate_distill._masked_per_token_loss(
    s: torch.Tensor,
    t: torch.Tensor,
    attention_mask: torch.Tensor | None,
    loss_type: str
) -> torch.Tensor

nemo_automodel.components.loss.intermediate_distill._pick_projector(
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None,
    s_idx: int
) -> collections.abc.Callable[[torch.Tensor], torch.Tensor] | None

nemo_automodel.components.loss.intermediate_distill.intermediate_loss_function(
    student_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    teacher_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    layer_pairs: collections.abc.Sequence[tuple[int, int]],
    attention_mask: torch.Tensor | None = None,
    projector: nemo_automodel.components.loss.intermediate_distill.ProjectorLike | None = None,
    loss_type: str = 'mse',
    normalize: bool = False,
    layer_weights: collections.abc.Sequence[float] | None = None,
    reduction: str = 'mean'
) -> torch.Tensor

Per-token feature distillation between matched student/teacher layers.

nemo_automodel.components.loss.intermediate_distill.intermediate_loss_pair(
    s_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_q_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    s_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    t_d_hidden_states: nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike,
    layer_pairs: collections.abc.Sequence[tuple[int, int]],
    attn_q: torch.Tensor | None = None,
    attn_d: torch.Tensor | None = None,
    kwargs = {}
) -> torch.Tensor

Sum of intermediate loss on query and doc sides.

nemo_automodel.components.loss.intermediate_distill.HiddenStatesLike = Sequence[torch.Tensor] | Mapping[int, torch.Tensor]

nemo_automodel.components.loss.intermediate_distill.ProjectorLike = Callable[[torch.Tensor], torch.Tensor] | Mapping[int, Callable[[torch.Tensor], t...