Add Composer MPT to FasterTransformer Conversion Script

[nik-mosaic]

Directly convert a Composer MPT model to the FasterTransformer format, without an intermediate conversion to HuggingFace.

Associated task: CO-2277

The amount of CPU RAM required to run this script is approximately equal to the size of the MPT checkpoint.pt file.

Example:

Command:python convert_composer_mpt_to_ft.py --composer_path=<folder> --ft_save_dir='ft' --infer_gpu_num=1 --output_precision='fp16'

Output:

=============== Argument ===============
composer_path: <folder>
local_checkpoint_save_location: None
ft_save_dir: ft
infer_gpu_num: 1
force: False
output_precision: fp16
========================================
Downloading checkpoint from <folder> -> /tmp/<folder>/local-composer-checkpoint.pt
Loading checkpoint into CPU RAM...
##############################
Extracting HF Tokenizer...
##############################
Saving FasterTransformer config...
##############################
Converting weights to FasterTransformer format...
Working on parameter transformer.blocks.0.attn.Wqkv.weight ...
zero bias for weight: layers.0.attention.query_key_value.weight
...<Many lines of logs>...
Working on parameter transformer.wte.weight ...
##############################
FasterTransformer checkpoint folder successfully created at ft/1-gpu.
Done.
##############################

<TODO: Add image of output comparison>

[margaretqian]

looks good to me, but would appreciate if someone more familiar with the conversion scripts could also take a look!

[dakinggg]

Can we pull out the shared code here between the two FT conversion scripts and put it in llmfoundry proper? This will also allow us to more easily write a callback that saves these checkpoints during training.

[dakinggg]

LGTM

[dakinggg]

@nik-mosaic want to merge?