src/utils/model_stats.py

"""
Model statistics: parameter counts and memory estimation.

Reports total and trainable parameter counts, and estimates GPU
memory consumption based on the specified dtype (default bfloat16,
2 bytes per parameter).  Useful for quick sanity checks when
loading a new model variant.
"""

import torch
import torch.nn as nn

def get_model_stats(model: nn.Module, dtype=torch.bfloat16):
    """
    Analyzes a PyTorch model to report parameter counts and estimated memory usage.
    
    Args:
        model: The PyTorch model.
        dtype: The target precision for memory estimation (default: bfloat16).
    """
    # 1. Parameter Counts
    # using set() automatically handles weight tying (deduplicates shared tensors)
    # unique_params = sum(p.numel() for p in set(model.parameters()))
    total_params = sum(p.numel() for p in model.parameters())
    unique_params = total_params - model.tok_emb.weight.numel()
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    # 2. Memory Usage Calculation
    # We estimate memory for: Parameters + Gradients + Buffers
    param_mem = total_params * dtype.itemsize
    grad_mem = trainable_params * dtype.itemsize 
    buffer_mem = sum(b.numel() * b.element_size() for b in model.buffers())
    
    total_mem_bytes = param_mem + grad_mem + buffer_mem
    total_mem_gb = total_mem_bytes / (1024**3)

    # 3. Print Report
    print(f"📊 Model Statistics")
    print(f"{'-'*30}")
    print(f"• Total Parameters:    {total_params:,}")
    if total_params != unique_params:
        print(f"• Unique Parameters:   {unique_params:,} (Weight Tying Detected)")
    print(f"• Trainable Params:    {trainable_params:,}")
    print(f"• Est. Memory ({dtype}): {total_mem_gb:.2f} GB")
    print(f"{'-'*30}")

    return {
        "total": total_params,
        "unique": unique_params,
        "memory_gb": total_mem_gb
    }