Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 11, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/paginate.py", line 272, in __iter__
    response = self._make_request(current_kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/context.py", line 123, in wrapper
    return func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/paginate.py", line 360, in _make_request
    return self._method(**current_kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 602, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/context.py", line 123, in wrapper
    return func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 1060, in _make_api_call
    http, parsed_response = self._make_request(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 1084, in _make_request
    return self._endpoint.make_request(operation_model, request_dict)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/endpoint.py", line 119, in make_request
    return self._send_request(request_dict, operation_model)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/endpoint.py", line 197, in _send_request
    success_response, exception = self._get_response(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/endpoint.py", line 239, in _get_response
    success_response, exception = self._do_get_response(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/endpoint.py", line 313, in _do_get_response
    parsed_response = parser.parse(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 267, in parse
    parsed = self._do_parse(response, shape)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 1119, in _do_parse
    self._add_modeled_parse(response, shape, final_parsed)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 1129, in _add_modeled_parse
    self._parse_payload(response, shape, member_shapes, final_parsed)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 1171, in _parse_payload
    body_parsed = self._parse_shape(shape, original_parsed)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 347, in _parse_shape
    return handler(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 476, in _handle_structure
    parsed[member_name] = self._parse_shape(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 347, in _parse_shape
    return handler(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 1230, in _handle_list
    return super()._handle_list(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 452, in _handle_list
    return super()._handle_list(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 355, in _handle_list
    parsed.append(self._parse_shape(member_shape, item))
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 347, in _parse_shape
    return handler(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 476, in _handle_structure
    parsed[member_name] = self._parse_shape(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 347, in _parse_shape
    return handler(shape, node)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 188, in _get_text_content
    return func(self, shape, text)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/parsers.py", line 573, in _handle_timestamp
    return self._timestamp_parser(text)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/utils.py", line 993, in parse_timestamp
    return _parse_timestamp_with_tzinfo(value, tzinfo)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/utils.py", line 973, in _parse_timestamp_with_tzinfo
    return dateutil.parser.parse(value, tzinfos={'GMT': tzutc()})
  File "/home/ubuntu/.local/lib/python3.10/site-packages/dateutil/parser.py", line 1182, in parse
    return DEFAULTPARSER.parse(timestr, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/dateutil/parser.py", line 587, in parse
    if (isinstance(tzinfos, collections.Callable) or
AttributeError: module 'collections' has no attribute 'Callable'


The boto3/dateutil has a Python 3.10 compatibility issue with listing objects that have timestamps. Let me try downloading directly with the known path.


Bash

python3 -c "
import boto3, os, collections.abc, collections
# Monkey-patch for Python 3.10+ compatibility
collections.Callable = collections.abc.Callable

from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
resp = s3.list_objects_v2(Bucket=bucket, Prefix='hybrid-tdt-gemma/03-30-2026/')
for obj in resp.get('Contents', []):
    print(f'{obj[\"Size\"]/1e9:.2f}GB  {obj[\"Key\"]}')
" 2>&1

List 03-30 files with dateutil fix


0.00GB  hybrid-tdt-gemma/03-30-2026/INFERENCE.md
13.75GB  hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt
0.00GB  hybrid-tdt-gemma/03-30-2026/tokenizer/tokenizer.model
0.00GB  hybrid-tdt-gemma/03-30-2026/tokenizer/vocab.txt


13.75GB checkpoint. Let me download it and inspect it, plus grab the INFERENCE.md for context.


python3 -c "
import boto3, os, collections.abc, collections
collections.Callable = collections.abc.Callable
from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
# Download the small INFERENCE.md first
s3.download_file(bucket, 'hybrid-tdt-gemma/03-30-2026/INFERENCE.md', '/tmp/INFERENCE.md')
print('Downloaded INFERENCE.md')
" 2>&1

Download INFERENCE.md


Downloaded INFERENCE.md


python3 -c "
import boto3, os, collections.abc, collections
collections.Callable = collections.abc.Callable
from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
target = '/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt'
print('Downloading 13.75GB checkpoint... this may take a while')
s3.download_file(bucket, 'hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt', target)
print('Done!')
" 2>&1

Download 13.75GB checkpoint from R2


Downloading 13.75GB checkpoint... this may take a while
Done!


1	# Maya ASR — Parakeet TDT 1.1B Inference Guide
2
3	## Model Architecture
4
5	- **Base**: NVIDIA Parakeet TDT 1.1B (FastConformer Hybrid RNNT-CTC)
6	- **Encoder**: 42-layer FastConformer, 1024-dim hidden, ~1.06B params
7	- **Decoder**: TDT (Token-and-Duration Transducer) + CTC head
8	- **Tokenizer**: 32k BPE SentencePiece (`tokenizer/tokenizer.model`)
9	- **Language Conditioning**: 12-language embedding (additive bias on encoder output)
10
11	## Supported Languages
12
13	| Language   | Code | Lang ID |
14	|------------|------|---------|
15	| Hindi      | hi   | 0       |
16	| Bengali    | bn   | 1       |
17	| Tamil      | ta   | 2       |
18	| Telugu     | te   | 3       |
19	| Marathi    | mr   | 4       |
20	| Gujarati   | gu   | 5       |
21	| Kannada    | kn   | 6       |
22	| Malayalam  | ml   | 7       |
23	| Punjabi    | pa   | 8       |
24	| Odia       | or   | 9       |
25	| Assamese   | as   | 10      |
26	| English    | en   | 11      |
27
28	## Language Embedding
29
30	The model uses a **mandatory** language embedding that is added to the encoder output at inference time. You **must** provide a language ID — the model was trained with language conditioning and will produce degraded results without it.
31
32	```
33	Architecture:
34	  audio → FastConformer Encoder → encoded (B, 1024, T)
35	                                        ↓
36	  lang_id → nn.Embedding(12, 1024) → lang_bias (B, 1024, 1)
37	                                        ↓
38	                              encoded + lang_bias → Decoder → text
39	```
40
41	The embedding is a small module (12×1024 = 12,288 params) initialized with `N(0, 0.02)` to minimally perturb the pretrained encoder.
42
43	## Inference — Standard NeMo API
44
45	```python
46	import torch
47	import nemo.collections.asr as nemo_asr
48
49	# ── Load model ──────────────────────────────────────────────────────
50	model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(
51	    "path/to/maya_asr.nemo"
52	)
53	model.eval()
54	model.cuda()
55
56	# ── Language ID mapping ─────────────────────────────────────────────
57	LANG_TO_ID = {
58	    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
59	    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
60	}
61
62	# ── Set language before transcription ───────────────────────────────
63	# The language embedding is injected via model._current_lang_ids
64	# Set it before calling transcribe(), it gets consumed (cleared) after use.
65
66	lang = "hi"  # Hindi
67	audio_files = ["audio1.wav", "audio2.wav"]
68
69	model._current_lang_ids = torch.tensor(
70	    [LANG_TO_ID[lang]] * len(audio_files), device="cuda"
71	)
72
73	transcriptions = model.transcribe(audio_files)
74	print(transcriptions)
75	```
76
77	## Inference — Batch with Mixed Languages
78
79	```python
80	# For a batch with different languages per sample:
81	audio_files = ["hindi_audio.wav", "tamil_audio.wav", "english_audio.wav"]
82	lang_codes = ["hi", "ta", "en"]
83
84	model._current_lang_ids = torch.tensor(
85	    [LANG_TO_ID[lc] for lc in lang_codes], device="cuda"
86	)
87
88	transcriptions = model.transcribe(audio_files)
89	for audio, lang, text in zip(audio_files, lang_codes, transcriptions):
90	    print(f"[{lang}] {text}")
91	```
92
93	## Inference — Re-patching Encoder (if needed)
94
95	If the encoder forward hook isn't already patched (e.g. loading a raw checkpoint rather than a `.nemo` file), you need to re-apply the language embedding patch:
96
97	```python
98	import torch.nn as nn
99
100	# Assuming model is loaded and has _lang_embed_module
101	lang_embed = model._lang_embed_module  # nn.Embedding(12, 1024)
102
103	# Save original encoder forward
104	original_forward = model.encoder.forward
105
106	def encoder_forward_with_lang(audio_signal, length):
107	    encoded, encoded_len = original_forward(
108	        audio_signal=audio_signal, length=length
109	    )
110	    if hasattr(model, '_current_lang_ids') and model._current_lang_ids is not None:
111	        lang_ids = model._current_lang_ids
112	        if lang_ids.shape[0] == encoded.shape[0]:
113	            lang_bias = lang_embed(lang_ids)          # (B, 1024)
114	            lang_bias = lang_bias.unsqueeze(-1)        # (B, 1024, 1)
115	            encoded = encoded + lang_bias
116	        model._current_lang_ids = None  # clear after use
117	    return encoded, encoded_len
118
119	model.encoder.forward = encoder_forward_with_lang
120	```
121
122	## Decoding Configuration
123
124	The model supports two decoding modes:
125
126	### TDT Greedy (Primary — recommended)
127	```python
128	# Already configured in the checkpoint. Duration-aware decoding.
129	# Durations: [0, 1, 2, 3, 4] frames
130	# This is the default mode.
131	```
132
133	### CTC (Fallback — faster)
134	```python
135	# For CTC-only decoding (faster but less accurate):
136	transcriptions = model.transcribe(audio_files, return_hypotheses=False)
137	```
138
139	## Tokenizer
140
141	- **Type**: SentencePiece BPE
142	- **Vocab size**: 32,768 tokens
143	- **Files**: `tokenizer/tokenizer.model`, `tokenizer/vocab.txt`
144	- **Coverage**: All 12 supported Indic scripts + English + numerals
145
146	## Notes
147
148	- **Language embedding is mandatory** — always set `_current_lang_ids` before transcription
149	- Audio must be **16kHz mono** — the model expects this sample rate
150	- For long audio (>30s), consider chunking with overlap for best results
151	- The model was fine-tuned from NVIDIA Parakeet TDT 1.1B on Indic language data
152

# Maya ASR — Parakeet TDT 1.1B Inference Guide

## Model Architecture

- **Base**: NVIDIA Parakeet TDT 1.1B (FastConformer Hybrid RNNT-CTC)
- **Encoder**: 42-layer FastConformer, 1024-dim hidden, ~1.06B params
- **Decoder**: TDT (Token-and-Duration Transducer) + CTC head
- **Tokenizer**: 32k BPE SentencePiece (`tokenizer/tokenizer.model`)
- **Language Conditioning**: 12-language embedding (additive bias on encoder output)

## Supported Languages

| Language   | Code | Lang ID |
|------------|------|---------|
| Hindi      | hi   | 0       |
| Bengali    | bn   | 1       |
| Tamil      | ta   | 2       |
| Telugu     | te   | 3       |
| Marathi    | mr   | 4       |
| Gujarati   | gu   | 5       |
| Kannada    | kn   | 6       |
| Malayalam  | ml   | 7       |
| Punjabi    | pa   | 8       |
| Odia       | or   | 9       |
| Assamese   | as   | 10      |
| English    | en   | 11      |

## Language Embedding

The model uses a **mandatory** language embedding that is added to the encoder output at inference time. You **must** provide a language ID — the model was trained with language conditioning and will produce degraded results without it.

```
Architecture:
  audio → FastConformer Encoder → encoded (B, 1024, T)
                                        ↓
  lang_id → nn.Embedding(12, 1024) → lang_bias (B, 1024, 1)
                                        ↓
                              encoded + lang_bias → Decoder → text
```

The embedding is a small module (12×1024 = 12,288 params) initialized with `N(0, 0.02)` to minimally perturb the pretrained encoder.

## Inference — Standard NeMo API

```python
import torch
import nemo.collections.asr as nemo_asr

# ── Load model ──────────────────────────────────────────────────────
model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(
    "path/to/maya_asr.nemo"
)
model.eval()
model.cuda()

# ── Language ID mapping ─────────────────────────────────────────────
LANG_TO_ID = {
    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
}

# ── Set language before transcription ───────────────────────────────
# The language embedding is injected via model._current_lang_ids
# Set it before calling transcribe(), it gets consumed (cleared) after use.

lang = "hi"  # Hindi
audio_files = ["audio1.wav", "audio2.wav"]

model._current_lang_ids = torch.tensor(
    [LANG_TO_ID[lang]] * len(audio_files), device="cuda"
)

transcriptions = model.transcribe(audio_files)
print(transcriptions)
```

## Inference — Batch with Mixed Languages

```python
# For a batch with different languages per sample:
audio_files = ["hindi_audio.wav", "tamil_audio.wav", "english_audio.wav"]
lang_codes = ["hi", "ta", "en"]

model._current_lang_ids = torch.tensor(
    [LANG_TO_ID[lc] for lc in lang_codes], device="cuda"
)

transcriptions = model.transcribe(audio_files)
for audio, lang, text in zip(audio_files, lang_codes, transcriptions):
    print(f"[{lang}] {text}")
```

## Inference — Re-patching Encoder (if needed)

If the encoder forward hook isn't already patched (e.g. loading a raw checkpoint rather than a `.nemo` file), you need to re-apply the language embedding patch:

```python
import torch.nn as nn

# Assuming model is loaded and has _lang_embed_module
lang_embed = model._lang_embed_module  # nn.Embedding(12, 1024)

# Save original encoder forward
original_forward = model.encoder.forward

def encoder_forward_with_lang(audio_signal, length):
    encoded, encoded_len = original_forward(
        audio_signal=audio_signal, length=length
    )
    if hasattr(model, '_current_lang_ids') and model._current_lang_ids is not None:
        lang_ids = model._current_lang_ids
        if lang_ids.shape[0] == encoded.shape[0]:
            lang_bias = lang_embed(lang_ids)          # (B, 1024)
            lang_bias = lang_bias.unsqueeze(-1)        # (B, 1024, 1)
            encoded = encoded + lang_bias
        model._current_lang_ids = None  # clear after use
    return encoded, encoded_len

model.encoder.forward = encoder_forward_with_lang
```

## Decoding Configuration

The model supports two decoding modes:

### TDT Greedy (Primary — recommended)
```python
# Already configured in the checkpoint. Duration-aware decoding.
# Durations: [0, 1, 2, 3, 4] frames
# This is the default mode.
```

### CTC (Fallback — faster)
```python
# For CTC-only decoding (faster but less accurate):
transcriptions = model.transcribe(audio_files, return_hypotheses=False)
```

## Tokenizer

- **Type**: SentencePiece BPE
- **Vocab size**: 32,768 tokens
- **Files**: `tokenizer/tokenizer.model`, `tokenizer/vocab.txt`
- **Coverage**: All 12 supported Indic scripts + English + numerals

## Notes

- **Language embedding is mandatory** — always set `_current_lang_ids` before transcription
- Audio must be **16kHz mono** — the model expects this sample rate
- For long audio (>30s), consider chunking with overlap for best results
- The model was fine-tuned from NVIDIA Parakeet TDT 1.1B on Indic language data


python3 << 'PYEOF'
import torch
import sys

ckpt_path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

print("Loading checkpoint (metadata only)...")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

print(f"\n=== TOP-LEVEL KEYS ===")
for k in sorted(ckpt.keys()):
    v = ckpt[k]
    if isinstance(v, dict):
        print(f"  {k}: dict with {len(v)} keys")
    elif isinstance(v, (int, float, str, bool)):
        print(f"  {k}: {v}")
    else:
        print(f"  {k}: {type(v).__name__}")

print(f"\n=== STATE DICT SUMMARY ===")
state_dict = ckpt.get("state_dict", {})
print(f"Total tensors: {len(state_dict)}")

# Group by top-level module
from collections import defaultdict
groups = defaultdict(lambda: {"count": 0, "params": 0, "shapes": []})
for k, v in state_dict.items():
    parts = k.split(".")
    # Group by first 2 levels
    if len(parts) >= 2:
        group = f"{parts[0]}.{parts[1]}"
    else:
        group = parts[0]
    groups[group]["count"] += 1
    groups[group]["params"] += v.numel()

print(f"\n=== MODULE GROUPS (top-2 level) ===")
total_params = 0
for group in sorted(groups.keys()):
    g = groups[group]
    total_params += g["params"]
    print(f"  {group:50s}  tensors={g['count']:4d}  params={g['params']/1e6:10.2f}M")

print(f"\n  {'TOTAL':50s}  params={total_params/1e6:10.2f}M")

# Now show the FULL key list for non-encoder parts (decoder, joint, CTC, lang embed)
print(f"\n=== NON-ENCODER LAYER DETAILS ===")
for k, v in sorted(state_dict.items()):
    if not k.startswith("encoder."):
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel()/1e6:.4f}M  dtype={v.dtype}")

# Show language embedding specifically
print(f"\n=== LANGUAGE EMBEDDING ===")
for k, v in sorted(state_dict.items()):
    if "lang" in k.lower():
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel():,} params")

# Show a few encoder layers to understand structure
print(f"\n=== ENCODER LAYER STRUCTURE (first layer sample) ===")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.layers.0."):
        print(f"  {k:70s}  {str(v.shape):20s}")

print(f"\n=== GLOBAL STEP & EPOCH ===")
print(f"  global_step: {ckpt.get('global_step', 'N/A')}")
print(f"  epoch: {ckpt.get('epoch', 'N/A')}")

# Hyper params if present
if 'hyper_parameters' in ckpt:
    print(f"\n=== HYPER PARAMETERS (keys) ===")
    hp = ckpt['hyper_parameters']
    if isinstance(hp, dict):
        for hk in sorted(hp.keys())[:30]:
            print(f"  {hk}")
PYEOF

Inspect checkpoint structure and layers


Error: Exit code 1
/home/ubuntu/.local/lib/python3.10/site-packages/torchmetrics/utilities/imports.py:22: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import DistributionNotFound, get_distribution
Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
Loading checkpoint (metadata only)...
Traceback (most recent call last):
  File "<stdin>", line 7, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1530, in load
    return _load(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2122, in _load
    result = unpickler.load()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2111, in find_class
    return super().find_class(mod_name, name)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/__init__.py", line 15, in <module>
    from nemo.collections.asr import data, losses, models, modules
  File "/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/__init__.py", line 22, in <module>
    from nemo.collections.asr.models.clustering_diarizer import ClusteringDiarizer
  File "/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/clustering_diarizer.py", line 42, in <module>
    from nemo.collections.asr.parts.utils.vad_utils import (
  File "/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/utils/vad_utils.py", line 29, in <module>
    import matplotlib.pyplot as plt
  File "/home/ubuntu/.local/lib/python3.10/site-packages/matplotlib/__init__.py", line 264, in <module>
    _check_versions()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/matplotlib/__init__.py", line 260, in _check_versions
    raise ImportError(f"Matplotlib requires {modname}>={minver}; "
ImportError: Matplotlib requires dateutil>=2.7; you have 2.6.1


pip install python-dateutil --upgrade -q 2>&1 | tail -3

Upgrade dateutil


tts 0.0.1 requires transformers==4.55.2, but you have transformers 4.57.6 which is incompatible.
vastai 0.5.0 requires python-dateutil==2.6.1, but you have python-dateutil 2.9.0.post0 which is incompatible.
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)


python3 << 'PYEOF'
import torch

ckpt_path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

print("Loading checkpoint...")
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

print(f"\n=== TOP-LEVEL KEYS ===")
for k in sorted(ckpt.keys()):
    v = ckpt[k]
    if isinstance(v, dict):
        print(f"  {k}: dict with {len(v)} keys")
    elif isinstance(v, (int, float, str, bool)):
        print(f"  {k}: {v}")
    else:
        print(f"  {k}: {type(v).__name__}")

print(f"\n=== STATE DICT SUMMARY ===")
state_dict = ckpt.get("state_dict", {})
print(f"Total tensors: {len(state_dict)}")

from collections import defaultdict
groups = defaultdict(lambda: {"count": 0, "params": 0})
for k, v in state_dict.items():
    parts = k.split(".")
    if len(parts) >= 2:
        group = f"{parts[0]}.{parts[1]}"
    else:
        group = parts[0]
    groups[group]["count"] += 1
    groups[group]["params"] += v.numel()

print(f"\n=== MODULE GROUPS (top-2 level) ===")
total_params = 0
for group in sorted(groups.keys()):
    g = groups[group]
    total_params += g["params"]
    print(f"  {group:50s}  tensors={g['count']:4d}  params={g['params']/1e6:10.2f}M")

print(f"\n  {'TOTAL':50s}  params={total_params/1e6:10.2f}M")

print(f"\n=== NON-ENCODER LAYER DETAILS ===")
for k, v in sorted(state_dict.items()):
    if not k.startswith("encoder."):
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel()/1e6:.4f}M  dtype={v.dtype}")

print(f"\n=== LANGUAGE EMBEDDING ===")
for k, v in sorted(state_dict.items()):
    if "lang" in k.lower():
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel():,} params")
        # Print the embedding values stats
        print(f"    mean={v.float().mean():.6f}  std={v.float().std():.6f}  min={v.float().min():.6f}  max={v.float().max():.6f}")

print(f"\n=== ENCODER LAYER 0 STRUCTURE ===")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.layers.0."):
        print(f"  {k:70s}  {str(v.shape):20s}")

# Count encoder layers
layer_nums = set()
for k in state_dict:
    if k.startswith("encoder.layers."):
        layer_num = int(k.split(".")[2])
        layer_nums.add(layer_num)
print(f"\n=== ENCODER: {len(layer_nums)} layers (0-{max(layer_nums) if layer_nums else 'N/A'}) ===")

print(f"\n=== GLOBAL STEP & EPOCH ===")
print(f"  global_step: {ckpt.get('global_step', 'N/A')}")
print(f"  epoch: {ckpt.get('epoch', 'N/A')}")

# Check for hyper_parameters
if 'hyper_parameters' in ckpt:
    hp = ckpt['hyper_parameters']
    print(f"\n=== HYPER PARAMETERS ===")
    if isinstance(hp, dict):
        for hk in sorted(hp.keys())[:20]:
            print(f"  {hk}: {str(hp[hk])[:100]}")
PYEOF

Inspect checkpoint layers and params


Error: Exit code 1
/home/ubuntu/.local/lib/python3.10/site-packages/torchmetrics/utilities/imports.py:22: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import DistributionNotFound, get_distribution
Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
Loading checkpoint...
Traceback (most recent call last):
  File "<stdin>", line 6, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1530, in load
    return _load(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2122, in _load
    result = unpickler.load()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2111, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'nemo.collections.asr.parts.context_biasing.boosting_graph_batched'


NeMo version mismatch. Let me try loading with `weights_only=True` or bypassing the module import.


python3 << 'PYEOF'
import torch
import pickle
import io

ckpt_path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

# Custom unpickler that handles missing modules
class SafeUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        try:
            return super().find_class(module, name)
        except (ModuleNotFoundError, AttributeError):
            # Return a dummy class for missing NeMo modules
            return type(name, (), {})

class SafeLoader:
    @staticmethod
    def load(f, **kwargs):
        # Read the zip file entries
        import zipfile
        if isinstance(f, str):
            f = open(f, 'rb')
        return torch.load(f, map_location="cpu", weights_only=True)

# Try weights_only=True first
try:
    print("Trying weights_only=True...")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    print("Success with weights_only=True")
except Exception as e:
    print(f"weights_only=True failed: {e}")
    print("\nTrying with safe_globals...")
    # Use torch safe loading with allowed classes
    import sys
    import types

    # Create stub modules for missing NeMo
    missing = 'nemo.collections.asr.parts.context_biasing.boosting_graph_batched'
    parts = missing.split('.')
    for i in range(len(parts)):
        mod_name = '.'.join(parts[:i+1])
        if mod_name not in sys.modules:
            mod = types.ModuleType(mod_name)
            sys.modules[mod_name] = mod
            if i > 0:
                parent = '.'.join(parts[:i])
                setattr(sys.modules[parent], parts[i], mod)

    # Add missing classes
    for cls_name in ['BoostingGraphBatched', 'BoostingGraph']:
        if not hasattr(sys.modules[missing], cls_name):
            setattr(sys.modules[missing], cls_name, type(cls_name, (), {}))

    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    print("Success with module stubs")
PYEOF

Load checkpoint with module stubs


Error: Exit code 1
Trying weights_only=True...
weights_only=True failed: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
	(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
	(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
	WeightsUnpickler error: Unsupported global: GLOBAL omegaconf.dictconfig.DictConfig was not an allowed global by default. Please use `torch.serialization.add_safe_globals([omegaconf.dictconfig.DictConfig])` or the `torch.serialization.safe_globals([omegaconf.dictconfig.DictConfig])` context manager to allowlist this global if you trust this class/function.

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

Trying with safe_globals...
Traceback (most recent call last):
  File "<stdin>", line 54, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1530, in load
    return _load(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2122, in _load
    result = unpickler.load()
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 2111, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'nemo.collections.asr.parts.submodules'; 'nemo.collections.asr.parts' is not a package


python3 << 'PYEOF'
import torch
import sys
import types

# Pre-create ALL missing nemo submodule stubs BEFORE import
# This prevents nemo's __init__ chain from breaking

# List of known missing modules - we'll add more dynamically
stubs_needed = [
    'nemo.collections.asr.parts',
    'nemo.collections.asr.parts.context_biasing',
    'nemo.collections.asr.parts.context_biasing.boosting_graph_batched',
    'nemo.collections.asr.parts.submodules',
    'nemo.collections.asr.parts.submodules.rnnt_greedy_decoding',
]

# Actually, let's just add safe globals and use weights_only approach
from omegaconf import DictConfig, ListConfig, OmegaConf

torch.serialization.add_safe_globals([DictConfig, ListConfig])

# Try to add more needed globals
try:
    from nemo.core.config import hydra_runner
except:
    pass

# Just allowlist everything we need
print("Loading with safe globals allowlisting...")
try:
    ckpt = torch.load(
        "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt",
        map_location="cpu",
        weights_only=True
    )
except Exception as e:
    err = str(e)
    print(f"Still failing: {err[:200]}")
    # Extract all needed globals from the error
    import re
    globals_needed = re.findall(r'GLOBAL (\S+)', err)
    print(f"Globals needed: {globals_needed}")
PYEOF

Try safe globals approach


/home/ubuntu/.local/lib/python3.10/site-packages/torchmetrics/utilities/imports.py:22: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import DistributionNotFound, get_distribution
Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
Loading with safe globals allowlisting...
Still failing: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
	(1) In PyTorch 2.6, we changed the defa
Globals needed: ['omegaconf.base.ContainerMetadata']


python3 << 'PYEOF'
import torch
from omegaconf import DictConfig, ListConfig
from omegaconf.base import ContainerMetadata

# Keep adding safe globals iteratively
safe_classes = [DictConfig, ListConfig, ContainerMetadata]

# Try to find all needed classes
try:
    from omegaconf.nodes import ValueNode, AnyNode
    safe_classes.extend([ValueNode, AnyNode])
except: pass
try:
    from omegaconf import ValueNode
    safe_classes.append(ValueNode)
except: pass
try:
    from omegaconf._utils import Metadata
    safe_classes.append(Metadata)
except: pass
try:
    from omegaconf.basecontainer import BaseContainer
    safe_classes.append(BaseContainer)
except: pass
try:
    from omegaconf import MISSING
    safe_classes.append(type(MISSING))
except: pass
try:
    import enum
    from omegaconf import ReadonlyMode, DictKeyType
    safe_classes.extend([ReadonlyMode, DictKeyType])
except: pass
try:
    from omegaconf.base import Metadata as BaseMeta
    safe_classes.append(BaseMeta)
except: pass

# Add all omegaconf classes we can find
import omegaconf
for attr in dir(omegaconf):
    obj = getattr(omegaconf, attr)
    if isinstance(obj, type):
        safe_classes.append(obj)
for attr in dir(omegaconf.base):
    obj = getattr(omegaconf.base, attr)
    if isinstance(obj, type):
        safe_classes.append(obj)

# Also add builtins that might be needed
import collections
safe_classes.append(collections.defaultdict)

torch.serialization.add_safe_globals(list(set(safe_classes)))

path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

max_attempts = 10
for attempt in range(max_attempts):
    try:
        ckpt = torch.load(path, map_location="cpu", weights_only=True)
        print("SUCCESS!")
        break
    except Exception as e:
        err = str(e)
        import re
        globals_needed = re.findall(r'GLOBAL (\S+)', err)
        if not globals_needed:
            print(f"Non-global error: {err[:300]}")
            break
        for g in globals_needed:
            parts = g.rsplit('.', 1)
            if len(parts) == 2:
                try:
                    mod = __import__(parts[0], fromlist=[parts[1]])
                    cls = getattr(mod, parts[1])
                    torch.serialization.add_safe_globals([cls])
                    print(f"  Added: {g}")
                except Exception as ie:
                    print(f"  Failed to import {g}: {ie}")
                    break
        else:
            continue
        break
else:
    print("Exceeded max attempts")

if 'ckpt' in dir():
    state_dict = ckpt.get("state_dict", {})
    print(f"\nTotal tensors: {len(state_dict)}")
PYEOF

Iteratively add safe globals to load ckpt


  Added: typing.Any
Exceeded max attempts

  Added: typing.Any
Exceeded max attempts


python3 << 'PYEOF'
# Brute force: extract just the state_dict from the checkpoint using zipfile
# PyTorch checkpoints are ZIP files containing pickled tensors
import zipfile
import io
import struct

path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

# Check if it's a zip file
with open(path, 'rb') as f:
    magic = f.read(4)
    print(f"File magic: {magic.hex()} (PK={magic[:2] == b'PK'})")

if magic[:2] == b'PK':
    with zipfile.ZipFile(path, 'r') as z:
        names = z.namelist()
        print(f"ZIP entries: {len(names)}")
        # Show first 20 and last 10
        for n in names[:20]:
            print(f"  {n}")
        if len(names) > 30:
            print(f"  ... ({len(names)-30} more)")
        for n in names[-10:]:
            print(f"  {n}")
PYEOF

Inspect checkpoint ZIP structure


File magic: 504b0304 (PK=True)
ZIP entries: 6466
  archive/data.pkl
  archive/.format_version
  archive/.storage_alignment
  archive/byteorder
  archive/data/0
  archive/data/1
  archive/data/2
  archive/data/3
  archive/data/4
  archive/data/5
  archive/data/6
  archive/data/7
  archive/data/8
  archive/data/9
  archive/data/10
  archive/data/11
  archive/data/12
  archive/data/13
  archive/data/14
  archive/data/15
  ... (6436 more)
  archive/data/6452
  archive/data/6453
  archive/data/6454
  archive/data/6455
  archive/data/6456
  archive/data/6457
  archive/data/6458
  archive/data/6459
  archive/version
  archive/.data/serialization_id


python3 << 'PYEOF'
import torch
import sys
import types
import importlib

# Strategy: create stub modules for everything NeMo needs during unpickling
# Then load with weights_only=False

path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

# Attempt: use a custom Unpickler
import pickle
import io
import zipfile

class LenientUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        try:
            mod = importlib.import_module(module)
            return getattr(mod, name)
        except (ModuleNotFoundError, AttributeError, ImportError):
            # Return a dummy
            dummy = type(name, (), {'__reduce__': lambda self: (type(self), ())})
            return dummy

# Read the pickle data from the ZIP
with zipfile.ZipFile(path, 'r') as z:
    with z.open('archive/data.pkl') as f:
        pkl_data = f.read()

print(f"Pickle data size: {len(pkl_data)} bytes")

# Try to unpickle with lenient unpickler
# But we also need to handle the tensor storage references
# Let's just get the key names from the pickle

# Alternative: use torch's internal mechanism but patch missing modules
missing_modules = set()

def create_stub_chain(fullname):
    """Create all parent stub modules for a given module path."""
    parts = fullname.split('.')
    for i in range(len(parts)):
        mod_name = '.'.join(parts[:i+1])
        if mod_name not in sys.modules:
            mod = types.ModuleType(mod_name)
            sys.modules[mod_name] = mod
            if i > 0:
                parent_name = '.'.join(parts[:i])
                if parent_name in sys.modules:
                    setattr(sys.modules[parent_name], parts[i], mod)

# First pass: try loading and collect all missing modules
for attempt in range(20):
    try:
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        print("Loaded successfully!")
        break
    except ModuleNotFoundError as e:
        missing = str(e).split("'")[1] if "'" in str(e) else str(e)
        print(f"  Stubbing: {missing}")
        create_stub_chain(missing)
        # Add dummy classes
        mod = sys.modules.get(missing)
        if mod:
            for cls_name in ['BoostingGraphBatched', 'BoostingGraph', 'ContextGraph',
                           'ContextGraphBatched', 'WordTokenizer', 'SentencePieceTokenizer']:
                if not hasattr(mod, cls_name):
                    setattr(mod, cls_name, type(cls_name, (), {}))
    except Exception as e:
        print(f"  Other error: {type(e).__name__}: {str(e)[:200]}")
        break

if 'ckpt' in dir():
    state_dict = ckpt.get("state_dict", {})
    from collections import defaultdict

    print(f"\n{'='*70}")
    print(f"CHECKPOINT LOADED SUCCESSFULLY")
    print(f"{'='*70}")
    print(f"\nTotal tensors in state_dict: {len(state_dict)}")
    print(f"Global step: {ckpt.get('global_step', 'N/A')}")
    print(f"Epoch: {ckpt.get('epoch', 'N/A')}")

    # Group by module
    groups = defaultdict(lambda: {"count": 0, "params": 0})
    for k, v in state_dict.items():
        parts = k.split(".")
        group = f"{parts[0]}.{parts[1]}" if len(parts) >= 2 else parts[0]
        groups[group]["count"] += 1
        groups[group]["params"] += v.numel()

    print(f"\n{'='*70}")
    print(f"MODULE GROUPS")
    print(f"{'='*70}")
    total_params = 0
    for group in sorted(groups.keys()):
        g = groups[group]
        total_params += g["params"]
        print(f"  {group:50s}  tensors={g['count']:4d}  params={g['params']/1e6:10.2f}M")
    print(f"  {'TOTAL':50s}  params={total_params/1e6:10.2f}M")

    # Non-encoder details
    print(f"\n{'='*70}")
    print(f"NON-ENCODER LAYERS (decoder, joint, CTC, lang embed)")
    print(f"{'='*70}")
    for k, v in sorted(state_dict.items()):
        if not k.startswith("encoder."):
            print(f"  {k:70s}  {str(v.shape):30s}  {v.numel()/1e6:.4f}M  {v.dtype}")

    # Language embedding
    print(f"\n{'='*70}")
    print(f"LANGUAGE EMBEDDING DETAILS")
    print(f"{'='*70}")
    for k, v in sorted(state_dict.items()):
        if "lang" in k.lower():
            print(f"  {k}")
            print(f"    shape: {v.shape}")
            print(f"    dtype: {v.dtype}")
            print(f"    params: {v.numel():,}")
            vf = v.float()
            print(f"    mean={vf.mean():.6f}  std={vf.std():.6f}  min={vf.min():.6f}  max={vf.max():.6f}")

    # Encoder structure sample
    print(f"\n{'='*70}")
    print(f"ENCODER LAYER 0 (structure sample)")
    print(f"{'='*70}")
    for k, v in sorted(state_dict.items()):
        if k.startswith("encoder.layers.0."):
            print(f"  {k:70s}  {str(v.shape):20s}")

    layer_nums = set()
    for k in state_dict:
        if k.startswith("encoder.layers."):
            layer_nums.add(int(k.split(".")[2]))
    print(f"\n  Total encoder layers: {len(layer_nums)} (0-{max(layer_nums)})")

    # Encoder non-layer params
    print(f"\n{'='*70}")
    print(f"ENCODER NON-LAYER PARAMS")
    print(f"{'='*70}")
    for k, v in sorted(state_dict.items()):
        if k.startswith("encoder.") and not k.startswith("encoder.layers."):
            print(f"  {k:70s}  {str(v.shape):20s}")
PYEOF

Load ckpt with module stubs iteratively


/home/ubuntu/.local/lib/python3.10/site-packages/torchmetrics/utilities/imports.py:22: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import DistributionNotFound, get_distribution
Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info
Pickle data size: 15272519 bytes
  Stubbing: nemo.collections.asr.parts.context_biasing.boosting_graph_batched
  Other error: AttributeError: Can't get attribute 'BoostingTreeModelConfig' on <module 'nemo.collections.asr.parts.context_biasing.boosting_graph_batched'>


python3 << 'PYEOF'
import torch
import sys
import types
import importlib

path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

def create_stub_chain(fullname):
    parts = fullname.split('.')
    for i in range(len(parts)):
        mod_name = '.'.join(parts[:i+1])
        if mod_name not in sys.modules:
            mod = types.ModuleType(mod_name)
            sys.modules[mod_name] = mod
            if i > 0:
                parent_name = '.'.join(parts[:i])
                if parent_name in sys.modules:
                    setattr(sys.modules[parent_name], parts[i], mod)

# Pre-stub the problematic module
create_stub_chain('nemo.collections.asr.parts.context_biasing.boosting_graph_batched')

# Add a catch-all __getattr__ to return dummy classes
class DummyModule(types.ModuleType):
    def __getattr__(self, name):
        # Return a dummy dataclass-like class
        cls = type(name, (), {
            '__init__': lambda self, **kwargs: self.__dict__.update(kwargs),
            '__reduce__': lambda self: (type(self), (), self.__dict__),
            '__setstate__': lambda self, state: self.__dict__.update(state) if isinstance(state, dict) else None,
        })
        setattr(self, name, cls)
        return cls

# Replace the stub with our DummyModule
mod_name = 'nemo.collections.asr.parts.context_biasing.boosting_graph_batched'
dummy = DummyModule(mod_name)
sys.modules[mod_name] = dummy
parent = 'nemo.collections.asr.parts.context_biasing'
if parent in sys.modules:
    setattr(sys.modules[parent], 'boosting_graph_batched', dummy)

# Try loading again, catching more missing modules
for attempt in range(30):
    try:
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        print("Loaded successfully!")
        break
    except ModuleNotFoundError as e:
        missing = str(e).split("'")[1] if "'" in str(e) else str(e)
        print(f"  Stubbing: {missing}")
        create_stub_chain(missing)
        # Replace with DummyModule
        dm = DummyModule(missing)
        sys.modules[missing] = dm
        parts = missing.rsplit('.', 1)
        if len(parts) == 2 and parts[0] in sys.modules:
            setattr(sys.modules[parts[0]], parts[1], dm)
    except AttributeError as e:
        print(f"  AttrError: {e}")
        # The module exists but missing class - DummyModule should handle it
        # Try reloading
        break
    except Exception as e:
        print(f"  Error: {type(e).__name__}: {str(e)[:300]}")
        break

if 'ckpt' not in dir():
    print("Failed to load. Exiting.")
    sys.exit(1)

state_dict = ckpt.get("state_dict", {})
from collections import defaultdict

print(f"\n{'='*70}")
print(f"CHECKPOINT LOADED SUCCESSFULLY")
print(f"{'='*70}")
print(f"Total tensors in state_dict: {len(state_dict)}")
print(f"Global step: {ckpt.get('global_step', 'N/A')}")
print(f"Epoch: {ckpt.get('epoch', 'N/A')}")

groups = defaultdict(lambda: {"count": 0, "params": 0})
for k, v in state_dict.items():
    parts = k.split(".")
    group = f"{parts[0]}.{parts[1]}" if len(parts) >= 2 else parts[0]
    groups[group]["count"] += 1
    groups[group]["params"] += v.numel()

print(f"\n{'='*70}")
print(f"MODULE GROUPS")
print(f"{'='*70}")
total_params = 0
for group in sorted(groups.keys()):
    g = groups[group]
    total_params += g["params"]
    print(f"  {group:50s}  tensors={g['count']:4d}  params={g['params']/1e6:10.2f}M")
print(f"  {'TOTAL':50s}  params={total_params/1e6:10.2f}M")

print(f"\n{'='*70}")
print(f"NON-ENCODER LAYERS")
print(f"{'='*70}")
for k, v in sorted(state_dict.items()):
    if not k.startswith("encoder."):
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel()/1e6:.4f}M  {v.dtype}")

print(f"\n{'='*70}")
print(f"LANGUAGE EMBEDDING")
print(f"{'='*70}")
for k, v in sorted(state_dict.items()):
    if "lang" in k.lower():
        print(f"  {k}")
        print(f"    shape={v.shape}  dtype={v.dtype}  params={v.numel():,}")
        vf = v.float()
        print(f"    mean={vf.mean():.6f}  std={vf.std():.6f}  min={vf.min():.6f}  max={vf.max():.6f}")
        # Print per-language norms
        for i in range(v.shape[0]):
            norm = vf[i].norm().item()
            print(f"    lang_id={i:2d}  norm={norm:.4f}")

print(f"\n{'='*70}")
print(f"ENCODER LAYER 0")
print(f"{'='*70}")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.layers.0."):
        print(f"  {k:70s}  {str(v.shape):20s}")

layer_nums = set()
for k in state_dict:
    if k.startswith("encoder.layers."):
        layer_nums.add(int(k.split(".")[2]))
print(f"\n  Total encoder layers: {len(layer_nums)} (0-{max(layer_nums)})")

print(f"\n{'='*70}")
print(f"ENCODER NON-LAYER PARAMS")
print(f"{'='*70}")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.") and not k.startswith("encoder.layers."):
        print(f"  {k:70s}  {str(v.shape):20s}")
PYEOF

Load checkpoint with DummyModule approach


Error: Exit code 1
  Stubbing: nemo.collections.asr.parts.submodules
  Error: TypeError: 'type' object is not iterable
Failed to load. Exiting.


python3 << 'PYEOF'
import torch
import sys
import types

path = "/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt"

# The issue is nemo.collections.asr.parts is being set to a non-package Module
# We need to intervene BEFORE nemo is imported at all

# Remove all nemo modules from sys.modules first
to_remove = [k for k in sys.modules if k.startswith('nemo')]
for k in to_remove:
    del sys.modules[k]

class AutoStubModule(types.ModuleType):
    """A module that auto-creates attributes (classes and submodules) on demand."""
    def __init__(self, name):
        super().__init__(name)
        self.__path__ = []  # Make it look like a package

    def __getattr__(self, name):
        fullname = f"{self.__name__}.{name}"
        if fullname in sys.modules:
            return sys.modules[fullname]
        # Check if it should be a submodule
        # Create a dummy class by default
        cls = type(name, (), {
            '__init__': lambda self, *a, **kw: self.__dict__.update(kw),
            '__reduce__': lambda self: (type(self), ()),
            '__setstate__': lambda self, state: (self.__dict__.update(state) if isinstance(state, dict) else None),
        })
        setattr(self, name, cls)
        return cls

# Install auto-stub for the problematic nemo subpackage
# Key: 'nemo.collections.asr.parts' and its children need to be packages
stubs = [
    'nemo', 'nemo.collections', 'nemo.collections.asr', 'nemo.collections.asr.parts',
    'nemo.collections.asr.parts.context_biasing',
    'nemo.collections.asr.parts.context_biasing.boosting_graph_batched',
    'nemo.collections.asr.parts.submodules',
    'nemo.collections.asr.parts.submodules.rnnt_greedy_decoding',
]
for s in stubs:
    sys.modules[s] = AutoStubModule(s)
    parts = s.rsplit('.', 1)
    if len(parts) == 2 and parts[0] in sys.modules:
        setattr(sys.modules[parts[0]], parts[1], sys.modules[s])

# Now let the real nemo import over these stubs where it can
# Actually, let's just try loading directly - torch.load will trigger imports
for attempt in range(50):
    try:
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        print("Loaded!")
        break
    except ModuleNotFoundError as e:
        missing = str(e).split("'")[1] if "'" in str(e) else str(e).split("No module named ")[-1].strip()
        stub = AutoStubModule(missing)
        sys.modules[missing] = stub
        parts = missing.rsplit('.', 1)
        if len(parts) == 2 and parts[0] in sys.modules:
            setattr(sys.modules[parts[0]], parts[1], stub)
    except Exception as e:
        print(f"Error [{attempt}]: {type(e).__name__}: {str(e)[:200]}")
        break

if 'ckpt' not in dir():
    sys.exit(1)

state_dict = ckpt.get("state_dict", {})
from collections import defaultdict

print(f"\nTotal tensors: {len(state_dict)}")
print(f"Global step: {ckpt.get('global_step', 'N/A')}")
print(f"Epoch: {ckpt.get('epoch', 'N/A')}")

groups = defaultdict(lambda: {"count": 0, "params": 0})
for k, v in state_dict.items():
    parts = k.split(".")
    group = f"{parts[0]}.{parts[1]}" if len(parts) >= 2 else parts[0]
    groups[group]["count"] += 1
    groups[group]["params"] += v.numel()

print(f"\nMODULE GROUPS:")
total_params = 0
for group in sorted(groups.keys()):
    g = groups[group]
    total_params += g["params"]
    print(f"  {group:50s}  tensors={g['count']:4d}  params={g['params']/1e6:10.2f}M")
print(f"  {'TOTAL':50s}  params={total_params/1e6:10.2f}M")

print(f"\nNON-ENCODER LAYERS:")
for k, v in sorted(state_dict.items()):
    if not k.startswith("encoder."):
        print(f"  {k:70s}  {str(v.shape):30s}  {v.numel()/1e6:.4f}M  {v.dtype}")

print(f"\nLANGUAGE EMBEDDING:")
for k, v in sorted(state_dict.items()):
    if "lang" in k.lower():
        vf = v.float()
        print(f"  {k}  shape={v.shape}  params={v.numel():,}")
        print(f"    mean={vf.mean():.6f}  std={vf.std():.6f}")
        for i in range(v.shape[0]):
            print(f"    lang_id={i:2d}  norm={vf[i].norm():.4f}")

print(f"\nENCODER LAYER 0:")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.layers.0."):
        print(f"  {k:70s}  {str(v.shape):20s}")

layer_nums = set()
for k in state_dict:
    if k.startswith("encoder.layers."):
        layer_nums.add(int(k.split(".")[2]))
print(f"\n  Encoder layers: {len(layer_nums)} (0-{max(layer_nums)})")

print(f"\nENCODER NON-LAYER:")
for k, v in sorted(state_dict.items()):
    if k.startswith("encoder.") and not k.startswith("encoder.layers."):
        print(f"  {k:70s}  {str(v.shape):20s}")
PYEOF

Load ckpt with full auto-stub system


Loaded!

Total tensors: 1713
Global step: 260000
Epoch: 0

MODULE GROUPS:
  _lang_embed.embed                                   tensors=   1  params=      0.01M
  _lang_embed_module.embed                            tensors=   1  params=      0.01M
  ctc_decoder.decoder_layers                          tensors=   2  params=     33.59M
  decoder.prediction                                  tensors=   9  params=     27.54M
  encoder.layers                                      tensors=1680  params=   1058.81M
  encoder.pre_encode                                  tensors=  12  params=      2.76M
  joint.enc                                           tensors=   2  params=      0.66M
  joint.joint_net                                     tensors=   2  params=     21.01M
  joint.pred                                          tensors=   2  params=      0.41M
  preprocessor.featurizer                             tensors=   2  params=      0.02M
  TOTAL                                               params=   1144.82M

NON-ENCODER LAYERS:
  _lang_embed.embed.weight                                                torch.Size([12, 1024])          0.0123M  torch.float32
  _lang_embed_module.embed.weight                                         torch.Size([12, 1024])          0.0123M  torch.float32
  ctc_decoder.decoder_layers.0.bias                                       torch.Size([32769])             0.0328M  torch.float32
  ctc_decoder.decoder_layers.0.weight                                     torch.Size([32769, 1024, 1])    33.5555M  torch.float32
  decoder.prediction.dec_rnn.lstm.bias_hh_l0                              torch.Size([2560])              0.0026M  torch.float32
  decoder.prediction.dec_rnn.lstm.bias_hh_l1                              torch.Size([2560])              0.0026M  torch.float32
  decoder.prediction.dec_rnn.lstm.bias_ih_l0                              torch.Size([2560])              0.0026M  torch.float32
  decoder.prediction.dec_rnn.lstm.bias_ih_l1                              torch.Size([2560])              0.0026M  torch.float32
  decoder.prediction.dec_rnn.lstm.weight_hh_l0                            torch.Size([2560, 640])         1.6384M  torch.float32
  decoder.prediction.dec_rnn.lstm.weight_hh_l1                            torch.Size([2560, 640])         1.6384M  torch.float32
  decoder.prediction.dec_rnn.lstm.weight_ih_l0                            torch.Size([2560, 640])         1.6384M  torch.float32
  decoder.prediction.dec_rnn.lstm.weight_ih_l1                            torch.Size([2560, 640])         1.6384M  torch.float32
  decoder.prediction.embed.weight                                         torch.Size([32769, 640])        20.9722M  torch.float32
  joint.enc.bias                                                          torch.Size([640])               0.0006M  torch.float32
  joint.enc.weight                                                        torch.Size([640, 1024])         0.6554M  torch.float32
  joint.joint_net.2.bias                                                  torch.Size([32774])             0.0328M  torch.float32
  joint.joint_net.2.weight                                                torch.Size([32774, 640])        20.9754M  torch.float32
  joint.pred.bias                                                         torch.Size([640])               0.0006M  torch.float32
  joint.pred.weight                                                       torch.Size([640, 640])          0.4096M  torch.float32
  preprocessor.featurizer.fb                                              torch.Size([1, 80, 257])        0.0206M  torch.float32
  preprocessor.featurizer.window                                          torch.Size([400])               0.0004M  torch.float32

LANGUAGE EMBEDDING:
  _lang_embed.embed.weight  shape=torch.Size([12, 1024])  params=12,288
    mean=0.003843  std=0.109853
    lang_id= 0  norm=2.3555
    lang_id= 1  norm=3.4264
    lang_id= 2  norm=3.1921
    lang_id= 3  norm=2.6004
    lang_id= 4  norm=3.3673
    lang_id= 5  norm=3.9847
    lang_id= 6  norm=3.9113
    lang_id= 7  norm=3.3156
    lang_id= 8  norm=3.3969
    lang_id= 9  norm=5.4129
    lang_id=10  norm=4.2885
    lang_id=11  norm=1.2159
  _lang_embed_module.embed.weight  shape=torch.Size([12, 1024])  params=12,288
    mean=0.003843  std=0.109853
    lang_id= 0  norm=2.3555
    lang_id= 1  norm=3.4264
    lang_id= 2  norm=3.1921
    lang_id= 3  norm=2.6004
    lang_id= 4  norm=3.3673
    lang_id= 5  norm=3.9847
    lang_id= 6  norm=3.9113
    lang_id= 7  norm=3.3156
    lang_id= 8  norm=3.3969
    lang_id= 9  norm=5.4129
    lang_id=10  norm=4.2885
    lang_id=11  norm=1.2159

ENCODER LAYER 0:
  encoder.layers.0.conv.batch_norm.bias                                   torch.Size([1024])
  encoder.layers.0.conv.batch_norm.num_batches_tracked                    torch.Size([])
  encoder.layers.0.conv.batch_norm.running_mean                           torch.Size([1024])
  encoder.layers.0.conv.batch_norm.running_var                            torch.Size([1024])
  encoder.layers.0.conv.batch_norm.weight                                 torch.Size([1024])
  encoder.layers.0.conv.depthwise_conv.bias                               torch.Size([1024])
  encoder.layers.0.conv.depthwise_conv.weight                             torch.Size([1024, 1, 9])
  encoder.layers.0.conv.pointwise_conv1.bias                              torch.Size([2048])
  encoder.layers.0.conv.pointwise_conv1.weight                            torch.Size([2048, 1024, 1])
  encoder.layers.0.conv.pointwise_conv2.bias                              torch.Size([1024])
  encoder.layers.0.conv.pointwise_conv2.weight                            torch.Size([1024, 1024, 1])
  encoder.layers.0.feed_forward1.linear1.bias                             torch.Size([4096])
  encoder.layers.0.feed_forward1.linear1.weight                           torch.Size([4096, 1024])
  encoder.layers.0.feed_forward1.linear2.bias                             torch.Size([1024])
  encoder.layers.0.feed_forward1.linear2.weight                           torch.Size([1024, 4096])
  encoder.layers.0.feed_forward2.linear1.bias                             torch.Size([4096])
  encoder.layers.0.feed_forward2.linear1.weight                           torch.Size([4096, 1024])
  encoder.layers.0.feed_forward2.linear2.bias                             torch.Size([1024])
  encoder.layers.0.feed_forward2.linear2.weight                           torch.Size([1024, 4096])
  encoder.layers.0.norm_conv.bias                                         torch.Size([1024])
  encoder.layers.0.norm_conv.weight                                       torch.Size([1024])
  encoder.layers.0.norm_feed_forward1.bias                                torch.Size([1024])
  encoder.layers.0.norm_feed_forward1.weight                              torch.Size([1024])
  encoder.layers.0.norm_feed_forward2.bias                                torch.Size([1024])
  encoder.layers.0.norm_feed_forward2.weight                              torch.Size([1024])
  encoder.layers.0.norm_out.bias                                          torch.Size([1024])
  encoder.layers.0.norm_out.weight                                        torch.Size([1024])
  encoder.layers.0.norm_self_att.bias                                     torch.Size([1024])
  encoder.layers.0.norm_self_att.weight                                   torch.Size([1024])
  encoder.layers.0.self_attn.linear_k.bias                                torch.Size([1024])
  encoder.layers.0.self_attn.linear_k.weight                              torch.Size([1024, 1024])
  encoder.layers.0.self_attn.linear_out.bias                              torch.Size([1024])
  encoder.layers.0.self_attn.linear_out.weight                            torch.Size([1024, 1024])
  encoder.layers.0.self_attn.linear_pos.weight                            torch.Size([1024, 1024])
  encoder.layers.0.self_attn.linear_q.bias                                torch.Size([1024])
  encoder.layers.0.self_attn.linear_q.weight                              torch.Size([1024, 1024])
  encoder.layers.0.self_attn.linear_v.bias                                torch.Size([1024])
  encoder.layers.0.self_attn.linear_v.weight                              torch.Size([1024, 1024])
  encoder.layers.0.self_attn.pos_bias_u                                   torch.Size([8, 128])
  encoder.layers.0.self_attn.pos_bias_v                                   torch.Size([8, 128])

  Encoder layers: 42 (0-41)

ENCODER NON-LAYER:
  encoder.pre_encode.conv.0.bias                                          torch.Size([256])
  encoder.pre_encode.conv.0.weight                                        torch.Size([256, 1, 3, 3])
  encoder.pre_encode.conv.2.bias                                          torch.Size([256])
  encoder.pre_encode.conv.2.weight                                        torch.Size([256, 1, 3, 3])
  encoder.pre_encode.conv.3.bias                                          torch.Size([256])
  encoder.pre_encode.conv.3.weight                                        torch.Size([256, 256, 1, 1])
  encoder.pre_encode.conv.5.bias                                          torch.Size([256])
  encoder.pre_encode.conv.5.weight                                        torch.Size([256, 1, 3, 3])
  encoder.pre_encode.conv.6.bias                                          torch.Size([256])
  encoder.pre_encode.conv.6.weight                                        torch.Size([256, 256, 1, 1])
  encoder.pre_encode.out.bias                                             torch.Size([1024])
  encoder.pre_encode.out.weight                                           torch.Size([1024, 2560])


Got everything. Here's the complete summary:

---

## Maya ASR Stage 1 — Model Structure Summary

**Checkpoint:** `maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt` (step 260,000, 13.75GB)

### Total: 1,144.82M parameters (1,713 tensors)

---

### Architecture Overview

```
Audio (16kHz mono)
   ↓
Mel Spectrogram (80-dim filterbank)          preprocessor.featurizer     0.02M
   ↓
Pre-Encode (DW-Striding 8× subsampling)     encoder.pre_encode          2.76M
   ↓
42-layer FastConformer Encoder               encoder.layers          1,058.81M
   ↓
 + Language Embedding (additive bias)        _lang_embed                 0.01M  ← ADDED
   ↓
   ├─── TDT Decoder (2-layer LSTM)           decoder.prediction         27.54M  ← NEW (random init)
   │        ↓
   │    Joint Network (640-dim)              joint.*                    22.07M  ← NEW (random init)
   │    (enc proj + pred proj + joint_net)
   │    + 5 duration heads (output=32774 = 32769 vocab + 5 durations)
   │
   └─── CTC Head (linear 1024→32769)        ctc_decoder                33.59M  ← NEW (random init)
```

---

### Key Modifications from Base Parakeet TDT 1.1B

#### 1. Language Embedding (the main addition)
- **Layer:** `_lang_embed.embed.weight` — `nn.Embedding(12, 1024)` = **12,288 params**
- **Mechanism:** Maps a language ID (0-11) to a 1024-dim vector, unsqueezed to `(B, 1024, 1)`, then **added to every frame** of the encoder output `(B, 1024, T)`
- **Injection point:** Monkey-patched into `encoder.forward()` — runs AFTER the full 42-layer encoder, BEFORE decoder/joint
- **Languages:** hi=0, bn=1, ta=2, te=3, mr=4, gu=5, kn=6, ml=7, pa=8, or=9, as=10, en=11

**Trained embedding norms** (higher = more learned divergence from base):

| Lang ID | Language   | Norm   |
|---------|------------|--------|
| 0       | Hindi      | 2.36   |
| 1       | Bengali    | 3.43   |
| 2       | Tamil      | 3.19   |
| 3       | Telugu     | 2.60   |
| 4       | Marathi    | 3.37   |
| 5       | Gujarati   | 3.98   |
| 6       | Kannada    | 3.91   |
| 7       | Malayalam  | 3.32   |
| 8       | Punjabi    | 3.40   |
| 9       | Odia       | **5.41** |
| 10      | Assamese   | 4.29   |
| 11      | English    | **1.22** |

English has the lowest norm (makes sense — the base Parakeet encoder was English-pretrained, so minimal correction needed). Odia and Assamese have the highest norms, indicating the model learned the most correction for these low-resource languages.

#### 2. New Decoder + Joint + CTC (83.2M params, all randomly initialized)

**Decoder (27.54M):** 2-layer LSTM prediction network
- `decoder.prediction.embed.weight` — `[32769, 640]` (token embedding, 21.0M)
- `decoder.prediction.dec_rnn.lstm` — 2-layer LSTM, input=640, hidden=640 (6.6M)

**Joint Network (22.07M):** Encoder-decoder fusion + output projection
- `joint.enc` — Linear(1024→640) projects encoder to joint space
- `joint.pred` — Linear(640→640) projects decoder to joint space
- `joint.joint_net.2` — Linear(640→**32774**) — note: 32769 vocab + **5 TDT duration heads**

**CTC Head (33.59M):** Conv1D(1024→32769) auxiliary decoder

#### 3. Training Strategy (from `scripts/train_prod.py`)

| Technique | Detail |
|-----------|--------|
| **Differential LR** | Encoder: 2e-5 (preserve pretrained), Heads: 1e-3 (fast adapt) |
| **Encoder Freeze** | First 5K steps frozen, then unfrozen |
| **CTC Warmup** | Loss weight ramped 0→0.3 over 3K steps |
| **TDT Loss** | Durations [0,1,2,3,4], sigma=0.02, omega=0.1 |
| **Language rebalancing** | Temperature T=0.3 (upweights low-resource langs) |
| **OOM protection** | Triple-cap: batch size, batch duration (120s), max tokens (400) |

#### 4. FastConformer Encoder Layer Structure (per layer, ×42)

Each of the 42 layers has these sublayers (40 tensors per layer):
- `self_attn` — Multi-head attention (Q/K/V/Out projections, 1024→1024, 8 heads × 128 dim) + relative positional biases
- `feed_forward1` — FFN (1024→4096→1024, ReLU)
- `conv` — Depthwise separable convolution (pointwise1 1024→2048, depthwise kernel=9, pointwise2 1024→1024, BatchNorm)
- `feed_forward2` — FFN (1024→4096→1024, ReLU)
- LayerNorms: `norm_self_att`, `norm_conv`, `norm_feed_forward1`, `norm_feed_forward2`, `norm_out`

#### 5. Pre-Encode (Subsampling)

8× downsampling via depthwise-strided convolutions:
- 3 depthwise conv layers (256 channels, 3×3 kernel) with pointwise projections
- Final linear projection: 2560→1024

---

### How Language Conditioning Works (code flow)

1. **Data loading** (`NeMoTarOffsetDataset`): Each sample returns 5-tuple: `(signal, signal_len, tokens, token_len, lang_id)`
2. **Training step patch** (`patch_training_step`): Extracts `lang_ids` from batch position [4], stores on `model._current_lang_ids`, passes standard 4-tuple to NeMo
3. **Encoder forward patch** (`patch_encoder_for_lang`): After the real encoder runs, adds `lang_embed(lang_ids)` as a broadcast bias to the encoded features
4. **Cleared after use** to prevent stale values leaking into validation

The language embedding is **mandatory at inference** — set `model._current_lang_ids` before calling `model.transcribe()`.


arent we doing TDT based trianing ? why do we have CTC ?


Good question. It's **both** — this is a **hybrid TDT+CTC** training setup, by design. Here's why:

### The CTC head is an auxiliary loss, not a separate decoder

From the config (`configs/train/stage1_prod_8xh200.yaml`):
```yaml
aux_ctc:
  ctc_loss_weight: 0.3   # secondary loss
```

The model class is literally `EncDecHybridRNNTCTCBPEModel` — NeMo's hybrid that trains **both objectives on the same encoder simultaneously**:

```
                    ┌─── TDT Decoder + Joint ──→ TDT Loss (primary, weight=0.7)
Encoder output ─────┤
                    └─── CTC Linear Head ──────→ CTC Loss (auxiliary, weight=0.3)
```

### Why keep CTC alongside TDT?

1. **Encoder regularization** — CTC provides a frame-level alignment signal that acts as a regularizer on the encoder. The TDT loss alone works through the joint network, so the encoder gets sparser gradients. CTC gives direct per-frame supervision, which helps the encoder learn better representations faster.

2. **Pretrained encoder stability** — The Parakeet TDT 1.1B encoder was originally trained with this same hybrid objective. Keeping the CTC loss maintains consistency with how the encoder was pretrained, reducing distribution shift.

3. **CTC warmup protects early training** — From `train_prod.py:245-263`, the `CTCWarmupCallback` ramps CTC weight from 0→0.3 over the first 3K steps. This prevents random CTC gradients (from the untrained CTC head) from corrupting the pretrained encoder early on, while still getting the regularization benefit once the head stabilizes.

4. **Fallback decoding option** — At inference you can optionally use CTC greedy decoding (faster, no autoregressive loop) if you need speed over accuracy. The primary path remains TDT greedy.

### At inference

Only the TDT decoder path is used by default. The CTC head just sits there unused unless you explicitly request CTC decoding. The 33.6M CTC params are essentially "free" regularization during training — they cost VRAM but significantly improve encoder quality.


is CTC here any advantage or TDT ? what other lossess can we add here if nemo is multi objective trainable


### CTC vs TDT — is it actually helping?

Honestly, the benefit is **marginal and debatable** at this point in training. Here's the real tradeoff:

**Where CTC helps:**
- Early training stabilization — gives the encoder a direct gradient signal while the TDT joint network is still random garbage. The encoder doesn't have to wait for the decoder+joint to converge before getting useful learning signal.
- Acts as a monotonic alignment prior — prevents the encoder from learning representations that are hard to align temporally

**Where CTC is redundant or even hurts:**
- TDT already models alignment explicitly via duration heads (0-4 frame skips). The alignment signal CTC provides is something TDT handles natively.
- CTC's **conditional independence assumption** (each frame predicted independently) conflicts with TDT's autoregressive decoder. You're training the same encoder with two contradictory assumptions about output dependency.
- After 260K steps the encoder is well past the stabilization phase — the CTC loss is mostly just adding noise to gradients at this point.
- **33.6M params** in the CTC head that serve no purpose at inference. That's VRAM during training for diminishing returns.

**Verdict:** CTC was useful for the first ~10-20K steps. At step 260K, you could probably drop it to 0.0 or remove it entirely without hurting — and might even see a small TDT improvement from cleaner gradients.

---

### Other losses NeMo supports for multi-objective training

Let me check what the framework actually supports.


Agent

Find NeMo multi-objective losses

Prompt:
Search the installed NeMo package and the local project for all available auxiliary losses, multi-objective training options, and loss functions that can be used with EncDecHybridRNNTCTCBPEModel or similar ASR models.

Look in:
1. The installed nemo package: search for loss classes, auxiliary objectives, multi-task losses in /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/
2. Specifically look at:
   - nemo/collections/asr/losses/ directory
   - nemo/collections/asr/models/ for any multi-objective or hybrid models
   - Any RNNT/TDT/CTC/transducer loss variants
   - Any language ID losses, contrastive losses, or self-supervised objectives
3. Also check if there are any auxiliary loss hooks or mixins

Report all loss classes found, their purpose, and whether they could be added as auxiliary objectives. Be thorough - read the actual loss files.


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/`

Prompt:
Search the installed NeMo package and the local project for all available auxiliary losses, multi-objective training options, and loss functions that can be used with EncDecHybridRNNTCTCBPEModel or similar ASR models.

Look in:
1. The installed nemo package: search for loss classes, auxiliary objectives, multi-task losses in /home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/
2. Specifically look at:
   - nemo/collections/asr/losses/ directory
   - nemo/collections/asr/models/ for any multi-objective or hybrid models
   - Any RNNT/TDT/CTC/transducer loss variants
   - Any language ID losses, contrastive losses, or self-supervised objectives
3. Also check if there are any auxiliary loss hooks or mixins

Report all loss classes found, their purpose, and whether they could be added as auxiliary objectives. Be thorough - read the actual loss files.


### Is CTC actually an advantage here alongside TDT?

**At step 260K — probably not anymore.** CTC was valuable for early stabilization, but at this stage:
- TDT's duration heads already provide explicit alignment modeling (superior to CTC's framewise conditional independence)
- The two objectives have conflicting assumptions about output dependency
- You're spending 33.6M params and gradient bandwidth on something unused at inference

You could safely anneal `ctc_loss_weight` to 0 for the remaining training.

---

### Other losses you can add in NeMo's multi-objective framework

Here's what's actually available and worth considering:

#### 1. InterCTC (best bang for buck for a 42-layer encoder)

This is the most relevant one. It applies CTC losses at **intermediate encoder layers**, not just the final output. For a 42-layer encoder, this dramatically improves gradient flow to early layers.

```yaml
interctc:
  loss_weights: [0.1, 0.1]      # auxiliary weights
  apply_at_layers: [14, 28]     # tap layers 14 and 28 of 42
```

**Why this matters for you:** Your encoder is 42 layers deep, pretrained, with a low LR of 2e-5. Gradients from the TDT loss have to backprop through the joint network, then through all 42 layers. InterCTC gives direct supervision at layers 14 and 28, helping the lower layers adapt to Indic languages faster without relying on gradients surviving the full depth.

This is strictly better than the final-layer CTC you currently have.

#### 2. FastEmit regularization (already available, currently disabled)

```yaml
loss:
  tdt_kwargs:
    fastemit_lambda: 0.0   # ← currently off, try 0.001-0.01
```

Penalizes the model for delaying emissions. Encourages faster, lower-latency transcription. Useful if you care about streaming or real-time ASR. It's a regularization term on the TDT loss itself, not a separate head — zero additional params.

#### 3. Multi-Blank RNNT (alternative to TDT)

NeMo supports `multiblank_rnnt` as a loss backend — similar concept to TDT (skip multiple frames) but models it as multiple blank symbols rather than duration heads. Not worth switching mid-training, but worth knowing about.

#### 4. AngularSoftmaxLoss (ArcFace) — for language discrimination

```python
# Could add as auxiliary on encoder output pooled representation
# Forces language embeddings to be more separable
AngularSoftmaxLoss(scale=20.0, margin=1.35)
```

Would require a custom pooling + classification head on the encoder output to predict language ID. This would make the encoder learn more language-discriminative features, which could help your language embedding be more effective. **Requires custom code** — not built into the hybrid model, but straightforward to add as a callback.

#### 5. Contrastive Loss (SSL-style)

Available in `nemo/collections/asr/losses/ssl_losses/`. Could be used to enforce that encoder representations of the same language/speaker cluster together. Unlikely to help at this stage of supervised training though.

---

### What I'd actually recommend

For your remaining training:

| Change | Effort | Impact |
|--------|--------|--------|
| **Replace final CTC with InterCTC at layers 14, 28** | Config change | High — better gradient flow for deep encoder |
| **Enable fastemit_lambda=0.005** | Config change | Low-medium — slightly sharper emissions |
| **Add LangID classification head** (ArcFace on pooled encoder) | ~50 lines in train_prod.py | Medium — stronger language conditioning signal |
| **Anneal CTC weight to 0** if keeping final CTC | Config change | Small — cleaner TDT gradients |

InterCTC is the clear winner — it solves the same problem as your current CTC (encoder regularization) but does it **much** better for deep encoders by injecting supervision at multiple depths rather than only at the output.


any of these loss functions focus on WER, semantics, robustness of WER/CER or performance or accuracy6 ?


Great question — and the honest answer is **no, none of them directly optimize WER/CER**. Here's why, and what actually can:

### Why none of the NeMo losses optimize WER directly

Every loss in NeMo (CTC, RNNT, TDT, InterCTC, contrastive) optimizes a **surrogate** — log-probability of the correct token sequence given audio. WER/CER are:
- **Non-differentiable** — they're based on edit distance (insertions, deletions, substitutions), which is a discrete counting operation
- **Not decomposable per-frame** — a single substitution error and a 10-word hallucination both backprop the same way through cross-entropy

So the training loop minimizes `P(correct_tokens | audio)` and *hopes* that lower token-level loss correlates with lower WER. It usually does, but not always — you can have low loss with high WER if the model is confidently wrong on specific patterns (code-switching, rare words, homophones).

### What actually targets WER/CER/accuracy directly

#### 1. Minimum Bayes Risk (MBR) / Minimum Word Error Rate (MWER) Training

This is the real answer to your question. Instead of maximizing log-likelihood, you:
- Generate N-best hypotheses from the model (beam search)
- Compute actual WER for each hypothesis against the reference
- Use WER as the loss signal via REINFORCE / policy gradient

```
L_MWER = Σ P(hyp_i | audio) × WER(hyp_i, reference)
```

**NeMo doesn't have this built-in**, but it's implementable:

```python
# Sketch — would go in a custom training step
with torch.no_grad():
    hypotheses = model.transcribe(audio, num_beams=8, return_hypotheses=True)

rewards = []
for hyp in hypotheses:
    wer = compute_wer(hyp.text, reference)  # actual edit distance
    rewards.append(-wer)  # negative because we minimize

# REINFORCE gradient
loss = -sum(hyp.score * reward for hyp, reward in zip(hypotheses, rewards))
```

**This is the only method that directly optimizes WER.** It's typically used as a fine-tuning stage after convergence with standard losses. Google, Meta, and NVIDIA all use MWER in production ASR.

#### 2. Sequence-Level Knowledge Distillation

Use a stronger teacher model (or ensemble) to generate soft targets:
- Teacher produces token probability distributions, not just hard labels
- Student trains on KL-divergence from teacher's distribution
- Captures "semantic" similarity between likely alternatives (e.g., "colour" and "color" are both reasonable)

This doesn't optimize WER directly but captures richer signal than hard labels.

#### 3. Lattice-Based MMI (Maximum Mutual Information)

NeMo actually has this — `LatticeLoss` with `criterion_type='map'`, `loss_type='mmi'`:

```yaml
loss:
  loss_type: mmi
  criterion_type: map
```

MMI maximizes the **margin** between the correct transcription and competing hypotheses in a lattice. It's the closest thing to WER optimization that's differentiable — it penalizes confusable alternatives more heavily. Requires k2 library.

#### 4. Focal Loss / Label Smoothing (robustness, not WER directly)

- **Focal loss**: Downweights easy tokens, focuses gradient on hard/confusable ones — the tokens that actually cause WER errors
- **Label smoothing**: Prevents overconfidence, improves generalization on unseen accents/noise

Neither is built into NeMo's RNNT/TDT path but both are trivial to add.

#### 5. Contrastive / Discriminative fine-tuning on error patterns

Mine your existing WER errors from benchmarks → create hard negative pairs:
- Audio where model says "कार" but truth is "कर"
- Fine-tune with extra weight on these confusable pairs

This is targeted WER reduction, not a loss function per se.

---

### What I'd recommend for your Stage 1 model

You're at step 260K with val_loss=164.48. The path to lower WER from here:

| Stage | Method | What it targets |
|-------|--------|-----------------|
| **Now** | InterCTC at layers 14, 28 | Better encoder features → indirect WER improvement |
| **After convergence (~300K)** | **MWER fine-tuning** (10-20K steps) | **Directly optimizes WER** |
| **Targeted fixes** | Error mining → weighted retraining on confusable pairs | Specific WER patterns (substitutions) |
| **If you have a teacher** | Sequence-level KD from ensemble | Robustness, especially on low-resource langs |

**MWER fine-tuning after convergence is the highest-impact thing you can do.** It typically drops WER by 5-15% relative on top of a converged cross-entropy model. It's expensive (need beam search in the training loop) but it's the only method that directly sees and optimizes the metric you care about.

Want me to implement the MWER training step for your `train_prod.py`?