| """Helper methods for the checkpoint conversion scripts. | |
| The checkpoint conversion scripts are located in the | |
| llmfoundry/scripts/inference/benchmarking/ folder. Users should run those | |
| scripts directly to convert between checkpoints; this file contains only common | |
| utility functions that are present in multiple scripts. | |
| """ | |
| import json | |
| import logging | |
| import os | |
| import random | |
| import string | |
| from pathlib import Path | |
| from typing import Any, Dict, Optional, Tuple, Union | |
| import numpy as np | |
| import sentencepiece as spm | |
| from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast | |
| log = logging.getLogger(__name__) | |
| def _get_weight_data_type(data_type: str): | |
| if data_type == 'fp32': | |
| return np.float32 | |
| elif data_type == 'fp16': | |
| return np.float16 | |
| else: | |
| raise RuntimeError('Unsupported data type: {data_type} for conversion.') | |
| def get_hf_tokenizer_from_composer_state_dict(state_dict: Dict[str, Any], trust_remote_code: bool, tokenizer_save_dir: Optional[str]=None) -> Optional[PreTrainedTokenizer]: | |
| if 'state' not in state_dict: | |
| raise RuntimeError('Unexpected composer state dictionary. Did you pass in a full composer checkpoint?') | |
| if 'integrations' not in state_dict['state'] or 'huggingface' not in state_dict['state']['integrations']: | |
| raise RuntimeError('Did not find HuggingFace related state (e.g., tokenizer) in the provided composer checkpoint!') | |
| hf_tokenizer_state = state_dict['state']['integrations']['huggingface']['tokenizer'] | |
| hf_tokenizer = None | |
| if hf_tokenizer_state != {}: | |
| if tokenizer_save_dir is None: | |
| unique_suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=6)) | |
| tokenizer_save_dir = os.path.join(os.getcwd(), f'tokenizer-save-dir-{unique_suffix}') | |
| os.makedirs(tokenizer_save_dir, exist_ok=True) | |
| for filename, saved_content in hf_tokenizer_state.items(): | |
| if filename.endswith(saved_content['file_extension']): | |
| tokenizer_file_name = filename | |
| else: | |
| tokenizer_file_name = filename + saved_content['file_extension'] | |
| tokenizer_file_path = Path(tokenizer_save_dir) / tokenizer_file_name | |
| if saved_content['file_extension'] == '.json': | |
| with open(tokenizer_file_path, 'w') as _tmp_file: | |
| json.dump(saved_content['content'], _tmp_file) | |
| elif saved_content['file_extension'] == '.txt': | |
| with open(tokenizer_file_path, 'w') as _tmp_file: | |
| for line in saved_content['content']: | |
| _tmp_file.write(line) | |
| _tmp_file.write('\n') | |
| elif saved_content['file_extension'] == '.py': | |
| with open(tokenizer_file_path, 'w') as _tmp_file: | |
| _tmp_file.write(saved_content['content']) | |
| elif saved_content['file_extension'] == '.model': | |
| s = spm.SentencePieceProcessor() | |
| s.load_from_serialized_proto(saved_content['content']) | |
| with open(tokenizer_file_path, 'wb') as _tmp_file: | |
| _tmp_file.write(s.serialized_model_proto()) | |
| hf_tokenizer = load_tokenizer(tokenizer_save_dir, trust_remote_code=trust_remote_code) | |
| hf_tokenizer.name_or_path = '' | |
| hf_tokenizer.init_kwargs['name_or_path'] = '' | |
| return hf_tokenizer | |
| def load_tokenizer(tokenizer_save_dir: str, trust_remote_code: bool) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: | |
| try: | |
| return AutoTokenizer.from_pretrained(tokenizer_save_dir, trust_remote_code=trust_remote_code) | |
| except ValueError as e: | |
| raise ValueError(f'Got error while loading tokenizer with trust_remote_code={trust_remote_code}: {e}. ' + 'If accessing a tokenizer defined outside of the transformers module,' + ' please use --trust_remote_code.') | |
| def _write_zero_bias(weight_name: str, weight_file_path: str, bias_shape: Union[Tuple[int, ...], int], np_data_type: np.dtype) -> None: | |
| """Write zeros for bias when converting MPT to FasterTransformer weights. | |
| MPT model might not have bias while FT expects bias. | |
| Args: | |
| weight_name (str): Name of the weight tensor. | |
| weight_file_path (str): Output path for storing the weight (NOT zero bias). | |
| bias_shape (Union[Tuple[int, ...], int]): Shape of the bias array. | |
| np_data_type (np.dtype): The data type for bias. | |
| """ | |
| if 'weight' not in weight_file_path: | |
| raise RuntimeError(f'Cannot write zero bias for {weight_name}. Input is not a weight tensor') | |
| log.debug(f'zero bias for weight: {weight_name}') | |
| bias_file_path = weight_file_path.replace('.weight', '.bias') | |
| bias = np.zeros(bias_shape, dtype=np_data_type) | |
| bias.tofile(bias_file_path) | |
| def _convert_weight_to_ft_each(save_dir: str, infer_gpu_num: int, tensor_name: str, config: Dict[str, Any], data: np.ndarray, np_weight_data_type: np.dtype) -> None: | |
| """Convert each MPT weight to a FasterTransformer compatible format. | |
| Args: | |
| save_dir (str): Path of the directory to save the weight in FT format. The directory must already exist. | |
| infer_gpu_num (int): The number of gpus you are planning to use for inference. | |
| tensor_name (str): Name of the weight tensor. Used in naming the output file. | |
| config (Dict[str, Any]): Configuration for the model. This is used in getting model specific parameters. | |
| data (np.ndarray): Tensor data in np.ndarray format. | |
| Returns: | |
| None: Writes to a file in `save_dir`. File name is based on the `tensor_name` | |
| """ | |
| if tensor_name.find('input_layernorm.weight') != -1 or tensor_name.find('input_layernorm.bias') != -1 or tensor_name.find('attention.dense.bias') != -1 or (tensor_name.find('post_attention_layernorm.weight') != -1) or (tensor_name.find('post_attention_layernorm.bias') != -1) or (tensor_name.find('mlp.dense_4h_to_h.bias') != -1) or (tensor_name.find('final_layernorm.weight') != -1) or (tensor_name.find('final_layernorm.bias') != -1): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.bin') | |
| data.tofile(save_path) | |
| if 'weight' in tensor_name and config['no_bias']: | |
| _write_zero_bias(tensor_name, save_path, data.shape[-1], np_weight_data_type) | |
| elif tensor_name.find('attention.dense.weight') != -1: | |
| assert data.shape == (config['d_model'], config['d_model']), f'unexpected dim for {tensor_name}' | |
| data = data.T | |
| split_vals = np.split(data, infer_gpu_num, axis=0) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| if config['no_bias']: | |
| fake_weight_path = os.path.join(save_dir, f'model.{tensor_name}.bin') | |
| _write_zero_bias(tensor_name, fake_weight_path, data.shape[-1], np_weight_data_type) | |
| elif tensor_name.find('mlp.dense_4h_to_h.weight') != -1: | |
| assert data.shape == (config['d_model'], config['expansion_ratio'] * config['d_model']), f'unexpected dim for {tensor_name}' | |
| data = data.T | |
| split_vals = np.split(data, infer_gpu_num, axis=0) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| if config['no_bias']: | |
| fake_weight_path = os.path.join(save_dir, f'model.{tensor_name}.bin') | |
| _write_zero_bias(tensor_name, fake_weight_path, data.shape[-1], np_weight_data_type) | |
| elif tensor_name.find('mlp.dense_h_to_4h.weight') != -1: | |
| assert data.shape == (config['expansion_ratio'] * config['d_model'], config['d_model']), f'unexpected dim for {tensor_name}' | |
| data = data.T | |
| split_vals = np.split(data, infer_gpu_num, axis=-1) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| if config['no_bias']: | |
| _write_zero_bias(tensor_name, save_path, split_vals[j].shape[-1], np_weight_data_type) | |
| elif tensor_name.find('mlp.dense_h_to_4h.bias') != -1: | |
| assert data.shape == (config['expansion_ratio'] * config['d_model'],), f'unexpected dim for {tensor_name}' | |
| split_vals = np.split(data, infer_gpu_num, axis=-1) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir + f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| elif tensor_name.find('attention.query_key_value.bias') != -1: | |
| assert data.shape == (3 * config['d_model'],), f'unexpected dim for {tensor_name}' | |
| data = data.reshape(3, config['d_model']) | |
| split_vals = np.split(data, infer_gpu_num, axis=-1) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| elif tensor_name.find('attention.query_key_value.weight') != -1: | |
| assert data.shape == (3 * config['d_model'], config['d_model']), f'unexpected dim for {tensor_name}' | |
| data = data.T | |
| data = data.reshape(config['d_model'], 3, config['d_model']) | |
| split_vals = np.split(data, infer_gpu_num, axis=-1) | |
| for j in range(infer_gpu_num): | |
| save_path = os.path.join(save_dir, f'model.{tensor_name}.{j}.bin') | |
| split_vals[j].tofile(save_path) | |
| if config['no_bias']: | |
| _write_zero_bias(tensor_name, save_path, (3, split_vals[j].shape[-1]), np_weight_data_type) | |
| else: | |
| raise RuntimeError(f'Tensor with name {tensor_name} is not handled') | |
| def convert_and_save_ft_weights(named_params: dict, config: dict, infer_gpu_num: int=1, weight_data_type: str='fp32', save_dir: str='') -> None: | |
| """Convert a Composer MPT checkpoint to a FasterTransformer format. | |
| Args: | |
| named_params (Dict[str, Parameter]): A dictionary containing the Composer MPT model's parameter names and data. | |
| config (Dict[str, Any]): Configuration for the model. This is used in getting model specific parameters. | |
| infer_gpu_num (int): The number of gpus you are planning to use for inference. | |
| weight_data_type (str): The dtype of the converted FasterTransformer model. | |
| save_dir (str): Path of the directory to save the weight in FT format. The directory must already exist. | |
| Returns: | |
| None: Writes to the `save_dir` folder. File names within this folder are based on the model parameter names. | |
| """ | |
| np_weight_data_type = _get_weight_data_type(weight_data_type) | |
| param_remapping = {'norm_1.bias': 'input_layernorm.bias', 'norm_1.weight': 'input_layernorm.weight', 'attn.Wqkv.bias': 'attention.query_key_value.bias', 'attn.Wqkv.weight': 'attention.query_key_value.weight', 'attn.out_proj.bias': 'attention.dense.bias', 'attn.out_proj.weight': 'attention.dense.weight', 'norm_2.bias': 'post_attention_layernorm.bias', 'norm_2.weight': 'post_attention_layernorm.weight', 'ffn.up_proj.bias': 'mlp.dense_h_to_4h.bias', 'ffn.up_proj.weight': 'mlp.dense_h_to_4h.weight', 'ffn.down_proj.bias': 'mlp.dense_4h_to_h.bias', 'ffn.down_proj.weight': 'mlp.dense_4h_to_h.weight'} | |
| for name, param in named_params.items(): | |
| log.debug(f'Working on parameter {name} ...') | |
| data = param.detach().cpu().numpy().astype(np_weight_data_type) | |
| if name.find('weight') == -1 and name.find('bias') == -1: | |
| log.debug(f'found a parameter name that is not handled: {name}') | |
| continue | |
| if name == 'transformer.wpe.weight': | |
| assert data.shape == (config['max_seq_len'], config['d_model']), f'unexpected dim for {name}' | |
| data.tofile(os.path.join(save_dir, 'model.wpe.bin')) | |
| elif name == 'transformer.wte.weight': | |
| assert data.shape == (config['vocab_size'], config['d_model']), f'unexpected dim for {name}' | |
| data.tofile(os.path.join(save_dir, 'model.wte.bin')) | |
| elif name == 'transformer.norm_f.bias': | |
| assert data.shape == (config['d_model'],), f'unexpected dim for {name}' | |
| data.tofile(os.path.join(save_dir, 'model.final_layernorm.bias.bin')) | |
| elif name == 'transformer.norm_f.weight': | |
| assert data.shape == (config['d_model'],), f'unexpected dim for {name}' | |
| save_path = os.path.join(save_dir, 'model.final_layernorm.weight.bin') | |
| data.tofile(save_path) | |
| if config['no_bias']: | |
| _write_zero_bias(name, save_path, data.shape[-1], np_weight_data_type) | |
| elif name == 'transformer.lm_head.weight': | |
| data.tofile(os.path.join(save_dir, 'model.lm_head.weight.bin')) | |
| else: | |
| for mpt_pattern, ft_pattern in param_remapping.items(): | |
| if name.find(mpt_pattern) != -1: | |
| new_name = name.replace('transformer.blocks.', 'layers.').replace(mpt_pattern, ft_pattern) | |
| _convert_weight_to_ft_each(save_dir, infer_gpu_num, new_name, config, data, np_weight_data_type) |