scripts/laya_memory_probe.py

#!/usr/bin/env python3
"""Measure full-parameter Laya supervised-CE memory on a single GPU.

This is an engineering probe only.  It uses synthetic token ids (every position is
attended, length 1024), so it does not produce a biological result.  The probe loads
the fixed local typed-decisions checkpoint, keeps parameters in FP32, uses BF16 CUDA
autocast, enables gradient checkpointing on the ModernBERT encoder, and runs AdamW
updates for micro-batches 4 and 8.  The decision head's ``act_head`` is frozen because
this CE probe has no action target.  Laya's ``head_checkpointing`` attribute is only a
configuration flag; it is deliberately not counted as checkpointing here.
"""

from __future__ import annotations

import argparse
import gc
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from safetensors.torch import load_file


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", default="artifacts/laya_model")
    p.add_argument("--laya-repo", default="vendor/laya")
    p.add_argument("--output", default="artifacts/laya_memory_probe.json")
    p.add_argument("--seq-len", type=int, default=1024)
    p.add_argument("--updates", type=int, default=3)
    p.add_argument("--lr", type=float, default=2e-5)
    p.add_argument("--seed", type=int, default=20260922)
    p.add_argument("--micro-batches", type=int, nargs="+", default=[4, 8])
    return p.parse_args()


def driver_info() -> dict[str, Any]:
    """Read driver details without making nvidia-smi a hard dependency."""
    info: dict[str, Any] = {}
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=index,name,memory.total,driver_version",
                "--format=csv,noheader,nounits",
            ],
            text=True,
            stderr=subprocess.STDOUT,
        ).strip()
        rows = []
        for line in out.splitlines():
            fields = [x.strip() for x in line.split(",")]
            if len(fields) >= 4:
                rows.append(
                    {
                        "index": fields[0],
                        "name": fields[1],
                        "memory_total_mib": fields[2],
                        "driver_version": fields[3],
                    }
                )
        info["nvidia_smi"] = rows
    except (OSError, subprocess.CalledProcessError) as exc:
        info["nvidia_smi_error"] = str(exc)
    return info


def import_laya(repo: str):
    repo_path = str(Path(repo).resolve())
    if repo_path not in sys.path:
        sys.path.insert(0, repo_path)
    from laya.common import build_model

    return build_model


def load_fixed_model(model_dir: Path, build_model, device: torch.device):
    """Build and load on CPU first, then move one copy to the GPU."""
    with (model_dir / "rl_agent_config.json").open(encoding="utf-8") as fh:
        cfg = json.load(fh)
    encoder_dir = model_dir / "encoder"
    model = build_model(cfg, encoder_dir=str(encoder_dir))
    # safetensors is explicitly read on CPU.  Loading into the CPU model first
    # avoids a transient second full model allocation on the GPU.
    state = load_file(str(model_dir / "model.safetensors"), device="cpu")
    model.load_state_dict(state, strict=True)
    del state
    gc.collect()

    # This is the real activation checkpointing switch for ModernBERT.  The
    # DecisionModel.head_checkpointing attribute is not used by its forward pass.
    model.encoder.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False}
    )
    model.to(device)
    model.train()
    # There is no action supervision in this CE objective.  Keep act_head out of
    # autograd and report that it was frozen rather than pretending it was trained.
    for parameter in model.act_head.parameters():
        parameter.requires_grad_(False)
    model.act_head.eval()
    return model, cfg


def memory_snapshot(device: torch.device) -> dict[str, float]:
    return {
        "allocated_gib": float(torch.cuda.memory_allocated(device) / 2**30),
        "reserved_gib": float(torch.cuda.memory_reserved(device) / 2**30),
        "max_allocated_gib": float(torch.cuda.max_memory_allocated(device) / 2**30),
        "max_reserved_gib": float(torch.cuda.max_memory_reserved(device) / 2**30),
    }


def assert_finite_gradients(model: torch.nn.Module) -> tuple[float, int]:
    finite_count = 0
    squared = None
    for parameter in model.parameters():
        if not parameter.requires_grad or parameter.grad is None:
            continue
        assert torch.isfinite(parameter.grad).all().item(), "non-finite gradient"
        finite_count += 1
        value = parameter.grad.detach().float().norm(2)
        squared = value * value if squared is None else squared + value * value
    assert finite_count > 0, "no trainable parameter received a gradient"
    grad_norm = float(torch.sqrt(squared).detach().cpu())
    assert torch.isfinite(torch.tensor(grad_norm)), "non-finite aggregate gradient norm"
    return grad_norm, finite_count


def run_trial(
    model: torch.nn.Module,
    device: torch.device,
    micro_batch: int,
    seq_len: int,
    updates: int,
    lr: float,
    seed: int,
) -> dict[str, Any]:
    # Parameters remain FP32; BF16 is used only by CUDA autocast.
    parameter_dtypes = sorted({str(p.dtype) for p in model.parameters()})
    assert parameter_dtypes == ["torch.float32"], parameter_dtypes
    assert torch.cuda.is_bf16_supported(), "the requested BF16 CUDA autocast is unavailable"

    model.train()
    model.act_head.eval()
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01
    )
    optimizer.zero_grad(set_to_none=True)

    vocab_size = int(model.encoder.config.vocab_size)
    rng = torch.Generator(device="cpu").manual_seed(seed)
    input_ids_cpu = torch.randint(
        0, vocab_size, (micro_batch, seq_len), generator=rng, dtype=torch.long
    )
    attention_cpu = torch.ones((micro_batch, seq_len), dtype=torch.long)
    # Two valid choice markers.  Their token values are irrelevant to the memory probe.
    marker_pos_cpu = torch.tensor([[0, 1]] * micro_batch, dtype=torch.long)
    marker_mask_cpu = torch.ones((micro_batch, 2), dtype=torch.bool)
    qtype_cpu = torch.zeros(micro_batch, dtype=torch.long)
    labels_cpu = torch.arange(micro_batch, dtype=torch.long) % 2

    torch.cuda.reset_peak_memory_stats(device)
    update_records: list[dict[str, Any]] = []
    sync_times: list[float] = []
    t_total_start = time.perf_counter()

    for update in range(updates):
        # Include host-to-device copies outside the measured synchronized model step.
        batch = {
            "input_ids": input_ids_cpu.to(device, non_blocking=True),
            "attention_mask": attention_cpu.to(device, non_blocking=True),
            "marker_pos": marker_pos_cpu.to(device, non_blocking=True),
            "marker_mask": marker_mask_cpu.to(device, non_blocking=True),
            "qtype": qtype_cpu.to(device, non_blocking=True),
            "labels": labels_cpu.to(device, non_blocking=True),
        }
        pre_t = time.perf_counter()
        torch.cuda.synchronize(device)
        sync_times.append(time.perf_counter() - pre_t)
        compute_t = time.perf_counter()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits, _ = model(
                batch["input_ids"],
                batch["attention_mask"],
                batch["marker_pos"],
                batch["marker_mask"],
                batch["qtype"],
            )
            logits = logits.masked_fill(~batch["marker_mask"], -1e4)
            loss = F.cross_entropy(logits, batch["labels"])
        assert torch.isfinite(loss).item(), "non-finite loss"
        loss.backward()
        grad_norm, finite_grad_tensors = assert_finite_gradients(model)
        clipped_norm = torch.nn.utils.clip_grad_norm_(
            [p for p in model.parameters() if p.requires_grad], 1.0
        )
        assert torch.isfinite(clipped_norm).item(), "non-finite clipped gradient norm"
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        post_t = time.perf_counter()
        torch.cuda.synchronize(device)
        sync_times.append(time.perf_counter() - post_t)
        wall_seconds = time.perf_counter() - compute_t

        state_tensors = sum(
            1 for state in optimizer.state.values() for key in ("exp_avg", "exp_avg_sq") if key in state
        )
        snap = memory_snapshot(device)
        update_records.append(
            {
                "optimizer_update": update + 1,
                "loss": float(loss.detach().cpu()),
                "grad_norm_before_clip": grad_norm,
                "clipped_grad_norm": float(clipped_norm.detach().cpu()),
                "finite_gradient_tensors": finite_grad_tensors,
                "adam_state_tensors": state_tensors,
                "wall_seconds_with_synchronize": wall_seconds,
                "memory": snap,
            }
        )
        print(
            f"[memory-probe] microbatch={micro_batch} update={update + 1}/{updates} "
            f"loss={float(loss.detach().cpu()):.5f} grad={grad_norm:.3f} "
            f"max_alloc={snap['max_allocated_gib']:.2f}GiB "
            f"max_reserved={snap['max_reserved_gib']:.2f}GiB",
            flush=True,
        )
        del batch, logits, loss

    total_seconds = time.perf_counter() - t_total_start
    final = memory_snapshot(device)
    return {
        "micro_batch": micro_batch,
        "seq_len": seq_len,
        "attention_all_ones": True,
        "updates": updates,
        "parameter_dtypes": parameter_dtypes,
        "autocast_dtype": "torch.bfloat16",
        "encoder_gradient_checkpointing": bool(
            getattr(model.encoder, "is_gradient_checkpointing", False)
        ),
        "head_checkpointing_attribute": bool(getattr(model, "head_checkpointing", False)),
        "head_checkpointing_used": False,
        "act_head_frozen": all(not p.requires_grad for p in model.act_head.parameters()),
        "total_seconds": total_seconds,
        "synchronize_seconds": {
            "count": len(sync_times),
            "sum": float(sum(sync_times)),
            "max": float(max(sync_times) if sync_times else 0.0),
        },
        "updates_detail": update_records,
        "final_memory": final,
        "peak_memory": {
            "max_allocated_gib": max(x["memory"]["max_allocated_gib"] for x in update_records),
            "max_reserved_gib": max(x["memory"]["max_reserved_gib"] for x in update_records),
        },
        "assertions": {
            "finite_loss_each_update": True,
            "finite_gradients_each_update": True,
            "adam_states_present_after_update_1": update_records[0]["adam_state_tensors"] > 0,
        },
    }


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable; the GPU should be enabled before this probe")
    assert args.seq_len == 1024, "the requested probe is defined for 1024-token inputs"
    device = torch.device("cuda:0")
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    model_dir = Path(args.model_dir)
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    build_model = import_laya(args.laya_repo)
    props = torch.cuda.get_device_properties(device)
    result: dict[str, Any] = {
        "probe": "laya_full_supervised_ce_memory",
        "synthetic_inputs_only": True,
        "model_dir": str(model_dir),
        "laya_repo": str(Path(args.laya_repo).resolve()),
        "gpu": {
            "index": 0,
            "name": props.name,
            "total_memory_gib": float(props.total_memory / 2**30),
            "cuda_runtime": torch.version.cuda,
            "torch": torch.__version__,
            **driver_info(),
        },
        "config": {
            "seq_len": args.seq_len,
            "updates": args.updates,
            "micro_batches": args.micro_batches,
            "lr": args.lr,
            "seed": args.seed,
            "optimizer": "AdamW",
            "parameters": "FP32",
            "autocast": "BF16",
            "encoder_gradient_checkpointing": True,
            "no_cpt": True,
        },
        "trials": [],
    }

    for micro_batch in args.micro_batches:
        # Each trial gets a fresh model and optimizer, ensuring the second trial's
        # peak is not contaminated by the first trial's Adam state.
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats(device)
        model = None
        try:
            model, cfg = load_fixed_model(model_dir, build_model, device)
            result["checkpoint"] = {
                "encoder": cfg.get("encoder"),
                "max_len": cfg.get("max_len"),
                "head_max_len": cfg.get("head_max_len"),
                "vocab_size": int(model.encoder.config.vocab_size),
            }
            loaded = memory_snapshot(device)
            trial = run_trial(
                model,
                device,
                micro_batch,
                args.seq_len,
                args.updates,
                args.lr,
                args.seed + micro_batch,
            )
            trial["model_loaded_memory"] = loaded
            result["trials"].append(trial)
        except torch.cuda.OutOfMemoryError as exc:
            torch.cuda.empty_cache()
            result["trials"].append(
                {
                    "micro_batch": micro_batch,
                    "status": "oom",
                    "error": str(exc),
                    "memory_at_oom": memory_snapshot(device),
                }
            )
            print(f"[memory-probe] microbatch={micro_batch} OOM", flush=True)
        finally:
            del model
            gc.collect()
            torch.cuda.empty_cache()
            torch.cuda.synchronize(device)

    result["status"] = "ok" if all(t.get("status", "ok") == "ok" for t in result["trials"]) else "partial"
    output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print("[memory-probe] summary", flush=True)
    for trial in result["trials"]:
        if trial.get("status") == "oom":
            print(f"  microbatch={trial['micro_batch']}: OOM", flush=True)
        else:
            peak = trial["peak_memory"]
            print(
                f"  microbatch={trial['micro_batch']}: "
                f"max_alloc={peak['max_allocated_gib']:.2f} GiB, "
                f"max_reserved={peak['max_reserved_gib']:.2f} GiB, "
                f"updates={trial['updates']}",
                flush=True,
            )


if __name__ == "__main__":
    main()