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.
        np_weight_data_type (np.dtype): Data type of the numpy array `data`.

    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,
            )  # pyright: ignore [reportGeneralTypeIssues]

    elif tensor_name.find('attention.dense.weight') != -1:
        assert data.shape == (
            config['d_model'],
            config['d_model'],
        ), f'unexpected dim for {tensor_name}'
        # nn.Linear weights are transposed
        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,
            )  # pyright: ignore [reportGeneralTypeIssues]

    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}'
        # nn.Linear weights are transposed
        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,
            )  # pyright: ignore [reportGeneralTypeIssues]

    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}'
        # nn.Linear weights are transposed
        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,
                )  # pyright: ignore [reportGeneralTypeIssues]

    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}'
        # nn.Linear weights are transposed
        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,
                )  # pyright: ignore [reportGeneralTypeIssues]

    else:
        raise RuntimeError(f'Tensor with name {tensor_name} is not handled')
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.
        np_weight_data_type (np.dtype): Data type of the numpy array `data`.

    Returns:
        None: Writes to a file in `save_dir`. File name is based on the `tensor_name`
    
llmfoundry/utils/checkpoint_conversion_helpers.py