wehub-resource-sync
736 lines
31 KiB
Python
736 lines
31 KiB
Python
# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
|
|
#
|
|
# This file implements the LitGPT Python API
|
|
import sys
|
|
import time
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
import lightning as L
|
|
import numpy as np
|
|
import torch
|
|
from lightning.fabric.accelerators import CUDAAccelerator
|
|
from lightning.fabric.plugins import BitsandbytesPrecision
|
|
from tqdm import tqdm
|
|
|
|
from litgpt.chat.base import generate as stream_generate_fn
|
|
from litgpt.config import Config, name_to_config
|
|
from litgpt.generate.base import generate as generate_fn
|
|
from litgpt.generate.sequentially import sequential
|
|
from litgpt.generate.tp import tensor_parallel
|
|
from litgpt.model import GPT
|
|
from litgpt.prompts import PromptStyle, has_prompt_style, load_prompt_style, save_prompt_style
|
|
from litgpt.tokenizer import Tokenizer
|
|
from litgpt.utils import (
|
|
auto_download_checkpoint,
|
|
check_file_size_on_cpu_and_warn,
|
|
check_nvlink_connectivity,
|
|
chunked_cross_entropy,
|
|
copy_config_files,
|
|
extend_checkpoint_dir,
|
|
get_default_supported_precision,
|
|
load_checkpoint,
|
|
save_config,
|
|
)
|
|
|
|
|
|
class LLM(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
model: GPT,
|
|
preprocessor=None,
|
|
prompt_style: PromptStyle = None,
|
|
devices: int | list[int] = None,
|
|
config: Config = None,
|
|
checkpoint_dir: Path = None,
|
|
fabric: L.Fabric = None,
|
|
generate_strategy: Literal["sequential", "tensor_parallel"] | None = None,
|
|
kv_cache_initialized: bool = False,
|
|
fixed_kv_cache_size: int | Literal["max_model_supported"] | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.model = model
|
|
self.preprocessor = preprocessor
|
|
self.devices = devices
|
|
self.prompt_style = prompt_style
|
|
self.config = config
|
|
self.checkpoint_dir = checkpoint_dir
|
|
self.fabric = fabric
|
|
self.generate_strategy = generate_strategy
|
|
self.kv_cache_initialized = kv_cache_initialized
|
|
self.fixed_kv_cache_size = fixed_kv_cache_size
|
|
self.prev_generated_seq_length = 0
|
|
|
|
"""
|
|
LLM model class for inference, pretraining, and finetuning.
|
|
|
|
Example:
|
|
from litgpt.api import LLM
|
|
|
|
llm = LLM.load("microsoft/phi-2")
|
|
text = llm.generate("What do Llamas eat?", top_k=1)
|
|
print(text)
|
|
"""
|
|
|
|
@property
|
|
def tokenizer(self):
|
|
return self.preprocessor.tokenizer
|
|
|
|
def state_dict(self, destination=None, prefix="", keep_vars=False):
|
|
return self.model.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
|
|
|
|
def load_state_dict(self, state_dict, strict=True):
|
|
return self.model.load_state_dict(state_dict, strict=strict)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
target_ids: torch.Tensor | None = None,
|
|
loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
|
logits = self.model(input_ids)
|
|
if target_ids is not None:
|
|
if loss_fn is None:
|
|
loss_fn = chunked_cross_entropy
|
|
loss = loss_fn(logits[..., :-1, :], target_ids[..., 1:])
|
|
return logits, loss
|
|
else:
|
|
return logits
|
|
|
|
def trainer_setup(self, trainer_ckpt: Path | None = None) -> None:
|
|
"""Initializes the model checkpoint for PyTorch Lightning Trainer contexts"""
|
|
self.model = GPT(self.config)
|
|
|
|
if trainer_ckpt is not None:
|
|
# strip the object name key from the state_dict
|
|
state_dict = torch.load(trainer_ckpt, weights_only=True)["state_dict"]
|
|
first_key = next(iter(state_dict))
|
|
prefix = first_key.split(".")[0] + "."
|
|
keys_to_modify = [key for key in state_dict if key.startswith(prefix)]
|
|
for key in keys_to_modify:
|
|
new_key = key.replace(prefix, "", 1)
|
|
state_dict[new_key] = state_dict.pop(key)
|
|
|
|
self.load_state_dict(state_dict, strict=True)
|
|
|
|
elif self.checkpoint_dir is not None:
|
|
state_dict = torch.load(self.checkpoint_dir / "lit_model.pth", weights_only=False)
|
|
self.load_state_dict(state_dict, strict=False)
|
|
|
|
else:
|
|
raise ValueError(
|
|
"No checkpoint found. Either provide a valid path via `trainer_ckpt` "
|
|
"or ensure that `self.checkpoint_dir` points to a folder containing a `lit_model.pth` weight file."
|
|
)
|
|
|
|
def save(self, out_dir: Path | None = None, prompt_style: PromptStyle | None = None) -> None:
|
|
out_dir = Path(out_dir)
|
|
save_path = out_dir / "lit_model.pth"
|
|
save_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if prompt_style is None:
|
|
prompt_style = PromptStyle.from_config(self.config)
|
|
if self.fabric is None:
|
|
torch.save(self.state_dict(), save_path)
|
|
else:
|
|
self.fabric.save(save_path, self.state_dict())
|
|
|
|
if self.fabric is None or self.fabric.global_rank == 0:
|
|
# If initialization a model with random weights, the checkpoint dir can be none
|
|
if self.checkpoint_dir is not None:
|
|
copy_config_files(Path(self.checkpoint_dir), save_path.parent)
|
|
else:
|
|
save_config(self.config, out_dir)
|
|
|
|
save_prompt_style(prompt_style, save_path.parent)
|
|
|
|
@classmethod
|
|
def load(
|
|
cls,
|
|
model: str,
|
|
init: Literal["pretrained", "random"] | None = "pretrained",
|
|
tokenizer_dir: Path | None = None,
|
|
access_token: str | None = None,
|
|
distribute: Literal["auto"] | None = "auto",
|
|
) -> "LLM":
|
|
"""
|
|
Loads the LLM from a local directory or model hub.
|
|
|
|
Arguments
|
|
model: A local path to a directory containing the model weights or a valid model name.
|
|
You can get a list of valid model names via the `litgpt download list` command line argument.
|
|
init: If "pretrained" (default), downloads the model from the HF Hub if a local model can't be found at the `model`
|
|
directory name; otherwise loads the model from the local directory.
|
|
If "random", initializes the `model` with random weights.
|
|
tokenizer_dir: An optional tokenizer directory if `model` is not a checkpoint directory, or if a user
|
|
wants to use a different tokenizer instead.
|
|
access_token: Optional API token to access models with restrictions when using `init="pretrained"`.
|
|
distribute: If "auto" (default), initializes the model on a single GPU if available and otherwise on the CPU.
|
|
To have more control over the model distribution strategy and utilize multiple GPUs, you can set
|
|
`llm = LLM.load(..., distribute=None)` and call `llm.distribute(...)` manually.
|
|
"""
|
|
|
|
allowed_init = {"pretrained", "random"}
|
|
|
|
if init == "pretrained":
|
|
checkpoint_dir = auto_download_checkpoint(
|
|
model_name=model, access_token=access_token, ignore_tokenizer_files=tokenizer_dir is not None
|
|
)
|
|
config = Config.from_file(checkpoint_dir / "model_config.yaml")
|
|
|
|
elif init == "random":
|
|
checkpoint_dir = None
|
|
try:
|
|
config = Config.from_name(model)
|
|
except ValueError:
|
|
print(f"Model name {model} is not supported.\n")
|
|
available_models = "\n".join(sorted(name_to_config))
|
|
print(f"Available values:\n{available_models}")
|
|
return
|
|
|
|
else:
|
|
raise ValueError(f"Invalid init option: {init}. Must be one of {allowed_init}")
|
|
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
if tokenizer_dir is not None:
|
|
tokenizer_dir = extend_checkpoint_dir(Path(tokenizer_dir))
|
|
tokenizer = Tokenizer(tokenizer_dir)
|
|
elif checkpoint_dir is not None:
|
|
tokenizer = Tokenizer(checkpoint_dir)
|
|
else:
|
|
raise ValueError("Provide a path to a tokenizer directory via the `tokenizer_dir` setting.")
|
|
|
|
if checkpoint_dir is not None:
|
|
prompt_style = (
|
|
load_prompt_style(checkpoint_dir)
|
|
if has_prompt_style(checkpoint_dir)
|
|
else PromptStyle.from_config(config)
|
|
)
|
|
else:
|
|
prompt_style = PromptStyle.from_config(config)
|
|
|
|
if distribute == "auto":
|
|
if torch.cuda.is_available():
|
|
accelerator = "cuda"
|
|
elif torch.backends.mps.is_available():
|
|
accelerator = "mps"
|
|
else:
|
|
accelerator = "cpu"
|
|
|
|
fabric = L.Fabric(
|
|
accelerator=accelerator,
|
|
devices=1,
|
|
precision=get_default_supported_precision(training=False),
|
|
)
|
|
|
|
with fabric.init_module(empty_init=False):
|
|
model = GPT(config)
|
|
model.eval()
|
|
preprocessor = Preprocessor(tokenizer, device=fabric.device)
|
|
|
|
if checkpoint_dir is not None:
|
|
checkpoint_path = checkpoint_dir / "lit_model.pth"
|
|
check_file_size_on_cpu_and_warn(checkpoint_path, fabric.device)
|
|
load_checkpoint(fabric, model, checkpoint_path)
|
|
|
|
model = fabric.setup_module(model)
|
|
|
|
else:
|
|
preprocessor = Preprocessor(tokenizer, device="cuda" if torch.cuda.is_available() else "cpu")
|
|
model = None
|
|
fabric = None
|
|
|
|
return cls(
|
|
model=model,
|
|
preprocessor=preprocessor,
|
|
prompt_style=prompt_style,
|
|
config=config,
|
|
checkpoint_dir=checkpoint_dir,
|
|
fabric=fabric,
|
|
generate_strategy=None,
|
|
kv_cache_initialized=False,
|
|
fixed_kv_cache_size=False,
|
|
)
|
|
|
|
def distribute(
|
|
self,
|
|
accelerator: Literal["cpu", "cuda", "auto"] = "auto",
|
|
devices: int | Literal["auto"] = "auto",
|
|
precision: Any | None = None,
|
|
quantize: Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"] | None = None,
|
|
generate_strategy: Literal["sequential", "tensor_parallel"] | None = None,
|
|
fixed_kv_cache_size: int | Literal["max_model_supported"] | None = None,
|
|
) -> None:
|
|
"""
|
|
Moves the model onto specified devices for single-GPU or multi-GPU inference
|
|
|
|
accelerator: Which device type to load the model on ("cpu", "gpu", "mps", "cuda", or "auto")
|
|
devices: The number of devices (1, 2, etc.) or "auto", which uses all available devices
|
|
quantize: Whether to quantize the model and using which method:
|
|
- bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes
|
|
- bnb.int8: 8-bit quantization from bitsandbytes
|
|
for more details, see https://github.com/Lightning-AI/litgpt/blob/main/tutorials/quantize.md
|
|
precision: Indicates the Fabric precision setting to use.
|
|
For instance, "32-true", "16-mixed", "16-true", "bf16-mixed", "bf16-true".
|
|
For more details, see https://lightning.ai/docs/fabric/stable/api/fabric_args.html#precision
|
|
generate_strategy: Whether to use a sequential model generation strategy. The "sequential" settings allows running
|
|
models that wouldn't fit in a single card by partitioning the transformer blocks across
|
|
all devices and running them sequentially. Sequential generation may be slower but allows using larger models.
|
|
Note that sequential generation sets `fixed_kv_cache_size="max_model_supported"`. You can set it to a lower integer
|
|
value, `fixed_kv_cache_size=256` to reduce memory. The `fixed_kv_cache_size` value determines the maximum number
|
|
of tokens that can be returned via `llm.generate(...)`.
|
|
fixed_kv_cache_size: If set to an integer value or "max_model_supported" is set, the kv-cache won't be resized dynamically
|
|
during `llm.generate` calls. Use this setting if you plan to compile the model or use `generate_strategy="sequential`.
|
|
Note that the chosen `fixed_kv_cache_size` value determines the maximum number of tokens that can be returned in `llm.generate(...)`.
|
|
"""
|
|
|
|
if self.checkpoint_dir is None:
|
|
raise NotImplementedError(
|
|
"The LLM was initialized with init='random' but .distribute() "
|
|
"currently only supports pretrained weights."
|
|
)
|
|
|
|
allowed_accelerators = {"cpu", "gpu", "cuda", "mps", "auto"}
|
|
if accelerator not in allowed_accelerators:
|
|
raise ValueError(f"Invalid accelerator: {accelerator}. Must be one of {allowed_accelerators}.")
|
|
|
|
if accelerator == "auto":
|
|
if torch.cuda.is_available():
|
|
accelerator = "cuda"
|
|
elif torch.backends.mps.is_available():
|
|
accelerator = "mps"
|
|
else:
|
|
accelerator = "cpu"
|
|
|
|
if generate_strategy in ("sequential", "tensor_parallel") and accelerator not in ("cuda", "gpu"):
|
|
raise NotImplementedError(
|
|
f"generate_strategy='{generate_strategy}' is only supported for accelerator='cuda'|'gpu'."
|
|
)
|
|
|
|
if devices == "auto":
|
|
if generate_strategy in ("sequential", "tensor_parallel"):
|
|
total_devices = CUDAAccelerator.auto_device_count()
|
|
else:
|
|
total_devices = 1
|
|
elif isinstance(devices, int) and accelerator == "cuda":
|
|
use_devices = calculate_number_of_devices(devices)
|
|
total_devices = CUDAAccelerator.auto_device_count()
|
|
if use_devices > total_devices:
|
|
raise ValueError(
|
|
f"You selected more devices ({use_devices}) than available in your system ({total_devices})."
|
|
)
|
|
else:
|
|
total_devices = use_devices
|
|
|
|
if total_devices > 1 and generate_strategy not in ("sequential", "tensor_parallel"):
|
|
raise NotImplementedError(
|
|
"Support for multiple devices is currently only implemented for generate_strategy='sequential'|'tensor_parallel'."
|
|
)
|
|
elif accelerator == "cpu" or accelerator == "mps":
|
|
total_devices = 1
|
|
|
|
else:
|
|
raise ValueError(f"devices argument must be an integer or 'auto', got {devices}")
|
|
|
|
print(f"Using {total_devices} device(s)", file=sys.stderr)
|
|
|
|
if precision is None:
|
|
precision = get_default_supported_precision(training=False)
|
|
|
|
print("Precision set", file=sys.stderr)
|
|
|
|
plugins = None
|
|
if quantize is not None and quantize.startswith("bnb."):
|
|
if "mixed" in precision:
|
|
raise ValueError("The combination of quantization and mixed precision is not supported.")
|
|
dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision]
|
|
plugins = BitsandbytesPrecision(quantize[4:], dtype)
|
|
precision = None
|
|
|
|
# set "ddp" as the strategy for the launching functionality, but there's no data-parallelism
|
|
if generate_strategy != "tensor_parallel":
|
|
fabric = L.Fabric(
|
|
accelerator=accelerator,
|
|
devices=1, # Otherwise sequential wouldn't work, see litgpt/generate/sequentially.py
|
|
# devices=devices,
|
|
precision=precision,
|
|
plugins=plugins,
|
|
)
|
|
else:
|
|
fabric = L.Fabric(
|
|
accelerator=accelerator, devices=total_devices, strategy="ddp", precision=precision, plugins=plugins
|
|
)
|
|
if torch.cuda.is_available() and fabric.accelerator.auto_device_count() > 1:
|
|
check_nvlink_connectivity(fabric)
|
|
fabric.launch()
|
|
|
|
print("Fabric launched", file=sys.stderr)
|
|
|
|
self.kv_cache_initialized = False
|
|
if generate_strategy is None:
|
|
with fabric.init_module(empty_init=(total_devices > 1)):
|
|
model = GPT(self.config)
|
|
model.eval()
|
|
|
|
if self.checkpoint_dir is not None:
|
|
load_checkpoint(fabric, model, self.checkpoint_dir / "lit_model.pth")
|
|
|
|
model = fabric.setup_module(model)
|
|
|
|
if fixed_kv_cache_size is not None:
|
|
if fixed_kv_cache_size is None or fixed_kv_cache_size == "max_model_supported":
|
|
kv_cache_size = model.max_seq_length
|
|
else:
|
|
kv_cache_size = fixed_kv_cache_size
|
|
model.set_kv_cache(batch_size=1, max_seq_length=kv_cache_size, device=fabric.device)
|
|
self.kv_cache_initialized = True
|
|
self.fixed_kv_cache_size = fixed_kv_cache_size
|
|
|
|
elif generate_strategy in ("sequential", "tensor_parallel"):
|
|
with fabric.init_tensor(), torch.device("meta"):
|
|
model = GPT(self.config)
|
|
model.eval()
|
|
|
|
if generate_strategy == "sequential":
|
|
state_dict = torch.load(
|
|
str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu", weights_only=False
|
|
)
|
|
model.load_state_dict(state_dict, assign=True)
|
|
model = fabric.setup_module(model, move_to_device=False)
|
|
|
|
if fixed_kv_cache_size is None:
|
|
fixed_kv_cache_size = "max_model_supported"
|
|
if fixed_kv_cache_size == "max_model_supported":
|
|
kv_cache_size = model.max_seq_length
|
|
else:
|
|
kv_cache_size = fixed_kv_cache_size
|
|
|
|
model = sequential(model, fabric.device, kv_cache_size, total_devices)
|
|
self.fixed_kv_cache_size = fixed_kv_cache_size
|
|
|
|
elif generate_strategy == "tensor_parallel":
|
|
if fabric.global_rank == 0:
|
|
pbar = tqdm(total=fabric.world_size, desc="Loading model weights")
|
|
for rank in range(fabric.world_size):
|
|
if fabric.global_rank == rank:
|
|
state_dict = torch.load(
|
|
str(self.checkpoint_dir / "lit_model.pth"),
|
|
mmap=True,
|
|
map_location="cpu",
|
|
weights_only=False,
|
|
)
|
|
model.load_state_dict(state_dict, assign=True)
|
|
|
|
# cannot use `.setup_module` because it will wrap with DDP
|
|
model = fabric._precision.convert_module(model)
|
|
model = tensor_parallel(fabric, model)
|
|
|
|
with fabric.init_tensor():
|
|
if fixed_kv_cache_size is None:
|
|
fixed_kv_cache_size = "max_model_supported"
|
|
if fixed_kv_cache_size == "max_model_supported":
|
|
kv_cache_size = model.max_seq_length
|
|
else:
|
|
kv_cache_size = fixed_kv_cache_size
|
|
model.max_seq_length = kv_cache_size
|
|
# the rope cache which is on meta device
|
|
model.cos, model.sin = model.rope_cache()
|
|
# enable the kv cache
|
|
model.set_kv_cache(batch_size=1)
|
|
model.eval()
|
|
model = fabric.to_device(model)
|
|
|
|
fabric.barrier()
|
|
if fabric.global_rank == 0:
|
|
pbar.update(1)
|
|
|
|
if fabric.global_rank == 0:
|
|
pbar.close()
|
|
|
|
self.kv_cache_initialized = True
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported generate_strategy: {generate_strategy}")
|
|
|
|
self.model = model
|
|
self.fabric = fabric
|
|
self.preprocessor.device = fabric.device
|
|
|
|
@torch.inference_mode()
|
|
def generate(
|
|
self,
|
|
prompt: str,
|
|
sys_prompt: str | None = None,
|
|
max_new_tokens: int = 50,
|
|
temperature: float = 1.0,
|
|
top_k: int | None = None,
|
|
top_p: float = 1.0,
|
|
return_as_token_ids: bool = False,
|
|
stream: bool = False,
|
|
) -> str | torch.Tensor:
|
|
"""
|
|
Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
|
|
|
|
Arguments:
|
|
model: The model to use.
|
|
prompt: The prompt string to use for generating the samples.
|
|
sys_prompt: The system prompt string to use for generating the samples.
|
|
The system prompt allows the user to provide additional instructions to shape all responses by providing additional context, behavioral guidelines, style, and constraints.
|
|
max_new_tokens: The maximum number of new tokens to return.
|
|
temperature: Scales the predicted logits by 1 / temperature.
|
|
top_k: If specified, only sample among the tokens with the k highest probabilities.
|
|
top_p: If specified, it represents the cumulative probability threshold to consider in the sampling process.
|
|
In top-p sampling, the next token is sampled from the highest probability tokens
|
|
whose cumulative probability exceeds the threshold `top_p`. When specified,
|
|
it must be `0 <= top_p <= 1`. Here, `top_p=0` is equivalent
|
|
to sampling the most probable token, while `top_p=1` samples from the whole distribution.
|
|
It can be used in conjunction with `top_k` and `temperature` with the following order
|
|
of application:
|
|
|
|
1. `top_k` sampling
|
|
2. `temperature` scaling
|
|
3. `top_p` sampling
|
|
|
|
For more details, see https://arxiv.org/abs/1904.09751
|
|
or https://huyenchip.com/2024/01/16/sampling.html#top_p
|
|
return_as_token_ids: If True, returns the token IDs as a torch.Tensor. Otherwise, returns the decoded text as a string.
|
|
stream: If True, returns a generator that yields tokens as they are generated.
|
|
At the moment, this setting is slower and may use more memory than the non-streaming version.
|
|
We plan to resolve this in the future.
|
|
"""
|
|
if self.model is None:
|
|
raise AttributeError(
|
|
"The model is not initialized yet; use the .distribute() "
|
|
"or .trainer_setup() method to initialize the model."
|
|
)
|
|
input_ids = self._text_to_token_ids(prompt, sys_prompt)
|
|
prompt_length = input_ids.size(0)
|
|
max_returned_tokens = prompt_length + max_new_tokens
|
|
|
|
if not self.kv_cache_initialized:
|
|
if self.fabric is not None:
|
|
device = self.fabric.device
|
|
else:
|
|
device = self.preprocessor.device
|
|
self.model.set_kv_cache(batch_size=1, max_seq_length=max_returned_tokens, device=device)
|
|
self.kv_cache_initialized = True
|
|
|
|
# Dynamically grow the kv cache size if necessary
|
|
if not self.fixed_kv_cache_size and self.prev_generated_seq_length < max_returned_tokens:
|
|
tmp_device = self.model.mask_cache.device
|
|
self.model.clear_kv_cache()
|
|
self.model.set_kv_cache(batch_size=1, max_seq_length=max_returned_tokens, device=tmp_device)
|
|
|
|
else:
|
|
for block in self.model.transformer.h:
|
|
block.attn.kv_cache.reset_parameters()
|
|
|
|
self.prev_generated_seq_length = max_returned_tokens
|
|
self.model.eval()
|
|
|
|
def iterator():
|
|
outputs = stream_generate_fn(
|
|
model=self.model,
|
|
prompt=input_ids,
|
|
max_returned_tokens=max_returned_tokens,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
stop_tokens=([self.preprocessor.tokenizer.eos_id],),
|
|
)
|
|
if return_as_token_ids:
|
|
yield from outputs
|
|
else:
|
|
for output in outputs:
|
|
yield self.preprocessor.decode(output)
|
|
return
|
|
|
|
if stream:
|
|
outputs = iterator()
|
|
else:
|
|
outputs = generate_fn(
|
|
model=self.model,
|
|
prompt=input_ids,
|
|
max_returned_tokens=max_returned_tokens,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
eos_id=self.preprocessor.tokenizer.eos_id,
|
|
include_prompt=False,
|
|
)
|
|
|
|
if stream:
|
|
return outputs
|
|
elif return_as_token_ids:
|
|
return outputs
|
|
else:
|
|
return self.preprocessor.decode(outputs)
|
|
|
|
def _text_to_token_ids(self, prompt: str, sys_prompt: str | None = None) -> torch.Tensor:
|
|
"""Utility method to convert a prompt text to token IDs"""
|
|
prompt = self.prompt_style.apply(prompt, sys_prompt=sys_prompt)
|
|
input_ids = self.preprocessor.encode(prompt)
|
|
return input_ids
|
|
|
|
def benchmark(self, num_iterations=1, **kwargs):
|
|
"""
|
|
A wrapper around the .generate() method to calculate runtime performance.
|
|
|
|
Arguments:
|
|
num_iterations: How often the `.generate()` call is repeated.
|
|
kwargs: Keyword arguments that are passed to the .generate() method.
|
|
"""
|
|
benchmark_dict = {}
|
|
|
|
for i in range(num_iterations):
|
|
time_to_first_token = None
|
|
t0 = time.perf_counter()
|
|
outputs = self.generate(**kwargs)
|
|
|
|
if kwargs.get("stream", False):
|
|
gen_outputs = []
|
|
for e in outputs:
|
|
if time_to_first_token is None:
|
|
t1 = time.perf_counter()
|
|
time_to_first_token = t1 - t0
|
|
gen_outputs.append(e)
|
|
outputs = "".join(gen_outputs)
|
|
else:
|
|
outputs = self.generate(
|
|
**kwargs,
|
|
)
|
|
benchmark_dict.setdefault("Seconds total", []).append(time.perf_counter() - t0)
|
|
|
|
benchmark_dict.setdefault("Seconds to first token", []).append(time_to_first_token)
|
|
tokens_generated = self.preprocessor.encode(outputs).size(0)
|
|
benchmark_dict.setdefault("Tokens generated", []).append(tokens_generated)
|
|
benchmark_dict.setdefault("Inference speed in tokens/sec", []).append(
|
|
benchmark_dict["Tokens generated"][-1] / benchmark_dict["Seconds total"][-1]
|
|
)
|
|
if self.fabric is not None and self.fabric.device.type == "cuda":
|
|
benchmark_dict.setdefault("Total GPU memory allocated in GB", []).append(
|
|
torch.cuda.max_memory_allocated() / 1e9
|
|
)
|
|
|
|
return outputs, benchmark_dict
|
|
|
|
|
|
class Preprocessor:
|
|
"""
|
|
Preprocessor class for tokenization and de-tokenization.
|
|
"""
|
|
|
|
def __init__(self, tokenizer: Tokenizer, device: str = "cpu") -> None:
|
|
self.tokenizer = tokenizer
|
|
self.device = device
|
|
|
|
def encode(self, text: str) -> torch.Tensor:
|
|
return self.tokenizer.encode(text, device=self.device)
|
|
|
|
def decode(self, token_ids: torch.Tensor) -> str:
|
|
return self.tokenizer.decode(token_ids)
|
|
|
|
|
|
def calculate_number_of_devices(devices):
|
|
"""
|
|
Utility function to calculate the number of devices.
|
|
"""
|
|
num_devices = devices if isinstance(devices, int) else len(devices) if isinstance(devices, list) else 0
|
|
return num_devices
|
|
|
|
|
|
def benchmark_dict_to_markdown_table(data):
|
|
"""
|
|
Converts .benchmark() outputs to a markdown table
|
|
"""
|
|
markdown_table = (
|
|
"| Metric | Mean | Std Dev |\n"
|
|
)
|
|
markdown_table += (
|
|
"|-------------------------------------|-----------------------------|-----------------------------|\n"
|
|
)
|
|
|
|
for key, values in data.items():
|
|
mean_value = np.mean(values)
|
|
std_dev_value = np.std(values, ddof=1)
|
|
|
|
formatted_mean = f"{mean_value:.2f}"
|
|
formatted_std_dev = f"{std_dev_value:.2f}"
|
|
|
|
markdown_table += f"| {key.ljust(35)} | {formatted_mean.ljust(27)} | {formatted_std_dev.ljust(27)} |\n"
|
|
|
|
return markdown_table
|
|
|
|
|
|
def pull_request_benchmark_util(model_name="microsoft/phi-2", num_iterations=6):
|
|
def print_table(header, data):
|
|
print(f"\n### {header}\n")
|
|
markdown_table = (
|
|
f"| Metric | First Iteration | "
|
|
f"Iter 2-{num_iterations} Mean | Iter 2-{num_iterations} Standard Dev. |\n"
|
|
f"|--------------------------------------|-----------------|"
|
|
f"-------------------|-------------------------|\n"
|
|
)
|
|
|
|
for key, value in data.items():
|
|
first_iteration = f"{value[0]:.2f}" if value[0] is not None else "N/A"
|
|
clean_values = [v for v in value[1:] if v is not None]
|
|
|
|
if clean_values:
|
|
mean_value = np.mean(clean_values)
|
|
std_dev_value = np.std(clean_values, ddof=1)
|
|
mean_str = f"{mean_value:.2f}"
|
|
std_dev_str = f"{std_dev_value:.2f}"
|
|
else:
|
|
mean_str = "N/A"
|
|
std_dev_str = "N/A"
|
|
|
|
markdown_table += f"| {key:<36} | {first_iteration:<15} | {mean_str:<17} | {std_dev_str:<23} |\n"
|
|
print(markdown_table)
|
|
|
|
import subprocess
|
|
|
|
try:
|
|
g_hash = subprocess.run(
|
|
["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, check=True
|
|
).stdout.strip()
|
|
print(f"Git Commit Hash: {g_hash}")
|
|
except subprocess.CalledProcessError:
|
|
print("Git Commit Hash: N/A")
|
|
print(f"PyTorch version: {torch.__version__}")
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}\n")
|
|
|
|
# 1st table
|
|
llm = LLM.load(
|
|
model=model_name,
|
|
)
|
|
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1)
|
|
print_table(f"Defaults ({model_name}), 1st time", bench_d)
|
|
del llm
|
|
|
|
# 2nd table
|
|
llm = LLM.load(
|
|
model=model_name,
|
|
)
|
|
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1)
|
|
print_table(f"Defaults ({model_name}), 2nd time", bench_d)
|
|
del llm
|
|
|
|
# 3rd table
|
|
llm = LLM.load(
|
|
model=model_name,
|
|
)
|
|
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1, stream=True)
|
|
print_table("stream=True", bench_d)
|
|
del llm
|
|
|
|
# 4th table
|
|
llm = LLM.load(model=model_name, distribute=None)
|
|
llm.distribute(fixed_kv_cache_size=500)
|
|
|
|
text, bench_d = llm.benchmark(num_iterations=num_iterations, prompt="What do llamas eat?", top_k=1, stream=True)
|
|
print_table("stream=True + fixed_kv_cache=500", bench_d)
|