# 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))