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