unslothai--unsloth/studio/backend/utils/hardware/vram_estimation.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.

"""
Training VRAM estimation.

Total VRAM = weights + LoRA adapters + optimizer states + gradients
           + activations + CUDA overhead.
Activation formula from unsloth_zoo/vllm_utils.py.
All constants empirically calibrated against Llama-3.2-1B on B200.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, Optional

QUANT_4BIT_FACTOR = 16 / 5
DOUBLE_QUANT_4BIT_FACTOR = 3.6  # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1
CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3)  # calibrated on RTX 5070 Ti
NON_FLASH_ATTENTION_FACTOR = (
    12.0  # eager attention score+workspace overhead; see VRAM_ESTIMATION.md section 5
)

LINEAR_ATTENTION_IMPLS = frozenset({"flash_attention_2", "sdpa", "flex_attention"})

_SKIP_MODULE_TEXT_PREFIXES = frozenset(
    {
        "model",
        "model.model",
        "language_model",
        "language_model.model",
        "model.language_model",
        "model.language_model.model",
    }
)

DEFAULT_TARGET_MODULES = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
]
ATTENTION_TARGET_MODULES = {"q_proj", "k_proj", "v_proj", "o_proj"}
MLP_TARGET_MODULES = {"gate_proj", "up_proj", "down_proj"}

# Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale.
OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = {
    "adamw_8bit": 4,  # BNB upcasts to fp32 during step
    "paged_adamw_8bit": 4,
    "adamw_bnb_8bit": 4,
    "paged_adamw_32bit": 8,
    "adamw_torch": 6,  # fused, no master copy
    "adamw_torch_fused": 6,
    "sgd": 4,
}

# (full_ft_multiplier, lora_multiplier) — fraction of num_layers.
# LoRA: frozen layers skip activation storage, but ~1 is in flight during backprop.
GC_LAYER_MULTIPLIERS = {
    "none": (None, None),
    "true": (2.0, 1.0),
    "unsloth": (1.5, 1.0),
}


@dataclass
class ModelArchConfig:
    hidden_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_key_value_heads: int
    intermediate_size: int
    vocab_size: int
    tie_word_embeddings: bool = True
    num_experts: Optional[int] = None
    moe_intermediate_size: Optional[int] = None
    n_shared_experts: int = 0
    shared_expert_intermediate_size: Optional[int] = None
    num_experts_per_tok: int = 1
    num_dense_layers: int = 0
    q_lora_rank: Optional[int] = None
    kv_lora_rank: Optional[int] = None
    qk_nope_head_dim: Optional[int] = None
    qk_rope_head_dim: Optional[int] = None
    v_head_dim: Optional[int] = None
    head_dim: Optional[int] = None
    global_head_dim: Optional[int] = None
    num_global_key_value_heads: Optional[int] = None
    attention_k_eq_v: bool = False
    layer_types: Optional[list] = None
    num_kv_shared_layers: int = 0
    use_double_wide_mlp: bool = False
    vocab_size_per_layer_input: int = 0
    hidden_size_per_layer_input: int = 0
    quantization_skip_modules: list = field(default_factory = list)
    quant_4bit_factor: float = QUANT_4BIT_FACTOR
    moe_has_dense_mlp: bool = False
    dense_layer_indices: tuple = ()
    dense_intermediate_size: Optional[int] = None


@dataclass
class TrainingVramConfig:
    training_method: str = "qlora"
    batch_size: int = 4
    max_seq_length: int = 2048
    lora_rank: int = 16
    target_modules: list = field(default_factory = lambda: list(DEFAULT_TARGET_MODULES))
    gradient_checkpointing: str = "unsloth"
    optimizer: str = "adamw_8bit"
    load_in_4bit: bool = True
    attention_implementation: str = "flash_attention_2"


@dataclass
class VramBreakdown:
    model_weights: int
    lora_adapters: int
    optimizer_states: int
    gradients: int
    activations: int
    cuda_overhead: int
    # Equals `activations`; kept for backward compat with field consumers.
    activations_computed: int = 0

    @property
    def total(self) -> int:
        return (
            self.model_weights
            + self.lora_adapters
            + self.optimizer_states
            + self.gradients
            + self.activations
            + self.cuda_overhead
        )

    def min_gpu_vram(self, n_gpus: int) -> int:
        """Min VRAM one GPU needs: its shard + non-shardable costs.

        Weights/LoRA/optimizer/gradients shard across GPUs; activations do
        NOT (the GPU running a layer holds them).
        """
        shardable = self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients
        per_gpu_fixed = self.activations + self.cuda_overhead
        return shardable // max(n_gpus, 1) + per_gpu_fixed

    def to_gb_dict(self) -> Dict[str, float]:
        return {
            "model_weights_gb": round(self.model_weights / (1024**3), 3),
            "lora_adapters_gb": round(self.lora_adapters / (1024**3), 3),
            "optimizer_states_gb": round(self.optimizer_states / (1024**3), 3),
            "gradients_gb": round(self.gradients / (1024**3), 3),
            "activations_gb": round(self.activations / (1024**3), 3),
            "cuda_overhead_gb": round(self.cuda_overhead / (1024**3), 3),
            "total_gb": round(self.total / (1024**3), 3),
        }


def _first_scalar(value):
    # ERNIE MoE ships moe_intermediate_size / moe_num_experts as
    # [routed, shared] lists; downstream arithmetic needs the routed scalar.
    if isinstance(value, (list, tuple)):
        return value[0] if value else None
    return value


def _max_scalar(value):
    # Hunyuan-V1-MoE moe_topk can be a per-layer list; activation accounting
    # uses max top-k as a conservative upper bound.
    if isinstance(value, (list, tuple)):
        items = [v for v in value if v is not None]
        return max(items) if items else None
    return value


def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple:
    """Layer indices that use dense MLP instead of MoE. Position matters."""
    # Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite / Ernie4_5_VL_MoE
    # prefer per-position `mlp_layer_types` over prefix `first_k_dense_replace`.
    layer_types = getattr(text_config, "mlp_layer_types", None)
    if layer_types:
        return tuple(
            i for i, t in enumerate(layer_types[:total_layers]) if str(t).lower() == "dense"
        )

    # Llama4TextConfig.__init__ auto-populates self.moe_layers from
    # interleave_moe_layer_step; Llama4TextDecoderLayer dispatches via
    # `layer_idx in config.moe_layers` (modeling_llama4.py).
    llama4_moe_layers = getattr(text_config, "moe_layers", None)
    if llama4_moe_layers is not None:
        moe_indices = {int(i) for i in llama4_moe_layers}
        return tuple(i for i in range(total_layers) if i not in moe_indices)

    # ERNIE 4.5 (VL) MoE: layers via moe_layer_start/end_index + interval;
    # per-layer guard `(layer_idx+1) % interval == 0` within [start, end]
    # (modeling_ernie4_5_moe.py).
    moe_start = getattr(text_config, "moe_layer_start_index", None)
    moe_interval = getattr(text_config, "moe_layer_interval", None)
    if moe_start is not None and moe_interval is not None and int(moe_interval) > 0:
        moe_end_raw = getattr(text_config, "moe_layer_end_index", None)
        end = (
            total_layers
            if moe_end_raw is None or int(moe_end_raw) == -1
            else min(int(moe_end_raw) + 1, total_layers)
        )
        start = max(0, int(moe_start))
        interval = int(moe_interval)
        moe_indices = {i for i in range(start, end) if (i + 1) % interval == 0}
        return tuple(i for i in range(total_layers) if i not in moe_indices)

    first_k = getattr(text_config, "first_k_dense_replace", None)
    if first_k is not None:
        return tuple(range(min(int(first_k), total_layers)))

    sparse_step = getattr(text_config, "decoder_sparse_step", None)
    mlp_only = getattr(text_config, "mlp_only_layers", None) or []
    if sparse_step is not None and sparse_step > 0:
        mlp_only_set = {int(i) for i in mlp_only}
        return tuple(
            i for i in range(total_layers) if i in mlp_only_set or (i + 1) % sparse_step != 0
        )
    return ()


def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
    text_config = getattr(hf_config, "text_config", None) or hf_config
    quantization_config = getattr(hf_config, "quantization_config", None) or {}
    if not isinstance(quantization_config, dict):
        quantization_config = getattr(quantization_config, "to_dict", lambda: {})()
    quant_4bit_factor = (
        DOUBLE_QUANT_4BIT_FACTOR
        if quantization_config.get("bnb_4bit_use_double_quant", False)
        else QUANT_4BIT_FACTOR
    )

    hidden_size = getattr(text_config, "hidden_size", None)
    num_layers = getattr(text_config, "num_hidden_layers", None)
    num_heads = getattr(text_config, "num_attention_heads", None)
    intermediate_size = getattr(text_config, "intermediate_size", None)
    vocab_size = getattr(text_config, "vocab_size", None)

    if isinstance(intermediate_size, (list, tuple)):
        intermediate_size = intermediate_size[0] if intermediate_size else None
    if intermediate_size is None and hidden_size is not None:
        intermediate_size = hidden_size * 4

    if not all(
        v is not None for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size)
    ):
        return None
    if num_heads <= 0:
        return None

    num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads)

    # DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe
    # ffn_config as a secondary source so DBRX isn't misclassified as dense.
    ffn_config = getattr(text_config, "ffn_config", None)

    def _moe_attr(name):
        value = getattr(text_config, name, None)
        if value is None and ffn_config is not None:
            value = getattr(ffn_config, name, None)
        return value

    num_experts = None
    for attr in (
        "num_local_experts",
        "num_experts",
        "n_routed_experts",
        "moe_num_experts",
    ):
        num_experts = _first_scalar(_moe_attr(attr))
        if num_experts is not None:
            break

    moe_intermediate_raw = _moe_attr("moe_intermediate_size")
    if moe_intermediate_raw is None:
        moe_intermediate_raw = _moe_attr("ffn_hidden_size")
    moe_intermediate = _first_scalar(moe_intermediate_raw)
    # Exaone-MoE / ERNIE alias num_shared_experts / moe_num_shared_experts
    # to the canonical n_shared_experts.
    n_shared_experts = (
        _first_scalar(_moe_attr("n_shared_experts"))
        or _first_scalar(_moe_attr("num_shared_experts"))
        or _first_scalar(_moe_attr("moe_num_shared_experts"))
        or 0
    )
    shared_expert_intermediate_size = _moe_attr("shared_expert_intermediate_size")
    if shared_expert_intermediate_size and n_shared_experts == 0:
        n_shared_experts = 1
    # DBRX moe_top_k; Hunyuan-V1-MoE moe_topk (may be a per-layer list).
    # _max_scalar normalizes lists to the worst case so int(...) can't crash.
    num_experts_per_tok = (
        _max_scalar(_moe_attr("num_experts_per_tok"))
        or _max_scalar(_moe_attr("top_k_experts"))
        or _max_scalar(_moe_attr("moe_top_k"))
        or _max_scalar(_moe_attr("moe_topk"))
        or 1
    )

    dense_layer_indices: tuple = ()
    if num_experts is not None and num_experts > 1:
        dense_layer_indices = _compute_dense_layer_indices(text_config, num_layers)
    num_dense_layers = len(dense_layer_indices)

    # Llama4 dense layers use intermediate_size_mlp; experts use
    # intermediate_size. One shared_expert per MoE layer (modeling_llama4.py).
    intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp"))
    dense_intermediate_size = (
        int(intermediate_size_mlp_raw) if intermediate_size_mlp_raw is not None else None
    )
    if (
        intermediate_size_mlp_raw is not None
        and num_experts is not None
        and num_experts > 1
        and shared_expert_intermediate_size is None
        and n_shared_experts == 0
    ):
        n_shared_experts = 1

    q_lora_rank = getattr(text_config, "q_lora_rank", None)
    kv_lora_rank = getattr(text_config, "kv_lora_rank", None)
    qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", None)
    qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", None)
    v_head_dim = getattr(text_config, "v_head_dim", None)

    return ModelArchConfig(
        hidden_size = hidden_size,
        num_hidden_layers = num_layers,
        num_attention_heads = num_heads,
        num_key_value_heads = num_kv_heads,
        intermediate_size = intermediate_size,
        vocab_size = vocab_size,
        tie_word_embeddings = getattr(text_config, "tie_word_embeddings", True),
        num_experts = num_experts,
        moe_intermediate_size = moe_intermediate,
        n_shared_experts = n_shared_experts,
        shared_expert_intermediate_size = shared_expert_intermediate_size,
        num_experts_per_tok = int(num_experts_per_tok),
        num_dense_layers = num_dense_layers,
        q_lora_rank = q_lora_rank,
        kv_lora_rank = kv_lora_rank,
        qk_nope_head_dim = qk_nope_head_dim,
        qk_rope_head_dim = qk_rope_head_dim,
        v_head_dim = v_head_dim,
        head_dim = getattr(text_config, "head_dim", None),
        global_head_dim = getattr(text_config, "global_head_dim", None),
        num_global_key_value_heads = getattr(
            text_config,
            "num_global_key_value_heads",
            None,
        ),
        attention_k_eq_v = bool(getattr(text_config, "attention_k_eq_v", False)),
        layer_types = getattr(text_config, "layer_types", None),
        num_kv_shared_layers = getattr(text_config, "num_kv_shared_layers", None) or 0,
        use_double_wide_mlp = bool(getattr(text_config, "use_double_wide_mlp", False)),
        vocab_size_per_layer_input = getattr(
            text_config,
            "vocab_size_per_layer_input",
            None,
        )
        or 0,
        hidden_size_per_layer_input = getattr(
            text_config,
            "hidden_size_per_layer_input",
            None,
        )
        or 0,
        quantization_skip_modules = list(quantization_config.get("llm_int8_skip_modules", []) or []),
        quant_4bit_factor = quant_4bit_factor,
        moe_has_dense_mlp = bool(getattr(text_config, "enable_moe_block", False)),
        dense_layer_indices = dense_layer_indices,
        dense_intermediate_size = dense_intermediate_size,
    )


def _targets_all_linear(target_modules) -> bool:
    # peft LoraConfig accepts target_modules="all-linear" as a bare string;
    # iterating a string yields chars and never matches the set.
    if isinstance(target_modules, str):
        target_modules = [target_modules]
    normalized = {str(module).lower().replace("_", "-") for module in target_modules}
    return normalized == {"all-linear"}


def _head_dim(arch: ModelArchConfig) -> int:
    return arch.head_dim or arch.hidden_size // arch.num_attention_heads


def _layer_types(arch: ModelArchConfig) -> list:
    if arch.layer_types and len(arch.layer_types) == arch.num_hidden_layers:
        return arch.layer_types
    return ["full_attention"] * arch.num_hidden_layers


def _uses_structured_layer_shapes(arch: ModelArchConfig) -> bool:
    # MLA configs have their own q/kv low-rank projection shape formulas in
    # _compute_attn_elements / _lora_attn_elements; do not let head_dim or
    # other structured fields override that path.
    if arch.q_lora_rank is not None:
        return False
    return bool(
        arch.layer_types
        or arch.head_dim is not None
        or arch.global_head_dim is not None
        or arch.num_global_key_value_heads is not None
        or arch.attention_k_eq_v
        or arch.num_kv_shared_layers > 0
        or arch.use_double_wide_mlp
    )


def _is_kv_shared_layer(arch: ModelArchConfig, layer_idx: int) -> bool:
    if arch.num_kv_shared_layers <= 0:
        return False
    first_shared = arch.num_hidden_layers - arch.num_kv_shared_layers
    # Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863) uses the same
    # `> 0` guard so a fully-shared config raises at model construction;
    # matching upstream avoids estimating a shape the model code rejects.
    return layer_idx >= first_shared > 0


def _is_dense_mlp_layer(arch: ModelArchConfig, layer_idx: int) -> bool:
    if arch.dense_layer_indices:
        return layer_idx in arch.dense_layer_indices
    return layer_idx < arch.num_dense_layers


def _per_layer_input_quantizable(arch: ModelArchConfig) -> int:
    # Gemma4 PLE block adds per_layer_model_projection (single Linear),
    # per_layer_input_gate (per layer), and per_layer_projection (per layer);
    # see gemma4/modular_gemma4.py:1077-1083 and :1247-1253.
    pli = arch.hidden_size_per_layer_input
    if pli <= 0:
        return 0
    n_layers = arch.num_hidden_layers
    hd = arch.hidden_size
    return hd * (n_layers * pli) + (hd * pli) * n_layers + (pli * hd) * n_layers


def _per_layer_input_norm_elements(arch: ModelArchConfig) -> int:
    pli = arch.hidden_size_per_layer_input
    if pli <= 0:
        return 0
    n_layers = arch.num_hidden_layers
    hd = arch.hidden_size
    return hd * n_layers + pli


def _per_layer_input_lora_params(arch: ModelArchConfig, r: int, target_modules) -> int:
    # get_peft_regex requires a component tag (mlp/attn/...); PLE names lack
    # one, so all-linear skips them. Count PLE LoRA only when named explicitly.
    pli = arch.hidden_size_per_layer_input
    if pli <= 0:
        return 0
    targets = {target_modules} if isinstance(target_modules, str) else set(target_modules or [])
    n_layers = arch.num_hidden_layers
    hd = arch.hidden_size
    total = 0
    if "per_layer_model_projection" in targets:
        total += hd * r + r * (n_layers * pli)
    if "per_layer_input_gate" in targets:
        total += (hd * r + r * pli) * n_layers
    if "per_layer_projection" in targets:
        total += (pli * r + r * hd) * n_layers
    return total


def _layer_attention_dims(arch: ModelArchConfig, layer_idx: int) -> tuple:
    layer_types = _layer_types(arch)
    layer_type = layer_types[layer_idx]
    is_sliding = layer_type == "sliding_attention"
    head_dim = arch.global_head_dim if not is_sliding and arch.global_head_dim else _head_dim(arch)
    use_alt_attention = arch.attention_k_eq_v and not is_sliding
    num_kv_heads = (
        arch.num_global_key_value_heads
        if use_alt_attention and arch.num_global_key_value_heads
        else arch.num_key_value_heads
    )
    q_size = arch.num_attention_heads * head_dim
    kv_size = num_kv_heads * head_dim
    has_k = not _is_kv_shared_layer(arch, layer_idx)
    has_v = has_k and not use_alt_attention
    return q_size, kv_size, has_k, has_v


def _layer_mlp_size(arch: ModelArchConfig, layer_idx: int) -> int:
    if arch.use_double_wide_mlp and _is_kv_shared_layer(arch, layer_idx):
        return _dense_mlp_size(arch) * 2
    return _dense_mlp_size(arch)


def _text_linear_dims(arch: ModelArchConfig, layer_idx: int) -> Dict[str, tuple[int, int]]:
    hd = arch.hidden_size
    if _uses_structured_layer_shapes(arch):
        q_size, kv_size, has_k, has_v = _layer_attention_dims(arch, layer_idx)
        mlp_size = _layer_mlp_size(arch, layer_idx)
    else:
        q_size = hd
        kv_size = _get_kv_size(arch)
        has_k = True
        has_v = True
        mlp_size = _get_mlp_size(arch)

    dims = {
        "q_proj": (hd, q_size),
        "o_proj": (q_size, hd),
    }
    if has_k:
        dims["k_proj"] = (hd, kv_size)
    if has_v:
        dims["v_proj"] = (hd, kv_size)

    dims.update(
        {
            "gate_proj": (hd, mlp_size),
            "up_proj": (hd, mlp_size),
            "down_proj": (mlp_size, hd),
        }
    )
    return dims


def _module_path_matches(skip_module: str, alias: str) -> bool:
    skip_parts = [part for part in skip_module.split(".") if part]
    alias_parts = [part for part in alias.split(".") if part]
    if not skip_parts or not alias_parts:
        return False
    if alias_parts[0] == "layers":
        return skip_parts == alias_parts
    if len(skip_parts) <= len(alias_parts):
        # BNB suffix-matches short skip entries (["q_proj"], ["lm_head"]) so a
        # skip shorter than the alias is a tail match.
        return alias_parts[-len(skip_parts) :] == skip_parts
    if skip_parts[-len(alias_parts) :] != alias_parts:
        return False
    prefix_parts = skip_parts[: len(skip_parts) - len(alias_parts)]
    if not prefix_parts:
        return True
    # Bound the prefix to text-tower roots so VLM skips like
    # vision_tower.model.layers... don't shadow the text alias.
    return ".".join(prefix_parts) in _SKIP_MODULE_TEXT_PREFIXES


def _add_module_aliases(aliases: Dict[str, str], canonical: str, suffix: str) -> None:
    for prefix in (
        "",
        "model",
        "model.model",
        "language_model",
        "language_model.model",
        "model.language_model",
        "model.language_model.model",
    ):
        alias = f"{prefix}.{suffix}" if prefix else suffix
        aliases[alias] = canonical


def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], Dict[str, str]]:
    elements: Dict[str, int] = {}
    aliases: Dict[str, str] = {}

    is_mla = arch.q_lora_rank is not None and not _uses_structured_layer_shapes(arch)
    pli = arch.hidden_size_per_layer_input
    hd_global = arch.hidden_size

    for layer_idx in range(arch.num_hidden_layers):
        layer_modules: Dict[str, int] = {}
        dims = _text_linear_dims(arch, layer_idx)
        attn_dims = {name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES}
        mlp_dims = {name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES}

        if is_mla:
            # MLA splits q/o into q_a/q_b/kv_a/kv_b; emit a single self_attn
            # aggregate at the authoritative MLA per-layer total.
            layer_modules["self_attn"] = _compute_attn_elements(arch)
        else:
            for name, (in_dim, out_dim) in attn_dims.items():
                layer_modules[f"self_attn.{name}"] = in_dim * out_dim

        if arch.num_experts and arch.num_experts > 1:
            if _is_dense_mlp_layer(arch, layer_idx):
                layer_modules.update(
                    {
                        f"mlp.{name}": in_dim * out_dim
                        for name, (in_dim, out_dim) in mlp_dims.items()
                    }
                )
            else:
                layer_modules["mlp.experts"] = _compute_routed_moe_elements(arch)
                shared_moe = _compute_shared_moe_elements(arch)
                if shared_moe:
                    # Qwen3.5-MoE: mlp.shared_expert; Exaone-MoE/Laguna/GLM:
                    # mlp.shared_experts. Register both so skip_modules match.
                    layer_modules["mlp.shared_expert"] = shared_moe
                if arch.moe_has_dense_mlp:
                    # enable_moe_block runs dense MLP and experts in parallel;
                    # register both. Non-structured _get_mlp_size prefers
                    # moe_intermediate_size, so rebuild dense dims directly.
                    if _uses_structured_layer_shapes(arch):
                        dense_dims = mlp_dims
                    else:
                        hd = arch.hidden_size
                        inter = arch.intermediate_size
                        dense_dims = {
                            "gate_proj": (hd, inter),
                            "up_proj": (hd, inter),
                            "down_proj": (inter, hd),
                        }
                    layer_modules.update(
                        {
                            f"mlp.{name}": in_dim * out_dim
                            for name, (in_dim, out_dim) in dense_dims.items()
                        }
                    )
        else:
            layer_modules.update(
                {f"mlp.{name}": in_dim * out_dim for name, (in_dim, out_dim) in mlp_dims.items()}
            )

        if pli > 0:
            # Register PLE per-layer linears so llm_int8_skip_modules entries
            # like model.layers.0.per_layer_input_gate match.
            layer_modules["per_layer_input_gate"] = hd_global * pli
            layer_modules["per_layer_projection"] = pli * hd_global

        attn_total = sum(
            value
            for name, value in layer_modules.items()
            if name == "self_attn" or name.startswith("self_attn.")
        )
        # gemma4 enable_moe_block puts routed experts at sibling
        # layers.<i>.experts, not under self.mlp; keep the "mlp" aggregate to
        # the dense path so a `model.layers.0.mlp` skip doesn't over-skip.
        is_sibling_experts = bool(arch.moe_has_dense_mlp)
        mlp_total = sum(
            value
            for name, value in layer_modules.items()
            if (
                name == "mlp"
                or (name.startswith("mlp.") and not (is_sibling_experts and name == "mlp.experts"))
            )
        )
        experts_total = layer_modules.get("mlp.experts", 0) if is_sibling_experts else 0
        layer_total = sum(layer_modules.values())

        aggregate_modules = {
            f"text.layers.{layer_idx}": layer_total,
            f"text.layers.{layer_idx}.self_attn": attn_total,
            f"text.layers.{layer_idx}.mlp": mlp_total,
        }
        if experts_total:
            aggregate_modules[f"text.layers.{layer_idx}.experts"] = experts_total
        elements.update(aggregate_modules)
        for canonical in aggregate_modules:
            suffix = canonical.removeprefix("text.")
            _add_module_aliases(aliases, canonical, suffix)

        for name, value in layer_modules.items():
            canonical = f"text.layers.{layer_idx}.{name}"
            elements[canonical] = value
            _add_module_aliases(aliases, canonical, canonical.removeprefix("text."))
            if name == "mlp.experts" and arch.moe_has_dense_mlp:
                # gemma4: routed experts at sibling layers.<i>.experts, not mlp.
                _add_module_aliases(aliases, canonical, f"layers.{layer_idx}.experts")
            elif name == "mlp.shared_expert":
                # Exaone-MoE/Laguna/GLM use plural `shared_experts`; add both.
                _add_module_aliases(
                    aliases,
                    canonical,
                    f"layers.{layer_idx}.mlp.shared_experts",
                )

    if pli > 0:
        canonical = "text.per_layer_model_projection"
        elements[canonical] = hd_global * (arch.num_hidden_layers * pli)
        _add_module_aliases(aliases, canonical, canonical.removeprefix("text."))

    return elements, aliases


def _compute_skipped_quantizable_elements(arch: ModelArchConfig) -> int:
    if not arch.quantization_skip_modules:
        return 0

    module_elements, aliases = _build_text_module_elements(arch)
    matched = set()
    for skip_module in arch.quantization_skip_modules:
        for alias, canonical in aliases.items():
            if _module_path_matches(skip_module, alias):
                matched.add(canonical)

    pruned = {
        canonical
        for canonical in matched
        if not any(canonical != parent and canonical.startswith(f"{parent}.") for parent in matched)
    }
    return sum(module_elements[canonical] for canonical in pruned)


def _get_kv_size(arch: ModelArchConfig) -> int:
    return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads


def _get_mlp_size(arch: ModelArchConfig) -> int:
    if arch.moe_intermediate_size is not None:
        return arch.moe_intermediate_size
    return arch.intermediate_size


def _dense_mlp_size(arch: ModelArchConfig) -> int:
    # Llama4 dense layers use intermediate_size_mlp; routed/shared experts use
    # intermediate_size. Other configs leave the field None.
    return arch.dense_intermediate_size or arch.intermediate_size


def _get_num_experts(arch: ModelArchConfig) -> int:
    return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1


def _compute_attn_elements(arch: ModelArchConfig) -> int:
    """Attention weight elements per layer."""
    hd = arch.hidden_size
    if arch.q_lora_rank is not None:
        nh = arch.num_attention_heads
        qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
        q_a = hd * arch.q_lora_rank
        q_b = arch.q_lora_rank * (nh * qk_head)
        kv_a = hd * (arch.kv_lora_rank + arch.qk_rope_head_dim)
        kv_b = arch.kv_lora_rank * (nh * (arch.qk_nope_head_dim + arch.v_head_dim))
        o = (nh * arch.v_head_dim) * hd
        norms = arch.q_lora_rank + arch.kv_lora_rank
        return q_a + q_b + kv_a + kv_b + o + norms
    kv_size = _get_kv_size(arch)
    return (hd + kv_size + kv_size + hd) * hd


def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int:
    return arch.hidden_size * _dense_mlp_size(arch) * 3


def _shared_expert_size(arch: ModelArchConfig) -> int:
    # Qwen3.5-MoE shared expert has its own intermediate_size (default 512)
    # distinct from moe_intermediate_size; fall back to routed mlp_size for
    # families that share it (deepseek-style configs).
    return arch.shared_expert_intermediate_size or _get_mlp_size(arch)


def _compute_routed_moe_elements(arch: ModelArchConfig) -> int:
    hd = arch.hidden_size
    n_experts = _get_num_experts(arch)
    return hd * _get_mlp_size(arch) * 3 * n_experts + n_experts * hd


def _compute_shared_moe_elements(arch: ModelArchConfig) -> int:
    if not arch.n_shared_experts:
        return 0
    hd = arch.hidden_size
    shared_size = _shared_expert_size(arch)
    total = hd * shared_size * 3 * arch.n_shared_experts
    # Only Qwen2/Qwen3.5-MoE add a shared_expert_gate Linear (hidden_size->1);
    # shared_expert_intermediate_size is the Qwen-style discriminator.
    if arch.shared_expert_intermediate_size:
        total += arch.n_shared_experts * hd
    return total


def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int:
    return _compute_routed_moe_elements(arch) + _compute_shared_moe_elements(arch)


def _compute_layer_elements(arch: ModelArchConfig):
    """Return (total_quantizable, layernorms_per_layer, embed, lm_head) element counts.

    total_quantizable is summed across ALL layers (not per-layer).
    """
    hd = arch.hidden_size
    n_layers = arch.num_hidden_layers
    n_experts = _get_num_experts(arch)

    if _uses_structured_layer_shapes(arch):
        attn_total = 0
        per_layer_dense_mlp = []
        for layer_idx in range(n_layers):
            layer_dense_mlp = 0
            for name, (in_dim, out_dim) in _text_linear_dims(
                arch,
                layer_idx,
            ).items():
                elements = in_dim * out_dim
                if name in ATTENTION_TARGET_MODULES:
                    attn_total += elements
                elif name in MLP_TARGET_MODULES:
                    layer_dense_mlp += elements
            per_layer_dense_mlp.append(layer_dense_mlp)
        if n_experts > 1:
            n_dense = arch.num_dense_layers
            n_moe = n_layers - n_dense
            moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe
            if arch.moe_has_dense_mlp:
                # enable_moe_block runs dense MLP and MoE experts in parallel;
                # count dense for every layer alongside MoE.
                mlp_total = sum(per_layer_dense_mlp) + moe_mlp_total
            else:
                dense_only_total = sum(
                    value
                    for i, value in enumerate(per_layer_dense_mlp)
                    if _is_dense_mlp_layer(arch, i)
                )
                mlp_total = moe_mlp_total + dense_only_total
        else:
            mlp_total = sum(per_layer_dense_mlp)
    elif n_experts > 1:
        attn_total = _compute_attn_elements(arch) * n_layers
        n_dense = arch.num_dense_layers
        n_moe = n_layers - n_dense
        moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe
        if arch.moe_has_dense_mlp:
            mlp_total = _compute_dense_mlp_elements(arch) * n_layers + moe_mlp_total
        else:
            mlp_total = moe_mlp_total + _compute_dense_mlp_elements(arch) * n_dense
    else:
        attn_total = _compute_attn_elements(arch) * n_layers
        mlp_total = _compute_dense_mlp_elements(arch) * n_layers

    layernorms = 2 * hd
    per_layer_embed = arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers
    ple_text_linear = _per_layer_input_quantizable(arch)
    ple_norms = _per_layer_input_norm_elements(arch)
    embed_tokens = arch.vocab_size * hd + per_layer_embed + ple_norms
    lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd
    return attn_total + mlp_total + ple_text_linear, layernorms, embed_tokens, lm_head


def compute_model_weights_bytes(
    arch: ModelArchConfig, training_method: str, load_in_4bit: bool
) -> int:
    total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
    n_layers = arch.num_hidden_layers
    non_quantizable = layernorms * n_layers + embed_tokens + lm_head

    if training_method == "qlora" and load_in_4bit:
        skipped_quantizable = min(
            _compute_skipped_quantizable_elements(arch),
            total_quantizable,
        )
        quantized = total_quantizable - skipped_quantizable
        return int(
            quantized * 2 / arch.quant_4bit_factor + skipped_quantizable * 2 + non_quantizable * 2
        )

    return int((total_quantizable + non_quantizable) * 2)


def compute_total_params(arch: ModelArchConfig) -> int:
    total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
    n_layers = arch.num_hidden_layers
    return total_quantizable + layernorms * n_layers + embed_tokens + lm_head


def _lora_attn_elements(arch: ModelArchConfig, r: int, target_modules: list) -> int:
    hd = arch.hidden_size
    if arch.q_lora_rank is not None:
        # MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o
        nh = arch.num_attention_heads
        qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
        kv_out = nh * (arch.qk_nope_head_dim + arch.v_head_dim)
        o_in = nh * arch.v_head_dim
        dims = {
            "q_proj": (arch.q_lora_rank, nh * qk_head),
            "k_proj": (hd, arch.kv_lora_rank + arch.qk_rope_head_dim),
            "v_proj": (arch.kv_lora_rank, kv_out),
            "o_proj": (o_in, hd),
        }
    else:
        kv_size = _get_kv_size(arch)
        dims = {
            "q_proj": (hd, hd),
            "k_proj": (hd, kv_size),
            "v_proj": (hd, kv_size),
            "o_proj": (hd, hd),
        }
    total = 0
    for name, (in_dim, out_dim) in dims.items():
        if name in target_modules:
            total += in_dim * r + r * out_dim
    return total


def _lora_mlp_elements(
    hd: int, mlp_size: int, r: int, target_modules: list, expert_mult: int
) -> int:
    module_ab = {
        "gate_proj": (hd * r, r * mlp_size),
        "up_proj": (hd * r, r * mlp_size),
        "down_proj": (mlp_size * r, r * hd),
    }
    total = 0
    for name, (a, b) in module_ab.items():
        if name in target_modules:
            total += (a + b) * expert_mult
    return total


def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: list) -> int:
    all_linear = _targets_all_linear(target_modules)
    selected_modules = list(DEFAULT_TARGET_MODULES) if all_linear else target_modules
    hd = arch.hidden_size
    r = lora_rank
    n_layers = arch.num_hidden_layers
    n_experts = _get_num_experts(arch)

    use_structured_shapes = _uses_structured_layer_shapes(arch)
    if use_structured_shapes:
        attn_total = 0
        structured_dense_mlp = 0
        per_layer_dense_mlp = []
        for layer_idx in range(n_layers):
            layer_dense = 0
            for name, (in_dim, out_dim) in _text_linear_dims(
                arch,
                layer_idx,
            ).items():
                if name not in selected_modules:
                    continue
                if name in ATTENTION_TARGET_MODULES:
                    attn_total += in_dim * r + r * out_dim
                elif name in MLP_TARGET_MODULES:
                    layer_dense += in_dim * r + r * out_dim
            per_layer_dense_mlp.append(layer_dense)
            structured_dense_mlp += layer_dense
        if n_experts > 1:
            n_dense = arch.num_dense_layers
            n_moe = n_layers - n_dense
            # peft "all-linear" attaches LoRA to nn.Linear only; routed experts
            # are nn.Parameter and need explicit gate_proj/up_proj/down_proj
            # naming via Unsloth's get_moe_target_parameters. Shared experts are
            # nn.Linear, picked up by get_peft_regex.
            routed_moe = (
                0
                if all_linear
                else _lora_mlp_elements(
                    hd,
                    _get_mlp_size(arch),
                    r,
                    selected_modules,
                    n_experts,
                )
            )
            shared_moe = _lora_mlp_elements(
                hd,
                _shared_expert_size(arch),
                r,
                selected_modules,
                arch.n_shared_experts,
            )
            moe_mlp = routed_moe + shared_moe
            if arch.moe_has_dense_mlp:
                # Parallel dense MLP coexists with MoE on every layer.
                mlp_total = structured_dense_mlp + moe_mlp * n_moe
            else:
                dense_only = sum(
                    value
                    for i, value in enumerate(per_layer_dense_mlp)
                    if _is_dense_mlp_layer(arch, i)
                )
                mlp_total = moe_mlp * n_moe + dense_only
        else:
            mlp_total = structured_dense_mlp
        return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules)
    elif n_experts > 1:
        attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers
        n_dense = arch.num_dense_layers
        n_moe = n_layers - n_dense
        # Routed and shared experts may use different intermediate sizes
        # (Qwen3.5-MoE: routed mlp_size != shared_expert_intermediate_size).
        # See structured branch for the all-linear exclusion rationale; only
        # routed (nn.Parameter) experts are excluded under all-linear.
        routed_moe = (
            0
            if all_linear
            else _lora_mlp_elements(
                hd,
                _get_mlp_size(arch),
                r,
                selected_modules,
                n_experts,
            )
        )
        shared_moe = _lora_mlp_elements(
            hd,
            _shared_expert_size(arch),
            r,
            selected_modules,
            arch.n_shared_experts,
        )
        moe_mlp = routed_moe + shared_moe
        dense_mlp = _lora_mlp_elements(
            hd,
            _dense_mlp_size(arch),
            r,
            selected_modules,
            1,
        )
        if arch.moe_has_dense_mlp:
            mlp_total = moe_mlp * n_moe + dense_mlp * n_layers
        else:
            mlp_total = moe_mlp * n_moe + dense_mlp * n_dense
    else:
        attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers
        mlp_total = (
            _lora_mlp_elements(
                hd,
                _dense_mlp_size(arch),
                r,
                selected_modules,
                1,
            )
            * n_layers
        )

    return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules)


def compute_lora_adapter_bytes(lora_params: int) -> int:
    return lora_params * 2


def compute_optimizer_bytes(trainable_params: int, optimizer: str) -> int:
    optimizer_key = optimizer.lower().replace("-", "_")
    bytes_per_param = OPTIMIZER_BYTES_PER_PARAM.get(optimizer_key, 4)
    return trainable_params * bytes_per_param


def compute_gradient_bytes(trainable_params: int) -> int:
    return trainable_params * 2


def _is_linear_attention(attention_implementation: Optional[str]) -> bool:
    # PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only
    # eager (and other non-flash impls) need the quadratic correction.
    return attention_implementation in LINEAR_ATTENTION_IMPLS


def _compute_non_flash_attention_bytes(
    arch: ModelArchConfig, batch_size: int, seq_len: int, effective_layers: float
) -> int:
    score_elements = batch_size * arch.num_attention_heads * seq_len * seq_len
    return int(score_elements * 2 * NON_FLASH_ATTENTION_FACTOR * effective_layers)


def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple:
    n_experts = _get_num_experts(arch)
    is_moe_layer = n_experts > 1 and not _is_dense_mlp_layer(arch, layer_idx)
    if _uses_structured_layer_shapes(arch):
        q_size, kv_size, _has_k, _has_v = _layer_attention_dims(arch, layer_idx)
        # KV-shared layers (Gemma4/Gemma3n) drop k/v WEIGHTS but the donor's
        # K/V tensors stay alive, so activations still pay kv_size; only the
        # weight path uses has_k/has_v.
        layer_type = _layer_types(arch)[layer_idx]
        use_alt_attention = arch.attention_k_eq_v and layer_type != "sliding_attention"
        kv_count = 1 if use_alt_attention else 2
        qkv_size = q_size + kv_size * kv_count
        if is_moe_layer:
            # Each token routes through num_experts_per_tok experts; all their
            # gate/up/down intermediates are live during MLP forward.
            mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok
            if arch.n_shared_experts:
                mlp_size += _shared_expert_size(arch) * arch.n_shared_experts
            if arch.moe_has_dense_mlp:
                mlp_size += _layer_mlp_size(arch, layer_idx)
        else:
            mlp_size = _layer_mlp_size(arch, layer_idx)
        return qkv_size, mlp_size
    kv_size = _get_kv_size(arch)
    if is_moe_layer:
        mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok
        if arch.n_shared_experts:
            mlp_size += _shared_expert_size(arch) * arch.n_shared_experts
        if arch.moe_has_dense_mlp:
            mlp_size += arch.intermediate_size
    else:
        mlp_size = _get_mlp_size(arch)
    return arch.hidden_size + kv_size + kv_size, mlp_size


def _per_layer_activation_bytes(
    arch: ModelArchConfig, layer_idx: int, batch_size: int, seq_len: int
) -> int:
    qkv_size, mlp_size = _layer_qkv_mlp_sizes(arch, layer_idx)
    activation_qkv = seq_len * batch_size * qkv_size
    residual_memory = (seq_len * batch_size) * 2
    activation_mlp = seq_len * batch_size * (mlp_size + mlp_size)
    # PLE gate (hd) + projection (pli) outputs materialize once per decoder
    # layer when hidden_size_per_layer_input is set (gemma4 modular:1141-1145).
    pli = arch.hidden_size_per_layer_input
    activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0
    return int((activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25)


def compute_activation_bytes(
    arch: ModelArchConfig,
    batch_size: int,
    seq_len: int,
    gradient_checkpointing: str,
    is_lora: bool = False,
    attention_implementation: Optional[str] = "flash_attention_2",
) -> int:
    n_layers = arch.num_hidden_layers

    gc_key = gradient_checkpointing.lower()
    gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None))
    full_ft_mult, lora_mult = gc_entry
    gc_multiplier = lora_mult if is_lora else full_ft_mult

    if gc_multiplier is None:
        effective_layers = n_layers
        linear_bytes = sum(
            _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers)
        )
    else:
        effective_layers = gc_multiplier
        max_layer_bytes = max(
            _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers)
        )
        linear_bytes = int(max_layer_bytes * effective_layers)

    # gemma4 per_layer_model_projection runs once outside the per-decoder loop
    # and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247.
    pli = arch.hidden_size_per_layer_input
    if pli > 0:
        linear_bytes += int(seq_len * batch_size * n_layers * pli * 2 * 1.25)

    if _is_linear_attention(attention_implementation):
        return linear_bytes
    return max(
        linear_bytes,
        _compute_non_flash_attention_bytes(
            arch,
            batch_size,
            seq_len,
            effective_layers,
        ),
    )


def estimate_training_vram(arch: ModelArchConfig, config: TrainingVramConfig) -> VramBreakdown:
    method = config.training_method.lower()
    is_lora = method in ("qlora", "lora")
    load_in_4bit = config.load_in_4bit or method == "qlora"

    model_weights = compute_model_weights_bytes(arch, method, load_in_4bit)

    lora_params = 0
    lora_adapter_bytes = 0
    if is_lora:
        lora_params = compute_lora_params(
            arch,
            config.lora_rank,
            config.target_modules,
        )
        lora_adapter_bytes = compute_lora_adapter_bytes(lora_params)

    trainable_params = lora_params if is_lora else compute_total_params(arch)
    optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer)
    activations_computed = compute_activation_bytes(
        arch,
        config.batch_size,
        config.max_seq_length,
        config.gradient_checkpointing,
        is_lora = is_lora,
        attention_implementation = config.attention_implementation,
    )
    raw_gradient_bytes = compute_gradient_bytes(trainable_params)
    gradient_floor = int(model_weights * 0.15)
    if is_lora:
        gradient_floor = min(
            gradient_floor,
            max(activations_computed, optimizer_bytes),
        )
    gradient_bytes = max(raw_gradient_bytes, gradient_floor)
    activation_bytes = activations_computed

    return VramBreakdown(
        model_weights = model_weights,
        lora_adapters = lora_adapter_bytes,
        optimizer_states = optimizer_bytes,
        gradients = gradient_bytes,
        activations = activation_bytes,
        cuda_overhead = CUDA_OVERHEAD_BYTES,
        activations_computed = activations_computed,
    )