linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Benchmark] Benchmark refactor for training and RoPE ops (#1193)

Browse Source 
## Summary
* Add `model_config` sweep support to training and Rope op benchmarks,
enabling evaluation across different model architectures at a fixed
sequence length
* Introduce a pre-sweep probe step to determine safe configurations and
prevent OOM during benchmarking

Loss kernels:
benchmark_cross_entropy.py
benchmark_fused_linear_cross_entropy.py
benchmark_jsd.py
benchmark_fused_linear_jsd.py
benchmark_kl_div.py
benchmark_tvd.py
benchmark_rope.py
benchmark_qwen2vl_mrope.py
benchmark_llama4_rope.py

`BT` denotes the product of batch size and sequence length (`batch_size`
× `sequence_length`).
Sweep mode `model_config`:
`python benchmark_cross_entropy.py --sweep-mode model_config --bt 2048`

Sweep mode `BT`:
`python benchmark_cross_entropy.py  --model llama_3_8b`

- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: noemotiovon <[email protected]>
This commit is contained in:
lowdy1
2026-04-18 15:27:22 +00:00
committed by GitHub
co-authored by noemotiovon
parent 1796cf0d4f
commit ef85f2469e
9 changed files with 1815 additions and 750 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/scripts/benchmark_cross_entropy.py
+195 -55
View File
@@ -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,
)
benchmark/scripts/benchmark_fused_linear_cross_entropy.py
+209 -66
View File
@@ -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,
)
benchmark/scripts/benchmark_fused_linear_jsd.py
+206 -75
View File
@@ -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,
)
benchmark/scripts/benchmark_jsd.py
+199 -58
View File
@@ -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,
)
benchmark/scripts/benchmark_kl_div.py
+201 -56
View File
@@ -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,
)
benchmark/scripts/benchmark_llama4_rope.py
+202 -129
View File
@@ -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,
)
benchmark/scripts/benchmark_qwen2vl_mrope.py
+210 -131
View File
@@ -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,
)
benchmark/scripts/benchmark_rope.py
+198 -118
View File
@@ -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,
)
benchmark/scripts/benchmark_tvd.py
+195 -62
View File
@@ -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,
)