heapr/model_utils.py

"""Model loading and Laguna sparse-layer discovery."""

from __future__ import annotations

from dataclasses import asdict
from pathlib import Path
from typing import Any

from heapr.constants import DEFAULT_PRUNE_MODEL, DEFAULT_SMOKE_MODEL
from heapr.typing import RunMetadata, SparseLayerInfo
from heapr.utils import (
    collect_dependency_versions,
    collect_hardware_metadata,
    ensure_dir,
    require_torch,
    write_json,
)


def load_tokenizer(model_id_or_path: str, revision: str | None = None):
    from transformers import AutoTokenizer

    return AutoTokenizer.from_pretrained(
        model_id_or_path,
        revision=revision,
        trust_remote_code=True,
    )


def load_causal_lm(
    model_id_or_path: str,
    *,
    revision: str | None = None,
    dtype: str = "bfloat16",
    device_map: str | dict[str, Any] = "auto",
    max_memory: dict[str, str] | None = None,
    offload_folder: str | None = None,
    attn_implementation: str | None = None,
    use_cache: bool = False,
    cache_implementation: str | None = None,
    output_router_logits: bool = False,
):
    """Load a Laguna causal LM through Transformers remote code."""

    torch = require_torch()
    from transformers import AutoModelForCausalLM

    dtype_value = getattr(torch, dtype) if isinstance(dtype, str) else dtype
    kwargs: dict[str, Any] = {
        "revision": revision,
        "trust_remote_code": True,
        "dtype": dtype_value,
        "device_map": device_map,
    }
    if max_memory:
        kwargs["max_memory"] = max_memory
    if offload_folder:
        kwargs["offload_folder"] = offload_folder
    if attn_implementation:
        kwargs["attn_implementation"] = attn_implementation

    model = AutoModelForCausalLM.from_pretrained(model_id_or_path, **kwargs)
    if max_memory and getattr(model, "hf_device_map", None) is None:
        try:
            from accelerate import dispatch_model, infer_auto_device_map

            device_map_inferred = infer_auto_device_map(
                model,
                max_memory=max_memory,
                no_split_module_classes=["LagunaDecoderLayer"],
            )
            model = dispatch_model(model, device_map=device_map_inferred, offload_dir=offload_folder)
        except Exception as exc:
            print(f"[heapr] accelerate dispatch fallback failed: {exc}")
    model.eval()
    if hasattr(model, "config"):
        model.config.use_cache = use_cache
        if cache_implementation is not None and hasattr(model.config, "cache_implementation"):
            model.config.cache_implementation = cache_implementation
        if hasattr(model.config, "output_router_logits"):
            model.config.output_router_logits = output_router_logits
    if cache_implementation is not None and hasattr(model, "generation_config"):
        model.generation_config.cache_implementation = cache_implementation
    return model


def build_max_memory(
    *,
    gpu_memory_per_device: str | None = None,
    max_gpu_memory: str | None = None,
    max_cpu_memory: str | None = None,
    allow_cpu_offload: bool = False,
) -> dict[int | str, str] | None:
    """Build an Accelerate max_memory map without CPU offload by default."""

    if gpu_memory_per_device and max_gpu_memory:
        raise ValueError("pass either --gpu-memory-per-device or --max-gpu-memory, not both")
    if max_cpu_memory and not allow_cpu_offload:
        raise ValueError("--max-cpu-memory requires --allow-cpu-offload")
    if not gpu_memory_per_device and not max_gpu_memory:
        if allow_cpu_offload and max_cpu_memory:
            return {"cpu": max_cpu_memory}
        return None

    torch = require_torch()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required when GPU max-memory limits are requested")

    if gpu_memory_per_device:
        max_memory: dict[int | str, str] = {
            device_idx: gpu_memory_per_device for device_idx in range(torch.cuda.device_count())
        }
    else:
        max_memory = {0: str(max_gpu_memory)}

    if allow_cpu_offload:
        if not max_cpu_memory:
            raise ValueError("--allow-cpu-offload requires --max-cpu-memory")
        max_memory["cpu"] = max_cpu_memory
    return max_memory


def _device_map_values(device_map: Any) -> list[str]:
    if not isinstance(device_map, dict):
        return []
    values: list[str] = []
    for value in device_map.values():
        if isinstance(value, dict):
            values.extend(_device_map_values(value))
        else:
            values.append(str(value))
    return values


def validate_model_device_placement(
    model: Any,
    *,
    allow_cpu_offload: bool = False,
    requested_gpu_count: int | None = None,
) -> None:
    """Fail fast when a loaded model spills outside the requested GPU placement."""

    device_map = getattr(model, "hf_device_map", None)
    map_values = _device_map_values(device_map)
    offloaded = sorted({value for value in map_values if value in {"cpu", "disk"}})
    if offloaded and not allow_cpu_offload:
        raise RuntimeError(
            "model was offloaded to CPU/disk while CPU offload is disabled: "
            f"{', '.join(offloaded)}"
        )

    parameter_devices = {str(parameter.device) for parameter in model.parameters()}
    offloaded_params = sorted(device for device in parameter_devices if device in {"cpu", "meta"})
    if offloaded_params and not allow_cpu_offload:
        raise RuntimeError(
            "model parameters are not fully on CUDA while CPU offload is disabled: "
            f"{', '.join(offloaded_params)}"
        )

    if requested_gpu_count is not None and requested_gpu_count > 1:
        cuda_param_devices = {device for device in parameter_devices if device.startswith("cuda")}
        if len(cuda_param_devices) <= 1:
            raise RuntimeError(
                f"multi-GPU placement was requested for {requested_gpu_count} GPUs, "
                f"but parameters landed on {sorted(parameter_devices)}"
            )


def default_model_id(stage: str) -> str:
    if stage in {"smoke", "quantized"}:
        return DEFAULT_SMOKE_MODEL
    if stage in {"prune", "score", "baseline"}:
        return DEFAULT_PRUNE_MODEL
    raise ValueError(f"unknown model stage: {stage}")


def resolve_model_revision(model_id_or_path: str, revision: str | None = None) -> str | None:
    """Resolve a Hugging Face model revision SHA when possible."""

    path = Path(model_id_or_path)
    if path.exists():
        return None
    try:
        from huggingface_hub import model_info

        return model_info(model_id_or_path, revision=revision).sha
    except Exception:
        return revision


def _module_path(root: Any, target: Any) -> str:
    for name, module in root.named_modules():
        if module is target:
            return name
    return "<unknown>"


def get_model_layers(model: Any) -> list[Any]:
    """Return decoder layers for common HF causal LM wrappers."""

    candidates = [
        ("model", "layers"),
        ("model", "model", "layers"),
        ("layers",),
    ]
    for path in candidates:
        obj = model
        for attr in path:
            obj = getattr(obj, attr, None)
            if obj is None:
                break
        if obj is not None:
            return list(obj)
    raise ValueError("could not find model decoder layers")


def get_expert_tensors(mlp: Any):
    """Return packed Laguna expert tensors from a sparse MLP module."""

    experts = getattr(mlp, "experts", None)
    if experts is None:
        return None
    gate_up = getattr(experts, "gate_up_proj", None)
    down = getattr(experts, "down_proj", None)
    if gate_up is None or down is None:
        return None
    return gate_up, down


def discover_sparse_layers(model: Any) -> list[SparseLayerInfo]:
    """Discover Laguna sparse MoE layers using packed expert tensor shapes."""

    layers = get_model_layers(model)
    sparse_layers: list[SparseLayerInfo] = []
    for model_layer_idx, layer in enumerate(layers):
        mlp = getattr(layer, "mlp", None)
        tensors = get_expert_tensors(mlp) if mlp is not None else None
        if tensors is None:
            continue
        gate_up, down = tensors
        if len(gate_up.shape) != 3 or len(down.shape) != 3:
            continue
        num_experts = int(gate_up.shape[0])
        routed_width = int(down.shape[2])
        hidden_size = int(down.shape[1])
        expected_gate_up_rows = routed_width * 2
        if int(gate_up.shape[1]) != expected_gate_up_rows:
            raise ValueError(
                f"layer {model_layer_idx} has gate_up rows {gate_up.shape[1]}, "
                f"expected {expected_gate_up_rows} for routed width {routed_width}"
            )
        if int(gate_up.shape[2]) != hidden_size:
            raise ValueError(
                f"layer {model_layer_idx} hidden mismatch: gate_up {gate_up.shape}, down {down.shape}"
            )
        sparse_layers.append(
            SparseLayerInfo(
                model_layer_idx=model_layer_idx,
                sparse_idx=len(sparse_layers),
                num_experts=num_experts,
                routed_width=routed_width,
                hidden_size=hidden_size,
                module_path=_module_path(model, mlp),
            )
        )
    if not sparse_layers:
        raise ValueError("no Laguna-style sparse MoE layers were discovered")
    return sparse_layers


def create_run_metadata(
    *,
    run_id: str,
    model_id_or_path: str,
    output_dir: str | Path,
    args: dict[str, Any],
    revision: str | None = None,
) -> RunMetadata:
    output_dir = ensure_dir(output_dir)
    resolved_revision = resolve_model_revision(model_id_or_path, revision)
    metadata = RunMetadata(
        run_id=run_id,
        model_id_or_path=model_id_or_path,
        revision=resolved_revision,
        output_dir=Path(output_dir),
        args=args,
        hardware=collect_hardware_metadata(),
        dependency_versions=collect_dependency_versions(),
    )
    payload = asdict(metadata)
    payload["output_dir"] = str(metadata.output_dir)
    write_json(Path(output_dir) / "summary.json", payload)
    return metadata


def model_device_summary(model: Any) -> dict[str, Any]:
    """Summarize where model parameters landed after loading."""

    summary: dict[str, Any] = {"hf_device_map": getattr(model, "hf_device_map", None)}
    counts: dict[str, int] = {}
    bytes_by_device: dict[str, int] = {}
    for parameter in model.parameters():
        device = str(parameter.device)
        counts[device] = counts.get(device, 0) + int(parameter.numel())
        bytes_by_device[device] = (
            bytes_by_device.get(device, 0) + int(parameter.numel()) * int(parameter.element_size())
        )
    summary["parameter_counts_by_device"] = counts
    summary["parameter_gib_by_device"] = {
        device: value / (1024**3) for device, value in bytes_by_device.items()
    }
    try:
        torch = require_torch()
    except RuntimeError:
        return summary
    if torch.cuda.is_available():
        summary["cuda_memory_allocated_gib_by_device"] = {
            str(device_idx): torch.cuda.memory_allocated(device_idx) / (1024**3)
            for device_idx in range(torch.cuda.device_count())
        }
        summary["cuda_memory_reserved_gib_by_device"] = {
            str(device_idx): torch.cuda.memory_reserved(device_idx) / (1024**3)
            for device_idx in range(torch.cuda.device_count())
        }
        summary["cuda_peak_memory_allocated_gib_by_device"] = {
            str(device_idx): torch.cuda.max_memory_allocated(device_idx) / (1024**3)
            for device_idx in range(torch.cuda.device_count())
        }
        summary["cuda_peak_memory_reserved_gib_by_device"] = {
            str(device_idx): torch.cuda.max_memory_reserved(device_idx) / (1024**3)
            for device_idx in range(torch.cuda.device_count())
        }
    return summary