code/scripts/bench_7b_end_to_end.py · viai957/align at main

align / code /scripts /bench_7b_end_to_end.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

10.8 kB

"""Benchmark #2: end-to-end step time on the REAL 7B, triton_fa2 vs SDPA.
Run (GPU 3 window; pin BEFORE torch import):
    CUDA_VISIBLE_DEVICES=3 TMPDIR=/apps/development/rlhf/nano-mistral-o1/.tmp \
        uv run python scripts/bench_7b_end_to_end.py
WHAT THIS MEASURES, AND WHAT IT DELIBERATELY DOES NOT
-----------------------------------------------------
A full-parameter 7B training step does NOT fit on one A10G: bf16 weights are
13.5 GiB and bf16 parameter gradients are another 13.5 GiB, so 27 GiB against a
23 GiB card. Real SFT gets around this with 4-way FSDP plus an offloaded
optimizer (section-07/11); that is not reproducible inside a single-GPU window,
and pretending otherwise would be measuring a different thing than the label.
So this benchmark freezes every parameter and requires grad on `inputs_embeds`
instead. Backward still flows through every attention and MLP in all 32 layers,
computing the full activation-gradient chain -- dQ/dK/dV included, which is
exactly the part the attention backend controls. What is skipped is the
per-parameter weight-gradient GEMM, which is IDENTICAL for both backends.
The consequence is worth stating precisely, because it decides how to read the
result: the absolute per-step delta (triton minus sdpa, in ms) is exactly what
it would be in a full step, since the omitted work is backend-independent and
additive. The delta as a PERCENTAGE is inflated, because the denominator is
missing that same backend-independent work. **This benchmark therefore reports
an UPPER BOUND on the true percentage cost of choosing triton_fa2.** A number
that bounds the answer from the safe side is worth more here than a number that
requires four GPUs and a training run to obtain.
Both no-checkpointing and activation-checkpointing configurations are measured:
under checkpointing the attention forward runs a second time during backward,
so the backend-attributable share is larger, and SFT is expected to checkpoint.
"""

from __future__ import annotations

import glob
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from align.utils.guard import apply_resource_guards

# Light job by the box rules: no heavy lock (other GPU jobs are live), no
# dataloaders, no multi-GB writes.
apply_resource_guards(cpu_threads=8, heavy=False, label="bench-7b-attn")

import torch
from torch.utils.checkpoint import checkpoint as ckpt_fn

from align.model import checkpoint as ckpt
from align.model.mistral import NmoMistralForCausalLM

REPO = "/apps/development/rlhf/nano-mistral-o1"
OUT = f"{REPO}/docs/bench_7b_end_to_end.json"
SNAPSHOT = (
    "/apps/voxtral/.cache/huggingface/hub/"
    "models--mistralai--Mistral-7B-v0.3/snapshots/"
    "caa1feb0e54d415e2df31207e5f4e273e33509b1"
)

SEQ = 1024
BATCH = 1
#: Hard wall-clock budget. The GPU window is shared and SFT launches the moment
#: the golden gate passes; a benchmark that overstays blocks a training run.
DEADLINE_S = float(os.environ.get("NMO_BENCH_DEADLINE_S", 1500))
_START = time.time()


def _left() -> float:
    return DEADLINE_S - (time.time() - _START)


def _load_7b(attn_impl: str, state_dict, cfg):
    model = NmoMistralForCausalLM.from_state_dict(cfg, state_dict, attn_implementation=attn_impl)
    model = model.to(torch.bfloat16).to("cuda").eval()
    # Freeze: no 13.5 GiB parameter-gradient buffer. See module docstring.
    for p in model.parameters():
        p.requires_grad_(False)
    return model


def _blocks(model):
    inner = model.model if hasattr(model, "model") else model
    for name in ("layers", "blocks", "h"):
        if hasattr(inner, name):
            return getattr(inner, name)
    raise AttributeError("cannot find the decoder block list on the model")


def _timed(fn, iters, warmup):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters * 1e3


def _run_config(model, embeds, labels, use_ckpt: bool, iters: int):
    """fwd-only and fwd+bwd timings for one model/checkpointing combination."""
    blocks = _blocks(model)
    originals = [b.forward for b in blocks]
    if use_ckpt:
        for b in blocks:
            fwd = b.forward
            b.forward = lambda *a, _f=fwd, **kw: ckpt_fn(_f, *a, use_reentrant=False, **kw)

    try:

        def fwd_only():
            with torch.no_grad():
                model(inputs_embeds=embeds.detach())

        def fwd_bwd():
            if embeds.grad is not None:
                embeds.grad = None
            logits = model(inputs_embeds=embeds)
            loss = torch.nn.functional.cross_entropy(
                logits.float().view(-1, logits.shape[-1]), labels.view(-1)
            )
            loss.backward()

        torch.cuda.reset_peak_memory_stats()
        f = _timed(fwd_only, iters, 2)
        fb = _timed(fwd_bwd, iters, 2)
        peak = torch.cuda.max_memory_allocated() / 2**30
    finally:
        # `b.forward = original_bound_method` would leave the instance holding a
        # bound method that holds the instance -- a reference cycle the caching
        # allocator cannot reclaim on `empty_cache()`, which OOM'd the second
        # model load. Deleting the instance attribute restores the class method
        # and leaves no cycle.
        for b in blocks:
            if "forward" in b.__dict__:
                del b.__dict__["forward"]
        del originals
    return {"fwd_ms": f, "fwd_bwd_ms": fb, "peak_gib": peak}


def main() -> int:
    if torch.cuda.device_count() != 1:
        print(
            f"{torch.cuda.device_count()} devices visible; pin exactly one "
            "(shared box, three other GPU jobs live)",
            file=sys.stderr,
        )
        return 1
    free_gib = torch.cuda.mem_get_info()[0] / 2**30
    print(f"device: {torch.cuda.get_device_name(0)}  free {free_gib:.1f} GiB")
    if free_gib < 20.0:
        print(
            "refusing: need ~20 GiB free for 13.5 GiB of bf16 weights plus "
            "activations; another job is resident on this device.",
            file=sys.stderr,
        )
        return 1

    from align.eval.golden import resolve_model_config

    cfg = resolve_model_config(SNAPSHOT)
    print(
        f"config: {cfg.n_layers}L hidden={cfg.hidden_size} "
        f"{cfg.n_heads}q:{cfg.n_kv_heads}kv head_dim={cfg.head_dim} "
        f"vocab={cfg.vocab_size} sliding_window={cfg.sliding_window}"
    )

    import safetensors.torch as st

    t0 = time.time()
    hf_sd = {}
    for shard in sorted(glob.glob(os.path.join(SNAPSHOT, "*.safetensors"))):
        hf_sd.update(st.load_file(shard))
    nmo_sd = ckpt.hf_to_nmo(hf_sd, expected_n_layers=cfg.n_layers)
    del hf_sd
    print(f"loaded + remapped {len(nmo_sd)} tensors in {time.time() - t0:.1f}s")

    torch.manual_seed(0)
    labels = torch.randint(0, cfg.vocab_size, (BATCH, SEQ), device="cuda")
    results = {}

    # One backend per PROCESS when an arg is given. 13.5 GiB of weights leaves
    # no room for a second model, and process exit is the only VRAM release
    # that cannot be defeated by a stray reference. Partials are merged below.
    only = sys.argv[1] if len(sys.argv) > 1 else None
    impls = (only,) if only else ("sdpa", "triton_fa2")

    for impl in impls:
        if _left() < 240:
            print(f"\ndeadline guard: {_left():.0f}s left, skipping {impl}", file=sys.stderr)
            break
        print(f"\n=== {impl} ===", flush=True)
        model = _load_7b(impl, nmo_sd, cfg)
        embeds = torch.randn(
            BATCH, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True
        )
        results[impl] = {}
        for use_ckpt in (False, True):
            if _left() < 120:
                print(
                    f"deadline guard: {_left():.0f}s left, skipping {impl}/ckpt={use_ckpt}",
                    file=sys.stderr,
                )
                break
            r = _run_config(model, embeds, labels, use_ckpt, iters=5)
            key = "activation_ckpt" if use_ckpt else "no_ckpt"
            results[impl][key] = r
            print(
                f"  {key:16s} fwd {r['fwd_ms']:8.1f} ms   "
                f"fwd+bwd {r['fwd_bwd_ms']:8.1f} ms   peak {r['peak_gib']:5.2f} GiB",
                flush=True,
            )
        del model, embeds
        import gc

        gc.collect()
        torch.cuda.empty_cache()

    # Merge any partial written by a sibling process, so running one backend
    # per process still produces a single complete receipt.
    partial = f"{REPO}/.tmp/bench7b_partial.json"
    if only:
        prev = {}
        if os.path.isfile(partial):
            with open(partial) as f:
                prev = json.load(f)
        prev.update(results)
        with open(partial, "w") as f:
            json.dump(prev, f, indent=2)
        results = prev
        print(f"\npartial results for {sorted(results)} -> {partial}")

    payload = {
        "note": "parameters FROZEN; grad flows via inputs_embeds. Absolute ms "
        "deltas are exact; percentages are an UPPER BOUND on the true "
        "cost, because the omitted weight-grad GEMMs are "
        "backend-independent and would enlarge the denominator only. "
        "See the module docstring.",
        "device": torch.cuda.get_device_name(0),
        "torch": torch.__version__,
        "model": "mistralai/Mistral-7B-v0.3 (real weights)",
        "batch": BATCH,
        "seq": SEQ,
        "dtype": "bfloat16",
        "results": results,
    }

    if "sdpa" in results and "triton_fa2" in results:
        payload["comparison"] = {}
        print("\n=== triton_fa2 vs sdpa on the real 7B ===")
        for key in ("no_ckpt", "activation_ckpt"):
            if key in results["sdpa"] and key in results["triton_fa2"]:
                s = results["sdpa"][key]["fwd_bwd_ms"]
                t = results["triton_fa2"][key]["fwd_bwd_ms"]
                c = {
                    "sdpa_ms": s,
                    "triton_ms": t,
                    "delta_ms": t - s,
                    "ratio": t / s,
                    "upper_bound_pct": 100.0 * (t - s) / s,
                }
                payload["comparison"][key] = c
                print(
                    f"  {key:16s} sdpa {s:8.1f} ms -> triton {t:8.1f} ms "
                    f"({c['delta_ms']:+.1f} ms, {c['upper_bound_pct']:+.1f}% "
                    "UPPER BOUND)"
                )

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "w") as f:
        json.dump(payload, f, indent=2)
        f.write("\n")
    print(f"\nwrote {OUT}  ({_left():.0f}s of window left)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

	"""Benchmark #2: end-to-end step time on the REAL 7B, triton_fa2 vs SDPA.

	Run (GPU 3 window; pin BEFORE torch import):
	CUDA_VISIBLE_DEVICES=3 TMPDIR=/apps/development/rlhf/nano-mistral-o1/.tmp \
	uv run python scripts/bench_7b_end_to_end.py

	WHAT THIS MEASURES, AND WHAT IT DELIBERATELY DOES NOT
	-----------------------------------------------------
	A full-parameter 7B training step does NOT fit on one A10G: bf16 weights are
	13.5 GiB and bf16 parameter gradients are another 13.5 GiB, so 27 GiB against a
	23 GiB card. Real SFT gets around this with 4-way FSDP plus an offloaded
	optimizer (section-07/11); that is not reproducible inside a single-GPU window,
	and pretending otherwise would be measuring a different thing than the label.

	So this benchmark freezes every parameter and requires grad on `inputs_embeds`
	instead. Backward still flows through every attention and MLP in all 32 layers,
	computing the full activation-gradient chain -- dQ/dK/dV included, which is
	exactly the part the attention backend controls. What is skipped is the
	per-parameter weight-gradient GEMM, which is IDENTICAL for both backends.

	The consequence is worth stating precisely, because it decides how to read the
	result: the absolute per-step delta (triton minus sdpa, in ms) is exactly what
	it would be in a full step, since the omitted work is backend-independent and
	additive. The delta as a PERCENTAGE is inflated, because the denominator is
	missing that same backend-independent work. **This benchmark therefore reports
	an UPPER BOUND on the true percentage cost of choosing triton_fa2.** A number
	that bounds the answer from the safe side is worth more here than a number that
	requires four GPUs and a training run to obtain.

	Both no-checkpointing and activation-checkpointing configurations are measured:
	under checkpointing the attention forward runs a second time during backward,
	so the backend-attributable share is larger, and SFT is expected to checkpoint.
	"""

	from __future__ import annotations

	import glob
	import json
	import os
	import sys
	import time

	sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

	from align.utils.guard import apply_resource_guards

	# Light job by the box rules: no heavy lock (other GPU jobs are live), no
	# dataloaders, no multi-GB writes.
	apply_resource_guards(cpu_threads=8, heavy=False, label="bench-7b-attn")

	import torch
	from torch.utils.checkpoint import checkpoint as ckpt_fn

	from align.model import checkpoint as ckpt
	from align.model.mistral import NmoMistralForCausalLM

	REPO = "/apps/development/rlhf/nano-mistral-o1"
	OUT = f"{REPO}/docs/bench_7b_end_to_end.json"
	SNAPSHOT = (
	"/apps/voxtral/.cache/huggingface/hub/"
	"models--mistralai--Mistral-7B-v0.3/snapshots/"
	"caa1feb0e54d415e2df31207e5f4e273e33509b1"
	)

	SEQ = 1024
	BATCH = 1
	#: Hard wall-clock budget. The GPU window is shared and SFT launches the moment
	#: the golden gate passes; a benchmark that overstays blocks a training run.
	DEADLINE_S = float(os.environ.get("NMO_BENCH_DEADLINE_S", 1500))
	_START = time.time()


	def _left() -> float:
	return DEADLINE_S - (time.time() - _START)


	def _load_7b(attn_impl: str, state_dict, cfg):
	model = NmoMistralForCausalLM.from_state_dict(cfg, state_dict, attn_implementation=attn_impl)
	model = model.to(torch.bfloat16).to("cuda").eval()
	# Freeze: no 13.5 GiB parameter-gradient buffer. See module docstring.
	for p in model.parameters():
	p.requires_grad_(False)
	return model


	def _blocks(model):
	inner = model.model if hasattr(model, "model") else model
	for name in ("layers", "blocks", "h"):
	if hasattr(inner, name):
	return getattr(inner, name)
	raise AttributeError("cannot find the decoder block list on the model")


	def _timed(fn, iters, warmup):
	for _ in range(warmup):
	fn()
	torch.cuda.synchronize()
	t0 = time.perf_counter()
	for _ in range(iters):
	fn()
	torch.cuda.synchronize()
	return (time.perf_counter() - t0) / iters * 1e3


	def _run_config(model, embeds, labels, use_ckpt: bool, iters: int):
	"""fwd-only and fwd+bwd timings for one model/checkpointing combination."""
	blocks = _blocks(model)
	originals = [b.forward for b in blocks]
	if use_ckpt:
	for b in blocks:
	fwd = b.forward
	b.forward = lambda a, _f=fwd, kw: ckpt_fn(_f, a, use_reentrant=False, **kw)

	try:

	def fwd_only():
	with torch.no_grad():
	model(inputs_embeds=embeds.detach())

	def fwd_bwd():
	if embeds.grad is not None:
	embeds.grad = None
	logits = model(inputs_embeds=embeds)
	loss = torch.nn.functional.cross_entropy(
	logits.float().view(-1, logits.shape[-1]), labels.view(-1)
	)
	loss.backward()

	torch.cuda.reset_peak_memory_stats()
	f = _timed(fwd_only, iters, 2)
	fb = _timed(fwd_bwd, iters, 2)
	peak = torch.cuda.max_memory_allocated() / 2**30
	finally:
	# `b.forward = original_bound_method` would leave the instance holding a
	# bound method that holds the instance -- a reference cycle the caching
	# allocator cannot reclaim on `empty_cache()`, which OOM'd the second
	# model load. Deleting the instance attribute restores the class method
	# and leaves no cycle.
	for b in blocks:
	if "forward" in b.__dict__:
	del b.__dict__["forward"]
	del originals
	return {"fwd_ms": f, "fwd_bwd_ms": fb, "peak_gib": peak}


	def main() -> int:
	if torch.cuda.device_count() != 1:
	print(
	f"{torch.cuda.device_count()} devices visible; pin exactly one "
	"(shared box, three other GPU jobs live)",
	file=sys.stderr,
	)
	return 1
	free_gib = torch.cuda.mem_get_info()[0] / 2**30
	print(f"device: {torch.cuda.get_device_name(0)} free {free_gib:.1f} GiB")
	if free_gib < 20.0:
	print(
	"refusing: need ~20 GiB free for 13.5 GiB of bf16 weights plus "
	"activations; another job is resident on this device.",
	file=sys.stderr,
	)
	return 1

	from align.eval.golden import resolve_model_config

	cfg = resolve_model_config(SNAPSHOT)
	print(
	f"config: {cfg.n_layers}L hidden={cfg.hidden_size} "
	f"{cfg.n_heads}q:{cfg.n_kv_heads}kv head_dim={cfg.head_dim} "
	f"vocab={cfg.vocab_size} sliding_window={cfg.sliding_window}"
	)

	import safetensors.torch as st

	t0 = time.time()
	hf_sd = {}
	for shard in sorted(glob.glob(os.path.join(SNAPSHOT, "*.safetensors"))):
	hf_sd.update(st.load_file(shard))
	nmo_sd = ckpt.hf_to_nmo(hf_sd, expected_n_layers=cfg.n_layers)
	del hf_sd
	print(f"loaded + remapped {len(nmo_sd)} tensors in {time.time() - t0:.1f}s")

	torch.manual_seed(0)
	labels = torch.randint(0, cfg.vocab_size, (BATCH, SEQ), device="cuda")
	results = {}

	# One backend per PROCESS when an arg is given. 13.5 GiB of weights leaves
	# no room for a second model, and process exit is the only VRAM release
	# that cannot be defeated by a stray reference. Partials are merged below.
	only = sys.argv[1] if len(sys.argv) > 1 else None
	impls = (only,) if only else ("sdpa", "triton_fa2")

	for impl in impls:
	if _left() < 240:
	print(f"\ndeadline guard: {_left():.0f}s left, skipping {impl}", file=sys.stderr)
	break
	print(f"\n=== {impl} ===", flush=True)
	model = _load_7b(impl, nmo_sd, cfg)
	embeds = torch.randn(
	BATCH, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True
	)
	results[impl] = {}
	for use_ckpt in (False, True):
	if _left() < 120:
	print(
	f"deadline guard: {_left():.0f}s left, skipping {impl}/ckpt={use_ckpt}",
	file=sys.stderr,
	)
	break
	r = _run_config(model, embeds, labels, use_ckpt, iters=5)
	key = "activation_ckpt" if use_ckpt else "no_ckpt"
	results[impl][key] = r
	print(
	f" {key:16s} fwd {r['fwd_ms']:8.1f} ms "
	f"fwd+bwd {r['fwd_bwd_ms']:8.1f} ms peak {r['peak_gib']:5.2f} GiB",
	flush=True,
	)
	del model, embeds
	import gc

	gc.collect()
	torch.cuda.empty_cache()

	# Merge any partial written by a sibling process, so running one backend
	# per process still produces a single complete receipt.
	partial = f"{REPO}/.tmp/bench7b_partial.json"
	if only:
	prev = {}
	if os.path.isfile(partial):
	with open(partial) as f:
	prev = json.load(f)
	prev.update(results)
	with open(partial, "w") as f:
	json.dump(prev, f, indent=2)
	results = prev
	print(f"\npartial results for {sorted(results)} -> {partial}")

	payload = {
	"note": "parameters FROZEN; grad flows via inputs_embeds. Absolute ms "
	"deltas are exact; percentages are an UPPER BOUND on the true "
	"cost, because the omitted weight-grad GEMMs are "
	"backend-independent and would enlarge the denominator only. "
	"See the module docstring.",
	"device": torch.cuda.get_device_name(0),
	"torch": torch.__version__,
	"model": "mistralai/Mistral-7B-v0.3 (real weights)",
	"batch": BATCH,
	"seq": SEQ,
	"dtype": "bfloat16",
	"results": results,
	}

	if "sdpa" in results and "triton_fa2" in results:
	payload["comparison"] = {}
	print("\n=== triton_fa2 vs sdpa on the real 7B ===")
	for key in ("no_ckpt", "activation_ckpt"):
	if key in results["sdpa"] and key in results["triton_fa2"]:
	s = results["sdpa"][key]["fwd_bwd_ms"]
	t = results["triton_fa2"][key]["fwd_bwd_ms"]
	c = {
	"sdpa_ms": s,
	"triton_ms": t,
	"delta_ms": t - s,
	"ratio": t / s,
	"upper_bound_pct": 100.0 * (t - s) / s,
	}
	payload["comparison"][key] = c
	print(
	f" {key:16s} sdpa {s:8.1f} ms -> triton {t:8.1f} ms "
	f"({c['delta_ms']:+.1f} ms, {c['upper_bound_pct']:+.1f}% "
	"UPPER BOUND)"
	)

	os.makedirs(os.path.dirname(OUT), exist_ok=True)
	with open(OUT, "w") as f:
	json.dump(payload, f, indent=2)
	f.write("\n")
	print(f"\nwrote {OUT} ({_left():.0f}s of window left)")
	return 0


	if __name__ == "__main__":
	raise SystemExit(main())