1	# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14	from typing import Dict, List, Optional, Tuple
15
16	import torch
17
18	from nemo.core.classes.mixins import AccessMixin
19	from nemo.utils import logging
20
21
22	class InterCTCMixin:
23	    """Adds utilities for computing interCTC loss from https://arxiv.org/abs/2102.03216.
24
25	    To use, make sure encoder accesses ``interctc['capture_layers']``
26	    property in the AccessMixin and registers ``interctc/layer_output_X`` and
27	    ``interctc/layer_length_X`` for all layers that we want to get loss from.
28	    Additionally, specify the following config parameters to set up loss::
29
30	        interctc:
31	            # can use different values
32	            loss_weights: [0.3]
33	            apply_at_layers: [8]
34
35	    Then call
36
37	        * ``self.setup_interctc(ctc_decoder_name, ctc_loss_name, ctc_wer_name)``
38	          in the init method
39	        * ``self.add_interctc_losses`` after computing regular loss.
40	        * ``self.finalize_interctc_metrics(metrics, outputs, prefix="val_")``
41	          in the `multi_validation_epoch_end` method.
42	        * ``self.finalize_interctc_metrics(metrics, outputs, prefix="test_")``
43	          in the `multi_test_epoch_end` method.
44	    """
45
46	    def _process_config_values(self, loss_weights: List[float], apply_at_layers: List[int]):
47	        self.set_interctc_param('intermediate_loss_weights', loss_weights)
48	        self.set_interctc_param('apply_at_layers', apply_at_layers)
49	        self.set_interctc_param('main_loss_weight', 1.0 - sum(loss_weights))
50	        if self.get_interctc_param('main_loss_weight') <= 0.0:
51	            raise ValueError(
52	                "Make sure that sum of intermediate loss weights is < 1.0. "
53	                "Note that we don't do any normalization and assign "
54	                "remaining weight to the regular model loss. "
55	                "E.g., if interctc.loss_weights = [0.1, 0.3], regular "
56	                "loss will have weight of 0.6"
57	            )
58	        self.set_interctc_param('enabled', len(loss_weights) > 0)
59
60	        if len(apply_at_layers) != len(loss_weights):
61	            raise ValueError('Length of interctc.apply_at_layers has to match interctc.loss_weights')
62
63	        # setting up config for AccessMixin that will be checked in encoders to
64	        # log the layers we need
65	        AccessMixin.update_access_cfg(
66	            {'interctc': {'capture_layers': apply_at_layers}}, guid=getattr(self, "model_guid", None)
67	        )
68	        if hasattr(self, "propagate_model_guid"):
69	            self.propagate_model_guid()
70	        else:
71	            logging.warning(
72	                f"Not able to propagate model_guid to the submodules. Make sure to call self.propagate_model_guid() in ModelPT class."
73	            )
74
75	    def setup_interctc(self, decoder_name, loss_name, wer_name):
76	        """Sets up all interctc-specific parameters and checks config consistency.
77
78	        Caller has to specify names of attributes to perform CTC-specific WER,
79	        decoder and loss computation. They will be looked up in the class
80	        state with ``getattr``.
81
82	        The reason we get the names and look up object later is because those
83	        objects might change without re-calling the setup of this class. So
84	        we always want to look up the most up-to-date object instead of
85	        "caching" it here.
86	        """
87	        # registering all parameters in a dictionary to avoid conflicts with
88	        # main class's names
89	        self._interctc_params = {}
90	        interctc_config = self.cfg.get("interctc")
91	        if interctc_config is not None:
92	            # if interctc is in the config, we want to check that it indeed defines
93	            # the required keys and nothing else - that's automatically done by
94	            # matching with keyword arguments in self._process_config_values
95	            self._process_config_values(**interctc_config)
96	            self._interctc_params['decoder_name'] = decoder_name
97	            self._interctc_params['loss_name'] = loss_name
98	            self._interctc_params['wer_name'] = wer_name
99	        else:
100	            self.set_interctc_param('enabled', False)
101
102	    def get_interctc_param(self, param_name):
103	        """Either directly get parameter from ``self._interctc_params`` or
104	        call getattr with the corresponding name.
105	        """
106	        if param_name in ['decoder', 'loss', 'wer']:
107	            return getattr(self, self._interctc_params[param_name + "_name"])
108	        return self._interctc_params[param_name]
109
110	    def set_interctc_param(self, param_name, param_value):
111	        """Setting the parameter to the ``self._interctc_params`` dictionary.
112
113	        Raises an error if trying to set decoder, loss or wer as those should
114	        always come from the main class.
115	        """
116	        if param_name in ['decoder', 'loss', 'wer']:
117	            raise ValueError(
118	                'Cannot set "decoder", "loss" or "wer" as parameters. '
119	                'They are always looked up in the main class state.'
120	            )
121	        self._interctc_params[param_name] = param_value
122
123	    def _verify_setup_was_called(self):
124	        """Can be used to verify if setup_interctc was called."""
125	        if not hasattr(self, '_interctc_params'):
126	            raise RuntimeError(
127	                'self.setup_interctc(ctc_decoder_name, ctc_loss_name, ctc_wer_name) has to be '
128	                'called before InterCTC loss can be used!'
129	            )
130
131	    def is_interctc_enabled(self) -> bool:
132	        """Returns whether interCTC loss is enabled."""
133	        self._verify_setup_was_called()
134	        return self.get_interctc_param('enabled')
135
136	    def set_interctc_enabled(self, enabled: bool):
137	        """Can be used to enable/disable InterCTC manually."""
138	        self._verify_setup_was_called()
139	        if enabled:  # checking if proper config parameters were specified
140	            if len(self.get_interctc_param('intermediate_loss_weights')) == 0:
141	                raise RuntimeError(
142	                    'InterCTC cannot be enabled since interctc.loss_weights was not specified in the config.'
143	                )
144	            if len(self.get_interctc_param('apply_at_layers')) != len(
145	                self.get_interctc_param('intermediate_loss_weights')
146	            ):
147	                raise RuntimeError(
148	                    'InterCTC cannot be enabled, since length of "loss_weights" does not match "apply_at_layers".'
149	                )
150	        self.set_interctc_param('enabled', enabled)
151
152	    def finalize_interctc_metrics(self, metrics: Dict, outputs: List[Dict], prefix: str):
153	        """Finalizes InterCTC WER and loss metrics for logging purposes.
154
155	        Should be called inside ``multi_validation_epoch_end`` (with ``prefix="val_"``) or
156	        ``multi_test_epoch_end`` (with ``prefix="test_"``).
157
158	        Note that ``metrics`` dictionary is going to be updated in-place.
159	        """
160	        if self.is_interctc_enabled():
161	            for layer_idx in self.get_interctc_param('apply_at_layers'):
162	                # assuming that if the first batch logged the metrics, then all batches did
163	                if f"{prefix}inter_ctc_loss_l{layer_idx}" in outputs[0]:
164	                    loss = torch.stack([x[f"{prefix}inter_ctc_loss_l{layer_idx}"] for x in outputs]).mean()
165	                    metrics["log"][f"{prefix}inter_ctc_loss_l{layer_idx}"] = loss
166
167	                if f"{prefix}inter_wer_num_l{layer_idx}" in outputs[0]:
168	                    wer_num = torch.stack([x[f"{prefix}inter_wer_num_l{layer_idx}"] for x in outputs]).sum()
169	                    wer_denom = torch.stack([x[f"{prefix}inter_wer_denom_l{layer_idx}"] for x in outputs]).sum()
170	                    metrics["log"][f"{prefix}inter_wer_l{layer_idx}"] = wer_num / wer_denom
171
172	            if f"{prefix}final_loss" in outputs[0]:
173	                metrics["log"][f"{prefix}final_loss"] = torch.stack([x[f"{prefix}final_loss"] for x in outputs]).mean()
174
175	    def get_captured_interctc_tensors(self) -> List[Tuple[torch.Tensor, torch.Tensor]]:
176	        """Returns a list of captured tensors from encoder: tuples of (output, length).
177
178	        Will additionally apply ``ctc_decoder`` to the outputs.
179	        """
180	        if not self.is_interctc_enabled():
181	            return []
182
183	        # note that we have a loop here, because tensors can be defined from
184	        # submodules of encoder (e.g., that's the case in Jasper)
185	        total_registry = {}
186	        for module_registry in AccessMixin.get_module_registry(self.encoder).values():
187	            for key in module_registry:
188	                if key.startswith("interctc/") and key in total_registry:
189	                    raise RuntimeError(f"layer {key} has been logged multiple times!")
190	            total_registry.update(module_registry)
191	        # if intermediate_loss_weights was set, the encoder has to register
192	        # interctc/layer_output_X and interctc/layer_length_X tensors.
193	        # We need to apply decoder to each of them and compute CTC loss.
194	        captured_tensors = []
195	        for layer_idx in self.get_interctc_param('apply_at_layers'):
196	            try:
197	                layer_outputs = total_registry[f"interctc/layer_output_{layer_idx}"]
198	                layer_lengths = total_registry[f"interctc/layer_length_{layer_idx}"]
199	            except KeyError:
200	                raise RuntimeError(


Bash

find /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr -type f -name "*.py" | xargs grep -l "InterCTC\|inter_ctc\|auxiliary.*loss\|multi.*objective" | head -15


Let me check the local project for specific implementations:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.py" | head -30


Let me search for more relevant model files and check the RNNTLoss variants:


grep -rn "FastemitLambda\|fastemit\|clamp\|sigma\|duration" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/ --include="*.py" | head -20


/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/bce_loss.py:108:                norm_weight_norm = torch.clamp(norm_weight_norm, min=0.05, max=1.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:264:        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:265:        clamp = loss_kwargs.pop('clamp', -1.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:266:        loss_func = RNNTLossNumba(blank=blank_idx, reduction='none', fastemit_lambda=fastemit_lambda, clamp=clamp)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:274:        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:275:        clamp = loss_kwargs.pop('clamp', -1.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:276:        big_blank_durations = loss_kwargs.pop('big_blank_durations', None)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:277:        sigma = loss_kwargs.pop('sigma', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:280:            big_blank_durations=big_blank_durations,
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:282:            fastemit_lambda=fastemit_lambda,
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:283:            clamp=clamp,
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:284:            sigma=sigma,
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:289:        big_blank_durations = loss_kwargs.pop('big_blank_durations', None)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:290:        sigma = loss_kwargs.pop('sigma', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:292:            blank=blank_idx, big_blank_durations=big_blank_durations, reduction='none', sigma=sigma
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:297:        fastemit_lambda = loss_kwargs.pop('fastemit_lambda', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:298:        clamp = loss_kwargs.pop('clamp', -1.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:299:        durations = loss_kwargs.pop('durations', None)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:300:        sigma = loss_kwargs.pop('sigma', 0.0)
/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py:304:            durations=durations,


Now let me search for any language ID or speaker verification losses:


find /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr -type f -name "*.py" | xargs grep -l "language\|speaker.*loss\|lid\|verification" | head -10


Let me check the models directory more thoroughly for loss implementations:


grep -rn "def.*loss\|self.loss\|auxiliary\|multi.*loss" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py | head -30


405:            loss_value = self.loss(
409:            # Add auxiliary losses, if registered
410:            loss_value = self.add_auxiliary_losses(loss_value)
439:            # Add auxiliary losses, if registered
440:            loss_value = self.add_auxiliary_losses(loss_value)
540:                loss_value = self.loss(


Perfect! Let me check the add_auxiliary_losses method and related code:


grep -n "add_auxiliary_losses\|auxiliary.*loss\|class.*Model" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py | head -20


1	# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14
15	import copy
16	from typing import Any, List, Optional, Union
17
18	import torch
19	from lightning.pytorch import Trainer
20	from omegaconf import DictConfig, OmegaConf, open_dict
21
22	from nemo.collections.asr.data.audio_to_text_dali import DALIOutputs
23	from nemo.collections.asr.losses.ctc import CTCLoss
24	from nemo.collections.asr.metrics.wer import WER
25	from nemo.collections.asr.models.rnnt_models import EncDecRNNTModel
26	from nemo.collections.asr.parts.mixins import ASRBPEMixin, ASRTranscriptionMixin, InterCTCMixin, TranscribeConfig
27	from nemo.collections.asr.parts.mixins.transcription import TranscriptionReturnType
28	from nemo.collections.asr.parts.preprocessing.segment import ChannelSelectorType
29	from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecoding, CTCDecodingConfig
30	from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
31	from nemo.collections.asr.parts.utils.timestamp_utils import process_timestamp_outputs
32	from nemo.core.classes.common import PretrainedModelInfo
33	from nemo.core.classes.mixins import AccessMixin
34	from nemo.utils import logging, model_utils
35
36
37	class EncDecHybridRNNTCTCModel(EncDecRNNTModel, ASRBPEMixin, InterCTCMixin, ASRTranscriptionMixin):
38	    """Base class for hybrid RNNT/CTC models."""
39
40	    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
41	        cfg = model_utils.convert_model_config_to_dict_config(cfg)
42	        cfg = model_utils.maybe_update_config_version(cfg)
43	        super().__init__(cfg=cfg, trainer=trainer)
44
45	        if 'aux_ctc' not in self.cfg:
46	            raise ValueError(
47	                "The config need to have a section for the CTC decoder named as aux_ctc for Hybrid models."
48	            )
49	        with open_dict(self.cfg.aux_ctc):
50	            if "feat_in" not in self.cfg.aux_ctc.decoder or (
51	                not self.cfg.aux_ctc.decoder.feat_in and hasattr(self.encoder, '_feat_out')
52	            ):
53	                self.cfg.aux_ctc.decoder.feat_in = self.encoder._feat_out
54	            if "feat_in" not in self.cfg.aux_ctc.decoder or not self.cfg.aux_ctc.decoder.feat_in:
55	                raise ValueError("param feat_in of the decoder's config is not set!")
56
57	            if self.cfg.aux_ctc.decoder.num_classes < 1 and self.cfg.aux_ctc.decoder.vocabulary is not None:
58	                logging.info(
59	                    "\nReplacing placeholder number of classes ({}) with actual number of classes - {}".format(
60	                        self.cfg.aux_ctc.decoder.num_classes, len(self.cfg.aux_ctc.decoder.vocabulary)
61	                    )
62	                )
63	                self.cfg.aux_ctc.decoder["num_classes"] = len(self.cfg.aux_ctc.decoder.vocabulary)
64
65	        self.ctc_decoder = EncDecRNNTModel.from_config_dict(self.cfg.aux_ctc.decoder)
66	        self.ctc_loss_weight = self.cfg.aux_ctc.get("ctc_loss_weight", 0.5)
67
68	        self.ctc_loss = CTCLoss(
69	            num_classes=self.ctc_decoder.num_classes_with_blank - 1,
70	            zero_infinity=True,
71	            reduction=self.cfg.aux_ctc.get("ctc_reduction", "mean_batch"),
72	        )
73
74	        ctc_decoding_cfg = self.cfg.aux_ctc.get('decoding', None)
75	        if ctc_decoding_cfg is None:
76	            ctc_decoding_cfg = OmegaConf.structured(CTCDecodingConfig)
77	            with open_dict(self.cfg.aux_ctc):
78	                self.cfg.aux_ctc.decoding = ctc_decoding_cfg
79
80	        self.ctc_decoding = CTCDecoding(self.cfg.aux_ctc.decoding, vocabulary=self.ctc_decoder.vocabulary)
81	        self.ctc_wer = WER(
82	            decoding=self.ctc_decoding,
83	            use_cer=self.cfg.aux_ctc.get('use_cer', False),
84	            dist_sync_on_step=True,
85	            log_prediction=self.cfg.get("log_prediction", False),
86	        )
87
88	        # setting the RNNT decoder as the default one
89	        self.cur_decoder = "rnnt"
90
91	        # setting up interCTC loss (from InterCTCMixin)
92	        self.setup_interctc(decoder_name='ctc_decoder', loss_name='ctc_loss', wer_name='ctc_wer')
93
94	    @torch.no_grad()
95	    def transcribe(
96	        self,
97	        audio: List[str],
98	        batch_size: int = 4,
99	        return_hypotheses: bool = False,
100	        partial_hypothesis: Optional[List['Hypothesis']] = None,
101	        num_workers: int = 0,
102	        channel_selector: Optional[ChannelSelectorType] = None,
103	        augmentor: DictConfig = None,
104	        verbose: bool = True,
105	        timestamps: bool = None,
106	        override_config: Optional[TranscribeConfig] = None,
107	    ) -> TranscriptionReturnType:
108	        """
109	        Uses greedy decoding to transcribe audio files. Use this method for debugging and prototyping.
110
111	        Args:
112
113	            audio: (a single or list) of paths to audio files or a np.ndarray audio array.
114	                Can also be a dataloader object that provides values that can be consumed by the model.
115	                Recommended length per file is between 5 and 25 seconds. \
116	                But it is possible to pass a few hours long file if enough GPU memory is available.
117	            batch_size: (int) batch size to use during inference. \
118	                Bigger will result in better throughput performance but would use more memory.
119	            return_hypotheses: (bool) Either return hypotheses or text
120	                With hypotheses can do some postprocessing like getting timestamp or rescoring
121	            num_workers: (int) number of workers for DataLoader
122	            channel_selector (int | Iterable[int] | str): select a single channel or a subset of
123	                channels from multi-channel audio. If set to `'average'`, it performs averaging across channels.
124	                Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing.
125	            augmentor: (DictConfig): Augment audio samples during transcription if augmentor is applied.
126	            timestamps: Optional(Bool): timestamps will be returned if set to True as part of hypothesis object
127	                (output.timestep['segment']/output.timestep['word']). Refer to `Hypothesis` class for more details.
128	                Default is None and would retain the previous state set by using self.change_decoding_strategy().
129	            verbose: (bool) whether to display tqdm progress bar
130	            logprobs: (bool) whether to return ctc logits insted of hypotheses
131
132	        Returns:
133	            Returns a tuple of 2 items -
134	            * A list of greedy transcript texts / Hypothesis
135	            * An optional list of beam search transcript texts / Hypothesis / NBestHypothesis.
136	        """
137
138	        if timestamps is not None:
139	            if self.cur_decoder not in ["ctc", "rnnt"]:
140	                raise ValueError(
141	                    f"{self.cur_decoder} is not supported for cur_decoder. Supported values are ['ctc', 'rnnt']"
142	                )
143	            decoding_cfg = self.cfg.aux_ctc.decoding if self.cur_decoder == "ctc" else self.cfg.decoding
144	            if timestamps or (override_config is not None and override_config.timestamps):
145	                logging.info(
146	                    "Timestamps requested, setting decoding timestamps to True. Capture them in Hypothesis object, \
147	                        with output[idx].timestep['word'/'segment'/'char']"
148	                )
149	                return_hypotheses = True
150	                with open_dict(decoding_cfg):
151	                    decoding_cfg.compute_timestamps = True
152	                    decoding_cfg.preserve_alignments = True
153	            else:
154	                with open_dict(decoding_cfg):
155	                    decoding_cfg.compute_timestamps = False
156	                    decoding_cfg.preserve_alignments = False
157	            self.change_decoding_strategy(decoding_cfg, decoder_type=self.cur_decoder, verbose=False)
158
159	        return ASRTranscriptionMixin.transcribe(
160	            self,
161	            audio=audio,
162	            batch_size=batch_size,
163	            return_hypotheses=return_hypotheses,
164	            partial_hypothesis=partial_hypothesis,
165	            num_workers=num_workers,
166	            channel_selector=channel_selector,
167	            augmentor=augmentor,
168	            verbose=verbose,
169	            timestamps=timestamps,
170	            override_config=override_config,
171	        )
172
173	    def _transcribe_on_begin(self, audio, trcfg: TranscribeConfig):
174	        super()._transcribe_on_begin(audio, trcfg)
175
176	        if hasattr(self, 'ctc_decoder'):
177	            self.ctc_decoder.freeze()
178
179	    def _transcribe_on_end(self, trcfg: TranscribeConfig):
180	        super()._transcribe_on_end(trcfg)
181
182	        if hasattr(self, 'ctc_decoder'):
183	            self.ctc_decoder.unfreeze(partial=True)
184
185	    def _transcribe_forward(self, batch: Any, trcfg: TranscribeConfig):
186	        if self.cur_decoder == "rnnt":
187	            return super()._transcribe_forward(batch, trcfg)
188
189	        # CTC Path
190	        encoded, encoded_len = self.forward(input_signal=batch[0], input_signal_length=batch[1])
191
192	        logits = self.ctc_decoder(encoder_output=encoded)
193	        output = dict(logits=logits, encoded_len=encoded_len)
194
195	        del encoded
196	        return output
197
198	    def _transcribe_output_processing(
199	        self, outputs, trcfg: TranscribeConfig
200	    ) -> Union[List['Hypothesis'], List[List['Hypothesis']]]:
201	        if self.cur_decoder == "rnnt":
202	            return super()._transcribe_output_processing(outputs, trcfg)
203
204	        # CTC Path
205	        logits = outputs.pop('logits')
206	        encoded_len = outputs.pop('encoded_len')
207
208	        hypotheses = self.ctc_decoding.ctc_decoder_predictions_tensor(
209	            logits,
210	            encoded_len,
211	            return_hypotheses=trcfg.return_hypotheses,
212	        )
213	        logits = logits.cpu()
214
215	        if trcfg.return_hypotheses:
216	            # dump log probs per file
217	            for idx in range(logits.shape[0]):
218	                hypotheses[idx].y_sequence = logits[idx][: encoded_len[idx]]
219	                if hypotheses[idx].alignments is None:
220	                    hypotheses[idx].alignments = hypotheses[idx].y_sequence
221
222	        # DEPRECATED?
223	        # if logprobs:
224	        #     for logit, elen in zip(logits, encoded_len):
225	        #         logits_list.append(logit[:elen])
226
227	        if trcfg.timestamps:
228	            hypotheses = process_timestamp_outputs(
229	                hypotheses, self.encoder.subsampling_factor, self.cfg['preprocessor']['window_stride']
230	            )
231
232	        del logits, encoded_len
233
234	        return hypotheses
235
236	    def change_vocabulary(
237	        self,
238	        new_vocabulary: List[str],
239	        decoding_cfg: Optional[DictConfig] = None,
240	        ctc_decoding_cfg: Optional[DictConfig] = None,
241	    ):
242	        """
243	        Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning a
244	        pre-trained model. This method changes only decoder and leaves encoder and pre-processing
245	        modules unchanged. For example, you would use it if you want to use pretrained encoder
246	        when fine-tuning on data in another language, or when you'd need model to learn capitalization,
247	        punctuation and/or special characters.
248
249	        Args:
250	            new_vocabulary: list with new vocabulary. Must contain at least 2 elements. Typically, \
251	                this is target alphabet.
252	            decoding_cfg: A config for the decoder, which is optional. If the decoding type
253	                needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
254	            ctc_decoding_cfg: A config for CTC decoding, which is optional and can be used to change decoding type.
255
256	        Returns: None
257
258	        """
259	        super().change_vocabulary(new_vocabulary=new_vocabulary, decoding_cfg=decoding_cfg)
260
261	        # set up the new tokenizer for the CTC decoder
262	        if hasattr(self, 'ctc_decoder'):
263	            if self.ctc_decoder.vocabulary == new_vocabulary:
264	                logging.warning(
265	                    f"Old {self.ctc_decoder.vocabulary} and new {new_vocabulary} match. Not changing anything."
266	                )
267	            else:
268	                if new_vocabulary is None or len(new_vocabulary) == 0:
269	                    raise ValueError(f'New vocabulary must be non-empty list of chars. But I got: {new_vocabulary}')
270	                decoder_config = self.ctc_decoder.to_config_dict()
271	                new_decoder_config = copy.deepcopy(decoder_config)
272	                new_decoder_config['vocabulary'] = new_vocabulary
273	                new_decoder_config['num_classes'] = len(new_vocabulary)
274
275	                del self.ctc_decoder
276	                self.ctc_decoder = EncDecHybridRNNTCTCModel.from_config_dict(new_decoder_config)
277	                del self.ctc_loss
278	                self.ctc_loss = CTCLoss(
279	                    num_classes=self.ctc_decoder.num_classes_with_blank - 1,
280	                    zero_infinity=True,
281	                    reduction=self.cfg.aux_ctc.get("ctc_reduction", "mean_batch"),
282	                )
283
284	                if ctc_decoding_cfg is None:
285	                    # Assume same decoding config as before
286	                    logging.info("No `ctc_decoding_cfg` passed when changing decoding strategy, using internal config")
287	                    ctc_decoding_cfg = self.cfg.aux_ctc.decoding
288
289	                # Assert the decoding config with all hyper parameters
290	                ctc_decoding_cls = OmegaConf.structured(CTCDecodingConfig)
291	                ctc_decoding_cls = OmegaConf.create(OmegaConf.to_container(ctc_decoding_cls))
292	                ctc_decoding_cfg = OmegaConf.merge(ctc_decoding_cls, ctc_decoding_cfg)
293
294	                self.ctc_decoding = CTCDecoding(decoding_cfg=ctc_decoding_cfg, vocabulary=self.ctc_decoder.vocabulary)
295
296	                self.ctc_wer = WER(
297	                    decoding=self.ctc_decoding,
298	                    use_cer=self.ctc_wer.use_cer,
299	                    log_prediction=self.ctc_wer.log_prediction,
300	                    dist_sync_on_step=True,


32:from nemo.core.classes.common import PretrainedModelInfo
37:class EncDecHybridRNNTCTCModel(EncDecRNNTModel, ASRBPEMixin, InterCTCMixin, ASRTranscriptionMixin):
409:            # Add auxiliary losses, if registered
410:            loss_value = self.add_auxiliary_losses(loss_value)
439:            # Add auxiliary losses, if registered
440:            loss_value = self.add_auxiliary_losses(loss_value)


Now let me check where add_auxiliary_losses is defined:


grep -rn "def add_auxiliary_losses" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/ --include="*.py" | head -5


/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/asr_model.py:101:    def add_auxiliary_losses(self, loss: torch.Tensor, reset_registry: bool = False) -> torch.Tensor:


Excellent! Let me read that:


1	# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
2	#
3	# Licensed under the Apache License, Version 2.0 (the "License");
4	# you may not use this file except in compliance with the License.
5	# You may obtain a copy of the License at
6	#
7	#     http://www.apache.org/licenses/LICENSE-2.0
8	#
9	# Unless required by applicable law or agreed to in writing, software
10	# distributed under the License is distributed on an "AS IS" BASIS,
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and
13	# limitations under the License.
14	import logging
15	from abc import ABC
16	from typing import List, Optional
17
18	import torch
19
20	from nemo.collections.common.parts.optional_cuda_graphs import WithOptionalCudaGraphs
21	from nemo.core.classes import ModelPT
22	from nemo.core.classes.common import PretrainedModelInfo
23	from nemo.core.classes.exportable import Exportable
24	from nemo.core.classes.mixins import AccessMixin
25	from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
26	from nemo.core.utils.neural_type_utils import get_io_names
27	from nemo.utils import logging, model_utils
28	from nemo.utils.cast_utils import cast_all
29
30	__all__ = ['ASRModel']
31
32
33	class ASRModel(ModelPT, ABC):
34	    def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0):
35	        val_loss = {}
36	        tensorboard_logs = {}
37
38	        if 'val_loss' in outputs[0]:
39	            val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
40	            val_loss = {'val_loss': val_loss_mean}
41
42	            tensorboard_logs.update(val_loss)
43
44	        if "val_wer_num" in outputs[0]:
45	            wer_num = torch.stack([x['val_wer_num'] for x in outputs]).sum()
46	            wer_denom = torch.stack([x['val_wer_denom'] for x in outputs]).sum()
47	            val_wer = {'val_wer': wer_num / wer_denom}
48
49	            tensorboard_logs.update(val_wer)
50
51	        if "val_bleu_num" in outputs[0]:
52	            bleu_pred_len = torch.stack([x[f"val_bleu_pred_len"] for x in outputs]).sum()
53	            bleu_target_len = torch.stack([x[f"val_bleu_target_len"] for x in outputs]).sum()
54	            bleu_num = torch.stack([x[f"val_bleu_num"] for x in outputs]).sum(dim=0)
55	            bleu_denom = torch.stack([x[f"val_bleu_denom"] for x in outputs]).sum(dim=0)
56	            val_bleu = {"val_bleu": self.bleu._compute_bleu(bleu_pred_len, bleu_target_len, bleu_num, bleu_denom)}
57
58	            tensorboard_logs.update(val_bleu)
59
60	        return {**val_loss, 'log': tensorboard_logs}
61
62	    def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0):
63	        val_loss = {}
64	        tensorboard_logs = {}
65
66	        if 'test_loss' in outputs[0]:
67	            val_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
68	            val_loss = {'test_loss': val_loss_mean}
69
70	            tensorboard_logs.update(val_loss)
71
72	        if "test_wer_num" in outputs[0]:
73	            wer_num = torch.stack([x['test_wer_num'] for x in outputs]).sum()
74	            wer_denom = torch.stack([x['test_wer_denom'] for x in outputs]).sum()
75	            val_wer = {'test_wer': wer_num / wer_denom}
76
77	            tensorboard_logs.update(val_wer)
78
79	        if "test_bleu_num" in outputs[0]:
80	            bleu_pred_len = torch.stack([x[f"test_bleu_pred_len"] for x in outputs]).sum()
81	            bleu_target_len = torch.stack([x[f"test_bleu_target_len"] for x in outputs]).sum()
82	            bleu_num = torch.stack([x[f"test_bleu_num"] for x in outputs]).sum()
83	            bleu_denom = torch.stack([x[f"test_bleu_denom"] for x in outputs]).sum()
84	            val_bleu = {"test_bleu": self.wer._compute_bleu(bleu_pred_len, bleu_target_len, bleu_num, bleu_denom)}
85
86	            tensorboard_logs.update(val_bleu)
87
88	        return {**val_loss, 'log': tensorboard_logs}
89
90	    @classmethod
91	    def list_available_models(cls) -> 'List[PretrainedModelInfo]':
92	        """
93	        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
94	        Returns:
95	            List of available pre-trained models.
96	        """
97	        # recursively walk the subclasses to generate pretrained model info
98	        list_of_models = model_utils.resolve_subclass_pretrained_model_info(cls)
99	        return list_of_models
100
101	    def add_auxiliary_losses(self, loss: torch.Tensor, reset_registry: bool = False) -> torch.Tensor:
102	        """
103	        Utility method to enable calculation of auxiliary losses for ASR training.
104
105	        Args:
106	            loss: The output loss value prior to addition with auxiliary losses.
107	            reset_registry: Bool, whether to reset the AccessMixin registry after adding auxiliary losses.
108
109	        Returns:
110	            Loss tensor used for back propagation.
111	        """
112	        # Add adapter auxiliary losses, if registered
113	        if AccessMixin.is_access_enabled(self.model_guid):
114	            registry = AccessMixin.get_module_registry(self)
115	            log_dict = {}
116
117	            for loss_key, loss_registry in registry.items():
118	                # Add auxiliary loss to total loss
119	                if 'adapter_loss' in loss_registry:
120	                    loss_list = loss_registry['adapter_loss']
121	                    loss_value = sum(loss_list)
122	                    loss += loss_value
123
124	                    # Log current loss name and value
125	                    keys = loss_key.split(".")
126	                    key = "/".join(keys)
127	                    key = "adapter_loss/" + key
128	                    log_dict[key] = loss_value.detach()
129
130	            if len(log_dict) > 0:
131	                self.log_dict(log_dict)
132
133	        if reset_registry:
134	            AccessMixin.reset_registry(self)
135
136	        # return total loss
137	        return loss
138
139	    def setup_optimization_flags(self):
140	        """
141	        Utility method that must be explicitly called by the subclass in order to support optional optimization flags.
142	        This method is the only valid place to access self.cfg prior to DDP training occurs.
143
144	        The subclass may chose not to support this method, therefore all variables here must be checked via hasattr()
145	        """
146	        # Skip update if nan/inf grads appear on any rank.
147	        self._skip_nan_grad = False
148	        if "skip_nan_grad" in self._cfg and self._cfg["skip_nan_grad"]:
149	            self._skip_nan_grad = self._cfg["skip_nan_grad"]
150
151	    def on_after_backward(self):
152	        """
153	        zero-out the gradients which any of them is NAN or INF
154	        """
155	        super().on_after_backward()
156
157	        if hasattr(self, '_skip_nan_grad') and self._skip_nan_grad:
158	            device = next(self.parameters()).device
159	            valid_gradients = torch.tensor([1], device=device, dtype=torch.float32)
160
161	            # valid_gradients = True
162	            for param_name, param in self.named_parameters():
163	                if param.grad is not None:
164	                    is_not_nan_or_inf = not (torch.isnan(param.grad).any() or torch.isinf(param.grad).any())
165	                    if not is_not_nan_or_inf:
166	                        valid_gradients = valid_gradients * 0
167	                        break
168
169	            if torch.distributed.is_initialized():
170	                torch.distributed.all_reduce(valid_gradients, op=torch.distributed.ReduceOp.MIN)
171
172	            if valid_gradients < 1:
173	                logging.warning(f'detected inf or nan values in gradients! Setting gradients to zero.')
174	                self.zero_grad()
175
176	    def on_train_epoch_start(self) -> None:
177	        """
178	        Decoder with CUDA graphs does not release memory, thus we disable it for training epoch.
179	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs
180	        """
181	        WithOptionalCudaGraphs.disable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
182
183	    def on_train_epoch_end(self) -> None:
184	        """
185	        After training, we can enable the decoder with CUDA graphs.
186	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs
187	        """
188	        WithOptionalCudaGraphs.enable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
189
190	    def on_validation_epoch_start(self) -> None:
191	        """
192	        For validation, we enable CUDA graphs to speedup validation.
193	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs.
194	        """
195	        WithOptionalCudaGraphs.enable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
196
197	    def on_validation_epoch_end(self) -> Optional[dict[str, dict[str, torch.Tensor]]]:
198	        """
199	        After validation, we disable CUDA graphs, since `validation` can be called in training loop, and
200	        training will continue after validation
201	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs.
202	        """
203	        WithOptionalCudaGraphs.disable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
204	        return super().on_validation_epoch_end(sync_metrics=True)
205
206	    def on_test_epoch_start(self) -> None:
207	        """
208	        For testing, we enable CUDA graphs to speedup validation.
209	        We do not need to disable CUDA graphs after testing, since `test` cannot be called in training loop.
210	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs.
211	        """
212	        WithOptionalCudaGraphs.enable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
213
214	    def on_predict_epoch_start(self) -> None:
215	        """
216	        For predicting, we enable CUDA graphs to speedup validation.
217	        We do not need to disable CUDA graphs after predicting, since `predict` cannot be called in training loop.
218	        EncDecRNNTModel.decoding.decoding is the inference class with CUDA graphs
219	        """
220	        WithOptionalCudaGraphs.enable_cuda_graphs_recursive(self, attribute_path="decoding.decoding")
221
222	    @property
223	    def oomptimizer_schema(self) -> dict:
224	        """
225	        Return a typing schema for optimal batch size calibration for various
226	        sequence lengths using OOMptimizer.
227	        """
228	        return {
229	            "cls": tuple,
230	            "inputs": [
231	                {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},
232	                {"type": NeuralType(("B",), LengthsType()), "seq_length": "input"},
233	                {
234	                    "type": NeuralType(("B", "T"), LabelsType()),
235	                    "seq_length": "output",
236	                    "vocab_size": self.tokenizer.vocab_size,
237	                },
238	                {"type": NeuralType(("B",), LengthsType()), "seq_length": "output"},
239	            ],
240	        }
241
242
243	class ExportableEncDecModel(Exportable):
244	    """
245	    Simple utiliy mix-in to export models that consist of encoder/decoder pair
246	    plus pre/post processor, but have to be exported as encoder/decoder pair only
247	    (covers most ASR classes)
248	    """
249
250	    @property
251	    def input_module(self):
252	        return self.encoder
253
254	    @property
255	    def output_module(self):
256	        return self.decoder
257
258	    @property
259	    def output_names(self):
260	        otypes = self.output_module.output_types
261	        if getattr(self.input_module, 'export_cache_support', False):
262	            in_types = self.input_module.output_types
263	            otypes = {n: t for (n, t) in list(otypes.items())[:1]}
264	            for n, t in list(in_types.items())[1:]:
265	                otypes[n] = t
266	        return get_io_names(otypes, self.disabled_deployment_output_names)
267
268	    def forward_for_export(
269	        self, audio_signal, length=None, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None
270	    ):
271	        """
272	        This forward is used when we need to export the model to ONNX format.
273	        Inputs cache_last_channel and cache_last_time are needed to be passed for exporting streaming models.
274
275	        Args:
276	            input: Tensor that represents a batch of raw audio signals of shape [B, T]. T here represents timesteps.
277	            length: Vector of length B, that contains the individual lengths of the audio sequences.
278	            cache_last_channel: Tensor of shape [N, B, T, H] which contains the cache for last channel layers
279	            cache_last_time: Tensor of shape [N, B, H, T] which contains the cache for last time layers
280	                N is the number of such layers which need caching, B is batch size, H is the hidden size of activations,
281	                and T is the length of the cache
282
283	        Returns:
284	            the output of the model
285	        """
286	        enc_fun = getattr(self.input_module, 'forward_for_export', self.input_module.forward)
287	        if cache_last_channel is None:
288	            encoder_output = enc_fun(audio_signal=audio_signal, length=length)
289	            if isinstance(encoder_output, tuple):
290	                encoder_output = encoder_output[0]
291	        else:
292	            encoder_output, length, cache_last_channel, cache_last_time, cache_last_channel_len = enc_fun(
293	                audio_signal=audio_signal,
294	                length=length,
295	                cache_last_channel=cache_last_channel,
296	                cache_last_time=cache_last_time,
297	                cache_last_channel_len=cache_last_channel_len,
298	            )
299
300	        dec_fun = getattr(self.output_module, 'forward_for_export', self.output_module.forward)


Perfect! Now let me check the InterCTC mixin more thoroughly and then compile the full list. Let me also search for any other auxiliary loss patterns:


grep -rn "add_interctc_losses\|compute.*aux\|auxiliary" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py | head -30


409:            # Add auxiliary losses, if registered
410:            loss_value = self.add_auxiliary_losses(loss_value)
439:            # Add auxiliary losses, if registered
440:            loss_value = self.add_auxiliary_losses(loss_value)
474:        loss_value, additional_logs = self.add_interctc_losses(
607:        loss_value, additional_logs = self.add_interctc_losses(


Let me check the add_interctc_losses method:


grep -n "def add_interctc_losses" /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py


214:    def add_interctc_losses(


214	    def add_interctc_losses(
215	        self,
216	        loss_value: torch.Tensor,
217	        transcript: torch.Tensor,
218	        transcript_len: torch.Tensor,
219	        compute_wer: bool,
220	        compute_loss: bool = True,
221	        log_wer_num_denom: bool = False,
222	        log_prefix: str = "",
223	    ) -> Tuple[Optional[torch.Tensor], Dict]:
224	        """Adding interCTC losses if required.
225
226	        Will also register loss/wer metrics in the returned dictionary.
227
228	        Args:
229	            loss_value (torch.Tensor): regular loss tensor (will add interCTC loss to it).
230	            transcript (torch.Tensor): current utterance transcript.
231	            transcript_len (torch.Tensor): current utterance transcript length.
232	            compute_wer (bool): whether to compute WER for the current utterance.
233	                Should typically be True for validation/test and only True for
234	                training if current batch WER should be logged.
235	            compute_loss (bool): whether to compute loss for the current utterance.
236	                Should always be True in training and almost always True in
237	                validation, unless all other losses are disabled as well.
238	                Defaults to True.
239	            log_wer_num_denom (bool): if True, will additionally log WER num/denom
240	                in the returned metrics dictionary. Should always be True for
241	                validation/test to allow correct metrics aggregation. Should
242	                always be False for training. Defaults to False.
243	            log_prefix (str): prefix added to all log values. Should be ``""`` for
244	                training and ``"val_"`` for validation. Defaults to "".
245
246	        Returns:
247	            tuple[Optional[torch.Tensor], Dict]: tuple of new loss tensor and dictionary with logged metrics.
248	        """
249	        if not self.is_interctc_enabled() or not AccessMixin.is_access_enabled(getattr(self, "model_guid", None)):
250	            return loss_value, {}
251	        metrics = {}
252	        if compute_loss:
253	            metrics[f"{log_prefix}final_loss"] = loss_value
254	        else:
255	            loss_value = None
256	        captured_tensors = self.get_captured_interctc_tensors()
257
258	        if compute_loss:
259	            loss_value *= self.get_interctc_param('main_loss_weight')
260
261	        for layer_idx, intermediate_result, loss_weight in zip(
262	            self.get_interctc_param('apply_at_layers'),
263	            captured_tensors,
264	            self.get_interctc_param('intermediate_loss_weights'),
265	        ):
266	            if compute_loss:
267	                inter_loss_value = self.get_interctc_param('loss')(
268	                    log_probs=intermediate_result[0],
269	                    targets=transcript,
270	                    target_lengths=transcript_len,
271	                    input_lengths=intermediate_result[1],
272	                )
273	                metrics[f"{log_prefix}inter_ctc_loss_l{layer_idx}"] = inter_loss_value.detach()
274	                loss_value += inter_loss_value * loss_weight
275	            if compute_wer:
276	                self.get_interctc_param('wer').update(
277	                    predictions=intermediate_result[0],
278	                    targets=transcript,
279	                    targets_lengths=transcript_len,
280	                    predictions_lengths=intermediate_result[1],
281	                )
282	                wer, wer_num, wer_denom = self.get_interctc_param('wer').compute()
283	                self.get_interctc_param('wer').reset()
284	                metrics.update({f'{log_prefix}inter_wer_l{layer_idx}': wer})
285	                if log_wer_num_denom:
286	                    metrics.update(
287	                        {
288	                            f'{log_prefix}inter_wer_num_l{layer_idx}': wer_num,
289	                            f'{log_prefix}inter_wer_denom_l{layer_idx}': wer_denom,
290	                        }
291	                    )
292
293	        # return total loss and dictionary of metrics
294	        return loss_value, metrics
295


Perfect! Now let me search for any other custom loss implementations in the local project:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.py" | xargs grep -l "Loss\|loss" | head -20


Now let me create a comprehensive summary by checking one more important thing - the config documentation:


find /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr -type f -name "*.yaml" | head -10


(Bash completed with no output)


Let me verify the local project structure for any config files:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.yaml" -o -name "*.yml" | head -20


Let me check one of the training configs to see the loss configuration:


1	# Stage 1 Smoke Config: FastConformer Hybrid CTC+RNNT on tiny subset
2	# Purpose: Validate training pipeline end-to-end before scaling up
3	# Expected: ~50-200 steps, <5 min on single GPU
4
5	name: maya-asr-stage1-smoke
6
7	trainer:
8	  devices: 1
9	  num_nodes: 1
10	  accelerator: gpu
11	  strategy: auto
12	  precision: bf16-mixed
13	  max_steps: 100
14	  val_check_interval: 50
15	  accumulate_grad_batches: 1
16	  gradient_clip_val: 1.0
17	  log_every_n_steps: 10
18	  enable_checkpointing: false
19	  default_root_dir: experiments/stage1_smoke
20
21	model:
22	  model_defaults:
23	    asr_enc_hidden: 256
24	    enc_hidden: 256
25	    dec_hidden: 256
26	    pred_hidden: 256
27
28	  preprocessor:
29	    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
30	    sample_rate: 16000
31	    normalize: per_feature
32	    window_size: 0.025
33	    window_stride: 0.01
34	    features: 80
35	    n_fft: 512
36	    dither: 0.00001
37
38	  encoder:
39	    _target_: nemo.collections.asr.modules.ConformerEncoder
40	    feat_in: 80
41	    n_layers: 4
42	    d_model: ${model.model_defaults.enc_hidden}
43	    subsampling: dw_striding
44	    subsampling_factor: 8
45	    subsampling_conv_channels: 256
46	    ff_expansion_factor: 4
47	    self_attention_model: rel_pos
48	    n_heads: 4
49	    conv_kernel_size: 9
50	    dropout: 0.1
51	    dropout_pre_encoder: 0.1
52	    dropout_emb: 0.0
53
54	  decoder:
55	    _target_: nemo.collections.asr.modules.RNNTDecoder
56	    prednet:
57	      pred_hidden: ${model.model_defaults.pred_hidden}
58	      pred_rnn_layers: 1
59	      dropout: 0.1
60
61	  joint:
62	    _target_: nemo.collections.asr.modules.RNNTJoint
63	    jointnet:
64	      joint_hidden: 256
65	      activation: relu
66	      dropout: 0.1
67	      encoder_hidden: ${model.model_defaults.enc_hidden}
68	      pred_hidden: ${model.model_defaults.pred_hidden}
69
70	  loss:
71	    loss_name: default
72	    warprnnt_numba_kwargs:
73	      fastemit_lambda: 0.0
74
75	  decoding:
76	    strategy: greedy_batch
77	    greedy:
78	      max_symbols_per_step: 10
79	      loop_labels: true
80	      use_cuda_graph_decoder: false
81
82	  # Auxiliary CTC decoder (required for hybrid model)
83	  aux_ctc:
84	    ctc_loss_weight: 0.3
85	    use_cer: false
86	    decoder:
87	      _target_: nemo.collections.asr.modules.ConvASRDecoder
88	      feat_in: null
89	      num_classes: -1
90
91	  tokenizer:
92	    dir: tokenizers/smoke_bpe
93	    type: bpe
94
95	  train_ds:
96	    manifest_filepath: data/manifests/smoke_train.jsonl
97	    # For multi-corpus training, replace manifest_filepath with:
98	    #   input_cfg: configs/data/stage1_smoke_input_cfg.yaml
99	    # Auto-generated by preflight; metadata at configs/data/input_cfg_metadata.json
100	    batch_size: 4


Perfect! Now I have a comprehensive understanding. Let me compile a final summary with all findings:

## Summary of Available Auxiliary Losses, Multi-Objective Training Options, and Loss Functions for EncDecHybridRNNTCTCBPEModel

Based on thorough analysis of the NeMo package at `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/`, here are all available loss functions and multi-objective training mechanisms:

### Core Loss Functions

#### **RNNT/Transducer Losses** (Primary for Hybrid Models)
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py` and `rnnt_pytorch.py`

1. **RNNTLoss** - Main RNNT Loss wrapper with multiple backends:
   - `warprnnt` - NVIDIA's optimized RNNT loss
   - `warprnnt_numba` - Numba-based implementation (default)
   - `pytorch` - Pure PyTorch implementation (debugging)
   - `multiblank_rnnt` - Multi-blank RNNT variant
   - `multiblank_rnnt_pytorch` - PyTorch multi-blank variant
   - `graph_rnnt` - K2 graph-based RNNT
   - `graph_w_transducer` - K2 graph-based W-Transducer
   - `tdt` - Time Duration Transducer loss
   - `tdt_pytorch` - PyTorch TDT variant

   **Configurable Parameters:**
   - `fastemit_lambda` (0.0-1.0) - Fast emission penalty regularization
   - `clamp` - Loss clamping value
   - `sigma` - Logit under-normalization parameter (for multi-blank and TDT)
   - `omega` - Duration variance parameter (for TDT)
   - `big_blank_durations` - Duration list for multi-blank variant
   - `durations` - Duration list for TDT

#### **CTC Losses**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/ctc.py`

2. **CTCLoss** - Standard Connectionist Temporal Classification Loss
   - Supports multiple reduction modes: `mean_batch`, `mean_volume`, `mean`, `sum`, `none`
   - Parameter: `zero_infinity` - Zero out infinite losses

#### **Angular/Metric Losses**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/angularloss.py`

3. **AngularSoftmaxLoss (ArcFace)** - Angular margin loss for speaker discrimination
   - Parameters: `scale` (default 20.0), `margin` (default 1.35)
   - Purpose: Speaker/language embeddings separation

#### **Binary Classification Loss**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/bce_loss.py`

4. **BCELoss** - Binary Cross Entropy Loss with class normalization support
   - Parameters:
     - `reduction`: 'mean', 'sum', 'none'
     - `class_normalization`: 'none', 'class', 'class_binary', 'binary'
     - `weight`: Class weights tensor
     - `sorted_preds` / `sorted_loss`: Optional sorting flags

#### **Lattice-Based Losses**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/lattice_losses.py`

5. **LatticeLoss** - K2-based lattice losses (requires k2 v1.14+)
   - Variants:
     - `criterion_type`: 'ml' (Maximum Likelihood) or 'map' (Maximum A Posteriori)
     - `loss_type`: 'ctc', 'rnnt' (for ml), or 'mmi' (for map)
   - Parameters:
     - `split_batch_size`: Memory optimization via batch splitting
     - `graph_module_cfg`: Topology and LM configuration

### Self-Supervised Learning (SSL) Losses
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/ssl_losses/`

6. **ContrastiveLoss** - SimCLR-style contrastive learning
   - Purpose: Self-supervised representation learning
   - Parameters:
     - `num_negatives`: Number of contrastive negatives
     - `quantized_targets`: Use quantization (Gumbel-VQ)
     - `codebook_size`: VQ codebook size
     - `prob_ppl_weight`: Perplexity loss weight (0.1)
     - `logit_temp`: Temperature scaling
     - `sample_from_same_utterance_only`: Negative sampling strategy
     - `mask_threshold`: Masked frame threshold (0.8)
     - `multiplier`: Loss scaling factor (16.0)

7. **MLMLoss** - Masked Language Modeling loss
   - Purpose: Predict tokens in masked regions
   - Parameters:
     - `combine_time_steps`: Temporal fusion
     - `mask_threshold`: Masking threshold (0.8)

8. **MultiMLMLoss** - MLM for multiple decoders/heads
   - Purpose: Multi-task masked language modeling
   - Parameters:
     - `num_decoders`: Number of independent decoders
     - `squeeze_single`: Squeeze output for single decoder

9. **CTCLossForSSL** - CTC loss wrapper for SSL
   - Wrapper around standard CTCLoss for self-supervised setups

10. **RNNTLossForSSL** - RNNT loss wrapper for SSL
    - Wrapper around standard RNNTLoss for self-supervised setups

### Multi-Objective Training Mechanisms

#### **InterCTC Mixin** (Intermediate CTC Losses)
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py`

**Purpose**: Auxiliary CTC losses from intermediate encoder layers to improve gradient flow and training dynamics.

**Implementation**: `InterCTCMixin` class providing:
- `setup_interctc()` - Initialize InterCTC configuration
- `add_interctc_losses()` - Compute and aggregate intermediate losses
- `get_captured_interctc_tensors()` - Extract intermediate layer outputs

**Configuration Parameters**:
```yaml
interctc:
  loss_weights: [0.1, 0.2]          # Weights for layers 8 and 10
  apply_at_layers: [8, 10]          # Which encoder layers to apply CTC
  main_loss_weight: 0.7             # Automatically computed as 1.0 - sum(loss_weights)
```

**How It Works**:
- Captures encoder outputs at specified layers
- Applies CTC decoder to each captured output
- Computes CTC loss with configurable weights
- Accumulates auxiliary losses into main loss
- Enables early gradient flow in deep encoders

#### **Auxiliary CTC Decoder**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py`

**Purpose**: Parallel CTC branch for hybrid RNNT/CTC models to provide additional training signal.

**Configuration** (in model config):
```yaml
aux_ctc:
  ctc_loss_weight: 0.3              # Balance with RNNT loss (total = 1.0)
  ctc_reduction: "mean_batch"       # Reduction strategy
  use_cer: false                     # Compute CER instead of WER
  decoder:
    _target_: nemo.collections.asr.modules.ConvASRDecoder
    feat_in: null                    # Auto-filled from encoder output
    num_classes: -1                  # Auto-filled from vocabulary
```

**Training Formula**: `total_loss = rnnt_loss + ctc_loss_weight * ctc_loss`

#### **Adapter Auxiliary Losses**
Located: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/asr_model.py`

**Method**: `add_auxiliary_losses(loss, reset_registry=False)`

**Purpose**: Register and aggregate adapter losses from modules via AccessMixin registry.

**Mechanism**:
- Checks AccessMixin registry for `adapter_loss` entries
- Sums all registered losses
- Logs each loss component separately
- Optional registry reset after aggregation

### Loss Configuration in EncDecHybridRNNTCTCBPEModel

**Key Loss-Related Model Attributes**:

1. **Main RNNT Loss**:
   ```python
   self.loss = RNNTLoss(num_classes, reduction, loss_name, loss_kwargs)
   ```

2. **Auxiliary CTC Loss**:
   ```python
   self.ctc_loss = CTCLoss(num_classes, zero_infinity=True, reduction)
   self.ctc_loss_weight = 0.5  # From config aux_ctc.ctc_loss_weight
   ```

3. **InterCTC Setup** (if enabled):
   ```python
   self.setup_interctc(decoder_name='ctc_decoder',
                       loss_name='ctc_loss',
                       wer_name='ctc_wer')
   ```

### Integration Points

**Forward Training Pass** (hybrid_rnnt_ctc_models.py lines 405-440):
```
1. Compute RNNT loss: loss_value = self.loss(...)
2. Compute auxiliary CTC loss: ctc_loss = self.ctc_loss(...)
3. Combine: total_loss = loss_value + self.ctc_loss_weight * ctc_loss
4. Add InterCTC losses if enabled: self.add_interctc_losses(...)
5. Add adapter losses: self.add_auxiliary_losses(...)
```

### Recommended Configurations for Your Use Case

For **EncDecHybridRNNTCTCBPEModel**, viable multi-objective training combinations:

1. **Basic Hybrid** (RNNT + CTC):
   - Primary: RNNT loss
   - Auxiliary: CTC loss (weight 0.3-0.5)

2. **Hybrid + InterCTC** (Recommended for deep models):
   - Primary: RNNT loss
   - Auxiliary-1: CTC loss (weight 0.3)
   - Auxiliary-2: InterCTC from layers 12, 18 (weights 0.15, 0.15)
   - Main encoder weight: 0.4

3. **Hybrid + Self-Supervised**:
   - Primary: RNNT + CTC losses
   - Auxiliary: ContrastiveLoss (for representation learning)
   - Additional: MLMLoss (for masked prediction)

4. **Fine-tuning with Language ID** (requires custom extension):
   - Could add AngularSoftmaxLoss on encoder outputs for language discrimination
   - Would require custom forward hook registration

### File Locations Summary

- **Core losses**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/`
- **Loss variants**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/rnnt.py`
- **SSL losses**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/losses/ssl_losses/`
- **Model implementation**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py`
- **Auxiliary loss integration**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/asr_model.py` (line 101)
- **InterCTC mixin**: `/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/mixins/interctc_mixin.py`
- **Local project configs**: `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/`