ModelSizerV2.py

import argparse
from transformers import AutoConfig

# ---------------------------
# Argument parsing (once)
# ---------------------------
parser = argparse.ArgumentParser(description="Model Sizer")
parser.add_argument("--model_repo", type=str, default="openai/gpt-oss-120b")
parser.add_argument("--kv_dtype", type=str, choices=["fp8", "fp16", "fp32"], default="fp8")
parser.add_argument("--context_window", type=int)
parser.add_argument("--trust_remote_code", action="store_true")
args = parser.parse_args()

model_repo = args.model_repo
kv_cache_dtype = args.kv_dtype
user_context_window = args.context_window
trust_remote_code = args.trust_remote_code

# ---------------------------
# Config loading
# ---------------------------
def getModelConfig(repo: str, trust: bool):
    raw_config = AutoConfig.from_pretrained(repo, trust_remote_code=trust)
    return raw_config.text_config if hasattr(raw_config, "text_config") else raw_config

config = getModelConfig(model_repo, trust_remote_code)

# ---------------------------
# Common getters
# ---------------------------
def getContextWindow():
    if user_context_window is not None:
        return int(user_context_window)
    context = (
        getattr(config, "n_ctx", None)  # GPT-2 style
        or getattr(config, "context_window", None)  # Falcon style
        or getattr(config, "max_position_embeddings", None)
    )
    return int(context)

def getAttentionHeads():
    return int(
        getattr(config, "num_attention_heads", None)
        or getattr(config, "n_head", None)
    )

def getKvHeads():
    kv_heads = (
        getattr(config, "num_key_value_heads", None)
        or getattr(config, "num_kv_heads", None)
        or getattr(config, "n_kv_heads", None)
        or getattr(config, "n_head_kv", None)
    )
    if kv_heads is not None:
        return int(kv_heads)
    # fallback: MQA or MHA
    if getattr(config, "multi_query", False):
        return 1
    return getAttentionHeads()

def getHeadDimension():
    head_dim = (
        getattr(config, "head_dim", None)
        or getattr(config, "attention_head_size", None)
        or getattr(config, "dim_head", None)
        or getattr(config, "headdim", None)
    )
    if head_dim is not None:
        return int(head_dim)
    # last resort
    hidden_size = getHiddenSize()
    n_heads = getAttentionHeads()
    return int(hidden_size // n_heads) if hidden_size and n_heads else None

def getHiddenSize():
    return int(
        getattr(config, "hidden_size", None)
        or getattr(config, "n_embd", None)
        or getattr(config, "d_model", None)
    )

def getNumLayers():
    return int(getattr(config, "num_hidden_layers", 0))

def getIntermediateSize():
    return int(getattr(config, "intermediate_size", 0))

def getVocabSize():
    return int(getattr(config, "vocab_size", 0))

def usesTiedEmbeddings():
    return bool(getattr(config, "tie_word_embeddings", True))

def getNumLocalExperts():
    return int(getattr(config, "num_local_experts", 0))

def isGatedMLP():
    # Identify GLU/SwiGLU even if hidden_act is not explicit
    act = (getattr(config, "hidden_act", "") or getattr(config, "activation_function", "") or "").lower()
    if "glu" in act:
        return True
    if hasattr(config, "swiglu_limit"):
        return True
    if getattr(config, "model_type", "").lower() in {"llama", "mistral", "qwen", "gemma", "phi3", "gpt_oss"}:
        return True
    return False

# ---------------------------
# Parameter counting (weights only)
# ---------------------------
def computeAttentionParameterCount():
    """
    Q: d x (H * d_k)
    K: d x (H_kv * d_k)
    V: d x (H_kv * d_k)
    O: d x (H * d_k)
    Total per layer: 2*d*(H*d_k) + 2*d*(H_kv*d_k)
    """
    hidden_size = getHiddenSize()
    n_heads = getAttentionHeads()
    n_kv_heads = getKvHeads()
    head_dim = getHeadDimension()

    if not (hidden_size and n_heads and n_kv_heads and head_dim):
        return 0

    attn_proj_dim = n_heads * head_dim
    kv_proj_dim = n_kv_heads * head_dim
    q_and_o = 2 * hidden_size * attn_proj_dim
    k_and_v = 2 * hidden_size * kv_proj_dim
    return q_and_o + k_and_v

def computeMlpParameterCountPerExpert():
    """
    Standard MLP: 2*d*d_ff
    Gated/SwiGLU: 3*d*d_ff
    """
    hidden_size = getHiddenSize()
    d_ff = getIntermediateSize()
    if not (hidden_size and d_ff):
        return 0
    return (3 if isGatedMLP() else 2) * hidden_size * d_ff

def computeRouterParameterCountPerLayer():
    """
    Simple token router: d * num_local_experts
    """
    hidden_size = getHiddenSize()
    num_local_experts = getNumLocalExperts()
    if not (hidden_size and num_local_experts):
        return 0
    return hidden_size * num_local_experts

def computePerLayerParameterCount():
    attn_params = computeAttentionParameterCount()
    mlp_params_one_expert = computeMlpParameterCountPerExpert()
    num_local_experts = getNumLocalExperts()

    if num_local_experts > 0:
        moe_mlp_params = num_local_experts * mlp_params_one_expert
        router_params = computeRouterParameterCountPerLayer()
        return attn_params + moe_mlp_params + router_params
    else:
        return attn_params + mlp_params_one_expert

def computeEmbeddingParameterCount():
    """
    Embedding and lm_head. If untied, count both.
    """
    vocab_size = getVocabSize()
    hidden_size = getHiddenSize()
    if not (vocab_size and hidden_size):
        return 0
    base = vocab_size * hidden_size
    return base if usesTiedEmbeddings() else 2 * base

def computeTotalParameterCount():
    num_layers = getNumLayers()
    per_layer = computePerLayerParameterCount()
    embeddings = computeEmbeddingParameterCount()
    return embeddings + num_layers * per_layer

# ---------------------------
# Bytes per parameter (weights)
# ---------------------------
def getBytesPerParameter():
    """
    Returns bytes/parameter for weights.
    Includes common FP formats and block-scaled FP4 variants.
    """
    quant_cfg = getattr(config, "quantization_config", {}) or {}
    dtype_str = str(getattr(config, "dtype", "")) if getattr(config, "dtype", None) else ""
    quant_method = (quant_cfg.get("quant_method") or dtype_str or "").lower()

    mapping = {
        # dense
        "fp32": 4.0, "float32": 4.0, "bf32": 4.0, "bfloat32": 4.0, "torch.bfloat32": 4.0,
        "fp16": 2.0, "float16": 2.0, "bf16": 2.0, "bfloat16": 2.0, "torch.bfloat16": 2.0,
        "fp8": 1.0, "e4m3": 1.0, "e5m2": 1.0,
        # block-scaled FP
        "mxfp8": 1.0 + 1/32,   # +1 byte per 32 weights
        "fp4": 0.5,
        "mxfp4": 0.5 + 1/32,   # +1 byte per 32 weights
        "nvfp4": 0.5 + 1/16,   # +1 byte per 16 weights
        "nf4":   0.5 + 32/64/8 # 4.5 bpp at g=64 without double-quant (≈0.5625 bytes)
    }

    if quant_method in mapping:
        return mapping[quant_method]

    # bitsandbytes 4-bit with optional double-quant
    if quant_cfg.get("load_in_4bit", False):
        group_size = quant_cfg.get("group_size", 64) or 64
        quant_type = (quant_cfg.get("bnb_4bit_quant_type") or "").lower()  # "nf4" or "fp4"
        use_double_quant = bool(quant_cfg.get("bnb_4bit_use_double_quant", False))
        if quant_type in {"nf4", "fp4"}:
            if use_double_quant and group_size == 64:
                return 0.5159  # typical NF4+DQ@g=64
            return 0.5 + 4.0 / group_size  # FP32 constants per group by default

    # Fallback to BF16 if unknown
    return 2.0

# ---------------------------
# KV cache size (bytes)
# ---------------------------
def getBytesPerElementForKvCache():
    if kv_cache_dtype == "fp8":
        return 1
    if kv_cache_dtype == "fp16":
        return 2
    if kv_cache_dtype == "fp32":
        return 4
    return 1

def computeKvCacheSizeBytes():
    """
    KV bytes ≈ L × (2 × num_layers × n_kv_heads × head_dim × bytes_per_elem)
    """
    tokens_in_cache = getContextWindow()
    num_layers = getNumLayers()
    n_kv_heads = getKvHeads()
    head_dim = getHeadDimension()
    bytes_per_elem = getBytesPerElementForKvCache()

    per_token_bytes = 2 * num_layers * n_kv_heads * head_dim * bytes_per_elem
    return int(tokens_in_cache * per_token_bytes)

# ---------------------------
# Helpers
# ---------------------------
def bytesToGiB(num_bytes: int) -> float:
    return num_bytes / (1024 ** 3)

# ---------------------------
# Compute and print
# ---------------------------
total_parameters = computeTotalParameterCount()
bytes_per_parameter = getBytesPerParameter()
model_size_bytes = int(total_parameters * bytes_per_parameter)
kv_cache_bytes = computeKvCacheSizeBytes()

print(
    f"Model: {model_repo}\n"
    f"Context Window: {getContextWindow()}\n"
    f"Number of Layers: {getNumLayers()}\n"
    f"Number of Attention Heads: {getAttentionHeads()}\n"
    f"Number of KV Heads: {getKvHeads()}\n"
    f"Head Dimension: {getHeadDimension()}\n"
    f"Hidden Size: {getHiddenSize()}\n"
    f"Intermediate Size: {getIntermediateSize()}\n"
    f"Local Experts: {getNumLocalExperts()}\n"
    f"Tied Embeddings: {usesTiedEmbeddings()}\n"
    f"Bytes per Weight Param (quant={getattr(getattr(config, 'quantization_config', {}), 'get', lambda k, d=None: None)('quant_method', getattr(config, 'dtype', 'unknown'))}): {bytes_per_parameter}\n"
    f"Total Parameters: {total_parameters/1e9:.2f} Billion\n"
    f"Estimated Model Size: {bytesToGiB(model_size_bytes):.2f} GiB\n"
    f"Estimated KV Cache Size: {bytesToGiB(kv_cache_bytes):.2f} GiB"
)