optimum_benchmark/scenarios/inference/scenario.py

import time
from contextlib import ExitStack, contextmanager

from transformers import LogitsProcessorList

from ...backends.base import Backend, BackendConfigT
from ...benchmark.report import BenchmarkReport
from ...generators.input_generator import InputGenerator
from ...task_utils import IMAGE_DIFFUSION_TASKS, TEXT_GENERATION_TASKS
from ...trackers.energy import Efficiency, EnergyTracker
from ...trackers.latency import (
    LatencySessionTracker,
    PerStepLatencySessionTrackerPipelineCallback,
    PerTokenLatencySessionTrackerLogitsProcessor,
    Throughput,
)
from ...trackers.memory import MemoryTracker
from ..base import Scenario
from .config import InferenceConfig

PER_TOKEN_BACKENDS = ["pytorch", "onnxruntime", "openvino", "neural-compressor", "ipex"]

TEXT_GENERATION_DEFAULT_KWARGS = {
    "num_return_sequences": 1,
    "max_new_tokens": 100,
    "min_new_tokens": 100,
    "do_sample": False,
    "use_cache": True,
    "num_beams": 1,
}
TEXT_GENERATION_PREFILL_OVERRIDES = {
    "max_new_tokens": 1,
    "min_new_tokens": 1,
}
TEXT_GENERATION_WARMUP_OVERRIDES = {
    "max_new_tokens": 2,
    "min_new_tokens": 2,
}

IMAGE_DIFFUSION_DEFAULT_KWARGS = {
    "num_inference_steps": 30,
    "num_images_per_prompt": 1,
}
IMAGE_DIFFUSION_WARMUP_OVERRIDES = {
    "num_inference_steps": 2,
}

GENERATE_THROUGHPUT_UNIT = "samples/s"
FORWARD_THROUGHPUT_UNIT = "samples/s"
PREFILL_THROUGHPUT_UNIT = "samples/s"
DECODE_THROUGHPUT_UNIT = "tokens/s"
CALL_THROUGHPUT_UNIT = "images/s"

GENERATE_EFFICIENCY_UNIT = "samples/kWh"
FORWARD_EFFICIENCY_UNIT = "samples/kWh"
PREFILL_EFFICIENCY_UNIT = "samples/kWh"
DECODE_EFFICIENCY_UNIT = "tokens/kWh"
CALL_EFFICIENCY_UNIT = "images/kWh"


class InferenceScenario(Scenario[InferenceConfig]):
    NAME = "inference"

    def __init__(self, config: InferenceConfig) -> None:
        super().__init__(config)

    def run(self, backend: Backend[BackendConfigT]) -> BenchmarkReport:
        self.backend = backend

        if self.backend.config.task in TEXT_GENERATION_TASKS:
            self.logger.info("\t+ Updating Text Generation kwargs with default values")
            self.config.generate_kwargs = {**TEXT_GENERATION_DEFAULT_KWARGS, **self.config.generate_kwargs}
        elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
            self.logger.info("\t+ Updating Image Diffusion kwargs with default values")
            self.config.call_kwargs = {**IMAGE_DIFFUSION_DEFAULT_KWARGS, **self.config.call_kwargs}

        if self.backend.config.task in TEXT_GENERATION_TASKS:
            self.logger.info("\t+ Initializing Text Generation targets list")
            targets = ["load_model", "first_generate", "generate", "prefill", "decode"]
            if self.backend.config.name in PER_TOKEN_BACKENDS:
                targets.append("per_token")
        elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
            self.logger.info("\t+ Initializing Image Diffusion targets list")
            targets = ["load_model", "first_call", "call", "per_step"]
        else:
            self.logger.info("\t+ Initializing Inference targets list")
            targets = ["load_model", "first_forward", "forward"]

        self.report = BenchmarkReport.from_list(targets=targets)

        if self.config.latency:
            self.logger.info("\t+ Initializing Latency tracker")
            self.latency_tracker = LatencySessionTracker(
                device=self.backend.config.device, backend=self.backend.config.name
            )
            if self.backend.config.task in TEXT_GENERATION_TASKS and self.backend.config.name in PER_TOKEN_BACKENDS:
                self.logger.info("\t+ Initializing Per-Token Latency tracker")
                self.per_token_latency_tracker = PerTokenLatencySessionTrackerLogitsProcessor(
                    device=self.backend.config.device, backend=self.backend.config.name
                )
                self.config.generate_kwargs["logits_processor"] = LogitsProcessorList([self.per_token_latency_tracker])
            elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
                self.logger.info("\t+ Initializing Diffusion Step Latency tracker")
                self.per_step_latency_tracker = PerStepLatencySessionTrackerPipelineCallback(
                    device=self.backend.config.device, backend=self.backend.config.name
                )
                self.config.call_kwargs["callback_on_step_end"] = self.per_step_latency_tracker

        if self.config.memory:
            self.logger.info("\t+ Initializing Memory tracker")
            self.memory_tracker = MemoryTracker(
                backend=self.backend.config.name,
                device=self.backend.config.device,
                device_ids=self.backend.config.device_ids,
            )

        if self.config.energy:
            self.logger.info("\t+ Initializing Energy tracker")
            self.energy_tracker = EnergyTracker(
                backend=self.backend.config.name,
                device=self.backend.config.device,
                device_ids=self.backend.config.device_ids,
            )

        self.logger.info(f"\t+ Generating inputs for task {self.backend.config.task}")
        self.inputs = InputGenerator(
            task=self.backend.config.task,
            model_shapes=self.backend.model_shapes,
            model_type=self.backend.config.model_type,
            input_shapes=self.config.input_shapes,
        )()

        self.run_model_loading_tracking()

        self.logger.info(f"\t+ Preparing inputs for backend {self.backend.config.name}")
        self.inputs = self.backend.prepare_inputs(inputs=self.inputs)

        if self.config.warmup_runs > 0:
            if self.backend.config.task in TEXT_GENERATION_TASKS:
                self.warmup_text_generation()
            elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
                self.warmup_image_diffusion()
            else:
                self.warmup_inference()

        if self.config.latency:
            if self.backend.config.task in TEXT_GENERATION_TASKS:
                if self.backend.config.name in PER_TOKEN_BACKENDS:
                    self.run_per_token_text_generation_latency_tracking()
                else:
                    self.run_text_generation_latency_tracking()
            elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
                self.run_image_diffusion_latency_tracking()
            else:
                self.run_inference_latency_tracking()

        if self.config.memory:
            if self.backend.config.task in TEXT_GENERATION_TASKS:
                self.run_text_generation_memory_tracking()
            elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
                self.run_image_diffusion_memory_tracking()
            else:
                self.run_inference_memory_tracking()

        if self.config.energy:
            if self.backend.config.task in TEXT_GENERATION_TASKS:
                self.run_text_generation_energy_tracking()
            elif self.backend.config.task in IMAGE_DIFFUSION_TASKS:
                self.run_image_diffusion_energy_tracking()
            else:
                self.run_inference_energy_tracking()

        return self.report

    @contextmanager
    def global_tracking(self, target: str):
        with ExitStack() as context_stack:
            if self.config.energy:
                context_stack.enter_context(self.energy_tracker.track(task_name=target))
            if self.config.memory:
                context_stack.enter_context(self.memory_tracker.track())
            if self.config.latency:
                context_stack.enter_context(self.latency_tracker.session())
                context_stack.enter_context(self.latency_tracker.track())

            yield

        if self.config.latency:
            getattr(self.report, target).latency = self.latency_tracker.get_latency()
        if self.config.memory:
            getattr(self.report, target).memory = self.memory_tracker.get_max_memory()
        if self.config.energy:
            getattr(self.report, target).energy = self.energy_tracker.get_energy()

    # Model loading
    def run_model_loading_tracking(self):
        self.logger.info("\t+ Running model loading tracking")
        with self.global_tracking(target="load_model"):
            self.backend.load()

    # Warmup & cold start
    def warmup_text_generation(self):
        self.logger.info("\t+ Warming up backend for Text Generation")
        self.backend.generate(self.inputs, self.config.generate_kwargs)

        with self.global_tracking(target="first_generate"):
            self.backend.generate(self.inputs, self.config.generate_kwargs)

        for _ in range(self.config.warmup_runs):
            self.backend.generate(self.inputs, {**self.config.generate_kwargs, **TEXT_GENERATION_WARMUP_OVERRIDES})

    def warmup_image_diffusion(self):
        self.logger.info("\t+ Warming up backend for Image Diffusion")

        with self.global_tracking(target="first_call"):
            self.backend.call(self.inputs, self.config.call_kwargs)

        for _ in range(self.config.warmup_runs):
            self.backend.call(self.inputs, {**self.config.call_kwargs, **IMAGE_DIFFUSION_WARMUP_OVERRIDES})

    def warmup_inference(self):
        self.logger.info("\t+ Warming up backend for Inference")

        with self.global_tracking(target="first_forward"):
            self.backend.forward(self.inputs, self.config.forward_kwargs)

        for _ in range(self.config.warmup_runs):
            self.backend.forward(self.inputs, self.config.forward_kwargs)

    # Memory tracking
    def run_text_generation_memory_tracking(self):
        prefill_kwargs = {**self.config.generate_kwargs, **TEXT_GENERATION_PREFILL_OVERRIDES}

        self.logger.info("\t+ Running Text Generation memory tracking")

        with self.memory_tracker.track():
            self.backend.prefill(self.inputs, prefill_kwargs)

        self.report.prefill.memory = self.memory_tracker.get_max_memory()

        with self.memory_tracker.track():
            self.backend.generate(self.inputs, self.config.generate_kwargs)

        self.report.decode.memory = self.memory_tracker.get_max_memory()

    def run_image_diffusion_memory_tracking(self):
        self.logger.info("\t+ Running Image Diffusion memory tracking")

        with self.memory_tracker.track():
            self.backend.call(self.inputs, self.config.call_kwargs)

        self.report.call.memory = self.memory_tracker.get_max_memory()

    def run_inference_memory_tracking(self):
        self.logger.info("\t+ Running Inference memory tracking")

        with self.memory_tracker.track():
            self.backend.forward(self.inputs, self.config.forward_kwargs)

        self.report.forward.memory = self.memory_tracker.get_max_memory()

    # Latency tracking
    def run_per_token_text_generation_latency_tracking(self):
        self.logger.info("\t+ Running Per-Token Text Generation latency tracking")

        with self.per_token_latency_tracker.session():
            while (
                self.per_token_latency_tracker.elapsed() < self.config.duration
                or self.per_token_latency_tracker.count() < self.config.iterations
            ):
                with self.per_token_latency_tracker.track():
                    self.backend.generate(self.inputs, self.config.generate_kwargs)

        per_token_latency = self.per_token_latency_tracker.get_per_token_latency()
        generate_latency = self.per_token_latency_tracker.get_generate_latency()
        prefill_latency = self.per_token_latency_tracker.get_prefill_latency()
        decode_latency = self.per_token_latency_tracker.get_decode_latency()

        self.report.per_token.latency = per_token_latency
        self.report.generate.latency = generate_latency
        self.report.prefill.latency = prefill_latency
        self.report.decode.latency = decode_latency

        # we don't register a per-token throughput,
        # it's a confusing metric and the same signal as the decode throughput
        self.report.generate.throughput = Throughput.from_latency(
            generate_latency, self.atomic_decode_volume, unit=GENERATE_THROUGHPUT_UNIT
        )
        self.report.prefill.throughput = Throughput.from_latency(
            prefill_latency, self.atomic_prefill_volume, unit=PREFILL_THROUGHPUT_UNIT
        )
        self.report.decode.throughput = Throughput.from_latency(
            decode_latency, self.atomic_decode_volume, unit=DECODE_THROUGHPUT_UNIT
        )

    def run_text_generation_latency_tracking(self):
        self.logger.info("\t+ Running Text Generation latency tracking")

        prefill_kwargs = {**self.config.generate_kwargs, **TEXT_GENERATION_PREFILL_OVERRIDES}

        with self.latency_tracker.session():
            while (
                self.latency_tracker.elapsed() < self.config.duration
                or self.latency_tracker.count() < self.config.iterations
            ):
                with self.latency_tracker.track():
                    self.backend.prefill(self.inputs, prefill_kwargs)

        prefill_latency = self.latency_tracker.get_latency()

        self.report.prefill.latency = prefill_latency
        self.report.prefill.throughput = Throughput.from_latency(
            prefill_latency, self.atomic_prefill_volume, unit=PREFILL_THROUGHPUT_UNIT
        )

        with self.latency_tracker.session():
            while (
                self.latency_tracker.elapsed() < self.config.duration
                or self.latency_tracker.count() < self.config.iterations
            ):
                with self.latency_tracker.track():
                    self.backend.generate(self.inputs, self.config.generate_kwargs)

        generate_latency = self.latency_tracker.get_latency()
        self.report.generate.latency = generate_latency
        self.report.generate.throughput = Throughput.from_latency(
            generate_latency, self.atomic_generate_volume, unit=GENERATE_THROUGHPUT_UNIT
        )

        decode_latency = generate_latency - prefill_latency
        self.report.decode.latency = decode_latency
        self.report.decode.throughput = Throughput.from_latency(
            decode_latency, self.atomic_decode_volume, unit=DECODE_THROUGHPUT_UNIT
        )

    def run_image_diffusion_latency_tracking(self):
        self.logger.info("\t+ Running Image Diffusion latency tracking")

        with self.per_step_latency_tracker.session():
            while (
                self.per_step_latency_tracker.elapsed() < self.config.duration
                or self.per_step_latency_tracker.count() < self.config.iterations
            ):
                with self.per_step_latency_tracker.track():
                    self.backend.call(self.inputs, self.config.call_kwargs)

        call_latency = self.per_step_latency_tracker.get_call_latency()
        self.report.call.latency = call_latency
        self.report.call.throughput = Throughput.from_latency(
            call_latency, self.atomic_call_volume, unit=CALL_THROUGHPUT_UNIT
        )

        per_step_latency = self.per_step_latency_tracker.get_step_latency()
        self.report.per_step.latency = per_step_latency

    def run_inference_latency_tracking(self):
        self.logger.info("\t+ Running Inference latency tracking")

        with self.latency_tracker.session():
            while (
                self.latency_tracker.elapsed() < self.config.duration
                or self.latency_tracker.count() < self.config.iterations
            ):
                with self.latency_tracker.track():
                    self.backend.forward(self.inputs, self.config.forward_kwargs)

        forward_latency = self.latency_tracker.get_latency()
        self.report.forward.latency = forward_latency
        self.report.forward.throughput = Throughput.from_latency(
            forward_latency, self.atomic_forward_volume, unit=FORWARD_THROUGHPUT_UNIT
        )

    # Energy tracking
    def run_text_generation_energy_tracking(self):
        self.logger.info("\t+ Running Text Generation energy tracking")
        prefill_kwargs = {**self.config.generate_kwargs, **TEXT_GENERATION_PREFILL_OVERRIDES}

        count, elapsed, start_time = 0, 0, time.perf_counter()

        with self.energy_tracker.track(task_name="prefill"):
            while elapsed < self.config.duration or count < self.config.iterations:
                self.backend.prefill(self.inputs, prefill_kwargs)
                elapsed = time.perf_counter() - start_time
                count += 1

        prefill_energy = self.energy_tracker.get_energy() / count

        self.report.prefill.energy = prefill_energy
        self.report.prefill.efficiency = Efficiency.from_energy(
            prefill_energy, self.atomic_prefill_volume, unit=PREFILL_EFFICIENCY_UNIT
        )

        count, elapsed, start_time = 0, 0, time.perf_counter()

        with self.energy_tracker.track(task_name="generate"):
            while elapsed < self.config.duration or count < self.config.iterations:
                self.backend.generate(self.inputs, self.config.generate_kwargs)
                elapsed = time.perf_counter() - start_time
                count += 1

        generate_energy = self.energy_tracker.get_energy() / count
        self.report.generate.energy = generate_energy
        self.report.generate.efficiency = Efficiency.from_energy(
            generate_energy, self.atomic_generate_volume, unit=GENERATE_EFFICIENCY_UNIT
        )

        decode_energy = generate_energy - prefill_energy
        self.report.decode.energy = decode_energy
        self.report.decode.efficiency = Efficiency.from_energy(
            decode_energy, self.atomic_decode_volume, unit=DECODE_EFFICIENCY_UNIT
        )

    def run_image_diffusion_energy_tracking(self):
        self.logger.info("\t+ Running Image Diffusion energy tracking")

        count = 0
        elapsed = 0
        start_time = time.perf_counter()

        with self.energy_tracker.track(task_name="call"):
            while elapsed < self.config.duration or count < self.config.iterations:
                self.backend.call(self.inputs, self.config.call_kwargs)
                elapsed = time.perf_counter() - start_time
                count += 1

        call_energy = self.energy_tracker.get_energy() / count

        self.report.call.energy = call_energy
        self.report.call.efficiency = Efficiency.from_energy(
            call_energy, self.atomic_call_volume, unit=CALL_EFFICIENCY_UNIT
        )

    def run_inference_energy_tracking(self):
        self.logger.info("\t+ Running energy tracking")

        count = 0
        elapsed = 0
        start_time = time.perf_counter()

        with self.energy_tracker.track(task_name="forward"):
            while elapsed < self.config.duration or count < self.config.iterations:
                self.backend.forward(self.inputs, self.config.forward_kwargs)
                elapsed = time.perf_counter() - start_time
                count += 1

        forward_energy = self.energy_tracker.get_energy() / count

        self.report.forward.energy = forward_energy
        self.report.forward.efficiency = Efficiency.from_energy(
            forward_energy, self.atomic_forward_volume, unit=FORWARD_EFFICIENCY_UNIT
        )

    @property
    def atomic_forward_volume(self) -> int:  # in terms of processed samples
        return self.config.input_shapes["batch_size"]

    @property
    def atomic_generate_volume(self) -> int:  # in terms of processed samples
        return self.config.input_shapes["batch_size"]

    @property
    def atomic_prefill_volume(self) -> int:  # in terms of processed samples
        return self.config.input_shapes["batch_size"]

    @property
    def atomic_decode_volume(self) -> int:  # in terms of generated tokens
        return (
            self.config.input_shapes["batch_size"]
            * self.config.generate_kwargs["num_beams"]  # at each beam stage there are num_beams tokens generated
            * (self.config.generate_kwargs["max_new_tokens"] - 1)  # 1 token is generated during prefill
        )

    @property
    def atomic_call_volume(self) -> int:  # in terms of generated images
        if self.backend.config.task == "text-to-image":
            return self.config.input_shapes["batch_size"] * self.config.call_kwargs["num_images_per_prompt"]
        else:
            return self.config.input_shapes["batch_size"]