bench_decompose.py

#!/usr/bin/env python3
"""Single-clock step decomposition: attention (FlashAttention kernels) vs everything else.

Splits the real DeepSpeed training step onto ONE clock -- measured wall-clock step time --
so buckets sum to the total by construction (unlike measuring attention in isolation and the
step end-to-end: different clocks that don't reconcile, which is the bug this replaces).

Method:
  1. Measure wall-clock step time T (fwd+bwd+optimizer), profiler off.
  2. torch.profiler one window; classify CUDA kernels by name:
       attention = name contains "flash" (FA3 on Hopper / FA4 on Blackwell)
       comm      = name contains "nccl" (ZeRO grad/param comms)
       compute   = everything else on the compute stream (dense GEMMs, norm/rope/swiglu, CE, optimizer)
  3. NCCL comm overlaps backward compute (overlap_comm=True) -> it's hidden behind the wall clock,
     so we split T by each bucket's share of *compute-stream* busy time (comm excluded):
       t_att  = T * att_time / (att_time + compute_time)
       t_rest = T - t_att            # dense compute + any non-overlapped comm tail
  4. bucket TFLOPS = bucket_FLOPs / bucket_time ; total = total_FLOPs / T (consistent blend).

Deep-dive config (right-sized 8k): ZeRO-2, no activation checkpointing, no offload, no Liger.
Run on both GPUs with the SAME command; each prints its own row:

    deepspeed --num_gpus=8 bench_decompose.py
"""

import importlib.metadata as md
import os
import platform
import time

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
os.environ.setdefault("HF_HOME", os.path.expanduser("~/.cache/huggingface"))

import torch
import deepspeed
import torch.distributed as dist
from torch.profiler import profile, ProfilerActivity

# ---- config (deep-dive: right-sized 8k) -------------------------------------------------
MODEL_ID = "meta-llama/Llama-3.1-8B"
BATCH_SIZE = 1
SEQ_LEN = int(os.environ.get("SEQ_LEN", 8192))
WARMUP_STEPS = int(os.environ.get("WARMUP_STEPS", 3))
TIME_STEPS = int(os.environ.get("TIME_STEPS", 10))
PROFILE_STEPS = int(os.environ.get("PROFILE_STEPS", 5))
ZERO_STAGE = int(os.environ.get("ZERO_STAGE", 2))
OFFLOAD_OPTIMIZER = int(os.environ.get("OFFLOAD_OPTIMIZER", 0))
GRAD_CHECKPOINT = int(os.environ.get("GRAD_CHECKPOINT", 0))
USE_LIGER = int(os.environ.get("USE_LIGER", 0))
LR = 1e-4
SEED = 42
# bf16 dense-matmul peak TFLOPS (for MFU), per GPU family
PEAK_TFLOPS = {9: 989.0, 10: 2250.0}  # H200 (sm90) / B200 (sm100)
# -----------------------------------------------------------------------------------------


def resolve_model(model_id: str) -> str:
    local = os.path.join(os.environ.get("DATA_FAST", os.path.expanduser("~/base-models")), model_id)
    return local if os.path.isfile(os.path.join(local, "config.json")) else model_id


def flops_split(config):
    """(attn_flops, dense_flops) per step in units of 1e12, causal (coeff 6) MFU convention."""
    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
    q_size = config.num_attention_heads * head_dim
    k_size = v_size = config.num_key_value_heads * head_dim
    mlp_N = config.hidden_size * config.intermediate_size * 3
    attn_linear_N = config.hidden_size * (q_size + k_size + v_size + q_size)
    dense_N = (mlp_N + attn_linear_N) * config.num_hidden_layers + config.vocab_size * config.hidden_size
    tokens = BATCH_SIZE * SEQ_LEN
    dense_flops = 6 * dense_N * tokens
    attn_flops = 6 * BATCH_SIZE * SEQ_LEN**2 * head_dim * config.num_attention_heads * config.num_hidden_layers
    return attn_flops / 1e12, dense_flops / 1e12


def detect_backend(device=0) -> str:
    major, _ = torch.cuda.get_device_capability(device)
    if major == 9:
        return "flash_attention_3"
    if major == 10:
        return "flash_attention_4"
    raise RuntimeError(f"unsupported compute capability {major}.x")


def build_model(backend: str):
    path = resolve_model(MODEL_ID)
    if USE_LIGER:
        from liger_kernel.transformers import AutoLigerKernelForCausalLM
        model = AutoLigerKernelForCausalLM.from_pretrained(path, dtype=torch.bfloat16, attn_implementation=backend)
    else:
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16, attn_implementation=backend)
    if GRAD_CHECKPOINT:
        model.config.use_cache = False
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.train()
    return model


def ds_config():
    zero = {"stage": ZERO_STAGE, "overlap_comm": True, "contiguous_gradients": True}
    if OFFLOAD_OPTIMIZER:
        zero["offload_optimizer"] = {"device": "cpu", "pin_memory": True}
    return {
        "train_micro_batch_size_per_gpu": BATCH_SIZE,
        "gradient_accumulation_steps": 1,
        "gradient_clipping": 1.0,
        "bf16": {"enabled": True},
        "optimizer": {"type": "AdamW", "params": {"lr": LR, "betas": [0.9, 0.999], "eps": 1e-8}},
        "zero_optimization": zero,
    }


def main():
    deepspeed.init_distributed()
    rank = dist.get_rank()
    device = f"cuda:{int(os.environ.get('LOCAL_RANK', 0))}"
    torch.cuda.set_device(device)
    torch.manual_seed(SEED)

    backend = detect_backend()
    model = build_model(backend)
    engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config())

    attn_tflop, dense_tflop = flops_split(model.config)
    total_tflop = attn_tflop + dense_tflop
    ids = torch.randint(0, model.config.vocab_size - 256, (BATCH_SIZE, SEQ_LEN), device=device)
    batch = {"input_ids": ids, "labels": ids}

    def one_step():
        loss = engine(**batch).loss
        engine.backward(loss)
        engine.step()

    for _ in range(WARMUP_STEPS):
        one_step()
    torch.cuda.synchronize()

    # (1) wall-clock step time, profiler OFF
    torch.cuda.reset_peak_memory_stats(device)
    t0 = time.perf_counter()
    for _ in range(TIME_STEPS):
        one_step()
    torch.cuda.synchronize()
    T_step = (time.perf_counter() - t0) / TIME_STEPS
    peak_mem = torch.tensor(torch.cuda.max_memory_reserved(device) / 2**30, device=device)
    dist.all_reduce(peak_mem, op=dist.ReduceOp.MAX)
    peak_mem = peak_mem.item()

    # (2) kernel breakdown, profiler ON
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        for _ in range(PROFILE_STEPS):
            one_step()
        torch.cuda.synchronize()

    if rank != 0:
        return

    ka = prof.key_averages()
    att_us = sum(e.self_device_time_total for e in ka if "flash" in e.key.lower())
    comm_us = sum(e.self_device_time_total for e in ka if "nccl" in e.key.lower())
    all_us = sum(e.self_device_time_total for e in ka)
    compute_us = all_us - comm_us  # compute-stream busy time (comm overlaps, excluded)

    # (3) split measured wall time by compute-stream share
    f_att = att_us / compute_us
    t_att = f_att * T_step
    t_rest = T_step - t_att

    major = torch.cuda.get_device_capability(device)[0]
    peak = PEAK_TFLOPS[major]

    def row(name, tflop, t):
        tflops = tflop / t
        return f"{name:16s} {tflop:7.1f} TFLOP  {t*1e3:8.2f} ms  {tflops:8.1f} TFLOPS  MFU {tflops/peak*100:5.1f}%"

    def ver(pkg):
        try:
            return md.version(pkg)
        except md.PackageNotFoundError:
            return "n/a"

    print("=" * 100)
    print(f"# gpu={torch.cuda.get_device_name()} backend={backend} "
          f"zero={ZERO_STAGE} gc={GRAD_CHECKPOINT} offload={OFFLOAD_OPTIMIZER} liger={USE_LIGER} "
          f"seq={SEQ_LEN} time_steps={TIME_STEPS} profile_steps={PROFILE_STEPS}")
    fa = "flash-attn-3" if backend == "flash_attention_3" else "flash-attn-4"
    print(f"# python={platform.python_version()} torch={torch.__version__} cuda={torch.version.cuda} "
          f"transformers={ver('transformers')} deepspeed={ver('deepspeed')} "
          f"liger-kernel={ver('liger-kernel')} {fa}={ver(fa)}")
    print(f"# wall step={T_step*1e3:.2f} ms   peak_mem={peak_mem:.1f} GiB   "
          f"attention FLOP share={attn_tflop/total_tflop*100:.1f}%   "
          f"attention wall share={f_att*100:.1f}%   "
          f"comm={comm_us/all_us*100:.1f}% of kernel time (overlapped, hidden in wall)")
    print("-" * 100)
    print(row("attention(FA)", attn_tflop, t_att))
    print(row("everything else", dense_tflop, t_rest))
    print(row("TOTAL", total_tflop, T_step))
    print("-" * 100)
    print("# top CUDA kernels by self time:")
    rows = sorted(ka, key=lambda e: -e.self_device_time_total)[:12]
    for e in rows:
        tag = "ATT" if "flash" in e.key.lower() else ("NCL" if "nccl" in e.key.lower() else "   ")
        print(f"#  {tag} {e.self_device_time_total/all_us*100:5.1f}%  {e.key[:80]}")
    print("=" * 100)


if __name__ == "__main__":
    main()