mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
[Benchmark] Benchmark refactor for loss ops (#1191)
## Summary * Add `model_config` sweep support to loss ops in Liger-Kernel benchmarks, enabling evaluation across different model architectures at a fixed sequence length * Introduce a pre-sweep probe step to determine safe configurations and prevent OOM during benchmarking Loss kernels: benchmark_cpo_loss.py benchmark_distill_cosine_loss.py benchmark_distill_jsd_loss.py benchmark_dpo_loss.py benchmark_grpo_loss.py benchmark_kto_loss.py benchmark_orpo_loss.py benchmark_simpo_loss.py `BT` denotes the product of batch size and sequence length (`batch_size` × `sequence_length`). Sweep mode `model_config`: `python benchmark_grpo_loss.py --sweep-mode model_config --bt 2048` Sweep mode `BT`: `python benchmark_grpo_loss.py --model llama_3_8b` Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <[email protected]>
This commit is contained in:
lowdy1
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -18,41 +24,72 @@ device = infer_device()
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Test the memory consumption of the linear fused cross entropy loss
|
||||
#############################################################################
|
||||
|
||||
|
||||
def bench_memory_fused_linear_cpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
def _setup_cpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and CPO loss from benchmark config."""
|
||||
from test.chunked_loss.test_cpo_loss import LigerLMHeadCPO
|
||||
from test.chunked_loss.test_cpo_loss import TorchLMHeadCPO
|
||||
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0]
|
||||
liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0]
|
||||
T = cfg["T"]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
elif input.kernel_provider == "huggingface":
|
||||
loss_module = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for CPOLoss")
|
||||
|
||||
fwd_fn = lambda: loss_module(_input, target)[0]
|
||||
return _input, fwd_fn
|
||||
|
||||
|
||||
def bench_speed_cpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd_fn = _setup_cpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd_fn,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd_fn()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd_fn()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_cpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd_fn = _setup_cpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y = fwd_fn()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
@@ -63,40 +100,28 @@ def bench_memory_fused_linear_cpo_loss(
|
||||
)
|
||||
|
||||
|
||||
# #############################################################################
|
||||
# # Test the speed of the fused linear cross entropy loss
|
||||
# #############################################################################
|
||||
def _resolve_model_config_cpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Resolve model-config-sweep input into standard setup args."""
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_cpo_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_fused_linear_cpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
from test.chunked_loss.test_cpo_loss import LigerLMHeadCPO
|
||||
from test.chunked_loss.test_cpo_loss import TorchLMHeadCPO
|
||||
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_cpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_cpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_cpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_cpo = LigerLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target: torch_lm_head_cpo(x, target)[0]
|
||||
liger_fwd = lambda x, target: liger_lm_head_cpo(x, target)[0]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -105,7 +130,6 @@ def bench_speed_fused_linear_cpo_loss(
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
@@ -123,6 +147,9 @@ def bench_speed_fused_linear_cpo_loss(
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
@@ -130,38 +157,142 @@ def bench_speed_fused_linear_cpo_loss(
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_cpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd_fn = _resolve_model_config_cpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd_fn()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_cpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "B",
|
||||
"x_values": [2**i for i in range(1, 5)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 1024,
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
T = 1024
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_fused_linear_cpo_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_fused_linear_cpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd_fn = _setup_cpo_loss(probe_input)
|
||||
return fwd_fn()
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
|
||||
model_configs_info = {
|
||||
cfg.name: {
|
||||
"hidden_size": cfg.hidden_size,
|
||||
"vocab_size": cfg.vocab_size,
|
||||
"dtype": cfg.dtype,
|
||||
}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_cpo_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"model_configs": model_configs_info,
|
||||
"B": B,
|
||||
"T": T,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_cpo_loss_model_config,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_cpo_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 1024
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd_fn = _setup_cpo_loss(probe_input)
|
||||
return fwd_fn()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_cpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_cpo_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_cpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
@@ -5,6 +6,11 @@ import torch
|
||||
import torch.nn as nn
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -84,106 +90,54 @@ class LigerCosineSimilarityLoss(nn.Module):
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
def _setup_distill_cosine_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and cosine similarity loss from benchmark config."""
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
bias = cfg["bias"]
|
||||
weight_hard_loss = cfg["weight_hard_loss"]
|
||||
weight_soft_loss = cfg["weight_soft_loss"]
|
||||
ignore_index = cfg["ignore_index"]
|
||||
BT = input.x
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"]
|
||||
weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
torch_cosine_loss = TorchCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
bias=bias,
|
||||
).to(device)
|
||||
liger_cosine_loss = LigerCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
|
||||
_tensor = torch.rand(BT, H // 2, device=device, dtype=dtype)
|
||||
student_input1 = _tensor.detach().clone().requires_grad_(True)
|
||||
student_input2 = _tensor.detach().clone().requires_grad_(True)
|
||||
|
||||
student_input = _tensor.detach().clone().requires_grad_(True)
|
||||
teacher_input = torch.rand(BT, H, device=device, dtype=dtype)
|
||||
|
||||
target = torch.randint(0, V, (BT,), device=device, dtype=torch.long)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_cosine_loss(student_input1, teacher_input, target)
|
||||
elif provider == "torch":
|
||||
return torch_cosine_loss(student_input2, teacher_input, target)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
elif input.kernel_provider == "torch":
|
||||
loss_module = TorchCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillCosineLoss")
|
||||
return student_input, teacher_input, target, loss_module
|
||||
|
||||
|
||||
def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
BT = input.x
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"]
|
||||
weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_distill_cosine_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
torch_cosine_loss = TorchCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
|
||||
liger_cosine_loss = LigerCosineSimilarityLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
|
||||
_tensor = torch.rand(BT, H // 2, device=device, dtype=dtype)
|
||||
student_input1 = _tensor.detach().clone().requires_grad_(True)
|
||||
student_input2 = _tensor.detach().clone().requires_grad_(True)
|
||||
|
||||
teacher_input = torch.rand(BT, H, device=device, dtype=dtype)
|
||||
|
||||
target = torch.randint(0, V, (BT,), device=device, dtype=torch.long)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_cosine_loss(student_input1, teacher_input, target)
|
||||
elif provider == "torch":
|
||||
return torch_cosine_loss(student_input2, teacher_input, target)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
@@ -191,18 +145,11 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[student_input1, student_input2],
|
||||
grad_to_none=[student_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
@@ -217,6 +164,8 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
@@ -225,42 +174,226 @@ def bench_speed_cosine_similarity_loss(input: SingleBenchmarkRunInput) -> Single
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_distill_cosine_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(input)
|
||||
|
||||
def full():
|
||||
y = loss_module(student_input, teacher_input, target)
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_model_config_distill_cosine_loss(input: SingleBenchmarkRunInput):
|
||||
"""Resolve model-config-sweep input into standard setup args."""
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_distill_cosine_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["BT"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"bias": cfg["bias"],
|
||||
"weight_hard_loss": cfg["weight_hard_loss"],
|
||||
"weight_soft_loss": cfg["weight_soft_loss"],
|
||||
"ignore_index": cfg["ignore_index"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_distill_cosine_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _resolve_model_config_distill_cosine_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
def fwd():
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[student_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
full,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_distill_cosine_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _resolve_model_config_distill_cosine_loss(input)
|
||||
|
||||
def full():
|
||||
y = loss_module(student_input, teacher_input, target)
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "distill_cosine_loss",
|
||||
"x_name": "BT",
|
||||
"x_label": "B x T",
|
||||
"x_values": [2**i for i in range(10, 14)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=probe_bt,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
},
|
||||
)
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(probe_input)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
|
||||
model_configs_info = {
|
||||
cfg.name: {
|
||||
"hidden_size": cfg.hidden_size,
|
||||
"vocab_size": cfg.vocab_size,
|
||||
"dtype": cfg.dtype,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_cosine_similarity_loss,
|
||||
kernel_operation_modes=["forward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
common_configs = {
|
||||
"kernel_name": "distill_cosine_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"model_configs": model_configs_info,
|
||||
"BT": sweep.bt,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_cosine_similarity_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_distill_cosine_loss_model_config,
|
||||
kernel_operation_modes=["forward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_distill_cosine_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=probe_bt,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
},
|
||||
)
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_cosine_loss(probe_input)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "distill_cosine_loss",
|
||||
"x_name": "BT",
|
||||
"x_label": "B * T",
|
||||
"x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_distill_cosine_loss,
|
||||
kernel_operation_modes=["forward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_distill_cosine_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -12,7 +18,6 @@ from utils import parse_benchmark_script_args
|
||||
from utils import run_benchmarks
|
||||
|
||||
from liger_kernel.chunked_loss.jsd_loss import LigerFusedLinearJSDFunction
|
||||
from liger_kernel.utils import get_total_gpu_memory
|
||||
from liger_kernel.utils import infer_device
|
||||
|
||||
device = infer_device()
|
||||
@@ -89,106 +94,54 @@ class LigerJSDLoss(torch.nn.Module):
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
def _setup_distill_jsd_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and JSD loss from benchmark config."""
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
bias = cfg["bias"]
|
||||
weight_hard_loss = cfg["weight_hard_loss"]
|
||||
weight_soft_loss = cfg["weight_soft_loss"]
|
||||
ignore_index = cfg["ignore_index"]
|
||||
BT = input.x
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"]
|
||||
weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
torch_jsd_loss = TorchJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
liger_jsd_loss = LigerJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
|
||||
_tensor = torch.rand(BT, H // 2, device=device, dtype=dtype)
|
||||
student_input1 = _tensor.detach().clone().requires_grad_(True)
|
||||
student_input2 = _tensor.detach().clone().requires_grad_(True)
|
||||
|
||||
student_input = _tensor.detach().clone().requires_grad_(True)
|
||||
teacher_input = torch.rand(BT, H, device=device, dtype=dtype)
|
||||
|
||||
target = torch.randint(0, V, (BT,), device=device, dtype=torch.long)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_jsd_loss(student_input1, teacher_input, target)
|
||||
elif provider == "torch":
|
||||
return torch_jsd_loss(student_input2, teacher_input, target)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
elif input.kernel_provider == "torch":
|
||||
loss_module = TorchJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for DistillJSDLoss")
|
||||
return student_input, teacher_input, target, loss_module
|
||||
|
||||
|
||||
def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
BT = input.x
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
weight_hard_loss = input.extra_benchmark_config["weight_hard_loss"]
|
||||
weight_soft_loss = input.extra_benchmark_config["weight_soft_loss"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_distill_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
torch_jsd_loss = TorchJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
liger_jsd_loss = LigerJSDLoss(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
ignore_index=ignore_index,
|
||||
bias=bias,
|
||||
weight_hard_loss=weight_hard_loss,
|
||||
weight_soft_loss=weight_soft_loss,
|
||||
).to(device)
|
||||
|
||||
_tensor = torch.rand(BT, H // 2, device=device, dtype=dtype)
|
||||
student_input1 = _tensor.detach().clone().requires_grad_(True)
|
||||
student_input2 = _tensor.detach().clone().requires_grad_(True)
|
||||
|
||||
teacher_input = torch.rand(BT, H, device=device, dtype=dtype)
|
||||
|
||||
target = torch.randint(0, V, (BT,), device=device, dtype=torch.long)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_jsd_loss(student_input1, teacher_input, target)
|
||||
elif provider == "torch":
|
||||
return torch_jsd_loss(student_input2, teacher_input, target)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
@@ -200,7 +153,81 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[student_input1, student_input2],
|
||||
grad_to_none=[student_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_distill_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(input)
|
||||
|
||||
def full():
|
||||
y = loss_module(student_input, teacher_input, target)
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_model_config_distill_jsd_loss(input: SingleBenchmarkRunInput):
|
||||
"""Resolve model-config-sweep input into standard setup args."""
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_distill_jsd_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["BT"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"bias": cfg["bias"],
|
||||
"weight_hard_loss": cfg["weight_hard_loss"],
|
||||
"weight_soft_loss": cfg["weight_soft_loss"],
|
||||
"ignore_index": cfg["ignore_index"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_distill_jsd_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _resolve_model_config_distill_jsd_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
def fwd():
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[student_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
@@ -215,6 +242,8 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
@@ -223,50 +252,148 @@ def bench_speed_jsd_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_distill_jsd_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
student_input, teacher_input, target, loss_module = _resolve_model_config_distill_jsd_loss(input)
|
||||
|
||||
def full():
|
||||
y = loss_module(student_input, teacher_input, target)
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
gpu_memory_gbs = get_total_gpu_memory()
|
||||
# We know that the full test will require 69GBs for vocab size 2^13 and 39GBs for vocab size 2^12 on torch
|
||||
if gpu_memory_gbs >= 69:
|
||||
x_max = 13
|
||||
elif gpu_memory_gbs >= 39:
|
||||
x_max = 12
|
||||
else:
|
||||
x_max = 11
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "distill_jsd_loss",
|
||||
"x_name": "BT",
|
||||
"x_label": "B x T",
|
||||
"x_values": [2**i for i in range(10, x_max + 1)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=probe_bt,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
},
|
||||
)
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(probe_input)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
|
||||
model_configs_info = {
|
||||
cfg.name: {
|
||||
"hidden_size": cfg.hidden_size,
|
||||
"vocab_size": cfg.vocab_size,
|
||||
"dtype": cfg.dtype,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_jsd_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
common_configs = {
|
||||
"kernel_name": "distill_jsd_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"model_configs": model_configs_info,
|
||||
"BT": sweep.bt,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_jsd_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_distill_jsd_loss_model_config,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_distill_jsd_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=probe_bt,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
},
|
||||
)
|
||||
student_input, teacher_input, target, loss_module = _setup_distill_jsd_loss(probe_input)
|
||||
return loss_module(student_input, teacher_input, target)
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "distill_jsd_loss",
|
||||
"x_name": "BT",
|
||||
"x_label": "B * T",
|
||||
"x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"bias": False,
|
||||
"weight_hard_loss": 0.5,
|
||||
"weight_soft_loss": 0.5,
|
||||
"ignore_index": -100,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_distill_jsd_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_distill_jsd_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -18,25 +24,20 @@ device = infer_device()
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
|
||||
def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
def _setup_dpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and DPO loss from benchmark config."""
|
||||
from test.chunked_loss.test_dpo_loss import LigerLMHeadDPO
|
||||
from test.chunked_loss.test_dpo_loss import TorchLMHeadDPO
|
||||
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
bias = cfg["bias"]
|
||||
beta = cfg["beta"]
|
||||
ignore_index = cfg["ignore_index"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
beta = input.extra_benchmark_config["beta"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0]
|
||||
liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0]
|
||||
T = cfg["T"]
|
||||
|
||||
# Input shape: [B, T, H]
|
||||
_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
@@ -49,62 +50,21 @@ def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunO
|
||||
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]
|
||||
target.view(-1)[indices_to_assign] = ignore_index
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, ref_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, ref_input, target)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
elif input.kernel_provider == "huggingface":
|
||||
loss_module = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for DPOLoss")
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
fwd = lambda: loss_module(_input, ref_input, target)[0]
|
||||
return _input, fwd
|
||||
|
||||
|
||||
def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
from test.chunked_loss.test_dpo_loss import LigerLMHeadDPO
|
||||
from test.chunked_loss.test_dpo_loss import TorchLMHeadDPO
|
||||
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
beta = input.extra_benchmark_config["beta"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
_input, fwd = _setup_dpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_dpo_loss = TorchLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
liger_dpo_loss = LigerLMHeadDPO(H=H, V=V, dtype=dtype, beta=beta, ignore_index=ignore_index, bias=bias).to(device)
|
||||
torch_fwd = lambda x, ref_x, target: torch_dpo_loss(x, ref_x, target)[0]
|
||||
liger_fwd = lambda x, ref_x, target: liger_dpo_loss(x, ref_x, target)[0]
|
||||
|
||||
# Input shape: [B, T, H]
|
||||
_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
ref_input = torch.randn(B, T, H, device=device, dtype=dtype, requires_grad=False)
|
||||
# Target shape: [B, T]
|
||||
target = torch.randint(V, (B, T), device=device, dtype=torch.long)
|
||||
|
||||
# Add ignore_index tokens
|
||||
num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item()
|
||||
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]
|
||||
target.view(-1)[indices_to_assign] = ignore_index
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, ref_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, ref_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -130,6 +90,8 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
@@ -138,42 +100,230 @@ def bench_speed_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_dpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_dpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_model_config_dpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Resolve model-config-sweep input into standard setup args."""
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_dpo_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
"bias": cfg["bias"],
|
||||
"beta": cfg["beta"],
|
||||
"ignore_index": cfg["ignore_index"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_dpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_dpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
full,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_dpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_dpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "dpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, 6)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 512,
|
||||
"H": 1024,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
T = 512
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_dpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
|
||||
model_configs_info = {
|
||||
cfg.name: {
|
||||
"hidden_size": cfg.hidden_size,
|
||||
"vocab_size": cfg.vocab_size,
|
||||
"dtype": cfg.dtype,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_dpo_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
T = 512
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_dpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
common_configs = {
|
||||
"kernel_name": "dpo_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"model_configs": model_configs_info,
|
||||
"B": B,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_dpo_loss_model_config,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_dpo_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 512
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_dpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "dpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(config.batch_size * config.seq_len // T)) + 1)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_dpo_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_dpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -18,113 +24,44 @@ device = infer_device()
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Test the memory consumption of the linear fused GRPO loss
|
||||
#############################################################################
|
||||
|
||||
|
||||
def bench_memory_fused_linear_grpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
def _setup_grpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and GRPO loss from benchmark config."""
|
||||
from test.chunked_loss.test_grpo_loss import LigerLMHeadGRPO
|
||||
from test.chunked_loss.test_grpo_loss import TorchLMHeadGRPO
|
||||
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
importance_sampling_level = cfg["importance_sampling_level"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
importance_sampling_level = input.extra_benchmark_config["importance_sampling_level"]
|
||||
provider = input.kernel_provider
|
||||
T = cfg["T"]
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_grpo = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
liger_lm_head_grpo = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
|
||||
# Create inputs
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
selected_token_ids = torch.randint(0, V, (B, T), dtype=torch.long, device=device)
|
||||
attention_mask = torch.ones(B, T, device=device)
|
||||
advantages = torch.randn(B, dtype=dtype, device=device)
|
||||
ref_input = torch.randn(B, T, H, dtype=dtype, device=device)
|
||||
|
||||
torch_fwd = lambda: torch_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[
|
||||
0
|
||||
]
|
||||
liger_fwd = lambda: liger_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[
|
||||
0
|
||||
]
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
elif input.kernel_provider == "torch":
|
||||
loss_module = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for GRPOLoss")
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd()
|
||||
elif provider == "torch":
|
||||
return torch_fwd()
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
fwd = lambda: loss_module(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[0]
|
||||
return _input, fwd
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Test the speed of the fused linear GRPO loss
|
||||
#############################################################################
|
||||
|
||||
|
||||
def bench_speed_fused_linear_grpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
from test.chunked_loss.test_grpo_loss import LigerLMHeadGRPO
|
||||
from test.chunked_loss.test_grpo_loss import TorchLMHeadGRPO
|
||||
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
importance_sampling_level = input.extra_benchmark_config["importance_sampling_level"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_grpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_grpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_grpo = TorchLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
liger_lm_head_grpo = LigerLMHeadGRPO(H=H, V=V, dtype=dtype, importance_sampling_level=importance_sampling_level).to(
|
||||
device
|
||||
)
|
||||
|
||||
# Create inputs
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
selected_token_ids = torch.randint(0, V, (B, T), dtype=torch.long, device=device)
|
||||
attention_mask = torch.ones(B, T, device=device)
|
||||
advantages = torch.randn(B, dtype=dtype, device=device)
|
||||
ref_input = torch.randn(B, T, H, dtype=dtype, device=device)
|
||||
|
||||
torch_fwd = lambda: torch_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[
|
||||
0
|
||||
]
|
||||
liger_fwd = lambda: liger_lm_head_grpo(_input, selected_token_ids, attention_mask, advantages, ref_input=ref_input)[
|
||||
0
|
||||
]
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd()
|
||||
elif provider == "torch":
|
||||
return torch_fwd()
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -133,7 +70,6 @@ def bench_speed_fused_linear_grpo_loss(
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
@@ -151,6 +87,8 @@ def bench_speed_fused_linear_grpo_loss(
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
@@ -158,77 +96,215 @@ def bench_speed_fused_linear_grpo_loss(
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_grpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_grpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
|
||||
|
||||
|
||||
def _resolve_model_config_grpo_loss(input: SingleBenchmarkRunInput):
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_grpo_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
"importance_sampling_level": cfg["importance_sampling_level"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_grpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_grpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
full,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_grpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_grpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
|
||||
|
||||
|
||||
def _run_grpo_benchmarks(args, importance_sampling_level, kernel_name_suffix):
|
||||
"""Run D1 or D2 benchmarks for a given importance_sampling_level."""
|
||||
kernel_name = f"fused_linear_grpo_loss_{kernel_name_suffix}"
|
||||
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
T = 1024
|
||||
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
"importance_sampling_level": importance_sampling_level,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_grpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
model_configs_info = {
|
||||
cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": kernel_name,
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"model_configs": model_configs_info,
|
||||
"B": B,
|
||||
"T": T,
|
||||
"importance_sampling_level": importance_sampling_level,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_grpo_loss_model_config,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_grpo_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 1024
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="torch",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"importance_sampling_level": importance_sampling_level,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_grpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": kernel_name,
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"importance_sampling_level": importance_sampling_level,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_grpo_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_grpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
# Benchmark token-level importance sampling (original GRPO)
|
||||
token_configs = {
|
||||
"kernel_name": "fused_linear_grpo_loss_token",
|
||||
"x_name": "B",
|
||||
"x_label": "B",
|
||||
"x_values": [2**i for i in range(1, 5)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 1024,
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"importance_sampling_level": "token",
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
print("Benchmarking GRPO (token-level importance sampling)...")
|
||||
_run_grpo_benchmarks(args, importance_sampling_level="token", kernel_name_suffix="token")
|
||||
|
||||
# Benchmark sequence-level importance sampling (GSPO)
|
||||
sequence_configs = {
|
||||
"kernel_name": "fused_linear_grpo_loss_sequence",
|
||||
"x_name": "B",
|
||||
"x_label": "B",
|
||||
"x_values": [2**i for i in range(1, 5)],
|
||||
"kernel_providers": ["liger", "torch"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 1024,
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"importance_sampling_level": "sequence",
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
# Run benchmarks for token-level (GRPO)
|
||||
print("Benchmarking GRPO (token-level importance sampling)...")
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_fused_linear_grpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**token_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_fused_linear_grpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**token_configs,
|
||||
)
|
||||
|
||||
# Run benchmarks for sequence-level (GSPO)
|
||||
print("Benchmarking GSPO (sequence-level importance sampling)...")
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_fused_linear_grpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**sequence_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_fused_linear_grpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**sequence_configs,
|
||||
)
|
||||
_run_grpo_benchmarks(args, importance_sampling_level="sequence", kernel_name_suffix="sequence")
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -88,158 +94,52 @@ class LigerLMHeadKTO(torch.nn.Module):
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
def _setup_kto_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and KTO loss from benchmark config."""
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
bias = cfg["bias"]
|
||||
beta = cfg["beta"]
|
||||
ignore_index = cfg["ignore_index"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
beta = input.extra_benchmark_config["beta"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
torch_kto_loss = TorchLMHeadKTO(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
use_bias=bias,
|
||||
use_ref_bias=bias,
|
||||
ignore_index=ignore_index,
|
||||
beta=beta,
|
||||
).to(device)
|
||||
|
||||
liger_kto_loss = LigerLMHeadKTO(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
use_bias=bias,
|
||||
use_ref_bias=bias,
|
||||
ignore_index=ignore_index,
|
||||
beta=beta,
|
||||
).to(device)
|
||||
T = cfg["T"]
|
||||
|
||||
# Input shape: [B, T, H]
|
||||
_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
|
||||
# Target shape: [B, T]
|
||||
ref_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
# Preference labels shape: [B]
|
||||
# Create binary preference labels (0 or 1) for each sequence in the batch
|
||||
# Used to indicate preferred sequences (1) vs non-preferred sequences (0)
|
||||
preference_labels = torch.randint(2, (B,), dtype=torch.bool, device=device)
|
||||
|
||||
# Precomputed KL divergence between policy and reference distributions
|
||||
kl = torch.randn(1, device=device, dtype=dtype)
|
||||
|
||||
# Add ignore_index tokens to simulate padding
|
||||
num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item()
|
||||
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]
|
||||
target.view(-1)[indices_to_assign] = ignore_index
|
||||
|
||||
# Add ref_x with the same shape as _input
|
||||
ref_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadKTO(
|
||||
H=H, V=V, dtype=dtype, use_bias=bias, use_ref_bias=bias, ignore_index=ignore_index, beta=beta
|
||||
).to(device)
|
||||
elif input.kernel_provider == "huggingface":
|
||||
loss_module = TorchLMHeadKTO(
|
||||
H=H, V=V, dtype=dtype, use_bias=bias, use_ref_bias=bias, ignore_index=ignore_index, beta=beta
|
||||
).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for KTOLoss")
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_kto_loss(
|
||||
x=_input,
|
||||
ref_x=ref_input,
|
||||
y=target,
|
||||
preference_labels=preference_labels,
|
||||
kl=kl,
|
||||
)[0]
|
||||
elif provider == "huggingface":
|
||||
return torch_kto_loss(
|
||||
x=_input,
|
||||
ref_x=ref_input,
|
||||
y=target,
|
||||
preference_labels=preference_labels,
|
||||
kl=kl,
|
||||
)[0]
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
fwd = lambda: loss_module(x=_input, ref_x=ref_input, y=target, preference_labels=preference_labels, kl=kl)[0]
|
||||
return _input, fwd
|
||||
|
||||
|
||||
def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
bias = input.extra_benchmark_config["bias"]
|
||||
beta = input.extra_benchmark_config["beta"]
|
||||
ignore_index = input.extra_benchmark_config["ignore_index"]
|
||||
provider = input.kernel_provider
|
||||
_input, fwd = _setup_kto_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
torch_kto_loss = TorchLMHeadKTO(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
beta=beta,
|
||||
ignore_index=ignore_index,
|
||||
use_bias=bias,
|
||||
).to(device)
|
||||
liger_kto_loss = LigerLMHeadKTO(
|
||||
H=H,
|
||||
V=V,
|
||||
dtype=dtype,
|
||||
beta=beta,
|
||||
ignore_index=ignore_index,
|
||||
use_bias=bias,
|
||||
).to(device)
|
||||
|
||||
# Input shape: [B, T, H]
|
||||
_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
|
||||
# Target shape: [B, T]
|
||||
target = torch.randint(V, (B, T), device=device, dtype=torch.long)
|
||||
|
||||
# Preference labels shape: [B]
|
||||
# Create binary preference labels (0 or 1) for each sequence in the batch
|
||||
# Used to indicate preferred sequences (1) vs non-preferred sequences (0)
|
||||
preference_labels = torch.randint(2, (B,), dtype=torch.bool, device=device)
|
||||
|
||||
# Precomputed KL divergence between policy and reference distributions
|
||||
kl = torch.randn(1, device=device, dtype=dtype)
|
||||
|
||||
# Add ignore_index tokens
|
||||
num_elements_to_assign = torch.randint(1, B * T // 2, (1,)).item()
|
||||
indices_to_assign = torch.randperm(B * T)[:num_elements_to_assign]
|
||||
target.view(-1)[indices_to_assign] = ignore_index
|
||||
|
||||
# Add ref_x with the same shape as _input
|
||||
ref_input = torch.randn(B, T, H, device=device, dtype=dtype)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_kto_loss(
|
||||
x=_input,
|
||||
ref_x=ref_input,
|
||||
y=target,
|
||||
preference_labels=preference_labels,
|
||||
kl=kl,
|
||||
)[0]
|
||||
elif provider == "huggingface":
|
||||
return torch_kto_loss(
|
||||
x=_input,
|
||||
ref_x=ref_input,
|
||||
y=target,
|
||||
preference_labels=preference_labels,
|
||||
kl=kl,
|
||||
)[0]
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -265,7 +165,8 @@ def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
@@ -273,42 +174,190 @@ def bench_speed_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOu
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_kto_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_kto_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
|
||||
|
||||
|
||||
def _resolve_model_config_kto_loss(input: SingleBenchmarkRunInput):
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_kto_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
"bias": cfg["bias"],
|
||||
"beta": cfg["beta"],
|
||||
"ignore_index": cfg["ignore_index"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_kto_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_kto_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(fwd, rep=100, quantiles=QUANTILES)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True), grad_to_none=[_input], rep=100, quantiles=QUANTILES
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(y_20=ms_20, y_50=ms_50, y_80=ms_80)
|
||||
|
||||
|
||||
def bench_memory_kto_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_kto_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "kto_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, 6)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 512,
|
||||
"H": 1024,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
T = 512
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_kto_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_kto_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_kto_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
model_configs_info = {
|
||||
cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "kto_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{"model_configs": model_configs_info, "B": B, "T": T, "bias": True, "beta": 0.1, "ignore_index": 42}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_kto_loss_model_config,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_kto_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 512
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_kto_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "kto_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
"bias": True,
|
||||
"beta": 0.1,
|
||||
"ignore_index": 42,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_kto_loss,
|
||||
kernel_operation_modes=["forward", "backward", "full"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_kto_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -18,39 +24,66 @@ device = infer_device()
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Test the memory consumption of the linear fused cross entropy loss
|
||||
#############################################################################
|
||||
|
||||
|
||||
def bench_memory_fused_linear_orpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
def _setup_orpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and ORPO loss from benchmark config."""
|
||||
from test.chunked_loss.test_orpo_loss import LigerLMHeadORPO
|
||||
from test.chunked_loss.test_orpo_loss import TorchLMHeadORPO
|
||||
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0]
|
||||
liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0]
|
||||
T = cfg["T"]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
nll_target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target, nll_target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target, nll_target)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
elif input.kernel_provider == "huggingface":
|
||||
loss_module = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for ORPOLoss")
|
||||
|
||||
fwd = lambda: loss_module(_input, target, nll_target)[0]
|
||||
return _input, fwd
|
||||
|
||||
|
||||
def bench_speed_orpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_orpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True), grad_to_none=[_input], rep=100, quantiles=QUANTILES
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(full, rep=100, quantiles=QUANTILES)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_orpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_orpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
@@ -64,41 +97,27 @@ def bench_memory_fused_linear_orpo_loss(
|
||||
)
|
||||
|
||||
|
||||
# #############################################################################
|
||||
# # Test the speed of the fused linear cross entropy loss
|
||||
# #############################################################################
|
||||
def _resolve_model_config_orpo_loss(input: SingleBenchmarkRunInput):
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_orpo_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_fused_linear_orpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
from test.chunked_loss.test_orpo_loss import LigerLMHeadORPO
|
||||
from test.chunked_loss.test_orpo_loss import TorchLMHeadORPO
|
||||
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_orpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_orpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_orpo = TorchLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_orpo = LigerLMHeadORPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target, nll_target: torch_lm_head_orpo(x, target, nll_target)[0]
|
||||
liger_fwd = lambda x, target, nll_target: liger_lm_head_orpo(x, target, nll_target)[0]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
nll_target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target, nll_target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target, nll_target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -107,7 +126,6 @@ def bench_speed_fused_linear_orpo_loss(
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
@@ -125,6 +143,8 @@ def bench_speed_fused_linear_orpo_loss(
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
@@ -132,38 +152,124 @@ def bench_speed_fused_linear_orpo_loss(
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_orpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_orpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_orpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "B",
|
||||
"x_values": [2**i for i in range(1, 5)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 1024,
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
T = 1024
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_fused_linear_orpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_fused_linear_orpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_orpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
model_configs_info = {
|
||||
cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_orpo_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [{"model_configs": model_configs_info, "B": B, "T": T}],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_orpo_loss_model_config,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_orpo_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 1024
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_orpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_orpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{"hidden_size": model.hidden_size, "vocab_size": model.vocab_size, "dtype": model.dtype, "T": T}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_orpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_orpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from benchmark_model_configs import MODEL_REGISTRY
|
||||
from benchmark_model_configs import compute_model_config_sweep_config
|
||||
from benchmark_model_configs import compute_seq_len_sweep_config
|
||||
from benchmark_model_configs import estimate_kernel_peak_memory
|
||||
from benchmark_model_configs import get_benchmark_model_config
|
||||
from utils import QUANTILES
|
||||
from utils import SingleBenchmarkRunInput
|
||||
from utils import SingleBenchmarkRunOutput
|
||||
@@ -18,85 +24,36 @@ device = infer_device()
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
|
||||
#############################################################################
|
||||
# Test the memory consumption of the linear fused cross entropy loss
|
||||
#############################################################################
|
||||
|
||||
|
||||
def bench_memory_fused_linear_simpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
def _setup_simpo_loss(input: SingleBenchmarkRunInput):
|
||||
"""Create input tensors and SimPO loss from benchmark config."""
|
||||
from test.chunked_loss.test_simpo_loss import LigerLMHeadSimPO
|
||||
from test.chunked_loss.test_simpo_loss import TorchLMHeadCPO
|
||||
|
||||
cfg = input.extra_benchmark_config
|
||||
H = cfg["hidden_size"]
|
||||
V = cfg["vocab_size"]
|
||||
dtype = cfg["dtype"]
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0]
|
||||
liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0]
|
||||
T = cfg["T"]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target)
|
||||
if input.kernel_provider == "liger":
|
||||
loss_module = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)
|
||||
elif input.kernel_provider == "huggingface":
|
||||
loss_module = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
else:
|
||||
raise ValueError(f"Invalid provider: {input.kernel_provider} for SimPOLoss")
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
fwd = lambda: loss_module(_input, target)[0]
|
||||
return _input, fwd
|
||||
|
||||
|
||||
# #############################################################################
|
||||
# # Test the speed of the fused linear cross entropy loss
|
||||
# #############################################################################
|
||||
|
||||
|
||||
def bench_speed_fused_linear_simpo_loss(
|
||||
input: SingleBenchmarkRunInput,
|
||||
) -> SingleBenchmarkRunOutput:
|
||||
from test.chunked_loss.test_simpo_loss import LigerLMHeadSimPO
|
||||
from test.chunked_loss.test_simpo_loss import TorchLMHeadCPO
|
||||
|
||||
B = input.x
|
||||
T = input.extra_benchmark_config["T"]
|
||||
H = input.extra_benchmark_config["H"]
|
||||
V = input.extra_benchmark_config["V"]
|
||||
dtype = input.extra_benchmark_config["dtype"]
|
||||
provider = input.kernel_provider
|
||||
def bench_speed_simpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_simpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
# Instantiate once and retrieve the first output only
|
||||
torch_lm_head_simpo = TorchLMHeadCPO(H=H, V=V, dtype=dtype).to(device)
|
||||
liger_lm_head_simpo = LigerLMHeadSimPO(H=H, V=V, dtype=dtype).to(device)
|
||||
torch_fwd = lambda x, target: torch_lm_head_simpo(x, target)[0]
|
||||
liger_fwd = lambda x, target: liger_lm_head_simpo(x, target)[0]
|
||||
|
||||
_input = torch.randn(B, T, H, requires_grad=True, dtype=dtype, device=device)
|
||||
target = torch.randint(V, (B, T), dtype=torch.long, device=device)
|
||||
|
||||
def fwd():
|
||||
if provider == "liger":
|
||||
return liger_fwd(_input, target)
|
||||
elif provider == "huggingface":
|
||||
return torch_fwd(_input, target)
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
@@ -105,7 +62,6 @@ def bench_speed_fused_linear_simpo_loss(
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
@@ -123,6 +79,8 @@ def bench_speed_fused_linear_simpo_loss(
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
@@ -130,38 +88,190 @@ def bench_speed_fused_linear_simpo_loss(
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_simpo_loss(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _setup_simpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=mem_20,
|
||||
y_50=mem_50,
|
||||
y_80=mem_80,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_model_config_simpo_loss(input: SingleBenchmarkRunInput):
|
||||
cfg = input.extra_benchmark_config
|
||||
model_info = cfg["model_configs"][input.x]
|
||||
return _setup_simpo_loss(
|
||||
SingleBenchmarkRunInput(
|
||||
x=cfg["B"],
|
||||
kernel_provider=input.kernel_provider,
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_info["hidden_size"],
|
||||
"vocab_size": model_info["vocab_size"],
|
||||
"dtype": model_info["dtype"],
|
||||
"T": cfg["T"],
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def bench_speed_simpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_simpo_loss(input)
|
||||
mode = input.kernel_operation_mode
|
||||
|
||||
if mode == "forward":
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
fwd,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "backward":
|
||||
y = fwd()
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
lambda: y.backward(retain_graph=True),
|
||||
grad_to_none=[_input],
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
elif mode == "full":
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
ms_50, ms_20, ms_80 = triton.testing.do_bench(
|
||||
full,
|
||||
rep=100,
|
||||
quantiles=QUANTILES,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {mode}")
|
||||
return SingleBenchmarkRunOutput(
|
||||
y_20=ms_20,
|
||||
y_50=ms_50,
|
||||
y_80=ms_80,
|
||||
)
|
||||
|
||||
|
||||
def bench_memory_simpo_loss_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
|
||||
_input, fwd = _resolve_model_config_simpo_loss(input)
|
||||
|
||||
def full():
|
||||
y = fwd()
|
||||
y.backward()
|
||||
|
||||
mem_50, mem_20, mem_80 = _test_memory(full, _iter=10, quantiles=QUANTILES)
|
||||
return SingleBenchmarkRunOutput(y_20=mem_20, y_50=mem_50, y_80=mem_80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_benchmark_script_args()
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_simpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "B",
|
||||
"x_values": [2**i for i in range(1, 5)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{
|
||||
"T": 1024,
|
||||
"H": 4096,
|
||||
"V": 128256,
|
||||
"mode": "forward",
|
||||
"dtype": torch.bfloat16,
|
||||
}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
if args.sweep_mode == "model_config":
|
||||
all_model_configs = list(MODEL_REGISTRY.values())
|
||||
T = 1024
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_fused_linear_simpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_fused_linear_simpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
def _probe_factory(model_cfg, probe_bt):
|
||||
def _probe():
|
||||
B = max(1, probe_bt // T)
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model_cfg.hidden_size,
|
||||
"vocab_size": model_cfg.vocab_size,
|
||||
"dtype": model_cfg.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_simpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
return _probe
|
||||
|
||||
sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)
|
||||
model_configs_info = {
|
||||
cfg.name: {"hidden_size": cfg.hidden_size, "vocab_size": cfg.vocab_size, "dtype": cfg.dtype}
|
||||
for cfg in sweep.model_configs
|
||||
}
|
||||
B = max(1, sweep.bt // T)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_simpo_loss",
|
||||
"x_name": "model_config",
|
||||
"x_label": "model configuration",
|
||||
"x_values": [cfg.name for cfg in sweep.model_configs],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [{"model_configs": model_configs_info, "B": B, "T": T}],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_simpo_loss_model_config,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_simpo_loss_model_config,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
else:
|
||||
model = get_benchmark_model_config(args.model)
|
||||
T = 1024
|
||||
probe_bt = 1024
|
||||
|
||||
def _probe():
|
||||
B = probe_bt // T
|
||||
probe_input = SingleBenchmarkRunInput(
|
||||
x=B,
|
||||
kernel_provider="huggingface",
|
||||
extra_benchmark_config={
|
||||
"hidden_size": model.hidden_size,
|
||||
"vocab_size": model.vocab_size,
|
||||
"dtype": model.dtype,
|
||||
"T": T,
|
||||
},
|
||||
)
|
||||
_, fwd = _setup_simpo_loss(probe_input)
|
||||
return fwd()
|
||||
|
||||
peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
|
||||
kernel_bpt = peak_bytes // probe_bt
|
||||
config = compute_seq_len_sweep_config(model, kernel_bytes_per_token=kernel_bpt)
|
||||
|
||||
common_configs = {
|
||||
"kernel_name": "fused_linear_simpo_loss",
|
||||
"x_name": "B",
|
||||
"x_label": "Batch Size (B)",
|
||||
"x_values": [2**i for i in range(1, int(math.log2(max(2, config.batch_size * config.seq_len // T))) + 1)],
|
||||
"kernel_providers": ["liger", "huggingface"],
|
||||
"extra_benchmark_configs": [
|
||||
{"hidden_size": model.hidden_size, "vocab_size": model.vocab_size, "dtype": model.dtype, "T": T}
|
||||
],
|
||||
"overwrite": args.overwrite,
|
||||
}
|
||||
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_speed_simpo_loss,
|
||||
kernel_operation_modes=["forward", "full", "backward"],
|
||||
metric_name="speed",
|
||||
metric_unit="ms",
|
||||
**common_configs,
|
||||
)
|
||||
run_benchmarks(
|
||||
bench_test_fn=bench_memory_simpo_loss,
|
||||
kernel_operation_modes=["full"],
|
||||
metric_name="memory",
|
||||
metric_unit="MB",
|
||||
**common_configs,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user