Perceptrix/finetune/inference/convert_composer_mpt_to_ft.py

# Copyright 2022 MosaicML LLM Foundry authors
# SPDX-License-Identifier: Apache-2.0

# Note: This script is specifically for converting MPT Composer checkpoints to FasterTransformer format.

import configparser
import os
import tempfile
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import Any, Dict, Optional, Union

import torch
from composer.utils import get_file, safe_torch_load
from transformers import PreTrainedTokenizer

from llmfoundry.utils import (convert_and_save_ft_weights,
                              get_hf_tokenizer_from_composer_state_dict)


def save_ft_config(composer_config: Dict[str, Any],
                   tokenizer: PreTrainedTokenizer,
                   save_dir: str,
                   infer_gpu_num: int = 1,
                   weight_data_type: str = 'fp32',
                   force: bool = False):

    config = configparser.ConfigParser()
    config['gpt'] = {}
    try:
        config['gpt']['model_name'] = 'mpt'
        config['gpt']['head_num'] = str(composer_config['n_heads'])
        n_embd = composer_config['d_model']
        config['gpt']['size_per_head'] = str(n_embd //
                                             composer_config['n_heads'])
        config['gpt']['inter_size'] = str(n_embd * composer_config['mlp_ratio'])
        config['gpt']['max_pos_seq_len'] = str(composer_config['max_seq_len'])
        config['gpt']['num_layer'] = str(composer_config['n_layers'])
        config['gpt']['vocab_size'] = str(composer_config['vocab_size'])
        config['gpt']['start_id'] = str(tokenizer.bos_token_id)
        config['gpt']['end_id'] = str(tokenizer.eos_token_id)
        config['gpt']['weight_data_type'] = weight_data_type
        config['gpt']['tensor_para_size'] = str(infer_gpu_num)
        # nn.LayerNorm default eps is 1e-5
        config['gpt']['layernorm_eps'] = str(1e-5)
        if composer_config['alibi']:
            config['gpt']['has_positional_encoding'] = str(False)
            config['gpt']['use_attention_linear_bias'] = str(True)
        if composer_config['attn_clip_qkv'] and not force:
            raise RuntimeError(
                'clip_qkv is enabled for this MPT model. This may not work as expected in FT. Use --force to force a conversion.'
            )
        if composer_config['attn_qk_ln'] and not force:
            raise RuntimeError(
                'qk_ln is enabled for this MPT model. This may not work as expected in FT. Use --force to force a conversion.'
            )

        with open(os.path.join(save_dir, 'config.ini'), 'w') as configfile:
            config.write(configfile)
        return config
    except:
        print(f'Failed to save the config in config.ini.')
        raise


def write_ft_checkpoint_from_composer_checkpoint(
        checkpoint_path: Union[Path, str],
        infer_gpu_num: int,
        save_dir: str,
        output_precision: str = 'fp32',
        local_checkpoint_save_location: Optional[Union[Path,
                                                       str]] = None) -> None:
    """Convert a Composer checkpoint to a FasterTransformer checkpoint folder.

    .. note:: This function may not work properly if you used surgery algorithms when you trained your model. In that case you may need to
        edit the parameter conversion methods to properly convert your custom model.

    Args:
        checkpoint_path (Union[Path, str]): Path to the composer checkpoint, can be a local path, or a remote path beginning with ``s3://``, or another backend
            supported by Composer.
        infer_gpu_num (int): The number of gpus you are planning to use for inference.
        save_dir (str): Path of the directory to save the checkpoint in FT format.
        output_precision (str, optional): The precision of the output weights saved to the FasterTransformer model. Can be either ``fp32`` or ``fp16``.
        local_checkpoint_save_location (Optional[Union[Path, str]], optional): If specified, where to save the checkpoint file to locally.
                                                                                If the input ``checkpoint_path`` is already a local path, this will be a symlink.
                                                                   Defaults to None, which will use a temporary file.
    """
    dtype = {
        'fp32': torch.float32,
        'fp16': torch.float16,
    }[output_precision]

    # default local path to a tempfile if path is not provided
    if local_checkpoint_save_location is None:
        tmp_dir = tempfile.TemporaryDirectory()
        local_checkpoint_save_location = Path(
            tmp_dir.name) / 'local-composer-checkpoint.pt'

    # download the checkpoint file
    print(
        f'Downloading checkpoint from {checkpoint_path} -> {local_checkpoint_save_location}'
    )
    get_file(str(checkpoint_path), str(local_checkpoint_save_location))

    # Load the Composer checkpoint. Use it to get the
    # Composer state dict and weights
    print('Loading checkpoint into CPU RAM...')
    composer_state_dict = safe_torch_load(local_checkpoint_save_location)

    # Extract Composer config from state dict
    if 'state' not in composer_state_dict:
        raise RuntimeError(
            f'"state" is not an available key in the provided composer checkpoint. Is {local_checkpoint_save_location} ill-formed?'
        )
    if 'integrations' not in composer_state_dict[
            'state'] or 'huggingface' not in composer_state_dict['state'][
                'integrations']:
        raise RuntimeError(
            'Did not find HuggingFace related state (e.g., tokenizer) in the provided composer checkpoint!'
        )
    composer_config = composer_state_dict['state']['integrations'][
        'huggingface']['model']['config']['content']

    # Extract the HF tokenizer
    print('#' * 30)
    print('Extracting HF Tokenizer...')
    hf_tokenizer = get_hf_tokenizer_from_composer_state_dict(
        composer_state_dict)
    if hf_tokenizer is None:
        print('Warning! No HF Tokenizer found!')

    # Extract the model weights
    weights_state_dict = composer_state_dict['state']['model']
    torch.nn.modules.utils.consume_prefix_in_state_dict_if_present(
        weights_state_dict, prefix='model.')

    # Converting weights to desired dtype
    for k, v in weights_state_dict.items():
        if isinstance(v, torch.Tensor):
            weights_state_dict[k] = v.to(dtype=dtype)

    # Convert the weights using the config and tokenizer to FasterTransformer format
    print('#' * 30)
    print('Saving FasterTransformer config...')
    save_ft_config(composer_config,
                   tokenizer=hf_tokenizer,
                   save_dir=save_dir,
                   weight_data_type=output_precision)
    print('#' * 30)
    print('Converting weights to FasterTransformer format...')
    convert_and_save_ft_weights(named_params=weights_state_dict,
                                config=composer_config,
                                infer_gpu_num=infer_gpu_num,
                                weight_data_type=output_precision,
                                save_dir=save_dir)

    print('#' * 30)
    print(
        f'FasterTransformer checkpoint folder successfully created at {save_dir}.'
    )

    print('Done.')
    print('#' * 30)


def parse_args() -> Namespace:
    """Parse commandline arguments."""
    parser = ArgumentParser(
        description=
        'Convert an MPT Composer checkpoint into a standard FasterTransformer checkpoint folder.'
    )
    parser.add_argument(
        '--composer_path',
        '-i',
        type=str,
        help='Composer checkpoint path. Can be a local file path or cloud URI',
        required=True)
    parser.add_argument(
        '--local_checkpoint_save_location',
        type=str,
        help='If specified, where to save the checkpoint file to locally. \
                            If the input ``checkpoint_path`` is already a local path, this will be a symlink. \
                            Defaults to None, which will use a temporary file.',
        default=None)
    parser.add_argument(
        '--ft_save_dir',
        '-o',
        type=str,
        help='Directory to save FasterTransformer converted checkpoint in',
        required=True)
    parser.add_argument('--infer_gpu_num',
                        '-i_g',
                        type=int,
                        help='How many gpus for inference?',
                        required=True)
    parser.add_argument(
        '--force',
        action='store_true',
        help=
        'Force conversion to FT even if some features may not work as expected in FT'
    )
    parser.add_argument(
        '--output_precision',
        type=str,
        help=
        'Data type of weights in the FasterTransformer output model. Input checkpoint weights will be converted to this dtype.',
        choices=['fp32', 'fp16'],
        default='fp32')

    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    print('\n=============== Argument ===============')
    for key in vars(args):
        print('{}: {}'.format(key, vars(args)[key]))
    print('========================================')

    save_dir = os.path.join(args.ft_save_dir, f'{args.infer_gpu_num}-gpu')

    if os.path.exists(save_dir) == False:
        os.makedirs(save_dir)
    else:
        raise RuntimeError(f'Output path {save_dir} already exists!')

    write_ft_checkpoint_from_composer_checkpoint(
        checkpoint_path=args.composer_path,
        infer_gpu_num=args.infer_gpu_num,
        save_dir=save_dir,
        output_precision=args.output_precision,
        local_checkpoint_save_location=args.local_checkpoint_save_location)