diff --git a/benchmark/scripts/benchmark_attn_res.py b/benchmark/scripts/benchmark_attn_res.py index b2fbb778..d939f537 100644 --- a/benchmark/scripts/benchmark_attn_res.py +++ b/benchmark/scripts/benchmark_attn_res.py @@ -4,7 +4,6 @@ AttnRes Benchmark: Liger (Triton) vs PyTorch Kimi Attention Residuals: softmax attention over depth blocks. """ -import math import os import sys @@ -13,15 +12,14 @@ import torch sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.ops import LigerAttnResFunction from liger_kernel.utils import infer_device @@ -29,23 +27,31 @@ from liger_kernel.utils import infer_device device = infer_device() -def _setup_attn_res(input: SingleBenchmarkRunInput): +def setup_attn_res(input: SingleBenchmarkRunInput): """Create input tensors for AttnRes from benchmark config.""" cfg = input.extra_benchmark_config - seq_len = input.x + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] # V: [N, B, T, D] V = torch.randn( cfg["N"], cfg["bsz"], seq_len, - cfg["hidden_size"], + hidden_size, device=device, - dtype=cfg["dtype"], + dtype=dtype, requires_grad=True, ) - w_query = torch.randn(cfg["hidden_size"], device=device, dtype=cfg["dtype"]) * 0.02 - w_norm = torch.ones(cfg["hidden_size"], device=device, dtype=cfg["dtype"]) + w_query = torch.randn(hidden_size, device=device, dtype=dtype) * 0.02 + w_norm = torch.ones(hidden_size, device=device, dtype=dtype) eps = cfg.get("eps", 1e-6) if input.kernel_provider == "liger": @@ -57,164 +63,58 @@ def _setup_attn_res(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider}") - return V, fn - - -def bench_speed_attn_res(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - V, fn = _setup_attn_res(input) - return run_speed_benchmark(fn, input.kernel_operation_mode, [V]) - - -def bench_memory_attn_res(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - V, fn = _setup_attn_res(input) - return run_memory_benchmark(fn, input.kernel_operation_mode) - - -def _resolve_model_config_attn_res(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_attn_res( - SingleBenchmarkRunInput( - x=cfg["seq_len"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "N": cfg["N"], - "bsz": cfg["bsz"], - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "eps": cfg.get("eps", 1e-6), - }, - ) - ) - - -def bench_speed_attn_res_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - V, fn = _resolve_model_config_attn_res(input) - return run_speed_benchmark(fn, input.kernel_operation_mode, [V]) - - -def bench_memory_attn_res_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - V, fn = _resolve_model_config_attn_res(input) - return run_memory_benchmark(fn, input.kernel_operation_mode) + return V, lambda _: fn() if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_seq_len): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_seq_len, - kernel_provider="pytorch", - extra_benchmark_config={ - "N": 8, - "bsz": 1, - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "eps": 1e-6, - }, - ) - V, fn = _setup_attn_res(probe_input) - return fn() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "attn_res", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "pytorch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "N": 8, - "bsz": sweep.batch_size, - "seq_len": sweep.seq_len, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_attn_res_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_attn_res_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="attn_res", + setup_fn=setup_attn_res, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "N": 8, + "bsz": 1, + "eps": 1e-6, + }, + probe_provider="pytorch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_seq_len, - kernel_provider="pytorch", - extra_benchmark_config={ - "N": 8, - "bsz": 1, - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - }, - ) - V, fn = _setup_attn_res(probe_input) - return fn() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_seq_len) - - common_configs = { - "kernel_name": "attn_res", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, int(math.log2(config.seq_len)) + 1)], - "kernel_providers": ["liger", "pytorch"], - "extra_benchmark_configs": [ - { - "N": 8, - "bsz": config.batch_size, - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_attn_res, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_attn_res, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="attn_res", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_attn_res, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "N": 8, + "bsz": 1, + "eps": 1e-6, + }, + probe_provider="pytorch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["pytorch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_attn_res), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_attn_res), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_cpo_loss.py b/benchmark/scripts/benchmark_cpo_loss.py index f2d89f5f..9c0d22b5 100644 --- a/benchmark/scripts/benchmark_cpo_loss.py +++ b/benchmark/scripts/benchmark_cpo_loss.py @@ -1,4 +1,3 @@ -import math import os import sys @@ -80,10 +79,7 @@ if __name__ == "__main__": model_keys=["hidden_size", "vocab_size", "dtype"], extra_configs={"T": T}, scale_dim="B", - probe_provider="huggingface", - x_values_fn=lambda config: [ - 2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1) - ], + x_label="batch size", overwrite=args.overwrite, ) diff --git a/benchmark/scripts/benchmark_cross_entropy.py b/benchmark/scripts/benchmark_cross_entropy.py index 14f55f58..6e1664eb 100644 --- a/benchmark/scripts/benchmark_cross_entropy.py +++ b/benchmark/scripts/benchmark_cross_entropy.py @@ -1,17 +1,18 @@ -import math +import os +import sys import torch -import triton + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from torch.nn import CrossEntropyLoss -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -21,242 +22,73 @@ from liger_kernel.utils import infer_device device = infer_device() -def _setup_cross_entropy(input: SingleBenchmarkRunInput): +def setup_cross_entropy(input: SingleBenchmarkRunInput): """Create input tensor, target, and CE loss from benchmark config.""" cfg = input.extra_benchmark_config - V = cfg["vocab_size"] - BT = input.x + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + V = model_cfg.vocab_size + BT = cfg["bsz"] * cfg["seq_len"] + else: + BT = input.x + V = cfg["vocab_size"] + _input = torch.randn(BT, V, requires_grad=True, device=device) - target = torch.randint(V, (BT, 1), device=device).squeeze(1) + target = torch.randint(V, (BT,), device=device) + if input.kernel_provider == "liger": loss_fn = LigerCrossEntropyLoss() - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": loss_fn = CrossEntropyLoss() else: raise ValueError(f"Invalid provider: {input.kernel_provider} for CrossEntropy") - return _input, target, loss_fn - -def bench_speed_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_cross_entropy(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "no-grad-forward": - with torch.no_grad(): - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_cross_entropy(input) - - def full(): - y = loss_fn(_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_cross_entropy(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_cross_entropy( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "vocab_size": model_info["vocab_size"], - }, - ) - ) - - -def bench_speed_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_cross_entropy(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "no-grad-forward": - with torch.no_grad(): - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "backward": - y = fwd() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_cross_entropy(input) - - def full(): - y = loss_fn(_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return _input, lambda x: loss_fn(x, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "vocab_size": model_cfg.vocab_size, - }, - ) - _input, target, loss_fn = _setup_cross_entropy(probe_input) - return loss_fn(_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "vocab_size": cfg.vocab_size, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "cross_entropy", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_cross_entropy_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_cross_entropy_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="cross_entropy", + setup_fn=setup_cross_entropy, + model_keys=["vocab_size"], + extra_configs={}, + probe_dim="BT", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "vocab_size": model.vocab_size, - }, - ) - _input, target, loss_fn = _setup_cross_entropy(probe_input) - return loss_fn(_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "cross_entropy", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "vocab_size": model.vocab_size, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_cross_entropy, - kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_cross_entropy, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="cross_entropy", + probe_x=1024, + model=model, + setup_fn=setup_cross_entropy, + model_keys=["vocab_size"], + extra_configs={}, + scale_dim="BT", + probe_provider="torch", + x_label="B * T", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_cross_entropy), + kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_cross_entropy), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_distill_cosine_loss.py b/benchmark/scripts/benchmark_distill_cosine_loss.py index 18564e2a..52f314ba 100644 --- a/benchmark/scripts/benchmark_distill_cosine_loss.py +++ b/benchmark/scripts/benchmark_distill_cosine_loss.py @@ -1,19 +1,18 @@ -import math import os import sys import torch import torch.nn as nn -import triton + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -22,9 +21,6 @@ from liger_kernel.utils import infer_device device = infer_device() -# Ensure the project root is in the path -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) - class TorchCosineSimilarityLoss(nn.Module): def __init__( @@ -89,17 +85,25 @@ class LigerCosineSimilarityLoss(nn.Module): ) -def _setup_distill_cosine_loss(input: SingleBenchmarkRunInput): +def setup_distill_cosine_loss(input: SingleBenchmarkRunInput): """Create input tensors and cosine similarity loss from benchmark config.""" cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + V = model_cfg.vocab_size + H = model_cfg.hidden_size + dtype = model_cfg.dtype + BT = cfg["bsz"] * cfg["seq_len"] + else: + BT = input.x + V = cfg["vocab_size"] + H = cfg["hidden_size"] + dtype = cfg["dtype"] + bias = cfg["bias"] weight_hard_loss = cfg["weight_hard_loss"] weight_soft_loss = cfg["weight_soft_loss"] ignore_index = cfg["ignore_index"] - BT = input.x _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) student_input = _tensor.detach().clone().requires_grad_(True) @@ -128,268 +132,62 @@ def _setup_distill_cosine_loss(input: SingleBenchmarkRunInput): ).to(device) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillCosineLoss") - return student_input, teacher_input, target, loss_module - - -def bench_speed_distill_cosine_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_module(student_input, teacher_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_distill_cosine_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(input) - - def full(): - y = loss_module(student_input, teacher_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_distill_cosine_loss(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_distill_cosine_loss( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "bias": cfg["bias"], - "weight_hard_loss": cfg["weight_hard_loss"], - "weight_soft_loss": cfg["weight_soft_loss"], - "ignore_index": cfg["ignore_index"], - }, - ) - ) - - -def bench_speed_distill_cosine_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _resolve_model_config_distill_cosine_loss(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_module(student_input, teacher_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_distill_cosine_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _resolve_model_config_distill_cosine_loss(input) - - def full(): - y = loss_module(student_input, teacher_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return student_input, lambda _: loss_module(student_input, teacher_input, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - }, - ) - student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(probe_input) - return loss_module(student_input, teacher_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "vocab_size": cfg.vocab_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "distill_cosine_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_distill_cosine_loss_model_config, - kernel_operation_modes=["forward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_distill_cosine_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="distill_cosine_loss", + setup_fn=setup_distill_cosine_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "bias": False, + "weight_hard_loss": 0.5, + "weight_soft_loss": 0.5, + "ignore_index": -100, + }, + probe_dim="BT", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - }, - ) - student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(probe_input) - return loss_module(student_input, teacher_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "distill_cosine_loss", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_distill_cosine_loss, - kernel_operation_modes=["forward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_distill_cosine_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="distill_cosine_loss", + probe_x=1024, + model=model, + setup_fn=setup_distill_cosine_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "bias": False, + "weight_hard_loss": 0.5, + "weight_soft_loss": 0.5, + "ignore_index": -100, + }, + scale_dim="BT", + x_label="B * T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_distill_cosine_loss), + kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_distill_cosine_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_distill_jsd_loss.py b/benchmark/scripts/benchmark_distill_jsd_loss.py index b56d0c12..6c56563d 100644 --- a/benchmark/scripts/benchmark_distill_jsd_loss.py +++ b/benchmark/scripts/benchmark_distill_jsd_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -93,17 +90,25 @@ class LigerJSDLoss(torch.nn.Module): ) -def _setup_distill_jsd_loss(input: SingleBenchmarkRunInput): +def setup_distill_jsd_loss(input: SingleBenchmarkRunInput): """Create input tensors and JSD loss from benchmark config.""" cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + V = model_cfg.vocab_size + H = model_cfg.hidden_size + dtype = model_cfg.dtype + BT = cfg["bsz"] * cfg["seq_len"] + else: + BT = input.x + V = cfg["vocab_size"] + H = cfg["hidden_size"] + dtype = cfg["dtype"] + bias = cfg["bias"] weight_hard_loss = cfg["weight_hard_loss"] weight_soft_loss = cfg["weight_soft_loss"] ignore_index = cfg["ignore_index"] - BT = input.x _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) student_input = _tensor.detach().clone().requires_grad_(True) @@ -132,264 +137,61 @@ def _setup_distill_jsd_loss(input: SingleBenchmarkRunInput): ).to(device) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillJSDLoss") - return student_input, teacher_input, target, loss_module - - -def bench_speed_distill_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_module(student_input, teacher_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_distill_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(input) - - def full(): - y = loss_module(student_input, teacher_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_distill_jsd_loss(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_distill_jsd_loss( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "bias": cfg["bias"], - "weight_hard_loss": cfg["weight_hard_loss"], - "weight_soft_loss": cfg["weight_soft_loss"], - "ignore_index": cfg["ignore_index"], - }, - ) - ) - - -def bench_speed_distill_jsd_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _resolve_model_config_distill_jsd_loss(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_module(student_input, teacher_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_distill_jsd_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, target, loss_module = _resolve_model_config_distill_jsd_loss(input) - - def full(): - y = loss_module(student_input, teacher_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return student_input, lambda _: loss_module(student_input, teacher_input, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - }, - ) - student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(probe_input) - return loss_module(student_input, teacher_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "vocab_size": cfg.vocab_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "distill_jsd_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_distill_jsd_loss_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_distill_jsd_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="distill_jsd_loss", + setup_fn=setup_distill_jsd_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "bias": False, + "weight_hard_loss": 0.5, + "weight_soft_loss": 0.5, + "ignore_index": -100, + }, + probe_dim="BT", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 - - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - }, - ) - student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(probe_input) - return loss_module(student_input, teacher_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "distill_jsd_loss", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_distill_jsd_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_distill_jsd_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="distill_jsd_loss", + probe_x=1024, + model=model, + setup_fn=setup_distill_jsd_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "bias": False, + "weight_hard_loss": 0.5, + "weight_soft_loss": 0.5, + "ignore_index": -100, + }, + scale_dim="BT", + x_label="B * T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_distill_jsd_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_distill_jsd_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_dpo_loss.py b/benchmark/scripts/benchmark_dpo_loss.py index a25cc46f..5aa59d84 100644 --- a/benchmark/scripts/benchmark_dpo_loss.py +++ b/benchmark/scripts/benchmark_dpo_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -23,20 +20,28 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_dpo_loss(input: SingleBenchmarkRunInput): +def setup_dpo_loss(input: SingleBenchmarkRunInput): """Create input tensors and DPO loss from benchmark config.""" from test.chunked_loss.test_dpo_loss import LigerLMHeadDPO from test.chunked_loss.test_dpo_loss import TorchLMHeadDPO cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] + T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + H = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + B = cfg["bsz"] + else: + B = input.x + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] + bias = cfg["bias"] beta = cfg["beta"] ignore_index = cfg["ignore_index"] - B = input.x - T = cfg["T"] # Input shape: [B, T, H] _input = torch.randn(B, T, H, device=device, dtype=dtype) @@ -51,275 +56,69 @@ def _setup_dpo_loss(input: SingleBenchmarkRunInput): if input.kernel_provider == "liger": loss_module = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": loss_module = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for DPOLoss") - fwd = lambda: loss_module(_input, ref_input, target)[0] + fwd = lambda x: loss_module(x, ref_input, target)[0] return _input, fwd -def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_dpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_dpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_dpo_loss(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_dpo_loss( - SingleBenchmarkRunInput( - x=cfg["B"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "T": cfg["T"], - "bias": cfg["bias"], - "beta": cfg["beta"], - "ignore_index": cfg["ignore_index"], - }, - ) - ) - - -def bench_speed_dpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_dpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_dpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_dpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - if __name__ == "__main__": args = parse_benchmark_script_args() + T = 512 if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - T = 512 - B = max(1, probe_bt // T) - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - }, - ) - _, fwd = _setup_dpo_loss(probe_input) - return fwd() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "vocab_size": cfg.vocab_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - T = 512 - B = max(1, sweep.bt // T) - - common_configs = { - "kernel_name": "dpo_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "B": B, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_dpo_loss_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_dpo_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="dpo_loss", + setup_fn=setup_dpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "bias": True, + "beta": 0.1, + "ignore_index": 42, + }, + probe_dim="B", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 512 - probe_bt = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - }, - ) - _, fwd = _setup_dpo_loss(probe_input) - return fwd() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "dpo_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, int(math.log2(config.batch_size * config.seq_len // T)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_dpo_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_dpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="dpo_loss", + probe_x=1, + model=model, + setup_fn=setup_dpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "bias": True, + "beta": 0.1, + "ignore_index": 42, + }, + scale_dim="B", + probe_provider="torch", + x_label="batch size", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_dpo_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_dpo_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_dyt.py b/benchmark/scripts/benchmark_dyt.py index e5ea479f..0271fcb2 100644 --- a/benchmark/scripts/benchmark_dyt.py +++ b/benchmark/scripts/benchmark_dyt.py @@ -1,19 +1,17 @@ -import math import os import sys import torch from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.utils import infer_device @@ -22,174 +20,83 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_dyt(input: SingleBenchmarkRunInput): +def setup_dyt(input: SingleBenchmarkRunInput): """Create input tensor and DyT layer from benchmark config.""" from test.transformers.test_dyt import LigerDyT from test.transformers.test_dyt import TorchDyT cfg = input.extra_benchmark_config - hidden_size = cfg["hidden_size"] - bt = input.x - x = torch.randn(bt, hidden_size, device=device, dtype=cfg["dtype"], requires_grad=True) + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + bt = cfg["seq_len"] * cfg["bsz"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + bt = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] + + beta = cfg["beta"] + x = torch.randn(bt, hidden_size, device=device, dtype=dtype, requires_grad=True) if input.kernel_provider == "liger": - layer = LigerDyT(hidden_size=hidden_size, beta=cfg["beta"]).to(device) + layer = LigerDyT(hidden_size=hidden_size, beta=beta).to(device) elif input.kernel_provider == "torch": - layer = TorchDyT(hidden_size=hidden_size, beta=cfg["beta"]).to(device) + layer = TorchDyT(hidden_size=hidden_size, beta=beta).to(device) elif input.kernel_provider == "torch_compile": - layer = torch.compile(TorchDyT(hidden_size=hidden_size, beta=cfg["beta"]).to(device)) + layer = torch.compile(TorchDyT(hidden_size=hidden_size, beta=beta).to(device)) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for DyT") return x, layer -def bench_speed_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_dyt(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_dyt(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_dyt(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - -def _resolve_model_config_dyt(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_dyt( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "beta": cfg["beta"], - }, - ) - ) - - -def bench_speed_dyt_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_dyt(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_dyt_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_dyt(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - if __name__ == "__main__": args = parse_benchmark_script_args() - if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - for beta in [False, True]: - - def _probe_factory(model_cfg, probe_bt, _beta=beta): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "beta": _beta, - }, - ) - x, layer = _setup_dyt(probe_input) - return layer(x) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - common_configs = { - "kernel_name": f"dyt_beta={beta}", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch", "torch_compile"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "beta": beta, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_dyt_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, + for beta in [False, True]: + if args.sweep_mode == "model_config": + common_configs = build_model_config_sweep( + kernel_name="dyt", + setup_fn=setup_dyt, + model_keys=["hidden_size", "dtype"], + probe_provider="torch", + extra_configs={ + "beta": beta, + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) - run_benchmarks( - bench_test_fn=bench_memory_dyt_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + else: + model = get_benchmark_model_config(args.model) + probe_seq_len = 1024 + + common_configs = build_token_length_sweep( + kernel_name="dyt", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_dyt, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "beta": beta, + }, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) - else: - model = get_benchmark_model_config(args.model) - probe_bt = 1024 - for beta in [False, True]: - - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "beta": beta, - }, - ) - x, layer = _setup_dyt(probe_input) - return layer(x) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": f"dyt_beta={beta}", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch", "torch_compile"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "beta": beta, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_dyt, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_dyt, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + common_configs["kernel_providers"] = ["liger", "torch", "torch_compile"] + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_dyt), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_dyt), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_embedding.py b/benchmark/scripts/benchmark_embedding.py index 0bca2111..fffefc44 100644 --- a/benchmark/scripts/benchmark_embedding.py +++ b/benchmark/scripts/benchmark_embedding.py @@ -1,19 +1,16 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from torch.nn import Embedding -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -28,20 +25,23 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../. # is needed. -def _setup_embedding(input: SingleBenchmarkRunInput): +def setup_embedding(input: SingleBenchmarkRunInput): """Create input tensors and embedding module from benchmark config.""" cfg = input.extra_benchmark_config - V = cfg.get("vocab_size", input.x) - D = cfg["hidden_size"] - dtype = cfg["dtype"] - BT = cfg.get("BT", input.x) - T = cfg.get("T", 512) - B = max(1, BT // T) if "BT" not in cfg else BT // T + T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + D = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + else: + BT = input.x + D = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] - # If BT is the x value, compute B from BT and T - if "BT" not in cfg: - B = max(1, input.x // T) - BT = B * T + B = max(1, BT // T) if "BT" not in cfg else BT // T input_ids = torch.randint(0, V, (B, T), device=device) @@ -49,217 +49,63 @@ def _setup_embedding(input: SingleBenchmarkRunInput): emb = LigerEmbedding(V, D).to(device).to(dtype) elif input.kernel_provider == "torch_compile": emb = torch.compile(Embedding(V, D).to(device).to(dtype)) - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": emb = Embedding(V, D).to(device).to(dtype) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for embedding") - fwd_fn = lambda: emb(input_ids) - return input_ids, fwd_fn - - -def bench_speed_embedding(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - input_ids, fwd = _setup_embedding(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - output = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: output.backward(torch.randn_like(output), retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[input_ids], - rep=100, - ) - elif mode == "full": - - def full(): - output = fwd() - output.backward(torch.randn_like(output)) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_embedding(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - input_ids, fwd_fn = _setup_embedding(input) - - def full(): - output = fwd_fn() - output.backward(torch.randn_like(output)) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_embedding(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_embedding( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "vocab_size": model_info["vocab_size"], - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "BT": cfg["BT"], - "T": cfg["T"], - }, - ) - ) - - -def bench_speed_embedding_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - input_ids, fwd_fn = _resolve_model_config_embedding(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, quantiles=QUANTILES, rep=100) - elif mode == "backward": - output = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: output.backward(torch.randn_like(output), retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[input_ids], - rep=100, - ) - elif mode == "full": - - def full(): - output = fwd_fn() - output.backward(torch.randn_like(output)) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_embedding_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - input_ids, fwd_fn = _resolve_model_config_embedding(input) - - def full(): - output = fwd_fn() - output.backward(torch.randn_like(output)) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return input_ids, emb if __name__ == "__main__": args = parse_benchmark_script_args() + T = 1024 if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - B = 2 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - T = max(1, probe_bt // B) - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "vocab_size": model_cfg.vocab_size, - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "BT": probe_bt, - "T": T, - }, - ) - _, fwd_fn = _setup_embedding(probe_input) - return fwd_fn() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"vocab_size": cfg.vocab_size, "hidden_size": cfg.hidden_size, "dtype": cfg.dtype} - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "embedding", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface", "torch_compile"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "BT": sweep.bt, "T": sweep.seq_len}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_embedding_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_embedding_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="embedding", + setup_fn=setup_embedding, + model_keys=["hidden_size", "vocab_size", "dtype"], + probe_provider="torch", + extra_configs={ + "T": T, + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 512 - probe_bt = 2048 + probe_seq_len = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "vocab_size": model.vocab_size, - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "BT": B * T, - "T": T, - }, - ) - _, fwd_fn = _setup_embedding(probe_input) - return fwd_fn() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "embedding", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.batch_size * config.seq_len))) + 1)], - "kernel_providers": ["liger", "huggingface", "torch_compile"], - "extra_benchmark_configs": [ - {"vocab_size": model.vocab_size, "hidden_size": model.hidden_size, "dtype": model.dtype, "T": T} - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_embedding, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_embedding, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="embedding", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_embedding, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + }, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "torch_compile", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_embedding), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_embedding), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_fused_add_rms_norm.py b/benchmark/scripts/benchmark_fused_add_rms_norm.py index 2ecf3903..f09b40dc 100644 --- a/benchmark/scripts/benchmark_fused_add_rms_norm.py +++ b/benchmark/scripts/benchmark_fused_add_rms_norm.py @@ -1,17 +1,13 @@ -import math - import torch import torch.nn as nn -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -62,267 +58,84 @@ class AddLigerRMSNorm(nn.Module): return self.weight * hidden_states.to(input_dtype), residual.to(input_dtype) -def _setup_fused_add_rms_norm(input: SingleBenchmarkRunInput): +def setup_fused_add_rms_norm(input: SingleBenchmarkRunInput): """Create input tensors and FusedAddRMSNorm layer from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg["hidden_size"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + bt = cfg["seq_len"] * cfg["bsz"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + bt = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] + eps = cfg["eps"] - x_shape = (input.x, hidden_size) - x = torch.randn(x_shape, dtype=cfg["dtype"], device=device, requires_grad=True) - r = torch.randn(x_shape, dtype=cfg["dtype"], device=device, requires_grad=True) + x_shape = (bt, hidden_size) + x = torch.randn(x_shape, dtype=dtype, device=device, requires_grad=True) + r = torch.randn(x_shape, dtype=dtype, device=device, requires_grad=True) if input.kernel_provider == "liger_fused_add_rms_norm": layer = LigerFusedAddRMSNorm(hidden_size=hidden_size, eps=eps).to(device) - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": layer = NaiveAddRMSNorm(hidden_size=hidden_size, eps=eps).to(device) elif input.kernel_provider == "liger_rms_norm": layer = AddLigerRMSNorm(hidden_size=hidden_size, eps=eps).to(device) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for FusedAddRMSNorm") - return x, r, layer - -def bench_speed_fused_add_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, r, layer = _setup_fused_add_rms_norm(input) - mode = input.kernel_operation_mode - dy = torch.randn_like(x) - ds = torch.randn_like(r) - - def y_fwd(): - return layer(x, r) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - y_fwd, - grad_to_none=[x, r], - rep=500, - quantiles=QUANTILES, - ) - elif mode == "backward": - y, s = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.backward((y, s), (dy, ds), retain_graph=True), - grad_to_none=[x, r], - rep=500, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y, s = y_fwd() - torch.autograd.backward((y, s), (dy, ds)) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - grad_to_none=[x, r], - rep=500, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_fused_add_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, r, layer = _setup_fused_add_rms_norm(input) - dy = torch.randn_like(x) - ds = torch.randn_like(r) - - def y_fwd(): - return layer(x, r) - - def full(): - y, s = y_fwd() - torch.autograd.backward((y, s), (dy, ds)) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_fused_add_rms_norm(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_fused_add_rms_norm( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "eps": cfg["eps"], - }, - ) - ) - - -def bench_speed_fused_add_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, r, layer = _resolve_model_config_fused_add_rms_norm(input) - mode = input.kernel_operation_mode - dy = torch.randn_like(x) - ds = torch.randn_like(r) - - def y_fwd(): - return layer(x, r) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(y_fwd, grad_to_none=[x, r], rep=100, quantiles=QUANTILES) - elif mode == "backward": - y, s = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.backward((y, s), (dy, ds), retain_graph=True), - grad_to_none=[x, r], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y, s = y_fwd() - torch.autograd.backward((y, s), (dy, ds)) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[x, r], rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_fused_add_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, r, layer = _resolve_model_config_fused_add_rms_norm(input) - dy = torch.randn_like(x) - ds = torch.randn_like(r) - - def y_fwd(): - return layer(x, r) - - def full(): - y, s = y_fwd() - torch.autograd.backward((y, s), (dy, ds)) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return x, lambda _: layer(x, r)[0] if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "eps": 1e-6, - }, - ) - x, r, layer = _setup_fused_add_rms_norm(probe_input) - y, s = layer(x, r) - return y + s # combine for backward probe - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "fused_add_rms_norm", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger_fused_add_rms_norm", "huggingface", "liger_rms_norm"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_add_rms_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_add_rms_norm_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="fused_add_rms_norm", + setup_fn=setup_fused_add_rms_norm, + model_keys=["hidden_size", "intermediate_size", "dtype"], + probe_provider="torch", + extra_configs={ + "eps": 1e-6, + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - }, - ) - x, r, layer = _setup_fused_add_rms_norm(probe_input) - y, s = layer(x, r) - return y + s - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "fused_add_rms_norm", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger_fused_add_rms_norm", "huggingface", "liger_rms_norm"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_add_rms_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_add_rms_norm, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="fused_add_rms_norm", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_fused_add_rms_norm, + model_keys=["hidden_size", "intermediate_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger_rms_norm", "liger_fused_add_rms_norm"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_fused_add_rms_norm), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_fused_add_rms_norm), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_fused_linear_cross_entropy.py index 0ff883ef..3bcf3307 100644 --- a/benchmark/scripts/benchmark_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_fused_linear_cross_entropy.py @@ -1,16 +1,12 @@ -import math - import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -51,13 +47,20 @@ class LigerLMHeadCE(torch.nn.Module): return self.ce_loss(self.lin.weight, x, y) -def _setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): +def setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): """Create input tensor, target, and fused linear CE from benchmark config.""" cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] - BT = input.x + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + V = model_cfg.vocab_size + H = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + BT = input.x + V = cfg["vocab_size"] + H = cfg["hidden_size"] + dtype = cfg["dtype"] _input = torch.randn(BT, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (BT, 1), dtype=torch.long, device=device).squeeze(1) @@ -68,256 +71,57 @@ def _setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, accum_dtype=torch.float32).to(device) else: lm_head_ce = TorchLMHeadCE(H=H, V=V, dtype=dtype).to(device) - return _input, target, lm_head_ce - - -def bench_speed_fused_linear_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(input) - mode = input.kernel_operation_mode - - def fwd(): - return lm_head_ce(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "no-grad-forward": - with torch.no_grad(): - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_fused_linear_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(input) - - def full(): - y = lm_head_ce(_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_fused_linear_cross_entropy( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - }, - ) - ) - - -def bench_speed_fused_linear_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, lm_head_ce = _resolve_model_config_fused_linear_cross_entropy(input) - mode = input.kernel_operation_mode - - def fwd(): - return lm_head_ce(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "no-grad-forward": - with torch.no_grad(): - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_fused_linear_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, lm_head_ce = _resolve_model_config_fused_linear_cross_entropy(input) - - def full(): - y = lm_head_ce(_input, target) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return _input, lambda _: lm_head_ce(_input, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - }, - ) - _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(probe_input) - return lm_head_ce(_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "vocab_size": cfg.vocab_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "fused_linear_cross_entropy", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "liger-fp32-accum", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_cross_entropy_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_cross_entropy_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="fused_linear_cross_entropy", + setup_fn=setup_fused_linear_cross_entropy, + model_keys=["hidden_size", "vocab_size", "dtype"], + probe_provider="torch", + extra_configs={ + "eps": 1e-6, + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - }, - ) - _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(probe_input) - return lm_head_ce(_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "fused_linear_cross_entropy", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "liger-fp32-accum", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_cross_entropy, - kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_cross_entropy, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="fused_linear_cross_entropy", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_fused_linear_cross_entropy, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger", "liger-fp32-accum"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_fused_linear_cross_entropy), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_fused_linear_cross_entropy), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_fused_linear_jsd.py b/benchmark/scripts/benchmark_fused_linear_jsd.py index b6fdf358..eefa49ef 100644 --- a/benchmark/scripts/benchmark_fused_linear_jsd.py +++ b/benchmark/scripts/benchmark_fused_linear_jsd.py @@ -1,16 +1,12 @@ -import math - import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -118,13 +114,20 @@ class LigerLMHeadJSD(torch.nn.Module): ) -def _setup_fused_linear_jsd(input: SingleBenchmarkRunInput): +def setup_fused_linear_jsd(input: SingleBenchmarkRunInput): """Create input tensors and fused linear JSD from benchmark config.""" cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] - BT = input.x + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + V = model_cfg.vocab_size + H = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + BT = input.x + V = cfg["vocab_size"] + H = cfg["hidden_size"] + dtype = cfg["dtype"] torch_lm_head_jsd = TorchLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) liger_lm_head_jsd = LigerLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) @@ -147,241 +150,49 @@ def _setup_fused_linear_jsd(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for FusedLinearJSD") - return student_input, teacher_input, lm_head - - -def bench_speed_fused_linear_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, lm_head = _setup_fused_linear_jsd(input) - mode = input.kernel_operation_mode - - def fwd(): - return lm_head(student_input, teacher_input) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_fused_linear_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, lm_head = _setup_fused_linear_jsd(input) - - def full(): - y = lm_head(student_input, teacher_input) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_fused_linear_jsd(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_fused_linear_jsd( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - }, - ) - ) - - -def bench_speed_fused_linear_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, lm_head = _resolve_model_config_fused_linear_jsd(input) - mode = input.kernel_operation_mode - - def fwd(): - return lm_head(student_input, teacher_input) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[student_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_fused_linear_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - student_input, teacher_input, lm_head = _resolve_model_config_fused_linear_jsd(input) - - def full(): - y = lm_head(student_input, teacher_input) - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return student_input, lambda _: lm_head(student_input, teacher_input) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - }, - ) - student_input, teacher_input, lm_head = _setup_fused_linear_jsd(probe_input) - return lm_head(student_input, teacher_input) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "vocab_size": cfg.vocab_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "fused_linear_jsd", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_jsd_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_jsd_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="fused_linear_jsd", + setup_fn=setup_fused_linear_jsd, + model_keys=["hidden_size", "vocab_size", "dtype"], + probe_dim="BT", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 - - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - }, - ) - student_input, teacher_input, lm_head = _setup_fused_linear_jsd(probe_input) - return lm_head(student_input, teacher_input) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "fused_linear_jsd", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_jsd, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_jsd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="fused_linear_jsd", + probe_x=1024, + model=model, + setup_fn=setup_fused_linear_jsd, + model_keys=["hidden_size", "vocab_size", "dtype"], + scale_dim="BT", + x_label="B * T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_fused_linear_jsd), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_fused_linear_jsd), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_geglu.py b/benchmark/scripts/benchmark_geglu.py index ecbe70c2..e862399b 100644 --- a/benchmark/scripts/benchmark_geglu.py +++ b/benchmark/scripts/benchmark_geglu.py @@ -1,19 +1,16 @@ -import math - import torch from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaMLP from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.transformers.geglu import LigerGEGLUMLP from liger_kernel.utils import infer_device @@ -21,186 +18,93 @@ from liger_kernel.utils import infer_device device = infer_device() -def _setup_geglu(input: SingleBenchmarkRunInput): +def setup_geglu(input: SingleBenchmarkRunInput): """Create input tensor and GEGLU layer from benchmark config.""" cfg = input.extra_benchmark_config + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + intermediate_size = model_cfg.intermediate_size + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + intermediate_size = cfg["intermediate_size"] + dtype = cfg["dtype"] + llama_config = LlamaConfig( - hidden_size=cfg["hidden_size"], - intermediate_size=cfg["intermediate_size"], + hidden_size=hidden_size, + intermediate_size=intermediate_size, hidden_act=cfg["hidden_act"], ) x = torch.randn( cfg["bsz"], - input.x, - cfg["hidden_size"], + seq_len, + hidden_size, device=device, - dtype=cfg["dtype"], + dtype=dtype, requires_grad=True, ) if input.kernel_provider == "liger": - layer = LigerGEGLUMLP(config=llama_config).to(device).to(cfg["dtype"]) + layer = LigerGEGLUMLP(config=llama_config).to(device).to(dtype) elif input.kernel_provider == "huggingface": - layer = LlamaMLP(config=llama_config).to(device).to(cfg["dtype"]) + layer = LlamaMLP(config=llama_config).to(device).to(dtype) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for GEGLU") return x, layer -def bench_speed_geglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_geglu(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_geglu(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_geglu(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - -def _resolve_model_config_geglu(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_geglu( - SingleBenchmarkRunInput( - x=cfg["seq_len"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "bsz": cfg["bsz"], - "hidden_size": model_info["hidden_size"], - "intermediate_size": model_info["intermediate_size"], - "hidden_act": cfg["hidden_act"], - "dtype": model_info["dtype"], - }, - ) - ) - - -def bench_speed_geglu_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_geglu(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_geglu_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_geglu(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_seq_len): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_seq_len, - kernel_provider="huggingface", - extra_benchmark_config={ - "bsz": 1, - "hidden_size": model_cfg.hidden_size, - "intermediate_size": model_cfg.intermediate_size, - "hidden_act": "gelu_pytorch_tanh", - "dtype": model_cfg.dtype, - }, - ) - x, layer = _setup_geglu(probe_input) - return layer(x) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "intermediate_size": cfg.intermediate_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "geglu", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "bsz": sweep.batch_size, - "seq_len": sweep.seq_len, - "hidden_act": "gelu_pytorch_tanh", - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_geglu_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_geglu_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="geglu", + setup_fn=setup_geglu, + model_keys=["hidden_size", "intermediate_size", "dtype"], + probe_provider="huggingface", + extra_configs={ + "bsz": 1, + "hidden_act": "gelu_pytorch_tanh", + }, + probe_dim="T", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_seq_len, - kernel_provider="huggingface", - extra_benchmark_config={ - "bsz": 1, - "hidden_size": model.hidden_size, - "intermediate_size": model.intermediate_size, - "hidden_act": "gelu_pytorch_tanh", - "dtype": model.dtype, - }, - ) - x, layer = _setup_geglu(probe_input) - return layer(x) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_seq_len) - - common_configs = { - "kernel_name": "geglu", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, int(math.log2(config.seq_len)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "bsz": config.batch_size, - "hidden_size": model.hidden_size, - "intermediate_size": model.intermediate_size, - "hidden_act": "gelu_pytorch_tanh", - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_geglu, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_geglu, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="geglu", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_geglu, + model_keys=["hidden_size", "intermediate_size", "dtype"], + extra_configs={ + "bsz": 1, + "hidden_act": "gelu_pytorch_tanh", + }, + scale_dim="T", + x_label="sequence length", + probe_provider="huggingface", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "huggingface"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_geglu), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_geglu), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_group_norm.py b/benchmark/scripts/benchmark_group_norm.py index b1c3d023..ae2b8063 100644 --- a/benchmark/scripts/benchmark_group_norm.py +++ b/benchmark/scripts/benchmark_group_norm.py @@ -1,17 +1,14 @@ -import math - import torch from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.transformers.group_norm import LigerGroupNorm from liger_kernel.utils import infer_device @@ -19,197 +16,88 @@ from liger_kernel.utils import infer_device device = infer_device() -def _setup_group_norm(input: SingleBenchmarkRunInput): +def setup_group_norm(input: SingleBenchmarkRunInput): """Create input tensor and GroupNorm layer from benchmark config.""" cfg = input.extra_benchmark_config - num_channels = cfg["num_channels"] - channels_per_group = cfg["channels_per_group"] H = cfg["H"] - eps = cfg["eps"] - num_groups = num_channels // channels_per_group - x = torch.randn( - input.x, - num_channels, - H, - device=device, - dtype=cfg["dtype"], - requires_grad=True, - ) - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + num_channels = model_cfg.hidden_size + dtype = model_cfg.dtype + M = max(1, cfg["bsz"] * cfg["seq_len"] // H) + else: + # input.x is BT; derive M from BT / H + M = max(1, input.x // H) + num_channels = cfg["hidden_size"] + dtype = cfg["dtype"] + + num_groups = num_channels // cfg["channels_per_group"] + x = torch.randn(M, num_channels, H, device=device, dtype=dtype, requires_grad=True) + if input.kernel_provider == "liger": - layer = LigerGroupNorm(num_channels=num_channels, num_groups=num_groups, eps=eps).to(device=device, dtype=dtype) - elif input.kernel_provider == "huggingface": - layer = torch.nn.GroupNorm(num_groups=num_groups, num_channels=num_channels, eps=eps).to( + layer = LigerGroupNorm(num_channels=num_channels, num_groups=num_groups, eps=cfg["eps"]).to( + device=device, dtype=dtype + ) + elif input.kernel_provider == "torch": + layer = torch.nn.GroupNorm(num_groups=num_groups, num_channels=num_channels, eps=cfg["eps"]).to( device=device, dtype=dtype ) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for GroupNorm") + return x, layer -def bench_speed_group_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_group_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_group_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_group_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - -def _resolve_model_config_group_norm(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_group_norm( - SingleBenchmarkRunInput( - x=cfg["M"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "num_channels": model_info["hidden_size"], - "channels_per_group": cfg["channels_per_group"], - "H": cfg["H"], - "dtype": model_info["dtype"], - "eps": cfg["eps"], - }, - ) - ) - - -def bench_speed_group_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_group_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_group_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_group_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - if __name__ == "__main__": args = parse_benchmark_script_args() + extra_configs = { + "channels_per_group": 4, + "eps": 1e-6, + "H": 512, + } + + probe_bt = 1024 + if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - channels_per_group = 4 - H = 512 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - M = max(1, probe_bt // H) - probe_input = SingleBenchmarkRunInput( - x=M, - kernel_provider="huggingface", - extra_benchmark_config={ - "num_channels": model_cfg.hidden_size, - "channels_per_group": channels_per_group, - "H": H, - "dtype": model_cfg.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_group_norm(probe_input) - return layer(x) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - M = max(1, sweep.bt // H) - - common_configs = { - "kernel_name": "group_norm", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "M": M, - "channels_per_group": channels_per_group, - "H": H, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_group_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_group_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="group_norm", + setup_fn=setup_group_norm, + model_keys=["hidden_size", "dtype"], + probe_provider="torch", + extra_configs=extra_configs, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - channels_per_group = 4 - H = 512 - probe_bt = 1024 - - def _probe(): - M = max(1, probe_bt // H) - probe_input = SingleBenchmarkRunInput( - x=M, - kernel_provider="huggingface", - extra_benchmark_config={ - "num_channels": model.hidden_size, - "channels_per_group": channels_per_group, - "H": H, - "dtype": model.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_group_norm(probe_input) - return layer(x) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "group_norm", - "x_name": "M", - "x_label": "batch size (M)", - "x_values": [2**i for i in range(2, int(math.log2(config.batch_size * config.seq_len // H)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "num_channels": model.hidden_size, - "channels_per_group": channels_per_group, - "H": H, - "dtype": model.dtype, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_group_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_group_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="group_norm", + probe_x=probe_bt, + model=model, + setup_fn=setup_group_norm, + model_keys=["hidden_size", "dtype"], + extra_configs=extra_configs, + scale_dim="BT", + x_label="B*T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_group_norm), + kernel_operation_modes=["full", "forward", "backward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_group_norm), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_grpo_loss.py b/benchmark/scripts/benchmark_grpo_loss.py index eee0ec88..9729759a 100644 --- a/benchmark/scripts/benchmark_grpo_loss.py +++ b/benchmark/scripts/benchmark_grpo_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -23,18 +20,26 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_grpo_loss(input: SingleBenchmarkRunInput): +def setup_grpo_loss(input: SingleBenchmarkRunInput): """Create input tensors and GRPO loss from benchmark config.""" from test.chunked_loss.test_grpo_loss import LigerLMHeadGRPO from test.chunked_loss.test_grpo_loss import TorchLMHeadGRPO cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] - importance_sampling_level = cfg["importance_sampling_level"] - B = input.x T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + H = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + B = cfg["bsz"] + else: + B = input.x + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] + + importance_sampling_level = cfg["importance_sampling_level"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) selected_token_ids = torch.randint(0, V, (B, T), dtype=torch.long, device=device) @@ -53,247 +58,66 @@ def _setup_grpo_loss(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for GRPOLoss") - fwd = lambda: loss_module(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[0] + fwd = lambda x: loss_module(x, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[0] return _input, fwd -def bench_speed_grpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_grpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_grpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_grpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_grpo_loss(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_grpo_loss( - SingleBenchmarkRunInput( - x=cfg["B"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "T": cfg["T"], - "importance_sampling_level": cfg["importance_sampling_level"], - }, - ) - ) - - -def bench_speed_grpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_grpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_grpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_grpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - def _run_grpo_benchmarks(args, importance_sampling_level, kernel_name_suffix): """Run D1 or D2 benchmarks for a given importance_sampling_level.""" kernel_name = f"fused_linear_grpo_loss_{kernel_name_suffix}" + T = 1024 + if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - T = 1024 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - B = max(1, probe_bt // T) - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "T": T, - "importance_sampling_level": importance_sampling_level, - }, - ) - _, fwd = _setup_grpo_loss(probe_input) - return fwd() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype} - for cfg in sweep.model_configs - } - B = max(1, sweep.bt // T) - - common_configs = { - "kernel_name": kernel_name, - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "B": B, - "T": T, - "importance_sampling_level": importance_sampling_level, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_grpo_loss_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_grpo_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name=kernel_name, + setup_fn=setup_grpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "importance_sampling_level": importance_sampling_level, + }, + probe_dim="B", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 1024 - probe_bt = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "importance_sampling_level": importance_sampling_level, - }, - ) - _, fwd = _setup_grpo_loss(probe_input) - return fwd() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": kernel_name, - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "importance_sampling_level": importance_sampling_level, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_grpo_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_grpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name=kernel_name, + probe_x=1, + model=model, + setup_fn=setup_grpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "importance_sampling_level": importance_sampling_level, + }, + scale_dim="B", + probe_provider="torch", + x_label="batch size", + overwrite=args.overwrite, ) + common_configs["kernel_providers"] = ["liger", "torch"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_grpo_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_grpo_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + if __name__ == "__main__": args = parse_benchmark_script_args() diff --git a/benchmark/scripts/benchmark_jsd.py b/benchmark/scripts/benchmark_jsd.py index b9af7fec..45b002d3 100644 --- a/benchmark/scripts/benchmark_jsd.py +++ b/benchmark/scripts/benchmark_jsd.py @@ -1,16 +1,12 @@ -import math - import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -58,11 +54,17 @@ class TorchJSD(torch.nn.Module): return loss.to(self.dtype) -def _setup_jsd(input: SingleBenchmarkRunInput): +def setup_jsd(input: SingleBenchmarkRunInput): """Create input tensors and JSD loss from benchmark config.""" cfg = input.extra_benchmark_config - V = cfg["vocab_size"] - BT = input.x + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + V = model_cfg.vocab_size + else: + BT = input.x + V = cfg["vocab_size"] + _input = torch.randn(BT, V, requires_grad=True, device=device).log_softmax(dim=-1) target = torch.randn(BT, V, device=device).log_softmax(dim=-1) @@ -72,223 +74,51 @@ def _setup_jsd(input: SingleBenchmarkRunInput): loss_fn = TorchJSD() else: raise ValueError(f"Invalid provider: {input.kernel_provider} for JSD") - return _input, target, loss_fn - - -def bench_speed_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_jsd(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_jsd(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_jsd(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_jsd( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "vocab_size": model_info["vocab_size"], - }, - ) - ) - - -def bench_speed_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_jsd(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_jsd(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return _input, lambda _: loss_fn(_input, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model_cfg.vocab_size, - }, - ) - _input, target, loss_fn = _setup_jsd(probe_input) - return loss_fn(_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "vocab_size": cfg.vocab_size, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "jsd", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_jsd_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_jsd_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="jsd", + setup_fn=setup_jsd, + model_keys=["vocab_size"], + probe_provider="torch", + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model.vocab_size, - }, - ) - _input, target, loss_fn = _setup_jsd(probe_input) - return loss_fn(_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "jsd", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "vocab_size": model.vocab_size, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_jsd, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_jsd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="jsd", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_jsd, + model_keys=["vocab_size"], + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_jsd), + kernel_operation_modes=["full", "forward", "backward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_jsd), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_kl_div.py b/benchmark/scripts/benchmark_kl_div.py index adb09311..42aa0729 100644 --- a/benchmark/scripts/benchmark_kl_div.py +++ b/benchmark/scripts/benchmark_kl_div.py @@ -1,17 +1,13 @@ -import math - import torch import torch.nn as nn -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -21,13 +17,18 @@ from liger_kernel.utils import infer_device device = infer_device() -def _setup_kl_div(input: SingleBenchmarkRunInput): +def setup_kl_div(input: SingleBenchmarkRunInput): """Create input tensors and KL div loss from benchmark config.""" cfg = input.extra_benchmark_config - V = cfg["vocab_size"] - BT = input.x - reduction = "batchmean" + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + V = model_cfg.vocab_size + else: + BT = input.x + V = cfg["vocab_size"] + reduction = cfg.get("reduction", "batchmean") _input = torch.randn(BT, V, requires_grad=True, device=device).log_softmax(dim=-1) target = torch.randn(BT, V, device=device).softmax(dim=-1) @@ -37,222 +38,57 @@ def _setup_kl_div(input: SingleBenchmarkRunInput): loss_fn = nn.KLDivLoss(reduction=reduction) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for KLDiv") - return _input, target, loss_fn - - -def bench_speed_kl_div(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_kl_div(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_kl_div(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_kl_div(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_kl_div(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_kl_div( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "vocab_size": model_info["vocab_size"], - }, - ) - ) - - -def bench_speed_kl_div_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_kl_div(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_kl_div_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_kl_div(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return _input, lambda _: loss_fn(_input, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model_cfg.vocab_size, - }, - ) - _input, target, loss_fn = _setup_kl_div(probe_input) - return loss_fn(_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: { - "vocab_size": cfg.vocab_size, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "kl_div", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_kl_div_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_kl_div_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="kl_div", + setup_fn=setup_kl_div, + model_keys=["vocab_size"], + probe_provider="torch", + extra_configs={ + "reduction": "batchmean", + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model.vocab_size, - }, - ) - _input, target, loss_fn = _setup_kl_div(probe_input) - return loss_fn(_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "kl_div", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "vocab_size": model.vocab_size, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_kl_div, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_kl_div, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="kl_div", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_kl_div, + model_keys=["vocab_size"], + extra_configs={ + "reduction": "batchmean", + }, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_kl_div), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_kl_div), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_kto_loss.py b/benchmark/scripts/benchmark_kto_loss.py index aff1c637..aa04ba9c 100644 --- a/benchmark/scripts/benchmark_kto_loss.py +++ b/benchmark/scripts/benchmark_kto_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -93,17 +90,25 @@ class LigerLMHeadKTO(torch.nn.Module): ) -def _setup_kto_loss(input: SingleBenchmarkRunInput): +def setup_kto_loss(input: SingleBenchmarkRunInput): """Create input tensors and KTO loss from benchmark config.""" cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] + T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + H = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + B = cfg["bsz"] + else: + B = input.x + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] + bias = cfg["bias"] beta = cfg["beta"] ignore_index = cfg["ignore_index"] - B = input.x - T = cfg["T"] # Input shape: [B, T, H] _input = torch.randn(B, T, H, device=device, dtype=dtype) @@ -131,230 +136,64 @@ def _setup_kto_loss(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for KTOLoss") - fwd = lambda: loss_module(x=_input, ref_x=ref_input, y=target, preference_labels=preference_labels, kl=kl)[0] + fwd = lambda _input: loss_module(x=_input, ref_x=ref_input, y=target, preference_labels=preference_labels, kl=kl)[0] return _input, fwd -def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_kto_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_kto_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_kto_loss(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_kto_loss( - SingleBenchmarkRunInput( - x=cfg["B"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "T": cfg["T"], - "bias": cfg["bias"], - "beta": cfg["beta"], - "ignore_index": cfg["ignore_index"], - }, - ) - ) - - -def bench_speed_kto_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_kto_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), grad_to_none=[_input], rep=100, quantiles=QUANTILES - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_kto_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_kto_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - if __name__ == "__main__": args = parse_benchmark_script_args() + T = 1024 + if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - T = 512 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - B = max(1, probe_bt // T) - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - }, - ) - _, fwd = _setup_kto_loss(probe_input) - return fwd() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype} - for cfg in sweep.model_configs - } - B = max(1, sweep.bt // T) - - common_configs = { - "kernel_name": "kto_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - {"model_configs": model_configs_info, "B": B, "T": T, "bias": True, "beta": 0.1, "ignore_index": 42} - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_kto_loss_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_kto_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="kto_loss", + setup_fn=setup_kto_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "bias": True, + "beta": 0.1, + "ignore_index": 42, + }, + probe_dim="B", + probe_provider="huggingface", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 512 - probe_bt = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - }, - ) - _, fwd = _setup_kto_loss(probe_input) - return fwd() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "kto_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_kto_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_kto_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="kto_loss", + probe_x=1, + model=model, + setup_fn=setup_kto_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={ + "T": T, + "bias": True, + "beta": 0.1, + "ignore_index": 42, + }, + scale_dim="B", + x_label="batch size", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "huggingface"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_kto_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_kto_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_layer_norm.py b/benchmark/scripts/benchmark_layer_norm.py index 7e8c5c2a..33f96b66 100644 --- a/benchmark/scripts/benchmark_layer_norm.py +++ b/benchmark/scripts/benchmark_layer_norm.py @@ -39,7 +39,7 @@ def setup_layer_norm(input: SingleBenchmarkRunInput): ) if input.kernel_provider == "liger": layer = LigerLayerNorm(hidden_size=hidden_size, eps=eps).to(device).to(dtype) - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": layer = torch.nn.LayerNorm(hidden_size, eps=eps).to(device).to(dtype) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for LayerNorm") @@ -57,8 +57,8 @@ if __name__ == "__main__": extra_configs={ "eps": 1e-6, }, - probe_dim="BT", - probe_provider="huggingface", + probe_dim="T", + probe_provider="torch", bt=args.bt, overwrite=args.overwrite, ) @@ -76,23 +76,24 @@ if __name__ == "__main__": extra_configs={ "eps": 1e-6, }, - scale_dim="BT", - probe_provider="huggingface", + scale_dim="T", + x_label="Sequence length", + probe_provider="torch", overwrite=args.overwrite, ) - common_configs["kernel_providers"] = ["liger", "huggingface"] + common_configs["kernel_providers"] = ["liger", "torch"] run_benchmarks( bench_test_fn=build_speed_bench_fn(setup_layer_norm), - kernel_operation_modes=["full", "forward", "backward"], + kernel_operation_modes=["forward", "backward", "full"], metric_name="speed", metric_unit="ms", **common_configs, ) run_benchmarks( bench_test_fn=build_memory_bench_fn(setup_layer_norm), - kernel_operation_modes=["full", "forward", "backward"], + kernel_operation_modes=["full"], metric_name="memory", metric_unit="MB", **common_configs, diff --git a/benchmark/scripts/benchmark_llama4_rope.py b/benchmark/scripts/benchmark_llama4_rope.py index 7a239b2b..3e71dceb 100644 --- a/benchmark/scripts/benchmark_llama4_rope.py +++ b/benchmark/scripts/benchmark_llama4_rope.py @@ -1,21 +1,15 @@ -import math -import os -import sys - import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from transformers.models.llama4.configuration_llama4 import Llama4TextConfig from transformers.models.llama4.modeling_llama4 import Llama4TextRotaryEmbedding from transformers.models.llama4.modeling_llama4 import apply_rotary_emb -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -25,21 +19,26 @@ from liger_kernel.utils import transformers_version_dispatch device = infer_device() -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) - -def _setup_llama4_rope(input: SingleBenchmarkRunInput): +def setup_llama4_rope(input: SingleBenchmarkRunInput): """Create input tensors and Llama4 RoPE embedding from benchmark config.""" cfg = input.extra_benchmark_config - num_q_heads = cfg["num_q_heads"] - num_kv_heads = cfg["num_kv_heads"] - dtype = cfg["dtype"] - hidden_size = cfg.get("hidden_size", input.x) - seq_len = cfg.get("seq_len", input.x) + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + num_q_heads = model_cfg.num_attention_heads + num_kv_heads = model_cfg.num_key_value_heads + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + num_q_heads = cfg["num_attention_heads"] + num_kv_heads = cfg["num_key_value_heads"] + dtype = cfg["dtype"] head_dim = hidden_size // num_q_heads - # Create Llama4TextConfig for the rotary embedding config = Llama4TextConfig( hidden_size=hidden_size, num_attention_heads=num_q_heads, @@ -68,10 +67,6 @@ def _setup_llama4_rope(input: SingleBenchmarkRunInput): requires_grad=True, dtype=dtype, ) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device), - ) pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) freqs_cis = rotary_emb(q, pos_ids) @@ -82,234 +77,51 @@ def _setup_llama4_rope(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for Llama4 RoPE embedding") - return q, k, dq, dk, fwd_fn - - -def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd = _setup_llama4_rope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "backward": - q_out, k_out = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _setup_llama4_rope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_llama4_rope(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_llama4_rope( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "num_q_heads": model_info["num_q_heads"], - "num_kv_heads": model_info["num_kv_heads"], - "dtype": model_info["dtype"], - "seq_len": cfg["seq_len"], - }, - ) - ) - - -def bench_speed_llama4_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_llama4_rope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - elif mode == "backward": - q_out, k_out = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_llama4_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_llama4_rope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory( - full, - quantiles=QUANTILES, - ) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return q, lambda _: fwd_fn()[0] if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "num_q_heads": model_cfg.num_attention_heads, - "num_kv_heads": model_cfg.num_key_value_heads, - "dtype": model_cfg.dtype, - "seq_len": probe_bt, - }, - ) - _, _, _, _, fwd_fn = _setup_llama4_rope(probe_input) - return fwd_fn()[0] - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "num_q_heads": cfg.num_attention_heads, - "num_kv_heads": cfg.num_key_value_heads, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "llama4_rope", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_llama4_rope_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_llama4_rope_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="llama4_rope", + setup_fn=setup_llama4_rope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + probe_provider="huggingface", + probe_dim="T", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) probe_seq_len = 2048 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - "seq_len": probe_seq_len, - }, - ) - _, _, _, _, fwd_fn = _setup_llama4_rope(probe_input) - return fwd_fn()[0] - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_seq_len) - - common_configs = { - "kernel_name": "llama4_rope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_llama4_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_llama4_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="llama4_rope", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_llama4_rope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + scale_dim="T", + x_label="sequence length", + probe_provider="huggingface", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["huggingface", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_llama4_rope), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_llama4_rope), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_model_configs.py b/benchmark/scripts/benchmark_model_configs.py index 40030167..192495c9 100644 --- a/benchmark/scripts/benchmark_model_configs.py +++ b/benchmark/scripts/benchmark_model_configs.py @@ -530,6 +530,7 @@ def build_token_length_sweep( model_keys: List[str], extra_configs: Optional[Dict] = None, scale_dim: Literal["T", "B", "BT"] = "T", + scaling_method: Literal["linear", "quadratic"] = "linear", forward_fn: Callable[..., torch.Tensor] = default_forward_fn, probe_provider: str = "huggingface", x_label: str = "sequence length", @@ -601,12 +602,15 @@ def build_token_length_sweep( probe_fn=probe_fn, probe_seq_len=probe_seq_len, probe_batch_size=probe_batch_size, + scaling_method=scaling_method, ) if x_values_fn is None: if scale_dim == "T": x_values_fn = lambda cfg: [2**i for i in range(10, int(math.log2(cfg.seq_len)) + 1)] elif scale_dim == "B": - x_values_fn = lambda cfg: [2**i for i in range(0, int(math.log2(cfg.batch_size)) + 1)] + x_values_fn = lambda cfg: [ + 2**i for i in range(0, int(math.log2(max(2, cfg.batch_size * cfg.seq_len // T))) + 1) + ] elif scale_dim == "BT": x_values_fn = lambda cfg: [2**i for i in range(10, int(math.log2(cfg.seq_len * cfg.batch_size)) + 1)] @@ -616,11 +620,15 @@ def build_token_length_sweep( "x_label": x_label, "x_values": x_values_fn(config), "extra_benchmark_configs": [ - build_extra_config( - model, - model_keys, - extra_configs=extra_configs, - ) + { + **build_extra_config( + model, + model_keys, + extra_configs=extra_configs, + ), + "bsz": config.batch_size, + "seq_len": config.seq_len, + } ], "overwrite": overwrite, } diff --git a/benchmark/scripts/benchmark_orpo_loss.py b/benchmark/scripts/benchmark_orpo_loss.py index 8de3b4d5..e7af278f 100644 --- a/benchmark/scripts/benchmark_orpo_loss.py +++ b/benchmark/scripts/benchmark_orpo_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -23,17 +20,24 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_orpo_loss(input: SingleBenchmarkRunInput): +def setup_orpo_loss(input: SingleBenchmarkRunInput): """Create input tensors and ORPO loss from benchmark config.""" from test.chunked_loss.test_orpo_loss import LigerLMHeadORPO from test.chunked_loss.test_orpo_loss import TorchLMHeadORPO cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] - B = input.x T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + H = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + B = cfg["bsz"] + else: + B = input.x + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (B, T), dtype=torch.long, device=device) @@ -46,227 +50,53 @@ def _setup_orpo_loss(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for ORPOLoss") - fwd = lambda: loss_module(_input, target, nll_target)[0] + fwd = lambda x: loss_module(x, target, nll_target)[0] return _input, fwd -def bench_speed_orpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_orpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), grad_to_none=[_input], rep=100, quantiles=QUANTILES - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_orpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_orpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_orpo_loss(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_orpo_loss( - SingleBenchmarkRunInput( - x=cfg["B"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "T": cfg["T"], - }, - ) - ) - - -def bench_speed_orpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_orpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_orpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_orpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - if __name__ == "__main__": args = parse_benchmark_script_args() + T = 1024 if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - T = 1024 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - B = max(1, probe_bt // T) - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "T": T, - }, - ) - _, fwd = _setup_orpo_loss(probe_input) - return fwd() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype} - for cfg in sweep.model_configs - } - B = max(1, sweep.bt // T) - - common_configs = { - "kernel_name": "fused_linear_orpo_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "B": B, "T": T}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_orpo_loss_model_config, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_orpo_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="orpo_loss", + setup_fn=setup_orpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={"T": T}, + probe_dim="B", + probe_provider="huggingface", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 1024 - probe_bt = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - }, - ) - _, fwd = _setup_orpo_loss(probe_input) - return fwd() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "fused_linear_orpo_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - {"hidden_size": model.hidden_size, "vocab_size": model.vocab_size, "dtype": model.dtype, "T": T} - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_orpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_orpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="orpo_loss", + probe_x=1, + model=model, + setup_fn=setup_orpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={"T": T}, + scale_dim="B", + x_label="batch size", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "huggingface"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_orpo_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_orpo_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_poly_norm.py b/benchmark/scripts/benchmark_poly_norm.py index c1d7bd2d..4c397554 100644 --- a/benchmark/scripts/benchmark_poly_norm.py +++ b/benchmark/scripts/benchmark_poly_norm.py @@ -1,18 +1,15 @@ -import math - import torch import torch.nn as nn from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.transformers.poly_norm import LigerPolyNorm from liger_kernel.utils import infer_device @@ -72,172 +69,85 @@ class NaivePolyNorm(nn.Module): return output.to(input_dtype) -def _setup_poly_norm(input: SingleBenchmarkRunInput): +def setup_poly_norm(input: SingleBenchmarkRunInput): """Create input tensor and PolyNorm layer from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg["hidden_size"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] + eps = cfg["eps"] x = torch.randn( - input.x, + seq_len, hidden_size, device=device, - dtype=cfg["dtype"], + dtype=dtype, requires_grad=True, ) if input.kernel_provider == "liger": layer = LigerPolyNorm(eps=eps).to(device) - elif input.kernel_provider == "huggingface": + elif input.kernel_provider == "torch": layer = NaivePolyNorm(eps=eps).to(device) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for PolyNorm") return x, layer -def bench_speed_poly_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_poly_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_poly_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_poly_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - -def _resolve_model_config_poly_norm(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_poly_norm( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "eps": cfg["eps"], - }, - ) - ) - - -def bench_speed_poly_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_poly_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_poly_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_poly_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_poly_norm(probe_input) - return layer(x) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "poly_norm", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_poly_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_poly_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="poly_norm", + setup_fn=setup_poly_norm, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + probe_dim="T", + probe_provider="torch", + bt=args.bt, + overwrite=args.overwrite, ) + else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_poly_norm(probe_input) - return layer(x) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "poly_norm", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_poly_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_poly_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="poly_norm", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_poly_norm, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + scale_dim="T", + x_label="Sequence length", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_poly_norm), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_poly_norm), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_qwen2vl_mrope.py b/benchmark/scripts/benchmark_qwen2vl_mrope.py index 869363a5..ed7cbfb8 100644 --- a/benchmark/scripts/benchmark_qwen2vl_mrope.py +++ b/benchmark/scripts/benchmark_qwen2vl_mrope.py @@ -1,21 +1,18 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLTextConfig from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLRotaryEmbedding from transformers.models.qwen2_vl.modeling_qwen2_vl import apply_multimodal_rotary_pos_emb -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -27,14 +24,22 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_qwen2vl_mrope(input: SingleBenchmarkRunInput): +def setup_qwen2vl_mrope(input: SingleBenchmarkRunInput): """Create input tensors and Qwen2VL M-RoPE embedding from benchmark config.""" cfg = input.extra_benchmark_config - num_q_heads = cfg["num_q_heads"] - num_kv_heads = cfg["num_kv_heads"] - dtype = cfg["dtype"] - hidden_size = cfg.get("hidden_size", input.x) - seq_len = cfg.get("seq_len", input.x) + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + num_q_heads = model_cfg.num_attention_heads + num_kv_heads = model_cfg.num_key_value_heads + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + num_q_heads = cfg["num_attention_heads"] + num_kv_heads = cfg["num_key_value_heads"] + dtype = cfg["dtype"] head_dim = hidden_size // num_q_heads mrope_section_hw = head_dim * 3 // 16 @@ -63,10 +68,7 @@ def _setup_qwen2vl_mrope(input: SingleBenchmarkRunInput): requires_grad=True, dtype=dtype, ).transpose(1, 2) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device, dtype=dtype), - ) + pos_ids = torch.arange(seq_len * 3, device=device, dtype=torch.long).view(3, 1, -1) cos, sin = rotary_emb(k, pos_ids) @@ -77,241 +79,51 @@ def _setup_qwen2vl_mrope(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for M-RoPE embedding") - return q, k, dq, dk, fwd_fn - - -def bench_speed_qwen2vl_mrope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd = _setup_qwen2vl_mrope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "backward": - q_out, k_out = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_qwen2vl_mrope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _setup_qwen2vl_mrope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory( - full, - quantiles=QUANTILES, - ) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_qwen2vl_mrope(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_qwen2vl_mrope( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "num_q_heads": model_info["num_q_heads"], - "num_kv_heads": model_info["num_kv_heads"], - "dtype": model_info["dtype"], - "seq_len": cfg["seq_len"], - }, - ) - ) - - -def bench_speed_qwen2vl_mrope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_qwen2vl_mrope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - elif mode == "backward": - q_out, k_out = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_qwen2vl_mrope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_qwen2vl_mrope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory( - full, - quantiles=QUANTILES, - ) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return q, lambda _: fwd_fn()[0] if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "num_q_heads": model_cfg.num_attention_heads, - "num_kv_heads": model_cfg.num_key_value_heads, - "dtype": model_cfg.dtype, - "seq_len": probe_bt, - }, - ) - _, _, _, _, fwd_fn = _setup_qwen2vl_mrope(probe_input) - return fwd_fn()[0] - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "num_q_heads": cfg.num_attention_heads, - "num_kv_heads": cfg.num_key_value_heads, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "qwen2vl_mrope", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_qwen2vl_mrope_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_qwen2vl_mrope_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="qwen2vl_mrope", + setup_fn=setup_qwen2vl_mrope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + probe_provider="huggingface", + probe_dim="T", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) probe_seq_len = 2048 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - "seq_len": probe_seq_len, - }, - ) - _, _, _, _, fwd_fn = _setup_qwen2vl_mrope(probe_input) - return fwd_fn()[0] - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_seq_len) - - common_configs = { - "kernel_name": "qwen2vl_mrope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_qwen2vl_mrope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_qwen2vl_mrope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="qwen2vl_mrope", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_qwen2vl_mrope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + scale_dim="T", + x_label="sequence length", + probe_provider="huggingface", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["huggingface", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_qwen2vl_mrope), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_qwen2vl_mrope), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_relu_squared.py b/benchmark/scripts/benchmark_relu_squared.py index a03cc87b..3ea88ac2 100644 --- a/benchmark/scripts/benchmark_relu_squared.py +++ b/benchmark/scripts/benchmark_relu_squared.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -30,15 +27,20 @@ class TorchReLUSquared(torch.nn.Module): return torch.square(relu_applied) -def _setup_relu_squared(input: SingleBenchmarkRunInput): - """Create input tensors and relu_squared module from benchmark config.""" +def setup_relu_squared(input: SingleBenchmarkRunInput): + """Create input tensor and relu_squared module from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg["hidden_size"] - M = cfg.get("M", input.x) - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + bt = cfg["seq_len"] * cfg["bsz"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + bt = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] - x = torch.randn(M, hidden_size, dtype=dtype, device=device, requires_grad=True) - dy = torch.randn_like(x) + x = torch.randn(bt, hidden_size, dtype=dtype, device=device, requires_grad=True) if input.kernel_provider == "liger": relu_sq = LigerReLUSquared().to(device) @@ -47,203 +49,50 @@ def _setup_relu_squared(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for relu_squared") - fwd_fn = lambda: relu_sq(x) - return x, dy, fwd_fn - - -def bench_speed_relu_squared(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _setup_relu_squared(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(y_fwd, quantiles=QUANTILES, grad_to_none=[x], rep=500) - elif mode == "backward": - y = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[x], - rep=500, - ) - elif mode == "full": - - def full(): - y = y_fwd() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, grad_to_none=[x], rep=500) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_relu_squared(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _setup_relu_squared(input) - - def full(): - y = y_fwd() - y.backward(torch.ones_like(y), retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_relu_squared(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_relu_squared( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "M": cfg["M"], - }, - ) - ) - - -def bench_speed_relu_squared_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _resolve_model_config_relu_squared(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, quantiles=QUANTILES, grad_to_none=[x], rep=500) - elif mode == "backward": - y = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[x], - rep=500, - ) - elif mode == "full": - - def full(): - y = fwd_fn() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, grad_to_none=[x], rep=500) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_relu_squared_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _resolve_model_config_relu_squared(input) - - def full(): - y = fwd_fn() - y.backward(torch.ones_like(y), retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return x, relu_sq if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "M": probe_bt, - }, - ) - _, _, fwd_fn = _setup_relu_squared(probe_input) - return fwd_fn() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "dtype": cfg.dtype} for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "relu_squared", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "M": sweep.bt}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_relu_squared_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_relu_squared_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="relu_squared", + setup_fn=setup_relu_squared, + model_keys=["hidden_size", "dtype"], + probe_provider="torch", + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 2048 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "M": probe_bt, - }, - ) - _, _, fwd_fn = _setup_relu_squared(probe_input) - return fwd_fn() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "relu_squared", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.batch_size * config.seq_len))) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"hidden_size": model.hidden_size, "dtype": model.dtype}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_relu_squared, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_relu_squared, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="relu_squared", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_relu_squared, + model_keys=["hidden_size", "dtype"], + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "torch"] + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_relu_squared), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_relu_squared), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_rms_norm.py b/benchmark/scripts/benchmark_rms_norm.py index ca8004fd..886bc1a6 100644 --- a/benchmark/scripts/benchmark_rms_norm.py +++ b/benchmark/scripts/benchmark_rms_norm.py @@ -1,18 +1,15 @@ -import math - import torch import torch.nn as nn from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks -from utils import run_memory_benchmark -from utils import run_speed_benchmark from liger_kernel.transformers.rms_norm import LigerRMSNorm from liger_kernel.utils import infer_device @@ -37,16 +34,25 @@ class LlamaRMSNorm(nn.Module): return self.weight * hidden_states.to(input_dtype) -def _setup_rms_norm(input: SingleBenchmarkRunInput): +def setup_rms_norm(input: SingleBenchmarkRunInput): """Create input tensor and RMSNorm layer from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg["hidden_size"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] + eps = cfg["eps"] x = torch.randn( - input.x, + seq_len, hidden_size, device=device, - dtype=cfg["dtype"], + dtype=dtype, requires_grad=True, ) if input.kernel_provider == "liger": @@ -58,151 +64,55 @@ def _setup_rms_norm(input: SingleBenchmarkRunInput): return x, layer -def bench_speed_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_rms_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _setup_rms_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - -def _resolve_model_config_rms_norm(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_rms_norm( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "eps": cfg["eps"], - }, - ) - ) - - -def bench_speed_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_rms_norm(input) - return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x]) - - -def bench_memory_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, layer = _resolve_model_config_rms_norm(input) - return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode) - - if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_rms_norm(probe_input) - return layer(x) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "rms_norm", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_rms_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_rms_norm_model_config, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="rms_norm", + setup_fn=setup_rms_norm, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + probe_dim="T", + probe_provider="huggingface", + bt=args.bt, + overwrite=args.overwrite, ) + else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - }, - ) - x, layer = _setup_rms_norm(probe_input) - return layer(x) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "rms_norm", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "eps": 1e-6, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_rms_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_rms_norm, - kernel_operation_modes=["full", "forward", "backward"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="rms_norm", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_rms_norm, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "eps": 1e-6, + }, + scale_dim="T", + x_label="Sequence length", + probe_provider="huggingface", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["huggingface", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_rms_norm), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_rms_norm), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_rope.py b/benchmark/scripts/benchmark_rope.py index a4670654..d7e1868c 100644 --- a/benchmark/scripts/benchmark_rope.py +++ b/benchmark/scripts/benchmark_rope.py @@ -1,21 +1,18 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding from transformers.models.llama.modeling_llama import apply_rotary_pos_emb -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -28,14 +25,22 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_rope(input: SingleBenchmarkRunInput): +def setup_rope(input: SingleBenchmarkRunInput): """Create input tensors and RoPE embedding from benchmark config.""" cfg = input.extra_benchmark_config - num_q_heads = cfg["num_q_heads"] - num_kv_heads = cfg["num_kv_heads"] - dtype = cfg["dtype"] - hidden_size = cfg.get("hidden_size", input.x) - seq_len = cfg.get("seq_len", input.x) + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + seq_len = cfg["seq_len"] + hidden_size = model_cfg.hidden_size + num_q_heads = model_cfg.num_attention_heads + num_kv_heads = model_cfg.num_key_value_heads + dtype = model_cfg.dtype + else: + seq_len = input.x + hidden_size = cfg["hidden_size"] + num_q_heads = cfg["num_attention_heads"] + num_kv_heads = cfg["num_key_value_heads"] + dtype = cfg["dtype"] head_dim = hidden_size // num_q_heads rotary_emb = transformers_version_dispatch( @@ -57,10 +62,7 @@ def _setup_rope(input: SingleBenchmarkRunInput): requires_grad=True, dtype=dtype, ).transpose(1, 2) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device), - ) + pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) cos, sin = rotary_emb(k, pos_ids) @@ -71,230 +73,51 @@ def _setup_rope(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for RoPE embedding") - return q, k, dq, dk, fwd_fn - - -def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd = _setup_rope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "backward": - q_out, k_out = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _setup_rope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_rope(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_rope( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "num_q_heads": model_info["num_q_heads"], - "num_kv_heads": model_info["num_kv_heads"], - "dtype": model_info["dtype"], - "seq_len": cfg["seq_len"], - }, - ) - ) - - -def bench_speed_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_rope(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - elif mode == "backward": - q_out, k_out = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), - grad_to_none=[q, k], - rep=400, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - q, k, dq, dk, fwd_fn = _resolve_model_config_rope(input) - - def full(): - q_out, k_out = fwd_fn() - torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory( - full, - quantiles=QUANTILES, - ) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return q, lambda _: fwd_fn()[0] if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "num_q_heads": model_cfg.num_attention_heads, - "num_kv_heads": model_cfg.num_key_value_heads, - "dtype": model_cfg.dtype, - "seq_len": probe_bt, - }, - ) - _, _, _, _, fwd_fn = _setup_rope(probe_input) - return fwd_fn()[0] # return q_out for memory estimation - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: { - "hidden_size": cfg.hidden_size, - "num_q_heads": cfg.num_attention_heads, - "num_kv_heads": cfg.num_key_value_heads, - "dtype": cfg.dtype, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "rope", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_rope_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_rope_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="rope", + setup_fn=setup_rope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + probe_provider="huggingface", + probe_dim="T", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) probe_seq_len = 2048 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - "seq_len": probe_seq_len, - }, - ) - _, _, _, _, fwd_fn = _setup_rope(probe_input) - return fwd_fn()[0] - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_seq_len) - - common_configs = { - "kernel_name": "rope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "hidden_size": model.hidden_size, - "num_q_heads": model.num_attention_heads, - "num_kv_heads": model.num_key_value_heads, - "dtype": model.dtype, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="rope", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_rope, + model_keys=["hidden_size", "num_attention_heads", "num_key_value_heads", "dtype"], + scale_dim="T", + x_label="sequence length", + probe_provider="huggingface", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["huggingface", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_rope), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_rope), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_simpo_loss.py b/benchmark/scripts/benchmark_simpo_loss.py index ce53bf33..86f1019b 100644 --- a/benchmark/scripts/benchmark_simpo_loss.py +++ b/benchmark/scripts/benchmark_simpo_loss.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -23,17 +20,24 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_simpo_loss(input: SingleBenchmarkRunInput): +def setup_simpo_loss(input: SingleBenchmarkRunInput): """Create input tensors and SimPO loss from benchmark config.""" from test.chunked_loss.test_simpo_loss import LigerLMHeadSimPO from test.chunked_loss.test_simpo_loss import TorchLMHeadCPO cfg = input.extra_benchmark_config - H = cfg["hidden_size"] - V = cfg["vocab_size"] - dtype = cfg["dtype"] - B = input.x T = cfg["T"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + H = model_cfg.hidden_size + V = model_cfg.vocab_size + dtype = model_cfg.dtype + B = cfg["bsz"] + else: + B = input.x + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (B, T), dtype=torch.long, device=device) @@ -45,230 +49,54 @@ def _setup_simpo_loss(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for SimPOLoss") - fwd = lambda: loss_module(_input, target)[0] + fwd = lambda x: loss_module(x, target)[0] return _input, fwd -def bench_speed_simpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_simpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_simpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _setup_simpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_simpo_loss(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_simpo_loss( - SingleBenchmarkRunInput( - x=cfg["B"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "vocab_size": model_info["vocab_size"], - "dtype": model_info["dtype"], - "T": cfg["T"], - }, - ) - ) - - -def bench_speed_simpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_simpo_loss(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - fwd, - rep=100, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - grad_to_none=[_input], - rep=100, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - rep=100, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_simpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, fwd = _resolve_model_config_simpo_loss(input) - - def full(): - y = fwd() - y.backward() - - mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - if __name__ == "__main__": args = parse_benchmark_script_args() + T = 1024 + if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - T = 1024 - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - B = max(1, probe_bt // T) - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "vocab_size": model_cfg.vocab_size, - "dtype": model_cfg.dtype, - "T": T, - }, - ) - _, fwd = _setup_simpo_loss(probe_input) - return fwd() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype} - for cfg in sweep.model_configs - } - B = max(1, sweep.bt // T) - - common_configs = { - "kernel_name": "fused_linear_simpo_loss", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "B": B, "T": T}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_simpo_loss_model_config, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_simpo_loss_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="simpo_loss", + setup_fn=setup_simpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={"T": T}, + probe_dim="B", + probe_provider="huggingface", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - T = 1024 - probe_bt = 1024 - def _probe(): - B = probe_bt // T - probe_input = SingleBenchmarkRunInput( - x=B, - kernel_provider="huggingface", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "vocab_size": model.vocab_size, - "dtype": model.dtype, - "T": T, - }, - ) - _, fwd = _setup_simpo_loss(probe_input) - return fwd() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "fused_linear_simpo_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - {"hidden_size": model.hidden_size, "vocab_size": model.vocab_size, "dtype": model.dtype, "T": T} - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_simpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_simpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="simpo_loss", + probe_x=1, + model=model, + setup_fn=setup_simpo_loss, + model_keys=["hidden_size", "vocab_size", "dtype"], + extra_configs={"T": T}, + scale_dim="B", + x_label="batch size", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["liger", "huggingface"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_simpo_loss), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_simpo_loss), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_softmax.py b/benchmark/scripts/benchmark_softmax.py index f1259654..be4b4ee3 100644 --- a/benchmark/scripts/benchmark_softmax.py +++ b/benchmark/scripts/benchmark_softmax.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -24,15 +21,20 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def _setup_softmax(input: SingleBenchmarkRunInput): +def setup_softmax(input: SingleBenchmarkRunInput): """Create input tensors and softmax module from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg.get("hidden_size", input.x) - bt = cfg.get("bt", input.x) - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + bt = cfg["seq_len"] * cfg["bsz"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + bt = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] x = torch.randn(bt, hidden_size, dtype=dtype, device=device, requires_grad=True) - dy = torch.randn_like(x) if input.kernel_provider == "liger": softmax = LigerSoftmax().to(device).to(dtype) @@ -41,208 +43,50 @@ def _setup_softmax(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for softmax") - fwd_fn = lambda: softmax(x) - return x, dy, fwd_fn - - -def bench_speed_softmax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _setup_softmax(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(y_fwd, quantiles=QUANTILES, grad_to_none=[x], rep=500) - elif mode == "backward": - y = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[x], - rep=500, - ) - elif mode == "full": - - def full(): - y = y_fwd() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - quantiles=QUANTILES, - grad_to_none=[x], - rep=500, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_softmax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _setup_softmax(input) - - def full(): - y = fwd_fn() - y.backward(torch.ones_like(y), retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) - - -def _resolve_model_config_softmax(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_softmax( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "bt": cfg["bt"], - }, - ) - ) - - -def bench_speed_softmax_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _resolve_model_config_softmax(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, quantiles=QUANTILES, grad_to_none=[x], rep=500) - elif mode == "backward": - y = fwd_fn() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[x], - rep=500, - ) - elif mode == "full": - - def full(): - y = fwd_fn() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, grad_to_none=[x], rep=500) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_softmax_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _resolve_model_config_softmax(input) - - def full(): - y = fwd_fn() - y.backward(torch.ones_like(y), retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return x, softmax if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "bt": probe_bt, - }, - ) - _, _, fwd_fn = _setup_softmax(probe_input) - return fwd_fn() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "dtype": cfg.dtype} for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "softmax", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "bt": sweep.bt}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_softmax_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_softmax_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="softmax", + setup_fn=setup_softmax, + model_keys=["hidden_size", "dtype"], + probe_provider="torch", + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 2048 + probe_seq_len = 2048 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "bt": probe_bt, - }, - ) - _, _, fwd_fn = _setup_softmax(probe_input) - return fwd_fn() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "softmax", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.batch_size * config.seq_len))) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"hidden_size": model.hidden_size, "dtype": model.dtype}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_softmax, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_softmax, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="softmax", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_softmax, + model_keys=["hidden_size", "dtype"], + scale_dim="BT", + x_label="B*T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_softmax), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_softmax), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_sparsemax.py b/benchmark/scripts/benchmark_sparsemax.py index 94a56a6e..5a90b72e 100644 --- a/benchmark/scripts/benchmark_sparsemax.py +++ b/benchmark/scripts/benchmark_sparsemax.py @@ -1,18 +1,15 @@ -import math import os import sys import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -52,16 +49,22 @@ class TorchSparsemax(torch.nn.Module): return torch_sparsemax(x, dim=self.dim) -def _setup_sparsemax(input: SingleBenchmarkRunInput): +def setup_sparsemax(input: SingleBenchmarkRunInput): """Create input tensors and sparsemax module from benchmark config.""" cfg = input.extra_benchmark_config - hidden_size = cfg.get("hidden_size", input.x) - bt = cfg.get("bt", input.x) - dtype = cfg["dtype"] + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + bt = cfg["seq_len"] * cfg["bsz"] + hidden_size = model_cfg.hidden_size + dtype = model_cfg.dtype + else: + bt = input.x + hidden_size = cfg["hidden_size"] + dtype = cfg["dtype"] + dim = cfg.get("dim", -1) x = torch.randn(bt, hidden_size, dtype=dtype, device=device, requires_grad=True) - dy = torch.randn_like(x) if input.kernel_provider == "liger": sparsemax_module = LigerSparsemax(dim=dim).to(device) @@ -70,220 +73,56 @@ def _setup_sparsemax(input: SingleBenchmarkRunInput): else: raise ValueError(f"Invalid provider: {input.kernel_provider} for sparsemax") - fwd_fn = lambda: sparsemax_module(x) - return x, dy, fwd_fn - - -def bench_speed_sparsemax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _setup_sparsemax(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench( - y_fwd, - grad_to_none=[x], - rep=500, - quantiles=QUANTILES, - ) - elif mode == "backward": - y = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - grad_to_none=[x], - rep=500, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = y_fwd() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - full, - grad_to_none=[x], - rep=500, - quantiles=QUANTILES, - ) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_sparsemax(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _setup_sparsemax(input) - - def full(): - y = y_fwd() - y.backward(dy, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_sparsemax(input: SingleBenchmarkRunInput): - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_sparsemax( - SingleBenchmarkRunInput( - x=input.x, - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "hidden_size": model_info["hidden_size"], - "dtype": model_info["dtype"], - "bt": cfg["bt"], - "dim": cfg.get("dim", -1), - }, - ) - ) - - -def bench_speed_sparsemax_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, y_fwd = _resolve_model_config_sparsemax(input) - mode = input.kernel_operation_mode - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(y_fwd, grad_to_none=[x], rep=500, quantiles=QUANTILES) - elif mode == "backward": - y = y_fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(dy, retain_graph=True), - grad_to_none=[x], - rep=500, - quantiles=QUANTILES, - ) - elif mode == "full": - - def full(): - y = y_fwd() - y.backward(dy, retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[x], rep=500, quantiles=QUANTILES) - else: - raise ValueError(f"Unsupported mode: {mode}") - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_sparsemax_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - x, dy, fwd_fn = _resolve_model_config_sparsemax(input) - - def full(): - y = fwd_fn() - y.backward(dy, retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return x, sparsemax_module if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model_cfg.hidden_size, - "dtype": model_cfg.dtype, - "bt": probe_bt, - "dim": -1, - }, - ) - _, _, fwd_fn = _setup_sparsemax(probe_input) - return fwd_fn() - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - model_configs_info = { - cfg.name: {"hidden_size": cfg.hidden_size, "dtype": cfg.dtype} for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "sparsemax", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"model_configs": model_configs_info, "bt": sweep.bt, "dim": -1}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_sparsemax_model_config, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_sparsemax_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="sparsemax", + setup_fn=setup_sparsemax, + model_keys=["hidden_size", "dtype"], + probe_provider="torch", + extra_configs={ + "dim": -1, + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 2048 + probe_seq_len = 2048 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=0, - kernel_provider="torch", - extra_benchmark_config={ - "hidden_size": model.hidden_size, - "dtype": model.dtype, - "bt": probe_bt, - "dim": -1, - }, - ) - _, _, fwd_fn = _setup_sparsemax(probe_input) - return fwd_fn() - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "sparsemax", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.batch_size * config.seq_len))) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"hidden_size": model.hidden_size, "dtype": model.dtype, "dim": -1}], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_sparsemax, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_sparsemax, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="sparsemax", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_sparsemax, + model_keys=["hidden_size", "dtype"], + extra_configs={ + "dim": -1, + }, + scale_dim="BT", + x_label="B*T", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_sparsemax), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_sparsemax), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_swiglu.py b/benchmark/scripts/benchmark_swiglu.py index c54f96a8..65b3351d 100644 --- a/benchmark/scripts/benchmark_swiglu.py +++ b/benchmark/scripts/benchmark_swiglu.py @@ -68,7 +68,7 @@ if __name__ == "__main__": "bsz": 1, "hidden_act": "silu", }, - probe_dim="BT", + probe_dim="T", bt=args.bt, overwrite=args.overwrite, ) @@ -82,17 +82,21 @@ if __name__ == "__main__": model=model, setup_fn=setup_swiglu, model_keys=["hidden_size", "intermediate_size", "dtype"], - extra_configs={"hidden_act": "silu", "bsz": 1}, - scale_dim="BT", + extra_configs={ + "bsz": 1, + "hidden_act": "silu", + }, + scale_dim="T", + x_label="total tokens", probe_provider="huggingface", overwrite=args.overwrite, ) - common_configs["kernel_providers"] = ["liger", "huggingface"] + common_configs["kernel_providers"] = ["huggingface", "liger"] run_benchmarks( bench_test_fn=build_speed_bench_fn(setup_swiglu), - kernel_operation_modes=["full", "forward", "backward"], + kernel_operation_modes=["forward", "backward", "full"], metric_name="speed", metric_unit="ms", **common_configs, diff --git a/benchmark/scripts/benchmark_tvd.py b/benchmark/scripts/benchmark_tvd.py index 1cf425d4..27aae052 100644 --- a/benchmark/scripts/benchmark_tvd.py +++ b/benchmark/scripts/benchmark_tvd.py @@ -1,16 +1,12 @@ -import math - import torch -import triton from benchmark_model_configs import MODEL_REGISTRY -from benchmark_model_configs import compute_model_config_sweep_config -from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import build_model_config_sweep +from benchmark_model_configs import build_token_length_sweep from benchmark_model_configs import get_benchmark_model_config -from utils import QUANTILES from utils import SingleBenchmarkRunInput -from utils import SingleBenchmarkRunOutput -from utils import _test_memory +from utils import build_memory_bench_fn +from utils import build_speed_bench_fn from utils import parse_benchmark_script_args from utils import run_benchmarks @@ -39,12 +35,18 @@ class TorchTVDLoss(torch.nn.Module): raise ValueError("Invalid reduction type.") -def _setup_tvd(input: SingleBenchmarkRunInput): +def setup_tvd(input: SingleBenchmarkRunInput): """Create input tensors and TVD loss from benchmark config.""" cfg = input.extra_benchmark_config - V = cfg["vocab_size"] - BT = input.x - reduction = "batchmean" + if isinstance(input.x, str): + model_cfg = MODEL_REGISTRY[input.x] + BT = cfg["seq_len"] * cfg["bsz"] + V = model_cfg.vocab_size + else: + BT = input.x + V = cfg["vocab_size"] + + reduction = cfg.get("reduction", "batchmean") _input = torch.randn(BT, V, requires_grad=True, device=device).softmax(dim=-1) target = torch.randn(BT, V, device=device).softmax(dim=-1) @@ -55,220 +57,55 @@ def _setup_tvd(input: SingleBenchmarkRunInput): loss_fn = TorchTVDLoss(reduction=reduction) else: raise ValueError(f"Invalid provider: {input.kernel_provider} for TVD") - return _input, target, loss_fn - - -def bench_speed_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_tvd(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput( - y_20=ms_20, - y_50=ms_50, - y_80=ms_80, - ) - - -def bench_memory_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _setup_tvd(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) - - -def _resolve_model_config_tvd(input: SingleBenchmarkRunInput): - """Resolve model-config-sweep input into standard setup args.""" - cfg = input.extra_benchmark_config - model_info = cfg["model_configs"][input.x] - return _setup_tvd( - SingleBenchmarkRunInput( - x=cfg["BT"], - kernel_provider=input.kernel_provider, - extra_benchmark_config={ - "vocab_size": model_info["vocab_size"], - }, - ) - ) - - -def bench_speed_tvd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_tvd(input) - mode = input.kernel_operation_mode - - def fwd(): - return loss_fn(_input, target) - - if mode == "forward": - ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif mode == "backward": - y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( - lambda: y.backward(retain_graph=True), - quantiles=QUANTILES, - grad_to_none=[_input], - rep=100, - ) - elif mode == "full": - - def full(): - y = fwd() - y.backward(retain_graph=True) - - ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) - else: - raise ValueError(f"Unsupported mode: {mode}") - - return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) - - -def bench_memory_tvd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - _input, target, loss_fn = _resolve_model_config_tvd(input) - - def full(): - y = loss_fn(_input, target) - y.backward(retain_graph=True) - - mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) - - return SingleBenchmarkRunOutput( - y_20=mem_20, - y_50=mem_50, - y_80=mem_80, - ) + return _input, lambda x: loss_fn(x, target) if __name__ == "__main__": args = parse_benchmark_script_args() if args.sweep_mode == "model_config": - all_model_configs = list(MODEL_REGISTRY.values()) - - def _probe_factory(model_cfg, probe_bt): - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model_cfg.vocab_size, - }, - ) - _input, target, loss_fn = _setup_tvd(probe_input) - return loss_fn(_input, target) - - return _probe - - sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) - - model_configs_info = { - cfg.name: { - "vocab_size": cfg.vocab_size, - } - for cfg in sweep.model_configs - } - - common_configs = { - "kernel_name": "tvd", - "x_name": "model_config", - "x_label": "model configuration", - "x_values": [cfg.name for cfg in sweep.model_configs], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "model_configs": model_configs_info, - "BT": sweep.bt, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_tvd_model_config, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_tvd_model_config, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_model_config_sweep( + kernel_name="tvd", + setup_fn=setup_tvd, + model_keys=["vocab_size"], + probe_provider="torch", + extra_configs={ + "reduction": "batchmean", + }, + probe_dim="BT", + bt=args.bt, + overwrite=args.overwrite, ) else: model = get_benchmark_model_config(args.model) - probe_bt = 1024 + probe_seq_len = 1024 - def _probe(): - probe_input = SingleBenchmarkRunInput( - x=probe_bt, - kernel_provider="torch", - extra_benchmark_config={ - "vocab_size": model.vocab_size, - }, - ) - _input, target, loss_fn = _setup_tvd(probe_input) - return loss_fn(_input, target) - - config = compute_seq_len_sweep_config(model, probe_fn=_probe, probe_seq_len=probe_bt) - - common_configs = { - "kernel_name": "tvd", - "x_name": "BT", - "x_label": "B * T", - "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "vocab_size": model.vocab_size, - } - ], - "overwrite": args.overwrite, - } - - run_benchmarks( - bench_test_fn=bench_speed_tvd, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_tvd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, + common_configs = build_token_length_sweep( + kernel_name="tvd", + probe_x=probe_seq_len, + model=model, + setup_fn=setup_tvd, + model_keys=["vocab_size"], + extra_configs={"reduction": "batchmean"}, + scale_dim="BT", + x_label="total tokens", + probe_provider="torch", + overwrite=args.overwrite, ) + + common_configs["kernel_providers"] = ["torch", "liger"] + + run_benchmarks( + bench_test_fn=build_speed_bench_fn(setup_tvd), + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=build_memory_bench_fn(setup_tvd), + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/utils.py b/benchmark/scripts/utils.py index f3774a13..16e6000d 100644 --- a/benchmark/scripts/utils.py +++ b/benchmark/scripts/utils.py @@ -157,6 +157,17 @@ def run_speed_benchmark( rep=rep, quantiles=QUANTILES, ) + elif mode == "no-grad-forward": + + def no_grad_forward(): + with torch.no_grad(): + fwd_fn() + + ms_50, ms_20, ms_80 = triton.testing.do_bench( + no_grad_forward, + rep=rep, + quantiles=QUANTILES, + ) else: raise ValueError(f"Unsupported mode: {mode}. Use 'forward', 'backward', or 'full'.") return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80)