Thoroughly explore the codebase at /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ to understand how to perform inference with a trained hybrid FastConformer RNNT/TDT model. I need to understand:

1. The model architecture - is it NeMo-based? What's "hybrid TDT"?
2. How to load a .ckpt checkpoint and run inference on audio files
3. Any existing inference/evaluation scripts
4. What dependencies are needed (NeMo version, etc.)
5. The config files for the model architecture
6. How language-guided / multilingual inference works (if applicable)

Look at: src/, scripts/, configs/, lf_asr/, tests/, and any inference/eval related files. Read the README.md and PLAN.md too. Be very thorough - read all relevant files.


I'll thoroughly explore the codebase to understand the hybrid FastConformer RNNT/TDT model and inference process. Let me start by examining the structure and key files.


Bash

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


Let me read the README and PLAN files first to get an overview:


     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→


total 144
drwxrwxr-x 10 ubuntu ubuntu  4096 Mar 27 05:10 .
drwxrwxr-x 18 ubuntu ubuntu  4096 Mar 27 02:56 ..
-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:10 .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  2 ubuntu ubuntu  4096 Mar 25 02:16 tools


Now let me explore the directory structure and key files:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f \( -name "*.py" -o -name "*.yaml" \) | head -50


No files found


Let me look at the scripts and configs directories more carefully:


total 208
drwxrwxr-x  2 ubuntu ubuntu  4096 Mar 27 05:10 .
drwxrwxr-x 10 ubuntu ubuntu  4096 Mar 27 05:10 ..
-rwxrwxr-x  1 ubuntu ubuntu  4061 Mar 27 05:10 auto_resume_train.sh
-rw-rw-r--  1 ubuntu ubuntu  6837 Mar 25 02:16 auto_upload_checkpoints.py
-rw-rw-r--  1 ubuntu ubuntu  9476 Mar 25 02:16 build_manifest.py
-rw-rw-r--  1 ubuntu ubuntu  5294 Mar 25 02:16 build_tokenizer.py
-rw-rw-r--  1 ubuntu ubuntu 11216 Mar 25 02:16 check_prod_readiness.py
-rwxrwxr-x  1 ubuntu ubuntu  6953 Mar 25 02:16 drill_resume_from_r2.sh
-rw-rw-r--  1 ubuntu ubuntu  5005 Mar 25 02:16 extract_smoke_audio.py
-rw-rw-r--  1 ubuntu ubuntu  4579 Mar 25 02:16 generate_input_cfg.py
-rw-rw-r--  1 ubuntu ubuntu     0 Mar 25 02:16 .gitkeep
-rwxrwxr-x  1 ubuntu ubuntu  1825 Mar 25 02:16 gpu_idle_monitor.sh
-rwxrwxr-x  1 ubuntu ubuntu  5205 Mar 25 02:16 prepare_stage1_prod.sh
-rw-rw-r--  1 ubuntu ubuntu   892 Mar 25 02:16 print_runtime_info.py
-rwxrwxr-x  1 ubuntu ubuntu  2511 Mar 25 02:16 profile_smoke_run.sh
-rwxrwxr-x  1 ubuntu ubuntu  2802 Mar 25 21:14 push_milestone_checkpoint.sh
-rw-rw-r--  1 ubuntu ubuntu  6331 Mar 25 02:16 restore_checkpoint.py
-rwxrwxr-x  1 ubuntu ubuntu  3853 Mar 25 02:16 run_stage1_prod.sh
-rwxrwxr-x  1 ubuntu ubuntu 14914 Mar 25 02:16 run_stage1_smoke.sh
-rw-rw-r--  1 ubuntu ubuntu  5884 Mar 25 02:16 split_manifest.py
-rw-rw-r--  1 ubuntu ubuntu 44190 Mar 27 05:10 train_prod.py
-rw-rw-r--  1 ubuntu ubuntu  3652 Mar 25 02:16 train_smoke.py
-rw-rw-r--  1 ubuntu ubuntu  7238 Mar 25 02:16 upload_checkpoint.py
-rw-rw-r--  1 ubuntu ubuntu  4797 Mar 25 02:16 validate_manifest.py
-rw-rw-r--  1 ubuntu ubuntu  4214 Mar 25 02:16 validate_r2_pipeline.py


total 16
drwxrwxr-x  4 ubuntu ubuntu 4096 Mar 25 02:16 .
drwxrwxr-x 10 ubuntu ubuntu 4096 Mar 27 05:10 ..
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 25 02:16 data
-rw-rw-r--  1 ubuntu ubuntu    0 Mar 25 02:16 .gitkeep
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 27 05:10 train


total 16
drwxrwxr-x 2 ubuntu ubuntu 4096 Mar 25 02:16 .
drwxrwxr-x 3 ubuntu ubuntu 4096 Mar 25 02:16 ..
-rw-rw-r-- 1 ubuntu ubuntu  856 Mar 25 02:16 config.py
-rw-rw-r-- 1 ubuntu ubuntu   93 Mar 25 02:16 __init__.py


total 32
drwxrwxr-x 2 ubuntu ubuntu 4096 Mar 27 05:10 .
drwxrwxr-x 4 ubuntu ubuntu 4096 Mar 25 02:16 ..
-rw-rw-r-- 1 ubuntu ubuntu 4583 Mar 27 05:10 stage1_prod_8xh200.yaml
-rw-rw-r-- 1 ubuntu ubuntu 4272 Mar 25 02:16 stage1_prod_smoke_tarred.yaml
-rw-rw-r-- 1 ubuntu ubuntu 2987 Mar 25 02:16 stage1_prod_smoke_test.yaml
-rw-rw-r-- 1 ubuntu ubuntu 3296 Mar 25 02:16 stage1_smoke.yaml


total 1260
drwxrwxr-x 2 ubuntu ubuntu    4096 Mar 25 02:16 .
drwxrwxr-x 4 ubuntu ubuntu    4096 Mar 25 02:16 ..
-rw-rw-r-- 1 ubuntu ubuntu     436 Mar 25 02:16 input_cfg_metadata.json
-rw-rw-r-- 1 ubuntu ubuntu     739 Mar 25 02:16 local_smoke.yaml
-rw-rw-r-- 1 ubuntu ubuntu 1265062 Mar 25 02:16 stage1_prod_input_cfg_v2.yaml
-rw-rw-r-- 1 ubuntu ubuntu    3491 Mar 25 02:16 stage1_prod_input_cfg.yaml
-rw-rw-r-- 1 ubuntu ubuntu     579 Mar 25 02:16 stage1_smoke_input_cfg.yaml


Now let me read the key training and config files:


     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:


     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→


     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 continue reading the train_prod.py script to understand inference:


   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:
   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


Let me check the lf_asr modules for data loading and see if there are any inference examples:


     1→"""Tar offset reader: zero-extraction audio loading via positional I/O.
     2→
     3→Reads FLAC bytes directly from tar files using os.pread() with byte offsets.
     4→No tar extraction, no temp files, no runtime resampling.
     5→
     6→Usage:
     7→    reader = TarOffsetReader(max_fd_cache=256)
     8→    waveform, sr = reader.read(tar_path, offset, nbytes)
     9→    assert sr == 16000
    10→"""
    11→
    12→import io
    13→import os
    14→import threading
    15→from collections import OrderedDict
    16→from pathlib import Path
    17→
    18→import soundfile
    19→
    20→
    21→class TarOffsetReader:
    22→    """Read audio from tar files via byte offsets using pread().
    23→
    24→    Thread-safe with per-thread fd cache.
    25→    """
    26→
    27→    def __init__(self, max_fd_cache: int = 256, max_retries: int = 2):
    28→        self.max_fd_cache = max_fd_cache
    29→        self.max_retries = max_retries
    30→        self._local = threading.local()
    31→
    32→    def _get_cache(self) -> OrderedDict:
    33→        if not hasattr(self._local, "fd_cache"):
    34→            self._local.fd_cache = OrderedDict()
    35→        return self._local.fd_cache
    36→
    37→    def _get_fd(self, tar_path: str) -> int:
    38→        cache = self._get_cache()
    39→        if tar_path in cache:
    40→            cache.move_to_end(tar_path)
    41→            return cache[tar_path]
    42→        # Evict oldest if cache full
    43→        while len(cache) >= self.max_fd_cache:
    44→            _, old_fd = cache.popitem(last=False)
    45→            os.close(old_fd)
    46→        fd = os.open(tar_path, os.O_RDONLY)
    47→        cache[tar_path] = fd
    48→        return fd
    49→
    50→    def read(self, tar_path: str, offset: int, nbytes: int) -> tuple:
    51→        """Read and decode audio from tar at given offset.
    52→
    53→        Returns: (waveform_1d_numpy, sample_rate)
    54→        Raises: RuntimeError on decode failure after retries.
    55→        """
    56→        last_err = None
    57→        for attempt in range(self.max_retries + 1):
    58→            try:
    59→                fd = self._get_fd(tar_path)
    60→                raw = os.pread(fd, nbytes, offset)
    61→                if len(raw) != nbytes:
    62→                    raise IOError(
    63→                        f"Short read: got {len(raw)} expected {nbytes} "
    64→                        f"at offset {offset} in {tar_path}"
    65→                    )
    66→                data, sr = soundfile.read(io.BytesIO(raw))
    67→                if sr != 16000:
    68→                    raise ValueError(f"Expected 16kHz, got {sr}Hz in {tar_path}")
    69→                if len(data.shape) > 1:
    70→                    data = data[:, 0]  # Take first channel if stereo
    71→                return data, sr
    72→            except Exception as e:
    73→                last_err = e
    74→                # Invalidate fd on error (might be stale)
    75→                cache = self._get_cache()
    76→                if tar_path in cache:
    77→                    try:
    78→                        os.close(cache.pop(tar_path))
    79→                    except OSError:
    80→                        pass
    81→        raise RuntimeError(
    82→            f"Failed to read after {self.max_retries + 1} attempts: "
    83→            f"{tar_path}@{offset}+{nbytes}: {last_err}"
    84→        )
    85→
    86→    def close(self):
    87→        """Close all cached file descriptors."""
    88→        cache = self._get_cache()
    89→        for fd in cache.values():
    90→            try:
    91→                os.close(fd)
    92→            except OSError:
    93→                pass
    94→        cache.clear()
    95→
    96→    def __del__(self):
    97→        try:
    98→            self.close()
    99→        except Exception:
   100→            pass


     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→        """


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f -name "*.py" | xargs grep -l "transcribe\|inference\|eval" 2>/dev/null | head -20


(Bash completed with no output)


Let me search for any jupyter notebooks or example scripts that show inference:


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -type f \( -name "*.ipynb" -o -name "*example*" -o -name "*demo*" \)


grep -r "restore_from\|transcribe" /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 --include="*.py" | head -30


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 -name "pyproject.toml" -o -name "requirements.txt" -o -name "setup.py"


     1→[build-system]
     2→requires = ["setuptools>=68.0", "wheel"]
     3→build-backend = "setuptools.build_meta"
     4→
     5→[project]
     6→name = "maya-asr"
     7→version = "0.1.0"
     8→description = "Multilingual Indic ASR: FastConformer encoder + Gemma decoder"
     9→readme = "README.md"
    10→requires-python = ">=3.10"
    11→license = {text = "MIT"}
    12→
    13→dependencies = [
    14→    "pyarrow>=14.0",
    15→    "pandas>=2.0",
    16→    "pyyaml>=6.0",
    17→    "omegaconf>=2.3",
    18→    "tqdm>=4.65",
    19→    "soundfile>=0.12",
    20→    "sentencepiece>=0.2",
    21→    "python-dotenv>=1.0",
    22→    "boto3>=1.28",
    23→]
    24→
    25→[project.optional-dependencies]
    26→train = [
    27→    "nemo_toolkit[asr]",
    28→    "transformers>=4.40",
    29→    "accelerate>=0.30",
    30→    "deepspeed>=0.14",
    31→]
    32→dev = [
    33→    "pytest>=8.0",
    34→    "ruff>=0.4",
    35→]
    36→
    37→[tool.setuptools.packages.find]
    38→where = ["src", "."]
    39→include = ["maya_asr*", "lf_asr*"]
    40→
    41→[tool.ruff]
    42→target-version = "py310"
    43→line-length = 100
    44→
    45→[tool.ruff.lint]
    46→select = ["E", "F", "I", "W"]
    47→
    48→[tool.pytest.ini_options]
    49→testpaths = ["tests"]
    50→markers = [
    51→    "integration: tests that require /root/sft_data (deselect with -m 'not integration')",
    52→]
    53→


grep -r "EncDec" /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 --include="*.py" | head -20


/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_smoke.py:"""Run training with NeMo EncDecHybridRNNTCTCBPEModel.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_smoke.py:    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py:Model: NeMo EncDecHybridRNNTCTCBPEModel (FastConformer + TDT + aux CTC).
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py:    matching NeMo's EncDecHybridRNNTCTCBPEModel expectations.
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py:    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/tools/calibrate_batch_size.py:    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/tools/calibrate_batch_size_full.py:    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)


Let me read the train_prod.py more to understand model restoration and look at the test files:


   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,
   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,