From 1796cf0d4ff3fbe78ec0ad0d39b4da05e30bb564 Mon Sep 17 00:00:00 2001 From: lowdy1 Date: Sat, 18 Apr 2026 23:25:56 +0800 Subject: [PATCH] [Benchmark] Benchmark refactor for loss ops (#1191) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary * Add `model_config` sweep support to loss ops in Liger-Kernel benchmarks, enabling evaluation across different model architectures at a fixed sequence length * Introduce a pre-sweep probe step to determine safe configurations and prevent OOM during benchmarking Loss kernels: benchmark_cpo_loss.py benchmark_distill_cosine_loss.py benchmark_distill_jsd_loss.py benchmark_dpo_loss.py benchmark_grpo_loss.py benchmark_kto_loss.py benchmark_orpo_loss.py benchmark_simpo_loss.py `BT` denotes the product of batch size and sequence length (`batch_size` × `sequence_length`). Sweep mode `model_config`: `python benchmark_grpo_loss.py --sweep-mode model_config --bt 2048` Sweep mode `BT`: `python benchmark_grpo_loss.py --model llama_3_8b` Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <757486878@qq.com> --- benchmark/scripts/benchmark_cpo_loss.py | 305 ++++++++++---- .../scripts/benchmark_distill_cosine_loss.py | 393 +++++++++++------ .../scripts/benchmark_distill_jsd_loss.py | 391 +++++++++++------ benchmark/scripts/benchmark_dpo_loss.py | 346 ++++++++++----- benchmark/scripts/benchmark_grpo_loss.py | 394 +++++++++++------- benchmark/scripts/benchmark_kto_loss.py | 379 +++++++++-------- benchmark/scripts/benchmark_orpo_loss.py | 280 +++++++++---- benchmark/scripts/benchmark_simpo_loss.py | 306 +++++++++----- 8 files changed, 1838 insertions(+), 956 deletions(-) diff --git a/benchmark/scripts/benchmark_cpo_loss.py b/benchmark/scripts/benchmark_cpo_loss.py index 8b10d518..07501275 100644 --- a/benchmark/scripts/benchmark_cpo_loss.py +++ b/benchmark/scripts/benchmark_cpo_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -18,41 +24,72 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -############################################################################# -# Test the memory consumption of the linear fused cross entropy loss -############################################################################# - - -def bench_memory_fused_linear_cpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +def _setup_cpo_loss(input: SingleBenchmarkRunInput): + """Create input tensors and CPO loss from benchmark config.""" from test.chunked_loss.test_cpo_loss import LigerLMHeadCPO from test.chunked_loss.test_cpo_loss import TorchLMHeadCPO + cfg = input.extra_benchmark_config + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider - - # Instantiate once and retrieve the first output only - torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0] - liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0] + T = cfg["T"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (B, T), dtype=torch.long, device=device) - def fwd(): - if provider == "liger": - return liger_fwd(_input, target) - elif provider == "huggingface": - return torch_fwd(_input, target) + if input.kernel_provider == "liger": + loss_module = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device) + elif input.kernel_provider == "huggingface": + loss_module = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for CPOLoss") + + fwd_fn = lambda: loss_module(_input, target)[0] + return _input, fwd_fn + + +def bench_speed_cpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd_fn = _setup_cpo_loss(input) + mode = input.kernel_operation_mode + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench( + fwd_fn, + rep=100, + quantiles=QUANTILES, + ) + elif mode == "backward": + y = fwd_fn() + 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_fn() + 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_cpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd_fn = _setup_cpo_loss(input) def full(): - y = fwd() + y = fwd_fn() y.backward() mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) @@ -63,40 +100,28 @@ def bench_memory_fused_linear_cpo_loss( ) -# ############################################################################# -# # Test the speed of the fused linear cross entropy loss -# ############################################################################# +def _resolve_model_config_cpo_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_cpo_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_fused_linear_cpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - from test.chunked_loss.test_cpo_loss import LigerLMHeadCPO - from test.chunked_loss.test_cpo_loss import TorchLMHeadCPO - - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider +def bench_speed_cpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd = _resolve_model_config_cpo_loss(input) mode = input.kernel_operation_mode - # Instantiate once and retrieve the first output only - torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0] - liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0] - - _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) - target = torch.randint(V, (B, T), dtype=torch.long, device=device) - - def fwd(): - if provider == "liger": - return liger_fwd(_input, target) - elif provider == "huggingface": - return torch_fwd(_input, target) - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -105,7 +130,6 @@ def bench_speed_fused_linear_cpo_loss( ) 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], @@ -123,6 +147,9 @@ def bench_speed_fused_linear_cpo_loss( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -130,38 +157,142 @@ def bench_speed_fused_linear_cpo_loss( ) +def bench_memory_cpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd_fn = _resolve_model_config_cpo_loss(input) + + def full(): + y = fwd_fn() + 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() - common_configs = { - "kernel_name": "fused_linear_cpo_loss", - "x_name": "B", - "x_label": "B", - "x_values": [2**i for i in range(1, 5)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "T": 1024, - "H": 4096, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - } - ], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + T = 1024 - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_cpo_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_cpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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_fn = _setup_cpo_loss(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, + "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_cpo_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_cpo_loss_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_cpo_loss_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + 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_fn = _setup_cpo_loss(probe_input) + return fwd_fn() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + common_configs = { + "kernel_name": "fused_linear_cpo_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_cpo_loss, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_cpo_loss, + 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 5cf12b49..12f9faff 100644 --- a/benchmark/scripts/benchmark_distill_cosine_loss.py +++ b/benchmark/scripts/benchmark_distill_cosine_loss.py @@ -1,3 +1,4 @@ +import math import os import sys @@ -5,6 +6,11 @@ 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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -84,106 +90,54 @@ class LigerCosineSimilarityLoss(nn.Module): ) -def bench_memory_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: +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"] + 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 - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"] - weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider - - torch_cosine_loss = TorchCosineSimilarityLoss( - H=H, - V=V, - dtype=dtype, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - bias=bias, - ).to(device) - liger_cosine_loss = LigerCosineSimilarityLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) - student_input1 = _tensor.detach().clone().requires_grad_(True) - student_input2 = _tensor.detach().clone().requires_grad_(True) - + student_input = _tensor.detach().clone().requires_grad_(True) teacher_input = torch.rand(BT, H, device=device, dtype=dtype) - target = torch.randint(0, V, (BT,), device=device, dtype=torch.long) - def fwd(): - if provider == "liger": - return liger_cosine_loss(student_input1, teacher_input, target) - elif provider == "torch": - return torch_cosine_loss(student_input2, teacher_input, target) - - 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 input.kernel_provider == "liger": + loss_module = LigerCosineSimilarityLoss( + H=H, + V=V, + dtype=dtype, + ignore_index=ignore_index, + bias=bias, + weight_hard_loss=weight_hard_loss, + weight_soft_loss=weight_soft_loss, + ).to(device) + elif input.kernel_provider == "torch": + loss_module = TorchCosineSimilarityLoss( + H=H, + V=V, + dtype=dtype, + ignore_index=ignore_index, + bias=bias, + weight_hard_loss=weight_hard_loss, + weight_soft_loss=weight_soft_loss, + ).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillCosineLoss") + return student_input, teacher_input, target, loss_module -def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - BT = input.x - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"] - weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider +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 - torch_cosine_loss = TorchCosineSimilarityLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) - - liger_cosine_loss = LigerCosineSimilarityLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) - - _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) - student_input1 = _tensor.detach().clone().requires_grad_(True) - student_input2 = _tensor.detach().clone().requires_grad_(True) - - teacher_input = torch.rand(BT, H, device=device, dtype=dtype) - - target = torch.randint(0, V, (BT,), device=device, dtype=torch.long) - def fwd(): - if provider == "liger": - return liger_cosine_loss(student_input1, teacher_input, target) - elif provider == "torch": - return torch_cosine_loss(student_input2, teacher_input, target) + return loss_module(student_input, teacher_input, target) if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -191,18 +145,11 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single rep=100, quantiles=QUANTILES, ) - elif mode == "backward": - y = fwd() - 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_input1, student_input2], + grad_to_none=[student_input], rep=100, quantiles=QUANTILES, ) @@ -217,6 +164,8 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, @@ -225,42 +174,226 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single ) +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, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs = { - "kernel_name": "distill_cosine_loss", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, 14)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "H": 4096, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, + 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, } - ], - "overwrite": args.overwrite, - } + for cfg in sweep.model_configs + } - run_benchmarks( - bench_test_fn=bench_speed_cosine_similarity_loss, - kernel_operation_modes=["forward", "full"], - metric_name="speed", - metric_unit="ms", - **common_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_memory_cosine_similarity_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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) + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) diff --git a/benchmark/scripts/benchmark_distill_jsd_loss.py b/benchmark/scripts/benchmark_distill_jsd_loss.py index 324418e1..7602a2ee 100644 --- a/benchmark/scripts/benchmark_distill_jsd_loss.py +++ b/benchmark/scripts/benchmark_distill_jsd_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -12,7 +18,6 @@ from utils import parse_benchmark_script_args from utils import run_benchmarks from liger_kernel.chunked_loss.jsd_loss import LigerFusedLinearJSDFunction -from liger_kernel.utils import get_total_gpu_memory from liger_kernel.utils import infer_device device = infer_device() @@ -89,106 +94,54 @@ class LigerJSDLoss(torch.nn.Module): ) -def bench_memory_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: +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"] + 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 - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"] - weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider - - torch_jsd_loss = TorchJSDLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) - liger_jsd_loss = LigerJSDLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) - student_input1 = _tensor.detach().clone().requires_grad_(True) - student_input2 = _tensor.detach().clone().requires_grad_(True) - + student_input = _tensor.detach().clone().requires_grad_(True) teacher_input = torch.rand(BT, H, device=device, dtype=dtype) - target = torch.randint(0, V, (BT,), device=device, dtype=torch.long) - def fwd(): - if provider == "liger": - return liger_jsd_loss(student_input1, teacher_input, target) - elif provider == "torch": - return torch_jsd_loss(student_input2, teacher_input, target) - - 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 input.kernel_provider == "liger": + loss_module = LigerJSDLoss( + H=H, + V=V, + dtype=dtype, + ignore_index=ignore_index, + bias=bias, + weight_hard_loss=weight_hard_loss, + weight_soft_loss=weight_soft_loss, + ).to(device) + elif input.kernel_provider == "torch": + loss_module = TorchJSDLoss( + H=H, + V=V, + dtype=dtype, + ignore_index=ignore_index, + bias=bias, + weight_hard_loss=weight_hard_loss, + weight_soft_loss=weight_soft_loss, + ).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillJSDLoss") + return student_input, teacher_input, target, loss_module -def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - BT = input.x - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"] - weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider +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 - torch_jsd_loss = TorchJSDLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) - liger_jsd_loss = LigerJSDLoss( - H=H, - V=V, - dtype=dtype, - ignore_index=ignore_index, - bias=bias, - weight_hard_loss=weight_hard_loss, - weight_soft_loss=weight_soft_loss, - ).to(device) - - _tensor = torch.rand(BT, H // 2, device=device, dtype=dtype) - student_input1 = _tensor.detach().clone().requires_grad_(True) - student_input2 = _tensor.detach().clone().requires_grad_(True) - - teacher_input = torch.rand(BT, H, device=device, dtype=dtype) - - target = torch.randint(0, V, (BT,), device=device, dtype=torch.long) - def fwd(): - if provider == "liger": - return liger_jsd_loss(student_input1, teacher_input, target) - elif provider == "torch": - return torch_jsd_loss(student_input2, teacher_input, target) + return loss_module(student_input, teacher_input, target) if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -200,7 +153,81 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu y = fwd() ms_50, ms_20, ms_80 = triton.testing.do_bench( lambda: y.backward(retain_graph=True), - grad_to_none=[student_input1, student_input2], + 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, ) @@ -215,6 +242,8 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, @@ -223,50 +252,148 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu ) +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, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() - gpu_memory_gbs = get_total_gpu_memory() - # We know that the full test will require 69GBs for vocab size 2^13 and 39GBs for vocab size 2^12 on torch - if gpu_memory_gbs >= 69: - x_max = 13 - elif gpu_memory_gbs >= 39: - x_max = 12 - else: - x_max = 11 - common_configs = { - "kernel_name": "distill_jsd_loss", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(10, x_max + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "H": 4096, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - "bias": False, - "weight_hard_loss": 0.5, - "weight_soft_loss": 0.5, - "ignore_index": -100, + 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, } - ], - "overwrite": args.overwrite, - } + for cfg in sweep.model_configs + } - run_benchmarks( - bench_test_fn=bench_speed_jsd_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_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_memory_jsd_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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) + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) diff --git a/benchmark/scripts/benchmark_dpo_loss.py b/benchmark/scripts/benchmark_dpo_loss.py index 228a228d..da241289 100644 --- a/benchmark/scripts/benchmark_dpo_loss.py +++ b/benchmark/scripts/benchmark_dpo_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -18,25 +24,20 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: +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"] + bias = cfg["bias"] + beta = cfg["beta"] + ignore_index = cfg["ignore_index"] B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - beta = input.extra_benchmark_config["beta"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider - - # Instantiate once and retrieve the first output only - torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) - liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) - torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0] - liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0] + T = cfg["T"] # Input shape: [B, T, H] _input = torch.randn(B, T, H, device=device, dtype=dtype) @@ -49,62 +50,21 @@ def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunO indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] target.view(-1)[indices_to_assign] = ignore_index - def fwd(): - if provider == "liger": - return liger_fwd(_input, ref_input, target) - elif provider == "huggingface": - return torch_fwd(_input, ref_input, target) + 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": + 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") - 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, - ) + fwd = lambda: loss_module(_input, ref_input, target)[0] + return _input, fwd def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - from test.chunked_loss.test_dpo_loss import LigerLMHeadDPO - from test.chunked_loss.test_dpo_loss import TorchLMHeadDPO - - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - beta = input.extra_benchmark_config["beta"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider + _input, fwd = _setup_dpo_loss(input) mode = input.kernel_operation_mode - # Instantiate once and retrieve the first output only - torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) - liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device) - torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0] - liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0] - - # Input shape: [B, T, H] - _input = torch.randn(B, T, H, device=device, dtype=dtype) - ref_input = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=False) - # Target shape: [B, T] - target = torch.randint(V, (B, T), device=device, dtype=torch.long) - - # Add ignore_index tokens - num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item() - indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] - target.view(-1)[indices_to_assign] = ignore_index - - def fwd(): - if provider == "liger": - return liger_fwd(_input, ref_input, target) - elif provider == "huggingface": - return torch_fwd(_input, ref_input, target) - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -130,6 +90,8 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, @@ -138,42 +100,230 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu ) +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() - common_configs = { - "kernel_name": "dpo_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, 6)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "T": 512, - "H": 1024, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - "bias": True, - "beta": 0.1, - "ignore_index": 42, + 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, } - ], - "overwrite": args.overwrite, - } + for cfg in sweep.model_configs + } - run_benchmarks( - bench_test_fn=bench_speed_dpo_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) + T = 512 + B = max(1, sweep.bt // T) - run_benchmarks( - bench_test_fn=bench_memory_dpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) diff --git a/benchmark/scripts/benchmark_grpo_loss.py b/benchmark/scripts/benchmark_grpo_loss.py index 497d8692..676a9a85 100644 --- a/benchmark/scripts/benchmark_grpo_loss.py +++ b/benchmark/scripts/benchmark_grpo_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -18,113 +24,44 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -############################################################################# -# Test the memory consumption of the linear fused GRPO loss -############################################################################# - - -def bench_memory_fused_linear_grpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +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 = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - importance_sampling_level = input.extra_benchmark_config["importance_sampling_level"] - provider = input.kernel_provider + T = cfg["T"] - # Instantiate once and retrieve the first output only - torch_lm_head_grpo = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( - device - ) - liger_lm_head_grpo = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( - device - ) - - # Create inputs _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) attention_mask = torch.ones(B, T, device=device) advantages = torch.randn(B, dtype=dtype, device=device) ref_input = torch.randn(B, T, H, dtype=dtype, device=device) - torch_fwd = lambda: torch_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[ - 0 - ] - liger_fwd = lambda: liger_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[ - 0 - ] + if input.kernel_provider == "liger": + loss_module = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( + device + ) + elif input.kernel_provider == "torch": + loss_module = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( + device + ) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for GRPOLoss") - def fwd(): - if provider == "liger": - return liger_fwd() - elif provider == "torch": - return torch_fwd() - - 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, - ) + fwd = lambda: loss_module(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[0] + return _input, fwd -############################################################################# -# Test the speed of the fused linear GRPO loss -############################################################################# - - -def bench_speed_fused_linear_grpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - from test.chunked_loss.test_grpo_loss import LigerLMHeadGRPO - from test.chunked_loss.test_grpo_loss import TorchLMHeadGRPO - - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - importance_sampling_level = input.extra_benchmark_config["importance_sampling_level"] - provider = input.kernel_provider +def bench_speed_grpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd = _setup_grpo_loss(input) mode = input.kernel_operation_mode - # Instantiate once and retrieve the first output only - torch_lm_head_grpo = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( - device - ) - liger_lm_head_grpo = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to( - device - ) - - # Create inputs - _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) - attention_mask = torch.ones(B, T, device=device) - advantages = torch.randn(B, dtype=dtype, device=device) - ref_input = torch.randn(B, T, H, dtype=dtype, device=device) - - torch_fwd = lambda: torch_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[ - 0 - ] - liger_fwd = lambda: liger_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[ - 0 - ] - - def fwd(): - if provider == "liger": - return liger_fwd() - elif provider == "torch": - return torch_fwd() - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -133,7 +70,6 @@ def bench_speed_fused_linear_grpo_loss( ) 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], @@ -151,6 +87,8 @@ def bench_speed_fused_linear_grpo_loss( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -158,77 +96,215 @@ def bench_speed_fused_linear_grpo_loss( ) +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}" + + 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, + ) + 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() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() # Benchmark token-level importance sampling (original GRPO) - token_configs = { - "kernel_name": "fused_linear_grpo_loss_token", - "x_name": "B", - "x_label": "B", - "x_values": [2**i for i in range(1, 5)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "T": 1024, - "H": 4096, - "V": 128256, - "importance_sampling_level": "token", - "dtype": torch.bfloat16, - } - ], - "overwrite": args.overwrite, - } + print("Benchmarking GRPO (token-level importance sampling)...") + _run_grpo_benchmarks(args, importance_sampling_level="token", kernel_name_suffix="token") # Benchmark sequence-level importance sampling (GSPO) - sequence_configs = { - "kernel_name": "fused_linear_grpo_loss_sequence", - "x_name": "B", - "x_label": "B", - "x_values": [2**i for i in range(1, 5)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [ - { - "T": 1024, - "H": 4096, - "V": 128256, - "importance_sampling_level": "sequence", - "dtype": torch.bfloat16, - } - ], - "overwrite": args.overwrite, - } - - # Run benchmarks for token-level (GRPO) - print("Benchmarking GRPO (token-level importance sampling)...") - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_grpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **token_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_grpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **token_configs, - ) - - # Run benchmarks for sequence-level (GSPO) print("Benchmarking GSPO (sequence-level importance sampling)...") - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_grpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **sequence_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_grpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **sequence_configs, - ) + _run_grpo_benchmarks(args, importance_sampling_level="sequence", kernel_name_suffix="sequence") diff --git a/benchmark/scripts/benchmark_kto_loss.py b/benchmark/scripts/benchmark_kto_loss.py index bbde1d5c..b0ca76d0 100644 --- a/benchmark/scripts/benchmark_kto_loss.py +++ b/benchmark/scripts/benchmark_kto_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -88,158 +94,52 @@ class LigerLMHeadKTO(torch.nn.Module): ) -def bench_memory_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: +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"] + bias = cfg["bias"] + beta = cfg["beta"] + ignore_index = cfg["ignore_index"] B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - beta = input.extra_benchmark_config["beta"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider - - torch_kto_loss = TorchLMHeadKTO( - H=H, - V=V, - dtype=dtype, - use_bias=bias, - use_ref_bias=bias, - ignore_index=ignore_index, - beta=beta, - ).to(device) - - liger_kto_loss = LigerLMHeadKTO( - H=H, - V=V, - dtype=dtype, - use_bias=bias, - use_ref_bias=bias, - ignore_index=ignore_index, - beta=beta, - ).to(device) + T = cfg["T"] # Input shape: [B, T, H] _input = torch.randn(B, T, H, device=device, dtype=dtype) - - # Target shape: [B, T] + ref_input = torch.randn(B, T, H, device=device, dtype=dtype) target = torch.randint(V, (B, T), dtype=torch.long, device=device) - # Preference labels shape: [B] # Create binary preference labels (0 or 1) for each sequence in the batch # Used to indicate preferred sequences (1) vs non-preferred sequences (0) preference_labels = torch.randint(2, (B,), dtype=torch.bool, device=device) - # Precomputed KL divergence between policy and reference distributions kl = torch.randn(1, device=device, dtype=dtype) - # Add ignore_index tokens to simulate padding num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item() indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] target.view(-1)[indices_to_assign] = ignore_index - # Add ref_x with the same shape as _input - ref_input = torch.randn(B, T, H, device=device, dtype=dtype) + if input.kernel_provider == "liger": + loss_module = LigerLMHeadKTO( + H=H, V=V, dtype=dtype, use_bias=bias, use_ref_bias=bias, ignore_index=ignore_index, beta=beta + ).to(device) + elif input.kernel_provider == "huggingface": + loss_module = TorchLMHeadKTO( + H=H, V=V, dtype=dtype, use_bias=bias, use_ref_bias=bias, ignore_index=ignore_index, beta=beta + ).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for KTOLoss") - def fwd(): - if provider == "liger": - return liger_kto_loss( - x=_input, - ref_x=ref_input, - y=target, - preference_labels=preference_labels, - kl=kl, - )[0] - elif provider == "huggingface": - return torch_kto_loss( - x=_input, - ref_x=ref_input, - y=target, - preference_labels=preference_labels, - kl=kl, - )[0] - - 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, - ) + fwd = lambda: 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: - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - bias = input.extra_benchmark_config["bias"] - beta = input.extra_benchmark_config["beta"] - ignore_index = input.extra_benchmark_config["ignore_index"] - provider = input.kernel_provider + _input, fwd = _setup_kto_loss(input) mode = input.kernel_operation_mode - torch_kto_loss = TorchLMHeadKTO( - H=H, - V=V, - dtype=dtype, - beta=beta, - ignore_index=ignore_index, - use_bias=bias, - ).to(device) - liger_kto_loss = LigerLMHeadKTO( - H=H, - V=V, - dtype=dtype, - beta=beta, - ignore_index=ignore_index, - use_bias=bias, - ).to(device) - - # Input shape: [B, T, H] - _input = torch.randn(B, T, H, device=device, dtype=dtype) - - # Target shape: [B, T] - target = torch.randint(V, (B, T), device=device, dtype=torch.long) - - # Preference labels shape: [B] - # Create binary preference labels (0 or 1) for each sequence in the batch - # Used to indicate preferred sequences (1) vs non-preferred sequences (0) - preference_labels = torch.randint(2, (B,), dtype=torch.bool, device=device) - - # Precomputed KL divergence between policy and reference distributions - kl = torch.randn(1, device=device, dtype=dtype) - - # Add ignore_index tokens - num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item() - indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign] - target.view(-1)[indices_to_assign] = ignore_index - - # Add ref_x with the same shape as _input - ref_input = torch.randn(B, T, H, device=device, dtype=dtype) - - def fwd(): - if provider == "liger": - return liger_kto_loss( - x=_input, - ref_x=ref_input, - y=target, - preference_labels=preference_labels, - kl=kl, - )[0] - elif provider == "huggingface": - return torch_kto_loss( - x=_input, - ref_x=ref_input, - y=target, - preference_labels=preference_labels, - kl=kl, - )[0] - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -265,7 +165,8 @@ def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu rep=100, quantiles=QUANTILES, ) - + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -273,42 +174,190 @@ def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu ) +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() - common_configs = { - "kernel_name": "kto_loss", - "x_name": "B", - "x_label": "Batch Size (B)", - "x_values": [2**i for i in range(1, 6)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "T": 512, - "H": 1024, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - "bias": True, - "beta": 0.1, - "ignore_index": 42, - } - ], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + T = 512 - run_benchmarks( - bench_test_fn=bench_speed_kto_loss, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) + 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() - run_benchmarks( - bench_test_fn=bench_memory_kto_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) diff --git a/benchmark/scripts/benchmark_orpo_loss.py b/benchmark/scripts/benchmark_orpo_loss.py index 30b308c4..663fa875 100644 --- a/benchmark/scripts/benchmark_orpo_loss.py +++ b/benchmark/scripts/benchmark_orpo_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -18,39 +24,66 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -############################################################################# -# Test the memory consumption of the linear fused cross entropy loss -############################################################################# - - -def bench_memory_fused_linear_orpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +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 = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider - - # Instantiate once and retrieve the first output only - torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0] - liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0] + T = cfg["T"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (B, T), dtype=torch.long, device=device) nll_target = torch.randint(V, (B, T), dtype=torch.long, device=device) - def fwd(): - if provider == "liger": - return liger_fwd(_input, target, nll_target) - elif provider == "huggingface": - return torch_fwd(_input, target, nll_target) + if input.kernel_provider == "liger": + loss_module = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device) + elif input.kernel_provider == "huggingface": + loss_module = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for ORPOLoss") + + fwd = lambda: loss_module(_input, 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() @@ -64,41 +97,27 @@ def bench_memory_fused_linear_orpo_loss( ) -# ############################################################################# -# # Test the speed of the fused linear cross entropy loss -# ############################################################################# +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_fused_linear_orpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - from test.chunked_loss.test_orpo_loss import LigerLMHeadORPO - from test.chunked_loss.test_orpo_loss import TorchLMHeadORPO - - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider +def bench_speed_orpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd = _resolve_model_config_orpo_loss(input) mode = input.kernel_operation_mode - # Instantiate once and retrieve the first output only - torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0] - liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0] - - _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) - target = torch.randint(V, (B, T), dtype=torch.long, device=device) - nll_target = torch.randint(V, (B, T), dtype=torch.long, device=device) - - def fwd(): - if provider == "liger": - return liger_fwd(_input, target, nll_target) - elif provider == "huggingface": - return torch_fwd(_input, target, nll_target) - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -107,7 +126,6 @@ def bench_speed_fused_linear_orpo_loss( ) 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], @@ -125,6 +143,8 @@ def bench_speed_fused_linear_orpo_loss( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -132,38 +152,124 @@ def bench_speed_fused_linear_orpo_loss( ) +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() - common_configs = { - "kernel_name": "fused_linear_orpo_loss", - "x_name": "B", - "x_label": "B", - "x_values": [2**i for i in range(1, 5)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "T": 1024, - "H": 4096, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - } - ], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + T = 1024 - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_orpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_orpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + ) diff --git a/benchmark/scripts/benchmark_simpo_loss.py b/benchmark/scripts/benchmark_simpo_loss.py index 148b8e3e..f9b90f71 100644 --- a/benchmark/scripts/benchmark_simpo_loss.py +++ b/benchmark/scripts/benchmark_simpo_loss.py @@ -1,9 +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 estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from utils import QUANTILES from utils import SingleBenchmarkRunInput from utils import SingleBenchmarkRunOutput @@ -18,85 +24,36 @@ device = infer_device() sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -############################################################################# -# Test the memory consumption of the linear fused cross entropy loss -############################################################################# - - -def bench_memory_fused_linear_simpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +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 = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider - - # Instantiate once and retrieve the first output only - torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0] - liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0] + T = cfg["T"] _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (B, T), dtype=torch.long, device=device) - def fwd(): - if provider == "liger": - return liger_fwd(_input, target) - elif provider == "huggingface": - return torch_fwd(_input, target) + if input.kernel_provider == "liger": + loss_module = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device) + elif input.kernel_provider == "huggingface": + loss_module = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for SimPOLoss") - 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, - ) + fwd = lambda: loss_module(_input, target)[0] + return _input, fwd -# ############################################################################# -# # Test the speed of the fused linear cross entropy loss -# ############################################################################# - - -def bench_speed_fused_linear_simpo_loss( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - from test.chunked_loss.test_simpo_loss import LigerLMHeadSimPO - from test.chunked_loss.test_simpo_loss import TorchLMHeadCPO - - B = input.x - T = input.extra_benchmark_config["T"] - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider +def bench_speed_simpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, fwd = _setup_simpo_loss(input) mode = input.kernel_operation_mode - # Instantiate once and retrieve the first output only - torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device) - liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device) - torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0] - liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0] - - _input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device) - target = torch.randint(V, (B, T), dtype=torch.long, device=device) - - def fwd(): - if provider == "liger": - return liger_fwd(_input, target) - elif provider == "huggingface": - return torch_fwd(_input, target) - if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( fwd, @@ -105,7 +62,6 @@ def bench_speed_fused_linear_simpo_loss( ) 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], @@ -123,6 +79,8 @@ def bench_speed_fused_linear_simpo_loss( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -130,38 +88,190 @@ def bench_speed_fused_linear_simpo_loss( ) +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() - common_configs = { - "kernel_name": "fused_linear_simpo_loss", - "x_name": "B", - "x_label": "B", - "x_values": [2**i for i in range(1, 5)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "T": 1024, - "H": 4096, - "V": 128256, - "mode": "forward", - "dtype": torch.bfloat16, - } - ], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + T = 1024 - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_simpo_loss, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_simpo_loss, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + 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, + ) + 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() + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_bt + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + 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, + )