Thoroughly explore /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ to understand:

1. The full directory structure (all files, configs, scripts, checkpoints)
2. All Python training scripts - especially anything related to model architecture, language conditioning, added layers
3. Any YAML/JSON config files that define the model architecture
4. Any checkpoint files (.ckpt, .pt, .bin) available locally
5. Any README or documentation

Focus especially on:
- How language is used as an input to the model
- Any custom layers added on top of a base model (like parakeet or fastconformer)
- The hybrid TDT architecture modifications
- Any Gemma-related components

Be very thorough - read all Python files, config files, and scripts. Report file paths and key code sections.


I'll thoroughly explore the directory structure and analyze all relevant files. Let me start with a comprehensive scan.


Bash

find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.py" -o -name "*.yaml" -o -name "*.json" -o -name "*.ckpt" -o -name "*.pt" -o -name "*.bin" -o -name "README*" -o -name "*.md" | head -100


Let me continue exploring and get the full directory structure:


ls -la /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ 2>&1 | head -40


total 148
drwxrwxr-x 11 ubuntu ubuntu  4096 Mar 27 05:26 .
drwxrwxr-x 19 ubuntu ubuntu  4096 Mar 31 16:48 ..
-rw-rw-r--  1 ubuntu ubuntu  4939 Mar 27 05:10 analyze_tokenizer.py
drwxrwxr-x  4 ubuntu ubuntu  4096 Mar 25 02:16 configs
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 25 02:16 docs
-rw-rw-r--  1 ubuntu ubuntu   863 Mar 25 02:16 .env
-rw-rw-r--  1 ubuntu ubuntu 44503 Mar 25 02:16 finalsftdata.json
drwxrwxr-x  8 ubuntu ubuntu  4096 Mar 27 05:25 .git
-rw-rw-r--  1 ubuntu ubuntu   482 Mar 25 02:16 .gitignore
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 25 02:16 lf_asr
-rw-rw-r--  1 ubuntu ubuntu  1305 Mar 25 02:16 Makefile
-rw-rw-r--  1 ubuntu ubuntu 22260 Mar 25 02:16 PLAN.md
-rw-rw-r--  1 ubuntu ubuntu  1043 Mar 25 02:16 pyproject.toml
-rw-rw-r--  1 ubuntu ubuntu 10872 Mar 25 21:14 README.md
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 27 05:10 scripts
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 25 02:16 src
drwxrwxr-x  4 ubuntu ubuntu  4096 Mar 25 02:16 tests
drwxrwxr-x  3 ubuntu ubuntu  4096 Mar 27 05:26 tokenizers
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 25 02:16 tools


Now let me read the key documentation and configuration files:


1	# Maya ASR
2
3	Multilingual Indic automatic speech recognition system built on NVIDIA NeMo. Uses a FastConformer Hybrid TDT+CTC architecture with a Parakeet TDT 1.1B pretrained encoder, fine-tuned on ~150K hours of speech across 12 languages (11 Indic + English).
4
5	## Architecture
6
7	Maya ASR is based on NeMo's `EncDecHybridRNNTCTCBPEModel` — a FastConformer encoder with dual decoding heads (TDT transducer + auxiliary CTC).
8
9	```
10	Audio (16 kHz) → Mel Spectrogram (80-dim) → FastConformer Encoder (42 layers, 1024-dim)
11	                                                      │
12	                                          ┌───────────┴───────────┐
13	                                          ▼                       ▼
14	                                   TDT Decoder+Joint         CTC Head
15	                                   (primary loss)          (auxiliary loss)
16	```
17
18	### Model Components
19
20	| Component | Architecture | Parameters |
21	|---|---|---|
22	| **Encoder** | 42-layer FastConformer (Parakeet TDT 1.1B pretrained) | ~1,062M |
23	| **TDT Decoder** | 2-layer RNN (pred_hidden=640) | ~7M |
24	| **Joint Network** | Joint hidden=640, 5 TDT duration heads | ~42M |
25	| **CTC Head** | Linear projection over encoder output | ~33M |
26	| **Language Embedding** | 12 languages → 1024-dim additive bias | ~12K |
27	| **Total** | | ~1,144M |
28
29	### Tokenizer
30
31	- 32K BPE (SentencePiece), trained on quality-filtered text from all 12 languages
32	- Located at `tokenizers/stage1_prod_bpe/`
33
34	### Supported Languages
35
36	Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en).
37
38	## Training
39
40	Training is split into two stages:
41
42	### Stage 1: Acoustic Model with TDT (current)
43
44	Fine-tunes the Parakeet TDT 1.1B encoder on multilingual Indic data with a randomly initialized decoder, joint network, and CTC head sized for the 32K multilingual vocabulary.
45
46	**Key optimizations:**
47
48	- **TDT loss** — matches the pretrained encoder's original objective (durations `[0,1,2,3,4]`, sigma=0.02, omega=0.1). TDT allows the model to skip multiple frames per token emission, improving both speed and accuracy over standard RNNT.
49	- **Differential learning rate** — encoder at 2e-5 (preserves pretrained features), decoder/joint/CTC heads at 1e-3 (fast adaptation for new vocabulary).
50	- **Encoder freeze** — encoder is frozen for the first 5K optimizer steps so randomly initialized heads can warm up without corrupting pretrained weights.
51	- **CTC warmup** — CTC loss weight ramps from 0 to 0.3 over the first 3K steps, preventing random CTC gradients from interfering with early encoder adaptation.
52	- **Language embedding** — a learned per-language additive bias (12 languages to 1024-dim) is injected into the encoder output. Initialized near-zero (std=0.02) so it does not disrupt pretrained features.
53	- **Fused joint** — `fuse_loss_wer: true` with `fused_batch_size: 4` computes the joint tensor in sub-batches to cap peak VRAM.
54	- **Triple-cap OOM protection** — each micro-batch is constrained by three independent limits: max audio duration (120s), max sample count (per-bucket caps), and max total tokens (400). The token cap is critical because the TDT joint tensor size is `B x T_enc x T_dec x (V + 1 + durations)` and `T_dec` depends on transcript length, not audio length.
55
56	**Training schedule:**
57
58	| Phase | Steps | What happens |
59	|---|---|---|
60	| 0 – 3K | 0–3,000 | Encoder frozen, CTC weight ramps 0→0.3 |
61	| 3K – 5K | 3,000–5,000 | Encoder still frozen, CTC at full weight |
62	| 5K+ | 5,000–200,000 | Encoder unfrozen at 2e-5 LR, full TDT+CTC training |
63
64	**Optimizer:** AdamW (betas=0.9/0.98, weight_decay=1e-4) with CosineAnnealing schedule (2K warmup steps, min_lr=1e-5).
65
66	### Stage 2: Gemma LLM Decoder (planned)
67
68	Replaces the RNN decoder with a Gemma LLM for improved language modeling, particularly for low-resource Indic languages.
69
70	## Running Training
71
72	### Prerequisites
73
74	- 8x NVIDIA H200 GPUs (140 GB HBM3 each)
75	- Python 3.10+
76	- NeMo toolkit with ASR dependencies
77
78	### Setup
79
80	```bash
81	# Install package and dev dependencies
82	make setup
83
84	# For training dependencies (NeMo, transformers, deepspeed)
85	pip install -e ".[train]"
86	```
87
88	### Production Training (8x H200)
89
90	```bash
91	# 1. Build production data artifacts (manifests, splits, tokenizer, input config)
92	make prod-artifacts
93
94	# 2. Run full readiness check (disk, GPUs, config, artifacts, R2 storage)
95	make prod-readiness
96
97	# 3. Launch training (runs readiness gate, starts checkpoint watcher, then trains)
98	make prod-launch
99	```
100
101	Or run the training script directly:
102
103	```bash
104	python3 scripts/train_prod.py \
105	  --config configs/train/stage1_prod_8xh200.yaml \
106	  --train-parquet artifacts/phase3/production_train_final.parquet \
107	  --val-manifest data/manifests/stage1_prod_val_v2.jsonl \
108	  --pretrained-encoder pretrained/parakeet_tdt_1.1b_encoder.pt
109	```
110
111	### Smoke Test (single GPU)
112
113	```bash
114	python3 scripts/train_prod.py \
115	  --config configs/train/stage1_prod_8xh200.yaml \
116	  --max-steps 100 --devices 1 --log-every 10 --val-every 50 --smoke
117	```
118
119	### Key Training Arguments
120
121	| Argument | Default | Description |
122	|---|---|---|
123	| `--encoder-lr` | 2e-5 | Learning rate for pretrained encoder |
124	| `--head-lr` | 1e-3 | Learning rate for decoder/joint/CTC |
125	| `--freeze-encoder-steps` | 5000 | Steps to freeze encoder |
126	| `--ctc-warmup-steps` | 3000 | Steps to ramp CTC loss weight |
127	| `--max-batch-dur` | 120.0 | Max audio seconds per micro-batch |
128	| `--max-batch-size` | 16 | Hard cap on samples per micro-batch |
129	| `--max-tokens-in-batch` | 400 | Max total tokens per batch (joint tensor guard) |
130	| `--temperature` | 0.3 | Language rebalancing temperature (0=uniform, 1=natural) |
131	| `--grad-accum` | 4 | Gradient accumulation steps |
132	| `--no-lang-embed` | false | Disable language embedding |
133
134	## Inference
135
136	Maya ASR models are standard NeMo ASR models. After training produces a `.nemo` checkpoint, use NeMo's transcribe API:
137
138	```python
139	import nemo.collections.asr as nemo_asr
140
141	# Load the trained model
142	model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from("path/to/maya_asr.nemo")
143	model.eval()
144	model.cuda()
145
146	# Transcribe audio files
147	transcriptions = model.transcribe(["audio1.wav", "audio2.wav"])
148	print(transcriptions)
149	```
150
151	For TDT-specific decoding (using duration-aware greedy search):
152
153	```python
154	from omegaconf import OmegaConf
155
156	# The model defaults to TDT greedy decoding as configured in training
157	# To adjust decoding parameters:
158	decoding_cfg = OmegaConf.create({
159	    "strategy": "greedy",
160	    "model_type": "tdt",
161	    "durations": [0, 1, 2, 3, 4],
162	    "greedy": {"max_symbols": 10},
163	})
164	model.change_decoding_strategy(decoding_cfg)
165	transcriptions = model.transcribe(["audio.wav"])
166	```
167
168	To use the CTC head instead of TDT (faster but slightly less accurate):
169
170	```python
171	transcriptions = model.transcribe(["audio.wav"], return_hypotheses=False, verbose=False)
172	# Or explicitly switch to CTC decoding:
173	model.cur_decoder = "ctc"
174	transcriptions = model.transcribe(["audio.wav"])
175	```
176
177	## Data Pipeline
178
179	The training data flows through a multi-stage pipeline:
180
181	```
182	Raw audio tars (per-language shards on disk)
183	  │
184	  ▼
185	build_manifest.py → Raw manifest JSONL (path, duration, transcript, language)
186	  │
187	  ▼
188	split_manifest.py → Train/Val splits (deterministic, stratified)
189	  │
190	  ▼
191	build_tokenizer.py → 32K BPE tokenizer (SentencePiece, quality-filtered)
192	  │
193	  ▼
194	Phase 3 Production Parquet (tar_path, tar_offset_data, tar_nbytes, transcript, duration_s, language)
195	  │
196	  ▼
197	TarOffsetReader (os.pread, zero-extraction from existing tars)
198	  │
199	  ▼
200	ProductionBatchSampler (duration bucketing + language rebalancing + triple-cap OOM protection)
201	  │
202	  ▼
203	NeMo training loop
204	```
205
206	**Key data features:**
207
208	- **Tar-offset loading** — reads audio directly from tar files using byte offsets (`os.pread`), avoiding tar extraction. Zero disk overhead.
209	- **Duration bucketing** — groups similar-length clips into buckets `[0-4, 4-8, 8-12, 12-16, 16-20]s` to minimize padding waste.
210	- **Language temperature rebalancing** (T=0.3) — upweights low-resource languages so they are seen more frequently during training.
211	- **Compute-cost DDP sharding** — distributes batches across GPUs balanced by compute cost, not just sample count.
212	- **Token-length-aware batching** — pre-computes token counts per sample and caps total tokens per batch at 400, directly controlling the dominant VRAM consumer (joint tensor `T_dec` dimension).
213
214	## Hardware Requirements
215
216	### Production Training (Stage 1)
217
218	- **GPUs:** 8x NVIDIA H200 (140 GB HBM3 each)
219	- **Precision:** BF16 mixed
220	- **VRAM breakdown (per GPU):**
221	  - Model weights (bf16): ~4.5 GB
222	  - Optimizer states (AdamW, fp32 copies): ~13.5 GB
223	  - Activations + joint tensor: ~75 GB budget (75% safety margin)
224	- **Effective batch:** 8 GPUs x micro-batch x 4 gradient accumulation steps
225	- **Training target:** 200K optimizer steps, ~5-7 days
226	- **Checkpoints:** saved every 5K steps, uploaded to R2 object storage, validated (size match) before local deletion
227
228	### Inference
229
230	- **Minimum:** 1x GPU with 8+ GB VRAM (model is ~2.3 GB in bf16)
231	- **Recommended:** 1x GPU with 16+ GB VRAM for comfortable batch inference
232
233	## Project Structure
234
235	```
236	configs/
237	  train/              # Training YAML configs
238	  data/               # Data pipeline configs
239	scripts/
240	  train_prod.py       # Production training script (TDT + DiffLR + lang embed)
241	  build_manifest.py   # Raw manifest builder
242	  split_manifest.py   # Train/val splitter
243	  build_tokenizer.py  # BPE tokenizer builder
244	  run_stage1_prod.sh  # Production launcher (readiness + watcher + train)
245	lf_asr/
246	  data/
247	    production_sampler.py  # Triple-cap batch sampler with language rebalancing
248	    tar_offset_reader.py   # Zero-extraction tar audio reader
249	src/maya_asr/         # Core library
250	tokenizers/           # Trained BPE tokenizer artifacts
251	pretrained/           # Pretrained encoder weights
252	artifacts/            # Data processing outputs (parquet)
253	experiments/          # Training outputs (gitignored)
254	tests/
255	  unit/               # Unit tests (CI-safe, no data dependency)
256	  integration/        # Integration tests (requires data)
257	```
258
259	## Make Targets
260
261	| Target | Description |
262	|---|---|
263	| `make setup` | Install package + dev deps |
264	| `make test` | Run unit tests only |
265	| `make test-all` | Run unit + integration tests |
266	| `make lint` | Ruff lint check |
267	| `make format` | Auto-format + fix |
268	| `make smoke` | Import check + unit suite |
269	| `make prod-artifacts` | Build production data artifacts |
270	| `make prod-readiness` | Production launch gate (disk, GPUs, config, artifacts, R2) |
271	| `make prod-launch` | Fail-fast production launcher (readiness + watcher + training) |
272	| `make validate-r2` | End-to-end R2 upload/download/integrity test |
273
274	## License
275
276	MIT
277


1	# Maya ASR: FastConformer + Gemma Transcription Build Guide
2
3	## Executive Summary
4
5	Build a speech-to-text transcription system using a FastConformer encoder coupled with a small Gemma 3 decoder, trained at full-parameter scale on ~150K hours of transcribed Indic+English data across 12 languages. **Task: transcription only** (no translation). Covers architecture design, tokenizer construction, multi-stage training recipes, distributed training on 8×H100 GPUs, and concrete NeMo scripts and config references.
6
7	---
8
9	## 1. Architecture Overview
10
11	### 1.1 High-Level Design
12
13	```
14	┌──────────────────────────────────────────────────────────────────┐
15	│                    Maya ASR (~800M–1.2B)                         │
16	│                                                                  │
17	│  ┌───────────────────────┐    ┌──────────┐    ┌──────────────┐   │
18	│  │  FastConformer Encoder │───►│Projection│───►│ Gemma 3      │   │
19	│  │  (IndicConformer-600M)│    │  Layer   │    │ (small)      │   │
20	│  │  32 layers, 1024 dim  │    │ Linear   │    │ Decoder      │   │
21	│  │  8x subsampling       │    │ MLP      │    │              │   │
22	│  └───────────────────────┘    └──────────┘    └──────────────┘   │
23	│                                                                  │
24	│  Language Prompting: <|source_lang|> token guides decoder        │
25	│  Task: <|transcribe|> only                                       │
26	└──────────────────────────────────────────────────────────────────┘
27	```
28
29	### 1.2 Component Breakdown
30
31	**Encoder: FastConformer (IndicConformer-600M base)**
32
33	Based on [AI4Bharat's IndicConformer-600M multilingual model](https://huggingface.co/ai4bharat/indic-conformer-600m-multilingual), a Conformer-based hybrid CTC+RNNT model supporting all 22 scheduled Indian languages.
34
35	- Architecture: FastConformer with 8x depthwise convolutional subsampling (256 channels)
36	- Reduced convolutional kernel size of 9 in conformer blocks
37	- 2.4x faster training/inference than standard Conformer without quality degradation
38	- Supports local attention for audio >1 hour
39	- NeMo-native: `EncDecHybridRNNTCTCBPEModel` class
40
41	Scale to **32 encoder layers at 1024 hidden dimension** (matching Canary-1B-Flash's encoder spec). Initialize from IndicConformer-600M checkpoint and expand layers if needed.
42
43	Reference architecture table:
44
45	| Model | Params | encoder.n_layers | enc_hidden |
46	|---|---|---|---|
47	| Canary-1B-Flash | 883M | 32 | 1024 |
48	| Canary-180M-Flash | 182M | 17 | 512 |
49	| IndicConformer-600M | 600M | — | — |
50
51	**Recommended encoder config**: 32 layers, 1024 hidden (Canary-1B-Flash pattern), initialized from IndicConformer-600M weights.
52
53	**Decoder: Gemma 3 (Small Variant)**
54
55	Use the smallest viable Gemma 3 variant as a decoder-only LLM. Candidates:
56
57	| Variant | Params | Hidden Dim | Layers | Notes |
58	|---|---|---|---|---|
59	| `google/gemma-3-1b-pt` | 1B | 1536 | 26 | Smallest official Gemma 3 |
60	| `google/gemma-2-2b` | 2.6B | 2304 | 26 | Gemma 2, good Indic but larger |
61
62	**Primary choice: Gemma 3 1B** — it is already the smallest Gemma 3 release. Its 262K SentencePiece tokenizer has excellent Indic coverage:
63
64	| Language | Gemma 3 Fertility | Dedicated Tokens |
65	|---|---|---|
66	| Hindi | 1.38 | 4,391 Devanagari tokens |
67	| Bengali | 1.74 | 541 tokens |
68	| Tamil | 2.42 | 724 tokens |
69	| Telugu | 2.93 | 404 tokens |
70	| Kannada | 3.15 | 231 tokens |
71	| Malayalam | 3.39 | ~200 tokens (est.) |
72
73	With 150K hours, full instruction tuning (not LoRA) is the right approach — the decoder needs to deeply learn acoustic-to-text alignment for all 12 languages.
74
75	If Gemma 3 1B proves too heavy, consider distilling down or using only a subset of its layers (e.g., first 16 of 26 layers) as a smaller decoder, keeping the tokenizer.
76
77	**Projection Layer: Linear MLP**
78
79	Linear projection outperforms Q-Former for speech-LLM coupling ([MLC-SLM 2025](https://arxiv.org/html/2601.01461v1): Linear 11.05% WER vs Q-Former 11.51%).
80
81	Design: `Conv1D(stride=2) → LayerNorm → Linear(1024 → 1536) → SiLU → Linear(1536 → 1536)`
82
83	This downsamples encoder output by 2x temporally (on top of 8x from FastConformer), then maps from encoder dimension (1024) to Gemma 3 1B's hidden_size (1536).
84
85	---
86
87	## 2. Local Data
88
89	### 2.1 Data Location
90
91	All data is pre-downloaded at `/root/sft_data/` (~13 TiB).
92
93	### 2.2 Data Sources
94
95	| Source | Path | Languages | Description |
96	|---|---|---|---|
97	| **final-export** | `/root/sft_data/final-export/production/shards/` | 12 (as, bn, en, gu, hi, kn, ml, mr, or, pa, ta, te) | Primary dataset, largest |
98	| **indicvoices** | `/root/sft_data/indicvoices/` | 11 Indic | AI4Bharat public benchmark |
99	| **indicvoices-r** | `/root/sft_data/indicvoices-r/` | 11 Indic | High-quality Indic speech |
100	| **josh** | `/root/sft_data/josh/` | 5 (bn, en, gu, hi, mr) | Indic content |
101	| **joshdelivery** | `/root/sft_data/joshdelivery/` | 5 | Additional Indic data |
102	| **globe** | `/root/sft_data/globe/` | Multi | Multilingual data |
103	| **ears** | `/root/sft_data/ears/` | en | English |
104	| **expresso** | `/root/sft_data/expresso/` | en | English |
105	| **librittsr** | `/root/sft_data/librittsr/` | en | English TTS-quality |
106	| **ljspeech** | `/root/sft_data/ljspeech/` | en | English speech |
107	| **vctk** | `/root/sft_data/vctk/` | en | Multi-speaker English |
108
109	### 2.3 Shard Structure (final-export)
110
111	Organized as `lang={code}/{shard_id}/` with each shard containing:
112	- `audio.tar` — compressed FLAC audio files (~15,000 segments per shard)
113	- `audio_index.parquet` — audio indexing metadata
114	- `manifest.json` — shard-level metadata (counts, checksums, language)
115	- `metadata.parquet` — per-segment metadata (transcription, duration, speaker, domain)
116	- `xcodec2_tokens.parquet` — pre-computed acoustic tokens (optional)
117
118	**Shard counts per language (final-export):**
119
120	| Language | Shards | Est. Hours |
121	|---|---|---|
122	| English (en) | 840 | ~28K |
123	| Hindi (hi) | 475 | ~16K |
124	| Telugu (te) | 400 | ~13K |
125	| Malayalam (ml) | 372 | ~12K |
126	| Punjabi (pa) | 343 | ~11K |
127	| Tamil (ta) | 312 | ~10K |
128	| Kannada (kn) | 200 | ~7K |
129	| Gujarati (gu) | 186 | ~6K |
130	| Bengali (bn) | 158 | ~5K |
131	| Marathi (mr) | 148 | ~5K |
132	| Odia (or) | 78 | ~3K |
133	| Assamese (as) | 39 | ~1K |
134	| **Total** | **3,551** | **~117K** |
135
136	### 2.4 Data Format Conversion
137
138	The local data needs to be converted to NeMo manifest format for training:
139
140	```json
141	{"audio_filepath": "/data/hindi/audio_001.wav", "text": "नमस्ते दुनिया", "duration": 5.2, "lang": "hi", "taskname": "asr", "source_lang": "hi", "target_lang": "hi"}
142	```
143
144	**Step 1: Extract per-segment manifests from metadata parquet files**
145
146	Write a script to iterate all shards, read `metadata.parquet`, and produce NeMo-compatible JSON-line manifests pointing into the `audio.tar` files (or extracted audio).
147
148	**Step 2: Convert to tarred WebDataset format** (recommended for 150K hours):
149
150	```bash
151	python scripts/speech_recognition/convert_to_tarred_audio_dataset.py \
152	  --manifest_path=train_manifest.json \
153	  --target_dir=/data/tarred/hindi/ \
154	  --num_shards=512 \
155	  --max_duration=40.0 \
156	  --min_duration=0.01
157	```
158
159	Since audio is already tarred per-shard, we may be able to adapt existing tars directly or write a custom `LhotseDataset` that reads from the existing shard format.
160
161	### 2.5 Multi-Corpus Input Configuration
162
163	Use NeMo's `input_cfg` YAML format:
164
165	```yaml
166	input_cfg:
167	  - corpus: final_export_hindi
168	    language: hi
169	    manifest_filepath: /data/nemo_manifests/hi/manifest.json
170	    tarred_audio_filepaths: /data/nemo_tarred/hi/audio_*.tar
171	    type: nemo_tarred
172	    weight: 10000
173	  - corpus: final_export_bengali
174	    language: bn
175	    manifest_filepath: /data/nemo_manifests/bn/manifest.json
176	    tarred_audio_filepaths: /data/nemo_tarred/bn/audio_*.tar
177	    type: nemo_tarred
178	    weight: 10000
179	  # ... repeat for all 12 languages + supplementary sources
180	```
181
182	### 2.6 Language Balancing
183
184	Use **two-stage deterministic upsampling** ([Polyglot-Lion, March 2026](https://arxiv.org/html/2603.16184v1)):
185
186	1. **Intra-language**: Upsample smaller corpora within each language to match the largest corpus
187	2. **Inter-language**: Upsample each language to match the largest language group
188
189	This produced 72% relative WER reduction for underrepresented languages without degrading high-resource ones.
190
191	Alternative: Temperature-based sampling (α = 0.3–0.5) as used by [Canary-1B-v2](https://arxiv.org/html/2509.14128v2).
192
193	### 2.7 Data Augmentation
194
195	At 150K hours, heavy augmentation is less critical but still beneficial:
196
197	- **Speed perturbation** (0.9x, 1.0x, 1.1x): Standard 3x multiplier
198	- **SpecAugment**: Frequency masking (F=27, num_masks=2) + time masking (T=0.05×duration, num_masks=10)
199	- **Noise augmentation**: Environmental noise at SNRs 5–20 dB (MUSAN or Indic-environment noise)
200
201	---
202
203	## 3. Tokenizer Design
204
205	### 3.1 Decoder Tokenizer: Gemma 3's Tokenizer
206
207	Use Gemma 3's existing 262K SentencePiece tokenizer directly. It already has strong Indic coverage with 4K+ dedicated Indic tokens. No custom tokenizer training needed for the decoder side.
208
209	Add special tokens:
210	- `<|transcribe|>` — task token
211	- Language ID tokens: `<|hi|>`, `<|bn|>`, `<|ta|>`, `<|te|>`, `<|gu|>`, `<|kn|>`, `<|ml|>`, `<|mr|>`, `<|pa|>`, `<|or|>`, `<|as|>`, `<|en|>`
212
213	### 3.2 Encoder-Side Tokenizer (for CTC Auxiliary Loss)
214
215	Train a smaller custom Indic BPE tokenizer (4K–8K vocab) for the CTC auxiliary loss on the encoder:
216
217	```bash
218	python scripts/tokenizers/process_asr_text_tokenizer.py \
219	  --manifest=balanced_text_corpus.json \
220	  --vocab_size=8192 \
221	  --tokenizer_type=bpe \
222	  --spe_type=unigram \
223	  --spe_character_coverage=0.9999 \
224	  --spe_split_by_unicode_script=true \
225	  --output_dir=/tokenizers/encoder_8k/
226	```
227
228	---
229
230	## 4. Training Recipe
231
232	### 4.1 Stage 0: Encoder Pretraining / Warm-Start
233
234	Start from [IndicConformer-600M multilingual](https://huggingface.co/ai4bharat/indic-conformer-600m-multilingual):
235
236	```python
237	import nemo.collections.asr as nemo_asr
238
239	model = nemo_asr.models.ASRModel.from_pretrained(
240	    "ai4bharat/indic-conformer-600m-multilingual"
241	)
242	```
243
244	If scaling encoder to 32 layers, initialize new layers from existing weights via interpolation or duplication, then continue pretraining with CTC+RNNT loss for 50K–100K steps.
245
246	Alternative: Initialize from [Canary-1B-Flash encoder](https://huggingface.co/nvidia/canary-1b-flash) (32 layers, 1024 dim):
247
248	```yaml
249	init_from_pretrained_model:
250	  model0:
251	    name: "nvidia/canary-1b-flash"
252	    include: ["encoder"]
253	    exclude: ["encoder.pre_encode.out"]
254	```
255
256	### 4.2 Stage 1: Encoder Fine-Tuning on Indic Data (CTC+RNNT)
257
258	Fine-tune the encoder on all ~150K hours using hybrid CTC+RNNT loss. This creates a strong Indic speech feature extractor before coupling with the LLM decoder.
259
260	```bash
261	python examples/asr/asr_hybrid_transducer_ctc/speech_to_text_hybrid_rnnt_ctc_bpe_prompt.py \
262	  --config-path=examples/asr/conf/fastconformer/hybrid_transducer_ctc/ \
263	  --config-name=fastconformer_hybrid_transducer_ctc_bpe_prompt.yaml \
264	  model.train_ds.manifest_filepath=/data/train_manifest.json \
265	  model.validation_ds.manifest_filepath=/data/val_manifest.json \
266	  model.tokenizer.dir=/tokenizers/encoder_8k/ \
267	  model.train_ds.use_lhotse=true \
268	  model.train_ds.input_cfg=/configs/input_cfg.yaml \
269	  trainer.devices=8 \
270	  trainer.num_nodes=1 \
271	  trainer.max_steps=200000 \
272	  trainer.val_check_interval=5000 \
273	  model.optim.name=adamw \
274	  model.optim.lr=0.001 \
275	  model.optim.betas=[0.9,0.98] \
276	  model.optim.weight_decay=0.0001 \
277	  model.optim.sched.name=CosineAnnealing \
278	  model.optim.sched.warmup_steps=5000 \
279	  trainer.precision=bf16-mixed
280	```
281
282	**Duration**: ~5–10 days on 8×H100 for 200K steps.
283
284	Key insight from MLC-SLM research: "fine-tuning speech encoders on in-domain data before integrating them into the LLM is more effective than the tri-stage pipeline" ([arXiv](https://arxiv.org/html/2601.01461v1)).
285
286	### 4.3 Stage 2: Attach Gemma Decoder + Joint Training
287
288	Two-phase approach:
289
290	**Stage 2a: Train projection only (encoder + LLM frozen), ~20K steps**
291
292	```python
293	import torch
294	from transformers import GemmaForCausalLM, GemmaTokenizer
295	import nemo.collections.asr as nemo_asr
296
297	# Load pretrained components
298	encoder = load_indicconformer_encoder("stage1_checkpoint.nemo")
299	llm = GemmaForCausalLM.from_pretrained("google/gemma-3-1b-pt")
300	tokenizer = GemmaTokenizer.from_pretrained("google/gemma-3-1b-pt")
301
302	# Define projection layer
303	projection = nn.Sequential(
304	    nn.Conv1d(1024, 1024, kernel_size=2, stride=2),  # 2x temporal downsampling
305	    nn.LayerNorm(1024),
306	    nn.Linear(1024, 1536),  # enc_dim → gemma-3-1b hidden_dim
307	    nn.SiLU(),
308	    nn.Linear(1536, 1536),
309	)
310
311	# Stage 2a: projection only
312	for batch in dataloader:
313	    audio_features = encoder(batch.audio)       # frozen
314	    projected = projection(audio_features)       # trainable
315
316	    # Prepend language + task tokens: <|hi|><|transcribe|>
317	    prompt_embeds = llm.embed_tokens(batch.prompt_ids)
318	    input_embeds = torch.cat([projected, prompt_embeds], dim=1)
319
320	    logits = llm(inputs_embeds=input_embeds)     # frozen
321	    loss = cross_entropy(logits, batch.target_ids)
322	    loss.backward()
323	```
324
325	**Stage 2b: Joint training (projection + full LLM, encoder frozen or LoRA), ~100K–200K steps**
326
327	- LR: 1e-5 for LLM, 1e-4 for projection
328	- Use DeepSpeed ZeRO-2 or FSDP for memory efficiency
329	- Cross-entropy loss on transcription tokens only
330
331	**SKIP-SALSA reference** ([Interspeech 2025](https://www.isca-archive.org/interspeech_2025/mittal25_interspeech.pdf)):
332
333	Designed for scenarios where the LLM has much better tokenization than the ASR decoder (exactly our case). Showed up to 20% WER improvement on IndicVoices data. At 150K hour scale, full parameter training is preferred over SKIP-SALSA's frozen-encoder approach, but the projection design is valuable inspiration.
334
335	Code reference: [github.com/csalt-research/salsa](https://github.com/csalt-research/salsa)
336
337	### 4.4 Stage 3: Quality Fine-Tuning
338
339	Fine-tune on a **language-balanced, high-quality subset**:
340
341	- Select cleanest 10% of data per language
342	- Train for 10K–20K additional steps
343	- Lower learning rate (0.1× of Stage 2 peak)
344	- This stage polishes accuracy and reduces hallucination
345
346	---
347
348	## 5. Distributed Training Configuration
349
350	### 5.1 Hardware: 8×H100 80GB
351
352	For a ~800M–1.2B parameter model (600M encoder + projection + Gemma decoder):
353	- **Single H100**: 80GB VRAM, sufficient for inference but tight for full training
354	- **8×H100**: DDP as default for encoder-only stages. DeepSpeed ZeRO-2 for joint training with Gemma
355
356	### 5.2 NeMo Trainer Config (Encoder Stages)
357
358	```yaml
359	trainer:
360	  devices: 8
361	  num_nodes: 1
362	  accelerator: gpu
363	  strategy: ddp
364	  precision: bf16-mixed
365	  max_steps: 200000
366	  val_check_interval: 5000
367	  accumulate_grad_batches: 4
368	  gradient_clip_val: 1.0
369	  log_every_n_steps: 100
370	```
371
372	### 5.3 DeepSpeed ZeRO-2 (Joint Training with Gemma)
373
374	```json
375	{
376	  "zero_optimization": {
377	    "stage": 2,
378	    "offload_optimizer": {"device": "none"},
379	    "allgather_partitions": true,
380	    "allgather_bucket_size": 2e8,
381	    "reduce_scatter": true,
382	    "reduce_bucket_size": 2e8,
383	    "overlap_comm": true
384	  },
385	  "bf16": {"enabled": true},
386	  "gradient_accumulation_steps": 4,
387	  "gradient_clipping": 1.0,
388	  "train_micro_batch_size_per_gpu": 8,
389	  "wall_clock_breakdown": false
390	}
391	```
392
393	ZeRO-2 shards optimizer + gradients (~3x memory reduction). ZeRO-3 only needed if >3B params.
394
395	### 5.4 Training Time Estimation
396
397	| Stage | Steps | Time/Step (8×H100) | Total |
398	|---|---|---|---|
399	| Stage 1: Encoder CTC+RNNT | 200K | ~2–3s | ~5–7 days |
400	| Stage 2a: Projection warmup | 20K | ~3s | ~1 day |
401	| Stage 2b: Full joint training | 150K | ~4–6s | ~7–10 days |
402	| Stage 3: Quality fine-tune | 20K | ~4s | ~1 day |
403	| **Total** | | | **~14–19 days** |
404
405	### 5.5 Training Optimizations
406
407	- **bf16 mixed precision**: Standard on H100, 2× memory reduction + TF32 tensor core acceleration
408	- **Activation checkpointing**: Essential for 32-layer encoder — rematerialize activations during backward pass
409	- **Lhotse dynamic bucketing**: Groups samples by duration, minimizing padding waste:
410	  ```yaml
411	  train_ds:
412	    use_lhotse: true
413	    use_bucketing: true
414	    num_buckets: 30
415	    bucket_buffer_size: 20000
416	    shuffle_buffer_size: 10000
417	    max_duration: 40.0
418	    min_duration: 0.01
419	  ```
420	  Use NeMo's [2D duration bin estimation script](https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speech_recognition/estimate_duration_bins_2d.py)
421	- **Tarred audio datasets**: Pre-tar into WebDataset shards (512–1024 per language). Our data is already shard-tarred.
422	- **Gradient accumulation**: accumulate_grad_batches=4–8 for large effective batch sizes
423	- **Cosine annealing scheduler**: Standard for FastConformer training
424
425	---
426
427	## 6. Evaluation
428
429	### 6.1 Benchmarks
430
431	| Benchmark | Languages | Domain | Source |
432	|---|---|---|---|
433	| [Vistaar](https://github.com/AI4Bharat/vistaar) | 12 Indic | 59 benchmarks: news, education, tourism | AI4Bharat |
434	| [IndicSUPERB](https://ai4bharat.iitm.ac.in/) | 22 Indic | Standardized clean + noisy splits | AI4Bharat |
435	| [FLEURS](https://huggingface.co/datasets/google/fleurs) | Indic subset | Read speech | Google |
436	| [CommonVoice](https://commonvoice.mozilla.org/) | Multiple Indic | Crowd-sourced read speech | Mozilla |
437	| [Kathbath](https://github.com/AI4Bharat/vistaar) | 12 Indic | Conversational | AI4Bharat |
438
439	### 6.2 Metrics
440
441	- **WER** (Word Error Rate): Primary metric for all languages
442	- **CER** (Character Error Rate): Secondary for morphologically rich languages (Tamil, Malayalam, Telugu)
443	- **RTFx** (Real-Time Factor, inverse): Throughput measurement
444	- Per-language breakdown — never average across languages without reporting individual scores
445
446	### 6.3 Baselines to Compare Against
447
448	| Model | Best Hindi WER | Notes |
449	|---|---|---|
450	| [Google STT](https://arxiv.org/html/2602.03868v1) | 16.2% (agricultural) | Commercial |
451	| [IndicWhisper](https://github.com/AI4Bharat/vistaar) | 13.6% avg (Vistaar) | Fine-tuned Whisper, 12 Indic |
452	| [IndicConformer-600M](https://huggingface.co/ai4bharat/indic-conformer-600m-multilingual) | Varies by lang | 22 Indic, NeMo-native |
453	| [Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B) | SOTA open-source | Hindi only among Indic |
454
455	---
456
457	## 7. Avoiding Catastrophic Forgetting
458
459	With 12 languages trained jointly, forgetting is less of a concern. Key strategies:
460
461	- **Joint training** (all languages simultaneously): Best approach. No forgetting by design.
462	- **Language-balanced sampling**: Critical — without it, high-resource languages dominate
463	- **If adding languages later**: Use Learning without Forgetting (LwF) — consistently low and stable WER, best BWT in noisy settings
464	- **CTC auxiliary loss**: Acts as implicit regularizer, stabilizing encoder representations across languages
465
466	Code: [github.com/FrozenWolf-Cyber/Indic-CL-ASR](https://github.com/FrozenWolf-Cyber/Indic-CL-ASR)
467
468	---
469
470	## 8. Key References and Code
471
472	### 8.1 Training Scripts
473
474	| Script | Path | Purpose |
475	|---|---|---|
476	| Hybrid CTC+RNNT Training | `examples/asr/asr_hybrid_transducer_ctc/speech_to_text_hybrid_rnnt_ctc_bpe_prompt.py` | Encoder training with prompt support |
477	| AED Multitask Training | `examples/asr/speech_multitask/speech_to_text_aed.py` | Canary-style training (reference) |
478	| CTC Language Fine-tuning | `tutorials/asr/ASR_CTC_Language_Finetuning.ipynb` | Tutorial for adding new languages |
479	| Tokenizer Training | `scripts/tokenizers/process_asr_text_tokenizer.py` | Build SentencePiece tokenizer |
480	| Duration Bin Estimation | `scripts/speech_recognition/estimate_duration_bins_2d.py` | Optimize Lhotse bucketing |
481	| SALSA/SKIP-SALSA | `github.com/csalt-research/salsa` | ASR-LLM synchronous coupling |
482
483	### 8.2 Config Files
484
485	| Config | Path | Purpose |
486	|---|---|---|
487	| FastConformer AED | `examples/asr/conf/speech_multitask/fast-conformer_aed.yaml` | Canary-style config (reference) |
488	| FastConformer CTC | `examples/asr/conf/fastconformer/fast-conformer_ctc_bpe.yaml` | CTC encoder training |
489	| FastConformer RNNT | `examples/asr/conf/fastconformer/fast-conformer_transducer_bpe.yaml` | Transducer training |
490	| Hybrid CTC+RNNT Prompt | `examples/asr/conf/fastconformer/hybrid_transducer_ctc/fastconformer_hybrid_transducer_ctc_bpe_prompt.yaml` | Hybrid with prompt support |
491	| Long-form CTC | `examples/asr/conf/fastconformer/fast-conformer-long_ctc_bpe.yaml` | Longformer attention for long audio |
492
493	### 8.3 Key Papers
494
495	- [Canary-1B-v2 & Parakeet-TDT-0.6B-v3 Technical Report](https://arxiv.org/html/2509.14128v2) — NVIDIA's training methodology
496	- [SALSA: Speedy ASR-LLM Synchronous Aggregation](https://arxiv.org/abs/2408.16542) — ASR+LLM coupling via projection layers
497	- [SKIP-SALSA](https://www.isca-archive.org/interspeech_2025/mittal25_interspeech.pdf) — Handles token fertility gap for Indic languages
498	- [Bridging the Gap: Speech-LLM vs E2E ASR](https://arxiv.org/html/2601.01461v1) — Two-stage training beats tri-stage
499	- [Polyglot-Lion](https://arxiv.org/html/2603.16184v1) — Balanced multilingual sampling
500	- [Indic CL-ASR](https://arxiv.org/abs/2508.06280) — Continual learning for Indian language ASR
501	- [Slam: Training SLMs on One GPU in a Day](https://arxiv.org/html/2502.15814v1) — Efficient speech-LLM training
502	- [Vistaar/IndicWhisper](https://arxiv.org/abs/2305.15386) — Indic training data and benchmarks
503
504	### 8.4 NeMo Installation
505
506	```bash
507	# Install AI4Bharat's NeMo fork (for IndicConformer compatibility)
508	git clone https://github.com/AI4Bharat/NeMo.git
509	cd NeMo && git checkout nemo-v2 && bash reinstall.sh
510
511	# Or install official NeMo (for Canary-style training)
512	pip install nemo_toolkit[asr]
513
514	# For Speech-LLM integration with Gemma
515	pip install transformers accelerate deepspeed
516	```
517


Now let me examine the training configurations and Python code:


1	# Stage 1 Production Config: FastConformer Hybrid TDT+CTC on full dataset
2	# Hardware: 8x NVIDIA H200 (140 GB HBM3 each)
3	# Target: ~200K steps, ~5-7 days
4	#
5	# Pretrained encoder: Parakeet TDT 1.1B (42-layer FastConformer, English)
6	# Tokenizer: encoder-side 32K BPE (quality-filtered text, 11 Indic + English)
7	# Decoder/joint: randomly initialized for 32K vocab, then finetuned
8	#
9	# Optimizations v2:
10	#   - TDT loss (matching pretrained encoder objective) instead of RNNT
11	#   - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
12	#   - Encoder frozen for first 5K steps, then unfrozen at low LR
13	#   - CTC loss warmup: 0 -> 0.3 over first 3K steps
14	#   - Language embedding conditioning (lang_id -> encoder)
15
16	name: maya-asr-stage1-prod-v2
17
18	trainer:
19	  devices: 8
20	  num_nodes: 1
21	  accelerator: gpu
22	  strategy: ddp_find_unused_parameters_true
23	  precision: bf16-mixed
24	  max_steps: 200000
25	  val_check_interval: 10000
26	  accumulate_grad_batches: 4  # effective = 8 GPUs x batch x 4 accum
27	  gradient_clip_val: 1.0
28	  log_every_n_steps: 100
29	  enable_checkpointing: false  # handled by exp_manager
30	  default_root_dir: experiments/stage1_prod
31
32	model:
33	  model_defaults:
34	    asr_enc_hidden: 1024
35	    enc_hidden: 1024
36	    dec_hidden: 640
37	    pred_hidden: 640
38
39	  preprocessor:
40	    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
41	    sample_rate: 16000
42	    normalize: per_feature
43	    window_size: 0.025
44	    window_stride: 0.01
45	    features: 80
46	    n_fft: 512
47	    dither: 0.00001
48
49	  encoder:
50	    _target_: nemo.collections.asr.modules.ConformerEncoder
51	    feat_in: 80
52	    n_layers: 42
53	    d_model: ${model.model_defaults.enc_hidden}
54	    subsampling: dw_striding
55	    subsampling_factor: 8
56	    subsampling_conv_channels: 256
57	    ff_expansion_factor: 4
58	    self_attention_model: rel_pos
59	    n_heads: 8
60	    conv_kernel_size: 9
61	    dropout: 0.1
62	    dropout_pre_encoder: 0.1
63	    dropout_emb: 0.0
64
65	  decoder:
66	    _target_: nemo.collections.asr.modules.RNNTDecoder
67	    prednet:
68	      pred_hidden: ${model.model_defaults.pred_hidden}
69	      pred_rnn_layers: 2
70	      dropout: 0.1
71
72	  joint:
73	    _target_: nemo.collections.asr.modules.RNNTJoint
74	    num_extra_outputs: 5  # TDT duration heads (durations=[0,1,2,3,4])
75	    fuse_loss_wer: true
76	    fused_batch_size: 2   # sub-batch joint computation → caps peak VRAM
77	    jointnet:
78	      joint_hidden: 640
79	      activation: relu
80	      dropout: 0.1
81	      encoder_hidden: ${model.model_defaults.enc_hidden}
82	      pred_hidden: ${model.model_defaults.pred_hidden}
83
84	  loss:
85	    loss_name: tdt
86	    tdt_kwargs:
87	      durations: [0, 1, 2, 3, 4]
88	      sigma: 0.02
89	      omega: 0.1
90	      fastemit_lambda: 0.0
91	      clamp: -1.0
92
93	  decoding:
94	    strategy: greedy
95	    model_type: tdt
96	    durations: [0, 1, 2, 3, 4]
97	    greedy:
98	      max_symbols: 10
99
100	  aux_ctc:
101	    ctc_loss_weight: 0.3  # warmed up from 0->0.3 over 3K steps in train_prod.py
102	    use_cer: false
103	    decoder:
104	      _target_: nemo.collections.asr.modules.ConvASRDecoder
105	      feat_in: null
106	      num_classes: -1
107
108	  # Encoder-side tokenizer (32K BPE, quality-filtered, 11 Indic + English)
109	  tokenizer:
110	    dir: tokenizers/stage1_prod_bpe
111	    type: bpe
112
113	  train_ds:
114	    # Production: tar-offset loading via train_prod.py (not NeMo tarred format)
115	    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl  # placeholder for model init
116	    batch_size: 4
117	    shuffle: true
118	    num_workers: 8
119	    max_duration: 20.0
120	    min_duration: 0.5
121	    sample_rate: 16000
122	    trim_silence: false
123
124	  validation_ds:
125	    manifest_filepath: data/manifests/stage1_prod_val_v2.jsonl
126	    batch_size: 4
127	    shuffle: false
128	    num_workers: 4
129	    sample_rate: 16000
130
131	  # NOTE: optim LR is overridden by differential LR in train_prod.py
132	  # This base LR is used for NeMo's scheduler init but actual param group
133	  # LRs are set in the training script.
134	  optim:
135	    name: adamw
136	    lr: 0.001  # decoder/joint/CTC LR (encoder uses 2e-5 via param groups)
137	    betas: [0.9, 0.98]
138	    weight_decay: 0.0001
139	    sched:
140	      name: CosineAnnealing
141	      warmup_steps: 2000
142	      min_lr: 0.00001
143
144	exp_manager:
145	  exp_dir: experiments/stage1_prod
146	  name: maya_asr_stage1_tdt_v2
147	  create_tensorboard_logger: true
148	  create_wandb_logger: true
149	  wandb_logger_kwargs:
150	    project: maya-asr
151	    name: tdt-parakeet-finetune-stage1-v2
152	    resume: allow
153	  resume_if_exists: false
154	  resume_ignore_no_checkpoint: true
155	  create_checkpoint_callback: true
156	  checkpoint_callback_params:
157	    monitor: val_loss
158	    mode: min
159	    save_top_k: 3
160	    save_last: true
161	    every_n_train_steps: 5000
162	    every_n_epochs: 0
163	    save_nemo_on_train_end: true
164


Error: File content (13522 tokens) exceeds maximum allowed tokens (10000). Use offset and limit parameters to read specific portions of the file, or search for specific content instead of reading the whole file.


Let me read the training script in sections:


1	#!/usr/bin/env python3
2	"""Production training script with optimized data loading + TDT + differential LR.
3
4	Data strategy (ported from parakeet pipeline):
5	  - Duration bucketing: groups similar-length clips → minimal padding waste
6	  - Language temperature rebalancing (T=0.3): upweights low-resource languages
7	  - Max-batch-duration packing: caps audio per micro-batch → stable VRAM
8	  - Compute-cost DDP sharding: balances load across GPUs
9	  - Grad accum alignment: prevents DDP sync bugs
10
11	Optimizations v2:
12	  - TDT loss (matching pretrained Parakeet encoder objective)
13	  - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
14	  - Encoder frozen for first 5K steps, then unfrozen
15	  - CTC loss warmup: 0→0.3 over first 3K steps
16	  - Language embedding conditioning
17	  - OOM-safe VRAM ceiling design
18
19	Audio I/O: TarOffsetReader (os.pread, zero extraction from existing tars).
20	Model: NeMo EncDecHybridRNNTCTCBPEModel (FastConformer + TDT + aux CTC).
21
22	Usage:
23	  # Full production run (8×H200)
24	  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml
25
26	  # Quick smoke test
27	  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml \
28	    --max-steps 100 --devices 1 --log-every 10 --val-every 50
29	"""
30
31	import argparse
32	import os
33	import sys
34	import time
35	from pathlib import Path
36
37	import lightning.pytorch as pl
38	import nemo.collections.asr as nemo_asr
39	import numpy as np
40	import pyarrow.parquet as pq
41	import torch
42	from nemo.utils.exp_manager import exp_manager
43	from omegaconf import OmegaConf, open_dict
44
45
46	# ---------------------------------------------------------------------------
47	# Language embedding module
48	# ---------------------------------------------------------------------------
49	LANG_TO_ID = {
50	    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
51	    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
52	}
53	NUM_LANGUAGES = len(LANG_TO_ID)
54
55
56	class LanguageEmbedding(torch.nn.Module):
57	    """Projects a language ID into a per-frame additive bias for the encoder."""
58
59	    def __init__(self, num_languages: int, embed_dim: int):
60	        super().__init__()
61	        self.embed = torch.nn.Embedding(num_languages, embed_dim)
62	        # Initialize near-zero so it doesn't disrupt pretrained encoder features
63	        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
64
65	    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
66	        """
67	        Args:
68	            lang_ids: (B,) integer tensor of language IDs
69	        Returns:
70	            (B, embed_dim, 1) — broadcast-addable to encoder output (B, D, T)
71	        """
72	        # Encoder output is (B, D, T), so we need (B, D, 1) for broadcasting
73	        return self.embed(lang_ids).unsqueeze(2)
74
75
76	# ---------------------------------------------------------------------------
77	# Callbacks
78	# ---------------------------------------------------------------------------
79	class ThroughputCallback(pl.Callback):
80	    """Logs VRAM, throughput (samples/s, audio-s/s), and MFU every log_every steps."""
81
82	    def __init__(self, log_every: int = 10, model_params: int = 0,
83	                 grad_accum: int = 1, world_size: int = 1):
84	        self.log_every = log_every
85	        self.model_params = model_params
86	        self.grad_accum = grad_accum
87	        self.world_size = world_size
88	        self._window_t0 = None
89	        self._window_samples = 0
90	        self._window_audio_s = 0.0
91	        self._window_micro_steps = 0
92	        # H200 specs: bf16 peak ~990 TFLOPS
93	        self._gpu_peak_tflops = 989.0
94
95	    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
96	        if self._window_t0 is None:
97	            self._window_t0 = time.time()
98
99	    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
100	        if batch is not None and len(batch) >= 2:
101	            sig_lens = batch[1]
102	            n_samples = sig_lens.shape[0]
103	            audio_s = sig_lens.float().sum().item() / 16000.0
104	            self._window_samples += n_samples
105	            self._window_audio_s += audio_s
106	        self._window_micro_steps += 1
107
108	        rank = int(os.environ.get("LOCAL_RANK", 0))
109	        global_step = trainer.global_step
110
111	        if rank == 0 and global_step > 0 and global_step % self.log_every == 0 and self._window_t0 is not None:
112	            elapsed = time.time() - self._window_t0
113	            if elapsed > 0:
114	                samples_per_s = self._window_samples / elapsed
115	                audio_s_per_s = self._window_audio_s / elapsed
116	                rtf_per_gpu = audio_s_per_s / max(self.world_size, 1)
117	                opt_steps_per_s = (self._window_micro_steps / max(self.grad_accum, 1)) / elapsed
118
119	                vram_alloc_gb = torch.cuda.memory_allocated() / 1e9
120	                vram_reserved_gb = torch.cuda.memory_reserved() / 1e9
121	                vram_peak_gb = torch.cuda.max_memory_allocated() / 1e9
122
123	                mel_frames_per_s = audio_s_per_s * 100
124	                estimated_tflops = (6 * self.model_params * mel_frames_per_s) / 1e12
125	                mfu = estimated_tflops / (self._gpu_peak_tflops * self.world_size) * 100
126
127	                print(f"  [step {global_step:>6d}] "
128	                      f"samples/s={samples_per_s:.1f}  "
129	                      f"audio-s/s={audio_s_per_s:.0f}  "
130	                      f"RTF/GPU={rtf_per_gpu:.1f}  "
131	                      f"opt-steps/s={opt_steps_per_s:.3f}  "
132	                      f"VRAM={vram_alloc_gb:.1f}/{vram_reserved_gb:.1f}GB  "
133	                      f"peak={vram_peak_gb:.1f}GB  "
134	                      f"MFU≈{mfu:.1f}%")
135
136	                # Log custom metrics to wandb/tensorboard via Lightning
137	                if pl_module is not None:
138	                    custom_metrics = {
139	                        "throughput/samples_per_s": samples_per_s,
140	                        "throughput/audio_s_per_s": audio_s_per_s,
141	                        "throughput/rtf_per_gpu": rtf_per_gpu,
142	                        "throughput/opt_steps_per_s": opt_steps_per_s,
143	                        "throughput/mfu_pct": mfu,
144	                        "memory/vram_alloc_gb": vram_alloc_gb,
145	                        "memory/vram_reserved_gb": vram_reserved_gb,
146	                        "memory/vram_peak_gb": vram_peak_gb,
147	                    }
148	                    # Log LR per param group
149	                    optimizer = trainer.optimizers[0] if trainer.optimizers else None
150	                    if optimizer:
151	                        for i, pg in enumerate(optimizer.param_groups):
152	                            name = pg.get('name', f'group_{i}')
153	                            custom_metrics[f"lr/{name}"] = pg['lr']
154	                    pl_module.log_dict(custom_metrics, prog_bar=False)
155
156	            self._window_t0 = time.time()
157	            self._window_samples = 0
158	            self._window_audio_s = 0.0
159	            self._window_micro_steps = 0
160
161
162	class EncoderFreezeCallback(pl.Callback):
163	    """Freezes encoder for the first N optimizer steps, then unfreezes."""
164
165	    def __init__(self, freeze_steps: int = 5000):
166	        self.freeze_steps = freeze_steps
167	        self._frozen = False
168	        self._unfrozen = False
169
170	    def on_train_start(self, trainer, pl_module):
171	        if self.freeze_steps > 0:
172	            self._freeze_encoder(pl_module)
173
174	    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
175	        if self._unfrozen:
176	            return
177	        if trainer.global_step >= self.freeze_steps and self._frozen:
178	            self._unfreeze_encoder(pl_module)
179
180	    def _freeze_encoder(self, model):
181	        frozen = 0
182	        for name, param in model.named_parameters():
183	            if name.startswith("encoder."):
184	                param.requires_grad = False
185	                frozen += 1
186	        self._frozen = True
187	        rank = int(os.environ.get("LOCAL_RANK", 0))
188	        if rank == 0:
189	            print(f"  [EncoderFreeze] Froze {frozen} encoder params for first {self.freeze_steps} steps")
190
191	    def _unfreeze_encoder(self, model):
192	        unfrozen = 0
193	        for name, param in model.named_parameters():
194	            if name.startswith("encoder."):
195	                param.requires_grad = True
196	                unfrozen += 1
197	        self._frozen = False
198	        self._unfrozen = True
199	        rank = int(os.environ.get("LOCAL_RANK", 0))
200	        if rank == 0:


6:  - Language temperature rebalancing (T=0.3): upweights low-resource languages
16:  - Language embedding conditioning
47:# Language embedding module
56:class LanguageEmbedding(torch.nn.Module):
57:    """Projects a language ID into a per-frame additive bias for the encoder."""
59:    def __init__(self, num_languages: int, embed_dim: int):
61:        self.embed = torch.nn.Embedding(num_languages, embed_dim)
65:    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
68:            lang_ids: (B,) integer tensor of language IDs
79:class ThroughputCallback(pl.Callback):
82:    def __init__(self, log_every: int = 10, model_params: int = 0,
95:    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
99:    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
162:class EncoderFreezeCallback(pl.Callback):
165:    def __init__(self, freeze_steps: int = 5000):
170:    def on_train_start(self, trainer, pl_module):
174:    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
180:    def _freeze_encoder(self, model):
191:    def _unfreeze_encoder(self, model):
204:class SamplerStateCallback(pl.Callback):
212:    def __init__(self, sampler, state_dir: str = "/tmp"):
216:    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
221:    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
224:    def _save_state(self, global_step):
236:    def load_state(state_dir: str = "/tmp") -> dict | None:
245:class CTCWarmupCallback(pl.Callback):
251:    def __init__(self, target_weight: float = 0.3, warmup_steps: int = 3000):
255:    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
266:class DifferentialLRCallback(pl.Callback):
276:    def __init__(self, encoder_lr: float = 2e-5, head_lr: float = 1e-3):
280:    def on_train_start(self, trainer, pl_module):
338:def compute_vram_safe_batch_params(
414:class NeMoTarOffsetDataset(torch.utils.data.Dataset):
420:    Also stores language IDs for language conditioning.
423:    def __init__(self, manifest_parquet, tokenizer, sample_rate=16000,
449:        self.languages = self.df["language"].values
451:        # Pre-compute language IDs
453:            [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"]) for lang in self.languages],
457:    def __len__(self):
460:    def __getitem__(self, idx):
490:def collate_nemo_batch(batch):
523:def get_ddp_info() -> tuple[int, int]:
533:# Language-conditioned forward hook
535:def install_language_conditioning(model, lang_embed: LanguageEmbedding):
536:    """Install a forward hook on the encoder to add language embeddings.
538:    After the encoder produces its output, we add the language embedding
539:    as a per-frame bias. This gives the encoder a language signal without
543:    model._lang_embed = lang_embed
547:    def encoder_forward_with_lang(audio_signal, length, lang_ids=None):
549:        if lang_ids is not None and hasattr(model, '_lang_embed'):
550:            lang_bias = model._lang_embed(lang_ids)  # (B, 1, D)
559:# Custom training step with language conditioning
561:def patch_training_step(model):
563:    and pass them through to the encoder via language conditioning.
567:    def patched_training_step(batch, batch_idx):
582:def patch_encoder_for_lang(model):
583:    """Patch encoder forward to inject language IDs stored on model."""
586:    def encoder_forward_with_lang(audio_signal, length):
588:        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
593:                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
605:def _tokenize_chunk_worker(texts, tokenizer_model_path):
613:def main():
654:                        help="Disable language embedding conditioning")
724:        print(f"  Lang embed:     {'yes' if not args.no_lang_embed else 'no'}")
800:    # Language conditioning
802:    if not args.no_lang_embed:
803:        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804:        model._lang_embed = lang_embed
806:        model.register_module('_lang_embed_module', lang_embed)
810:            lang_params = sum(p.numel() for p in lang_embed.parameters())
811:            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
855:        # Language distribution
856:        unique_langs, lang_counts = np.unique(train_dataset.languages, return_counts=True)
857:        print(f"Languages:     {len(unique_langs)} ({', '.join(sorted(unique_langs))})")
916:        languages=train_dataset.languages,
945:        lang_dist = sampler.get_language_distribution()
946:        print("Language rebalancing (T={:.1f}):".format(args.temperature))


Now let me look at the data loading and model architecture more closely:


1	"""Production batch sampler: duration bucketing + language temperature rebalancing.
2
3	OOM-safe design (3-layer protection):
4	  1. Per-bucket BS caps — limits samples per batch based on duration bucket
5	  2. max_batch_duration — caps total audio seconds per batch
6	  3. max_tokens_in_batch — caps total token count per batch (controls T_dec in joint)
7
8	Vectorized with numpy for fast iteration over 77M+ samples (~30s vs 8min).
9
10	Usage:
11	    sampler = ProductionBatchSampler(
12	        durations=dataset.durations,
13	        languages=dataset.languages,
14	        token_lengths=dataset.token_lengths,
15	        max_batch_duration=90.0,
16	        max_tokens_in_batch=400,
17	        temperature=0.3,
18	        rank=rank, world_size=world_size,
19	        grad_accum=4,
20	    )
21	    dl = DataLoader(dataset, batch_sampler=sampler, ...)
22	"""
23
24	import math
25	import os
26	import time
27	from collections import defaultdict
28	from typing import Iterator
29
30	import numpy as np
31	from torch.utils.data import Sampler
32
33	# Default per-bucket BS caps
34	DEFAULT_BUCKET_BS_CAPS = {
35	    0: 24,   # 0-4s
36	    1: 16,   # 4-8s
37	    2: 12,   # 8-12s
38	    3: 8,    # 12-16s
39	    4: 6,    # 16-20s
40	}
41
42
43	class ProductionBatchSampler(Sampler[list[int]]):
44	    """DDP-aware batch sampler with triple-cap OOM protection.
45
46	    Vectorized with numpy — iterates 77M samples in ~30s instead of 8min.
47	    """
48
49	    def __init__(
50	        self,
51	        durations: np.ndarray,
52	        languages: np.ndarray,
53	        token_lengths: np.ndarray | None = None,
54	        bucket_boundaries: list[float] | None = None,
55	        max_batch_duration: float = 90.0,
56	        max_batch_size: int = 16,
57	        max_tokens_in_batch: int = 400,
58	        bucket_bs_caps: dict[int, int] | None = None,
59	        temperature: float = 0.3,
60	        rank: int = 0,
61	        world_size: int = 1,
62	        grad_accum: int = 1,
63	        shuffle: bool = True,
64	        seed: int = 42,
65	        drop_last: bool = True,
66	    ):
67	        self.durations = np.asarray(durations, dtype=np.float32)
68	        self.languages = np.asarray(languages, dtype=object)
69	        self.max_batch_duration = max_batch_duration
70	        self.max_batch_size = max_batch_size
71	        self.max_tokens_in_batch = max_tokens_in_batch
72	        self.temperature = temperature
73	        self.rank = rank
74	        self.world_size = world_size
75	        self.grad_accum = grad_accum
76	        self.shuffle = shuffle
77	        self.seed = seed
78	        self.drop_last = drop_last
79	        self.epoch = 0
80
81	        if token_lengths is not None:
82	            self.token_lengths = np.asarray(token_lengths, dtype=np.int32)
83	        else:
84	            self.token_lengths = np.maximum(1, (self.durations * 7).astype(np.int32))
85
86	        if bucket_boundaries is None:
87	            bucket_boundaries = [4.0, 8.0, 12.0, 16.0]
88	        self.boundaries = sorted(bucket_boundaries)
89
90	        self.bucket_bs_caps = bucket_bs_caps or DEFAULT_BUCKET_BS_CAPS
91
92	        # Pre-compute bucket assignments (vectorized)
93	        self._bucket_ids = np.searchsorted(self.boundaries, self.durations)
94
95	        # Pre-compute language sample weights
96	        self._sample_weights = self._compute_language_weights()
97
98	        self._last_stats = {}
99
100	    def _compute_language_weights(self) -> np.ndarray:
101	        lang_counts = defaultdict(int)
102	        for lang in self.languages:
103	            lang_counts[lang] += 1
104
105	        total = len(self.languages)
106	        if self.temperature <= 0:
107	            lang_weights = {lang: 1.0 for lang in lang_counts}
108	        else:
109	            lang_weights = {
110	                lang: math.pow(count, self.temperature)
111	                for lang, count in lang_counts.items()
112	            }
113
114	        weight_sum = sum(lang_weights.values())
115	        lang_target_props = {lang: w / weight_sum for lang, w in lang_weights.items()}
116
117	        sample_weights = np.ones(total, dtype=np.float64)
118	        for lang, count in lang_counts.items():
119	            actual_prop = count / total
120	            target_prop = lang_target_props[lang]
121	            weight = target_prop / actual_prop if actual_prop > 0 else 0.0
122	            mask = self.languages == lang
123	            sample_weights[mask] = weight
124
125	        return sample_weights
126
127	    def _get_bucket_bs(self, bucket_id: int) -> int:
128	        bucket_cap = self.bucket_bs_caps.get(bucket_id, self.max_batch_size)
129	        return min(bucket_cap, self.max_batch_size)
130
131	    def set_epoch(self, epoch: int):
132	        self.epoch = epoch
133
134	    def set_resume_state(self, consumed_batches: int = 0):
135	        """Set number of batches to skip on resume."""
136	        self._skip_batches = consumed_batches
137
138	    def get_state(self) -> dict:
139	        """Get sampler state for checkpointing."""
140	        return {
141	            "epoch": self.epoch,
142	            "consumed_batches": getattr(self, "_consumed_batches", 0),
143	        }
144
145	    def _pack_bucket_vectorized(self, indices: np.ndarray, bucket_bs: int,
146	                                 rng_state: np.random.RandomState) -> list[list[int]]:
147	        """Pack a single bucket into batches using vectorized operations.
148
149	        Uses numpy for the weighted permutation and greedy packing.
150	        """


1	"""Central configuration constants for Maya ASR."""
2
3	import hashlib
4	from pathlib import Path
5
6	# Paths
7	PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
8	SFT_DATA_ROOT = Path("/root/sft_data")
9	FINAL_EXPORT_ROOT = SFT_DATA_ROOT / "final-export" / "production" / "shards"
10	EXPERIMENTS_DIR = PROJECT_ROOT / "experiments"
11
12	# Languages
13	LANGUAGES = ["as", "bn", "en", "gu", "hi", "kn", "ml", "mr", "or", "pa", "ta", "te"]
14
15	# Shard structure
16	SHARD_FILES = ["audio.tar", "audio_index.parquet", "manifest.json", "metadata.parquet"]
17
18	# Model
19	ENCODER_DIM = 1024
20	ENCODER_LAYERS = 32
21	GEMMA_HIDDEN_DIM = 1536
22
23
24	def file_sha256(path: Path) -> str:
25	    """Compute SHA-256 hex digest of a file."""
26	    h = hashlib.sha256()
27	    with open(path, "rb") as f:
28	        for chunk in iter(lambda: f.read(1 << 20), b""):
29	            h.update(chunk)
30	    return h.hexdigest()
31


Let me examine more of the training script and check for model architecture details:


200	        if rank == 0:
201	            print(f"  [EncoderFreeze] Unfroze {unfrozen} encoder params at step {model.trainer.global_step}")
202
203
204	class SamplerStateCallback(pl.Callback):
205	    """Saves sampler state (consumed batches) alongside checkpoints.
206
207	    On resume, the sampler can skip already-trained batches to avoid
208	    reprocessing the same data.
209	    """
210	    STATE_FILE = "sampler_state.json"
211
212	    def __init__(self, sampler, state_dir: str = "/tmp"):
213	        self.sampler = sampler
214	        self.state_dir = state_dir
215
216	    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
217	        # Save state every 1000 optimizer steps
218	        if trainer.global_step > 0 and trainer.global_step % 1000 == 0:
219	            self._save_state(trainer.global_step)
220
221	    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
222	        self._save_state(trainer.global_step)
223
224	    def _save_state(self, global_step):
225	        import json
226	        state = {
227	            "global_step": global_step,
228	            "epoch": self.sampler.epoch,
229	            "consumed_batches": getattr(self.sampler, "_consumed_batches", 0),
230	        }
231	        state_path = os.path.join(self.state_dir, self.STATE_FILE)
232	        with open(state_path, "w") as f:
233	            json.dump(state, f)
234
235	    @staticmethod
236	    def load_state(state_dir: str = "/tmp") -> dict | None:
237	        import json
238	        state_path = os.path.join(state_dir, SamplerStateCallback.STATE_FILE)
239	        if os.path.exists(state_path):
240	            with open(state_path) as f:
241	                return json.load(f)
242	        return None
243
244
245	class CTCWarmupCallback(pl.Callback):
246	    """Warms up CTC loss weight from 0 to target over N steps.
247
248	    Prevents random CTC gradients from corrupting pretrained encoder in early training.
249	    """
250
251	    def __init__(self, target_weight: float = 0.3, warmup_steps: int = 3000):
252	        self.target_weight = target_weight
253	        self.warmup_steps = warmup_steps
254
255	    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
256	        step = trainer.global_step
257	        if step < self.warmup_steps:
258	            weight = self.target_weight * (step / self.warmup_steps)
259	        else:
260	            weight = self.target_weight
261
262	        if hasattr(pl_module, 'ctc_loss_weight'):
263	            pl_module.ctc_loss_weight = weight
264
265
266	class DifferentialLRCallback(pl.Callback):
267	    """Sets differential learning rates for encoder vs decoder/joint/CTC param groups.
268
269	    After NeMo constructs the optimizer (which uses a single LR), this callback
270	    modifies the param groups to use different LRs.
271
272	    Encoder (pretrained): low LR to preserve features
273	    Decoder/Joint/CTC (random init): high LR for fast adaptation
274	    """
275
276	    def __init__(self, encoder_lr: float = 2e-5, head_lr: float = 1e-3):
277	        self.encoder_lr = encoder_lr
278	        self.head_lr = head_lr
279
280	    def on_train_start(self, trainer, pl_module):
281	        optimizer = trainer.optimizers[0]
282
283	        # NeMo creates a single param group. We need to split it into
284	        # encoder (low LR) and head (high LR) groups.
285	        # Include ALL params regardless of requires_grad (freeze callback
286	        # will toggle requires_grad but params must be in optimizer).
287
288	        encoder_params = []
289	        head_params = []
290
291	        for name, param in pl_module.named_parameters():
292	            if name.startswith("encoder."):
293	                encoder_params.append(param)
294	            else:
295	                head_params.append(param)
296
297	        # Preserve optimizer settings from existing group
298	        base_group = optimizer.param_groups[0].copy()
299
300	        optimizer.param_groups.clear()
301
302	        encoder_group = base_group.copy()
303	        encoder_group['params'] = encoder_params
304	        encoder_group['lr'] = self.encoder_lr
305	        encoder_group['name'] = 'encoder'
306
307	        head_group = base_group.copy()
308	        head_group['params'] = head_params
309	        head_group['lr'] = self.head_lr
310	        head_group['name'] = 'decoder_joint_ctc'
311
312	        optimizer.param_groups.append(encoder_group)
313	        optimizer.param_groups.append(head_group)
314
315	        # Fix LR scheduler: it was initialized with 1 group, now we have 2.
316	        # Update base_lrs and last_epoch tracking for the scheduler.
317	        if trainer.lr_scheduler_configs:
318	            for sched_config in trainer.lr_scheduler_configs:
319	                scheduler = sched_config.scheduler
320	                # Update base_lrs to match new param groups
321	                scheduler.base_lrs = [self.encoder_lr, self.head_lr]
322	                # Update _last_lr if it exists
323	                if hasattr(scheduler, '_last_lr'):
324	                    scheduler._last_lr = [self.encoder_lr, self.head_lr]
325
326	        rank = int(os.environ.get("LOCAL_RANK", 0))
327	        if rank == 0:
328	            enc_count = sum(p.numel() for p in encoder_params)
329	            head_count = sum(p.numel() for p in head_params)
330	            enc_trainable = sum(p.numel() for p in encoder_params if p.requires_grad)
331	            print(f"  [DiffLR] Encoder: {enc_count/1e6:.0f}M params @ lr={self.encoder_lr} ({enc_trainable/1e6:.0f}M trainable)")
332	            print(f"  [DiffLR] Heads:   {head_count/1e6:.0f}M params @ lr={self.head_lr}")
333
334
335	# ---------------------------------------------------------------------------
336	# OOM-safe VRAM budget calculator
337	# ---------------------------------------------------------------------------
338	def compute_vram_safe_batch_params(
339	    vocab_size: int = 32769,
340	    enc_hidden: int = 1024,
341	    joint_hidden: int = 640,
342	    num_durations: int = 5,
343	    max_audio_dur: float = 20.0,
344	    subsampling_factor: int = 8,
345	    sample_rate: int = 16000,
346	    window_stride: float = 0.01,
347	    gpu_memory_gb: float = 140.0,
348	    model_params_gb: float = 4.5,  # ~1.1B params in bf16
349	    safety_margin: float = 0.75,  # use at most 75% of available
350	    grad_accum: int = 4,
351	) -> dict:
352	    """Compute OOM-safe batch parameters based on VRAM budget.
353
354	    The dominant VRAM consumer is the TDT joint tensor:
355	      joint_size = B × T_enc × T_dec × (vocab_size + 1 + num_durations)
356	    For TDT, T_dec is much smaller than RNNT because of duration skipping.
357
358	    Returns dict with recommended max_batch_dur, max_batch_size.
359	    """
360	    # Available VRAM for activations (after model + optimizer + gradients)
361	    optimizer_gb = model_params_gb * 3  # AdamW: params + momentum + variance (bf16 + fp32)
362	    available_gb = gpu_memory_gb - model_params_gb - optimizer_gb
363	    activation_budget_gb = available_gb * safety_margin
364
365	    # Worst-case encoder frames for max_audio_dur
366	    mel_frames = max_audio_dur / window_stride  # 2000 for 20s
367	    enc_frames = mel_frames / subsampling_factor  # 250 for 8x sub
368
369	    # TDT joint tensor: B × T_enc × T_dec × (V + 1 + num_durations)
370	    # For TDT, T_dec ≈ max_tokens (32K vocab = ~50 tokens for 20s audio)
371	    # Conservative: assume T_dec = 80 (worst case for long utterance)
372	    max_t_dec = 80
373	    output_dim = vocab_size + 1 + num_durations
374
375	    # Joint tensor size per sample in GB (bf16 = 2 bytes)
376	    joint_per_sample_gb = (enc_frames * max_t_dec * output_dim * 2) / 1e9
377	    # With gradient: 2x (forward + backward)
378	    joint_per_sample_gb *= 2
379
380	    # Add encoder activations per sample (~0.5 GB for 42-layer conformer on 20s)
381	    enc_activations_per_sample_gb = 0.5
382
383	    total_per_sample_gb = joint_per_sample_gb + enc_activations_per_sample_gb
384
385	    # Max batch size (samples)
386	    max_safe_batch_size = max(1, int(activation_budget_gb / total_per_sample_gb))
387
388	    # Max batch duration (seconds)
389	    # Average audio in a batch = max_batch_dur, average sample ~8s
390	    avg_sample_dur = 8.0
391	    avg_batch_size = max_safe_batch_size
392	    max_safe_batch_dur = avg_batch_size * avg_sample_dur
393
394	    return {
395	        "max_batch_size": min(max_safe_batch_size, 16),  # cap at 16 for safety
396	        "max_batch_dur": min(max_safe_batch_dur, 120.0),  # cap at 120s
397	        "activation_budget_gb": round(activation_budget_gb, 1),
398	        "joint_per_sample_gb": round(joint_per_sample_gb, 3),
399	        "total_per_sample_gb": round(total_per_sample_gb, 3),
400	        "gpu_memory_gb": gpu_memory_gb,
401	        "model_params_gb": round(model_params_gb, 1),
402	        "optimizer_gb": round(optimizer_gb, 1),
403	    }
404
405
406	# ---------------------------------------------------------------------------
407	# Dataset and collate
408	# ---------------------------------------------------------------------------
409	sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
410	from lf_asr.data.production_sampler import ProductionBatchSampler
411	from lf_asr.data.tar_offset_reader import TarOffsetReader
412
413
414	class NeMoTarOffsetDataset(torch.utils.data.Dataset):
415	    """NeMo-compatible dataset backed by TarOffsetReader.
416
417	    Returns tuples: (signal, signal_length, tokens, token_length)
418	    matching NeMo's EncDecHybridRNNTCTCBPEModel expectations.
419
420	    Also stores language IDs for language conditioning.
421	    """
422
423	    def __init__(self, manifest_parquet, tokenizer, sample_rate=16000,
424	                 max_duration=30.0, min_duration=0.5):
425	        df = pq.read_table(manifest_parquet).to_pandas()
426	        mask = (
427	            (df["duration_s"] >= min_duration)
428	            & (df["duration_s"] <= max_duration)
429	            & (df["transcript"].str.strip().str.len() > 0)
430	        )
431	        n_blank = (~(df["transcript"].str.strip().str.len() > 0)).sum()
432	        if n_blank > 0:
433	            print(f"  WARNING: Dropped {n_blank:,} blank transcripts at load time")
434	        n_dur = ((df["duration_s"] < min_duration) | (df["duration_s"] > max_duration)).sum()
435	        if n_dur > 0:
436	            print(f"  Filtered {n_dur:,} samples outside [{min_duration}, {max_duration}]s")
437
438	        self.df = df[mask].reset_index(drop=True)
439	        self.reader = TarOffsetReader(max_fd_cache=512)
440	        self.tokenizer = tokenizer
441	        self.sample_rate = sample_rate
442
443	        # Pre-extract numpy arrays for fast indexing
444	        self.tar_paths = self.df["tar_path"].values
445	        self.offsets = self.df["tar_offset_data"].values.astype(np.int64)
446	        self.nbytes = self.df["tar_nbytes"].values.astype(np.int64)
447	        self.transcripts = self.df["transcript"].values
448	        self.durations = self.df["duration_s"].values.astype(np.float32)
449	        self.languages = self.df["language"].values


maya-asr-hybrid-fastconformer-rnnt-stage1/src/maya_asr/__init__.py
maya-asr-hybrid-fastconformer-rnnt-stage1/src/maya_asr/config.py


Let me check the complete training script for the main function and model initialization:


450
451	        # Pre-compute language IDs
452	        self.lang_ids = np.array(
453	            [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"]) for lang in self.languages],
454	            dtype=np.int64,
455	        )
456
457	    def __len__(self):
458	        return len(self.df)
459
460	    def __getitem__(self, idx):
461	        tar_path = str(self.tar_paths[idx])
462	        offset = int(self.offsets[idx])
463	        nbytes = int(self.nbytes[idx])
464
465	        try:
466	            waveform, sr = self.reader.read(tar_path, offset, nbytes)
467	        except Exception as e:
468	            print(f"  WARNING: Bad audio idx={idx} tar={tar_path} offset={offset}: {e}")
469	            return None
470	        if sr != self.sample_rate:
471	            print(f"  WARNING: Sample rate mismatch idx={idx}: {sr} != {self.sample_rate}")
472	            return None
473
474	        transcript = str(self.transcripts[idx]).strip()
475	        if not transcript:
476	            transcript = "<blank>"
477	        tokens = self.tokenizer.text_to_ids(transcript)
478
479	        lang_id = int(self.lang_ids[idx])
480
481	        return (
482	            torch.from_numpy(waveform).float(),
483	            torch.tensor(len(waveform), dtype=torch.long),
484	            torch.tensor(tokens, dtype=torch.long),
485	            torch.tensor(len(tokens), dtype=torch.long),
486	            torch.tensor(lang_id, dtype=torch.long),
487	        )
488
489
490	def collate_nemo_batch(batch):
491	    """Collate with padding — handles None samples from bad audio.
492
493	    Returns (signals, sig_lens, targets, tgt_lens, lang_ids).
494	    """
495	    batch = [b for b in batch if b is not None]
496	    if not batch:
497	        return (
498	            torch.zeros(1, 1),
499	            torch.tensor([1], dtype=torch.long),
500	            torch.zeros(1, 1, dtype=torch.long),
501	            torch.tensor([0], dtype=torch.long),
502	            torch.tensor([0], dtype=torch.long),
503	        )
504
505	    signals, sig_lens, targets, tgt_lens, lang_ids = zip(*batch)
506
507	    max_sig_len = max(s.shape[0] for s in signals)
508	    max_tgt_len = max(t.shape[0] for t in targets)
509
510	    padded_signals = torch.zeros(len(signals), max_sig_len)
511	    padded_targets = torch.zeros(len(targets), max_tgt_len, dtype=torch.long)
512	    signal_lengths = torch.stack(list(sig_lens))
513	    target_lengths = torch.stack(list(tgt_lens))
514	    lang_id_tensor = torch.stack(list(lang_ids))
515
516	    for i, (sig, tgt) in enumerate(zip(signals, targets)):
517	        padded_signals[i, :sig.shape[0]] = sig
518	        padded_targets[i, :tgt.shape[0]] = tgt
519
520	    return padded_signals, signal_lengths, padded_targets, target_lengths, lang_id_tensor
521
522
523	def get_ddp_info() -> tuple[int, int]:
524	    """Get (rank, world_size) from DDP environment."""
525	    if torch.distributed.is_available() and torch.distributed.is_initialized():
526	        return torch.distributed.get_rank(), torch.distributed.get_world_size()
527	    rank = int(os.environ.get("LOCAL_RANK", os.environ.get("RANK", 0)))
528	    world_size = int(os.environ.get("WORLD_SIZE", 1))
529	    return rank, world_size
530
531
532	# ---------------------------------------------------------------------------
533	# Language-conditioned forward hook
534	# ---------------------------------------------------------------------------
535	def install_language_conditioning(model, lang_embed: LanguageEmbedding):
536	    """Install a forward hook on the encoder to add language embeddings.
537
538	    After the encoder produces its output, we add the language embedding
539	    as a per-frame bias. This gives the encoder a language signal without
540	    modifying the encoder architecture (preserving pretrained weights).
541	    """
542	    # Store reference on model so it moves to GPU with model
543	    model._lang_embed = lang_embed
544
545	    original_forward = model.encoder.forward
546
547	    def encoder_forward_with_lang(audio_signal, length, lang_ids=None):
548	        encoded, encoded_len = original_forward(audio_signal=audio_signal, length=length)
549	        if lang_ids is not None and hasattr(model, '_lang_embed'):
550	            lang_bias = model._lang_embed(lang_ids)  # (B, 1, D)
551	            encoded = encoded + lang_bias
552	        return encoded, encoded_len
553
554	    model.encoder.forward = encoder_forward_with_lang
555	    return model
556
557
558	# ---------------------------------------------------------------------------
559	# Custom training step with language conditioning
560	# ---------------------------------------------------------------------------
561	def patch_training_step(model):
562	    """Patch the model's training_step to extract lang_ids from batch
563	    and pass them through to the encoder via language conditioning.
564	    """
565	    original_training_step = model.training_step
566
567	    def patched_training_step(batch, batch_idx):
568	        # NeMo expects batch = (signal, signal_len, tokens, token_len)
569	        # We added lang_ids as 5th element
570	        if isinstance(batch, (list, tuple)) and len(batch) >= 5:
571	            signal, signal_len, tokens, token_len, lang_ids = batch[:5]
572	            # Store lang_ids on model for encoder hook to pick up
573	            model._current_lang_ids = lang_ids.to(signal.device)
574	            # Pass standard 4-tuple to NeMo
575	            nemo_batch = (signal, signal_len, tokens, token_len)
576	            return original_training_step(nemo_batch, batch_idx)
577	        return original_training_step(batch, batch_idx)
578
579	    model.training_step = patched_training_step
580
581
582	def patch_encoder_for_lang(model):
583	    """Patch encoder forward to inject language IDs stored on model."""
584	    original_encoder_forward = model.encoder.forward
585
586	    def encoder_forward_with_lang(audio_signal, length):
587	        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal, length=length)
588	        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
589	                and model._current_lang_ids is not None):
590	            lang_ids = model._current_lang_ids
591	            # Safety: only apply if batch sizes match (skip during validation)
592	            if lang_ids.shape[0] == encoded.shape[0]:
593	                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
594	                encoded = encoded + lang_bias
595	            # Clear after use to prevent stale values during validation
596	            model._current_lang_ids = None
597	        return encoded, encoded_len
598
599	    model.encoder.forward = encoder_forward_with_lang
600
601
602	# ---------------------------------------------------------------------------
603	# Main
604	# ---------------------------------------------------------------------------
605	def _tokenize_chunk_worker(texts, tokenizer_model_path):
606	    """Tokenize a chunk of texts in a worker process. Module-level for pickling."""
607	    import sentencepiece as spm
608	    sp = spm.SentencePieceProcessor()
609	    sp.Load(tokenizer_model_path)
610	    return [len(sp.EncodeAsIds(str(t).strip() or "<blank>")) for t in texts]
611
612
613	def main():
614	    parser = argparse.ArgumentParser(
615	        description="Production NeMo ASR training with TDT + differential LR"
616	    )
617	    parser.add_argument("--config", type=str, required=True)
618	    parser.add_argument("--max-steps", type=int, default=None)
619	    parser.add_argument("--devices", type=int, default=None)
620	    parser.add_argument("--train-parquet", type=str,
621	                        default="artifacts/phase3/production_train_final.parquet")
622	    parser.add_argument("--val-manifest", type=str,
623	                        default="data/manifests/stage1_prod_val_v2.jsonl")
624	    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
625	    parser.add_argument("--resume-weights-only", action="store_true", default=False,
626	                        help="Resume model weights from checkpoint but reset optimizer (for config changes)")
627	    parser.add_argument("--pretrained-encoder", type=str, default=None,
628	                        help="Path to pretrained encoder weights (.pt file)")
629	    parser.add_argument("--smoke", action="store_true", default=False)
630	    parser.add_argument("--val-every", type=int, default=None)
631	    parser.add_argument("--log-every", type=int, default=None)
632
633	    # Data strategy args
634	    parser.add_argument("--max-batch-dur", type=float, default=120.0,
635	                        help="Max audio seconds per micro-batch (default: 120s)")
636	    parser.add_argument("--max-batch-size", type=int, default=16,
637	                        help="Global hard cap on samples per micro-batch")
638	    parser.add_argument("--max-tokens-in-batch", type=int, default=400,
639	                        help="Max total tokens per batch (controls joint tensor T_dec)")
640	    parser.add_argument("--temperature", type=float, default=0.3)
641	    parser.add_argument("--grad-accum", type=int, default=None)
642	    parser.add_argument("--num-workers", type=int, default=None)
643
644	    # v2 optimization args
645	    parser.add_argument("--encoder-lr", type=float, default=2e-5,
646	                        help="Learning rate for pretrained encoder")
647	    parser.add_argument("--head-lr", type=float, default=1e-3,
648	                        help="Learning rate for decoder/joint/CTC heads")
649	    parser.add_argument("--freeze-encoder-steps", type=int, default=5000,
650	                        help="Freeze encoder for this many optimizer steps (0=no freeze)")
651	    parser.add_argument("--ctc-warmup-steps", type=int, default=3000,
652	                        help="Steps to warm up CTC loss from 0 to target weight")
653	    parser.add_argument("--no-lang-embed", action="store_true", default=False,
654	                        help="Disable language embedding conditioning")
655	    args = parser.parse_args()
656
657	    cfg = OmegaConf.load(args.config)
658
659	    with open_dict(cfg):
660	        if args.max_steps is not None:
661	            cfg.trainer.max_steps = args.max_steps
662	        if args.devices is not None:
663	            cfg.trainer.devices = args.devices
664	            cfg.trainer.strategy = "ddp_find_unused_parameters_true" if args.devices > 1 else "auto"
665	        if args.grad_accum is not None:
666	            cfg.trainer.accumulate_grad_batches = args.grad_accum
667	        if args.val_every is not None:
668	            cfg.trainer.val_check_interval = args.val_every
669	            cfg.exp_manager.checkpoint_callback_params.every_n_train_steps = args.val_every
670	        if args.log_every is not None:
671	            cfg.trainer.log_every_n_steps = args.log_every
672	        if args.smoke and args.val_every is None:
673	            cfg.trainer.val_check_interval = 1.0
674	        if args.resume_from_checkpoint is not None:
675	            cfg.exp_manager.resume_from_checkpoint = args.resume_from_checkpoint
676	            cfg.exp_manager.resume_if_exists = True
677
678	        cfg.model.train_ds.is_tarred = False
679	        for key in ("tarred_audio_filepaths", "shuffle_n", "input_cfg"):
680	            if key in cfg.model.train_ds:
681	                del cfg.model.train_ds[key]
682	        cfg.model.train_ds.manifest_filepath = args.val_manifest
683	        cfg.model.validation_ds.manifest_filepath = args.val_manifest
684
685	    grad_accum = cfg.trainer.get("accumulate_grad_batches", 1)
686	    devices = cfg.trainer.get("devices", 1)
687	    max_dur = cfg.model.train_ds.get("max_duration", 20.0)
688	    min_dur = cfg.model.train_ds.get("min_duration", 0.5)
689	    num_workers = args.num_workers or cfg.model.train_ds.get("num_workers", 8)
690
691	    rank = int(os.environ.get("LOCAL_RANK", 0))
692	    world_size = devices if isinstance(devices, int) else 1
693
694	    # -----------------------------------------------------------------------
695	    # OOM-safe VRAM budget
696	    # -----------------------------------------------------------------------
697	    vram_budget = compute_vram_safe_batch_params(
698	        gpu_memory_gb=140.0,
699	        max_audio_dur=max_dur,


1023 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py


Let me get the rest of the main function:


700	        grad_accum=grad_accum,
701	    )
702	    max_batch_dur = args.max_batch_dur
703	    max_batch_size = args.max_batch_size
704
705	    if rank == 0:
706	        print("=" * 60)
707	        print("  Maya ASR — Production Training v2 (TDT + DiffLR)")
708	        print("=" * 60)
709	        print()
710	        print(f"Config:           {args.config}")
711	        print(f"Max steps:        {cfg.trainer.max_steps}")
712	        print(f"Devices:          {devices}")
713	        print(f"Grad accum:       {grad_accum}")
714	        print(f"Precision:        {cfg.trainer.get('precision', 'bf16-mixed')}")
715	        print(f"Train parquet:    {args.train_parquet}")
716	        print(f"Val manifest:     {args.val_manifest}")
717	        print()
718	        print("v2 Optimizations:")
719	        print(f"  Loss:           TDT (durations=[0,1,2,3,4])")
720	        print(f"  Encoder LR:     {args.encoder_lr}")
721	        print(f"  Head LR:        {args.head_lr}")
722	        print(f"  Encoder freeze: {args.freeze_encoder_steps} steps")
723	        print(f"  CTC warmup:     0→target over {args.ctc_warmup_steps} steps")
724	        print(f"  Lang embed:     {'yes' if not args.no_lang_embed else 'no'}")
725	        print()
726	        print("VRAM Budget (OOM-safe):")
727	        print(f"  GPU memory:     {vram_budget['gpu_memory_gb']:.0f} GB")
728	        print(f"  Model (bf16):   {vram_budget['model_params_gb']:.1f} GB")
729	        print(f"  Optimizer:      {vram_budget['optimizer_gb']:.1f} GB")
730	        print(f"  Activation:     {vram_budget['activation_budget_gb']:.1f} GB (75% safety)")
731	        print(f"  Joint/sample:   {vram_budget['joint_per_sample_gb']:.3f} GB")
732	        print(f"  Total/sample:   {vram_budget['total_per_sample_gb']:.3f} GB")
733	        print(f"  → max_batch_size: {max_batch_size}")
734	        print(f"  → max_batch_dur:  {max_batch_dur:.0f}s")
735	        print()
736	        print("Data strategy:")
737	        print(f"  Max batch dur:  {max_batch_dur}s")
738	        print(f"  Max batch size: {max_batch_size}")
739	        print(f"  Temperature:    {args.temperature}")
740	        print(f"  Duration range: [{min_dur}, {max_dur}]s")
741	        print(f"  Num workers:    {num_workers}")
742	        print()
743
744	    # -----------------------------------------------------------------------
745	    # Build trainer and model
746	    # -----------------------------------------------------------------------
747	    trainer = pl.Trainer(**cfg.trainer, logger=False, use_distributed_sampler=False)
748	    exp_manager(trainer, cfg.get("exp_manager", None))
749
750	    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
751	    params = sum(p.numel() for p in model.parameters())
752	    if rank == 0:
753	        print(f"Model parameters: {params:,}")
754
755	    # Resume weights from checkpoint (without optimizer state)
756	    if args.resume_weights_only and args.resume_from_checkpoint:
757	        ckpt_path = args.resume_from_checkpoint
758	        if rank == 0:
759	            print(f"\nResuming weights only from: {ckpt_path}")
760	        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
761	        model.load_state_dict(ckpt["state_dict"], strict=False)
762	        resume_step = ckpt.get("global_step", 0)
763	        if rank == 0:
764	            print(f"  Loaded model weights from step {resume_step}")
765	            print(f"  Optimizer state NOT restored (fresh optimizer with DiffLR)")
766	        # Clear resume flag so Lightning doesn't also try to restore
767	        args.resume_from_checkpoint = None
768	        with open_dict(cfg):
769	            cfg.exp_manager.resume_from_checkpoint = None
770	            cfg.exp_manager.resume_if_exists = False
771	        del ckpt
772	        print()
773
774	    # Load pretrained encoder weights
775	    if args.pretrained_encoder is not None and not args.resume_weights_only:
776	        if rank == 0:
777	            print(f"\nLoading pretrained encoder from: {args.pretrained_encoder}")
778	        pretrained_state = torch.load(args.pretrained_encoder, map_location="cpu", weights_only=True)
779	        model_state = model.state_dict()
780
781	        loaded, skipped = 0, 0
782	        for key, value in pretrained_state.items():
783	            if key in model_state and model_state[key].shape == value.shape:
784	                model_state[key] = value
785	                loaded += 1
786	            else:
787	                skipped += 1
788	                if rank == 0 and skipped <= 10:
789	                    expected = model_state[key].shape if key in model_state else "NOT FOUND"
790	                    print(f"  Skip: {key} pretrained={value.shape} model={expected}")
791
792	        model.load_state_dict(model_state)
793	        enc_params = sum(pretrained_state[k].numel() for k in pretrained_state if k in model_state and model_state[k].shape == pretrained_state[k].shape)
794	        if rank == 0:
795	            print(f"  Loaded {loaded} tensors ({enc_params/1e6:.0f}M params), skipped {skipped}")
796	            print(f"  Decoder + joint + CTC head: randomly initialized for {model.tokenizer.vocab_size} vocab")
797	        print()
798
799	    # -----------------------------------------------------------------------
800	    # Language conditioning
801	    # -----------------------------------------------------------------------
802	    if not args.no_lang_embed:
803	        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804	        model._lang_embed = lang_embed
805	        # Register as a proper submodule so it moves to GPU and is saved
806	        model.register_module('_lang_embed_module', lang_embed)
807	        patch_encoder_for_lang(model)
808	        patch_training_step(model)
809	        if rank == 0:
810	            lang_params = sum(p.numel() for p in lang_embed.parameters())
811	            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
812
813	    # -----------------------------------------------------------------------
814	    # Callbacks
815	    # -----------------------------------------------------------------------
816	    log_interval = args.log_every or cfg.trainer.get("log_every_n_steps", 100)
817	    throughput_cb = ThroughputCallback(
818	        log_every=log_interval,
819	        model_params=params,
820	        grad_accum=grad_accum,
821	        world_size=world_size,
822	    )
823	    trainer.callbacks.append(throughput_cb)
824
825	    freeze_cb = EncoderFreezeCallback(freeze_steps=args.freeze_encoder_steps)
826	    trainer.callbacks.append(freeze_cb)
827
828	    ctc_warmup_cb = CTCWarmupCallback(
829	        target_weight=cfg.model.aux_ctc.get("ctc_loss_weight", 0.3),
830	        warmup_steps=args.ctc_warmup_steps,
831	    )
832	    trainer.callbacks.append(ctc_warmup_cb)
833
834	    diff_lr_cb = DifferentialLRCallback(
835	        encoder_lr=args.encoder_lr,
836	        head_lr=args.head_lr,
837	    )
838	    trainer.callbacks.append(diff_lr_cb)
839
840	    # -----------------------------------------------------------------------
841	    # Load train dataset
842	    # -----------------------------------------------------------------------
843	    if rank == 0:
844	        print(f"Loading train parquet: {args.train_parquet}")
845	    train_dataset = NeMoTarOffsetDataset(
846	        manifest_parquet=args.train_parquet,
847	        tokenizer=model.tokenizer,
848	        sample_rate=cfg.model.train_ds.get("sample_rate", 16000),
849	        max_duration=max_dur,
850	        min_duration=min_dur,
851	    )
852	    if rank == 0:
853	        print(f"Train samples: {len(train_dataset):,}")
854	        print(f"Train audio:   {train_dataset.durations.sum() / 3600:.0f} hours")
855	        # Language distribution
856	        unique_langs, lang_counts = np.unique(train_dataset.languages, return_counts=True)
857	        print(f"Languages:     {len(unique_langs)} ({', '.join(sorted(unique_langs))})")
858	        print()
859
860	    # Create production batch sampler
861	    # Pre-compute token lengths for VRAM-safe batching
862	    # Token length controls T_dec in the joint tensor — the KEY factor for OOM
863	    # Cached to .npy file so subsequent runs skip tokenization entirely
864	    cache_path = Path(args.train_parquet).with_suffix(".token_lengths.npy")
865	    if cache_path.exists():
866	        if rank == 0:
867	            print(f"Loading cached token lengths from {cache_path}")
868	        token_lengths = np.load(cache_path)
869	        if len(token_lengths) != len(train_dataset):
870	            if rank == 0:
871	                print(f"  Cache size mismatch ({len(token_lengths)} vs {len(train_dataset)}), recomputing...")
872	            cache_path.unlink()
873	            token_lengths = None
874	        else:
875	            if rank == 0:
876	                print(f"  Loaded {len(token_lengths):,} cached token lengths")
877
878	    if not cache_path.exists():
879	        if rank == 0:
880	            print("Pre-computing token lengths (parallel, 32 workers)...")
881	        import multiprocessing as mp
882
883	        tok_dir = cfg.model.tokenizer.dir
884	        tok_model = str(Path(tok_dir) / "tokenizer.model")
885
886	        transcripts = list(train_dataset.transcripts)
887	        n_workers = min(32, mp.cpu_count())
888	        chunk_size = (len(transcripts) + n_workers - 1) // n_workers
889	        chunks = [transcripts[i:i + chunk_size] for i in range(0, len(transcripts), chunk_size)]
890
891	        if rank == 0:
892	            t_start = time.time()
893
894	        with mp.Pool(n_workers) as pool:
895	            results = pool.starmap(
896	                _tokenize_chunk_worker,
897	                [(chunk, tok_model) for chunk in chunks],
898	            )
899


Now let me check the tar offset reader and look at config files: