Fine-tuning ASR Lightning Error

[gcervantes8]

Describe the bug

I'm getting this error when trying to fine-tune the Parakeet model using the /examples/asr/speech_to_text_finetune.py file.

Exception has occurred: TypeError
`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `EncDecCTCModelBPE`
  File "/Users/gerrycervantes/code/NeMo/examples/asr/speech_to_text_finetune.py", line 221, in main
    trainer.fit(asr_model)
  File "/Users/gerrycervantes/code/NeMo/examples/asr/speech_to_text_finetune.py", line 225, in <module>
    main()  # noqa pylint: disable=no-value-for-parameter
    ^^^^^^
TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `EncDecCTCModelBPE`

Steps/Code to reproduce bug

Command
This is the command I use for fine-tuning

python examples/asr/speech_to_text_finetune.py --config-path=conf/asr_finetune --config-name=asr_custom

Data
So I've downloaded the Librispeech Datasets and processed it using /scripts/dataset_processing/get_librispeech_data.py

Model
I've also downloaded the models from HuggingFace. https://huggingface.co/nvidia/parakeet-ctc-0.6b

Config
I've done a few changes to the configuration file so that it fits my environment. I'm on a Mac computer, so I've tried running with CPU and MPS, but get the error with both ways.

This is the corresponding config YAML file, resembling the example config file.

name: "ASR_Custom_Finetuning"

# use `init_from_nemo_model` or `init_from_pretrained_model` to initialize the model
# We do not currently support `init_from_ptl_ckpt` to create a single script for all types of models.
init_from_nemo_model: models/parakeet-ctc-1.1b.nemo # path to nemo model
# init_from_pretrained_model: nvidia/parakeet-ctc-0.6b # path to nemo model

model:
  sample_rate: 16000

  train_ds:
    manifest_filepath: data/TRAIN_CLEAN_100.json
    sample_rate: ${model.sample_rate}
    batch_size: 16 # you may increase batch_size if your memory allows
    shuffle: true
    num_workers: 8
    pin_memory: true
    max_duration: 20
    min_duration: 0.1
    # tarred datasets
    is_tarred: false
    tarred_audio_filepaths: null
    shuffle_n: 2048
    # bucketing params
    bucketing_strategy: "fully_randomized"
    bucketing_batch_size: null

  validation_ds:
    manifest_filepath: data/dev_clean.json
    sample_rate: ${model.sample_rate}
    batch_size: 16
    shuffle: false
    use_start_end_token: false
    num_workers: 8
    pin_memory: true

  test_ds:
    manifest_filepath: null
    sample_rate: ${model.sample_rate}
    batch_size: 16
    shuffle: false
    use_start_end_token: false
    num_workers: 8
    pin_memory: true
  
  char_labels: # use for char based models
    update_labels: false
    labels: null # example list config: \[' ', 'a', 'b', 'c'\]

  tokenizer: # use for spe/bpe based tokenizer models
    update_tokenizer: false
    dir: null  # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe)
    type: bpe  # Can be either bpe (SentencePiece tokenizer) or wpe (WordPiece tokenizer)

  spec_augment:
    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
    freq_masks: 2 # set to zero to disable it
    time_masks: 10 # set to zero to disable it
    freq_width: 27
    time_width: 0.05

  optim:
    name: adamw
    lr: 1e-4
    # optimizer arguments
    betas: [0.9, 0.98]
    weight_decay: 1e-3

    # scheduler setup
    sched:
      name: CosineAnnealing
      # scheduler config override
      warmup_steps: 5000
      warmup_ratio: null
      min_lr: 5e-6

trainer:
  devices: 4 # number of GPUs, -1 would use all available GPUs
  num_nodes: 1
  max_epochs: 50
  max_steps: -1 # computed at runtime if not set
  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
  accelerator: cpu
  strategy:
    _target_: lightning.pytorch.strategies.SingleDeviceStrategy
    # _target_: lightning.pytorch.strategies.DDPStrategy
    gradient_as_bucket_view: true
  accumulate_grad_batches: 1
  gradient_clip_val: 0.0
  precision: 32 # 16, 32, or bf16
  log_every_n_steps: 10  # Interval of logging.
  enable_progress_bar: True
  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
  sync_batchnorm: true
  enable_checkpointing: False  # Provided by exp_manager
  logger: false  # Provided by exp_manager
  benchmark: false # needs to be false for models with variable-length speech input as it slows down training


exp_manager:
  exp_dir: null
  name: ${name}
  create_tensorboard_logger: true
  create_checkpoint_callback: true
  checkpoint_callback_params:
    # in case of multiple validation sets, first one is used
    monitor: "val_wer"
    mode: "min"
    save_top_k: 5
    always_save_nemo: True # saves the checkpoints as nemo files along with PTL checkpoints
  resume_if_exists: false
  resume_ignore_no_checkpoint: false

  create_wandb_logger: false
  wandb_logger_kwargs:
    name: null
    project: null

Environment overview (please complete the following information)

• Environment location: Bare-metal
• Method of NeMo install: From source (git commit 34f7408)

conda init --all

conda create --name nemo python==3.11.1
conda activate nemo
conda install pytorch::pytorch torchvision torchaudio -c pytorch

brew install mecab
conda install -c conda-forge pynini
pip install cython packaging

cd ~/code
git clone https://github.com/NVIDIA/NeMo.git
cd NeMo
pip install 'nemo_toolkit[asr]'

Environment details

• macOS: Sonoma 14.6.1
• PyTorch: 2.5.1
• Python version: 3.11.8
• lightning 2.4.0 pyhd8ed1ab_0 conda-forge
• lightning-utilities 0.11.9 pyhff2d567_0 conda-forge

[gcervantes8]

I was able to solve this by pulling the Nemo code from the 2.0.0 tag, and using that. I believe a mismatch in source code and the nemo library caused this. Probably because of the recent commit that changed from 'pytorch_lightning' -> 'lightning.pytorch'.