Skip to content

llmcompressor.modifiers.utils.pytorch_helpers

PyTorch-specific helper functions for model compression.

Provides utility functions for PyTorch model operations including batch processing, padding mask application, and model architecture detection. Supports MoE (Mixture of Experts) models and specialized tensor operations for compression workflows.

Functions:

apply_pad_mask_to_batch 

apply_pad_mask_to_batch(
    batch: Dict[str, Tensor]
) -> Dict[str, torch.Tensor]

Apply a mask to the input ids of a batch. This is used to zero out padding tokens so they do not contribute to the hessian calculation in the GPTQ and SparseGPT algorithms

Assumes that attention_mask only contains zeros and ones

Parameters:

  • batch

    (Dict[str, Tensor]) –

    batch to apply padding to if it exists

Returns:

  • Dict[str, Tensor] –

    batch with padding zeroed out in the input_ids

Source code in llmcompressor/modifiers/utils/pytorch_helpers.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def apply_pad_mask_to_batch(batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """
    Apply a mask to the input ids of a batch. This is used to zero out
    padding tokens so they do not contribute to the hessian calculation in the
    GPTQ and SparseGPT algorithms

    Assumes that `attention_mask` only contains zeros and ones

    :param batch: batch to apply padding to if it exists
    :return: batch with padding zeroed out in the input_ids
    """
    if "attention_mask" in batch:
        for key in ("input_ids", "decoder_input_ids"):
            if key in batch:
                batch[key] = batch[key] * batch["attention_mask"]

    return batch

is_moe_model 

is_moe_model(model: Module) -> bool

Check if the model is a mixture of experts model

Parameters:

  • model

    (Module) –

    the model to check

Returns:

  • bool –

    True if the model is a mixture of experts model

Source code in llmcompressor/modifiers/utils/pytorch_helpers.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def is_moe_model(model: Module) -> bool:
    """
    Check if the model is a mixture of experts model

    :param model: the model to check
    :return: True if the model is a mixture of experts model
    """

    # Check for MoE components
    for _, module in model.named_modules():
        module_name = module.__class__.__name__
        if "MoE" in module_name or "Expert" in module_name:
            return True

    # Check config for MoE attributes
    if hasattr(model, "config"):
        if any(
            "moe" in attr.lower() or "expert" in attr.lower()
            for attr in dir(model.config)
        ):
            return True

    return False