Training conformer interCTC with init_from_nemo_model argument is broke 
State: closed | Pull request: no
Labels: name: bug

name: stale

Hi folks!

I have an already trained model without interCTC, which I want to initialize from and continue training, this time enabling interCTC loss.
Here is my training command:
```
CUDA_VISIBLE_DEVICES=0 python speech_to_text_ctc_bpe.py \
--config-path=conf/conformer/ \
--config-name=conformer_ctc_bpe exp_manager.name="conformer-inter-ctc" \
exp_manager.exp_dir=${EXP} \
exp_manager.resume_ignore_no_checkpoint=true \
exp_manager.resume_if_exists=true \
model.tokenizer.dir="${TOKENIZER}" \
model.train_ds.manifest_filepath="${TRAIN_MANIFEST}" \
model.validation_ds.manifest_filepath="${VAL_MANIFEST}" \
init_from_nemo_model="${INIT_CKPT}"
```
Here is an exception I get:
```
RuntimeError: Intermediate layer 8 was not captured! Check if length of model.encoder.captured_layer_outputs matches length of model.intermediate_loss_weights properties.
```

This looks like a bug.

Any suggestions on how to get this working without manually hacking model initialization process?

When I launch training from scratch with interCTC everything works fine!

`nemo-toolkit==1.21.0`


which `${INIT_CKPT}` did you use?  could you also provide the full error log?

for `${INIT_CKPT}` I use `.nemo` checkpoint trained with same model config, the only difference in the config is that now iterCTC is enabled.
```
interctc:
     loss_weights: [0.3]
     apply_at_layers: [8]
```

here is the full stack trace:

```l
Traceback (most recent call last):
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py", line 188, in get_captured_interctc_tensors
    layer_outputs = total_registry[f"interctc/layer_output_{layer_idx}"]
KeyError: 'interctc/layer_output_8'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ubuntu/repos/transcriber/speech_to_text_ctc_bpe.py", line 114, in <module>
    main()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/core/config/hydra_runner.py", line 126, in wrapper
    _run_hydra(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/home/ubuntu/repos/transcriber/speech_to_text_ctc_bpe.py", line 106, in main
    trainer.fit(asr_model)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1023, in _run_stage
    self.fit_loop.run()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
    self.advance()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 355, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
    self.advance(data_fetcher)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 219, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 188, in run
    self._optimizer_step(kwargs.get("batch_idx", 0), closure)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 266, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 145, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1270, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 161, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 257, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 231, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 73, in optimizer_step
    return super().optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 116, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 68, in wrapper
    return wrapped(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 373, in wrapper
    out = func(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/optim/optimizer.py", line 76, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/optim/adamw.py", line 161, in step
    loss = closure()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 103, in _wrap_closure
    closure_result = closure()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 142, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 128, in closure
    step_output = self._step_fn()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 315, in _training_step
    training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 293, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 330, in training_step
    return self.model(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1519, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1355, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/pytorch_lightning/overrides/base.py", line 90, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/utils/model_utils.py", line 381, in wrap_training_step
    output_dict = wrapped(*args, **kwargs)
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/collections/asr/models/ctc_models.py", line 584, in training_step
    loss_value, tensorboard_logs = self.add_interctc_losses(
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py", line 247, in add_interctc_losses
    captured_tensors = self.get_captured_interctc_tensors()
  File "/home/ubuntu/repos/transcriber/venv/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py", line 191, in get_captured_interctc_tensors
    raise RuntimeError(
RuntimeError: Intermediate layer 8 was not captured! Check if length of model.encoder.captured_layer_outputs matches length of model.intermediate_loss_weights properties.
```

Alternatively I can initialize  just fine, even with interCTC changes in the config with `+init_from_ptl_model` argument when I pass equivalent checkpoint but saved in pytorch lightning format, i.e `.ckpt` instead of `.nemo`

Thanks for the information, I'll take a look and get back to you~

@osyka-m Thanks for pointing out the bug~! We have fixed it in this PR: https://github.com/NVIDIA/NeMo/pull/8536

Great to hear the issue if fixed!
Haven't expected that it affects 19 sources though.
Anyways, thank you for addressing the issue promptly!

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.