import copy import functools import os import time from collections import defaultdict from typing import Dict, Optional import numpy as np import safetensors import torch from tqdm import tqdm from transformers import AutoModel, AutoTokenizer from tensorrt_llm._utils import pad_vocab_size from tensorrt_llm.models import ChatGLMConfig from tensorrt_llm.models.convert_utils import (generate_int8, get_weight, get_weight_and_bias, load_calib_dataset, smooth_gemm) from tensorrt_llm.quantization import QuantAlgo from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS 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 tile_kv_weight_bias(v: torch.Tensor, kv_num_head: int, tp_size: int): head_size = v.shape[0] // kv_num_head reps = tp_size // kv_num_head if v.ndim == 1: v = v.reshape(kv_num_head, head_size)[:, None, :] v = v.expand(kv_num_head, reps, head_size).reshape(-1).clone() else: hidden_size = v.shape[1] v = v.reshape(kv_num_head, head_size, hidden_size)[:, None, :, :] v = v.expand(kv_num_head, reps, head_size, hidden_size).reshape(-1, hidden_size).clone() return v def split_qkv(v: torch.Tensor, tp_size: int, rank: int, hidden_size: int, num_heads: int, num_kv_heads: int): head_size = hidden_size // num_heads if tp_size == 1: return v assert v.shape[0] == hidden_size + head_size * num_kv_heads * 2 query = v[:hidden_size] key = v[hidden_size:hidden_size + head_size * num_kv_heads] value = v[hidden_size + head_size * num_kv_heads:hidden_size + head_size * num_kv_heads * 2] if num_kv_heads < tp_size: key = tile_kv_weight_bias(key, num_kv_heads, tp_size) value = tile_kv_weight_bias(value, num_kv_heads, tp_size) assert (key.shape[0] % (tp_size * head_size)) == 0 assert (value.shape[0] % (tp_size * head_size)) == 0 q_tmp = torch.chunk(query, tp_size, dim=0)[rank] k_tmp = torch.chunk(key, tp_size, dim=0)[rank] v_tmp = torch.chunk(value, tp_size, dim=0)[rank] return torch.concatenate([q_tmp, k_tmp, v_tmp], dim=0).contiguous() def split_embedding( param: torch.Tensor, tp_size: int, tp_rank: int, use_parallel_embedding: bool = False, sharding_dim: int = 0, ) -> torch.Tensor: if param is None: return None if not use_parallel_embedding: return param vocab_size, hidden_size = param.size() if sharding_dim == 0: if vocab_size % tp_size != 0: vocab_size_padded = pad_vocab_size(vocab_size, tp_size) pad_width = vocab_size_padded - vocab_size param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), value=0) else: assert hidden_size % tp_size == 0 return split(param, tp_size, tp_rank, dim=sharding_dim) def swap_and_split_mlp(weight: torch.Tensor, tp_size: int, tp_rank: int, dim: int = 0) -> torch.Tensor: """Swap the positions of gate and fc weights, and split weights for tensor parallel. """ gate_weight, fc_weight = torch.chunk(weight, 2, dim=dim) fc_w = split(fc_weight, tp_size, tp_rank, dim=dim) gate_w = split(gate_weight, tp_size, tp_rank, dim=dim) return torch.cat([fc_w, gate_w], dim=dim).contiguous() 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.trtllm.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 @torch.no_grad() def capture_activation_range( model, tokenizer, dataset, num_samples=64, seq_len=512, ): model.eval() device = next(model.parameters()).device scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) def stat_tensor(name, tensor, key): tensor = tensor.view(-1, tensor.shape[-1]).detach() comming_max = tensor.abs().max(dim=0)[0].float() if scales[name][key] is None: scales[name][key] = comming_max else: scales[name][key] = torch.max(scales[name][key], comming_max) def stat_input_hook(m, x, y, name): if isinstance(x, tuple): x = x[0] stat_tensor(name, x, "x") stat_tensor(name, y, "y") # TODO: we don't need to do it every forward because inference does not change weight if scales[name]["w"] is None: scales[name]["w"] = m.weight.abs().clip(1e-8, None).max(dim=1)[0] hooks = [] for name, m in model.named_modules(): if isinstance(m, torch.nn.Linear): hooks.append( m.register_forward_hook( functools.partial(stat_input_hook, name=name))) for i in tqdm(range(num_samples), desc="Calibration"): input_ids = tokenizer( dataset[i], return_tensors="pt", max_length=seq_len, truncation=True, ) model(input_ids.input_ids.to(device)) for h in hooks: h.remove() return scales @torch.no_grad() def smooth_chatglm_model( model, act_range, alpha, model_smoother, ): for name, module in model.named_modules(): if not module._get_name() == "GLMBlock": continue # QKV multiplication weight layer_name = name + '.self_attention.query_key_value' print(f'Smoothing module: {layer_name}') weight = module.self_attention.query_key_value.weight smoother = smooth_gemm( weight, act_range[layer_name]["x"], module.input_layernorm.weight, None, alpha, ) act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] # Dense multiplication weight layer_name = name + ".self_attention.dense" print(f'Smoothing module: {layer_name}') weight = module.self_attention.dense.weight smoother = smooth_gemm( weight, act_range[layer_name]["x"], None, None, alpha, ) model_smoother[layer_name] = smoother.float() act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] # Multilayer perceptron h -> 4h weight layer_name = name + ".mlp.dense_h_to_4h" print(f'Smoothing module: {layer_name}') weight = module.mlp.dense_h_to_4h.weight smoother = smooth_gemm( weight, act_range[layer_name]["x"], module.post_attention_layernorm.weight, None, alpha, ) act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] # Multilayer perceptron 4h -> h weight layer_name = name + ".mlp.dense_4h_to_h" print(f'Smoothing module: {layer_name}') weight = module.mlp.dense_4h_to_h.weight smoother = smooth_gemm( weight, act_range[layer_name]["x"], None, None, alpha, ) model_smoother[layer_name] = smoother.float() act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] def get_tllm_linear_sq_weight(vals, prefix, shape, is_qkv=False, per_token=False, per_channel=False, last_prefix=None, smoother_value=None, smoother_shape=None): results = {} col_shape = shape if (is_qkv or per_channel) else [1, 1] if per_token: if per_channel: original_weights = np.array(vals["weight.int8.col"]) else: original_weights = np.array(vals["weight.int8"]) cur_weights = original_weights if is_qkv: hidden_dim = cur_weights.shape[0] cur_weights = cur_weights.reshape(hidden_dim, -1) results[prefix + 'weight'] = torch.from_numpy(cur_weights).t().contiguous() if smoother_value is None: results[last_prefix] = torch.from_numpy( np.array([1.0], dtype=np.float16)) if per_channel: cur_per_channel_value = vals["scale_w_quant_orig.col"] else: cur_per_channel_value = vals["scale_w_quant_orig"] results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) else: if per_channel: original_weights = np.array(vals["weight.int8.col"]) else: original_weights = np.array(vals["weight.int8"]) cur_weights = original_weights if is_qkv: hidden_dim = cur_weights.shape[0] cur_weights = cur_weights.reshape(hidden_dim, -1) results[prefix + 'weight'] = torch.from_numpy(cur_weights).t().contiguous() if per_channel: cur_per_channel_value = vals["scale_y_accum_quant.col"] else: cur_per_channel_value = vals["scale_y_accum_quant"] results[prefix + 'per_channel_scale'] = torch.from_numpy( np.array([cur_per_channel_value], dtype=np.float32).reshape(col_shape)).contiguous() results[last_prefix] = torch.from_numpy( np.array([vals['scale_x_orig_quant']], dtype=np.float32)).contiguous() results[prefix + 'act_scale'] = torch.from_numpy( np.array([[vals["scale_y_quant_orig"]]], dtype=np.float32)).contiguous() if smoother_value is not None: results[prefix + 'smoother'] = smoother_value.reshape( smoother_shape).contiguous().to(torch.float32) return results def load_weights_from_hf_model(hf_model: AutoModel, config: ChatGLMConfig, act_range: Optional[dict] = None, smoother: Optional[dict] = None): weights = {} tik = time.time() model_params = dict(hf_model.named_parameters()) dtype = getattr(torch, config.dtype) num_attention_heads = config.num_attention_heads hidden_size = config.hidden_size num_kv_heads = config.num_key_value_heads num_hidden_layers = config.num_hidden_layers chatglm_version = config.chatglm_version mapping = config.mapping use_parallel_embedding = config.use_parallel_embedding sharding_dim = config.embedding_sharding_dim quant_algo = config.quantization.quant_algo use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] use_smooth_quant = config.quantization._use_plugin_sq per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 if use_weight_only: if quant_algo == QuantAlgo.W8A16: plugin_weight_only_quant_type = torch.int8 elif quant_algo == QuantAlgo.W4A16: plugin_weight_only_quant_type = torch.quint4x2 else: plugin_weight_only_quant_type = None layers_range = mapping.pp_layers(num_hidden_layers) for l in layers_range: if chatglm_version in GLM_ARCH1_VERSIONS: prefix = f'transformer.layers.{l}' elif chatglm_version in GLM_ARCH2_VERSIONS: prefix = f'transformer.encoder.layers.{l}' tllm_prex = f'transformer.layers.{l-layers_range[0]}' # Attention QKV attention_attr_name = '' if chatglm_version in GLM_ARCH1_VERSIONS: attention_attr_name = 'attention' elif chatglm_version in GLM_ARCH2_VERSIONS: attention_attr_name = 'self_attention' qkv_weight, qkv_bias = get_weight_and_bias( model_params, f'{prefix}.{attention_attr_name}.query_key_value', dtype) if use_smooth_quant: qkv_act_range = act_range.get( f'{prefix}.{attention_attr_name}.query_key_value') qkv_vals_int8 = generate_int8(qkv_weight.t(), qkv_act_range, is_qkv=True, multi_query_mode=True) weights.update( get_tllm_linear_sq_weight( vals=qkv_vals_int8, prefix=f'{tllm_prex}.attention.qkv.', shape=[1, qkv_weight.size(0)], is_qkv=True, per_token=per_token, per_channel=per_channel, last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int', smoother_value=None, smoother_shape=None)) if qkv_bias is not None: qkv_b = split_qkv(qkv_bias, mapping.tp_size, mapping.tp_rank, hidden_size, num_attention_heads, num_kv_heads=num_kv_heads) weights[f'{tllm_prex}.attention.qkv.bias'] = qkv_b else: qkv_w = split_qkv(qkv_weight, mapping.tp_size, mapping.tp_rank, hidden_size, num_attention_heads, num_kv_heads=num_kv_heads) if qkv_bias is None: qkv_b = None else: qkv_b = split_qkv(qkv_bias, mapping.tp_size, mapping.tp_rank, hidden_size, num_attention_heads, 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)) if int8_kv_cache: qkv_act_range = act_range.get( f'{prefix}.{attention_attr_name}.query_key_value') qkv_vals_int8 = generate_int8(qkv_weight.t(), qkv_act_range, is_qkv=True, multi_query_mode=True) weights[ f'{tllm_prex}.attention.kv_cache_scaling_factor'] = qkv_vals_int8[ 'scale_y_quant_orig'].contiguous() # Attention dense attn_dense_weight, attn_dense_bias = get_weight_and_bias( model_params, f'{prefix}.{attention_attr_name}.dense', dtype) if use_smooth_quant: dense_act_range = act_range.get( f'{prefix}.{attention_attr_name}.dense') dense_smoother = smoother.get( f'{prefix}.{attention_attr_name}.dense') dense_vals_int8 = generate_int8(attn_dense_weight.t(), dense_act_range, is_qkv=False, multi_query_mode=True) weights.update( get_tllm_linear_sq_weight( vals=dense_vals_int8, prefix=f'{tllm_prex}.attention.dense.', shape=[1, hidden_size], is_qkv=False, per_token=per_token, per_channel=per_channel, last_prefix= f'{tllm_prex}.attention.quantization_scaling_factor', smoother_value=dense_smoother, smoother_shape=[1, hidden_size])) if attn_dense_bias is not None: weights[f'{tllm_prex}.attention.dense.bias'] = attn_dense_bias else: attn_dense_w = split(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 mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) if use_smooth_quant: fc_act_range = act_range.get(f'{prefix}.mlp.dense_h_to_4h') fc_vals_int8 = generate_int8(mlp_fc_weight.t(), fc_act_range, is_qkv=False, multi_query_mode=True) cur_weights = get_tllm_linear_sq_weight( vals=fc_vals_int8, prefix=f'{tllm_prex}.mlp.fc.', shape=[1, mlp_fc_weight.size(0)], is_qkv=False, per_token=per_token, per_channel=per_channel, last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int', smoother_value=None, smoother_shape=None, ) cur_weights[f'{tllm_prex}.mlp.fc.weight'] = swap_and_split_mlp( cur_weights[f'{tllm_prex}.mlp.fc.weight'], mapping.tp_size, mapping.tp_rank, dim=0, ) if per_channel: cur_weights[ f'{tllm_prex}.mlp.fc.per_channel_scale'] = swap_and_split_mlp( cur_weights[f'{tllm_prex}.mlp.fc.per_channel_scale'], mapping.tp_size, mapping.tp_rank, dim=1, ) weights.update(cur_weights) if chatglm_version in GLM_ARCH1_VERSIONS: if mlp_fc_bias is not None: mlp_fc_b = split(mlp_fc_bias, mapping.tp_size, mapping.tp_rank, dim=0) weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b elif chatglm_version in GLM_ARCH2_VERSIONS: if mlp_fc_bias is not None: mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, mapping.tp_rank) weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b else: if chatglm_version in GLM_ARCH1_VERSIONS: mlp_fc_w = split(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(mlp_fc_bias, mapping.tp_size, mapping.tp_rank, dim=0) elif chatglm_version in GLM_ARCH2_VERSIONS: mlp_fc_w = swap_and_split_mlp(mlp_fc_weight, mapping.tp_size, mapping.tp_rank) if mlp_fc_bias is None: mlp_fc_b = None else: mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, mapping.tp_rank) 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 mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) if use_smooth_quant: proj_act_range = act_range.get(f'{prefix}.mlp.dense_4h_to_h') proj_smoother = smoother.get(f'{prefix}.mlp.dense_4h_to_h') proj_vals_int8 = generate_int8(mlp_proj_weight.t(), proj_act_range, is_qkv=False, multi_query_mode=True) weights.update( get_tllm_linear_sq_weight( vals=proj_vals_int8, prefix=f'{tllm_prex}.mlp.proj.', shape=[1, hidden_size], is_qkv=False, per_token=per_token, per_channel=per_channel, last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor', smoother_value=proj_smoother, smoother_shape=[1, config.intermediate_size])) if mlp_proj_bias is not None: weights[f'{tllm_prex}.mlp.proj.bias'] = mlp_proj_bias else: mlp_proj_w = split(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)) 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 post_ln_weight, post_ln_bias = get_weight_and_bias( model_params, f'{prefix}.post_attention_layernorm', dtype) 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 if mapping.is_first_pp_rank(): if chatglm_version == 'glm': embed_w = get_weight(model_params, 'word_embeddings', dtype) pos_embed_w = get_weight(model_params, 'transformer.position_embeddings', dtype) weights['transformer.position_embedding.weight'] = split_embedding( pos_embed_w, tp_size=mapping.tp_size, tp_rank=mapping.tp_rank, use_parallel_embedding=use_parallel_embedding, sharding_dim=sharding_dim) block_embed_w = get_weight(model_params, 'transformer.block_position_embeddings', dtype) weights['transformer.block_embedding.weight'] = split_embedding( block_embed_w, tp_size=mapping.tp_size, tp_rank=mapping.tp_rank, use_parallel_embedding=use_parallel_embedding, sharding_dim=sharding_dim) elif chatglm_version == 'chatglm': embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) elif chatglm_version in GLM_ARCH2_VERSIONS: embed_w = get_weight(model_params, 'transformer.embedding.word_embeddings', dtype) weights['transformer.vocab_embedding.weight'] = split_embedding( embed_w, tp_size=mapping.tp_size, tp_rank=mapping.tp_rank, use_parallel_embedding=use_parallel_embedding, sharding_dim=sharding_dim) if mapping.is_last_pp_rank(): if chatglm_version == 'glm': lm_head_weight = get_weight(model_params, 'word_embeddings', dtype).clone() elif chatglm_version == 'chatglm': lm_head_weight = get_weight(model_params, 'transformer.word_embeddings', dtype).clone() elif chatglm_version in GLM_ARCH2_VERSIONS: lm_head_weight = get_weight(model_params, 'transformer.output_layer', dtype) weights['lm_head.weight'] = split(lm_head_weight, mapping.tp_size, mapping.tp_rank, dim=0) if chatglm_version in GLM_ARCH1_VERSIONS: ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.final_layernorm', dtype) elif chatglm_version in GLM_ARCH2_VERSIONS: ln_f_w, ln_f_b = get_weight_and_bias( model_params, 'transformer.encoder.final_layernorm', 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 quantize(hf_model_dir: str, output_dir: str, config: ChatGLMConfig, calib_dataset: str = 'cnn_dailymail', device: str = 'auto', trust_remote_code: bool = True): ''' Quantize the save the model as TRT-LLM checkpoint to output_dir ''' os.makedirs(output_dir, exist_ok=True) config.to_json_file(os.path.join(output_dir, 'config.json')) mapping = config.mapping assert mapping.rank == 0, "quantize should be called at rank 0 only" quant_config = config.quantization use_smooth_quant = quant_config._use_plugin_sq int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" assert hf_model_dir is not None ## only load and call smooth quant routine once for all ranks if config.chatglm_version == 'glm': device_map = 'cuda' if device != "cpu" else 'cpu' else: device_map = 'auto' if device != "cpu" else 'cpu' hf_model = AutoModel.from_pretrained( hf_model_dir, trust_remote_code=trust_remote_code, dtype='auto' if config.chatglm_version != 'glm' else getattr( torch, config.dtype), device_map=device_map) os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( "TOKENIZERS_PARALLELISM", "false") tokenizer = AutoTokenizer.from_pretrained( hf_model_dir, trust_remote_code=trust_remote_code, ) dataset = load_calib_dataset(calib_dataset) act_range = capture_activation_range(hf_model, tokenizer, dataset, num_samples=64) smoother = {} if use_smooth_quant: smooth_chatglm_model(hf_model, act_range, quant_config.smoothquant_val, smoother) for rank in range(mapping.world_size): # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank config = copy.deepcopy(config) config.set_rank(rank) weights = load_weights_from_hf_model( hf_model, config=config, act_range=act_range, smoother=smoother, ) safetensors.torch.save_file( weights, os.path.join(output_dir, f'rank{rank}.safetensors')) del weights