diff --git a/benchmark/scripts/benchmark_cross_entropy.py b/benchmark/scripts/benchmark_cross_entropy.py index cdd61814..7c305e40 100644 --- a/benchmark/scripts/benchmark_cross_entropy.py +++ b/benchmark/scripts/benchmark_cross_entropy.py @@ -1,6 +1,13 @@ +import math + import torch import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import estimate_kernel_peak_memory +from benchmark_model_configs import get_benchmark_model_config from torch.nn import CrossEntropyLoss from utils import QUANTILES from utils import SingleBenchmarkRunInput @@ -15,28 +22,64 @@ from liger_kernel.utils import infer_device device = infer_device() -def bench_memory_cross_entropy( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - torch_ce = CrossEntropyLoss() - liger_ce = LigerCrossEntropyLoss() +def _setup_cross_entropy(input: SingleBenchmarkRunInput): + """Create input tensor, target, and CE loss from benchmark config.""" + cfg = input.extra_benchmark_config + V = cfg["vocab_size"] + BT = input.x + _input = torch.randn(BT, V, requires_grad=True, device=device) + target = torch.randint(V, (BT, 1), device=device).squeeze(1) + if input.kernel_provider == "liger": + loss_fn = LigerCrossEntropyLoss() + elif input.kernel_provider == "huggingface": + loss_fn = CrossEntropyLoss() + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for CrossEntropy") + return _input, target, loss_fn - V = input.x - provider = input.kernel_provider - B = input.extra_benchmark_config["B"] - T = input.extra_benchmark_config["T"] - _input = torch.randn(B * T, V, requires_grad=True, device=device) - target = torch.randint(V, (B * T, 1), device=device).squeeze(1) +def bench_speed_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _setup_cross_entropy(input) + mode = input.kernel_operation_mode def fwd(): - if provider == "liger": - return liger_ce(_input, target) - else: - return torch_ce(_input, target) + return loss_fn(_input, target) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) + elif mode == "no-grad-forward": + with torch.no_grad(): + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) + elif mode == "backward": + y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + grad_to_none=[_input], + rep=100, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward() + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _setup_cross_entropy(input) def full(): - y = fwd() + y = loss_fn(_input, target) y.backward() mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) @@ -47,26 +90,27 @@ def bench_memory_cross_entropy( ) -def bench_speed_cross_entropy( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - torch_ce = CrossEntropyLoss() - liger_ce = LigerCrossEntropyLoss() +def _resolve_model_config_cross_entropy(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_cross_entropy( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "vocab_size": model_info["vocab_size"], + }, + ) + ) - V = input.x - provider = input.kernel_provider + +def bench_speed_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_cross_entropy(input) mode = input.kernel_operation_mode - B = input.extra_benchmark_config["B"] - T = input.extra_benchmark_config["T"] - - _input = torch.randn(B * T, V, requires_grad=True, device=device) - target = torch.randint(V, (B * T, 1), device=device).squeeze(1) def fwd(): - if provider == "liger": - return liger_ce(_input, target) - else: - return torch_ce(_input, target) + return loss_fn(_input, target) if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) @@ -89,6 +133,8 @@ def bench_speed_cross_entropy( 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, @@ -97,30 +143,124 @@ def bench_speed_cross_entropy( ) +def bench_memory_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_cross_entropy(input) + + def full(): + y = loss_fn(_input, target) + y.backward() + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs = { - "kernel_name": "cross_entropy", - "x_name": "V", - "x_label": "vocab size", - "x_values": [2**i for i in range(12, 18)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [{"B": 8, "T": 2048}], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - run_benchmarks( - bench_test_fn=bench_speed_cross_entropy, - kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_cross_entropy, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="huggingface", + extra_benchmark_config={ + "vocab_size": model_cfg.vocab_size, + }, + ) + _input, target, loss_fn = _setup_cross_entropy(probe_input) + return loss_fn(_input, target) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + + model_configs_info = { + cfg.name: { + "vocab_size": cfg.vocab_size, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "cross_entropy", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_cross_entropy_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_cross_entropy_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_bt = 1024 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="huggingface", + extra_benchmark_config={ + "vocab_size": model.vocab_size, + }, + ) + _input, target, loss_fn = _setup_cross_entropy(probe_input) + return loss_fn(_input, target) + + 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": "cross_entropy", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [ + { + "vocab_size": model.vocab_size, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_cross_entropy, + kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_cross_entropy, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_fused_linear_cross_entropy.py b/benchmark/scripts/benchmark_fused_linear_cross_entropy.py index 4d36a66a..542e9988 100644 --- a/benchmark/scripts/benchmark_fused_linear_cross_entropy.py +++ b/benchmark/scripts/benchmark_fused_linear_cross_entropy.py @@ -1,6 +1,13 @@ +import math + import torch import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import 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 @@ -45,36 +52,72 @@ class LigerLMHeadCE(torch.nn.Module): return self.ce_loss(self.lin.weight, x, y) -############################################################################# -# Test the memory consumption of the linear fused cross entropy loss -############################################################################# - - -def bench_memory_fused_linear_cross_entropy( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +def _setup_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): + """Create input tensor, target, and fused linear CE from benchmark config.""" + cfg = input.extra_benchmark_config + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] BT = input.x - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider - - lm_head_ce = None - if provider == "liger": - lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype).to(device) - elif provider == "liger-fp32-accum": - lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, accum_dtype=torch.float32).to(device) - else: - lm_head_ce = TorchLMHeadCE(H=H, V=V, dtype=dtype).to(device) _input = torch.randn(BT, H, requires_grad=True, dtype=dtype, device=device) target = torch.randint(V, (BT, 1), dtype=torch.long, device=device).squeeze(1) + if input.kernel_provider == "liger": + lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype).to(device) + elif input.kernel_provider == "liger-fp32-accum": + lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, accum_dtype=torch.float32).to(device) + else: + lm_head_ce = TorchLMHeadCE(H=H, V=V, dtype=dtype).to(device) + return _input, target, lm_head_ce + + +def bench_speed_fused_linear_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(input) + mode = input.kernel_operation_mode + def fwd(): return lm_head_ce(_input, target) - def full(): + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench( + fwd, + rep=100, + quantiles=QUANTILES, + ) + elif mode == "no-grad-forward": + with torch.no_grad(): + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) + elif mode == "backward": y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + grad_to_none=[_input], + rep=100, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward() + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_fused_linear_cross_entropy(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(input) + + def full(): + y = lm_head_ce(_input, target) y.backward() mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) @@ -86,32 +129,27 @@ def bench_memory_fused_linear_cross_entropy( ) -# ############################################################################# -# # Test the speed of the fused linear cross entropy loss -# ############################################################################# +def _resolve_model_config_fused_linear_cross_entropy(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_fused_linear_cross_entropy( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "hidden_size": model_info["hidden_size"], + "vocab_size": model_info["vocab_size"], + "dtype": model_info["dtype"], + }, + ) + ) -def bench_speed_fused_linear_cross_entropy( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - BT = input.x - 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_fused_linear_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, lm_head_ce = _resolve_model_config_fused_linear_cross_entropy(input) mode = input.kernel_operation_mode - lm_head_ce = None - if provider == "liger": - lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype).to(device) - elif provider == "liger-fp32-accum": - lm_head_ce = LigerLMHeadCE(H=H, V=V, dtype=dtype, accum_dtype=torch.float32).to(device) - else: - lm_head_ce = TorchLMHeadCE(H=H, V=V, dtype=dtype).to(device) - - _input = torch.randn(BT, H, requires_grad=True, dtype=dtype, device=device) - target = torch.randint(V, (BT, 1), dtype=torch.long, device=device).squeeze(1) - def fwd(): return lm_head_ce(_input, target) @@ -148,6 +186,9 @@ def bench_speed_fused_linear_cross_entropy( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -155,30 +196,132 @@ def bench_speed_fused_linear_cross_entropy( ) +def bench_memory_fused_linear_cross_entropy_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, lm_head_ce = _resolve_model_config_fused_linear_cross_entropy(input) + + def full(): + y = lm_head_ce(_input, target) + y.backward() + + mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs = { - "kernel_name": "fused_linear_cross_entropy", - "x_name": "BT", - "x_label": "B x T", - "x_values": [2**i for i in range(12, 16)], - "kernel_providers": ["liger", "liger-fp32-accum", "huggingface"], - "extra_benchmark_configs": [{"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()) - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_cross_entropy, - kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_cross_entropy, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model_cfg.hidden_size, + "vocab_size": model_cfg.vocab_size, + "dtype": model_cfg.dtype, + }, + ) + _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(probe_input) + return lm_head_ce(_input, target) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + + model_configs_info = { + cfg.name: { + "hidden_size": cfg.hidden_size, + "vocab_size": cfg.vocab_size, + "dtype": cfg.dtype, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "fused_linear_cross_entropy", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "liger-fp32-accum", "huggingface"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_fused_linear_cross_entropy_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_fused_linear_cross_entropy_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_bt = 1024 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model.hidden_size, + "vocab_size": model.vocab_size, + "dtype": model.dtype, + }, + ) + _input, target, lm_head_ce = _setup_fused_linear_cross_entropy(probe_input) + return lm_head_ce(_input, target) + + 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_cross_entropy", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "liger-fp32-accum", "huggingface"], + "extra_benchmark_configs": [ + { + "hidden_size": model.hidden_size, + "vocab_size": model.vocab_size, + "dtype": model.dtype, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_fused_linear_cross_entropy, + kernel_operation_modes=["forward", "backward", "full", "no-grad-forward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_fused_linear_cross_entropy, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_fused_linear_jsd.py b/benchmark/scripts/benchmark_fused_linear_jsd.py index ac62863b..65d72c51 100644 --- a/benchmark/scripts/benchmark_fused_linear_jsd.py +++ b/benchmark/scripts/benchmark_fused_linear_jsd.py @@ -1,6 +1,13 @@ +import math + import torch import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import 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 @@ -112,19 +119,13 @@ class LigerLMHeadJSD(torch.nn.Module): ) -############################################################################# -# Test the memory consumption of the fused linear JSD -############################################################################# - - -def bench_memory_fused_linear_jsd( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: +def _setup_fused_linear_jsd(input: SingleBenchmarkRunInput): + """Create input tensors and fused linear JSD from benchmark config.""" + cfg = input.extra_benchmark_config + H = cfg["hidden_size"] + V = cfg["vocab_size"] + dtype = cfg["dtype"] BT = input.x - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider torch_lm_head_jsd = TorchLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) liger_lm_head_jsd = LigerLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) @@ -140,14 +141,55 @@ def bench_memory_fused_linear_jsd( student_input = torch.rand(BT, H, requires_grad=True, dtype=dtype, device=device) teacher_input = torch.rand(BT, H, dtype=dtype, device=device) + if input.kernel_provider == "liger": + lm_head = liger_lm_head_jsd + elif input.kernel_provider == "torch": + lm_head = torch_lm_head_jsd + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for FusedLinearJSD") + + return student_input, teacher_input, lm_head + + +def bench_speed_fused_linear_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + student_input, teacher_input, lm_head = _setup_fused_linear_jsd(input) + mode = input.kernel_operation_mode + def fwd(): - if provider == "liger": - return liger_lm_head_jsd(student_input, teacher_input) - elif provider == "torch": - return torch_lm_head_jsd(student_input, teacher_input) + return lm_head(student_input, teacher_input) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES) + elif mode == "backward": + y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + grad_to_none=[student_input], + rep=100, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward() + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_fused_linear_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + student_input, teacher_input, lm_head = _setup_fused_linear_jsd(input) def full(): - y = fwd() + y = lm_head(student_input, teacher_input) y.backward() mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) @@ -158,41 +200,29 @@ def bench_memory_fused_linear_jsd( ) -# ############################################################################# -# # Test the speed of the fused linear JSD -# ############################################################################# +def _resolve_model_config_fused_linear_jsd(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_fused_linear_jsd( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "hidden_size": model_info["hidden_size"], + "vocab_size": model_info["vocab_size"], + "dtype": model_info["dtype"], + }, + ) + ) -def bench_speed_fused_linear_jsd( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - BT = input.x - H = input.extra_benchmark_config["H"] - V = input.extra_benchmark_config["V"] +def bench_speed_fused_linear_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + student_input, teacher_input, lm_head = _resolve_model_config_fused_linear_jsd(input) mode = input.kernel_operation_mode - dtype = input.extra_benchmark_config["dtype"] - provider = input.kernel_provider - - torch_lm_head_jsd = TorchLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) - liger_lm_head_jsd = LigerLMHeadJSD(H=H, V=V, dtype=dtype, device=device).to(device) - - # init the linear in all FusedLinearJSDs with the same weights - torch_lm_head_jsd.student_lin.weight.data = liger_lm_head_jsd.student_lin.weight.data = torch.rand( - V, H, device=device, dtype=dtype - ) - torch_lm_head_jsd.teacher_lin.weight.data = liger_lm_head_jsd.teacher_lin.weight.data = torch.rand( - V, H, device=device, dtype=dtype - ) - - student_input = torch.rand(BT, H, requires_grad=True, dtype=dtype, device=device) - teacher_input = torch.rand(BT, H, dtype=dtype, device=device) - def fwd(): - if provider == "liger": - return liger_lm_head_jsd(student_input, teacher_input) - elif provider == "torch": - return torch_lm_head_jsd(student_input, teacher_input) + return lm_head(student_input, teacher_input) if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -205,11 +235,7 @@ def bench_speed_fused_linear_jsd( ms_50, ms_20, ms_80 = triton.testing.do_bench( lambda: y.backward(retain_graph=True), - grad_to_none=[ - student_input, - torch_lm_head_jsd.student_lin.weight, - torch_lm_head_jsd.teacher_lin.weight, - ], + grad_to_none=[student_input], rep=100, quantiles=QUANTILES, ) @@ -224,6 +250,9 @@ def bench_speed_fused_linear_jsd( rep=100, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -231,30 +260,132 @@ def bench_speed_fused_linear_jsd( ) +def bench_memory_fused_linear_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + student_input, teacher_input, lm_head = _resolve_model_config_fused_linear_jsd(input) + + def full(): + y = lm_head(student_input, teacher_input) + y.backward() + + mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs = { - "kernel_name": "fused_linear_jsd", - "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}], - "overwrite": args.overwrite, - } + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - run_benchmarks( - bench_test_fn=bench_speed_fused_linear_jsd, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs, - ) - run_benchmarks( - bench_test_fn=bench_memory_fused_linear_jsd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs, - ) + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "hidden_size": model_cfg.hidden_size, + "vocab_size": model_cfg.vocab_size, + "dtype": model_cfg.dtype, + }, + ) + student_input, teacher_input, lm_head = _setup_fused_linear_jsd(probe_input) + return lm_head(student_input, teacher_input) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + + model_configs_info = { + cfg.name: { + "hidden_size": cfg.hidden_size, + "vocab_size": cfg.vocab_size, + "dtype": cfg.dtype, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "fused_linear_jsd", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_fused_linear_jsd_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_fused_linear_jsd_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_bt = 1024 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "hidden_size": model.hidden_size, + "vocab_size": model.vocab_size, + "dtype": model.dtype, + }, + ) + student_input, teacher_input, lm_head = _setup_fused_linear_jsd(probe_input) + return lm_head(student_input, teacher_input) + + 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_jsd", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "hidden_size": model.hidden_size, + "vocab_size": model.vocab_size, + "dtype": model.dtype, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_fused_linear_jsd, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_fused_linear_jsd, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_jsd.py b/benchmark/scripts/benchmark_jsd.py index 16d71eac..d0e983c5 100644 --- a/benchmark/scripts/benchmark_jsd.py +++ b/benchmark/scripts/benchmark_jsd.py @@ -1,6 +1,13 @@ +import math + import torch import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import 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 @@ -9,7 +16,6 @@ from utils import parse_benchmark_script_args from utils import run_benchmarks from liger_kernel.transformers.jsd import LigerJSD -from liger_kernel.utils import get_total_gpu_memory from liger_kernel.utils import infer_device device = infer_device() @@ -53,39 +59,50 @@ class TorchJSD(torch.nn.Module): return loss.to(self.dtype) -def bench_speed_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - torch_jsd = TorchJSD() - liger_jsd = LigerJSD() +def _setup_jsd(input: SingleBenchmarkRunInput): + """Create input tensors and JSD loss from benchmark config.""" + cfg = input.extra_benchmark_config + V = cfg["vocab_size"] + BT = input.x + _input = torch.randn(BT, V, requires_grad=True, device=device).log_softmax(dim=-1) + target = torch.randn(BT, V, device=device).log_softmax(dim=-1) - _input = torch.randn(B * T, V, requires_grad=True, device=device).log_softmax(dim=-1) - target = torch.randn(B * T, V, device=device).log_softmax(dim=-1) + if input.kernel_provider == "liger": + loss_fn = LigerJSD() + elif input.kernel_provider == "torch": + loss_fn = TorchJSD() + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for JSD") + return _input, target, loss_fn + + +def bench_speed_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _setup_jsd(input) + mode = input.kernel_operation_mode def fwd(): - if input.kernel_provider == "liger": - return liger_jsd(_input, target) - else: - return torch_jsd(_input, target) + return loss_fn(_input, target) - if input.kernel_operation_mode == "forward": + if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif input.kernel_operation_mode == "backward": + elif mode == "backward": y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( lambda: y.backward(retain_graph=True), quantiles=QUANTILES, grad_to_none=[_input], rep=100, ) - elif input.kernel_operation_mode == "full": + elif mode == "full": def full(): y = fwd() y.backward(retain_graph=True) ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -94,23 +111,74 @@ def bench_speed_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: def bench_memory_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - torch_jsd = TorchJSD() - liger_jsd = LigerJSD() - - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - - _input = torch.randn(B * T, V, requires_grad=True, device=device).log_softmax(dim=-1) - target = torch.randn(B * T, V, device=device).log_softmax(dim=-1) - - def fwd(): - if input.kernel_provider == "liger": - return liger_jsd(_input, target) - else: - return torch_jsd(_input, target) + _input, target, loss_fn = _setup_jsd(input) def full(): + y = loss_fn(_input, target) + y.backward(retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + +def _resolve_model_config_jsd(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_jsd( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "vocab_size": model_info["vocab_size"], + }, + ) + ) + + +def bench_speed_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_jsd(input) + mode = input.kernel_operation_mode + + def fwd(): + return loss_fn(_input, target) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) + elif mode == "backward": y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + quantiles=QUANTILES, + grad_to_none=[_input], + rep=100, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward(retain_graph=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_jsd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_jsd(input) + + def full(): + y = loss_fn(_input, target) y.backward(retain_graph=True) mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) @@ -124,34 +192,107 @@ def bench_memory_jsd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput if __name__ == "__main__": args = parse_benchmark_script_args() - gpu_memory_gbs = get_total_gpu_memory() - # We know that the full test will require 54GBs for vocab size 2^17 on torch - if gpu_memory_gbs >= 54: - x_max = 17 + + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "vocab_size": model_cfg.vocab_size, + }, + ) + _input, target, loss_fn = _setup_jsd(probe_input) + return loss_fn(_input, target) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + + model_configs_info = { + cfg.name: { + "vocab_size": cfg.vocab_size, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "jsd", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_jsd_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_jsd_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) else: - x_max = 16 - common_args = { - "kernel_name": "jsd", - "x_name": "V", - "x_label": "vocab size", - "x_values": [2**i for i in range(12, x_max + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"B": 4, "T": 2048}], - "overwrite": args.overwrite, - } + model = get_benchmark_model_config(args.model) + probe_bt = 1024 - run_benchmarks( - bench_test_fn=bench_memory_jsd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_args, - ) + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "vocab_size": model.vocab_size, + }, + ) + _input, target, loss_fn = _setup_jsd(probe_input) + return loss_fn(_input, target) - run_benchmarks( - bench_test_fn=bench_speed_jsd, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_args, - ) + 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": "jsd", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "vocab_size": model.vocab_size, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_jsd, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_jsd, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_kl_div.py b/benchmark/scripts/benchmark_kl_div.py index 09948c38..3d4d8850 100644 --- a/benchmark/scripts/benchmark_kl_div.py +++ b/benchmark/scripts/benchmark_kl_div.py @@ -1,7 +1,14 @@ +import math + import torch import torch.nn as nn import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import 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 @@ -14,43 +21,53 @@ from liger_kernel.utils import infer_device device = infer_device() -S, E = 12, 18 - -def bench_speed_kldiv(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: +def _setup_kl_div(input: SingleBenchmarkRunInput): + """Create input tensors and KL div loss from benchmark config.""" + cfg = input.extra_benchmark_config + V = cfg["vocab_size"] + BT = input.x reduction = "batchmean" - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - torch_kl_div = nn.KLDivLoss(reduction=reduction) - liger_kl_div = LigerKLDIVLoss(reduction=reduction) - _input = torch.randn(B * T, V, requires_grad=True, device=device).log_softmax(dim=-1) - target = torch.randn(B * T, V, device=device).softmax(dim=-1) + _input = torch.randn(BT, V, requires_grad=True, device=device).log_softmax(dim=-1) + target = torch.randn(BT, V, device=device).softmax(dim=-1) + + if input.kernel_provider == "liger": + loss_fn = LigerKLDIVLoss(reduction=reduction) + elif input.kernel_provider == "torch": + loss_fn = nn.KLDivLoss(reduction=reduction) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for KLDiv") + return _input, target, loss_fn + + +def bench_speed_kl_div(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _setup_kl_div(input) + mode = input.kernel_operation_mode def fwd(): - if input.kernel_provider == "liger": - return liger_kl_div(_input, target) - else: - return torch_kl_div(_input, target) + return loss_fn(_input, target) - if input.kernel_operation_mode == "forward": + if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif input.kernel_operation_mode == "backward": + elif mode == "backward": y = fwd() - ms_50, ms_20, ms_80 = triton.testing.do_bench( lambda: y.backward(retain_graph=True), quantiles=QUANTILES, grad_to_none=[_input], rep=100, ) - elif input.kernel_operation_mode == "full": + elif mode == "full": def full(): y = fwd() y.backward(retain_graph=True) ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -58,25 +75,75 @@ def bench_speed_kldiv(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutpu ) -def bench_memory_kldiv(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - reduction = "batchmean" - torch_kl_div = nn.KLDivLoss(reduction=reduction) - liger_kl_div = LigerKLDIVLoss(reduction=reduction) - - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - - _input = torch.randn(B * T, V, requires_grad=True, device=device).log_softmax(dim=-1) - target = torch.randn(B * T, V, device=device).softmax(dim=-1) - - def fwd(): - if input.kernel_provider == "liger": - return liger_kl_div(_input, target) - else: - return torch_kl_div(_input, target) +def bench_memory_kl_div(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _setup_kl_div(input) def full(): + y = loss_fn(_input, target) + y.backward(retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + +def _resolve_model_config_kl_div(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_kl_div( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "vocab_size": model_info["vocab_size"], + }, + ) + ) + + +def bench_speed_kl_div_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_kl_div(input) + mode = input.kernel_operation_mode + + def fwd(): + return loss_fn(_input, target) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) + elif mode == "backward": y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + quantiles=QUANTILES, + grad_to_none=[_input], + rep=100, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward(retain_graph=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_kl_div_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_kl_div(input) + + def full(): + y = loss_fn(_input, target) y.backward(retain_graph=True) mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) @@ -90,28 +157,106 @@ def bench_memory_kldiv(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutp if __name__ == "__main__": args = parse_benchmark_script_args() - common_args = { - "kernel_name": "kl_div", - "x_name": "V", - "x_label": "vocab size", - "x_values": [2**i for i in range(12, 18)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"B": 8, "T": 512}], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_memory_kldiv, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_args, - ) + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - run_benchmarks( - bench_test_fn=bench_speed_kldiv, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_args, - ) + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "vocab_size": model_cfg.vocab_size, + }, + ) + _input, target, loss_fn = _setup_kl_div(probe_input) + return loss_fn(_input, target) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + model_configs_info = { + cfg.name: { + "vocab_size": cfg.vocab_size, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "kl_div", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_kl_div_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_kl_div_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + 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={ + "vocab_size": model.vocab_size, + }, + ) + _input, target, loss_fn = _setup_kl_div(probe_input) + return loss_fn(_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": "kl_div", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "vocab_size": model.vocab_size, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_kl_div, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_kl_div, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_llama4_rope.py b/benchmark/scripts/benchmark_llama4_rope.py index 47d06051..47a913a7 100644 --- a/benchmark/scripts/benchmark_llama4_rope.py +++ b/benchmark/scripts/benchmark_llama4_rope.py @@ -1,6 +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 transformers.models.llama4.configuration_llama4 import Llama4TextConfig from transformers.models.llama4.modeling_llama4 import Llama4TextRotaryEmbedding from transformers.models.llama4.modeling_llama4 import apply_rotary_emb @@ -17,19 +26,17 @@ from liger_kernel.utils import transformers_version_dispatch device = infer_device() +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - mode = input.kernel_operation_mode - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x +def _setup_llama4_rope(input: SingleBenchmarkRunInput): + """Create input tensors and Llama4 RoPE embedding from benchmark config.""" + cfg = input.extra_benchmark_config + num_q_heads = cfg["num_q_heads"] + num_kv_heads = cfg["num_kv_heads"] + dtype = cfg["dtype"] + hidden_size = cfg.get("hidden_size", input.x) + seq_len = cfg.get("seq_len", input.x) head_dim = hidden_size // num_q_heads @@ -69,13 +76,19 @@ def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRu pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) freqs_cis = rotary_emb(q, pos_ids) - def fwd(): - if provider == "liger": - return liger_llama4_text_rotary_pos_emb(q, k, freqs_cis) - elif provider == "huggingface": - return apply_rotary_emb(q, k, freqs_cis) - else: - raise ValueError(f"Invalid provider: {provider} for Llama4 RoPE embedding") + if input.kernel_provider == "liger": + fwd_fn = lambda: liger_llama4_text_rotary_pos_emb(q, k, freqs_cis) + elif input.kernel_provider == "huggingface": + fwd_fn = lambda: apply_rotary_emb(q, k, freqs_cis) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for Llama4 RoPE embedding") + + return q, k, dq, dk, fwd_fn + + +def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd = _setup_llama4_rope(input) + mode = input.kernel_operation_mode if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -104,6 +117,8 @@ def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRu rep=400, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -112,60 +127,69 @@ def bench_speed_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRu def bench_memory_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x - - head_dim = hidden_size // num_q_heads - - # Create Llama4TextConfig for the rotary embedding - config = Llama4TextConfig( - hidden_size=hidden_size, - num_attention_heads=num_q_heads, - num_key_value_heads=num_kv_heads, - head_dim=head_dim, - max_position_embeddings=seq_len, - ) - - rotary_emb = transformers_version_dispatch( - "4.48.0", - Llama4TextRotaryEmbedding, - Llama4TextRotaryEmbedding, - before_kwargs={"config": config, "device": device}, - after_kwargs={"config": config, "device": device}, - ) - - q = torch.randn( - (1, seq_len, num_q_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ) - k = torch.randn( - (1, seq_len, num_kv_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device), - ) - pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) - freqs_cis = rotary_emb(q, pos_ids) + q, k, dq, dk, fwd_fn = _setup_llama4_rope(input) def full(): - if provider == "liger": - q_out, k_out = liger_llama4_text_rotary_pos_emb(q, k, freqs_cis) - else: - q_out, k_out = apply_rotary_emb(q, k, freqs_cis) + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + +def _resolve_model_config_llama4_rope(input: SingleBenchmarkRunInput): + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_llama4_rope( + SingleBenchmarkRunInput( + x=input.x, + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "hidden_size": model_info["hidden_size"], + "num_q_heads": model_info["num_q_heads"], + "num_kv_heads": model_info["num_kv_heads"], + "dtype": model_info["dtype"], + "seq_len": cfg["seq_len"], + }, + ) + ) + + +def bench_speed_llama4_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_llama4_rope(input) + mode = input.kernel_operation_mode + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + elif mode == "backward": + q_out, k_out = fwd_fn() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), + grad_to_none=[q, k], + rep=400, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) + + +def bench_memory_llama4_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_llama4_rope(input) + + def full(): + q_out, k_out = fwd_fn() torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) mem_50, mem_20, mem_80 = _test_memory( @@ -182,64 +206,113 @@ def bench_memory_llama4_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkR if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs_varying_hidden_size = { - "kernel_name": "llama4_rope", - "x_name": "H", - "x_label": "hidden size", - "x_values": [32 * (2**i) for i in range(4, 10, 2)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "seq_len": 2048, - "num_q_heads": 32, - "num_kv_heads": 8, - } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_llama4_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_hidden_size, - ) - run_benchmarks( - bench_test_fn=bench_memory_llama4_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_hidden_size, - ) + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - common_configs_varying_seq_len = { - "kernel_name": "llama4_rope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, 15)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "hidden_size": 8192, - "num_q_heads": 32, - "num_kv_heads": 8, + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model_cfg.hidden_size, + "num_q_heads": model_cfg.num_attention_heads, + "num_kv_heads": model_cfg.num_key_value_heads, + "dtype": model_cfg.dtype, + "seq_len": probe_bt, + }, + ) + _, _, _, _, fwd_fn = _setup_llama4_rope(probe_input) + return fwd_fn()[0] + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + model_configs_info = { + cfg.name: { + "hidden_size": cfg.hidden_size, + "num_q_heads": cfg.num_attention_heads, + "num_kv_heads": cfg.num_key_value_heads, + "dtype": cfg.dtype, } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_llama4_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_seq_len, - ) - run_benchmarks( - bench_test_fn=bench_memory_llama4_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_seq_len, - ) + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "llama4_rope", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_llama4_rope_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_llama4_rope_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_seq_len = 2048 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + "seq_len": probe_seq_len, + }, + ) + _, _, _, _, fwd_fn = _setup_llama4_rope(probe_input) + return fwd_fn()[0] + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_seq_len + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + common_configs = { + "kernel_name": "llama4_rope", + "x_name": "T", + "x_label": "sequence length", + "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [ + { + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_llama4_rope, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_llama4_rope, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_qwen2vl_mrope.py b/benchmark/scripts/benchmark_qwen2vl_mrope.py index ec1c53b8..16b12547 100644 --- a/benchmark/scripts/benchmark_qwen2vl_mrope.py +++ b/benchmark/scripts/benchmark_qwen2vl_mrope.py @@ -1,6 +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 transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLTextConfig from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLRotaryEmbedding from transformers.models.qwen2_vl.modeling_qwen2_vl import apply_multimodal_rotary_pos_emb @@ -16,21 +25,17 @@ from liger_kernel.utils import infer_device device = infer_device() +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def bench_speed_qwen2vl_mrope( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - mode = input.kernel_operation_mode - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x +def _setup_qwen2vl_mrope(input: SingleBenchmarkRunInput): + """Create input tensors and Qwen2VL M-RoPE embedding from benchmark config.""" + cfg = input.extra_benchmark_config + num_q_heads = cfg["num_q_heads"] + num_kv_heads = cfg["num_kv_heads"] + dtype = cfg["dtype"] + hidden_size = cfg.get("hidden_size", input.x) + seq_len = cfg.get("seq_len", input.x) head_dim = hidden_size // num_q_heads mrope_section_hw = head_dim * 3 // 16 @@ -66,13 +71,19 @@ def bench_speed_qwen2vl_mrope( pos_ids = torch.arange(seq_len * 3, device=device, dtype=torch.long).view(3, 1, -1) cos, sin = rotary_emb(k, pos_ids) - def fwd(): - if provider == "liger": - return liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) - elif provider == "huggingface": - return apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) - else: - raise ValueError(f"Invalid provider: {provider} for M-RoPE embedding") + if input.kernel_provider == "liger": + fwd_fn = lambda: liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) + elif input.kernel_provider == "huggingface": + fwd_fn = lambda: apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for M-RoPE embedding") + + return q, k, dq, dk, fwd_fn + + +def bench_speed_qwen2vl_mrope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd = _setup_qwen2vl_mrope(input) + mode = input.kernel_operation_mode if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -101,6 +112,8 @@ def bench_speed_qwen2vl_mrope( rep=400, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -108,60 +121,77 @@ def bench_speed_qwen2vl_mrope( ) -def bench_memory_qwen2vl_mrope( - input: SingleBenchmarkRunInput, -) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x - - head_dim = hidden_size // num_q_heads - - mrope_section_hw = head_dim * 3 // 16 - mrope_section = [ - head_dim // 2 - 2 * mrope_section_hw, - mrope_section_hw, - mrope_section_hw, - ] - config = Qwen2VLTextConfig( - hidden_size=hidden_size, - num_attention_heads=num_q_heads, - num_key_value_heads=num_kv_heads, - rope_theta=1000000.0, - mrope_section=mrope_section, - ) - rotary_emb = Qwen2VLRotaryEmbedding(config, device=device) - q = torch.randn( - (1, seq_len, num_q_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ).transpose(1, 2) - k = torch.randn( - (1, seq_len, num_kv_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ).transpose(1, 2) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device, dtype=dtype), - ) - pos_ids = torch.arange(seq_len * 3, device=device, dtype=torch.long).view(3, 1, -1) - cos, sin = rotary_emb(k, pos_ids) +def bench_memory_qwen2vl_mrope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _setup_qwen2vl_mrope(input) def full(): - if provider == "liger": - q_out, k_out = liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) - else: - q_out, k_out = apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory( + full, + quantiles=QUANTILES, + ) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + +def _resolve_model_config_qwen2vl_mrope(input: SingleBenchmarkRunInput): + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_qwen2vl_mrope( + SingleBenchmarkRunInput( + x=input.x, + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "hidden_size": model_info["hidden_size"], + "num_q_heads": model_info["num_q_heads"], + "num_kv_heads": model_info["num_kv_heads"], + "dtype": model_info["dtype"], + "seq_len": cfg["seq_len"], + }, + ) + ) + + +def bench_speed_qwen2vl_mrope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_qwen2vl_mrope(input) + mode = input.kernel_operation_mode + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + elif mode == "backward": + q_out, k_out = fwd_fn() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), + grad_to_none=[q, k], + rep=400, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( + y_20=ms_20, + y_50=ms_50, + y_80=ms_80, + ) + + +def bench_memory_qwen2vl_mrope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_qwen2vl_mrope(input) + + def full(): + q_out, k_out = fwd_fn() torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) mem_50, mem_20, mem_80 = _test_memory( @@ -178,64 +208,113 @@ def bench_memory_qwen2vl_mrope( if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs_varying_hidden_size = { - "kernel_name": "qwen2vl_mrope", - "x_name": "H", - "x_label": "hidden size", - "x_values": [32 * (2**i) for i in range(4, 10, 2)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "seq_len": 2048, - "num_q_heads": 32, - "num_kv_heads": 8, - } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_qwen2vl_mrope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_hidden_size, - ) - run_benchmarks( - bench_test_fn=bench_memory_qwen2vl_mrope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_hidden_size, - ) + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - common_configs_varying_seq_len = { - "kernel_name": "qwen2vl_mrope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, 15)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "hidden_size": 8192, - "num_q_heads": 32, - "num_kv_heads": 8, + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model_cfg.hidden_size, + "num_q_heads": model_cfg.num_attention_heads, + "num_kv_heads": model_cfg.num_key_value_heads, + "dtype": model_cfg.dtype, + "seq_len": probe_bt, + }, + ) + _, _, _, _, fwd_fn = _setup_qwen2vl_mrope(probe_input) + return fwd_fn()[0] + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + model_configs_info = { + cfg.name: { + "hidden_size": cfg.hidden_size, + "num_q_heads": cfg.num_attention_heads, + "num_kv_heads": cfg.num_key_value_heads, + "dtype": cfg.dtype, } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_qwen2vl_mrope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_seq_len, - ) - run_benchmarks( - bench_test_fn=bench_memory_qwen2vl_mrope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_seq_len, - ) + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "qwen2vl_mrope", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_qwen2vl_mrope_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_qwen2vl_mrope_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_seq_len = 2048 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + "seq_len": probe_seq_len, + }, + ) + _, _, _, _, fwd_fn = _setup_qwen2vl_mrope(probe_input) + return fwd_fn()[0] + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_seq_len + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + common_configs = { + "kernel_name": "qwen2vl_mrope", + "x_name": "T", + "x_label": "sequence length", + "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [ + { + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_qwen2vl_mrope, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_qwen2vl_mrope, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_rope.py b/benchmark/scripts/benchmark_rope.py index 1951c3c2..f59671c0 100644 --- a/benchmark/scripts/benchmark_rope.py +++ b/benchmark/scripts/benchmark_rope.py @@ -1,6 +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 transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding from transformers.models.llama.modeling_llama import apply_rotary_pos_emb @@ -17,19 +26,17 @@ from liger_kernel.utils import transformers_version_dispatch device = infer_device() +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) -def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - mode = input.kernel_operation_mode - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x +def _setup_rope(input: SingleBenchmarkRunInput): + """Create input tensors and RoPE embedding from benchmark config.""" + cfg = input.extra_benchmark_config + num_q_heads = cfg["num_q_heads"] + num_kv_heads = cfg["num_kv_heads"] + dtype = cfg["dtype"] + hidden_size = cfg.get("hidden_size", input.x) + seq_len = cfg.get("seq_len", input.x) head_dim = hidden_size // num_q_heads rotary_emb = transformers_version_dispatch( @@ -58,13 +65,19 @@ def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) cos, sin = rotary_emb(k, pos_ids) - def fwd(): - if provider == "liger": - return liger_rotary_pos_emb(q, k, cos, sin, pos_ids) - elif provider == "huggingface": - return apply_rotary_pos_emb(q, k, cos, sin) - else: - raise ValueError(f"Invalid provider: {provider} for RoPE embedding") + if input.kernel_provider == "liger": + fwd_fn = lambda: liger_rotary_pos_emb(q, k, cos, sin, pos_ids) + elif input.kernel_provider == "huggingface": + fwd_fn = lambda: apply_rotary_pos_emb(q, k, cos, sin) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for RoPE embedding") + + return q, k, dq, dk, fwd_fn + + +def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd = _setup_rope(input) + mode = input.kernel_operation_mode if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -93,6 +106,8 @@ def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput rep=400, quantiles=QUANTILES, ) + else: + raise ValueError(f"Unsupported mode: {mode}") return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -101,49 +116,65 @@ def bench_speed_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput def bench_memory_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - provider = input.kernel_provider - - extra_benchmark_config = input.extra_benchmark_config - num_q_heads = extra_benchmark_config["num_q_heads"] - num_kv_heads = extra_benchmark_config["num_kv_heads"] - dtype = extra_benchmark_config["dtype"] - - # x can be either hidden_size or seq_len - hidden_size = extra_benchmark_config["hidden_size"] if "hidden_size" in extra_benchmark_config else input.x - seq_len = extra_benchmark_config["seq_len"] if "seq_len" in extra_benchmark_config else input.x - - head_dim = hidden_size // num_q_heads - rotary_emb = transformers_version_dispatch( - "4.48.0", - LlamaRotaryEmbedding, - LlamaRotaryEmbedding, - before_kwargs={"dim": head_dim, "device": device}, - after_kwargs={"config": LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), "device": device}, - ) - q = torch.randn( - (1, seq_len, num_q_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ).transpose(1, 2) - k = torch.randn( - (1, seq_len, num_kv_heads, head_dim), - device=device, - requires_grad=True, - dtype=dtype, - ).transpose(1, 2) - dq, dk = ( - torch.randn_like(q, device=device, dtype=dtype), - torch.randn_like(k, device=device), - ) - pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) - cos, sin = rotary_emb(k, pos_ids) + q, k, dq, dk, fwd_fn = _setup_rope(input) def full(): - if provider == "liger": - q_out, k_out = liger_rotary_pos_emb(q, k, cos, sin, pos_ids) - else: - q_out, k_out = apply_rotary_pos_emb(q, k, cos, sin) + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80) + + +def _resolve_model_config_rope(input: SingleBenchmarkRunInput): + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_rope( + SingleBenchmarkRunInput( + x=input.x, + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "hidden_size": model_info["hidden_size"], + "num_q_heads": model_info["num_q_heads"], + "num_kv_heads": model_info["num_kv_heads"], + "dtype": model_info["dtype"], + "seq_len": cfg["seq_len"], + }, + ) + ) + + +def bench_speed_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_rope(input) + mode = input.kernel_operation_mode + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd_fn, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + elif mode == "backward": + q_out, k_out = fwd_fn() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True), + grad_to_none=[q, k], + rep=400, + quantiles=QUANTILES, + ) + elif mode == "full": + + def full(): + q_out, k_out = fwd_fn() + torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, grad_to_none=[q, k], rep=400, quantiles=QUANTILES) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) + + +def bench_memory_rope_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + q, k, dq, dk, fwd_fn = _resolve_model_config_rope(input) + + def full(): + q_out, k_out = fwd_fn() torch.autograd.grad((q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True) mem_50, mem_20, mem_80 = _test_memory( @@ -160,64 +191,113 @@ def bench_memory_rope(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutpu if __name__ == "__main__": args = parse_benchmark_script_args() - common_configs_varying_hidden_size = { - "kernel_name": "rope", - "x_name": "H", - "x_label": "hidden size", - "x_values": [32 * (2**i) for i in range(4, 10, 2)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "seq_len": 2048, - "num_q_heads": 32, - "num_kv_heads": 8, - } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_hidden_size, - ) - run_benchmarks( - bench_test_fn=bench_memory_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_hidden_size, - ) + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) - common_configs_varying_seq_len = { - "kernel_name": "rope", - "x_name": "T", - "x_label": "sequence length", - "x_values": [2**i for i in range(10, 15)], - "kernel_providers": ["liger", "huggingface"], - "extra_benchmark_configs": [ - { - "dtype": torch.bfloat16, - "hidden_size": 8192, - "num_q_heads": 32, - "num_kv_heads": 8, + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model_cfg.hidden_size, + "num_q_heads": model_cfg.num_attention_heads, + "num_kv_heads": model_cfg.num_key_value_heads, + "dtype": model_cfg.dtype, + "seq_len": probe_bt, + }, + ) + _, _, _, _, fwd_fn = _setup_rope(probe_input) + return fwd_fn()[0] # return q_out for memory estimation + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + model_configs_info = { + cfg.name: { + "hidden_size": cfg.hidden_size, + "num_q_heads": cfg.num_attention_heads, + "num_kv_heads": cfg.num_key_value_heads, + "dtype": cfg.dtype, } - ], - "overwrite": args.overwrite, - } - run_benchmarks( - bench_test_fn=bench_speed_rope, - kernel_operation_modes=["forward", "backward", "full"], - metric_name="speed", - metric_unit="ms", - **common_configs_varying_seq_len, - ) - run_benchmarks( - bench_test_fn=bench_memory_rope, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_configs_varying_seq_len, - ) + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "rope", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [{"model_configs": model_configs_info, "seq_len": sweep.seq_len}], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_rope_model_config, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_rope_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) + else: + model = get_benchmark_model_config(args.model) + probe_seq_len = 2048 + + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=0, + kernel_provider="huggingface", + extra_benchmark_config={ + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + "seq_len": probe_seq_len, + }, + ) + _, _, _, _, fwd_fn = _setup_rope(probe_input) + return fwd_fn()[0] + + peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe) + kernel_bpt = peak_bytes // probe_seq_len + config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt) + + common_configs = { + "kernel_name": "rope", + "x_name": "T", + "x_label": "sequence length", + "x_values": [2**i for i in range(10, int(math.log2(max(1024, config.seq_len))) + 1)], + "kernel_providers": ["liger", "huggingface"], + "extra_benchmark_configs": [ + { + "hidden_size": model.hidden_size, + "num_q_heads": model.num_attention_heads, + "num_kv_heads": model.num_key_value_heads, + "dtype": model.dtype, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_rope, + kernel_operation_modes=["forward", "backward", "full"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_rope, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) diff --git a/benchmark/scripts/benchmark_tvd.py b/benchmark/scripts/benchmark_tvd.py index ef76380a..124bee0b 100644 --- a/benchmark/scripts/benchmark_tvd.py +++ b/benchmark/scripts/benchmark_tvd.py @@ -1,6 +1,13 @@ +import math + import torch import triton +from benchmark_model_configs import MODEL_REGISTRY +from benchmark_model_configs import compute_model_config_sweep_config +from benchmark_model_configs import compute_seq_len_sweep_config +from benchmark_model_configs import 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 @@ -9,7 +16,6 @@ from utils import parse_benchmark_script_args from utils import run_benchmarks from liger_kernel.transformers.tvd import LigerTVDLoss -from liger_kernel.utils import get_total_gpu_memory from liger_kernel.utils import infer_device device = infer_device() @@ -34,28 +40,35 @@ class TorchTVDLoss(torch.nn.Module): raise ValueError("Invalid reduction type.") -S, E = 12, 18 +def _setup_tvd(input: SingleBenchmarkRunInput): + """Create input tensors and TVD loss from benchmark config.""" + cfg = input.extra_benchmark_config + V = cfg["vocab_size"] + BT = input.x + reduction = "batchmean" + + _input = torch.randn(BT, V, requires_grad=True, device=device).softmax(dim=-1) + target = torch.randn(BT, V, device=device).softmax(dim=-1) + + if input.kernel_provider == "liger": + loss_fn = LigerTVDLoss(reduction=reduction) + elif input.kernel_provider == "torch": + loss_fn = TorchTVDLoss(reduction=reduction) + else: + raise ValueError(f"Invalid provider: {input.kernel_provider} for TVD") + return _input, target, loss_fn def bench_speed_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - reduction = "batchmean" - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - torch_tvd = TorchTVDLoss(reduction=reduction) - liger_tvd = LigerTVDLoss(reduction=reduction) - - _input = torch.randn(B * T, V, requires_grad=True, device=device).softmax(dim=-1) - target = torch.randn(B * T, V, device=device).softmax(dim=-1) + _input, target, loss_fn = _setup_tvd(input) + mode = input.kernel_operation_mode def fwd(): - if input.kernel_provider == "liger": - return liger_tvd(_input, target) - else: - return torch_tvd(_input, target) + return loss_fn(_input, target) - if input.kernel_operation_mode == "forward": + if mode == "forward": ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) - elif input.kernel_operation_mode == "backward": + elif mode == "backward": y = fwd() ms_50, ms_20, ms_80 = triton.testing.do_bench( @@ -64,13 +77,16 @@ def bench_speed_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: grad_to_none=[_input], rep=100, ) - elif input.kernel_operation_mode == "full": + elif mode == "full": def full(): y = fwd() y.backward(retain_graph=True) ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + return SingleBenchmarkRunOutput( y_20=ms_20, y_50=ms_50, @@ -79,24 +95,70 @@ def bench_speed_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: def bench_memory_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: - reduction = "batchmean" - torch_tvd = TorchTVDLoss(reduction=reduction) - liger_tvd = LigerTVDLoss(reduction=reduction) - - V = input.x - B, T = input.extra_benchmark_config["B"], input.extra_benchmark_config["T"] - - _input = torch.randn(B * T, V, requires_grad=True, device=device).softmax(dim=-1) - target = torch.randn(B * T, V, device=device).softmax(dim=-1) - - def fwd(): - if input.kernel_provider == "liger": - return liger_tvd(_input, target) - else: - return torch_tvd(_input, target) + _input, target, loss_fn = _setup_tvd(input) def full(): + y = loss_fn(_input, target) + y.backward(retain_graph=True) + + mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) + return SingleBenchmarkRunOutput( + y_20=mem_20, + y_50=mem_50, + y_80=mem_80, + ) + + +def _resolve_model_config_tvd(input: SingleBenchmarkRunInput): + """Resolve model-config-sweep input into standard setup args.""" + cfg = input.extra_benchmark_config + model_info = cfg["model_configs"][input.x] + return _setup_tvd( + SingleBenchmarkRunInput( + x=cfg["BT"], + kernel_provider=input.kernel_provider, + extra_benchmark_config={ + "vocab_size": model_info["vocab_size"], + }, + ) + ) + + +def bench_speed_tvd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_tvd(input) + mode = input.kernel_operation_mode + + def fwd(): + return loss_fn(_input, target) + + if mode == "forward": + ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, quantiles=QUANTILES, rep=100) + elif mode == "backward": y = fwd() + ms_50, ms_20, ms_80 = triton.testing.do_bench( + lambda: y.backward(retain_graph=True), + quantiles=QUANTILES, + grad_to_none=[_input], + rep=100, + ) + elif mode == "full": + + def full(): + y = fwd() + y.backward(retain_graph=True) + + ms_50, ms_20, ms_80 = triton.testing.do_bench(full, quantiles=QUANTILES, rep=100) + else: + raise ValueError(f"Unsupported mode: {mode}") + + return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80) + + +def bench_memory_tvd_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput: + _input, target, loss_fn = _resolve_model_config_tvd(input) + + def full(): + y = loss_fn(_input, target) y.backward(retain_graph=True) mem_50, mem_20, mem_80 = _test_memory(full, quantiles=QUANTILES) @@ -110,36 +172,107 @@ def bench_memory_tvd(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput if __name__ == "__main__": args = parse_benchmark_script_args() - gpu_memory_gbs = get_total_gpu_memory() - # We know that the full test will require 66GBs for vocab size 2^17 - if gpu_memory_gbs >= 66: - x_max = 17 - elif gpu_memory_gbs >= 32: - x_max = 16 + + if args.sweep_mode == "model_config": + all_model_configs = list(MODEL_REGISTRY.values()) + + def _probe_factory(model_cfg, probe_bt): + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "vocab_size": model_cfg.vocab_size, + }, + ) + _input, target, loss_fn = _setup_tvd(probe_input) + return loss_fn(_input, target) + + return _probe + + sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt) + + model_configs_info = { + cfg.name: { + "vocab_size": cfg.vocab_size, + } + for cfg in sweep.model_configs + } + + common_configs = { + "kernel_name": "tvd", + "x_name": "model_config", + "x_label": "model configuration", + "x_values": [cfg.name for cfg in sweep.model_configs], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "model_configs": model_configs_info, + "BT": sweep.bt, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_tvd_model_config, + kernel_operation_modes=["forward", "full", "backward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_tvd_model_config, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + ) else: - x_max = 15 - common_args = { - "kernel_name": "tvd", - "x_name": "V", - "x_label": "vocab size", - "x_values": [2**i for i in range(12, x_max + 1)], - "kernel_providers": ["liger", "torch"], - "extra_benchmark_configs": [{"B": 8, "T": 2048}], - "overwrite": args.overwrite, - } + model = get_benchmark_model_config(args.model) + probe_bt = 1024 - run_benchmarks( - bench_test_fn=bench_memory_tvd, - kernel_operation_modes=["full"], - metric_name="memory", - metric_unit="MB", - **common_args, - ) + def _probe(): + probe_input = SingleBenchmarkRunInput( + x=probe_bt, + kernel_provider="torch", + extra_benchmark_config={ + "vocab_size": model.vocab_size, + }, + ) + _input, target, loss_fn = _setup_tvd(probe_input) + return loss_fn(_input, target) - run_benchmarks( - bench_test_fn=bench_speed_tvd, - kernel_operation_modes=["forward", "full", "backward"], - metric_name="speed", - metric_unit="ms", - **common_args, - ) + 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": "tvd", + "x_name": "BT", + "x_label": "B * T", + "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)], + "kernel_providers": ["liger", "torch"], + "extra_benchmark_configs": [ + { + "vocab_size": model.vocab_size, + } + ], + "overwrite": args.overwrite, + } + + run_benchmarks( + bench_test_fn=bench_speed_tvd, + kernel_operation_modes=["forward", "full", "backward"], + metric_name="speed", + metric_unit="ms", + **common_configs, + ) + run_benchmarks( + bench_test_fn=bench_memory_tvd, + kernel_operation_modes=["full"], + metric_name="memory", + metric_unit="MB", + **common_configs, + )