mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
UserChen666
e45594124c
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> The original benchmar_grpo_loss.py is aiming to test chunked_loss/grpo_loss.py, so this Pr wants to add benchmark for ops/grpo_loss.py or some other grpo_loss kernels like _ascend/ops/grpo_loss.py. <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <img width="858" height="513" alt="image" src="https://github.com/user-attachments/assets/ca2f4a7b-d4fa-43bd-b58b-a435f6c0046e" /> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: Atlas A800I - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence
344 lines
11 KiB
Python
344 lines
11 KiB
Python
"""
|
|
Benchmark for GRPO loss kernel (ops layer).
|
|
|
|
This benchmark tests the performance of the low-level GRPO loss kernel
|
|
(liger_kernel.ops.grpo_loss.GrpoLossFunction) against a pure PyTorch baseline.
|
|
|
|
Unlike benchmark_grpo_loss.py which tests the chunked_loss layer (fused linear + loss),
|
|
this benchmark focuses on the kernel-level implementation without the linear layer fusion.
|
|
|
|
Usage:
|
|
python benchmark_grpo_loss_kernel.py [--overwrite]
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import triton
|
|
|
|
from utils import QUANTILES
|
|
from utils import SingleBenchmarkRunInput
|
|
from utils import SingleBenchmarkRunOutput
|
|
from utils import _test_memory
|
|
from utils import parse_benchmark_script_args
|
|
from utils import run_benchmarks
|
|
|
|
from liger_kernel.utils import infer_device
|
|
|
|
device = infer_device()
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
|
|
|
|
|
#############################################################################
|
|
# Module wrappers for GRPO loss kernel
|
|
#############################################################################
|
|
|
|
|
|
class TorchGRPOLoss(nn.Module):
|
|
"""Torch baseline module for GRPO loss kernel."""
|
|
|
|
def __init__(self, temperature=1.0, beta=0.0, eps_low=0.2, eps_high=0.2):
|
|
super().__init__()
|
|
self.temperature = temperature
|
|
self.beta = beta
|
|
self.eps_low = eps_low
|
|
self.eps_high = eps_high
|
|
|
|
def forward(self, logits, old_logp, ref_logp, completion_ids, advantages, completion_mask):
|
|
B, L_ADD_1, N = logits.shape
|
|
L = L_ADD_1 - 1
|
|
|
|
# Compute log probabilities
|
|
logits_for_loss = logits[:, :-1, :] / self.temperature
|
|
log_probs = torch.nn.functional.log_softmax(logits_for_loss, dim=-1)
|
|
|
|
# Gather log probs for selected tokens
|
|
completion_ids_expanded = completion_ids.unsqueeze(-1)
|
|
logp = log_probs.gather(dim=-1, index=completion_ids_expanded).squeeze(-1)
|
|
|
|
# Compute importance ratio
|
|
if old_logp is None:
|
|
old_logp_val = logp
|
|
else:
|
|
old_logp_val = old_logp
|
|
|
|
coef_1 = torch.exp(logp - old_logp_val)
|
|
|
|
# Compute clipped coefficient
|
|
coef_2 = torch.clamp(coef_1, 1 - self.eps_low, 1 + self.eps_high)
|
|
|
|
# Expand advantages to per-token
|
|
advantages_expanded = advantages.unsqueeze(-1).expand(-1, L)
|
|
|
|
# Compute per-token loss
|
|
per_token_loss1 = coef_1 * advantages_expanded
|
|
per_token_loss2 = coef_2 * advantages_expanded
|
|
per_token_loss = -torch.minimum(per_token_loss1, per_token_loss2)
|
|
|
|
# Add KL penalty if beta > 0
|
|
if self.beta != 0.0 and ref_logp is not None:
|
|
kl = torch.exp(ref_logp - logp) - (ref_logp - logp) - 1
|
|
per_token_loss += self.beta * kl
|
|
|
|
# Apply mask and reduce
|
|
if completion_mask is not None:
|
|
mask = completion_mask.float()
|
|
else:
|
|
mask = torch.ones(B, L, device=logits.device)
|
|
|
|
# Reduce loss (GRPO uses per-sequence mean)
|
|
loss = ((per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean()
|
|
|
|
return loss
|
|
|
|
|
|
class LigerGRPOLoss(nn.Module):
|
|
"""Liger module wrapper for GRPO loss kernel."""
|
|
|
|
def __init__(self, temperature=1.0, beta=0.0, eps_low=0.2, eps_high=0.2, importance_sampling_level="token"):
|
|
super().__init__()
|
|
self.temperature = temperature
|
|
self.beta = beta
|
|
self.eps_low = eps_low
|
|
self.eps_high = eps_high
|
|
self.importance_sampling_level = importance_sampling_level
|
|
|
|
def forward(self, logits, old_logp, ref_logp, completion_ids, advantages, completion_mask):
|
|
from liger_kernel.ops import GrpoLossFunction
|
|
|
|
return GrpoLossFunction.apply(
|
|
logits,
|
|
old_logp,
|
|
ref_logp,
|
|
completion_ids,
|
|
advantages,
|
|
completion_mask,
|
|
self.temperature,
|
|
self.beta,
|
|
self.eps_low,
|
|
self.eps_high,
|
|
False, # inplace
|
|
"grpo", # loss_type
|
|
None, # max_completion_length
|
|
True, # reduce
|
|
self.importance_sampling_level,
|
|
1.0, # sapo_temperature_pos
|
|
1.05, # sapo_temperature_neg
|
|
None, # vllm_is_ratio
|
|
None, # delta
|
|
False, # use_bias_correction_kl
|
|
)[0]
|
|
|
|
|
|
#############################################################################
|
|
# Test the memory consumption of the GRPO loss kernel
|
|
#############################################################################
|
|
|
|
|
|
def bench_memory_grpo_loss_kernel(
|
|
input: SingleBenchmarkRunInput,
|
|
) -> SingleBenchmarkRunOutput:
|
|
B = input.x
|
|
T = input.extra_benchmark_config["T"]
|
|
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
|
|
|
|
temperature = 1.0
|
|
beta = input.extra_benchmark_config.get("beta", 0.0)
|
|
eps_low = 0.2
|
|
eps_high = 0.2
|
|
|
|
# Instantiate modules
|
|
torch_grpo = TorchGRPOLoss(temperature=temperature, beta=beta, eps_low=eps_low, eps_high=eps_high).to(device)
|
|
liger_grpo = LigerGRPOLoss(
|
|
temperature=temperature,
|
|
beta=beta,
|
|
eps_low=eps_low,
|
|
eps_high=eps_high,
|
|
importance_sampling_level=importance_sampling_level,
|
|
).to(device)
|
|
|
|
# Create inputs
|
|
logits = torch.randn(B, T + 1, V, requires_grad=True, dtype=dtype, device=device)
|
|
completion_ids = torch.randint(0, V, (B, T), dtype=torch.long, device=device)
|
|
advantages = torch.randn(B, dtype=dtype, device=device)
|
|
completion_mask = torch.ones(B, T, dtype=torch.bool, device=device)
|
|
old_logp = None # On-policy case
|
|
ref_logp = torch.randn(B, T, dtype=torch.float32, device=device) if beta > 0 else None
|
|
|
|
def fwd():
|
|
if provider == "liger":
|
|
return liger_grpo(logits, old_logp, ref_logp, completion_ids, advantages, completion_mask)
|
|
elif provider == "torch":
|
|
return torch_grpo(logits, old_logp, ref_logp, completion_ids, advantages, completion_mask)
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
#############################################################################
|
|
# Test the speed of the GRPO loss kernel
|
|
#############################################################################
|
|
|
|
|
|
def bench_speed_grpo_loss_kernel(
|
|
input: SingleBenchmarkRunInput,
|
|
) -> SingleBenchmarkRunOutput:
|
|
B = input.x
|
|
T = input.extra_benchmark_config["T"]
|
|
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
|
|
mode = input.kernel_operation_mode
|
|
|
|
temperature = 1.0
|
|
beta = input.extra_benchmark_config.get("beta", 0.0)
|
|
eps_low = 0.2
|
|
eps_high = 0.2
|
|
|
|
# Instantiate modules
|
|
torch_grpo = TorchGRPOLoss(temperature=temperature, beta=beta, eps_low=eps_low, eps_high=eps_high).to(device)
|
|
liger_grpo = LigerGRPOLoss(
|
|
temperature=temperature,
|
|
beta=beta,
|
|
eps_low=eps_low,
|
|
eps_high=eps_high,
|
|
importance_sampling_level=importance_sampling_level,
|
|
).to(device)
|
|
|
|
# Create inputs
|
|
logits = torch.randn(B, T + 1, V, requires_grad=True, dtype=dtype, device=device)
|
|
completion_ids = torch.randint(0, V, (B, T), dtype=torch.long, device=device)
|
|
advantages = torch.randn(B, dtype=dtype, device=device)
|
|
completion_mask = torch.ones(B, T, dtype=torch.bool, device=device)
|
|
old_logp = None # On-policy case
|
|
ref_logp = torch.randn(B, T, dtype=torch.float32, device=device) if beta > 0 else None
|
|
|
|
def fwd():
|
|
if provider == "liger":
|
|
return liger_grpo(logits, old_logp, ref_logp, completion_ids, advantages, completion_mask)
|
|
elif provider == "torch":
|
|
return torch_grpo(logits, old_logp, ref_logp, completion_ids, advantages, completion_mask)
|
|
|
|
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=[logits],
|
|
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,
|
|
)
|
|
return SingleBenchmarkRunOutput(
|
|
y_20=ms_20,
|
|
y_50=ms_50,
|
|
y_80=ms_80,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_benchmark_script_args()
|
|
|
|
# Benchmark token-level importance sampling (original GRPO)
|
|
token_configs = {
|
|
"kernel_name": "grpo_loss_kernel_token",
|
|
"x_name": "B",
|
|
"x_label": "Batch Size",
|
|
"x_values": [2**i for i in range(1, 5)],
|
|
"kernel_providers": ["liger", "torch"],
|
|
"extra_benchmark_configs": [
|
|
{
|
|
"T": 512,
|
|
"V": 32000,
|
|
"beta": 0.1,
|
|
"dtype": torch.bfloat16,
|
|
"importance_sampling_level": "token",
|
|
}
|
|
],
|
|
"overwrite": args.overwrite,
|
|
}
|
|
|
|
# Benchmark sequence-level importance sampling (GSPO)
|
|
sequence_configs = {
|
|
"kernel_name": "grpo_loss_kernel_sequence",
|
|
"x_name": "B",
|
|
"x_label": "Batch Size",
|
|
"x_values": [2**i for i in range(1, 5)],
|
|
"kernel_providers": ["liger", "torch"],
|
|
"extra_benchmark_configs": [
|
|
{
|
|
"T": 512,
|
|
"V": 32000,
|
|
"beta": 0.1,
|
|
"dtype": torch.bfloat16,
|
|
"importance_sampling_level": "sequence",
|
|
}
|
|
],
|
|
"overwrite": args.overwrite,
|
|
}
|
|
|
|
# Run benchmarks for token-level (GRPO)
|
|
print("Benchmarking GRPO (token-level importance sampling)...")
|
|
run_benchmarks(
|
|
bench_test_fn=bench_speed_grpo_loss_kernel,
|
|
kernel_operation_modes=["forward", "full", "backward"],
|
|
metric_name="speed",
|
|
metric_unit="ms",
|
|
**token_configs,
|
|
)
|
|
run_benchmarks(
|
|
bench_test_fn=bench_memory_grpo_loss_kernel,
|
|
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_grpo_loss_kernel,
|
|
kernel_operation_modes=["forward", "full", "backward"],
|
|
metric_name="speed",
|
|
metric_unit="ms",
|
|
**sequence_configs,
|
|
)
|
|
run_benchmarks(
|
|
bench_test_fn=bench_memory_grpo_loss_kernel,
|
|
kernel_operation_modes=["full"],
|
|
metric_name="memory",
|
|
metric_unit="MB",
|
|
**sequence_configs,
|
|
)
|