def write_ft_checkpoint_from_composer_checkpoint(
    checkpoint_path: Union[Path, str],
    infer_gpu_num: int,
    save_dir: str,
    trust_remote_code: bool,
    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.
        trust_remote_code (bool): Whether or not to use code outside of the transformers module.
        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(  # pyright: ignore
        composer_state_dict,
        trust_remote_code,
    )
    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(  # pyright: ignore
        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...')
    assert hf_tokenizer is not None
    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)
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.
        trust_remote_code (bool): Whether or not to use code outside of the transformers module.
        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.
    
scripts/inference/convert_composer_mpt_to_ft.py