chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,622 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# ruff: noqa: F401
|
||||
|
||||
"""
|
||||
.. _opt_llm:
|
||||
|
||||
Optimize Large Language Model
|
||||
=============================
|
||||
As large language models (LLMs) have become a popular research topic in many different fields,
|
||||
deploying them on cloud and edge devices has become a challenging task. In this tutorial, we will
|
||||
demonstrate how to optimize a large language model using Apache TVM. We will use a pre-trained
|
||||
TinyLlama model from Hugging Face and deploy it on various devices.
|
||||
"""
|
||||
|
||||
######################################################################
|
||||
# Review Overall Flow
|
||||
# -------------------
|
||||
# .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
|
||||
# :align: center
|
||||
# :width: 80%
|
||||
#
|
||||
# The overall flow consists of the following steps:
|
||||
#
|
||||
# - **Construct or Import a Model**: Construct a neural network model or import a pre-trained
|
||||
# model from other frameworks (e.g. PyTorch, ONNX), and create the TVM IRModule, which contains
|
||||
# all the information needed for compilation, including high-level Relax functions for
|
||||
# computational graph, and low-level TensorIR functions for tensor program.
|
||||
# - **Perform Composable Optimizations**: Perform a series of optimization transformations,
|
||||
# such as graph optimizations, tensor program optimizations, and library dispatching.
|
||||
# - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
|
||||
# universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
|
||||
#
|
||||
|
||||
|
||||
######################################################################
|
||||
# Construct the model architecture
|
||||
# --------------------------------
|
||||
# We will use a pre-trained TinyLlama model from Hugging Face. However, usually we only load the
|
||||
# pre-trained weight from Hugging Face but not the model architecture. We need to construct the
|
||||
# model architecture by ourselves. Apache TVM prepares a PyTorch-liked API to construct the model
|
||||
# architecture. We can use the API to construct the model architecture.
|
||||
|
||||
import dataclasses
|
||||
import enum
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pprint import pprint
|
||||
|
||||
from tvm_ffi import Shape
|
||||
|
||||
import tvm
|
||||
from tvm import relax, te, tirx
|
||||
from tvm.relax import register_pipeline
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.relax.frontend.nn import Tensor, op
|
||||
from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache, TIRPagedKVCache
|
||||
from tvm.s_tir import dlight
|
||||
|
||||
######################################################################
|
||||
# First, we need to define the model configuration. The configuration includes the key parameters
|
||||
# of the model, such as hidden size, intermediate size, etc. Here for convenience, we define a
|
||||
# constant config specially for the TinyLlama model.
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LlamaConfig:
|
||||
hidden_size: int = 2048
|
||||
intermediate_size: int = 5632
|
||||
num_attention_heads: int = 32
|
||||
num_hidden_layers: int = 22
|
||||
rms_norm_eps: float = 1e-05
|
||||
vocab_size: int = 32000
|
||||
rope_theta: int = 10000
|
||||
context_window_size: int = 2048
|
||||
prefill_chunk_size: int = 2048
|
||||
num_key_value_heads: int = 4
|
||||
head_dim: int = 64 # hidden_size // num_attention_heads
|
||||
|
||||
|
||||
dev = tvm.device("cuda", 0)
|
||||
target = tvm.target.Target.from_device(dev)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Next, we define the RoPE mode of the Paged KV cache. The RoPE mode is used to apply the
|
||||
# Relative Positional Encoding (RoPE) to the query and key tensors. The RoPE mode can be set to
|
||||
# `NONE`, `NORMAL`, or `INLINE`. If the RoPE mode is `NONE`, the KV cache will not apply RoPE to
|
||||
# the query and key tensors. If the RoPE mode is `NORMAL`, RoPE will be applied to the key tensor
|
||||
# before adding the key tensor to the cache. If the RoPE mode is `INLINE`, RoPE will be applied to
|
||||
# the query and key tensors in the attention kernel on-the-fly.
|
||||
|
||||
|
||||
class RopeMode(enum.IntEnum):
|
||||
"""The RoPE mode of the Paged KV cache.
|
||||
If it is none, the KV cache will not apply RoPE to q and k.
|
||||
If it is normal, RoPE will be applied to k before adding k to cache.
|
||||
Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
|
||||
"""
|
||||
|
||||
NONE = 0
|
||||
NORMAL = 1
|
||||
INLINE = 2
|
||||
|
||||
|
||||
######################################################################
|
||||
# Secondly, we define the model architecture. The model architecture consists of three parts:
|
||||
#
|
||||
# - Embedding layer: The embedding layer converts the input token IDs to the hidden states.
|
||||
# - Decoder layers: The decoder layers are the core of the model. Each decoder layer consists of
|
||||
# a self-attention layer and a feed-forward network (FFN) layer.
|
||||
# - Output layer: The output layer converts the hidden states to the logits.
|
||||
#
|
||||
# First we define the FFN layer. Note that the following FFN layer is optimized implementation
|
||||
# where we fuse the gate and up projection into one kernel.
|
||||
# The naive implementation of FFN layer is: ``FFN(x) = down_proj(silu(gate(x)) * up(x))``
|
||||
# We could combine the ``gate`` and ``up`` projection into one kernel for better performance.
|
||||
# The optimized implementation is:
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# concat_x = gate_up(x)
|
||||
# gate_x, up_x = split(concat_x, 2, axis=-1)
|
||||
# FFN(x) = down_proj(silu(gate_x) * up_x)
|
||||
#
|
||||
|
||||
|
||||
class LlamaFFN(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
super().__init__()
|
||||
self.gate_up_proj = nn.Linear(
|
||||
in_features=config.hidden_size,
|
||||
out_features=2 * config.intermediate_size,
|
||||
bias=False,
|
||||
)
|
||||
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
concat_x1_x2 = self.gate_up_proj(x)
|
||||
x1, x2 = op.split(concat_x1_x2, 2, axis=-1)
|
||||
return self.down_proj(op.silu(x1) * x2)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Then we define the self-attention layer. The self-attention layer consists of three parts:
|
||||
#
|
||||
# - QKV projection: The QKV projection converts the input hidden states to the query, key, and
|
||||
# value tensors.
|
||||
# - Attention: The attention layer computes the attention scores and applies the softmax
|
||||
# operation.
|
||||
# - Output projection: The output projection converts the attention output to the hidden states.
|
||||
#
|
||||
# We perform optimizations on the different parts of the self-attention layer:
|
||||
#
|
||||
# - QKV projection: We leverage the horizontal fusion on QKV projection and fuse them into one
|
||||
# kernel.
|
||||
# - Attention: We leverage the horizontal fusion on attention and fuse the QKV projection and
|
||||
|
||||
|
||||
class LlamaAttention(nn.Module): # pylint: disable=too-many-instance-attributes
|
||||
def __init__(self, config: LlamaConfig):
|
||||
self.head_dim = config.head_dim
|
||||
self.num_q_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
# horizontal fusion on QKV projection
|
||||
self.qkv_proj = nn.Linear(
|
||||
in_features=config.hidden_size,
|
||||
out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
|
||||
d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads
|
||||
b, s, _ = hidden_states.shape
|
||||
# QKV Projection
|
||||
qkv = self.qkv_proj(hidden_states)
|
||||
qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d))
|
||||
# Attention
|
||||
output = op.reshape(
|
||||
paged_kv_cache.attention_with_fused_qkv(
|
||||
layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5
|
||||
),
|
||||
(b, s, h_q * d),
|
||||
)
|
||||
# Output Projection
|
||||
return self.o_proj(output)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Finally, we define the model architecture with FFN and self-attention layers.
|
||||
|
||||
|
||||
class LlamaDecoderLayer(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
rms_norm_eps = config.rms_norm_eps
|
||||
self.self_attn = LlamaAttention(config)
|
||||
self.mlp = LlamaFFN(config)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
||||
self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False)
|
||||
|
||||
def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int):
|
||||
hidden_states += self.self_attn(
|
||||
self.input_layernorm(hidden_states), paged_kv_cache, layer_id
|
||||
)
|
||||
hidden_states += self.mlp(self.post_attention_layernorm(hidden_states))
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LlamaModel(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
assert config.hidden_size % config.num_attention_heads == 0
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.layers = nn.ModuleList(
|
||||
[LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)]
|
||||
)
|
||||
self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False)
|
||||
|
||||
def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
hidden_states = input_embed
|
||||
for layer_id, layer in enumerate(self.layers):
|
||||
hidden_states = layer(hidden_states, paged_kv_cache, layer_id)
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class LlamaForCausalLM(nn.Module):
|
||||
def __init__(self, config: LlamaConfig):
|
||||
self.model = LlamaModel(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
self.num_hidden_layers = config.num_hidden_layers
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.num_key_value_heads = config.num_key_value_heads
|
||||
self.head_dim = config.head_dim
|
||||
self.hidden_size = config.hidden_size
|
||||
self.vocab_size = config.vocab_size
|
||||
self.rope_theta = config.rope_theta
|
||||
self.dtype = "float32"
|
||||
|
||||
def to(self, dtype: str | None = None):
|
||||
super().to(dtype=dtype)
|
||||
if dtype is not None:
|
||||
self.dtype = dtype
|
||||
|
||||
def embed(self, input_ids: Tensor):
|
||||
return self.model.embed_tokens(input_ids)
|
||||
|
||||
def get_logits(self, hidden_states: Tensor):
|
||||
logits = self.lm_head(hidden_states)
|
||||
if logits.dtype != "float32":
|
||||
logits = logits.astype("float32")
|
||||
return logits
|
||||
|
||||
def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
def _index(x: te.Tensor): # x[:-1,:]
|
||||
b, s, d = x.shape
|
||||
return te.compute((b, 1, d), lambda i, _, k: x[i, s - 1, k], name="index")
|
||||
|
||||
hidden_states = self.model(input_embed, paged_kv_cache)
|
||||
hidden_states = op.tensor_expr_op(_index, name_hint="index", args=[hidden_states])
|
||||
logits = self.get_logits(hidden_states)
|
||||
return logits, paged_kv_cache
|
||||
|
||||
def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache):
|
||||
hidden_states = self.model(input_embed, paged_kv_cache)
|
||||
logits = self.get_logits(hidden_states)
|
||||
return logits, paged_kv_cache
|
||||
|
||||
def create_tir_paged_kv_cache(
|
||||
self,
|
||||
max_batch_size: tirx.Var,
|
||||
max_total_seq_len: tirx.Var,
|
||||
prefill_chunk_size: tirx.Var,
|
||||
page_size: tirx.Var,
|
||||
) -> PagedKVCache:
|
||||
return TIRPagedKVCache(
|
||||
attn_kind="mha",
|
||||
max_batch_size=max_batch_size,
|
||||
max_total_seq_len=max_total_seq_len,
|
||||
prefill_chunk_size=prefill_chunk_size,
|
||||
page_size=page_size,
|
||||
support_sliding_window=0,
|
||||
layer_partition=relax.ShapeExpr([0, self.num_hidden_layers]),
|
||||
num_hidden_layers=self.num_hidden_layers,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
num_key_value_heads=self.num_key_value_heads,
|
||||
qk_head_dim=self.head_dim,
|
||||
v_head_dim=self.head_dim,
|
||||
mla_original_qk_head_dim=0,
|
||||
mla_original_v_head_dim=0,
|
||||
rope_mode=RopeMode.NORMAL,
|
||||
rope_scale=1,
|
||||
rope_theta=self.rope_theta,
|
||||
rope_scaling={},
|
||||
rope_ext_factors=tirx.IntImm("int64", 0),
|
||||
rotary_dim=self.head_dim,
|
||||
dtype=self.dtype,
|
||||
target=target,
|
||||
enable_disaggregation=False,
|
||||
)
|
||||
|
||||
def get_default_spec(self):
|
||||
mod_spec = {
|
||||
"embed": {
|
||||
"input_ids": nn.spec.Tensor(["seq_len"], "int32"),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"prefill": {
|
||||
"input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"decode": {
|
||||
"input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype),
|
||||
"paged_kv_cache": nn.spec.Object(object_type=PagedKVCache),
|
||||
"$": {
|
||||
"param_mode": "packed",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
"create_tir_paged_kv_cache": {
|
||||
"max_batch_size": int,
|
||||
"max_total_seq_len": int,
|
||||
"prefill_chunk_size": int,
|
||||
"page_size": int,
|
||||
"$": {
|
||||
"param_mode": "none",
|
||||
"effect_mode": "none",
|
||||
},
|
||||
},
|
||||
}
|
||||
return nn.spec.ModuleSpec.from_raw(mod_spec, self)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Export the model to Relax IRModule
|
||||
# ----------------------------------
|
||||
# After defining the model architecture, we can export the model to the Relax IRModule.
|
||||
# For demonstration, we only show the part of the model architecture. and parameters.
|
||||
|
||||
model_config = LlamaConfig()
|
||||
model = LlamaForCausalLM(model_config)
|
||||
model.to("float16")
|
||||
mod, named_params = model.export_tvm(spec=model.get_default_spec())
|
||||
prefill_str = mod["prefill"].script()
|
||||
print(*prefill_str.split("\n")[3:20], sep="\n") # Only show the first 10 lines for demonstration
|
||||
print(" ...")
|
||||
|
||||
print("\nParameters:")
|
||||
pprint(named_params[:5]) # Only show the first 5 parameters for demonstration
|
||||
|
||||
######################################################################
|
||||
# Define Optimization Pipeline
|
||||
# ----------------------------
|
||||
# We define a series of optimization passes to optimize the model. The optimization pipeline
|
||||
# is designed specifically for the LLMs.
|
||||
|
||||
|
||||
@register_pipeline("opt_llm")
|
||||
def _pipeline( # pylint: disable=too-many-arguments
|
||||
ext_mods: list[nn.ExternModule] | None = None,
|
||||
):
|
||||
ext_mods = ext_mods or []
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0)
|
||||
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
||||
seq = tvm.transform.Sequential(
|
||||
[
|
||||
# Phase 1. Passes on high-level operator graph
|
||||
# We can enable cublas for further optimization
|
||||
relax.transform.FuseTransposeMatmul(),
|
||||
# Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline
|
||||
relax.transform.LegalizeOps(),
|
||||
relax.transform.AnnotateTIROpPattern(),
|
||||
relax.transform.FoldConstant(),
|
||||
relax.transform.FuseOps(),
|
||||
relax.transform.FuseTIR(),
|
||||
# Phase 3. Passes on TIR
|
||||
relax.transform.DeadCodeElimination(),
|
||||
# Phase 4. Low-level Optimizations
|
||||
dlight.ApplyDefaultSchedule(
|
||||
dlight.gpu.Matmul(),
|
||||
dlight.gpu.GEMV(),
|
||||
dlight.gpu.Reduction(),
|
||||
dlight.gpu.GeneralReduction(),
|
||||
dlight.gpu.Fallback(),
|
||||
),
|
||||
# Phase 5. Lowering to VM bytecode
|
||||
relax.transform.RewriteDataflowReshape(),
|
||||
relax.transform.ToNonDataflow(),
|
||||
relax.transform.RemovePurityChecking(),
|
||||
relax.transform.CallTIRRewrite(),
|
||||
relax.transform.StaticPlanBlockMemory(),
|
||||
relax.transform.RewriteCUDAGraph(),
|
||||
relax.transform.LowerAllocTensor(),
|
||||
relax.transform.KillAfterLastUse(),
|
||||
relax.transform.LowerRuntimeBuiltin(),
|
||||
relax.transform.VMShapeLower(),
|
||||
relax.transform.AttachGlobalSymbol(),
|
||||
relax.transform.AttachExternModules(ext_mods),
|
||||
]
|
||||
)
|
||||
mod = seq(mod)
|
||||
return mod
|
||||
|
||||
return _pipeline
|
||||
|
||||
|
||||
with target:
|
||||
ex = tvm.compile(mod, target, relax_pipeline=relax.get_pipeline("opt_llm"))
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Prepare the model weights
|
||||
# -------------------------
|
||||
# We load the pre-trained weights from Hugging Face and prepare the model weights.
|
||||
# The pre-trained weights are stored in the Hugging Face format. We need to load the weights
|
||||
# and prepare the model parameters.
|
||||
#
|
||||
# .. note::
|
||||
#
|
||||
# Note that we won't execute the following code in this tutorial because the pre-trained weights
|
||||
# are not available in the CI environment.
|
||||
#
|
||||
|
||||
|
||||
IS_IN_CI = os.getenv("CI", "") == "true"
|
||||
|
||||
HF_WEIGHT_PATH = None
|
||||
# HF_WEIGHT_PATH = Path("/path/to/TinyLlama-1.1B-Chat-v1.0/")
|
||||
|
||||
if not IS_IN_CI:
|
||||
import numpy as np
|
||||
import safetensors.torch
|
||||
import torch
|
||||
|
||||
if HF_WEIGHT_PATH is None or not HF_WEIGHT_PATH.exists():
|
||||
raise ValueError("Please set the HF_WEIGHT_PATH to the path of the pre-trained weights.")
|
||||
|
||||
# Torch format weights
|
||||
param_dict = safetensors.torch.load_file(HF_WEIGHT_PATH / "model.safetensors", device="cpu")
|
||||
# Numpy format weights
|
||||
param_dict = {
|
||||
k: v.half().numpy() if v.dtype == torch.bfloat16 else v.numpy()
|
||||
for k, v in param_dict.items()
|
||||
}
|
||||
|
||||
named_params = dict(named_params)
|
||||
for i in range(model_config.num_hidden_layers):
|
||||
# Add QKV in self attention
|
||||
attn = f"model.layers.{i}.self_attn"
|
||||
param_dict[f"{attn}.qkv_proj.weight"] = np.concatenate(
|
||||
[
|
||||
param_dict.pop(f"{attn}.q_proj.weight"), # Pop the old parameters to save memory
|
||||
param_dict.pop(f"{attn}.k_proj.weight"),
|
||||
param_dict.pop(f"{attn}.v_proj.weight"),
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
# Add gates in MLP
|
||||
mlp = f"model.layers.{i}.mlp"
|
||||
param_dict[f"{mlp}.gate_up_proj.weight"] = np.concatenate(
|
||||
[
|
||||
param_dict.pop(f"{mlp}.gate_proj.weight"),
|
||||
param_dict.pop(f"{mlp}.up_proj.weight"),
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
|
||||
# Convert params into ndarray
|
||||
params = [
|
||||
tvm.runtime.tensor(param_dict[k].astype("float16"), device=dev) for k in named_params.keys()
|
||||
]
|
||||
|
||||
|
||||
######################################################################
|
||||
# Deploy the compiled model
|
||||
# -------------------------
|
||||
# After the model and weights are ready, we can deploy the compiled model on the target device.
|
||||
# The language models inference includes two steps: prefill and decode. The prefill step is
|
||||
# used to process the input tokens and store the KVCache. The decode step is used to generate
|
||||
# the token until the end token is generated.
|
||||
|
||||
|
||||
######################################################################
|
||||
# Tokenization
|
||||
# ~~~~~~~~~~~~
|
||||
# The first step is to tokenize the input prompt and embed the tokens into the hidden states.
|
||||
# The tokenization and embedding are the same as the original model. We use the HF tokenizer
|
||||
# to tokenize the input prompt and embed the tokens into the hidden states.
|
||||
# Note that different models require different tokenization and prompt format, please refer to
|
||||
# the model documentation for the correct tokenization and prompt format.
|
||||
|
||||
|
||||
if not IS_IN_CI:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(HF_WEIGHT_PATH)
|
||||
messages = [
|
||||
{"role": "user", "content": "What's your name?"},
|
||||
]
|
||||
prompt = tokenizer.apply_chat_template(messages)
|
||||
input_len = len(prompt)
|
||||
|
||||
# Load prompt tokens into TVM ndarray on the target device
|
||||
tokens = tvm.runtime.tensor(np.array(prompt).astype("int32"), device=dev)
|
||||
|
||||
######################################################################
|
||||
# Create the KVCache
|
||||
# ~~~~~~~~~~~~~~~~~~
|
||||
# Before starting the inference, we need to create the KVCache. The KVCache is used to store the
|
||||
# key and value tensors for the attention layer. Apache TVM provides a PagedKVCache to store the
|
||||
# key and value tensors. We create the PagedKVCache with the specified parameters.
|
||||
|
||||
if not IS_IN_CI:
|
||||
kv_cache = vm["create_tir_paged_kv_cache"](
|
||||
Shape([1]), # max_batch_size=1
|
||||
Shape([2048]), # max_total_seq_len=2048
|
||||
Shape([2048]), # prefill_chunk_size=2048
|
||||
Shape([16]), # page_size=16
|
||||
)
|
||||
|
||||
|
||||
######################################################################
|
||||
# Embedding
|
||||
# ~~~~~~~~~
|
||||
# The next step is to embed the tokens into the hidden states. We use the `embed` function
|
||||
# compiled in the Relax IRModule to embed the tokens into the hidden states.
|
||||
|
||||
nd_view_func = tvm.get_global_func("vm.builtin.reshape")
|
||||
|
||||
|
||||
def embed(tokens, params):
|
||||
_embed = vm["embed"](tokens, params)
|
||||
# Reshape hidden from [seq_len, hidden_size] to [1, seq_len, hidden_size]
|
||||
_embed = nd_view_func(_embed, Shape([1, _embed.shape[0], _embed.shape[1]]))
|
||||
return _embed
|
||||
|
||||
|
||||
######################################################################
|
||||
# Prefill
|
||||
# ~~~~~~~
|
||||
# Before running the forward pass, we first get some help functions for preparation.
|
||||
|
||||
add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
||||
begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
||||
end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
||||
|
||||
######################################################################
|
||||
# As we are creating a new sequence, we need to call `add_sequence_func` to initialize
|
||||
# the request. Additionally, we need to call `begin_forward_func` to start the forward pass,
|
||||
# and `end_forward_func` to end the forward pass.
|
||||
|
||||
if not IS_IN_CI:
|
||||
seq_id = 0
|
||||
add_sequence_func(kv_cache, seq_id)
|
||||
hidden_states = embed(tokens, params)
|
||||
begin_forward_func(kv_cache, Shape([seq_id]), Shape([input_len]))
|
||||
logits, kv_cache = vm["prefill"](hidden_states, kv_cache, params)
|
||||
end_forward_func(kv_cache)
|
||||
|
||||
######################################################################
|
||||
# Now we have the output logits from the prefill step. The logits are used to generate the token
|
||||
# via sampling. Let's sample the token from the logits.
|
||||
#
|
||||
# In this tutorial, we simplify the sampling process and pick the token with the highest
|
||||
# probability. In practice, we should sample the token based on the probability distribution.
|
||||
# Also, to make the tutorial concise, we execute the sample process on CPU.
|
||||
|
||||
|
||||
def sample_token(logits):
|
||||
logits_np = logits.numpy()
|
||||
return np.argmax(logits_np)
|
||||
|
||||
|
||||
if not IS_IN_CI:
|
||||
last_token = sample_token(logits)
|
||||
output_tokens = [last_token]
|
||||
|
||||
|
||||
######################################################################
|
||||
# Decode
|
||||
# ~~~~~~
|
||||
# After the prefill step, we can start the decode step. The decode step is used to generate the
|
||||
# token until the end token is generated. We use the `decode` function compiled in the Relax
|
||||
# IRModule to generate the token.
|
||||
|
||||
if not IS_IN_CI:
|
||||
print("The generated token:")
|
||||
|
||||
while last_token != tokenizer.eos_token_id:
|
||||
tokens = tvm.runtime.tensor(np.array([last_token]).astype("int32"), device=dev)
|
||||
hidden_states = embed(tokens, params)
|
||||
begin_forward_func(kv_cache, Shape([seq_id]), Shape([1]))
|
||||
logits, kv_cache = vm["decode"](hidden_states, kv_cache, params)
|
||||
|
||||
end_forward_func(kv_cache)
|
||||
last_token = sample_token(logits)
|
||||
output_tokens.append(last_token)
|
||||
|
||||
print(tokenizer.decode(output_tokens))
|
||||
Reference in New Issue
Block a user