compare_module_outputs

Compare module output tensors from different model variants.

This module provides: 1. OutputSaveHook - A PyTorch hook to capture module outputs during forward pass 2. Comparison utilities - Compute RMSE and cosine similarity between saved outputs

Usage Example:

Step 1: Capture outputs from multiple layers:

from modelopt.torch.prune.importance_hooks.compare_module_outputs import (
    OutputSaveHook,
    save_multi_layer_outputs,
)

# Register hooks on all target layers
hooks = {}
for name, module in model.named_modules():
    if name.endswith('mlp.linear_fc2'):
        hook = OutputSaveHook(layer_name=name)
        module.register_forward_hook(hook)
        hooks[name] = hook

# Run inference/training
model(input_data)

# Save all layer outputs
save_multi_layer_outputs(hooks, "output_unpruned.pt")

Step 2: Compare outputs from different model variants:

python compare_module_outputs.py \
    --reference output_unpruned.pt \
    --compare output_l2norm.pt \
    --output-json comparison_stats.json

The saved file format:

{
    'decoder.layers.0.mlp.linear_fc2': Tensor([steps, seq_len, batch, hidden]),
    'decoder.layers.1.mlp.linear_fc2': Tensor([...]),
    ...
    'metadata': {'num_layers': N, 'num_steps': M, 'layer_names': [...]}
}

Classes

OutputSaveHook

Hook to capture and save module outputs during forward pass.

Functions

compare_multi_layer	Compare multi-layer outputs.
compute_cosine_similarity	Compute average cosine similarity between two tensors.
compute_layer_metrics	Compute RMSE and cosine similarity for a layer's outputs.
compute_rmse	Compute Root Mean Square Error between two tensors.
main	Compare module output tensors from different model variants.
save_multi_layer_outputs	Save outputs from multiple layers to a single file.

class OutputSaveHook

Bases: object

Hook to capture and save module outputs during forward pass.

__init__(layer_name)

Initialize the output save hook.

Parameters:

layer_name (str) – Hierarchical name of the layer (e.g., ‘decoder.layers.0.mlp.linear_fc2’).

Return type:

None

get_outputs_list()

Return saved outputs as a list.

Return type:

list[Tensor]

compare_multi_layer(ref_data, comp_data, output_json=None)

Compare multi-layer outputs.

Parameters:

• ref_data (dict)

• comp_data (dict)

• output_json (str | None)

compute_cosine_similarity(tensor1, tensor2)

Compute average cosine similarity between two tensors.

Parameters:

• tensor1 (Tensor)

• tensor2 (Tensor)

Return type:

dict

compute_layer_metrics(ref_data, comp_data)

Compute RMSE and cosine similarity for a layer’s outputs.

Parameters:

• ref_data (list) – List of reference tensors.

• comp_data (list) – List of comparison tensors.

Returns:

Dictionary with metrics.

Raises:

ValueError – If lengths don’t match or tensor shapes don’t match.

Return type:

dict

compute_rmse(tensor1, tensor2)

Compute Root Mean Square Error between two tensors.

Parameters:

• tensor1 (Tensor)

• tensor2 (Tensor)

Return type:

float

main()

Compare module output tensors from different model variants.

save_multi_layer_outputs(hooks, path)

Save outputs from multiple layers to a single file.

Parameters:

• hooks (dict[str, OutputSaveHook]) – Dictionary mapping layer names to their hooks.

• path (str) – Path to save the outputs.

Return type:

None