import argparse import json import os import time from concurrent.futures import ThreadPoolExecutor, wait from typing import Dict, List, Optional, Tuple import numpy as np import safetensors import torch from transformers import AutoModelForCausalLM, FalconConfig, FalconForCausalLM import tensorrt_llm from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.llama.utils import ( # TODO: move the utils to common dir shared by models iterate_shard_files, load_state_dict, retrieved_layer_index_from_name) def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--model_dir', type=str, default=None) parser.add_argument('--world_size', type=int, default=1) parser.add_argument('--tp_size', type=int, default=1) parser.add_argument('--pp_size', type=int, default=1) parser.add_argument('--dtype', type=str, default='float16', choices=['float32', 'bfloat16', 'float16']) parser.add_argument('--logits_dtype', type=str, default='float32', choices=['float16', 'float32']) parser.add_argument( '--ammo_quant_ckpt_path', type=str, default=None, help='Path of a quantized model checkpoint in .npz format') parser.add_argument( '--per_group', default=False, action="store_true", help= 'By default, we use a single static scaling factor to scale weights in the int4 range. ' 'per_group chooses at run time, and for each group, a custom scaling factor. ' 'The flag is built for GPTQ/AWQ quantization.') parser.add_argument('--group_size', type=int, default=128, help='Group size used in GPTQ/AWQ quantization.') parser.add_argument( '--use_weight_only', default=False, action="store_true", help='Quantize weights for the various GEMMs to INT4/INT8.' 'See --weight_only_precision to set the precision') parser.add_argument( '--weight_only_precision', const='int8', type=str, nargs='?', default='int8', choices=['int8', 'int4', 'int4_awq'], help= 'Define the precision for the weights when using weight-only quantization.' 'You must also use --use_weight_only for that argument to have an impact.' ) parser.add_argument( '--enable_fp8', default=False, action='store_true', help='Use FP8 Linear layer for Attention QKV/Dense and MLP.') parser.add_argument( '--fp8_kv_cache', default=False, action="store_true", help='By default, we use dtype for KV cache. fp8_kv_cache chooses int8 ' 'quantization for KV') parser.add_argument('--load_by_shard', action='store_true', help='Load a pretrained model shard-by-shard.') parser.add_argument('--output_dir', type=str, default='tllm_checkpoint', help='The path to save the TensorRT-LLM checkpoint') parser.add_argument( '--workers', type=int, default=1, help='The number of workers for converting checkpoint in parallel') args = parser.parse_args() return args def load_falcon_config(model_dir: str) -> FalconConfig: """ Helper utility to load FalconConfig. A pretrained checkpoint from modeling_RW.py has a different structure and is not compatible with `transformers.FalconConfig` and `transformers.FalconModel`. We need to manually set the config values. """ config = FalconConfig.from_pretrained(model_dir) config.architectures = ["FalconForCausalLM"] if config.model_type not in ['RefinedWebModel', 'RefinedWeb']: return config if config.model_type == 'RefinedWeb': # Case 1. Falcon-40B / Falcon-40B-instruct # https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json config.num_hidden_layers = config.n_layer config.num_attention_heads = config.n_head config.num_kv_heads = config.n_head_kv config.new_decoder_architecture = True elif config.model_type == 'RefinedWebModel': # Case 2. Falcon-7B / Falcon-7B-instruct # https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json config.num_hidden_layers = config.n_layer config.num_attention_heads = config.n_head config.num_kv_heads = 1 if config.multi_query else config.n_head config.new_decoder_architecture = False else: raise ValueError("Shouldn't reach here.") config.model_type = 'falcon' return config def split(weight: torch.Tensor, tp_size: int, rank: int = 0, dim: int = 0) -> torch.Tensor: if tp_size == 1: return weight elif weight.ndim == 1: return torch.chunk(weight, tp_size)[rank].contiguous() else: return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous() def reorder_qkv_weight_or_bias(weight: torch.Tensor, head_dim: int, num_heads: int, num_kv_heads: Optional[int] = None, tp_size: int = 1, is_bias: bool = False) -> torch.Tensor: """ Reorder the qkv weight for TRT-LLM use. The shape of the fused QKV weights in HF is different from the shape that TRT-LLM requires. In particular, the weight of HF consists of interleaved q, k, v head weights, while that of TRT-LLM is contiguous. HF : [q1, k1, v1, ..., qh, kh, vh] TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, vh] where qi, vi, ki are weight vectors corresponding to attention head i. It's similar to multi/grouped query attention cases. We reorder and split the weight of an attention layer to fit into TRT-LLM. The reordered weight and bias will be weight: (T, Qh * D + 2 * KVh * D, H) bias : (T, Qh * D + 2 * KVh * D) where T=tp_size, Qh=local_num_q_heads, KVh=local_num_kv_heads, D=head_dim, H=hidden_dim. In the multi/grouped query attention, the number of K/V attention heads are less than that of Q attention, so that K/V attention heads may be shared across different ranks if necessary. For tensor parallelism, we use the first dimension to select the corresponding weights. """ # Query types and expected kv heads. # - Conventional MHA: num_heads = num_kv_heads # - Multi-Query Attention: num_kv_heads = 1 # - Grouped-Query Attention: num_heads % num_kv_heads = 0 num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads assert num_heads % num_kv_heads == 0, \ f'num_heads({num_heads}) must be divisible by '\ f'num_kv_heads({num_kv_heads})).' # The number of attention heads per group: N q head + 1 k head + 1 v head. num_group_heads = num_heads // num_kv_heads + 2 assert weight.shape[0] == num_kv_heads * num_group_heads * head_dim, \ f'{weight.shape[0]} != {num_kv_heads} * {num_group_heads} * {head_dim}' qkv_in = num_heads * head_dim if not is_bias else 1 # Split Q/K/V weights weight = weight.reshape(num_kv_heads, num_heads // num_kv_heads + 2, head_dim, qkv_in) q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, head_dim, qkv_in) k_w = weight[:, -2:-1, ...] # (nKV, 1, head_dim, qkv_in) v_w = weight[:, -1:, ...] # (nKV, 1, head_dim, qkv_in) if num_kv_heads < num_heads and num_kv_heads < tp_size: # Duplicate K/V heads to make sure that each rank has at least one # K/V heads. For instance, num_heads=8, num_kv_heads=2, tp_size=4, # we will make the qkv weight as below. # Orig: [q0 q1 q2 q3 k0 v0 q4 q5 q6 q7 k1 v0 v1] # >>>> [[q0 q1 k0 v0], [q2 q3 k0 v0], [q4 q5 k1 v1], [q6 q7 k1 v1]] assert tp_size % num_kv_heads == 0 num_dups = tp_size // num_kv_heads # k_w and v_w have the same shape. new_shape = (num_kv_heads, num_dups) + k_w.shape[2:] k_w = torch.broadcast_to(k_w, size=new_shape) v_w = torch.broadcast_to(v_w, size=new_shape) # Update the number of kv heads. num_kv_heads = tp_size reordered = torch.concat( [ q_w.reshape(tp_size, num_heads // tp_size, head_dim, qkv_in), k_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), v_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), ], dim=1, ) qkv_out = (num_heads + 2 * num_kv_heads) // tp_size * head_dim return reordered.reshape((tp_size, qkv_out, -1)) def split_qkv_weight(weight: torch.Tensor, hidden_size: int, num_heads: int, tp_size: int, rank: int, is_bias: bool, num_kv_heads: Optional[int] = None) -> torch.Tensor: """ Splits the QKV matrix according to tensor parallelism """ head_dim = hidden_size // num_heads weight = reorder_qkv_weight_or_bias(weight, head_dim=head_dim, num_heads=num_heads, num_kv_heads=num_kv_heads, tp_size=tp_size, is_bias=is_bias) # Copy a sliced tensor to prevent memory leak. A sliced tensor shares the # memory buffer of the original tensor. So, returning without copying makes # the buffer of a loaded "qkv" be referenced, resulting GC can't release # those weights until the whole process ends. if not is_bias: return weight[rank, ...].clone().contiguous() else: return weight[rank, ...].ravel().clone().contiguous() def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, dim: int) -> torch.Tensor: return split(weight, tp_size, rank, dim=dim) def get_weight(params: Dict[str, torch.Tensor], prefix: str, dtype: torch.dtype) -> torch.Tensor: if f'{prefix}.weight' not in params: return None return params[f'{prefix}.weight'].to(dtype).detach().cpu() def get_bias(params: Dict[str, torch.Tensor], prefix: str, dtype: torch.dtype) -> torch.Tensor: if f'{prefix}.bias' not in params: return None return params[f'{prefix}.bias'].to(dtype).detach().cpu() def get_weight_and_bias(params: Dict[str, torch.Tensor], prefix: str, dtype: torch.dtype) -> Tuple[torch.Tensor]: return get_weight(params, prefix, dtype), get_bias(params, prefix, dtype) def get_tllm_linear_weight( weight: torch.Tensor, prefix: str, bias: Optional[torch.Tensor] = None, use_weight_only: bool = False, plugin_weight_only_quant_type: torch.dtype = torch.int8 ) -> Dict[str, torch.Tensor]: results = {} if use_weight_only: v = weight.t().contiguous() processed_torch_weights, torch_weight_scales = \ torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( v, plugin_weight_only_quant_type) results[f'{prefix}.weight'] = processed_torch_weights results[f'{prefix}.per_channel_scale'] = torch_weight_scales else: results[f'{prefix}.weight'] = weight.contiguous() if bias is not None: results[f'{prefix}.bias'] = bias return results def get_tllm_param( param: torch.Tensor, name: str, use_weight_only: bool = False, plugin_weight_only_quant_type: torch.dtype = torch.int8 ) -> Dict[str, torch.Tensor]: results = {} if name.endswith('.weight') and use_weight_only: v = param.t().contiguous() processed_torch_weights, torch_weight_scales = \ torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( v, plugin_weight_only_quant_type) results[name] = processed_torch_weights results[name.replace('weight', 'per_channel_scale')] = torch_weight_scales else: results[name] = param return results def convert_hf_falcon(hf_model: FalconForCausalLM, hf_config: FalconConfig, mapping: Mapping, dtype: str = 'float32', use_weight_only: bool = False, plugin_weight_only_quant_type: torch.dtype = torch.int8): weights = {} tik = time.time() model_params = dict(hf_model.named_parameters()) dtype = getattr(torch, dtype) num_attention_heads = hf_config.num_attention_heads hidden_size = hf_config.hidden_size num_kv_heads = getattr(hf_config, 'num_kv_heads', num_attention_heads) num_hidden_layers = hf_config.num_hidden_layers parallel_attention = hf_config.parallel_attn new_decoder_architecture = hf_config.new_decoder_architecture layers_range = mapping.pp_layers(num_hidden_layers) for l in layers_range: prefix = f'transformer.h.{l}' tllm_prex = f'transformer.layers.{l-layers_range[0]}' qkv_weight, qkv_bias = get_weight_and_bias( model_params, f'{prefix}.self_attention.query_key_value', dtype) qkv_w = split_qkv_weight(qkv_weight, hidden_size, num_attention_heads, mapping.tp_size, mapping.tp_rank, is_bias=False, num_kv_heads=num_kv_heads) if qkv_bias is None: qkv_b = None else: qkv_b = split_qkv_weight(qkv_bias, hidden_size, num_attention_heads, mapping.tp_size, mapping.tp_rank, is_bias=True, num_kv_heads=num_kv_heads) weights.update( get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', qkv_b, use_weight_only, plugin_weight_only_quant_type)) attn_dense_weight, attn_dense_bias = get_weight_and_bias( model_params, f'{prefix}.self_attention.dense', dtype) attn_dense_w = split_matrix(attn_dense_weight, mapping.tp_size, mapping.tp_rank, dim=1) weights.update( get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', attn_dense_bias, use_weight_only, plugin_weight_only_quant_type)) mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) mlp_fc_w = split_matrix(mlp_fc_weight, mapping.tp_size, mapping.tp_rank, dim=0) if mlp_fc_bias is None: mlp_fc_b = None else: mlp_fc_b = split_matrix(mlp_fc_bias, mapping.tp_size, mapping.tp_rank, dim=0) weights.update( get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, use_weight_only, plugin_weight_only_quant_type)) mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) mlp_proj_w = split_matrix(mlp_proj_weight, mapping.tp_size, mapping.tp_rank, dim=1) weights.update( get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', mlp_proj_bias, use_weight_only, plugin_weight_only_quant_type)) if new_decoder_architecture: input_ln_weight, input_ln_bias = get_weight_and_bias( model_params, f'{prefix}.ln_attn', dtype) weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight if input_ln_bias is not None: weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias mlp_ln_weight, mlp_ln_bias = get_weight_and_bias( model_params, f'{prefix}.ln_mlp', dtype) weights[f'{tllm_prex}.mlp_layernorm.weight'] = mlp_ln_weight if mlp_ln_bias is not None: weights[f'{tllm_prex}.mlp_layernorm.bias'] = mlp_ln_bias else: input_ln_weight, input_ln_bias = get_weight_and_bias( model_params, f'{prefix}.input_layernorm', dtype) weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight if input_ln_bias is not None: weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias if not parallel_attention: post_ln_weight, post_ln_bias = get_weight_and_bias( model_params, f'{prefix}.post_attention_layernorm', dtype) if post_ln_weight is not None: weights[ f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight if post_ln_bias is not None: weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) if mapping.is_first_pp_rank(): weights['transformer.embedding.weight'] = embed_w if mapping.is_last_pp_rank(): weights['lm_head.weight'] = split_matrix(embed_w.clone(), mapping.tp_size, mapping.tp_rank, dim=0) ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', dtype) weights['transformer.ln_f.weight'] = ln_f_w if ln_f_b is not None: weights['transformer.ln_f.bias'] = ln_f_b tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) print(f'Weights loaded. Total time: {t}') return weights def load_from_hf_falcon_checkpoint( hf_model_dir: str, hf_config: FalconConfig, mapping: Mapping, dtype: str = 'float32', use_weight_only: bool = False, plugin_weight_only_quant_type: torch.dtype = torch.int8): weights = {} tik = time.time() dtype = getattr(torch, dtype) num_attention_heads = hf_config.num_attention_heads hidden_size = hf_config.hidden_size num_kv_heads = getattr(hf_config, 'num_kv_heads', num_attention_heads) num_hidden_layers = hf_config.num_hidden_layers layers_range = mapping.pp_layers(num_hidden_layers) for model_file in iterate_shard_files(hf_model_dir, mapping.tp_rank): state_dict = load_state_dict(model_file, dtype) for name, param in state_dict.items(): param = param.detach().cpu() l = retrieved_layer_index_from_name(name) if l is not None: if l not in layers_range: continue tllm_prex = f'transformer.layers.{l-layers_range[0]}' if 'self_attention.query_key_value' in name: if name.endswith('weight'): qkv_w = split_qkv_weight(param, hidden_size, num_attention_heads, mapping.tp_size, mapping.tp_rank, is_bias=False, num_kv_heads=num_kv_heads) weights.update( get_tllm_param(qkv_w, f'{tllm_prex}.attention.qkv.weight', use_weight_only, plugin_weight_only_quant_type)) else: qkv_b = split_qkv_weight(param, hidden_size, num_attention_heads, mapping.tp_size, mapping.tp_rank, is_bias=True, num_kv_heads=num_kv_heads) weights.update( get_tllm_param(qkv_b, f'{tllm_prex}.attention.qkv.bias', use_weight_only, plugin_weight_only_quant_type)) elif 'self_attention.dense' in name: if name.endswith('weight'): attn_dense_w = split_matrix(param, mapping.tp_size, mapping.tp_rank, dim=1) weights.update( get_tllm_param( attn_dense_w, f'{tllm_prex}.attention.dense.weight', use_weight_only, plugin_weight_only_quant_type)) else: weights.update( get_tllm_param(param, f'{tllm_prex}.attention.dense.bias', use_weight_only, plugin_weight_only_quant_type)) elif 'mlp.dense_h_to_4h' in name: if name.endswith('weight'): mlp_fc_w = split_matrix(param, mapping.tp_size, mapping.tp_rank, dim=0) weights.update( get_tllm_param(mlp_fc_w, f'{tllm_prex}.mlp.fc.weight', use_weight_only, plugin_weight_only_quant_type)) else: mlp_fc_b = split_matrix(param, mapping.tp_size, mapping.tp_rank, dim=0) weights.update( get_tllm_param(mlp_fc_b, f'{tllm_prex}.mlp.fc.bias', use_weight_only, plugin_weight_only_quant_type)) elif 'mlp.dense_4h_to_h' in name: if name.endswith('weight'): mlp_proj_w = split_matrix(param, mapping.tp_size, mapping.tp_rank, dim=1) weights.update( get_tllm_param(mlp_proj_w, f'{tllm_prex}.mlp.proj.weight', use_weight_only, plugin_weight_only_quant_type)) else: weights.update( get_tllm_param(param, f'{tllm_prex}.mlp.proj.bias', use_weight_only, plugin_weight_only_quant_type)) elif 'ln_attn' in name or 'input_layernorm' in name: if name.endswith('weight'): weights[f'{tllm_prex}.input_layernorm.weight'] = param else: weights[f'{tllm_prex}.input_layernorm.bias'] = param elif 'ln_mlp' in name: if name.endswith('weight'): weights[f'{tllm_prex}.mlp_layernorm.weight'] = param else: weights[f'{tllm_prex}.mlp_layernorm.bias'] = param elif 'post_attention_layernorm' in name: if name.endswith('weight'): weights[f'{tllm_prex}.post_layernorm.weight'] = param else: weights[f'{tllm_prex}.post_layernorm.bias'] = param else: if 'word_embeddings' in name: if mapping.is_first_pp_rank(): weights['transformer.embedding.weight'] = param if mapping.is_last_pp_rank(): weights['lm_head.weight'] = split_matrix( param.clone(), mapping.tp_size, mapping.tp_rank, dim=0) elif 'ln_f' in name: if mapping.is_last_pp_rank(): if name.endswith('weight'): weights['transformer.ln_f.weight'] = param else: weights['transformer.ln_f.bias'] = param del state_dict tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) print(f'Weights loaded. Total time: {t}') return weights def load_from_awq_falcon(quant_ckpt_path: str, hf_config: FalconConfig, mapping: Mapping, dtype: str = "float16"): weights = {} tik = time.time() num_hidden_layers = hf_config.num_hidden_layers parallel_attention = hf_config.parallel_attn new_decoder_architecture = hf_config.new_decoder_architecture packer = torch.ops.fastertransformer.pack_int8_tensor_to_packed_int4 preprocessor = torch.ops.fastertransformer.preprocess_weights_for_mixed_gemm torch_dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) if not quant_ckpt_path.endswith(".npz"): raise ValueError("Unsupported AWQ quantized checkpoint format") awq_falcon = np.load(quant_ckpt_path) awq_prefix = "_np:" awq_suffix_list = [ ":weight", ":weights_scaling_factor", ":prequant_scaling_factor", ] awq_key_list = [ "vocab_embedding:weight", # embedding "lm_head", # lm_head "final_layernorm", # ln_f "attention:qkv:", # attention.qkv "attention:dense", # attention.dense "mlp:proj", # mlp.proj "mlp:fc", # mlp.fc "input_layernorm", # input_layernorm.weight "mlp_layernorm", # mlp_layernorm.weight ] split_sym = ":" AMMO_WEIGHT_SCALING_FACTOR_COEFF = 7 def load(key): if awq_prefix + key not in awq_falcon: return None v = torch.from_numpy(awq_falcon[awq_prefix + key]).to(torch_dtype) if "weights_scaling_factor" in key: v *= AMMO_WEIGHT_SCALING_FACTOR_COEFF # For AMMO *.npz checkpoints return v group_size = load("layers:0:attention:dense:weight").numel() // load( "layers:0:attention:dense:weights_scaling_factor").numel() def torch_split(v, dim): if v.shape[dim] % mapping.tp_size != 0: tensorrt_llm.logger.error( "Current weight shape is invalid for mapping.tp_size=" + str(mapping.tp_size)) raise ValueError("Invalid TP size") return v.split(v.shape[dim] // mapping.tp_size, dim=dim)[mapping.tp_rank] def AWQ_quantize_pack_preprocess(weight, scale): weight /= scale.repeat_interleave(group_size, dim=0) qweight_int8 = torch.clamp(torch.round(weight.cuda()).char(), -8, 7) int4_weight = preprocessor(packer(qweight_int8.cpu()), torch.quint4x2) return int4_weight.view(torch.int8) def get_tllm_weight_from_awq(v: List[torch.Tensor], tllm_prex: str, tp_dim: int = 0): weight = v[0].T.contiguous() [k, n] = weight.shape weight = torch_split(weight, tp_dim) amax = v[1].reshape((n, k // group_size)).T.contiguous() amax = torch_split(amax, tp_dim) pre_quant_scale = v[2].reshape((1, k)) if tp_dim == 0: pre_quant_scale = torch_split(pre_quant_scale, 1) scale = amax / 8.0 results = { f'{tllm_prex}.weight': AWQ_quantize_pack_preprocess(weight, scale), f'{tllm_prex}.weights_scaling_factor': scale.to(torch_dtype), f'{tllm_prex}.prequant_scaling_factor': pre_quant_scale.to(torch_dtype), } return results def get_scale(weight): [k, n] = weight.shape weight_t = weight.T.contiguous() weight_t = weight_t.reshape(n, k // group_size, group_size) weight_t = torch.abs(weight_t.reshape(-1, group_size)) amax, idx = weight_t.max(1) amax = amax.reshape(n, k // group_size).T.contiguous() scale = amax / 8 return scale def get_tllm_qkv_weight_from_awq(prefix, tllm_prex: str): q_weight = load(prefix + "q" + awq_suffix_list[0]).T.contiguous() k_weight = load(prefix + "k" + awq_suffix_list[0]).T.contiguous() v_weight = load(prefix + "v" + awq_suffix_list[0]).T.contiguous() dim_k = q_weight.shape[0] q_weight = torch_split(q_weight, 1) k_weight = torch_split(k_weight, 1) v_weight = torch_split(v_weight, 1) qkv_pre_quant_scale = load(prefix + "q" + awq_suffix_list[2]).reshape( (1, dim_k)) qkv_weights = torch.cat((q_weight, k_weight, v_weight), dim=1) qkv_scale = get_scale(qkv_weights) results = { f'{tllm_prex}.prequant_scaling_factor': qkv_pre_quant_scale.to(torch_dtype), f'{tllm_prex}.weight': AWQ_quantize_pack_preprocess(qkv_weights, qkv_scale), f'{tllm_prex}.weights_scaling_factor': qkv_scale.to(torch_dtype), } return results # Load weights from AWQ checkpoint into TRT-LLM module # 1. embedding v = load(awq_key_list[0]) # TRT-LLM requires vocab_size to be multiple of 64 for successful GEMM if v.shape[0] % 64 != 0: v = torch.nn.functional.pad(v, [0, 0, 0, 64 - v.shape[0] % 64]) if mapping.is_first_pp_rank(): # tensorrt_llm_falcon.embedding.weight.value = v.to(torch_dtype) weights['transformer.embedding.weight'] = v.to(torch_dtype) # 2. lm_head v = [load(awq_key_list[1] + suf) for suf in awq_suffix_list] if v[0].shape[0] % 64 != 0: v[0] = torch.nn.functional.pad(v[0], [0, 0, 0, 64 - v[0].shape[0] % 64]) v[1] = torch.nn.functional.pad(v[1], [0, 0, 0, 64 - v[1].shape[0] % 64], value=1) if mapping.is_last_pp_rank(): # process_and_assign_weight(tensorrt_llm_falcon.lm_head, v, 1) weights.update(get_tllm_weight_from_awq(v, 'lm_head', 1)) # 3. ln_f v_weight = load(awq_key_list[2] + split_sym + "weight") v_bias = load(awq_key_list[2] + split_sym + "bias") if mapping.is_last_pp_rank(): # tensorrt_llm_falcon.ln_f.weight.value = v_weight.to(torch_dtype) # tensorrt_llm_falcon.ln_f.bias.value = v_bias.to(torch_dtype) weights['transformer.ln_f.weight'] = v_weight.to(torch_dtype) weights['transformer.ln_f.bias'] = v_bias.to(torch_dtype) # 4. Weights inside each layer layers_range = mapping.pp_layers(num_hidden_layers) for l in layers_range: # layer_idx = l - mapping.pp_rank * layers_per_pipeline_stage # prefix = "layers" + split_sym + str(layer_idx) + split_sym # tensorrt_llm.logger.info(f'Process weights in layer: {layer_idx}') # layer = tensorrt_llm_falcon.layers[layer_idx] prefix = f'layers{split_sym}{l}{split_sym}' tllm_prex = f'transformer.layers.{l-layers_range[0]}' # 4.1 attention.qkv # process_and_assign_qkv_weight(prefix + awq_key_list[3], # layer.attention.qkv) weights.update( get_tllm_qkv_weight_from_awq(prefix + awq_key_list[3], f'{tllm_prex}.attention.qkv')) q_b = load(prefix + awq_key_list[3] + 'q:bias') k_b = load(prefix + awq_key_list[3] + 'k:bias') v_b = load(prefix + awq_key_list[3] + 'v:bias') if q_b is not None: q_b = torch_split(q_b, dim=0) k_b = torch_split(k_b, dim=0) v_b = torch_split(v_b, dim=0) qkv_b = torch.cat((q_b, k_b, v_b), dim=0) weights[f'{tllm_prex}.attention.qkv.bias'] = qkv_b # 4.2 attention.dense v = [load(prefix + awq_key_list[4] + suf) for suf in awq_suffix_list] # process_and_assign_weight(layer.attention.dense, v, 0) weights.update( get_tllm_weight_from_awq(v, f'{tllm_prex}.attention.dense', tp_dim=0)) b = load(prefix + awq_key_list[4] + ':bias') if b is not None: if mapping.tp_rank > 0: b = torch.zeros_like(b) weights[f'{tllm_prex}.attention.dense.bias'] = b # 4.4 mlp.fc v = [load(prefix + awq_key_list[6] + suf) for suf in awq_suffix_list] # process_and_assign_weight(layer.mlp.fc, v, 1) weights.update( get_tllm_weight_from_awq(v, f'{tllm_prex}.mlp.fc', tp_dim=1)) b = load(prefix + awq_key_list[6] + ':bias') if b is not None: b = torch_split(b, dim=0) weights[f'{tllm_prex}.mlp.fc.bias'] = b # 4.3 mlp.proj v = [load(prefix + awq_key_list[5] + suf) for suf in awq_suffix_list] # process_and_assign_weight(layer.mlp.proj, v, 0) weights.update( get_tllm_weight_from_awq(v, f'{tllm_prex}.mlp.proj', tp_dim=0)) b = load(prefix + awq_key_list[5] + ':bias') if b is not None: if mapping.tp_rank > 0: b = torch.zeros_like(b) weights[f'{tllm_prex}.mlp.proj.bias'] = b if new_decoder_architecture: # 4.5 input_layernorm v = load(prefix + awq_key_list[7] + split_sym + "weight") # layer.input_layernorm.weight.value = v.to(torch_dtype) weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) v = load(prefix + awq_key_list[7] + split_sym + "bias") # layer.input_layernorm.bias.value = v.to(torch_dtype) weights[f'{tllm_prex}.input_layernorm.bias'] = v.to(torch_dtype) # 4.6 mlp_layernorm v = load(prefix + awq_key_list[8] + split_sym + "weight") # layer.mlp_layernorm.weight.value = v.to(torch_dtype) weights[f'{tllm_prex}.mlp_layernorm.weight'] = v.to(torch_dtype) v = load(prefix + awq_key_list[8] + split_sym + "bias") # layer.mlp_layernorm.bias.value = v.to(torch_dtype) weights[f'{tllm_prex}.mlp_layernorm.bias'] = v.to(torch_dtype) else: # 4.5 input_layernorm v = load(prefix + awq_key_list[7] + split_sym + "weight") # layer.input_layernorm.weight.value = v.to(torch_dtype) weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) v = load(prefix + awq_key_list[7] + split_sym + "bias") # layer.input_layernorm.bias.value = v.to(torch_dtype) weights[f'{tllm_prex}.input_layernorm.bias'] = v.to(torch_dtype) if not parallel_attention: # 4.6 post_layernorm v = load(prefix + 'post_layernorm' + split_sym + "weight") weights[f'{tllm_prex}.post_layernorm.weight'] = v.to( torch_dtype) v = load(prefix + 'post_layernorm' + split_sym + "bias") weights[f'{tllm_prex}.post_layernorm.bias'] = v.to(torch_dtype) tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) tensorrt_llm.logger.info(f'Weights loaded. Elapsed time: {t}') return weights def load_from_fp8_falcon(quant_ckpt_path: str, hf_config: FalconConfig, mapping: Mapping, fp8_kv_cache: bool): """ Get the fp8 scaling factors for Falcon model. """ fake_fp8_sf_dt = torch.float32 fp8_falcon = np.load(quant_ckpt_path) weights = {} num_hidden_layers = hf_config.num_hidden_layers layers_range = mapping.pp_layers(num_hidden_layers) for l in layers_range: prefix = f'_np:layers:{l}' tllm_prex = f'transformer.layers.{l-layers_range[0]}' weights[f'{tllm_prex}.attention.qkv.activation_scaling_factor'] = torch.tensor( [ max( fp8_falcon[ f'{prefix}:attention:qkv:q:activation_scaling_factor']. item(), fp8_falcon[ f'{prefix}:attention:qkv:k:activation_scaling_factor']. item(), fp8_falcon[ f'{prefix}:attention:qkv:v:activation_scaling_factor']. item()) ], dtype=fake_fp8_sf_dt) weights[ f'{tllm_prex}.attention.qkv.weights_scaling_factor'] = torch.tensor( [ max( fp8_falcon[ f'{prefix}:attention:qkv:q:weights_scaling_factor']. item(), fp8_falcon[ f'{prefix}:attention:qkv:k:weights_scaling_factor']. item(), fp8_falcon[ f'{prefix}:attention:qkv:v:weights_scaling_factor']. item()) ], dtype=fake_fp8_sf_dt) weights[ f'{tllm_prex}.attention.dense.activation_scaling_factor'] = torch.tensor( [ fp8_falcon[ f'{prefix}:attention:dense:activation_scaling_factor']. item() ], dtype=fake_fp8_sf_dt) weights[ f'{tllm_prex}.attention.dense.weights_scaling_factor'] = torch.tensor( [ fp8_falcon[ f'{prefix}:attention:dense:weights_scaling_factor']. item() ], dtype=fake_fp8_sf_dt) weights[f'{tllm_prex}.mlp.fc.activation_scaling_factor'] = torch.tensor( [fp8_falcon[f'{prefix}:mlp:fc:activation_scaling_factor'].item()], dtype=fake_fp8_sf_dt) weights[f'{tllm_prex}.mlp.fc.weights_scaling_factor'] = torch.tensor( [fp8_falcon[f'{prefix}:mlp:fc:weights_scaling_factor'].item()], dtype=fake_fp8_sf_dt) weights[ f'{tllm_prex}.mlp.proj.activation_scaling_factor'] = torch.tensor( [ fp8_falcon[f'{prefix}:mlp:proj:activation_scaling_factor']. item() ], dtype=fake_fp8_sf_dt) weights[f'{tllm_prex}.mlp.proj.weights_scaling_factor'] = torch.tensor( [fp8_falcon[f'{prefix}:mlp:proj:weights_scaling_factor'].item()], dtype=fake_fp8_sf_dt) if fp8_kv_cache: # Not calibrarting KV cache. scaling_factor = 1.0 weights[ f'{tllm_prex}.attention.kv_orig_quant_scale'] = torch.tensor( [scaling_factor], dtype=fake_fp8_sf_dt) weights[ f'{tllm_prex}.attention.kv_quant_orig_scale'] = torch.tensor( [1.0 / scaling_factor], dtype=fake_fp8_sf_dt) return weights if __name__ == '__main__': # TODO(qijun): Currently, the convert script depends on a torch op: # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, # which is included in tensorrt_llm Python package. Otherwise, the convert # script does not need to import tensorrt_llm. Will remove it after reimplementing # the op with PyTorch. print(tensorrt_llm.__version__) args = parse_arguments() tik = time.time() if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) hf_config = load_falcon_config(args.model_dir) config = { 'architecture': hf_config.architectures[0], 'dtype': args.dtype, 'logits_dtype': args.logits_dtype, 'num_hidden_layers': hf_config.num_hidden_layers, 'num_attention_heads': hf_config.num_attention_heads, 'num_key_value_heads': hf_config.num_kv_heads, 'hidden_size': hf_config.hidden_size, 'norm_epsilon': hf_config.layer_norm_epsilon, 'vocab_size': hf_config.vocab_size, 'position_embedding_type': 'alibi_with_scale' if hf_config.alibi else 'rope_gpt_neox', 'max_position_embeddings': hf_config.max_position_embeddings, 'hidden_act': 'gelu', 'quantization': { 'use_weight_only': args.use_weight_only, 'weight_only_precision': args.weight_only_precision, 'per_group': args.per_group, 'group_size': args.group_size, 'enable_fp8': args.enable_fp8, 'fp8_kv_cache': args.fp8_kv_cache, }, 'mapping': { 'world_size': args.world_size, 'tp_size': args.tp_size, 'pp_size': args.pp_size, }, 'bias': hf_config.bias, 'parallel_attention': hf_config.parallel_attn, 'new_decoder_architecture': hf_config.new_decoder_architecture, } if args.weight_only_precision == 'int4_awq': config['quantization'].update({ 'zero': False, 'pre_quant_scale': True, 'exclude_modules': [], }) with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: json.dump(config, f, indent=4) if args.weight_only_precision == 'int8': plugin_weight_only_quant_type = torch.int8 elif args.weight_only_precision == 'int4': plugin_weight_only_quant_type = torch.quint4x2 def covert_and_save(rank): mapping = Mapping(world_size=args.tp_size * args.pp_size, rank=rank, tp_size=args.tp_size, pp_size=args.pp_size) if args.use_weight_only and args.weight_only_precision == 'int4_awq': weights = load_from_awq_falcon(args.ammo_quant_ckpt_path, hf_config, mapping, dtype=args.dtype) else: if args.load_by_shard: weights = load_from_hf_falcon_checkpoint( args.model_dir, hf_config, mapping, dtype=args.dtype, use_weight_only=args.use_weight_only, plugin_weight_only_quant_type=plugin_weight_only_quant_type) else: hf_model = AutoModelForCausalLM.from_pretrained( args.model_dir, trust_remote_code=True, torch_dtype="auto") weights = convert_hf_falcon( hf_model, hf_config, mapping, dtype=args.dtype, use_weight_only=args.use_weight_only, plugin_weight_only_quant_type=plugin_weight_only_quant_type) del hf_model if args.enable_fp8 or args.fp8_kv_cache: scales = load_from_fp8_falcon(args.ammo_quant_ckpt_path, hf_config, mapping, args.fp8_kv_cache) weights.update(scales) safetensors.torch.save_file( weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) if args.workers == 1: for rank in range(args.world_size): covert_and_save(rank) else: with ThreadPoolExecutor(max_workers=args.workers) as p: futures = [ p.submit(covert_and_save, rank) for rank in range(args.world_size) ] wait(futures) tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) print(f'Total time of converting checkpoints: {t}')