How can I start finetuing/training from checkpoint?

[QXmX29]

First I finetuned basic model omniASR_LLM_1B_v2 like this:

cd /path/to/omnilingual-asr
CUDA_VISIBLE_DEVICES=2,3,4 torchrun --nproc_per_node=3 \
-m workflows.recipes.wav2vec2.asr $OUTPUT_DIR \
--config-file $CONFIG_HOME/my-llm-finetune.yaml

After that I get checkpoints located in path/to/checkpoints, like step_800/model/pp_00/tp_00 contains sdp_00.pt, sdp_01.pt and sdp_02.pt.

Now I want to start a new round of finetuning based on a certain checkpoint, rather than the original model like omniASR_LLM_1B_v2. So I modifed the my-llm-finetune.yaml as below. I made 2 attempts, one adding model.path while the other adding pretrained_encoder.name/path/family/dtype.

model:
  name: "omniASR_LLM_1B_v2"
  path: /path/to/checkpoints/step_800/consolidated.pt
  family: wav2vec2_llama
  dtype: bfloat16  # float32 # bfloat16

pretrained_encoder:
  # name: "LLM_1B_v2_800step-260113_1807"
  # path: /path/to/checkpoints/step_800/consolidated.pt
  # family: wav2vec2_llama
  # dtype: bfloat16
  mmap: true

Complete config.yaml: config-onlymodifiedmodel.yaml

However, I got error information:

WARNING  fairseq2 - Both `model.name` and `model.path` are specified. `model.path` takes         
                             precedence.                                                                             
ERROR    fairseq2 - Recipe failed due to an unexpected error. See logged stack trace for details 
         and file a bug report to the corresponding author.                                      
         Traceback (most recent call last):                                                      
           File                                                                                  
         "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/recipe/internal/
         model.py", line 212, in _load_custom_model                                              
             config = family.config_kls()                                                        
         TypeError: Wav2Vec2LlamaConfig.__init__() missing 2 required keyword-only arguments:    
         'wav2vec2_asr_config' and 'llama_config'                                                
                                                                                                 
         The above exception was the direct cause of the following exception:                    
                                                                                                                     
         Traceback (most recent call last):                                                      
           ...                                         
           File                                                                                  
         "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/recipe/compositi
         on/model.py", line 32, in get_model                                                     
             return model_provider.get()                                                         
           File                                                                                  
         "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/recipe/internal/
         model.py", line 63, in get                                                              
             model = self._bootstrapper.bootstrap()                                              
           File                                                                                  
         "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/recipe/internal/
         model.py", line 111, in bootstrap                                                       
             return self._load_custom_model()                                                    
           File                                                                                  
         "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/recipe/internal/
         model.py", line 214, in _load_custom_model                                              
             raise InternalError(                                                                
         fairseq2.error.InternalError: Default configuration of the                              
         <fairseq2.models.family.StandardModelFamily object at 0x7fb98bd02560> model family      
         cannot be constructed. 

More details: rank_0.log

---

Maybe it's worth noting that I had successfully merged them into a single file step_800/consolidated.pt, loaded this model and got ASRInferencePipeline using my code load_model_from.py, which proves it workable to load model from user's checkpoints:

# ...
tk = load_tokenizer("omniASR_tokenizer_written_v2")
pipeline = ASRInferencePipeline(model_card=None, model=model, tokenizer=tk)

My codes:

merge_sharded_ckpt.py

load_model_from.py

[cirquit]

Hi,

If I understand correctly, you want to keep training an already fine-tuned model. I think you might be calling the wrong recipe configuration, as the pretrained_encoder configuration should be entirely disabled in this scenario, similar to llm-finetune.py.

To load an already trained model and use as a starting point, you should be able to point the model configuration to the model checkpoint as follows:

model: 
  path: "path-to-checkpoint-dir/checkpoints/step_1000/model/"
  family: wav2vec2_llama
  model_arch: 1b

You can also do this without consolidation for inference, for example:

from omnilingual_asr.models.wav2vec2_llama import get_wav2vec2_llama_model_hub
from pathlib import Path

hub = get_wav2vec2_llama_model_hub()
config = hub.get_arch_config("1b")

# Point to the model/ directory containing pp_*/tp_*/sdp_*.pt structure
checkpoint_path = Path("checkpoints/step_1000/model/")

# Load for training/eval - format is auto-detected
model = hub.load_custom_model(checkpoint_path, config)

# Load in the inference pipeline
pipeline = ASRInferencePipeline(model_card=None, model=model, ...)

[QXmX29]

@cirquit Thanks for your patient explanation! But it doesn't work...

Traceback (most recent call last):
  File "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/models/family.py", line 444, in _do_load_model
    set_model_state(model, checkpoint, pr)
  File "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/models/utils/checkpoint.py", line 94, in set_model_state
    raise ModelCheckpointMismatchError(" ".join(errors))
fairseq2.models.utils.checkpoint.ModelCheckpointMismatchError: `text_frontend.weight` has a shape of (10289, 4096) in the checkpoint, but has a shape of (9813, 4096) in the model. `final_proj.weight` has a shape of (10288, 4096) in the checkpoint, but has a shape of (9812, 4096) in the model.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/xxx/projects/pashto/infer_ckpt.py", line 13, in <module>
    model = hub.load_custom_model(checkpoint_path, config)
  File "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/models/hub.py", line 456, in load_custom_model
    model = self._family.load_custom_model(
  File "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/models/family.py", line 411, in load_custom_model
    return self._do_load_model(path, config, gangs, dtype, mmap, restrict, progress)
  File "/root/anaconda3/envs/omnilingual/lib/python3.10/site-packages/fairseq2/models/family.py", line 448, in _do_load_model
    raise ModelCheckpointError(path, msg) from ex
fairseq2.model_checkpoint.loader.ModelCheckpointError: Checkpoint at /path/to/checkpoints/step_600/model is not compatible with the wav2vec2_llama model.

[cirquit]

The error points directly to the issue - fairseq2.models.utils.checkpoint.ModelCheckpointMismatchError: final_proj.weight has a shape of (10288, 4096) in the checkpoint, but has a shape of (9812, 4096) in the model.

This points to a wrong architecture being used, which you need to modify in the recipe preset: model_arch: 1b. Please pick the same model architecture that you trained with, you can find it here: v1 or v2. The architecture definition for the llama-based models is here.

[QXmX29]

I modified the config.yaml (arch: 1b_v2), because the base model is omniASR_LLM_1B_v2.
But this time I got the error as below:

 File "/path/to/omnilingual-asr/workflows/recipes/wav2vec2/asr/recipe.py", line 110, in    
prepare_model                                                                                           
   raise ValueError("Recipe configuration is invalid. Supported configurations: [model.name |          
model.arch + pretrained_encoder.name]")                                                                 
ValueError: Recipe configuration is invalid. Supported configurations: [model.name | model.arch +       
pretrained_encoder.name]

The related codes is here.
It warns that the configuration only support [model.name | model.arch + pretrained_encoder.name]...