How can I convert .nemo file trained with TP parallel enabled into Hugging Face format?

[Lagyu]

I have a .nemo file that cannot be converted into Hugging Face model with convert_llama_nemo_to_hf.py (with this PR change).

The .nemo file was trained like this:

torchrun --nproc_per_node=4 \
    /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py \
    exp_manager.exp_dir=${EXP_DIR} \
    exp_manager.name=${EXP_NAME} \
    exp_manager.create_wandb_logger=${WANDB} \
    exp_manager.wandb_logger_kwargs.project=${PJ_NAME} \
    exp_manager.wandb_logger_kwargs.name=${EXP_NAME} \
    exp_manager.early_stopping_callback_params.patience=5 \
    trainer.precision=bf16 \
    trainer.devices=4 \
    trainer.num_nodes=1 \
    trainer.max_epochs=-1 \
    trainer.val_check_interval=0.05 \
    trainer.gradient_clip_val=1.0 \
    model.megatron_amp_O2=False \
    ++model.mcore_gpt=True \
    model.restore_from_path=${MODEL} \
    model.peft.peft_scheme=lora \
    model.data.train_ds.file_names=${TRAIN_DS} \
    model.data.train_ds.concat_sampling_probabilities=[1.0] \
    model.data.validation_ds.file_names=${VALID_DS} \
    model.tensor_model_parallel_size=2 \
    model.pipeline_model_parallel_size=2 \
    model.global_batch_size=16 \
    model.micro_batch_size=4 \
    model.data.validation_ds.tokens_to_generate=20 \
    model.data.validation_ds.metric.name="loss" \
    model.data.train_ds.num_workers=0 \
    model.data.validation_ds.num_workers=0 \
    model.optim.lr=1e-5

And can be extracted like this:

$ mkdir extracted
$ tar -xf lora-finetuned-model.nemo -C ./extracted
$ cd extracted
model_config.yaml  tp_rank_00_pp_rank_000  tp_rank_00_pp_rank_001  tp_rank_01_pp_rank_000  tp_rank_01_pp_rank_001

The error I get when I run convert_llama_nemo_to_hf.py is like this:

$ python "/opt/NeMo/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py" --input_name_or_path="${NEMO_FILE_DIR}" --output_path="${OUTPUT_DIR}/${EXPERIMENT_NAME}.bin" --hf_input_path="${BASE_LLAMA3_MODEL}" --hf_output_path="${OUTPUT_DIR}/${EXPERIMENT_NAME}" --input_tokenizer="${BASE_LLAMA3_MODEL}" --hf_output_tokenizer="${OUTPUT_DIR}/${EXPERIMENT_NAME}"
[NeMo I 2024-08-21 08:55:57 megatron_init:269] Rank 0 has data parallel group : [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:275] Rank 0 has combined group of data parallel and context parallel : [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:280] All data parallel group ranks with context parallel combined: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:283] Ranks 0 has data parallel rank: 0
[NeMo I 2024-08-21 08:55:57 megatron_init:291] Rank 0 has context parallel group: [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:294] All context parallel group ranks: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:295] Ranks 0 has context parallel rank: 0
[NeMo I 2024-08-21 08:55:57 megatron_init:302] Rank 0 has model parallel group: [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:303] All model parallel group ranks: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:312] Rank 0 has tensor model parallel group: [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:316] All tensor model parallel group ranks: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:317] Rank 0 has tensor model parallel rank: 0
[NeMo I 2024-08-21 08:55:57 megatron_init:337] Rank 0 has pipeline model parallel group: [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:349] Rank 0 has embedding group: [0]
[NeMo I 2024-08-21 08:55:57 megatron_init:355] All pipeline model parallel group ranks: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:356] Rank 0 has pipeline model parallel rank 0
[NeMo I 2024-08-21 08:55:57 megatron_init:357] All embedding group ranks: [[0]]
[NeMo I 2024-08-21 08:55:57 megatron_init:358] Rank 0 has embedding rank: 0

<omit many warnings about keys...>

Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1
----------------------------------------------------------------------------------------------------
distributed_backend=gloo
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

[rank0]: Traceback (most recent call last):
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py", line 289, in <module>
[rank0]:     dtype = convert(args.input_name_or_path, args.output_path, precision=args.precision, cpu_only=args.cpu_only)
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py", line 127, in convert
[rank0]:     model = MegatronGPTModel.restore_from(
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/models/nlp_model.py", line 478, in restore_from
[rank0]:     return super().restore_from(
[rank0]:   File "/opt/NeMo/nemo/core/classes/modelPT.py", line 468, in restore_from
[rank0]:     instance = cls._save_restore_connector.restore_from(
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1322, in restore_from
[rank0]:     assert os.path.isdir(tmp_model_weights_dir), f'Expected {tmp_model_weights_dir} to be a directory.'
[rank0]: AssertionError: Expected /tmp/tmpy207fno1/model_weights to be a directory.

How can I convert this .nemo file into HF format?
Thank you in advance for your cooperation.

[Lagyu]

I found the script, convert_llama_nemo_to_hf.py, is only for sft nemo file, and we have convert_nemo_to_canonical.py for lora trained .nemo file. (and only tp=1 and pp=1 is supported for now)
I will try to use it.