/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 can you explain the model structure or the changes made to the model training. you can also check the available ckpts to open and get me a summary of the model and it's layers and params. if more details needed or mdoel is not local, use this, download and add to summary. ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/maya_asr_stage1_tdt_v2--val_loss=164.4795-epoch=0-last.ckpt i'm more specifically interested in the added layers and language as input. check the code training as well.


Let me explore the directory structure and training code, then examine any available checkpoints.


Agent

Explore maya-asr model structure

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

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

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

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


Perfect! Now I have comprehensive information. Let me create a detailed summary report:

## Comprehensive Analysis: Maya ASR Hybrid FastConformer-RNNT-Stage1

I've completed a thorough exploration of the Maya ASR project. Here's the complete breakdown:

### 1. DIRECTORY STRUCTURE & PROJECT OVERVIEW

**Root Directory:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/`

**Key Directories:**
- `configs/` — Training and data YAML configurations
- `lf_asr/data/` — Custom data loading and sampling modules
- `src/maya_asr/` — Core library with configuration constants
- `scripts/` — Training, data pipeline, checkpoint management scripts
- `tools/` — Data preparation and validation utilities
- `tests/` — Unit and integration tests
- `tokenizers/stage1_prod_bpe/` — 32K BPE SentencePiece tokenizer (1.5MB model + vocab)

---

### 2. ARCHITECTURE OVERVIEW

**Model Type:** `EncDecHybridRNNTCTCBPEModel` (NeMo native)

**Architecture Diagram:**
```
Audio (16 kHz)
    ↓
Mel Spectrogram (80-dim, 16000 Hz sample rate)
    ↓
FastConformer Encoder (42 layers, 1024-dim, 8x subsampling)
    ├─ Pretrained from Parakeet TDT 1.1B encoder
    └─ + Language Embedding (12 languages → 1024-dim additive bias)
    ↓
    ├─────────────────┬─────────────────┐
    ↓                 ↓                 ↓
TDT Decoder      Joint Network      CTC Head
(2-layer RNN)    (640-dim hidden)    (Linear projection)
(640-dim)        + 5 duration heads  (~33M params)
(~7M params)     (~42M params)
    ↓                 ↓
    └─────────┬───────┘
              ↓
         [Primary + Auxiliary Loss]
         TDT (transducer) + CTC
```

**Total Parameters:** ~1,144M (1.1B)

**Component Breakdown:**
- **Encoder:** 42-layer FastConformer, 1024-dim hidden, 8x subsampling (dw_striding), ReLU activation
  - ~1,062M parameters
  - Pretrained from Parakeet TDT 1.1B (frozen first 5K steps, then fine-tuned at 2e-5 LR)

- **Decoder:** 2-layer RNN with prediction network
  - pred_hidden=640, pred_rnn_layers=2
  - ~7M parameters
  - Randomly initialized, trained at 1e-3 LR

- **Joint Network:** RNNTJoint with extra duration outputs
  - joint_hidden=640
  - num_extra_outputs=5 (TDT duration heads for durations [0,1,2,3,4])
  - ~42M parameters
  - Fused computation with fused_batch_size=2 to cap VRAM

- **CTC Head:** Linear projection over encoder output
  - ~33M parameters
  - Auxiliary loss weighted 0.3 (warmed up from 0→0.3 over 3K steps)

- **Language Embedding (NEW):** Per-language additive bias
  - 12 languages → 1024-dim embeddings
  - ~12K parameters
  - Initialized with std=0.02 to avoid disrupting pretrained weights

---

### 3. LANGUAGE CONDITIONING & INPUT

**Language as Model Input:**

1. **Language Embedding Module** (`scripts/train_prod.py`, lines 56-73):
```python
class LanguageEmbedding(torch.nn.Module):
    """Projects a language ID into a per-frame additive bias for the encoder."""
    def __init__(self, num_languages: int, embed_dim: int):
        super().__init__()
        self.embed = torch.nn.Embedding(num_languages, embed_dim)
        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)

    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
        # lang_ids: (B,) → returns (B, embed_dim, 1) for broadcasting
        return self.embed(lang_ids).unsqueeze(2)
```

2. **Language ID Mapping** (lines 49-53):
```python
LANG_TO_ID = {
    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
}
```

3. **Data Loading with Language IDs** (`scripts/train_prod.py`, lines 451-455):
```python
self.lang_ids = np.array(
    [LANG_TO_ID.get(str(lang).strip().lower(), LANG_TO_ID["en"])
     for lang in self.languages],
    dtype=np.int64,
)
```

4. **Application in Encoder** (`scripts/train_prod.py`, lines 582-599):
```python
def patch_encoder_for_lang(model):
    """Patch encoder forward to inject language IDs stored on model."""
    original_encoder_forward = model.encoder.forward

    def encoder_forward_with_lang(audio_signal, length):
        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal,
                                                        length=length)
        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
                and model._current_lang_ids is not None):
            lang_ids = model._current_lang_ids
            if lang_ids.shape[0] == encoded.shape[0]:
                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
                encoded = encoded + lang_bias  # (B, D, T) + (B, D, 1)
        return encoded, encoded_len

    model.encoder.forward = encoder_forward_with_lang
```

**Supported Languages:** Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en)

---

### 4. CUSTOM LAYERS & MODIFICATIONS

**LanguageEmbedding Module** (NEW custom layer):
- **File:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py` (lines 56-73)
- **Purpose:** Injects language signal via per-frame additive bias
- **Architecture:** Embedding layer (12 → 1024-dim)
- **Initialization:** Normal(mean=0, std=0.02) to preserve pretrained encoder
- **Integration:** Hooked into encoder.forward via patch

**TDT (Temporal Duration Transducer) Modifications:**
- **Joint Network:** Added `num_extra_outputs: 5` for duration predictions (0, 1, 2, 3, 4 frames)
- **Loss:** Changed from RNNT to TDT loss with:
  - durations: [0, 1, 2, 3, 4]
  - sigma: 0.02 (softness parameter)
  - omega: 0.1 (regularization)
  - fastemit_lambda: 0.0

**Fused Joint Computation:**
- Uses `fuse_loss_wer: true` with `fused_batch_size: 2` to split joint tensor computation into sub-batches

---

### 5. TRAINING SCRIPTS & OPTIMIZATIONS

**Main Training Script:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/train_prod.py` (1,023 lines)

**Key Callbacks & Optimizations:**

1. **EncoderFreezeCallback** (lines 162-201):
   - Freezes encoder for first 5,000 optimizer steps
   - Allows decoder/joint/CTC heads to warm up without corrupting pretrained weights
   - Unfreezes encoder at 2e-5 LR thereafter

2. **CTCWarmupCallback** (lines 245-263):
   - Ramps CTC loss weight from 0 → 0.3 over first 3,000 steps
   - Prevents random CTC gradients from disrupting early training

3. **DifferentialLRCallback** (lines 266-332):
   - **Encoder:** 2e-5 LR (pretrained, low learning rate)
   - **Decoder/Joint/CTC:** 1e-3 LR (random init, fast adaptation)
   - Splits optimizer param groups and syncs with LR scheduler

4. **ThroughputCallback** (lines 79-159):
   - Logs VRAM, samples/s, audio-s/s, RTF/GPU, MFU every N steps
   - Per-param-group LR logging to wandb/tensorboard

5. **SamplerStateCallback** (lines 204-242):
   - Saves sampler state for resumption without reprocessing data

**Training Data Strategy:**

- **ProductionBatchSampler** (`/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/lf_asr/data/production_sampler.py`):
  - **Duration Bucketing:** Groups samples into 5 buckets ([0-4], [4-8], [8-12], [12-16], [16-20]s)
  - **Language Temperature Rebalancing (T=0.3):** Upweights underrepresented languages
  - **Triple-Cap OOM Protection:**
    1. Per-bucket BS caps (24→6 samples depending on bucket)
    2. max_batch_duration (120s audio per micro-batch)
    3. max_tokens_in_batch (400 tokens, controls T_dec in joint tensor)
  - **Compute-Cost DDP Sharding:** Distributes batches balanced by `len(batch) * max_duration_in_batch`
  - **Vectorized with numpy:** ~30s to iterate 77M+ samples (vs 8min with loops)

---

### 6. CONFIGURATION FILES

**Training Config:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/train/stage1_prod_8xh200.yaml` (164 lines)

**Key Settings:**
```yaml
trainer:
  devices: 8
  precision: bf16-mixed
  max_steps: 200000
  val_check_interval: 10000
  accumulate_grad_batches: 4
  gradient_clip_val: 1.0

model:
  encoder:
    n_layers: 42
    d_model: 1024
    subsampling: dw_striding
    subsampling_factor: 8
    conv_kernel_size: 9

  decoder:
    pred_hidden: 640
    pred_rnn_layers: 2

  joint:
    num_extra_outputs: 5  # TDT duration heads
    fuse_loss_wer: true
    fused_batch_size: 2
    joint_hidden: 640

  loss:
    loss_name: tdt
    tdt_kwargs:
      durations: [0, 1, 2, 3, 4]
      sigma: 0.02
      omega: 0.1

  aux_ctc:
    ctc_loss_weight: 0.3  # warmed up from 0→0.3

  tokenizer:
    dir: tokenizers/stage1_prod_bpe
    type: bpe

  optim:
    name: adamw
    betas: [0.9, 0.98]
    weight_decay: 0.0001
    sched:
      name: CosineAnnealing
      warmup_steps: 2000
      min_lr: 0.00001
```

**Data Configs:**
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg.yaml` (3.5K)
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/stage1_prod_input_cfg_v2.yaml` (1.3M, per-language dataset configs)
- `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/configs/data/local_smoke.yaml` (smoke test setup)

---

### 7. DATA PIPELINE SCRIPTS

**Key Scripts:**

1. **`scripts/build_manifest.py`** — Builds NeMo JSONL manifests from shard metadata
   - Reads metadata.parquet + audio_index.parquet
   - Outputs: `audio_filepath, text, duration, lang, taskname, source_lang, target_lang`
   - Supports hive-partitioned and flat directory layouts

2. **`scripts/split_manifest.py`** — Stratified train/val splits

3. **`scripts/build_tokenizer.py`** — Trains 32K SentencePiece BPE tokenizer
   - Quality-filtered text from all 12 languages
   - Located at: `tokenizers/stage1_prod_bpe/`
   - tokenizer.model: 962KB, vocab.txt: 472KB

4. **`scripts/generate_input_cfg.py`** — Generates NeMo input_cfg YAML

5. **Data Reading:** `lf_asr/data/tar_offset_reader.py`
   - Zero-extraction tar loading via `os.pread()` with byte offsets
   - Thread-safe fd cache (max 256 open tar files)
   - Returns (waveform, 16kHz sample rate)

6. **Dataset:** `lf_asr/data/tar_offset_dataset.py` + custom `NeMoTarOffsetDataset`
   - Wraps TarOffsetReader for NeMo training
   - Returns: (signal, signal_len, tokens, token_len, lang_id)
   - Pre-computes language IDs in constructor

---

### 8. TOKENIZER

**File Location:** `/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/tokenizers/stage1_prod_bpe/`

**Details:**
- Type: BPE (SentencePiece)
- Vocabulary Size: 32K (32,769 including special tokens)
- Coverage: 11 Indic languages + English
- Quality filtering applied during training
- Model: `tokenizer.model` (962KB)
- Vocabulary: `vocab.txt` (472KB)

**Usage:**
```python
from sentencepiece import SentencePieceProcessor
sp = SentencePieceProcessor()
sp.Load('tokenizers/stage1_prod_bpe/tokenizer.model')
token_ids = sp.EncodeAsIds(text)
```

---

### 9. HYBRID TDT ARCHITECTURE MODIFICATIONS

**TDT Loss Configuration** (`configs/train/stage1_prod_8xh200.yaml`, lines 85-91):
```yaml
loss:
  loss_name: tdt
  tdt_kwargs:
    durations: [0, 1, 2, 3, 4]
    sigma: 0.02
    omega: 0.1
    fastemit_lambda: 0.0
    clamp: -1.0
```

**Joint Network with Duration Heads** (line 74):
```yaml
joint:
  num_extra_outputs: 5  # Predicts duration 0, 1, 2, 3, or 4
```

**Why TDT over RNNT?**
- Matches Parakeet TDT 1.1B pretraining objective
- Allows skipping multiple frames per emission (more efficient)
- Better speed/accuracy tradeoff for multilingual data

---

### 10. PRETRAINED ENCODER & CHECKPOINT HANDLING

**Pretrained Encoder:**
- Parakeet TDT 1.1B (42-layer FastConformer, English pretrained)
- Loaded via `--pretrained-encoder` argument in train_prod.py
- Expected path: `pretrained/parakeet_tdt_1.1b_encoder.pt` (not present locally)

**Checkpoint Scripts:**
- `scripts/restore_checkpoint.py` — Restores from Cloudflare R2
- `scripts/upload_checkpoint.py` — Uploads to R2
- `scripts/auto_upload_checkpoints.py` — Automatic checkpoint watcher

**Checkpoint Format:**
- `.ckpt` (PyTorch Lightning checkpoint with optimizer state)
- `.nemo` (NeMo native format with config + weights)
- Saved every 5K steps with `save_top_k: 3`

---

### 11. PLANNED STAGE 2: GEMMA INTEGRATION

**From README.md & PLAN.md:**

**Stage 2: Gemma LLM Decoder (planned)** — Replaces RNN decoder with Gemma 3 1B for improved language modeling

**Architecture (Stage 2):**
```
FastConformer Encoder (1024-dim)
    ↓
Projection Layer (Conv1D stride=2 → Linear 1024→1536)
    ↓
Gemma 3 1B Decoder (1536-dim hidden, 26 layers)
```

**Gemma Details:**
- Model: `google/gemma-3-1b-pt` (1B params, 1536-dim hidden, 26 layers)
- Tokenizer: 262K SentencePiece with 4K+ Indic tokens
- Task: Transcription only (no translation)
- Language tokens: `<|hi|>`, `<|bn|>`, ..., `<|en|>`
- Task token: `<|transcribe|>`

**Training Plan:**
- Stage 2a: Train projection only (encoder + LLM frozen) — 20K steps
- Stage 2b: Full joint training (projection + LLM, encoder frozen or LoRA) — 100K-200K steps
- Stage 3: Quality fine-tuning on high-quality subset — 10K-20K steps

---

### 12. KEY FILES SUMMARY

**Core Training:**
- `scripts/train_prod.py` (1,023 lines) — Main training with TDT, DiffLR, language embedding
- `configs/train/stage1_prod_8xh200.yaml` (164 lines) — 8×H200 training config
- `lf_asr/data/production_sampler.py` — Duration bucketing + language rebalancing
- `lf_asr/data/tar_offset_reader.py` — Zero-extraction tar I/O

**Data Pipeline:**
- `scripts/build_manifest.py` — Build manifests from shards
- `scripts/split_manifest.py` — Stratified splits
- `scripts/build_tokenizer.py` — Train BPE tokenizer
- `tools/phase3_*` — Phase 3 processing (tar offset building, splits, validation)

**Configuration & Constants:**
- `src/maya_asr/config.py` — Central config (languages, paths, dims)
- `configs/data/stage1_prod_input_cfg*.yaml` — Dataset configs
- `configs/data/local_smoke.yaml` — Smoke test config

**Documentation:**
- `README.md` (10.8K) — Architecture, training, inference overview
- `PLAN.md` (22.3K) — Detailed Stage 2 planning (Gemma integration)
- `docs/phase2_runbook.md` — Data processing runbook

**Tokenizer:**
- `tokenizers/stage1_prod_bpe/tokenizer.model` — SentencePiece model
- `tokenizers/stage1_prod_bpe/vocab.txt` — 32K vocabulary

---

### 13. HARDWARE REQUIREMENTS & OPTIMIZATION

**Production Training (8×H200):**
- GPUs: 8× NVIDIA H200 (140GB HBM3)
- Precision: BF16 mixed
- Max steps: 200K (~5-7 days)
- Effective batch size: 8 GPUs × micro-batch × 4 grad accum

**VRAM Budget (per H200):**
- Model weights (bf16): ~4.5GB
- Optimizer states (AdamW fp32): ~13.5GB
- Activations + joint tensor: ~75GB budget (75% safety margin)
- Joint tensor dominates: `B × T_enc × T_dec × (V + 1 + durations)`

**OOM Protection (Triple-Cap):**
1. Per-bucket BS caps (24→6 samples)
2. max_batch_duration (120s audio)
3. max_tokens_in_batch (400 tokens) — **CRITICAL** for controlling T_dec

---

### 14. SUPPORTED LANGUAGES

12 total: Hindi (hi), Bengali (bn), Tamil (ta), Telugu (te), Marathi (mr), Gujarati (gu), Kannada (kn), Malayalam (ml), Punjabi (pa), Odia (or), Assamese (as), English (en)

---

### 15. KEY ARGUMENTS & TRAINING COMMAND

```bash
python3 scripts/train_prod.py \
  --config configs/train/stage1_prod_8xh200.yaml \
  --train-parquet artifacts/phase3/production_train_final.parquet \
  --val-manifest data/manifests/stage1_prod_val_v2.jsonl \
  --pretrained-encoder pretrained/parakeet_tdt_1.1b_encoder.pt \
  --encoder-lr 2e-5 \
  --head-lr 1e-3 \
  --freeze-encoder-steps 5000 \
  --ctc-warmup-steps 3000 \
  --max-batch-dur 120.0 \
  --max-tokens-in-batch 400 \
  --temperature 0.3
```

**Key Arguments:**
- `--encoder-lr`: 2e-5 (preserve pretrained features)
- `--head-lr`: 1e-3 (fast adaptation for new vocab)
- `--freeze-encoder-steps`: 5K (warm up heads before encoder adaptation)
- `--ctc-warmup-steps`: 3K (ramp aux loss gradually)
- `--temperature`: 0.3 (aggressive language rebalancing)
- `--no-lang-embed`: Disable language conditioning (default: enabled)

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

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

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

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


Now let me check if the checkpoint exists locally and try to inspect it for layer details.


Bash

ls -lah /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/ 2>/dev/null || echo "CHECKPOINT_NOT_LOCAL"

Check if checkpoint exists locally


find /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ -name "*.ckpt" -o -name "*.pt" -o -name "*.nemo" 2>/dev/null | head -20

Find any local checkpoints


(Bash completed with no output)


No checkpoints locally. Let me download the specified one and inspect it, and also read the key training code sections.


ls /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/ 2>/dev/null || echo "NO_PTCHECKPOINTS_DIR"; ls /home/ubuntu/training/ptcheckpoints/ 2>/dev/null || echo "NO_ROOT_PTCHECKPOINTS"

Check checkpoint directories


NO_PTCHECKPOINTS_DIR
NO_ROOT_PTCHECKPOINTS


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

#!/usr/bin/env python3
"""Production training script with optimized data loading + TDT + differential LR.

Data strategy (ported from parakeet pipeline):
  - Duration bucketing: groups similar-length clips → minimal padding waste
  - Language temperature rebalancing (T=0.3): upweights low-resource languages
  - Max-batch-duration packing: caps audio per micro-batch → stable VRAM
  - Compute-cost DDP sharding: balances load across GPUs
  - Grad accum alignment: prevents DDP sync bugs

Optimizations v2:
  - TDT loss (matching pretrained Parakeet encoder objective)
  - Differential LR: encoder=2e-5, decoder/joint/CTC=1e-3
  - Encoder frozen for first 5K steps, then unfrozen
  - CTC loss warmup: 0→0.3 over first 3K steps
  - Language embedding conditioning
  - OOM-safe VRAM ceiling design

Audio I/O: TarOffsetReader (os.pread, zero extraction from existing tars).
Model: NeMo EncDecHybridRNNTCTCBPEModel (FastConformer + TDT + aux CTC).

Usage:
  # Full production run (8×H200)
  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml

  # Quick smoke test
  python3 scripts/train_prod.py --config configs/train/stage1_prod_8xh200.yaml \
    --max-steps 100 --devices 1 --log-every 10 --val-every 50
"""

import argparse
import os
import sys
import time
from pathlib import Path

import lightning.pytorch as pl
import nemo.collections.asr as nemo_asr
import numpy as np
import pyarrow.parquet as pq
import torch
from nemo.utils.exp_manager import exp_manager
from omegaconf import OmegaConf, open_dict

# ---------------------------------------------------------------------------
# Language embedding module
# ---------------------------------------------------------------------------
LANG_TO_ID = {
    "hi": 0, "bn": 1, "ta": 2, "te": 3, "mr": 4, "gu": 5,
    "kn": 6, "ml": 7, "pa": 8, "or": 9, "as": 10, "en": 11,
}
NUM_LANGUAGES = len(LANG_TO_ID)

class LanguageEmbedding(torch.nn.Module):
    """Projects a language ID into a per-frame additive bias for the encoder."""

    def __init__(self, num_languages: int, embed_dim: int):
        super().__init__()
        self.embed = torch.nn.Embedding(num_languages, embed_dim)
        # Initialize near-zero so it doesn't disrupt pretrained encoder features
        torch.nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)

    def forward(self, lang_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            lang_ids: (B,) integer tensor of language IDs
        Returns:
            (B, embed_dim, 1) — broadcast-addable to encoder output (B, D, T)
        """
        # Encoder output is (B, D, T), so we need (B, D, 1) for broadcasting
        return self.embed(lang_ids).unsqueeze(2)

# ---------------------------------------------------------------------------
# Callbacks
# ---------------------------------------------------------------------------
class ThroughputCallback(pl.Callback):
    """Logs VRAM, throughput (samples/s, audio-s/s), and MFU every log_every steps."""

    def __init__(self, log_every: int = 10, model_params: int = 0,
                 grad_accum: int = 1, world_size: int = 1):
        self.log_every = log_every
        self.model_params = model_params
        self.grad_accum = grad_accum
        self.world_size = world_size
        self._window_t0 = None
        self._window_samples = 0
        self._window_audio_s = 0.0
        self._window_micro_steps = 0
        # H200 specs: bf16 peak ~990 TFLOPS
        self._gpu_peak_tflops = 989.0

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
        if self._window_t0 is None:
            self._window_t0 = time.time()

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if batch is not None and len(batch) >= 2:
            sig_lens = batch[1]
            n_samples = sig_lens.shape[0]
            audio_s = sig_lens.float().sum().item() / 16000.0
            self._window_samples += n_samples
            self._window_audio_s += audio_s
        self._window_micro_steps += 1

        rank = int(os.environ.get("LOCAL_RANK", 0))
        global_step = trainer.global_step

        if rank == 0 and global_step > 0 and global_step % self.log_every == 0 and self._window_t0 is not None:
            elapsed = time.time() - self._window_t0
            if elapsed > 0:
                samples_per_s = self._window_samples / elapsed
                audio_s_per_s = self._window_audio_s / elapsed
                rtf_per_gpu = audio_s_per_s / max(self.world_size, 1)
                opt_steps_per_s = (self._window_micro_steps / max(self.grad_accum, 1)) / elapsed

                vram_alloc_gb = torch.cuda.memory_allocated() / 1e9
                vram_reserved_gb = torch.cuda.memory_reserved() / 1e9


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:

            self._window_t0 = time.time()
            self._window_samples = 0
            self._window_audio_s = 0.0
            self._window_micro_steps = 0

class EncoderFreezeCallback(pl.Callback):
    """Freezes encoder for the first N optimizer steps, then unfreezes."""

    def __init__(self, freeze_steps: int = 5000):
        self.freeze_steps = freeze_steps
        self._frozen = False
        self._unfrozen = False

    def on_train_start(self, trainer, pl_module):
        if self.freeze_steps > 0:
            self._freeze_encoder(pl_module)

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if self._unfrozen:
            return
        if trainer.global_step >= self.freeze_steps and self._frozen:
            self._unfreeze_encoder(pl_module)

    def _freeze_encoder(self, model):
        frozen = 0
        for name, param in model.named_parameters():
            if name.startswith("encoder."):
                param.requires_grad = False
                frozen += 1
        self._frozen = True
        rank = int(os.environ.get("LOCAL_RANK", 0))
        if rank == 0:
            print(f"  [EncoderFreeze] Froze {frozen} encoder params for first {self.freeze_steps} steps")

    def _unfreeze_encoder(self, model):
        unfrozen = 0
        for name, param in model.named_parameters():
            if name.startswith("encoder."):
                param.requires_grad = True
                unfrozen += 1
        self._frozen = False
        self._unfrozen = True
        rank = int(os.environ.get("LOCAL_RANK", 0))
        if rank == 0:
            print(f"  [EncoderFreeze] Unfroze {unfrozen} encoder params at step {model.trainer.global_step}")

class SamplerStateCallback(pl.Callback):
    """Saves sampler state (consumed batches) alongside checkpoints.

    On resume, the sampler can skip already-trained batches to avoid
    reprocessing the same data.
    """
    STATE_FILE = "sampler_state.json"

    def __init__(self, sampler, state_dir: str = "/tmp"):
        self.sampler = sampler
        self.state_dir = state_dir

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        # Save state every 1000 optimizer steps
        if trainer.global_step > 0 and trainer.global_step % 1000 == 0:
            self._save_state(trainer.global_step)

    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
        self._save_state(trainer.global_step)

    def _save_state(self, global_step):
        import json
        state = {
            "global_step": global_step,
            "epoch": self.sampler.epoch,
            "consumed_batches": getattr(self.sampler, "_consumed_batches", 0),
        }
        state_path = os.path.join(self.state_dir, self.STATE_FILE)
        with open(state_path, "w") as f:
            json.dump(state, f)

    @staticmethod
    def load_state(state_dir: str = "/tmp") -> dict | None:
        import json
        state_path = os.path.join(state_dir, SamplerStateCallback.STATE_FILE)
        if os.path.exists(state_path):
            with open(state_path) as f:
                return json.load(f)
        return None

class CTCWarmupCallback(pl.Callback):
    """Warms up CTC loss weight from 0 to target over N steps.

    Prevents random CTC gradients from corrupting pretrained encoder in early training.
    """

    def __init__(self, target_weight: float = 0.3, warmup_steps: int = 3000):
        self.target_weight = target_weight
        self.warmup_steps = warmup_steps

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
        step = trainer.global_step
        if step < self.warmup_steps:
            weight = self.target_weight * (step / self.warmup_steps)
        else:
            weight = self.target_weight

        if hasattr(pl_module, 'ctc_loss_weight'):
            pl_module.ctc_loss_weight = weight

class DifferentialLRCallback(pl.Callback):
    """Sets differential learning rates for encoder vs decoder/joint/CTC param groups.

    After NeMo constructs the optimizer (which uses a single LR), this callback
    modifies the param groups to use different LRs.

    Encoder (pretrained): low LR to preserve features
    Decoder/Joint/CTC (random init): high LR for fast adaptation
    """

    def __init__(self, encoder_lr: float = 2e-5, head_lr: float = 1e-3):
        self.encoder_lr = encoder_lr
        self.head_lr = head_lr

    def on_train_start(self, trainer, pl_module):
        optimizer = trainer.optimizers[0]

        # NeMo creates a single param group. We need to split it into
        # encoder (low LR) and head (high LR) groups.
        # Include ALL params regardless of requires_grad (freeze callback
        # will toggle requires_grad but params must be in optimizer).

        encoder_params = []
        head_params = []

        for name, param in pl_module.named_parameters():
            if name.startswith("encoder."):
                encoder_params.append(param)
            else:
                head_params.append(param)

        # Preserve optimizer settings from existing group
        base_group = optimizer.param_groups[0].copy()

        optimizer.param_groups.clear()

        encoder_group = base_group.copy()
        encoder_group['params'] = encoder_params
        encoder_group['lr'] = self.encoder_lr
        encoder_group['name'] = 'encoder'

        head_group = base_group.copy()
        head_group['params'] = head_params
        head_group['lr'] = self.head_lr
        head_group['name'] = 'decoder_joint_ctc'

        optimizer.param_groups.append(encoder_group)
        optimizer.param_groups.append(head_group)

        # Fix LR scheduler: it was initialized with 1 group, now we have 2.
        # Update base_lrs and last_epoch tracking for the scheduler.
        if trainer.lr_scheduler_configs:
            for sched_config in trainer.lr_scheduler_configs:
                scheduler = sched_config.scheduler
                # Update base_lrs to match new param groups
                scheduler.base_lrs = [self.encoder_lr, self.head_lr]
                # Update _last_lr if it exists
                if hasattr(scheduler, '_last_lr'):
                    scheduler._last_lr = [self.encoder_lr, self.head_lr]

        rank = int(os.environ.get("LOCAL_RANK", 0))
        if rank == 0:
            enc_count = sum(p.numel() for p in encoder_params)
            head_count = sum(p.numel() for p in head_params)
            enc_trainable = sum(p.numel() for p in encoder_params if p.requires_grad)
            print(f"  [DiffLR] Encoder: {enc_count/1e6:.0f}M params @ lr={self.encoder_lr} ({enc_trainable/1e6:.0f}M trainable)")
            print(f"  [DiffLR] Heads:   {head_count/1e6:.0f}M params @ lr={self.head_lr}")

# ---------------------------------------------------------------------------
# OOM-safe VRAM budget calculator
# ---------------------------------------------------------------------------
def compute_vram_safe_batch_params(
    vocab_size: int = 32769,
    enc_hidden: int = 1024,
    joint_hidden: int = 640,
    num_durations: int = 5,
    max_audio_dur: float = 20.0,
    subsampling_factor: int = 8,
    sample_rate: int = 16000,
    window_stride: float = 0.01,
    gpu_memory_gb: float = 140.0,
    model_params_gb: float = 4.5,  # ~1.1B params in bf16
    safety_margin: float = 0.75,  # use at most 75% of available
    grad_accum: int = 4,
) -> dict:
    """Compute OOM-safe batch parameters based on VRAM budget.

    The dominant VRAM consumer is the TDT joint tensor:


find /home/ubuntu/training/ -maxdepth 1 -name "ptcheckpoints" -o -name "checkpoints" 2>/dev/null; ls /home/ubuntu/training/ptcheckpoints 2>/dev/null || true

Search for checkpoint directories


ls -la /home/ubuntu/training/checkpoints/ 2>/dev/null | head -20

List checkpoints directory


total 72
drwxrwxr-x 18 ubuntu ubuntu 4096 Apr  2 00:55 .
drwxrwxr-x 19 ubuntu ubuntu 4096 Mar 31 16:48 ..
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 30 16:54 gemma3n-e2b-ckpt-70000
drwxrwxr-x  5 ubuntu ubuntu 4096 Mar 29 07:33 llama3.2-1b-tts
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 25 23:02 qwen3-asr
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 26 02:30 qwen3-asr-ckpt-100000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 26 19:35 qwen3-asr-ckpt-170000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 27 02:50 qwen3-asr-ckpt-200000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 25 21:23 qwen3-asr-ckpt-24000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 27 19:30 qwen3-asr-ckpt-250000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 28 06:35 qwen3-asr-ckpt-300000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 29 07:09 qwen3-asr-ckpt-400000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 25 21:21 qwen3-asr-ckpt-72000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 31 15:22 qwen3-asr-mixed-ckpt-120000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 31 22:28 qwen3-asr-mixed-ckpt-150000
drwxrwxr-x  2 ubuntu ubuntu 4096 Apr  2 00:55 qwen3-asr-mixed-ckpt-246000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 30 20:19 qwen3-asr-mixed-ckpt-50000
drwxrwxr-x  2 ubuntu ubuntu 4096 Mar 31 02:52 qwen3-asr-mixed-ckpt-76000


560	# ---------------------------------------------------------------------------
561	def patch_training_step(model):
562	    """Patch the model's training_step to extract lang_ids from batch
563	    and pass them through to the encoder via language conditioning.
564	    """
565	    original_training_step = model.training_step
566
567	    def patched_training_step(batch, batch_idx):
568	        # NeMo expects batch = (signal, signal_len, tokens, token_len)
569	        # We added lang_ids as 5th element
570	        if isinstance(batch, (list, tuple)) and len(batch) >= 5:
571	            signal, signal_len, tokens, token_len, lang_ids = batch[:5]
572	            # Store lang_ids on model for encoder hook to pick up
573	            model._current_lang_ids = lang_ids.to(signal.device)
574	            # Pass standard 4-tuple to NeMo
575	            nemo_batch = (signal, signal_len, tokens, token_len)
576	            return original_training_step(nemo_batch, batch_idx)
577	        return original_training_step(batch, batch_idx)
578
579	    model.training_step = patched_training_step
580
581
582	def patch_encoder_for_lang(model):
583	    """Patch encoder forward to inject language IDs stored on model."""
584	    original_encoder_forward = model.encoder.forward
585
586	    def encoder_forward_with_lang(audio_signal, length):
587	        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal, length=length)
588	        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
589	                and model._current_lang_ids is not None):
590	            lang_ids = model._current_lang_ids
591	            # Safety: only apply if batch sizes match (skip during validation)
592	            if lang_ids.shape[0] == encoded.shape[0]:
593	                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
594	                encoded = encoded + lang_bias
595	            # Clear after use to prevent stale values during validation
596	            model._current_lang_ids = None
597	        return encoded, encoded_len
598
599	    model.encoder.forward = encoder_forward_with_lang
600
601
602	# ---------------------------------------------------------------------------
603	# Main
604	# ---------------------------------------------------------------------------
605	def _tokenize_chunk_worker(texts, tokenizer_model_path):
606	    """Tokenize a chunk of texts in a worker process. Module-level for pickling."""
607	    import sentencepiece as spm
608	    sp = spm.SentencePieceProcessor()
609	    sp.Load(tokenizer_model_path)
610	    return [len(sp.EncodeAsIds(str(t).strip() or "<blank>")) for t in texts]
611
612
613	def main():
614	    parser = argparse.ArgumentParser(
615	        description="Production NeMo ASR training with TDT + differential LR"
616	    )
617	    parser.add_argument("--config", type=str, required=True)
618	    parser.add_argument("--max-steps", type=int, default=None)
619	    parser.add_argument("--devices", type=int, default=None)
620	    parser.add_argument("--train-parquet", type=str,
621	                        default="artifacts/phase3/production_train_final.parquet")
622	    parser.add_argument("--val-manifest", type=str,
623	                        default="data/manifests/stage1_prod_val_v2.jsonl")
624	    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
625	    parser.add_argument("--resume-weights-only", action="store_true", default=False,
626	                        help="Resume model weights from checkpoint but reset optimizer (for config changes)")
627	    parser.add_argument("--pretrained-encoder", type=str, default=None,
628	                        help="Path to pretrained encoder weights (.pt file)")
629	    parser.add_argument("--smoke", action="store_true", default=False)
630	    parser.add_argument("--val-every", type=int, default=None)
631	    parser.add_argument("--log-every", type=int, default=None)
632
633	    # Data strategy args
634	    parser.add_argument("--max-batch-dur", type=float, default=120.0,
635	                        help="Max audio seconds per micro-batch (default: 120s)")
636	    parser.add_argument("--max-batch-size", type=int, default=16,
637	                        help="Global hard cap on samples per micro-batch")
638	    parser.add_argument("--max-tokens-in-batch", type=int, default=400,
639	                        help="Max total tokens per batch (controls joint tensor T_dec)")

# ---------------------------------------------------------------------------
def patch_training_step(model):
    """Patch the model's training_step to extract lang_ids from batch
    and pass them through to the encoder via language conditioning.
    """
    original_training_step = model.training_step

    def patched_training_step(batch, batch_idx):
        # NeMo expects batch = (signal, signal_len, tokens, token_len)
        # We added lang_ids as 5th element
        if isinstance(batch, (list, tuple)) and len(batch) >= 5:
            signal, signal_len, tokens, token_len, lang_ids = batch[:5]
            # Store lang_ids on model for encoder hook to pick up
            model._current_lang_ids = lang_ids.to(signal.device)
            # Pass standard 4-tuple to NeMo
            nemo_batch = (signal, signal_len, tokens, token_len)
            return original_training_step(nemo_batch, batch_idx)
        return original_training_step(batch, batch_idx)

    model.training_step = patched_training_step

def patch_encoder_for_lang(model):
    """Patch encoder forward to inject language IDs stored on model."""
    original_encoder_forward = model.encoder.forward

    def encoder_forward_with_lang(audio_signal, length):
        encoded, encoded_len = original_encoder_forward(audio_signal=audio_signal, length=length)
        if (hasattr(model, '_current_lang_ids') and hasattr(model, '_lang_embed')
                and model._current_lang_ids is not None):
            lang_ids = model._current_lang_ids
            # Safety: only apply if batch sizes match (skip during validation)
            if lang_ids.shape[0] == encoded.shape[0]:
                lang_bias = model._lang_embed(lang_ids)  # (B, D, 1)
                encoded = encoded + lang_bias
            # Clear after use to prevent stale values during validation
            model._current_lang_ids = None
        return encoded, encoded_len

    model.encoder.forward = encoder_forward_with_lang

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def _tokenize_chunk_worker(texts, tokenizer_model_path):
    """Tokenize a chunk of texts in a worker process. Module-level for pickling."""
    import sentencepiece as spm
    sp = spm.SentencePieceProcessor()
    sp.Load(tokenizer_model_path)
    return [len(sp.EncodeAsIds(str(t).strip() or "<blank>")) for t in texts]

def main():
    parser = argparse.ArgumentParser(
        description="Production NeMo ASR training with TDT + differential LR"
    )
    parser.add_argument("--config", type=str, required=True)
    parser.add_argument("--max-steps", type=int, default=None)
    parser.add_argument("--devices", type=int, default=None)
    parser.add_argument("--train-parquet", type=str,
                        default="artifacts/phase3/production_train_final.parquet")
    parser.add_argument("--val-manifest", type=str,
                        default="data/manifests/stage1_prod_val_v2.jsonl")
    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
    parser.add_argument("--resume-weights-only", action="store_true", default=False,
                        help="Resume model weights from checkpoint but reset optimizer (for config changes)")
    parser.add_argument("--pretrained-encoder", type=str, default=None,
                        help="Path to pretrained encoder weights (.pt file)")
    parser.add_argument("--smoke", action="store_true", default=False)
    parser.add_argument("--val-every", type=int, default=None)
    parser.add_argument("--log-every", type=int, default=None)

    # Data strategy args
    parser.add_argument("--max-batch-dur", type=float, default=120.0,
                        help="Max audio seconds per micro-batch (default: 120s)")
    parser.add_argument("--max-batch-size", type=int, default=16,
                        help="Global hard cap on samples per micro-batch")
    parser.add_argument("--max-tokens-in-batch", type=int, default=400,
                        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}")

    parser.add_argument("--temperature", type=float, default=0.3)
    parser.add_argument("--grad-accum", type=int, default=None)
    parser.add_argument("--num-workers", type=int, default=None)

    # v2 optimization args
    parser.add_argument("--encoder-lr", type=float, default=2e-5,
                        help="Learning rate for pretrained encoder")
    parser.add_argument("--head-lr", type=float, default=1e-3,
                        help="Learning rate for decoder/joint/CTC heads")
    parser.add_argument("--freeze-encoder-steps", type=int, default=5000,
                        help="Freeze encoder for this many optimizer steps (0=no freeze)")
    parser.add_argument("--ctc-warmup-steps", type=int, default=3000,
                        help="Steps to warm up CTC loss from 0 to target weight")
    parser.add_argument("--no-lang-embed", action="store_true", default=False,
                        help="Disable language embedding conditioning")
    args = parser.parse_args()

    cfg = OmegaConf.load(args.config)

    with open_dict(cfg):
        if args.max_steps is not None:
            cfg.trainer.max_steps = args.max_steps
        if args.devices is not None:
            cfg.trainer.devices = args.devices
            cfg.trainer.strategy = "ddp_find_unused_parameters_true" if args.devices > 1 else "auto"
        if args.grad_accum is not None:
            cfg.trainer.accumulate_grad_batches = args.grad_accum
        if args.val_every is not None:
            cfg.trainer.val_check_interval = args.val_every
            cfg.exp_manager.checkpoint_callback_params.every_n_train_steps = args.val_every
        if args.log_every is not None:
            cfg.trainer.log_every_n_steps = args.log_every
        if args.smoke and args.val_every is None:
            cfg.trainer.val_check_interval = 1.0
        if args.resume_from_checkpoint is not None:
            cfg.exp_manager.resume_from_checkpoint = args.resume_from_checkpoint
            cfg.exp_manager.resume_if_exists = True

        cfg.model.train_ds.is_tarred = False
        for key in ("tarred_audio_filepaths", "shuffle_n", "input_cfg"):
            if key in cfg.model.train_ds:
                del cfg.model.train_ds[key]
        cfg.model.train_ds.manifest_filepath = args.val_manifest
        cfg.model.validation_ds.manifest_filepath = args.val_manifest

    grad_accum = cfg.trainer.get("accumulate_grad_batches", 1)
    devices = cfg.trainer.get("devices", 1)
    max_dur = cfg.model.train_ds.get("max_duration", 20.0)
    min_dur = cfg.model.train_ds.get("min_duration", 0.5)
    num_workers = args.num_workers or cfg.model.train_ds.get("num_workers", 8)

    rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = devices if isinstance(devices, int) else 1

    # -----------------------------------------------------------------------
    # OOM-safe VRAM budget
    # -----------------------------------------------------------------------
    vram_budget = compute_vram_safe_batch_params(
        gpu_memory_gb=140.0,
        max_audio_dur=max_dur,
        grad_accum=grad_accum,
    )
    max_batch_dur = args.max_batch_dur
    max_batch_size = args.max_batch_size

    if rank == 0:
        print("=" * 60)
        print("  Maya ASR — Production Training v2 (TDT + DiffLR)")
        print("=" * 60)
        print()
        print(f"Config:           {args.config}")
        print(f"Max steps:        {cfg.trainer.max_steps}")
        print(f"Devices:          {devices}")
        print(f"Grad accum:       {grad_accum}")
        print(f"Precision:        {cfg.trainer.get('precision', 'bf16-mixed')}")
        print(f"Train parquet:    {args.train_parquet}")
        print(f"Val manifest:     {args.val_manifest}")
        print()
        print("v2 Optimizations:")
        print(f"  Loss:           TDT (durations=[0,1,2,3,4])")
        print(f"  Encoder LR:     {args.encoder_lr}")
        print(f"  Head LR:        {args.head_lr}")
        print(f"  Encoder freeze: {args.freeze_encoder_steps} steps")
        print(f"  CTC warmup:     0→target over {args.ctc_warmup_steps} steps")
        print(f"  Lang embed:     {'yes' if not args.no_lang_embed else 'no'}")
        print()
        print("VRAM Budget (OOM-safe):")
        print(f"  GPU memory:     {vram_budget['gpu_memory_gb']:.0f} GB")
        print(f"  Model (bf16):   {vram_budget['model_params_gb']:.1f} GB")
        print(f"  Optimizer:      {vram_budget['optimizer_gb']:.1f} GB")
        print(f"  Activation:     {vram_budget['activation_budget_gb']:.1f} GB (75% safety)")
        print(f"  Joint/sample:   {vram_budget['joint_per_sample_gb']:.3f} GB")
        print(f"  Total/sample:   {vram_budget['total_per_sample_gb']:.3f} GB")
        print(f"  → max_batch_size: {max_batch_size}")
        print(f"  → max_batch_dur:  {max_batch_dur:.0f}s")
        print()
        print("Data strategy:")
        print(f"  Max batch dur:  {max_batch_dur}s")
        print(f"  Max batch size: {max_batch_size}")
        print(f"  Temperature:    {args.temperature}")
        print(f"  Duration range: [{min_dur}, {max_dur}]s")
        print(f"  Num workers:    {num_workers}")
        print()

    # -----------------------------------------------------------------------
    # Build trainer and model
    # -----------------------------------------------------------------------
    trainer = pl.Trainer(**cfg.trainer, logger=False, use_distributed_sampler=False)
    exp_manager(trainer, cfg.get("exp_manager", None))

    model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel(cfg=cfg.model, trainer=trainer)
    params = sum(p.numel() for p in model.parameters())
    if rank == 0:
        print(f"Model parameters: {params:,}")

    # Resume weights from checkpoint (without optimizer state)
    if args.resume_weights_only and args.resume_from_checkpoint:
        ckpt_path = args.resume_from_checkpoint
        if rank == 0:
            print(f"\nResuming weights only from: {ckpt_path}")


760	        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
761	        model.load_state_dict(ckpt["state_dict"], strict=False)
762	        resume_step = ckpt.get("global_step", 0)
763	        if rank == 0:
764	            print(f"  Loaded model weights from step {resume_step}")
765	            print(f"  Optimizer state NOT restored (fresh optimizer with DiffLR)")
766	        # Clear resume flag so Lightning doesn't also try to restore
767	        args.resume_from_checkpoint = None
768	        with open_dict(cfg):
769	            cfg.exp_manager.resume_from_checkpoint = None
770	            cfg.exp_manager.resume_if_exists = False
771	        del ckpt
772	        print()
773
774	    # Load pretrained encoder weights
775	    if args.pretrained_encoder is not None and not args.resume_weights_only:
776	        if rank == 0:
777	            print(f"\nLoading pretrained encoder from: {args.pretrained_encoder}")
778	        pretrained_state = torch.load(args.pretrained_encoder, map_location="cpu", weights_only=True)
779	        model_state = model.state_dict()
780
781	        loaded, skipped = 0, 0
782	        for key, value in pretrained_state.items():
783	            if key in model_state and model_state[key].shape == value.shape:
784	                model_state[key] = value
785	                loaded += 1
786	            else:
787	                skipped += 1
788	                if rank == 0 and skipped <= 10:
789	                    expected = model_state[key].shape if key in model_state else "NOT FOUND"
790	                    print(f"  Skip: {key} pretrained={value.shape} model={expected}")
791
792	        model.load_state_dict(model_state)
793	        enc_params = sum(pretrained_state[k].numel() for k in pretrained_state if k in model_state and model_state[k].shape == pretrained_state[k].shape)
794	        if rank == 0:
795	            print(f"  Loaded {loaded} tensors ({enc_params/1e6:.0f}M params), skipped {skipped}")
796	            print(f"  Decoder + joint + CTC head: randomly initialized for {model.tokenizer.vocab_size} vocab")
797	        print()
798
799	    # -----------------------------------------------------------------------
800	    # Language conditioning
801	    # -----------------------------------------------------------------------
802	    if not args.no_lang_embed:
803	        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
804	        model._lang_embed = lang_embed
805	        # Register as a proper submodule so it moves to GPU and is saved
806	        model.register_module('_lang_embed_module', lang_embed)
807	        patch_encoder_for_lang(model)
808	        patch_training_step(model)
809	        if rank == 0:
810	            lang_params = sum(p.numel() for p in lang_embed.parameters())
811	            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")
812
813	    # -----------------------------------------------------------------------
814	    # Callbacks
815	    # -----------------------------------------------------------------------
816	    log_interval = args.log_every or cfg.trainer.get("log_every_n_steps", 100)
817	    throughput_cb = ThroughputCallback(
818	        log_every=log_interval,
819	        model_params=params,
820	        grad_accum=grad_accum,
821	        world_size=world_size,
822	    )
823	    trainer.callbacks.append(throughput_cb)
824
825	    freeze_cb = EncoderFreezeCallback(freeze_steps=args.freeze_encoder_steps)
826	    trainer.callbacks.append(freeze_cb)
827
828	    ctc_warmup_cb = CTCWarmupCallback(
829	        target_weight=cfg.model.aux_ctc.get("ctc_loss_weight", 0.3),
830	        warmup_steps=args.ctc_warmup_steps,
831	    )
832	    trainer.callbacks.append(ctc_warmup_cb)
833
834	    diff_lr_cb = DifferentialLRCallback(
835	        encoder_lr=args.encoder_lr,
836	        head_lr=args.head_lr,
837	    )
838	    trainer.callbacks.append(diff_lr_cb)
839
840	    # -----------------------------------------------------------------------
841	    # Load train dataset
842	    # -----------------------------------------------------------------------
843	    if rank == 0:
844	        print(f"Loading train parquet: {args.train_parquet}")
845	    train_dataset = NeMoTarOffsetDataset(
846	        manifest_parquet=args.train_parquet,
847	        tokenizer=model.tokenizer,
848	        sample_rate=cfg.model.train_ds.get("sample_rate", 16000),
849	        max_duration=max_dur,
850	        min_duration=min_dur,
851	    )
852	    if rank == 0:
853	        print(f"Train samples: {len(train_dataset):,}")
854	        print(f"Train audio:   {train_dataset.durations.sum() / 3600:.0f} hours")
855	        # Language distribution
856	        unique_langs, lang_counts = np.unique(train_dataset.languages, return_counts=True)
857	        print(f"Languages:     {len(unique_langs)} ({', '.join(sorted(unique_langs))})")
858	        print()
859
860	    # Create production batch sampler
861	    # Pre-compute token lengths for VRAM-safe batching
862	    # Token length controls T_dec in the joint tensor — the KEY factor for OOM
863	    # Cached to .npy file so subsequent runs skip tokenization entirely
864	    cache_path = Path(args.train_parquet).with_suffix(".token_lengths.npy")
865	    if cache_path.exists():
866	        if rank == 0:
867	            print(f"Loading cached token lengths from {cache_path}")
868	        token_lengths = np.load(cache_path)
869	        if len(token_lengths) != len(train_dataset):
870	            if rank == 0:
871	                print(f"  Cache size mismatch ({len(token_lengths)} vs {len(train_dataset)}), recomputing...")
872	            cache_path.unlink()
873	            token_lengths = None
874	        else:
875	            if rank == 0:
876	                print(f"  Loaded {len(token_lengths):,} cached token lengths")
877
878	    if not cache_path.exists():
879	        if rank == 0:

        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        model.load_state_dict(ckpt["state_dict"], strict=False)
        resume_step = ckpt.get("global_step", 0)
        if rank == 0:
            print(f"  Loaded model weights from step {resume_step}")
            print(f"  Optimizer state NOT restored (fresh optimizer with DiffLR)")
        # Clear resume flag so Lightning doesn't also try to restore
        args.resume_from_checkpoint = None
        with open_dict(cfg):
            cfg.exp_manager.resume_from_checkpoint = None
            cfg.exp_manager.resume_if_exists = False
        del ckpt
        print()

    # Load pretrained encoder weights
    if args.pretrained_encoder is not None and not args.resume_weights_only:
        if rank == 0:
            print(f"\nLoading pretrained encoder from: {args.pretrained_encoder}")
        pretrained_state = torch.load(args.pretrained_encoder, map_location="cpu", weights_only=True)
        model_state = model.state_dict()

        loaded, skipped = 0, 0
        for key, value in pretrained_state.items():
            if key in model_state and model_state[key].shape == value.shape:
                model_state[key] = value
                loaded += 1
            else:
                skipped += 1
                if rank == 0 and skipped <= 10:
                    expected = model_state[key].shape if key in model_state else "NOT FOUND"
                    print(f"  Skip: {key} pretrained={value.shape} model={expected}")

        model.load_state_dict(model_state)
        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)
        if rank == 0:
            print(f"  Loaded {loaded} tensors ({enc_params/1e6:.0f}M params), skipped {skipped}")
            print(f"  Decoder + joint + CTC head: randomly initialized for {model.tokenizer.vocab_size} vocab")
        print()

    # -----------------------------------------------------------------------
    # Language conditioning
    # -----------------------------------------------------------------------
    if not args.no_lang_embed:
        lang_embed = LanguageEmbedding(NUM_LANGUAGES, cfg.model.model_defaults.enc_hidden)
        model._lang_embed = lang_embed
        # Register as a proper submodule so it moves to GPU and is saved
        model.register_module('_lang_embed_module', lang_embed)
        patch_encoder_for_lang(model)
        patch_training_step(model)
        if rank == 0:
            lang_params = sum(p.numel() for p in lang_embed.parameters())
            print(f"  [LangEmbed] {NUM_LANGUAGES} languages → {cfg.model.model_defaults.enc_hidden}D ({lang_params:,} params)")

    # -----------------------------------------------------------------------
    # Callbacks
    # -----------------------------------------------------------------------
    log_interval = args.log_every or cfg.trainer.get("log_every_n_steps", 100)
    throughput_cb = ThroughputCallback(
        log_every=log_interval,
        model_params=params,
        grad_accum=grad_accum,
        world_size=world_size,
    )
    trainer.callbacks.append(throughput_cb)

    freeze_cb = EncoderFreezeCallback(freeze_steps=args.freeze_encoder_steps)
    trainer.callbacks.append(freeze_cb)

    ctc_warmup_cb = CTCWarmupCallback(
        target_weight=cfg.model.aux_ctc.get("ctc_loss_weight", 0.3),
        warmup_steps=args.ctc_warmup_steps,
    )
    trainer.callbacks.append(ctc_warmup_cb)

    diff_lr_cb = DifferentialLRCallback(
        encoder_lr=args.encoder_lr,
        head_lr=args.head_lr,
    )
    trainer.callbacks.append(diff_lr_cb)

    # -----------------------------------------------------------------------
    # Load train dataset
    # -----------------------------------------------------------------------
    if rank == 0:
        print(f"Loading train parquet: {args.train_parquet}")
    train_dataset = NeMoTarOffsetDataset(
        manifest_parquet=args.train_parquet,
        tokenizer=model.tokenizer,
        sample_rate=cfg.model.train_ds.get("sample_rate", 16000),
        max_duration=max_dur,
        min_duration=min_dur,
    )
    if rank == 0:
        print(f"Train samples: {len(train_dataset):,}")
        print(f"Train audio:   {train_dataset.durations.sum() / 3600:.0f} hours")
        # Language distribution
        unique_langs, lang_counts = np.unique(train_dataset.languages, return_counts=True)
        print(f"Languages:     {len(unique_langs)} ({', '.join(sorted(unique_langs))})")
        print()

    # Create production batch sampler
    # Pre-compute token lengths for VRAM-safe batching
    # Token length controls T_dec in the joint tensor — the KEY factor for OOM
    # Cached to .npy file so subsequent runs skip tokenization entirely
    cache_path = Path(args.train_parquet).with_suffix(".token_lengths.npy")
    if cache_path.exists():
        if rank == 0:
            print(f"Loading cached token lengths from {cache_path}")
        token_lengths = np.load(cache_path)
        if len(token_lengths) != len(train_dataset):
            if rank == 0:
                print(f"  Cache size mismatch ({len(token_lengths)} vs {len(train_dataset)}), recomputing...")
            cache_path.unlink()
            token_lengths = None
        else:
            if rank == 0:
                print(f"  Loaded {len(token_lengths):,} cached token lengths")

    if not cache_path.exists():
        if rank == 0:


Now let me download and inspect the checkpoint.


mkdir -p /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/ptcheckpoints/hybrid-tdt-gemma/03-30-2026/ckpt-260000/ && ls /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/restore_checkpoint.py 2>/dev/null && head -50 /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/restore_checkpoint.py

Check restore checkpoint script


/home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1/scripts/restore_checkpoint.py
#!/usr/bin/env python3
"""Restore a training checkpoint from Cloudflare R2.

Downloads .ckpt/.nemo files and verifies integrity via SHA-256.

Usage:
  # Restore latest step
  python3 scripts/restore_checkpoint.py \
    --model-name maya-asr-stage1 --target-dir /tmp/restored

  # Restore specific step
  python3 scripts/restore_checkpoint.py \
    --model-name maya-asr-stage1 --step 50 --target-dir /tmp/restored

  # List available steps
  python3 scripts/restore_checkpoint.py \
    --model-name maya-asr-stage1 --list
"""

import argparse
import hashlib
import json
import os
import re
import sys
from pathlib import Path

import boto3
from dotenv import load_dotenv

ENV_FILE = Path(__file__).resolve().parent.parent / ".env"
if ENV_FILE.exists():
    load_dotenv(ENV_FILE)

def get_s3_client():
    endpoint = os.environ.get("R2_ENDPOINT_URL")
    key_id = os.environ.get("R2_ACCESS_KEY_ID")
    secret = os.environ.get("R2_SECRET_ACCESS_KEY")
    if not all([endpoint, key_id, secret]):
        print("ERROR: R2 credentials not set.", file=sys.stderr)
        sys.exit(1)
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=key_id,
        aws_secret_access_key=secret,
        region_name="auto",
    )


50
51
52	def list_steps(s3, bucket: str, model_name: str) -> list[int]:
53	    """List all available step numbers for a model on R2."""
54	    prefix = f"{model_name}/"
55	    paginator = s3.get_paginator("list_objects_v2")
56	    steps = set()
57	    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
58	        for cp in page.get("CommonPrefixes", []):
59	            # cp["Prefix"] = "model/step_00000050/"
60	            m = re.search(r"step_(\d+)", cp["Prefix"])
61	            if m:
62	                steps.add(int(m.group(1)))
63	    return sorted(steps)
64
65
66	def _sha256(path: Path) -> str:
67	    h = hashlib.sha256()
68	    with open(path, "rb") as f:
69	        for chunk in iter(lambda: f.read(1 << 20), b""):
70	            h.update(chunk)
71	    return h.hexdigest()
72
73
74	def main():
75	    parser = argparse.ArgumentParser(description="Restore checkpoint from R2")
76	    parser.add_argument("--model-name", type=str, required=True)
77	    parser.add_argument("--step", type=int, default=None, help="Step to restore (default: latest)")
78	    parser.add_argument("--target-dir", type=Path, default=None)
79	    parser.add_argument("--bucket", type=str, default=None)
80	    parser.add_argument("--list", action="store_true", default=False, help="List available steps")
81	    args = parser.parse_args()
82
83	    bucket = args.bucket or os.environ.get("R2_BUCKET_CHECKPOINTS", "ptcheckpoints")
84	    s3 = get_s3_client()
85
86	    # List mode
87	    if args.list:
88	        steps = list_steps(s3, bucket, args.model_name)
89	        if not steps:
90	            print(f"No checkpoints found for {args.model_name} in {bucket}")
91	            sys.exit(1)
92	        print(f"Available steps for {args.model_name}:")
93	        for step in steps:
94	            print(f"  step_{step:08d}")
95	        return
96
97	    # Determine step
98	    if args.step is not None:
99	        step = args.step
100	    else:
101	        steps = list_steps(s3, bucket, args.model_name)
102	        if not steps:
103	            print(f"ERROR: No checkpoints for {args.model_name}", file=sys.stderr)
104	            sys.exit(1)
105	        step = steps[-1]
106	        print(f"Latest step: {step}")
107
108	    if args.target_dir is None:
109	        print("ERROR: --target-dir required for restore", file=sys.stderr)
110	        sys.exit(1)
111
112	    prefix = f"{args.model_name}/step_{step:08d}"
113	    target = args.target_dir
114	    target.mkdir(parents=True, exist_ok=True)
115
116	    print(f"Restoring {args.model_name} step {step}")
117	    print(f"  Source: s3://{bucket}/{prefix}/")
118	    print(f"  Target: {target}")
119	    print()
120
121	    # Download metadata first
122	    meta_key = f"{prefix}/upload_metadata.json"
123	    meta_local = target / "upload_metadata.json"
124	    try:
125	        s3.download_file(bucket, meta_key, str(meta_local))
126	        metadata = json.loads(meta_local.read_text())
127	        print(f"  Metadata: {len(metadata.get('files', []))} files recorded")
128	    except Exception:
129	        print("  WARN: upload_metadata.json not found, downloading without verification")
130	        metadata = None
131
132	    # List all objects under prefix
133	    resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix + "/")
134	    objects = resp.get("Contents", [])
135	    ckpt_objects = [o for o in objects if o["Key"].endswith((".ckpt", ".nemo"))]
136
137	    if not ckpt_objects:
138	        print(f"ERROR: No .ckpt/.nemo files at {prefix}/", file=sys.stderr)
139	        sys.exit(1)
140
141	    # Build expected checksums from metadata
142	    expected = {}
143	    if metadata:
144	        for f in metadata.get("files", []):
145	            expected[f["name"]] = {
146	                "size_bytes": f["size_bytes"],
147	                "sha256": f.get("sha256"),
148	            }
149

def list_steps(s3, bucket: str, model_name: str) -> list[int]:
    """List all available step numbers for a model on R2."""
    prefix = f"{model_name}/"
    paginator = s3.get_paginator("list_objects_v2")
    steps = set()
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
        for cp in page.get("CommonPrefixes", []):
            # cp["Prefix"] = "model/step_00000050/"
            m = re.search(r"step_(\d+)", cp["Prefix"])
            if m:
                steps.add(int(m.group(1)))
    return sorted(steps)

def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()

def main():
    parser = argparse.ArgumentParser(description="Restore checkpoint from R2")
    parser.add_argument("--model-name", type=str, required=True)
    parser.add_argument("--step", type=int, default=None, help="Step to restore (default: latest)")
    parser.add_argument("--target-dir", type=Path, default=None)
    parser.add_argument("--bucket", type=str, default=None)
    parser.add_argument("--list", action="store_true", default=False, help="List available steps")
    args = parser.parse_args()

    bucket = args.bucket or os.environ.get("R2_BUCKET_CHECKPOINTS", "ptcheckpoints")
    s3 = get_s3_client()

    # List mode
    if args.list:
        steps = list_steps(s3, bucket, args.model_name)
        if not steps:
            print(f"No checkpoints found for {args.model_name} in {bucket}")
            sys.exit(1)
        print(f"Available steps for {args.model_name}:")
        for step in steps:
            print(f"  step_{step:08d}")
        return

    # Determine step
    if args.step is not None:
        step = args.step
    else:
        steps = list_steps(s3, bucket, args.model_name)
        if not steps:
            print(f"ERROR: No checkpoints for {args.model_name}", file=sys.stderr)
            sys.exit(1)
        step = steps[-1]
        print(f"Latest step: {step}")

    if args.target_dir is None:
        print("ERROR: --target-dir required for restore", file=sys.stderr)
        sys.exit(1)

    prefix = f"{args.model_name}/step_{step:08d}"
    target = args.target_dir
    target.mkdir(parents=True, exist_ok=True)

    print(f"Restoring {args.model_name} step {step}")
    print(f"  Source: s3://{bucket}/{prefix}/")
    print(f"  Target: {target}")
    print()

    # Download metadata first
    meta_key = f"{prefix}/upload_metadata.json"
    meta_local = target / "upload_metadata.json"
    try:
        s3.download_file(bucket, meta_key, str(meta_local))
        metadata = json.loads(meta_local.read_text())
        print(f"  Metadata: {len(metadata.get('files', []))} files recorded")
    except Exception:
        print("  WARN: upload_metadata.json not found, downloading without verification")
        metadata = None

    # List all objects under prefix
    resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix + "/")
    objects = resp.get("Contents", [])
    ckpt_objects = [o for o in objects if o["Key"].endswith((".ckpt", ".nemo"))]

    if not ckpt_objects:
        print(f"ERROR: No .ckpt/.nemo files at {prefix}/", file=sys.stderr)
        sys.exit(1)

    # Build expected checksums from metadata
    expected = {}
    if metadata:
        for f in metadata.get("files", []):
            expected[f["name"]] = {
                "size_bytes": f["size_bytes"],
                "sha256": f.get("sha256"),
            }


The checkpoint uses R2 cloud storage. Let me try to download it.


cd /home/ubuntu/training/maya-asr-hybrid-fastconformer-rnnt-stage1 && python3 scripts/restore_checkpoint.py --model-name hybrid-tdt-gemma --list 2>&1 | head -30

List available checkpoints on R2


No checkpoints found for hybrid-tdt-gemma in ptcheckpoints


python3 scripts/restore_checkpoint.py --model-name maya-asr-stage1 --list 2>&1 | head -30

List maya-asr-stage1 checkpoints


Available steps for maya-asr-stage1:
  step_00000000


python3 -c "
import boto3, os
from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
resp = s3.list_objects_v2(Bucket=bucket, Delimiter='/')
for p in resp.get('CommonPrefixes', []):
    print(p['Prefix'])
" 2>&1 | head -20

List all model prefixes in R2 bucket


gemma3n-E2B-asr/
hybrid-tdt-gemma/
llama3.1-1b-instruct/
llama3.2-1b/
maya-asr-stage1/
qwen3-0.6b/
qwen3-asr-1.7B-mixed/
qwen3-asr-1.7B/


python3 -c "
import boto3, os
from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
resp = s3.list_objects_v2(Bucket=bucket, Prefix='hybrid-tdt-gemma/', Delimiter='/')
for p in resp.get('CommonPrefixes', []):
    print(p['Prefix'])
" 2>&1

List hybrid-tdt-gemma subfolders


hybrid-tdt-gemma/03-26-2026/
hybrid-tdt-gemma/03-27-2026/
hybrid-tdt-gemma/03-28-2026/
hybrid-tdt-gemma/03-29-2026/
hybrid-tdt-gemma/03-30-2026/


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

List 03-30 checkpoint files


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


python3 -c "
import boto3, os
from dotenv import load_dotenv
from pathlib import Path
load_dotenv(Path('.env'))
s3 = boto3.client('s3', endpoint_url=os.environ['R2_ENDPOINT_URL'],
    aws_access_key_id=os.environ['R2_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['R2_SECRET_ACCESS_KEY'], region_name='auto')
bucket = 'ptcheckpoints'
paginator = s3.get_paginator('list_objects_v2')
for page in paginator.paginate(Bucket=bucket, Prefix='hybrid-tdt-gemma/03-30-2026/', Delimiter='/'):
    for p in page.get('CommonPrefixes', []):
        print(p['Prefix'])
" 2>&1

List subdirs of 03-30 date