Liger-Kernel/benchmark/scripts/benchmark_modulated_rms_norm.py

import math
import os
import sys

import torch

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))

from benchmark_model_configs import MODEL_REGISTRY
from benchmark_model_configs import compute_model_config_sweep_config
from benchmark_model_configs import compute_seq_len_sweep_config
from benchmark_model_configs import estimate_kernel_peak_memory
from benchmark_model_configs import get_benchmark_model_config
from test.transformers.test_modulated_rms_norm import ModulatedRMSNormReference
from utils import SingleBenchmarkRunInput
from utils import SingleBenchmarkRunOutput
from utils import parse_benchmark_script_args
from utils import run_benchmarks
from utils import run_memory_benchmark
from utils import run_speed_benchmark

from liger_kernel.transformers.modulated_rms_norm import LigerModulatedRMSNorm
from liger_kernel.utils import infer_device

device = infer_device()


def _setup_modulated_rms_norm(input: SingleBenchmarkRunInput):
    """Create input tensors and ModulatedRMSNorm layer from benchmark config."""
    cfg = input.extra_benchmark_config
    hidden_size = cfg["hidden_size"]
    eps = cfg["eps"]
    has_shift = cfg["has_shift"]
    x = torch.randn(
        input.x,
        hidden_size,
        device=device,
        dtype=cfg["dtype"],
        requires_grad=True,
    )
    scale = (torch.randn_like(x) * 0.1).requires_grad_(True)
    shift = (torch.randn_like(x) * 0.1).requires_grad_(True) if has_shift else None

    if input.kernel_provider == "liger":
        layer = LigerModulatedRMSNorm(hidden_size=hidden_size, eps=eps, in_place=False).to(device)
    elif input.kernel_provider == "huggingface":
        layer = ModulatedRMSNormReference(hidden_size=hidden_size, eps=eps).to(device)
    else:
        raise ValueError(f"Invalid provider: {input.kernel_provider} for ModulatedRMSNorm")
    return x, scale, shift, layer


def _input_tensors(x, scale, shift):
    tensors = [x, scale]
    if shift is not None:
        tensors.append(shift)
    return tensors


def bench_speed_modulated_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
    x, scale, shift, layer = _setup_modulated_rms_norm(input)
    return run_speed_benchmark(
        lambda: layer(x, scale, shift), input.kernel_operation_mode, _input_tensors(x, scale, shift)
    )


def bench_memory_modulated_rms_norm(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
    x, scale, shift, layer = _setup_modulated_rms_norm(input)
    return run_memory_benchmark(lambda: layer(x, scale, shift), input.kernel_operation_mode)


def _resolve_model_config_modulated_rms_norm(input: SingleBenchmarkRunInput):
    cfg = input.extra_benchmark_config
    model_info = cfg["model_configs"][input.x]
    return _setup_modulated_rms_norm(
        SingleBenchmarkRunInput(
            x=cfg["BT"],
            kernel_provider=input.kernel_provider,
            extra_benchmark_config={
                "hidden_size": model_info["hidden_size"],
                "dtype": model_info["dtype"],
                "eps": cfg["eps"],
                "has_shift": cfg["has_shift"],
            },
        )
    )


def bench_speed_modulated_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
    x, scale, shift, layer = _resolve_model_config_modulated_rms_norm(input)
    return run_speed_benchmark(
        lambda: layer(x, scale, shift), input.kernel_operation_mode, _input_tensors(x, scale, shift)
    )


def bench_memory_modulated_rms_norm_model_config(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
    x, scale, shift, layer = _resolve_model_config_modulated_rms_norm(input)
    return run_memory_benchmark(lambda: layer(x, scale, shift), input.kernel_operation_mode)


if __name__ == "__main__":
    args = parse_benchmark_script_args()

    if args.sweep_mode == "model_config":
        all_model_configs = list(MODEL_REGISTRY.values())

        def _probe_factory(model_cfg, probe_bt):
            def _probe():
                probe_input = SingleBenchmarkRunInput(
                    x=probe_bt,
                    kernel_provider="huggingface",
                    extra_benchmark_config={
                        "hidden_size": model_cfg.hidden_size,
                        "dtype": model_cfg.dtype,
                        "eps": 1e-6,
                        "has_shift": True,
                    },
                )
                x, scale, shift, layer = _setup_modulated_rms_norm(probe_input)
                return layer(x, scale, shift)

            return _probe

        sweep = compute_model_config_sweep_config(all_model_configs, probe_fn_factory=_probe_factory, bt=args.bt)

        model_configs_info = {
            cfg.name: {
                "hidden_size": cfg.hidden_size,
                "dtype": cfg.dtype,
            }
            for cfg in sweep.model_configs
        }

        common_configs = {
            "kernel_name": "modulated_rms_norm",
            "x_name": "model_config",
            "x_label": "model configuration",
            "x_values": [cfg.name for cfg in sweep.model_configs],
            "kernel_providers": ["liger", "huggingface"],
            "extra_benchmark_configs": [
                {
                    "model_configs": model_configs_info,
                    "BT": sweep.bt,
                    "eps": 1e-6,
                    "has_shift": True,
                }
            ],
            "overwrite": args.overwrite,
        }

        run_benchmarks(
            bench_test_fn=bench_speed_modulated_rms_norm_model_config,
            kernel_operation_modes=["full", "forward", "backward"],
            metric_name="speed",
            metric_unit="ms",
            **common_configs,
        )
        run_benchmarks(
            bench_test_fn=bench_memory_modulated_rms_norm_model_config,
            kernel_operation_modes=["full"],
            metric_name="memory",
            metric_unit="MB",
            **common_configs,
        )
    else:
        model = get_benchmark_model_config(args.model)
        probe_bt = 1024

        def _probe():
            probe_input = SingleBenchmarkRunInput(
                x=probe_bt,
                kernel_provider="huggingface",
                extra_benchmark_config={
                    "hidden_size": model.hidden_size,
                    "dtype": model.dtype,
                    "eps": 1e-6,
                    "has_shift": True,
                },
            )
            x, scale, shift, layer = _setup_modulated_rms_norm(probe_input)
            return layer(x, scale, shift)

        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": "modulated_rms_norm",
            "x_name": "BT",
            "x_label": "B * T",
            "x_values": [2**i for i in range(10, int(math.log2(config.batch_size * config.seq_len)) + 1)],
            "kernel_providers": ["liger", "huggingface"],
            "extra_benchmark_configs": [
                {
                    "hidden_size": model.hidden_size,
                    "dtype": model.dtype,
                    "eps": 1e-6,
                    "has_shift": True,
                }
            ],
            "overwrite": args.overwrite,
        }

        run_benchmarks(
            bench_test_fn=bench_speed_modulated_rms_norm,
            kernel_operation_modes=["full", "forward", "backward"],
            metric_name="speed",
            metric_unit="ms",
            **common_configs,
        )
        run_benchmarks(
            bench_test_fn=bench_memory_modulated_rms_norm,
            kernel_operation_modes=["full"],
            metric_name="memory",
            metric_unit="MB",
            **common_configs,
        )