feat: Refactor training with SOLID principles and add optimizations

Training refactor:
- Modular training package (config, data, models, trainer, checkpoint)
- Auto GPU detection (H200 141GB, A100, RTX 4090)
- Flash Attention 2 with fallback to SDPA
- OOM recovery with proper cleanup (PyTorch best practices)
- NaN/Inf detection and gradient monitoring
- Sequence length bucketing for memory efficiency
- BF16 for numerical stability

Stage 1 improvements:
- Added --no_decay flag (fixed text_ratio for adapter warmup)
- Dynamic decay optional via --use_decay

Dataset tools:
- analyze_dataset.py for dataset inspection
- validate_dataset.py for format validation

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

datasets/analyze_dataset.py

@@ -0,0 +1,575 @@
+#!/usr/bin/env python3
+"""
+Deep Dataset Analyzer - Comprehensive analysis of generated dataset.
+
+Analyzes:
+1. Structure and field presence
+2. Whisper features (shape, statistics, quality)
+3. SNAC tokens (offsets, distribution, frame integrity)
+4. Text content (questions, answers, tokenization)
+5. Word alignments (timing, coverage, token mapping)
+6. Cross-field consistency
+7. Training readiness checks
+"""
+
+import argparse
+import sys
+from pathlib import Path
+from collections import defaultdict, Counter
+import json
+
+import torch
+import numpy as np
+
+# SNAC constants
+SNAC_BASE = 128266
+SNAC_LAYERS = 7
+SNAC_VOCAB_PER_LAYER = 4096
+WHISPER_DIM = 1280
+EOS_TOKEN = 128009
+
+
+def print_header(title):
+    print(f"\n{'='*70}")
+    print(f" {title}")
+    print('='*70)
+
+
+def print_subheader(title):
+    print(f"\n{'-'*50}")
+    print(f" {title}")
+    print('-'*50)
+
+
+def analyze_whisper_features(dataset):
+    """Deep analysis of Whisper features."""
+    print_header("WHISPER FEATURES ANALYSIS")
+
+    lengths = []
+    dims = []
+    min_vals = []
+    max_vals = []
+    mean_vals = []
+    std_vals = []
+    has_nan = 0
+    has_inf = 0
+    dtypes = Counter()
+
+    for i, sample in enumerate(dataset):
+        if "whisper_features" not in sample:
+            continue
+
+        wf = sample["whisper_features"]
+
+        if not isinstance(wf, torch.Tensor):
+            print(f"  [WARN] Sample {i}: whisper_features is {type(wf).__name__}, not Tensor")
+            continue
+
+        dtypes[str(wf.dtype)] += 1
+        lengths.append(wf.shape[0])
+
+        if wf.dim() >= 2:
+            dims.append(wf.shape[1])
+
+        if torch.isnan(wf).any():
+            has_nan += 1
+        if torch.isinf(wf).any():
+            has_inf += 1
+
+        min_vals.append(wf.min().item())
+        max_vals.append(wf.max().item())
+        mean_vals.append(wf.mean().item())
+        std_vals.append(wf.std().item())
+
+    if not lengths:
+        print("  No whisper_features found!")
+        return
+
+    lengths = np.array(lengths)
+
+    print(f"\n  Samples with whisper_features: {len(lengths)}/{len(dataset)}")
+    print(f"  Data types: {dict(dtypes)}")
+
+    print_subheader("Sequence Length")
+    print(f"  Min: {lengths.min()}, Max: {lengths.max()}")
+    print(f"  Mean: {lengths.mean():.1f}, Std: {lengths.std():.1f}")
+    print(f"  Median: {np.median(lengths):.1f}")
+
+    # Length distribution
+    percentiles = [10, 25, 50, 75, 90, 95, 99]
+    print(f"  Percentiles: {dict(zip(percentiles, [int(np.percentile(lengths, p)) for p in percentiles]))}")
+
+    if dims:
+        dims = np.array(dims)
+        unique_dims = np.unique(dims)
+        print(f"\n  Feature dimension: {unique_dims}")
+        if len(unique_dims) == 1 and unique_dims[0] == WHISPER_DIM:
+            print(f"  ✅ Correct Whisper dimension ({WHISPER_DIM})")
+        else:
+            print(f"  ⚠️ Expected dimension {WHISPER_DIM}")
+
+    print_subheader("Value Statistics")
+    print(f"  Min value: {np.min(min_vals):.4f}")
+    print(f"  Max value: {np.max(max_vals):.4f}")
+    print(f"  Mean of means: {np.mean(mean_vals):.4f}")
+    print(f"  Mean of stds: {np.mean(std_vals):.4f}")
+
+    print_subheader("Quality Checks")
+    print(f"  Samples with NaN: {has_nan} {'⚠️' if has_nan > 0 else '✅'}")
+    print(f"  Samples with Inf: {has_inf} {'⚠️' if has_inf > 0 else '✅'}")
+
+    # Check for constant/zero features
+    zero_features = sum(1 for s in std_vals if s < 0.001)
+    print(f"  Samples with near-zero std: {zero_features} {'⚠️' if zero_features > 0 else '✅'}")
+
+
+def analyze_snac_tokens(dataset):
+    """Deep analysis of SNAC tokens."""
+    print_header("SNAC TOKENS ANALYSIS")
+
+    lengths = []
+    frame_counts = []
+    token_values = []
+    offset_distribution = defaultdict(int)
+    raw_token_distribution = defaultdict(int)
+    incomplete_frames = 0
+    has_offset = 0
+    no_offset = 0
+    dtypes = Counter()
+
+    for i, sample in enumerate(dataset):
+        if "snac_tokens" not in sample:
+            continue
+
+        st = sample["snac_tokens"]
+
+        if not isinstance(st, torch.Tensor):
+            print(f"  [WARN] Sample {i}: snac_tokens is {type(st).__name__}, not Tensor")
+            continue
+
+        dtypes[str(st.dtype)] += 1
+        length = len(st)
+        lengths.append(length)
+        frame_counts.append(length // SNAC_LAYERS)
+
+        if length % SNAC_LAYERS != 0:
+            incomplete_frames += 1
+
+        # Analyze token values
+        for j, tok in enumerate(st.tolist()):
+            token_values.append(tok)
+
+            if tok >= SNAC_BASE:
+                has_offset += 1
+                # Determine which position offset
+                offset_idx = (tok - SNAC_BASE) // SNAC_VOCAB_PER_LAYER
+                offset_distribution[offset_idx] += 1
+                # Get raw token value
+                raw_tok = (tok - SNAC_BASE) % SNAC_VOCAB_PER_LAYER
+                raw_token_distribution[raw_tok] += 1
+            else:
+                no_offset += 1
+
+    if not lengths:
+        print("  No snac_tokens found!")
+        return
+
+    lengths = np.array(lengths)
+    frame_counts = np.array(frame_counts)
+    token_values = np.array(token_values)
+
+    print(f"\n  Samples with snac_tokens: {len(lengths)}/{len(dataset)}")
+    print(f"  Data types: {dict(dtypes)}")
+
+    print_subheader("Token Counts")
+    print(f"  Total tokens: {len(token_values):,}")
+    print(f"  Min per sample: {lengths.min()}, Max: {lengths.max()}")
+    print(f"  Mean: {lengths.mean():.1f}, Std: {lengths.std():.1f}")
+
+    print_subheader("Frame Analysis (7 tokens per frame)")
+    print(f"  Min frames: {frame_counts.min()}, Max: {frame_counts.max()}")
+    print(f"  Mean frames: {frame_counts.mean():.1f}")
+    print(f"  Incomplete frames (not multiple of 7): {incomplete_frames} {'⚠️' if incomplete_frames > 0 else '✅'}")
+
+    # Duration estimation (75 frames/second for SNAC 24kHz)
+    durations = frame_counts / 75.0
+    print(f"\n  Audio duration (estimated):")
+    print(f"    Min: {durations.min():.2f}s, Max: {durations.max():.2f}s")
+    print(f"    Mean: {durations.mean():.2f}s, Total: {durations.sum()/60:.1f} min")
+
+    print_subheader("Token Offset Analysis")
+    print(f"  Tokens with offset (>= {SNAC_BASE}): {has_offset:,} ({100*has_offset/len(token_values):.1f}%)")
+    print(f"  Tokens without offset: {no_offset:,} ({100*no_offset/len(token_values):.1f}%)")
+
+    if has_offset > 0:
+        print(f"\n  Offset distribution by position (0-6):")
+        for pos in range(SNAC_LAYERS):
+            count = offset_distribution.get(pos, 0)
+            expected_pct = 100 / SNAC_LAYERS
+            actual_pct = 100 * count / has_offset if has_offset > 0 else 0
+            status = "✅" if abs(actual_pct - expected_pct) < 5 else "⚠️"
+            print(f"    Position {pos}: {count:,} ({actual_pct:.1f}%) {status}")
+
+    print_subheader("Token Value Range")
+    print(f"  Min token: {token_values.min()}")
+    print(f"  Max token: {token_values.max()}")
+
+    if has_offset > 0:
+        expected_min = SNAC_BASE
+        expected_max = SNAC_BASE + (SNAC_LAYERS * SNAC_VOCAB_PER_LAYER) - 1
+        print(f"  Expected range: [{expected_min}, {expected_max}]")
+        if token_values.min() >= expected_min and token_values.max() <= expected_max:
+            print(f"  ✅ Token range valid")
+        else:
+            print(f"  ⚠️ Some tokens outside expected range")
+
+
+def analyze_text_content(dataset):
+    """Analyze text fields (question, answer, text_tokens)."""
+    print_header("TEXT CONTENT ANALYSIS")
+
+    questions = []
+    answers = []
+    text_token_lengths = []
+    has_text = 0
+    has_answer = 0
+    has_text_tokens = 0
+
+    for sample in dataset:
+        if "text" in sample and sample["text"]:
+            has_text += 1
+            questions.append(sample["text"])
+
+        if "answer" in sample and sample["answer"]:
+            has_answer += 1
+            answers.append(sample["answer"])
+
+        if "text_tokens" in sample:
+            has_text_tokens += 1
+            tt = sample["text_tokens"]
+            if isinstance(tt, torch.Tensor):
+                text_token_lengths.append(len(tt))
+            elif isinstance(tt, list):
+                text_token_lengths.append(len(tt))
+
+    print(f"\n  Field presence:")
+    print(f"    text (question): {has_text}/{len(dataset)} ({100*has_text/len(dataset):.1f}%)")
+    print(f"    answer: {has_answer}/{len(dataset)} ({100*has_answer/len(dataset):.1f}%)")
+    print(f"    text_tokens: {has_text_tokens}/{len(dataset)} ({100*has_text_tokens/len(dataset):.1f}%)")
+
+    if questions:
+        print_subheader("Questions Analysis")
+        q_lengths = [len(q) for q in questions]
+        q_words = [len(q.split()) for q in questions]
+        print(f"  Character length: min={min(q_lengths)}, max={max(q_lengths)}, mean={np.mean(q_lengths):.1f}")
+        print(f"  Word count: min={min(q_words)}, max={max(q_words)}, mean={np.mean(q_words):.1f}")
+
+        # Sample questions
+        print(f"\n  Sample questions:")
+        for q in questions[:3]:
+            print(f"    - {q[:80]}{'...' if len(q) > 80 else ''}")
+
+    if answers:
+        print_subheader("Answers Analysis")
+        a_lengths = [len(a) for a in answers]
+        a_words = [len(a.split()) for a in answers]
+        print(f"  Character length: min={min(a_lengths)}, max={max(a_lengths)}, mean={np.mean(a_lengths):.1f}")
+        print(f"  Word count: min={min(a_words)}, max={max(a_words)}, mean={np.mean(a_words):.1f}")
+
+        # Sample answers
+        print(f"\n  Sample answers:")
+        for a in answers[:3]:
+            print(f"    - {a[:100]}{'...' if len(a) > 100 else ''}")
+
+    if text_token_lengths:
+        print_subheader("Pre-tokenized Answer Tokens")
+        tl = np.array(text_token_lengths)
+        print(f"  Token count: min={tl.min()}, max={tl.max()}, mean={tl.mean():.1f}")
+
+
+def analyze_word_alignments(dataset):
+    """Analyze word alignments for IST-LM interleaving."""
+    print_header("WORD ALIGNMENTS ANALYSIS")
+
+    has_alignments = 0
+    word_counts = []
+    frame_coverages = []
+    has_tokens = 0
+    token_counts = []
+    timing_issues = 0
+
+    for i, sample in enumerate(dataset):
+        if "word_alignments" not in sample or not sample["word_alignments"]:
+            continue
+
+        has_alignments += 1
+        wa = sample["word_alignments"]
+        word_counts.append(len(wa))
+
+        # Check alignment quality
+        last_end = 0
+        sample_has_tokens = True
+        sample_token_count = 0
+
+        for j, alignment in enumerate(wa):
+            # Check required fields
+            if 'word' not in alignment or 'start_frame' not in alignment or 'end_frame' not in alignment:
+                continue
+
+            start = alignment['start_frame']
+            end = alignment['end_frame']
+
+            # Check timing consistency
+            if start < last_end - 1:  # Allow small overlap
+                timing_issues += 1
+            last_end = end
+
+            # Check token presence
+            if 'tokens' in alignment and alignment['tokens']:
+                sample_token_count += len(alignment['tokens'])
+            else:
+                sample_has_tokens = False
+
+        if sample_has_tokens and sample_token_count > 0:
+            has_tokens += 1
+            token_counts.append(sample_token_count)
+
+        # Calculate frame coverage
+        if "snac_tokens" in sample:
+            snac_len = len(sample["snac_tokens"])
+            total_frames = snac_len // SNAC_LAYERS
+            if total_frames > 0 and wa:
+                covered_frames = wa[-1]['end_frame'] if wa[-1]['end_frame'] else 0
+                frame_coverages.append(min(100, 100 * covered_frames / total_frames))
+
+    print(f"\n  Samples with word_alignments: {has_alignments}/{len(dataset)} ({100*has_alignments/len(dataset):.1f}%)")
+
+    if not has_alignments:
+        print("  No word alignments found!")
+        return
+
+    print(f"  Samples with pre-computed tokens: {has_tokens}/{has_alignments} ({100*has_tokens/has_alignments:.1f}%)")
+
+    word_counts = np.array(word_counts)
+    print_subheader("Word Count per Sample")
+    print(f"  Min: {word_counts.min()}, Max: {word_counts.max()}, Mean: {word_counts.mean():.1f}")
+
+    if token_counts:
+        token_counts = np.array(token_counts)
+        print_subheader("Token Count per Sample (from alignments)")
+        print(f"  Min: {token_counts.min()}, Max: {token_counts.max()}, Mean: {token_counts.mean():.1f}")
+
+    if frame_coverages:
+        fc = np.array(frame_coverages)
+        print_subheader("Frame Coverage")
+        print(f"  Min: {fc.min():.1f}%, Max: {fc.max():.1f}%, Mean: {fc.mean():.1f}%")
+
+    print_subheader("Quality Checks")
+    print(f"  Timing issues (overlaps): {timing_issues} {'⚠️' if timing_issues > 0 else '✅'}")
+
+
+def analyze_cross_field_consistency(dataset):
+    """Check consistency between related fields."""
+    print_header("CROSS-FIELD CONSISTENCY")
+
+    issues = defaultdict(int)
+
+    for i, sample in enumerate(dataset):
+        # Check answer vs text_tokens consistency
+        if "answer" in sample and "text_tokens" in sample:
+            answer = sample["answer"]
+            text_tokens = sample["text_tokens"]
+            if isinstance(text_tokens, torch.Tensor):
+                token_len = len(text_tokens)
+            else:
+                token_len = len(text_tokens) if text_tokens else 0
+
+            # Rough check: tokens should be roughly proportional to text length
+            if answer and token_len > 0:
+                chars_per_token = len(answer) / token_len
+                if chars_per_token < 1 or chars_per_token > 10:
+                    issues["answer_token_mismatch"] += 1
+
+        # Check word_alignments vs answer consistency
+        if "word_alignments" in sample and "answer" in sample:
+            wa = sample["word_alignments"]
+            answer = sample["answer"]
+            if wa and answer:
+                wa_words = len(wa)
+                answer_words = len(answer.split())
+                if abs(wa_words - answer_words) > answer_words * 0.3:  # 30% tolerance
+                    issues["alignment_word_mismatch"] += 1
+
+        # Check snac_tokens frame count vs alignments
+        if "snac_tokens" in sample and "word_alignments" in sample:
+            snac = sample["snac_tokens"]
+            wa = sample["word_alignments"]
+            if isinstance(snac, torch.Tensor) and wa:
+                snac_frames = len(snac) // SNAC_LAYERS
+                last_frame = max((a.get('end_frame', 0) for a in wa), default=0)
+                if last_frame > snac_frames * 1.1:  # 10% tolerance
+                    issues["frame_overflow"] += 1
+
+    print(f"\n  Consistency checks:")
+    print(f"    Answer/token length mismatch: {issues['answer_token_mismatch']} {'⚠️' if issues['answer_token_mismatch'] > 0 else '✅'}")
+    print(f"    Alignment/answer word mismatch: {issues['alignment_word_mismatch']} {'⚠️' if issues['alignment_word_mismatch'] > 0 else '✅'}")
+    print(f"    Frame overflow in alignments: {issues['frame_overflow']} {'⚠️' if issues['frame_overflow'] > 0 else '✅'}")
+
+
+def analyze_training_readiness(dataset):
+    """Check if dataset is ready for training."""
+    print_header("TRAINING READINESS CHECK")
+
+    required_fields = ["whisper_features", "snac_tokens"]
+    optional_fields = ["text", "answer", "text_tokens", "word_alignments"]
+
+    field_presence = defaultdict(int)
+    complete_samples = 0
+
+    issues = []
+
+    for sample in dataset:
+        has_required = all(f in sample for f in required_fields)
+        if has_required:
+            complete_samples += 1
+
+        for field in required_fields + optional_fields:
+            if field in sample and sample[field] is not None:
+                if isinstance(sample[field], (torch.Tensor, list, str)):
+                    if len(sample[field]) > 0:
+                        field_presence[field] += 1
+                else:
+                    field_presence[field] += 1
+
+    print(f"\n  Total samples: {len(dataset)}")
+    print(f"  Complete samples (with required fields): {complete_samples} ({100*complete_samples/len(dataset):.1f}%)")
+
+    print(f"\n  Field presence:")
+    for field in required_fields:
+        count = field_presence[field]
+        status = "✅" if count == len(dataset) else "❌"
+        print(f"    {field}: {count}/{len(dataset)} {status} (required)")
+
+    for field in optional_fields:
+        count = field_presence[field]
+        status = "✅" if count == len(dataset) else "⚠️"
+        print(f"    {field}: {count}/{len(dataset)} {status} (optional)")
+
+    # Overall assessment
+    print_subheader("Overall Assessment")
+
+    ready = True
+
+    if complete_samples < len(dataset):
+        print(f"  ❌ {len(dataset) - complete_samples} samples missing required fields")
+        ready = False
+    else:
+        print(f"  ✅ All samples have required fields")
+
+    if field_presence.get("text_tokens", 0) == len(dataset):
+        print(f"  ✅ Pre-tokenized text available (faster training)")
+    else:
+        print(f"  ⚠️ text_tokens missing in some samples (will tokenize on-the-fly)")
+
+    if field_presence.get("word_alignments", 0) == len(dataset):
+        print(f"  ✅ Word alignments available (semantic interleaving)")
+    else:
+        print(f"  ⚠️ word_alignments missing (will use positional interleaving)")
+
+    if ready:
+        print(f"\n  🎉 DATASET READY FOR TRAINING!")
+    else:
+        print(f"\n  ⚠️ Dataset has issues that need to be fixed")
+
+    return ready
+
+
+def print_sample_inspection(dataset, num_samples=3):
+    """Print detailed inspection of sample items."""
+    print_header(f"SAMPLE INSPECTION (first {num_samples} samples)")
+
+    for i in range(min(num_samples, len(dataset))):
+        sample = dataset[i]
+        print(f"\n  Sample {i}:")
+
+        for key, value in sample.items():
+            if isinstance(value, torch.Tensor):
+                print(f"    {key}: Tensor {list(value.shape)} {value.dtype}")
+                if key == "snac_tokens" and len(value) > 0:
+                    print(f"      First 7 tokens: {value[:7].tolist()}")
+                    print(f"      Last 7 tokens: {value[-7:].tolist()}")
+            elif isinstance(value, str):
+                preview = value[:60] + "..." if len(value) > 60 else value
+                print(f"    {key}: '{preview}'")
+            elif isinstance(value, list):
+                print(f"    {key}: list[{len(value)}]")
+                if key == "word_alignments" and len(value) > 0:
+                    print(f"      First: {value[0]}")
+                    if len(value) > 1:
+                        print(f"      Last: {value[-1]}")
+            else:
+                print(f"    {key}: {type(value).__name__}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Deep dataset analysis")
+    parser.add_argument("--path", type=str, required=True, help="Path to dataset .pt file")
+    parser.add_argument("--max-samples", type=int, default=None, help="Max samples to analyze")
+    parser.add_argument("--quick", action="store_true", help="Quick analysis (skip detailed checks)")
+    args = parser.parse_args()
+
+    print("\n" + "="*70)
+    print(" DEEP DATASET ANALYZER")
+    print("="*70)
+    print(f"\n  File: {args.path}")
+
+    # Check file exists
+    if not Path(args.path).exists():
+        print(f"\n  ❌ File not found: {args.path}")
+        sys.exit(1)
+
+    # Load dataset
+    print(f"  Loading dataset...")
+    try:
+        dataset = torch.load(args.path, map_location="cpu", weights_only=False)
+    except Exception as e:
+        print(f"\n  ❌ Failed to load: {e}")
+        sys.exit(1)
+
+    if not isinstance(dataset, list):
+        print(f"\n  ❌ Dataset should be list, got {type(dataset).__name__}")
+        sys.exit(1)
+
+    print(f"  Loaded {len(dataset):,} samples")
+
+    # Limit samples if requested
+    if args.max_samples and args.max_samples < len(dataset):
+        print(f"  Analyzing first {args.max_samples} samples")
+        dataset = dataset[:args.max_samples]
+
+    # File size
+    file_size = Path(args.path).stat().st_size / (1024**3)
+    print(f"  File size: {file_size:.2f} GB")
+    print(f"  Avg per sample: {file_size*1024/len(dataset):.2f} MB")
+
+    # Run analyses
+    analyze_whisper_features(dataset)
+    analyze_snac_tokens(dataset)
+    analyze_text_content(dataset)
+
+    if not args.quick:
+        analyze_word_alignments(dataset)
+        analyze_cross_field_consistency(dataset)
+
+    ready = analyze_training_readiness(dataset)
+    print_sample_inspection(dataset)
+
+    print("\n" + "="*70)
+    print(" ANALYSIS COMPLETE")
+    print("="*70 + "\n")
+
+    sys.exit(0 if ready else 1)
+
+
+if __name__ == "__main__":
+    main()

datasets/create_dataset.py

@@ -8,6 +8,7 @@ import sys
 import re
 import time
 import gc
+import logging
 import multiprocessing as mp
 from pathlib import Path
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -15,6 +16,43 @@ import numpy as np
 import requests
 import torch
 
+# Global logger
+logger = None
+
+def setup_logging(log_file=None):
+    """Setup logging to both console and file."""
+    global logger
+    logger = logging.getLogger("dataset_generator")
+    logger.setLevel(logging.INFO)
+    logger.handlers.clear()
+
+    # Console handler
+    console = logging.StreamHandler(sys.stdout)
+    console.setLevel(logging.INFO)
+    console.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(console)
+
+    # File handler (if specified)
+    if log_file:
+        file_handler = logging.FileHandler(log_file, mode='a')
+        file_handler.setLevel(logging.INFO)
+        file_handler.setFormatter(logging.Formatter('%(asctime)s | %(message)s', datefmt='%H:%M:%S'))
+        logger.addHandler(file_handler)
+
+    return logger
+
+def log(msg):
+    """Log message to both console and file."""
+    global logger
+    if logger:
+        logger.info(msg)
+        # Force flush all handlers
+        for handler in logger.handlers:
+            handler.flush()
+    else:
+        print(msg)
+        sys.stdout.flush()
+
 
 def load_dotenv(env_path=None):
     """Load environment variables from .env file."""
@@ -31,7 +69,7 @@ def load_dotenv(env_path=None):
                 if line and not line.startswith("#") and "=" in line:
                     key, value = line.split("=", 1)
                     os.environ.setdefault(key.strip(), value.strip())
-        print(f"[ENV] Loaded from {env_path}")
+        log(f"[ENV] Loaded from {env_path}")
 
 
 # Load .env file
@@ -40,7 +78,7 @@ load_dotenv()
 # Configuration
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
 GROQ_MODEL = "openai/gpt-oss-20b"
-GROQ_PARALLEL_REQUESTS = 3
+GROQ_PARALLEL_REQUESTS = 10
 QA_PER_REQUEST = 100
 
 
@@ -64,7 +102,7 @@ MEMORY_PER_ITEM = {
 # Default batch sizes (will be auto-adjusted based on VRAM)
 DEFAULT_BATCH_SIZES = {
     "tts": 200,        # TTS batch (processes questions + answers)
-    "whisper": 8,      # Whisper parallel threads (CPU-bound, low GPU mem)
+    "whisper": 8,      # Whisper parallel workers (each uses GPU)
     "snac": 50,        # SNAC encoding batch
 }
 
@@ -84,13 +122,13 @@ def calculate_batch_sizes(vram_gb: float = None, shared_gpu: bool = True) -> dic
     so batch sizes must be more conservative to avoid OOM.
 
     VRAM tiers (with shared GPU adjustment):
-    - 80GB+ (H100/H200): Full TTS, conservative SNAC
-    - 40-80GB (A100/H100-64GB): 75% TTS, conservative SNAC
-    - 24-40GB (RTX 4090): 50% TTS, minimal SNAC
-    - 16-24GB (RTX 3090/4080): 33% TTS, minimal SNAC
+    - 80GB+ (H100/H200): Large batches
+    - 40-80GB (A100/H100-64GB): Medium-large batches
+    - 24-40GB (RTX 4090): Medium batches
+    - 16-24GB (RTX 3090/4080): Smaller batches
     - <16GB: Minimum safe values
 
-    Returns dict with: tts, whisper, snac batch sizes
+    Returns dict with: tts, whisper (workers), snac batch sizes
     """
     if vram_gb is None:
         vram_gb = get_gpu_vram_gb()
@@ -98,29 +136,35 @@ def calculate_batch_sizes(vram_gb: float = None, shared_gpu: bool = True) -> dic
     # Determine scale factor based on VRAM
     if vram_gb >= 80:
         tts_scale = 1.0
-        snac_scale = 0.5    # SNAC needs to be conservative when sharing GPU
+        whisper_workers = 8  # Parallel Whisper threads
+        snac_scale = 0.6
     elif vram_gb >= 40:
         tts_scale = 0.75
-        snac_scale = 0.4    # More conservative for 64GB GPUs
+        whisper_workers = 6
+        snac_scale = 0.5
     elif vram_gb >= 24:
         tts_scale = 0.5
-        snac_scale = 0.3
+        whisper_workers = 4  # RTX 4090 - moderate parallelism
+        snac_scale = 0.4
     elif vram_gb >= 16:
         tts_scale = 0.33
-        snac_scale = 0.2
+        whisper_workers = 2
+        snac_scale = 0.25
     else:
         tts_scale = 0.2
+        whisper_workers = 1
         snac_scale = 0.15
 
-    # If not sharing GPU, SNAC can use more memory
+    # If not sharing GPU, can use more memory
     if not shared_gpu:
-        snac_scale = min(1.0, snac_scale * 2)
+        snac_scale = min(1.0, snac_scale * 1.5)
+        whisper_workers = min(8, whisper_workers + 2)
 
     # Calculate batch sizes
     batch_sizes = {
         "tts": max(10, int(DEFAULT_BATCH_SIZES["tts"] * tts_scale)),
-        "whisper": max(2, int(DEFAULT_BATCH_SIZES["whisper"] * tts_scale)),
-        "snac": max(4, int(DEFAULT_BATCH_SIZES["snac"] * snac_scale)),
+        "whisper": whisper_workers,  # Thread count for parallel Whisper
+        "snac": max(8, int(DEFAULT_BATCH_SIZES["snac"] * snac_scale)),
     }
 
     return batch_sizes
@@ -128,8 +172,8 @@ def calculate_batch_sizes(vram_gb: float = None, shared_gpu: bool = True) -> dic
 
 def print_batch_config(batch_sizes: dict, vram_gb: float):
     """Print batch configuration for transparency."""
-    print(f"[Config] GPU VRAM: {vram_gb:.1f}GB")
-    print(f"[Config] Batch sizes: TTS={batch_sizes['tts']}, Whisper={batch_sizes['whisper']}, SNAC={batch_sizes['snac']}")
+    log(f"[Config] GPU VRAM: {vram_gb:.1f}GB")
+    log(f"[Config] Batch sizes: TTS={batch_sizes['tts']}, Whisper={batch_sizes['whisper']}, SNAC={batch_sizes['snac']}")
 
 
 def parse_qa(content):
@@ -217,14 +261,14 @@ def qa_producer(target_count: int, tts_queue: mp.Queue, batch_size: int, status_
     status_queue.put(("qa_done", len(pairs), batch_idx))
 
 
-def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1):
+def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1, gpu_offset: int = 0):
     """TTS worker - processes batches with VRAM-adjusted batch sizes on specific GPU."""
     import torch
     _orig_load = torch.load
     torch.load = lambda *a, **kw: _orig_load(*a, **{**kw, 'weights_only': False})
 
-    # Set specific GPU for this worker (distribute across available GPUs)
-    actual_gpu = gpu_id % num_gpus
+    # Set specific GPU for this worker (distribute across available GPUs with offset)
+    actual_gpu = gpu_offset + (gpu_id % num_gpus)
     torch.cuda.set_device(actual_gpu)
     print(f"[TTS-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}")
 
@@ -240,19 +284,30 @@ def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_qu
     from soprano import SopranoTTS
     # Use lmdeploy backend for 2000x real-time speed (much faster than transformers)
     # Scale decoder_batch_size based on VRAM
-    try:
-        vram_gb = torch.cuda.get_device_properties(actual_gpu).total_memory / 1024**3
-        dec_batch = 32 if vram_gb >= 80 else (16 if vram_gb >= 40 else 8)
-        tts = SopranoTTS(
-            backend="lmdeploy",  # Fastest backend
-            device=f"cuda:{actual_gpu}",  # Specific GPU
-            cache_size_mb=4000 if vram_gb >= 24 else 2000,  # More cache = faster
-            decoder_batch_size=dec_batch,  # Parallel decoding based on VRAM
-        )
-        print(f"[TTS-GPU{gpu_id}] Using lmdeploy backend (decoder_batch={dec_batch}, cache={4000 if vram_gb >= 24 else 2000}MB)")
-    except Exception as e:
-        print(f"[TTS-GPU{gpu_id}] lmdeploy failed ({e}), falling back to transformers")
-        tts = SopranoTTS(backend="transformers", device=f"cuda:{actual_gpu}")
+    # Note: Soprano TTS only accepts "cuda", not "cuda:N"
+    # torch.cuda.set_device() already selected the correct GPU above
+    vram_gb = torch.cuda.get_device_properties(actual_gpu).total_memory / 1024**3
+    gpu_name = torch.cuda.get_device_properties(actual_gpu).name
+
+    # Check if GPU supports lmdeploy (Blackwell/RTX 50xx not supported yet)
+    use_lmdeploy = "5090" not in gpu_name and "5080" not in gpu_name and "B100" not in gpu_name
+
+    if use_lmdeploy:
+        try:
+            dec_batch = 32 if vram_gb >= 80 else (16 if vram_gb >= 40 else 8)
+            tts = SopranoTTS(
+                backend="lmdeploy",  # Fastest backend
+                device="cuda",  # Uses current device set by torch.cuda.set_device()
+                cache_size_mb=4000 if vram_gb >= 24 else 2000,  # More cache = faster
+                decoder_batch_size=dec_batch,  # Parallel decoding based on VRAM
+            )
+            print(f"[TTS-GPU{gpu_id}] Using lmdeploy backend (decoder_batch={dec_batch})")
+        except Exception as e:
+            print(f"[TTS-GPU{gpu_id}] lmdeploy failed ({e}), falling back to transformers")
+            tts = SopranoTTS(backend="transformers", device="cuda")
+    else:
+        print(f"[TTS-GPU{gpu_id}] Blackwell GPU detected ({gpu_name}), using transformers backend")
+        tts = SopranoTTS(backend="transformers", device="cuda")
 
     # Use centralized batch size
     tts_batch = batch_sizes.get("tts", 200)
@@ -262,8 +317,20 @@ def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_qu
 
     processed = 0
     t_start = time.time()
+    last_heartbeat = time.time()
+
     while True:
-        item = tts_queue.get()
+        # Send heartbeat every 30 seconds to show worker is alive
+        if time.time() - last_heartbeat > 30:
+            status_queue.put(("tts_heartbeat", gpu_id, processed))
+            last_heartbeat = time.time()
+
+        # Use timeout to allow heartbeat even when queue is empty
+        try:
+            item = tts_queue.get(timeout=5)
+        except:
+            continue  # Timeout, send heartbeat and retry
+
         if item is None:
             break
 
@@ -291,7 +358,10 @@ def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_qu
                 })
 
         except Exception as e:
-            status_queue.put(("tts_error", gpu_id, str(e)))
+            import traceback
+            status_queue.put(("tts_error", gpu_id, f"{str(e)}\n{traceback.format_exc()}"))
+            # Continue processing other batches instead of crashing
+            continue
 
         if all_results:
             feat_queue.put((batch_idx, all_results))
@@ -304,7 +374,7 @@ def tts_worker(gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_qu
     status_queue.put(("tts_done", gpu_id, processed))
 
 
-def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1):
+def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: dict, num_gpus: int = 1, gpu_offset: int = 0):
     """Features worker - VRAM-adjusted batch sizes for Whisper + SNAC + NeMo NFA on specific GPU."""
     import torch
     _orig = torch.load
@@ -317,8 +387,8 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
     from transformers import AutoTokenizer, WhisperModel, WhisperFeatureExtractor
     from huggingface_hub import login
 
-    # Set specific GPU for this worker (distribute across available GPUs)
-    actual_gpu = gpu_id % num_gpus
+    # Set specific GPU for this worker (distribute across available GPUs with offset)
+    actual_gpu = gpu_offset + (gpu_id % num_gpus)
     torch.cuda.set_device(actual_gpu)
     device = f"cuda:{actual_gpu}"
     print(f"[Features-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}")
@@ -335,6 +405,9 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
     print("[Features] Loading Whisper large-v3-turbo (transformers)...")
     whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3-turbo", torch_dtype=torch.float16).to(device).eval()
     whisper_feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3-turbo")
+
+    # Note: torch.compile is skipped for Whisper as it causes issues with conv1d
+    # The batched processing already provides significant speedup
     print("[Features] Whisper Turbo loaded successfully")
 
     snac_model = snac.SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
@@ -371,20 +444,26 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
     status_queue.put(("feat_ready", gpu_id))
 
     def process_whisper(audio_data):
-        # Resample to 16kHz for Whisper
-        audio_16k = torchaudio.functional.resample(
-            torch.from_numpy(audio_data), 32000, 16000
-        ).numpy().astype(np.float32)
+        """Process single audio with Whisper. GPU-accelerated resampling."""
+        # GPU-accelerated resampling
+        audio_tensor = torch.from_numpy(audio_data).to(device)
+        audio_16k = torchaudio.functional.resample(audio_tensor, 32000, 16000)
+
         # Truncate to max 30 seconds (480000 samples at 16kHz) - Whisper limit
         max_samples = 480000
-        if len(audio_16k) > max_samples:
+        if audio_16k.shape[0] > max_samples:
             audio_16k = audio_16k[:max_samples]
+
+        audio_16k_np = audio_16k.cpu().numpy().astype(np.float32)
+
         # Extract features using Whisper feature extractor
-        inputs = whisper_feature_extractor(audio_16k, sampling_rate=16000, return_tensors="pt")
+        inputs = whisper_feature_extractor(audio_16k_np, sampling_rate=16000, return_tensors="pt")
         input_features = inputs.input_features.to(device, dtype=torch.float16)
+
         # Encode with Whisper
         with torch.no_grad():
             encoder_outputs = whisper_model.encoder(input_features)
+
         # Return encoder hidden states [seq_len, 1280]
         return encoder_outputs.last_hidden_state.squeeze(0).cpu().float()
 
@@ -493,28 +572,55 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
         return None
 
     def get_word_alignments_proportional(audio_data, text, sample_rate=32000):
-        """Fallback: proportional word alignment based on character count."""
+        """Fallback: proportional word alignment based on character count.
+
+        Calculates frame indices that match actual SNAC output:
+        - SNAC operates at 24kHz with ~320 samples per frame (75 fps)
+        - Audio is resampled from sample_rate to 24kHz before SNAC
+        """
         words = text.split()
         if not words:
             return []
 
-        audio_duration = len(audio_data) / sample_rate
+        # Calculate actual SNAC frame count after resampling to 24kHz
+        # SNAC uses ~320 samples per frame at 24kHz
+        audio_24k_samples = len(audio_data) * 24000 / sample_rate
+        total_snac_frames = int(audio_24k_samples / 320)
+
+        if total_snac_frames == 0:
+            return []
+
         total_chars = sum(len(w) for w in words)
         if total_chars == 0:
             return []
 
+        audio_duration = len(audio_data) / sample_rate
         word_alignments = []
-        current_time = 0.0
-        for word in words:
-            word_duration = (len(word) / total_chars) * audio_duration
-            start_time = current_time
-            end_time = current_time + word_duration
-            start_frame = int(start_time * 75)
-            end_frame = int(end_time * 75)
-
-            word_tokens = []
-            if tokenizer is not None:
-                word_tokens = tokenizer.encode(word, add_special_tokens=False)
+        current_frame = 0
+
+        # Pre-tokenize all words in batch for efficiency
+        all_word_tokens = []
+        if tokenizer is not None:
+            for word in words:
+                all_word_tokens.append(tokenizer.encode(word, add_special_tokens=False))
+        else:
+            all_word_tokens = [[] for _ in words]
+
+        for i, word in enumerate(words):
+            # Distribute frames proportionally based on character count
+            word_frames = int((len(word) / total_chars) * total_snac_frames)
+
+            # Ensure last word gets remaining frames
+            if i == len(words) - 1:
+                end_frame = total_snac_frames
+            else:
+                end_frame = min(current_frame + max(1, word_frames), total_snac_frames)
+
+            start_frame = current_frame
+
+            # Calculate time from frames (for compatibility)
+            start_time = start_frame / 75.0
+            end_time = end_frame / 75.0
 
             word_alignments.append({
                 'word': word,
@@ -522,17 +628,28 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
                 'end': end_time,
                 'start_frame': start_frame,
                 'end_frame': end_frame,
-                'tokens': word_tokens
+                'tokens': all_word_tokens[i]
             })
-            current_time = end_time
+            current_frame = end_frame
 
         return word_alignments
 
     processed = 0
     t_start = time.time()
+    last_heartbeat = time.time()
 
     while True:
-        item = feat_queue.get()
+        # Send heartbeat every 30 seconds to show worker is alive
+        if time.time() - last_heartbeat > 30:
+            status_queue.put(("feat_heartbeat", gpu_id, processed))
+            last_heartbeat = time.time()
+
+        # Use timeout to allow heartbeat even when queue is empty
+        try:
+            item = feat_queue.get(timeout=5)
+        except:
+            continue  # Timeout, send heartbeat and retry
+
         if item is None:
             break
 
@@ -540,12 +657,12 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
         t0 = time.time()
 
         try:
-            # 1. Parallel Whisper encoding - CPU-bound, low GPU memory
+            # 1. Parallel Whisper encoding with GPU-accelerated resampling
             q_audios = [ad["q_audio"] for ad in audio_batch]
             with ThreadPoolExecutor(max_workers=whisper_workers) as ex:
                 whisper_features = list(ex.map(process_whisper, q_audios))
 
-            # 2. SNAC encoding - batch size from centralized config
+            # 2. SNAC encoding - GPU-batched with GPU resampling
             a_audios = [ad["a_audio"] for ad in audio_batch]
             all_tokens = []
 
@@ -555,8 +672,12 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
 
                 max_len = max(a.shape[0] for a in mini_audios)
                 padded = [np.pad(a, (0, max_len - len(a))) for a in mini_audios]
+
+                # GPU-accelerated resampling for SNAC
                 audios_24k = torch.stack([
-                    torchaudio.functional.resample(torch.from_numpy(a), 32000, 24000)
+                    torchaudio.functional.resample(
+                        torch.from_numpy(a).to(device), 32000, 24000
+                    ).cpu()
                     for a in padded
                 ])
 
@@ -568,6 +689,11 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
 
             torch.cuda.synchronize()
 
+            # Periodic GPU memory cleanup every 100 batches to prevent fragmentation
+            if processed > 0 and processed % (100 * len(audio_batch)) == 0:
+                torch.cuda.empty_cache()
+                gc.collect()
+
             # 3. Build results with pre-computed text tokens and word alignments
             results = []
             for i, ad in enumerate(audio_batch):
@@ -596,15 +722,48 @@ def features_worker(gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, s
             batch_time = time.time() - t0
             batch_rate = len(results) / batch_time if batch_time > 0 else 0
 
-            print(f"[DEBUG] Putting {len(results)} results to queue (batch {batch_idx})...")
-            result_queue.put((batch_idx, results))
-            print(f"[DEBUG] Put complete for batch {batch_idx}")
+            # Put results with timeout to prevent indefinite blocking
+            put_start = time.time()
+            max_put_attempts = 10
+            for attempt in range(max_put_attempts):
+                try:
+                    result_queue.put((batch_idx, results), timeout=30)
+                    break
+                except Exception as put_err:
+                    if attempt < max_put_attempts - 1:
+                        status_queue.put(("feat_warn", gpu_id, f"Queue full, retry {attempt+1}"))
+                        time.sleep(1)
+                    else:
+                        status_queue.put(("feat_error", gpu_id, f"Failed to put results after {max_put_attempts} attempts"))
+                        raise put_err
+
             processed += len(results)
             elapsed = time.time() - t_start
             status_queue.put(("feat", gpu_id, processed, processed/elapsed, batch_rate))
 
+            # Clear intermediate tensors to prevent memory accumulation
+            del whisper_features, all_tokens, results
+            if processed % 500 == 0:  # More aggressive cleanup every 500 items
+                torch.cuda.empty_cache()
+                gc.collect()
+
         except Exception as e:
-            status_queue.put(("feat_error", gpu_id, str(e)))
+            import traceback
+            error_msg = str(e)
+            status_queue.put(("feat_error", gpu_id, f"{error_msg}\n{traceback.format_exc()}"))
+
+            # Clear GPU memory aggressively
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+            gc.collect()
+
+            # If OOM, try to recover by reducing batch sizes
+            if "out of memory" in error_msg.lower() or "OOM" in error_msg:
+                status_queue.put(("feat_warn", gpu_id, "OOM detected, clearing memory..."))
+                time.sleep(2)  # Give GPU time to recover
+                torch.cuda.empty_cache()
+
+            continue
 
     status_queue.put(("feat_done", gpu_id, processed))
 
@@ -619,9 +778,12 @@ def main():
 
     import argparse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--count", type=int, default=100)
+    parser.add_argument("--count", "--num_samples", type=int, default=100, dest="count")
     parser.add_argument("--output", type=str, default="./data/dataset.pt")
     parser.add_argument("--gpus", type=int, default=NUM_GPUS)
+    parser.add_argument("--resume", action="store_true", help="Resume from existing checkpoint")
+    parser.add_argument("--checkpoint-interval", type=int, default=1000, help="Save checkpoint every N items")
+    parser.add_argument("--log-file", type=str, default=None, help="Log file path (default: output.log)")
     # Optional overrides for batch sizes (if not set, auto-calculated from VRAM)
     parser.add_argument("--tts-batch", type=int, default=None)
     parser.add_argument("--snac-batch", type=int, default=None)
@@ -630,6 +792,11 @@ def main():
 
     Path(args.output).parent.mkdir(parents=True, exist_ok=True)
 
+    # Setup logging (default to output path + .log)
+    log_file = args.log_file or (args.output + ".log")
+    setup_logging(log_file)
+    log(f"[Log] Writing to {log_file}")
+
     # Calculate batch sizes based on GPU VRAM
     vram_gb = get_gpu_vram_gb()
     batch_sizes = calculate_batch_sizes(vram_gb)
@@ -642,42 +809,115 @@ def main():
     if args.whisper_workers is not None:
         batch_sizes["whisper"] = args.whisper_workers
 
-    print("=" * 60)
-    print("Dataset Generator - Fully Async Pipeline")
-    print(f"Target: {args.count} items, GPUs: {args.gpus}")
+    log("=" * 60)
+    log("Dataset Generator - Fully Async Pipeline")
+    log(f"Target: {args.count} items, GPUs: {args.gpus}")
     print_batch_config(batch_sizes, vram_gb)
-    print("=" * 60)
+    log("=" * 60)
+
+    # Resume from checkpoint or output file if exists
+    checkpoint_path = args.output + ".checkpoint"
+    existing_items = []
+    start_count = 0
+
+    # Track resume state
+    resume_from_path = None
+
+    if args.resume:
+        # Check for existing data to resume from
+        # Priority: .new checkpoint (partial new items) + output file, or just output file
+
+        new_checkpoint = checkpoint_path + ".new"
+        new_items_count = 0
+        base_count = 0
+
+        # Check if we have new items checkpoint
+        if Path(new_checkpoint).exists():
+            try:
+                data = torch.load(new_checkpoint, map_location="cpu", weights_only=False, mmap=True)
+                new_items_count = len(data)
+                del data
+                log(f"[Resume] Found {new_items_count} new items in checkpoint")
+            except Exception as e:
+                log(f"[Resume] Failed to read {new_checkpoint}: {e}")
+
+        # Check output file for base items
+        if Path(args.output).exists():
+            try:
+                data = torch.load(args.output, map_location="cpu", weights_only=False, mmap=True)
+                base_count = len(data)
+                del data
+                resume_from_path = args.output
+                log(f"[Resume] Found {base_count} base items in {args.output}")
+            except Exception as e:
+                log(f"[Resume] Failed to read {args.output}: {e}")
+
+        # Total count is base + new
+        start_count = base_count + new_items_count
+        if start_count > 0:
+            log(f"[Resume] Total: {start_count} items ({base_count} base + {new_items_count} new), need {args.count - start_count} more")
+        else:
+            log("[Resume] No valid resume file found, starting fresh")
+
+    if start_count >= args.count:
+        log(f"[Resume] Already have {start_count} items, saving final...")
+        torch.save(existing_items[:args.count], args.output)
+        log(f"COMPLETE: {args.count} items saved to {args.output}")
+        return
+
+    remaining_count = args.count - start_count
+    log(f"[Main] Generating {remaining_count} new items...")
 
     total_start = time.time()
 
-    # Use Manager for more reliable cross-process queue communication
-    manager = mp.Manager()
-    tts_queue = mp.Queue(maxsize=100)
-    feat_queue = mp.Queue(maxsize=100)
-    result_queue = manager.Queue()  # Manager queue is more reliable with spawn
+    # Use regular mp.Queue with NO size limits to prevent deadlocks
+    # Workers will block on put() if queue is full, causing stalls
+    # Memory is managed by batch sizes instead
+    tts_queue = mp.Queue()  # No maxsize - prevents TTS blocking
+    feat_queue = mp.Queue()  # No maxsize - prevents Features blocking
+    result_queue = mp.Queue()  # No maxsize - prevents result collection blocking
     status_queue = mp.Queue()
 
     workers = []
 
     # Get actual GPU count for worker assignment
     actual_num_gpus = get_num_gpus() if torch.cuda.is_available() else 1
-    print(f"\n[Main] Detected {actual_num_gpus} GPUs, spawning {args.gpus} workers per stage")
 
-    qa_proc = mp.Process(target=qa_producer, args=(args.count, tts_queue, batch_sizes["tts"], status_queue, args.gpus))
+    # With 4+ GPUs, separate TTS and Features to avoid OOM
+    # TTS uses GPUs 0 to (N/2-1), Features uses GPUs (N/2) to (N-1)
+    if actual_num_gpus >= 4:
+        tts_gpus = actual_num_gpus // 2  # First half for TTS
+        feat_gpus = actual_num_gpus - tts_gpus  # Second half for Features
+        tts_gpu_offset = 0
+        feat_gpu_offset = tts_gpus
+        log(f"\n[Main] Detected {actual_num_gpus} GPUs - Separating: TTS on GPUs 0-{tts_gpus-1}, Features on GPUs {feat_gpu_offset}-{actual_num_gpus-1}")
+    else:
+        tts_gpus = actual_num_gpus
+        feat_gpus = actual_num_gpus
+        tts_gpu_offset = 0
+        feat_gpu_offset = 0
+        log(f"\n[Main] Detected {actual_num_gpus} GPUs, sharing between TTS and Features")
+
+    # Adjust number of workers based on available GPUs
+    num_tts_workers = min(args.gpus, tts_gpus)
+    num_feat_workers = min(args.gpus, feat_gpus)
+    log(f"[Main] Spawning {num_tts_workers} TTS workers, {num_feat_workers} Features workers")
+
+    qa_proc = mp.Process(target=qa_producer, args=(remaining_count, tts_queue, batch_sizes["tts"], status_queue, num_tts_workers))
     qa_proc.start()
     workers.append(qa_proc)
 
-    for gpu_id in range(args.gpus):
-        p = mp.Process(target=tts_worker, args=(gpu_id, tts_queue, feat_queue, status_queue, batch_sizes, actual_num_gpus))
+    for gpu_id in range(num_tts_workers):
+        p = mp.Process(target=tts_worker, args=(gpu_id, tts_queue, feat_queue, status_queue, batch_sizes, tts_gpus, tts_gpu_offset))
         p.start()
         workers.append(p)
 
-    for gpu_id in range(args.gpus):
-        p = mp.Process(target=features_worker, args=(gpu_id, feat_queue, result_queue, status_queue, batch_sizes, actual_num_gpus))
+    for gpu_id in range(num_feat_workers):
+        p = mp.Process(target=features_worker, args=(gpu_id, feat_queue, result_queue, status_queue, batch_sizes, feat_gpus, feat_gpu_offset))
         p.start()
         workers.append(p)
 
-    print("[Pipeline] All workers started, monitoring...")
+    log("[Pipeline] All workers started, monitoring...")
 
     results = {}
     tts_done_count = 0
@@ -686,124 +926,311 @@ def main():
     t0 = time.time()
     tts_ready = 0
     feat_ready = 0
-    expected_from_feat = {}  # Track expected items per GPU
-    feat_queue_closed = False  # Track if we've sent None to feat workers
+    expected_from_feat = {}
+    feat_queue_closed = False
+    last_checkpoint_count = 0
 
-    # Main loop: alternate between checking status and collecting results
+    # Main loop with improved stall detection
     last_result_time = time.time()
+    last_status_time = time.time()
+    last_heartbeat_time = {f"tts_{i}": time.time() for i in range(num_tts_workers)}
+    last_heartbeat_time.update({f"feat_{i}": time.time() for i in range(num_feat_workers)})
+    stall_warning_shown = False
+    errors = []
+
     while True:
-        # Process status messages (non-blocking)
-        status_processed = False
-        while True:
+        # Process ALL pending status messages
+        for _ in range(100):  # Limit to prevent infinite loop
             try:
                 msg = status_queue.get_nowait()
-                status_processed = True
+                last_status_time = time.time()
                 msg_type = msg[0]
 
                 if msg_type == "tts_ready":
                     tts_ready += 1
-                    print(f"[TTS-GPU{msg[1]}] Ready ({tts_ready}/{args.gpus})")
+                    log(f"[TTS-GPU{msg[1]}] Ready ({tts_ready}/{num_tts_workers})")
                 elif msg_type == "feat_ready":
                     feat_ready += 1
-                    print(f"[Features-GPU{msg[1]}] Ready ({feat_ready}/{args.gpus})")
+                    log(f"[Features-GPU{msg[1]}] Ready ({feat_ready}/{num_feat_workers})")
                 elif msg_type == "qa":
-                    print(f"[Q&A] {msg[1]}/{msg[2]} | {msg[3]:.1f}/s")
+                    log(f"[Q&A] {msg[1]}/{msg[2]} | {msg[3]:.1f}/s")
                 elif msg_type == "qa_done":
-                    print(f"[Q&A] Done: {msg[1]} pairs, {msg[2]} batches")
+                    log(f"[Q&A] Done: {msg[1]} pairs, {msg[2]} batches")
                 elif msg_type == "tts":
-                    print(f"[TTS-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
+                    log(f"[TTS-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
+                    last_heartbeat_time[f"tts_{msg[1]}"] = time.time()
                 elif msg_type == "tts_done":
                     tts_done_count += 1
-                    print(f"[TTS-GPU{msg[1]}] Done: {msg[2]} items")
+                    log(f"[TTS-GPU{msg[1]}] Done: {msg[2]} items")
                 elif msg_type == "feat":
-                    print(f"[Feat-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
+                    log(f"[Feat-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s | batch {msg[4]:.1f}/s")
+                    last_heartbeat_time[f"feat_{msg[1]}"] = time.time()
                 elif msg_type == "feat_done":
                     feat_done_count += 1
                     expected_from_feat[msg[1]] = msg[2]
-                    print(f"[Features-GPU{msg[1]}] Done: {msg[2]} items")
+                    log(f"[Features-GPU{msg[1]}] Done: {msg[2]} items")
+                elif msg_type == "tts_heartbeat":
+                    last_heartbeat_time[f"tts_{msg[1]}"] = time.time()
+                elif msg_type == "feat_heartbeat":
+                    last_heartbeat_time[f"feat_{msg[1]}"] = time.time()
+                elif msg_type == "feat_warn":
+                    log(f"[WARN] Features-GPU{msg[1]}: {msg[2]}")
                 elif "error" in msg_type:
-                    print(f"[Error] {msg}")
+                    log(f"[Error] {msg}")
+                    errors.append(msg)
             except:
                 break
 
-        # TTS workers now send None to feat_queue directly when done
-        # This ensures all batches are sent before close signal
-        if tts_done_count >= args.gpus and not feat_queue_closed:
-            print(f"[Main] TTS done, waiting for features to finish...")
+        if tts_done_count >= num_tts_workers and not feat_queue_closed:
+            log(f"[Main] TTS done, waiting for features to finish...")
             feat_queue_closed = True
 
-        # Try to collect results (blocking with short timeout)
-        try:
-            batch_idx, items = result_queue.get(timeout=0.1)
-            results[batch_idx] = items
-            total_items += len(items)
-            last_result_time = time.time()
+        # Collect ALL available results (drain aggressively with non-blocking gets)
+        collected_this_round = 0
+        drain_start = time.time()
+        while time.time() - drain_start < 2.0:  # Spend up to 2s draining
+            try:
+                batch_idx, items = result_queue.get_nowait()
+                results[batch_idx] = items
+                total_items += len(items)
+                collected_this_round += len(items)
+                last_result_time = time.time()
+                stall_warning_shown = False
+            except:
+                # No more items immediately available, wait briefly then check again
+                time.sleep(0.05)
+                try:
+                    batch_idx, items = result_queue.get_nowait()
+                    results[batch_idx] = items
+                    total_items += len(items)
+                    collected_this_round += len(items)
+                    last_result_time = time.time()
+                    stall_warning_shown = False
+                except:
+                    break  # Queue truly empty
+
+        if collected_this_round > 0:
             elapsed = time.time() - t0
-            print(f"[Results] {total_items}/{args.count} | {total_items/elapsed:.1f}/s")
-        except:
-            pass  # No results available right now
+            log(f"[Results] {total_items}/{remaining_count} | {total_items/elapsed:.1f}/s")
+
+        # Save checkpoint periodically (combine with any existing checkpoint)
+        if total_items - last_checkpoint_count >= args.checkpoint_interval:
+            log(f"[Checkpoint] Saving {total_items} new items...")
+            # Collect items from this run
+            items_this_run = []
+            for i in sorted(results.keys()):
+                items_this_run.extend(results[i])
+
+            # Load and combine with existing checkpoint if present
+            checkpoint_new_path = checkpoint_path + ".new"
+            all_checkpoint_items = []
+            if Path(checkpoint_new_path).exists():
+                try:
+                    prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
+                    all_checkpoint_items = list(prev_items)
+                    del prev_items
+                except:
+                    pass
+
+            # Only add items not already in checkpoint
+            items_to_add = items_this_run[len(all_checkpoint_items):]
+            all_checkpoint_items.extend(items_to_add)
+
+            torch.save(all_checkpoint_items, checkpoint_new_path)
+            last_checkpoint_count = total_items
+            log(f"[Checkpoint] Saved {len(all_checkpoint_items)} total to {checkpoint_new_path}")
 
         # Check exit conditions
-        if total_items >= args.count:
-            print(f"[Main] Target reached: {total_items}/{args.count}")
+        if total_items >= remaining_count:
+            log(f"[Main] Target reached: {total_items}/{remaining_count}")
             break
 
-        # If all workers done, keep draining with longer timeout
-        if feat_done_count >= args.gpus:
-            drain_timeout = 5.0
-            print(f"[Main] All workers done, draining queue (timeout={drain_timeout}s)...")
+        # If all feature workers done, drain remaining results
+        if feat_done_count >= num_feat_workers:
+            log(f"[Main] All workers done, draining queue...")
             drain_start = time.time()
-            while time.time() - drain_start < drain_timeout:
+            while time.time() - drain_start < 60.0:
                 try:
                     batch_idx, items = result_queue.get(timeout=0.5)
                     results[batch_idx] = items
                     total_items += len(items)
-                    print(f"[Results] {total_items}/{args.count} | (drained)")
+                    log(f"[Results] {total_items}/{remaining_count} | (drained)")
                 except:
-                    # No more items, check if we should keep waiting
-                    if time.time() - drain_start > 1.0:
+                    time.sleep(0.2)
+                    try:
+                        batch_idx, items = result_queue.get_nowait()
+                        results[batch_idx] = items
+                        total_items += len(items)
+                        log(f"[Results] {total_items}/{remaining_count} | (drained)")
+                    except:
                         break
             break
 
-        # Safety timeout: if no progress for 60 seconds, exit
-        if time.time() - last_result_time > 60 and total_items > 0:
-            print(f"[WARN] No progress for 60s, stopping with {total_items} items")
+        # Check if workers are still alive and responding
+        alive_workers = sum(1 for p in workers if p.is_alive())
+        time_since_result = time.time() - last_result_time
+        time_since_status = time.time() - last_status_time
+
+        # Check heartbeats - detect stuck workers
+        stuck_workers = []
+        for worker_id, last_hb in last_heartbeat_time.items():
+            if time.time() - last_hb > 120:  # No heartbeat for 2 minutes
+                stuck_workers.append(worker_id)
+
+        if time_since_result > 30 and not stall_warning_shown:
+            queue_size = 0
+            try:
+                queue_size = result_queue.qsize()
+            except:
+                pass
+            log(f"[WARN] No results for 30s, {alive_workers} workers alive, queue ~{queue_size} items")
+            if stuck_workers:
+                log(f"[WARN] Stuck workers (no heartbeat >2min): {stuck_workers}")
+            stall_warning_shown = True
+
+        # Save checkpoint on stall detection (combine with existing checkpoint)
+        if time_since_result > 60 and total_items > last_checkpoint_count:
+            log(f"[Checkpoint] Stall detected, saving {total_items} items...")
+            items_this_run = []
+            for i in sorted(results.keys()):
+                items_this_run.extend(results[i])
+
+            # Load and combine with existing checkpoint
+            checkpoint_new_path = checkpoint_path + ".new"
+            all_checkpoint_items = []
+            if Path(checkpoint_new_path).exists():
+                try:
+                    prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
+                    all_checkpoint_items = list(prev_items)
+                    del prev_items
+                except:
+                    pass
+
+            items_to_add = items_this_run[len(all_checkpoint_items):]
+            all_checkpoint_items.extend(items_to_add)
+            torch.save(all_checkpoint_items, checkpoint_new_path)
+            last_checkpoint_count = total_items
+
+        # Exit if truly stalled - but be smarter about it
+        # Only exit if: no results for 3min AND no heartbeats AND workers dead
+        all_workers_stuck = len(stuck_workers) >= (num_tts_workers + num_feat_workers)
+        if time_since_result > 180 and all_workers_stuck:
+            log(f"[WARN] All workers stuck, stopping with {total_items} items")
+            log(f"[WARN] Errors encountered: {len(errors)}")
+            for err in errors[-5:]:  # Show last 5 errors
+                log(f"  {err}")
+            break
+
+        if time_since_result > 600:
+            log(f"[WARN] No progress for 10min, stopping with {total_items} items")
             break
 
+    # Save checkpoint before cleanup (combine with existing checkpoint)
+    if total_items > last_checkpoint_count:
+        log(f"[Checkpoint] Final save before cleanup: {total_items} items...")
+        items_this_run = []
+        for i in sorted(results.keys()):
+            items_this_run.extend(results[i])
+
+        # Load and combine with existing checkpoint
+        checkpoint_new_path = checkpoint_path + ".new"
+        all_checkpoint_items = []
+        if Path(checkpoint_new_path).exists():
+            try:
+                prev_items = torch.load(checkpoint_new_path, map_location="cpu", weights_only=False)
+                all_checkpoint_items = list(prev_items)
+                del prev_items
+            except:
+                pass
+
+        items_to_add = items_this_run[len(all_checkpoint_items):]
+        all_checkpoint_items.extend(items_to_add)
+        torch.save(all_checkpoint_items, checkpoint_new_path)
+        last_checkpoint_count = total_items
+
     # Wait for workers to finish
-    print("[Main] Waiting for workers to finish...")
+    log("[Main] Waiting for workers to finish...")
     for p in workers:
         p.join(timeout=5)
         if p.is_alive():
             p.terminate()
 
-    # Final drain (quick)
-    while True:
+    # Final drain
+    for _ in range(100):
         try:
             batch_idx, items = result_queue.get_nowait()
             results[batch_idx] = items
             total_items += len(items)
-            print(f"[Results] {total_items}/{args.count} | (final)")
         except:
             break
 
-    items = []
+    # Collect new results from this run
+    new_items_this_run = []
     for i in sorted(results.keys()):
-        items.extend(results[i])
-    items = items[:args.count]
+        new_items_this_run.extend(results[i])
+
+    log(f"[Main] Generated {len(new_items_this_run)} new items this run")
 
-    if len(items) == 0:
-        print("[ERROR] No items collected! Check worker logs.")
+    # Load any previously checkpointed new items and combine with this run's items
+    new_checkpoint = checkpoint_path + ".new"
+    new_items = []
+
+    if Path(new_checkpoint).exists():
+        try:
+            prev_new = torch.load(new_checkpoint, map_location="cpu", weights_only=False)
+            new_items = list(prev_new)
+            log(f"[Main] Loaded {len(new_items)} items from previous checkpoint")
+            del prev_new
+        except Exception as e:
+            log(f"[Main] Failed to load previous checkpoint: {e}")
+
+    # Add this run's items to the checkpoint items
+    new_items.extend(new_items_this_run)
+    log(f"[Main] Total new items: {len(new_items)} (checkpoint: {len(new_items) - len(new_items_this_run)}, this run: {len(new_items_this_run)})")
+
+    total_new = len(new_items)
+
+    if total_new == 0 and (not resume_from_path or start_count == 0):
+        log("[ERROR] No items collected! Check worker logs.")
         sys.exit(1)
 
-    torch.save(items, args.output)
+    # Combine with base items if resuming
+    if resume_from_path and Path(resume_from_path).exists():
+        log(f"[Main] Loading base items from {resume_from_path}...")
+        base_data = torch.load(resume_from_path, map_location="cpu", weights_only=False, mmap=True)
+        base_count = len(base_data)
+
+        # Calculate how many base items to keep
+        items_needed_from_base = min(base_count, args.count - total_new)
+        log(f"[Main] Combining {items_needed_from_base} base + {total_new} new items...")
+
+        final_items = list(base_data[:items_needed_from_base]) + new_items
+        del base_data
+    else:
+        final_items = new_items
+
+    # Trim to target count
+    final_items = final_items[:args.count]
+
+    # Save final dataset
+    log(f"[Main] Saving {len(final_items)} items to {args.output}...")
+    torch.save(final_items, args.output)
+
+    # Remove checkpoint files if complete
+    if len(final_items) >= args.count:
+        for cp in [checkpoint_path, checkpoint_path + ".new"]:
+            if Path(cp).exists():
+                Path(cp).unlink()
+                log(f"[Cleanup] Removed {cp}")
 
     total_time = time.time() - total_start
-    print("\n" + "=" * 60)
-    print(f"COMPLETE: {len(items)} items saved to {args.output}")
-    print(f"Total time: {total_time:.1f}s ({total_time/60:.1f}m)")
-    print(f"Throughput: {len(items)/total_time:.2f} items/s")
-    print("=" * 60)
+    log("\n" + "=" * 60)
+    log(f"COMPLETE: {len(final_items)} items saved to {args.output}")
+    if start_count > 0:
+        log(f"  (resumed from {start_count}, added {len(final_items) - start_count} new)")
+    log(f"Total time: {total_time:.1f}s ({total_time/60:.1f}m)")
+    log(f"Throughput: {(len(final_items) - start_count)/total_time:.2f} items/s")
+    log("=" * 60)
 
 
 if __name__ == "__main__":

datasets/validate_dataset.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""
+Dataset Validator - Validates the generated dataset for training.
+
+Checks:
+1. Structure: list of dicts with required fields
+2. Whisper features: shape [seq_len, 1280]
+3. SNAC tokens: multiple of 7, valid range
+4. Optional fields: text, answer, text_tokens, word_alignments
+5. Statistics and distribution
+"""
+
+import argparse
+import sys
+from pathlib import Path
+from collections import defaultdict
+
+import torch
+import numpy as np
+
+
+# SNAC token constants (from CLAUDE.md)
+SNAC_BASE = 128266
+SNAC_LAYERS = 7
+SNAC_VOCAB_PER_LAYER = 4096
+WHISPER_DIM = 1280
+
+
+def validate_sample(idx: int, sample: dict, verbose: bool = False) -> tuple[bool, list[str]]:
+    """Validate a single sample. Returns (is_valid, list of errors)."""
+    errors = []
+    warnings = []
+
+    # Check required fields
+    if "whisper_features" not in sample:
+        errors.append("Missing 'whisper_features'")
+    if "snac_tokens" not in sample:
+        errors.append("Missing 'snac_tokens'")
+
+    if errors:
+        return False, errors
+
+    # Validate whisper_features
+    wf = sample["whisper_features"]
+    if not isinstance(wf, torch.Tensor):
+        errors.append(f"whisper_features should be Tensor, got {type(wf).__name__}")
+    else:
+        if wf.dim() != 2:
+            errors.append(f"whisper_features should be 2D [seq_len, 1280], got {wf.dim()}D")
+        elif wf.shape[1] != WHISPER_DIM:
+            errors.append(f"whisper_features dim should be {WHISPER_DIM}, got {wf.shape[1]}")
+        if wf.shape[0] == 0:
+            errors.append("whisper_features has 0 length")
+        if torch.isnan(wf).any():
+            errors.append("whisper_features contains NaN values")
+        if torch.isinf(wf).any():
+            errors.append("whisper_features contains Inf values")
+
+    # Validate snac_tokens
+    st = sample["snac_tokens"]
+    if not isinstance(st, torch.Tensor):
+        errors.append(f"snac_tokens should be Tensor, got {type(st).__name__}")
+    else:
+        if st.dim() != 1:
+            errors.append(f"snac_tokens should be 1D, got {st.dim()}D")
+        if len(st) == 0:
+            errors.append("snac_tokens has 0 length")
+        elif len(st) % SNAC_LAYERS != 0:
+            errors.append(f"snac_tokens length ({len(st)}) not multiple of {SNAC_LAYERS}")
+
+        # Check token range (raw tokens before offset, should be 0-4095)
+        if len(st) > 0:
+            min_tok = st.min().item()
+            max_tok = st.max().item()
+            # Tokens could be raw (0-4095) or with offset applied (128266+)
+            if max_tok < SNAC_BASE:
+                # Raw tokens
+                if min_tok < 0 or max_tok >= SNAC_VOCAB_PER_LAYER:
+                    warnings.append(f"snac_tokens range [{min_tok}, {max_tok}] outside [0, {SNAC_VOCAB_PER_LAYER-1}]")
+            else:
+                # Tokens with offset
+                expected_max = SNAC_BASE + (SNAC_LAYERS * SNAC_VOCAB_PER_LAYER)
+                if min_tok < SNAC_BASE or max_tok >= expected_max:
+                    warnings.append(f"snac_tokens with offset range [{min_tok}, {max_tok}] unexpected")
+
+    # Validate optional fields
+    if "text" in sample and not isinstance(sample["text"], str):
+        warnings.append(f"text should be str, got {type(sample['text']).__name__}")
+
+    if "answer" in sample and not isinstance(sample["answer"], str):
+        warnings.append(f"answer should be str, got {type(sample['answer']).__name__}")
+
+    if "text_tokens" in sample:
+        tt = sample["text_tokens"]
+        if not isinstance(tt, torch.Tensor):
+            warnings.append(f"text_tokens should be Tensor, got {type(tt).__name__}")
+        elif tt.dim() != 1:
+            warnings.append(f"text_tokens should be 1D, got {tt.dim()}D")
+
+    if "word_alignments" in sample:
+        wa = sample["word_alignments"]
+        if not isinstance(wa, list):
+            warnings.append(f"word_alignments should be list, got {type(wa).__name__}")
+
+    if verbose and warnings:
+        for w in warnings:
+            print(f"  [WARN] Sample {idx}: {w}")
+
+    return len(errors) == 0, errors
+
+
+def compute_statistics(dataset: list) -> dict:
+    """Compute dataset statistics."""
+    stats = {
+        "total_samples": len(dataset),
+        "whisper_lengths": [],
+        "snac_lengths": [],
+        "snac_frames": [],
+        "has_text": 0,
+        "has_answer": 0,
+        "has_text_tokens": 0,
+        "has_word_alignments": 0,
+        "text_lengths": [],
+        "answer_lengths": [],
+    }
+
+    for sample in dataset:
+        if "whisper_features" in sample and isinstance(sample["whisper_features"], torch.Tensor):
+            stats["whisper_lengths"].append(sample["whisper_features"].shape[0])
+
+        if "snac_tokens" in sample and isinstance(sample["snac_tokens"], torch.Tensor):
+            length = len(sample["snac_tokens"])
+            stats["snac_lengths"].append(length)
+            stats["snac_frames"].append(length // SNAC_LAYERS)
+
+        if "text" in sample:
+            stats["has_text"] += 1
+            if isinstance(sample["text"], str):
+                stats["text_lengths"].append(len(sample["text"]))
+
+        if "answer" in sample:
+            stats["has_answer"] += 1
+            if isinstance(sample["answer"], str):
+                stats["answer_lengths"].append(len(sample["answer"]))
+
+        if "text_tokens" in sample:
+            stats["has_text_tokens"] += 1
+
+        if "word_alignments" in sample:
+            stats["has_word_alignments"] += 1
+
+    return stats
+
+
+def print_statistics(stats: dict):
+    """Print dataset statistics."""
+    print("\n" + "=" * 60)
+    print("DATASET STATISTICS")
+    print("=" * 60)
+
+    print(f"\nTotal samples: {stats['total_samples']}")
+
+    # Whisper features
+    if stats["whisper_lengths"]:
+        wl = np.array(stats["whisper_lengths"])
+        print(f"\nWhisper features length:")
+        print(f"  Min: {wl.min()}, Max: {wl.max()}, Mean: {wl.mean():.1f}, Std: {wl.std():.1f}")
+
+    # SNAC tokens
+    if stats["snac_lengths"]:
+        sl = np.array(stats["snac_lengths"])
+        sf = np.array(stats["snac_frames"])
+        print(f"\nSNAC tokens:")
+        print(f"  Tokens - Min: {sl.min()}, Max: {sl.max()}, Mean: {sl.mean():.1f}")
+        print(f"  Frames - Min: {sf.min()}, Max: {sf.max()}, Mean: {sf.mean():.1f}")
+
+        # Duration estimate (24kHz, 512 samples per frame = 21.3ms per frame)
+        duration_sec = sf * 0.0213
+        print(f"  Duration - Min: {duration_sec.min():.1f}s, Max: {duration_sec.max():.1f}s, Mean: {duration_sec.mean():.1f}s")
+
+    # Optional fields
+    print(f"\nOptional fields present:")
+    print(f"  text: {stats['has_text']}/{stats['total_samples']} ({100*stats['has_text']/stats['total_samples']:.1f}%)")
+    print(f"  answer: {stats['has_answer']}/{stats['total_samples']} ({100*stats['has_answer']/stats['total_samples']:.1f}%)")
+    print(f"  text_tokens: {stats['has_text_tokens']}/{stats['total_samples']} ({100*stats['has_text_tokens']/stats['total_samples']:.1f}%)")
+    print(f"  word_alignments: {stats['has_word_alignments']}/{stats['total_samples']} ({100*stats['has_word_alignments']/stats['total_samples']:.1f}%)")
+
+    # Text lengths
+    if stats["text_lengths"]:
+        tl = np.array(stats["text_lengths"])
+        print(f"\nText (question) length (chars):")
+        print(f"  Min: {tl.min()}, Max: {tl.max()}, Mean: {tl.mean():.1f}")
+
+    if stats["answer_lengths"]:
+        al = np.array(stats["answer_lengths"])
+        print(f"\nAnswer length (chars):")
+        print(f"  Min: {al.min()}, Max: {al.max()}, Mean: {al.mean():.1f}")
+
+
+def validate_dataset(path: str, max_samples: int = None, verbose: bool = False) -> bool:
+    """Validate the dataset file."""
+    print(f"\nValidating: {path}")
+    print("=" * 60)
+
+    # Check file exists
+    if not Path(path).exists():
+        print(f"[ERROR] File not found: {path}")
+        return False
+
+    # Load dataset
+    print("Loading dataset...")
+    try:
+        dataset = torch.load(path, map_location="cpu", weights_only=False)
+    except Exception as e:
+        print(f"[ERROR] Failed to load dataset: {e}")
+        return False
+
+    # Check type
+    if not isinstance(dataset, list):
+        print(f"[ERROR] Dataset should be a list, got {type(dataset).__name__}")
+        return False
+
+    total = len(dataset)
+    print(f"Loaded {total} samples")
+
+    if total == 0:
+        print("[ERROR] Dataset is empty")
+        return False
+
+    # Validate samples
+    if max_samples and max_samples < total:
+        print(f"Validating first {max_samples} samples...")
+        samples_to_check = dataset[:max_samples]
+    else:
+        print(f"Validating all {total} samples...")
+        samples_to_check = dataset
+
+    valid_count = 0
+    error_counts = defaultdict(int)
+
+    for idx, sample in enumerate(samples_to_check):
+        if not isinstance(sample, dict):
+            print(f"[ERROR] Sample {idx}: should be dict, got {type(sample).__name__}")
+            error_counts["not_dict"] += 1
+            continue
+
+        is_valid, errors = validate_sample(idx, sample, verbose=verbose)
+        if is_valid:
+            valid_count += 1
+        else:
+            for err in errors:
+                error_counts[err] += 1
+                if verbose:
+                    print(f"[ERROR] Sample {idx}: {err}")
+
+    # Summary
+    checked = len(samples_to_check)
+    invalid = checked - valid_count
+
+    print(f"\n{'=' * 60}")
+    print("VALIDATION SUMMARY")
+    print("=" * 60)
+    print(f"Samples checked: {checked}/{total}")
+    print(f"Valid: {valid_count} ({100*valid_count/checked:.1f}%)")
+    print(f"Invalid: {invalid} ({100*invalid/checked:.1f}%)")
+
+    if error_counts:
+        print(f"\nError breakdown:")
+        for err, count in sorted(error_counts.items(), key=lambda x: -x[1]):
+            print(f"  {count:5d}x {err}")
+
+    # Compute and print statistics
+    stats = compute_statistics(samples_to_check)
+    print_statistics(stats)
+
+    # Check a sample for inspection
+    if verbose and valid_count > 0:
+        print(f"\n{'=' * 60}")
+        print("SAMPLE INSPECTION (first valid sample)")
+        print("=" * 60)
+        for idx, sample in enumerate(samples_to_check):
+            is_valid, _ = validate_sample(idx, sample)
+            if is_valid:
+                print(f"Sample {idx}:")
+                for key, value in sample.items():
+                    if isinstance(value, torch.Tensor):
+                        print(f"  {key}: Tensor {value.shape} {value.dtype}")
+                    elif isinstance(value, str):
+                        preview = value[:100] + "..." if len(value) > 100 else value
+                        print(f"  {key}: '{preview}'")
+                    elif isinstance(value, list):
+                        print(f"  {key}: list[{len(value)}]")
+                    else:
+                        print(f"  {key}: {type(value).__name__}")
+                break
+
+    print("\n" + "=" * 60)
+    if invalid == 0:
+        print("RESULT: PASSED - All samples valid")
+        return True
+    else:
+        print(f"RESULT: FAILED - {invalid} invalid samples")
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Validate dataset for training")
+    parser.add_argument("--path", type=str, required=True, help="Path to dataset .pt file")
+    parser.add_argument("--max-samples", type=int, default=None, help="Max samples to validate (default: all)")
+    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
+    args = parser.parse_args()
+
+    success = validate_dataset(args.path, args.max_samples, args.verbose)
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

passo0_setup.py

@@ -5,7 +5,13 @@ Instala todas as dependências para:
 - Geração de dataset (Soprano TTS, Whisper, SNAC, NeMo NFA)
 - Treinamento do modelo Speech-to-Speech
 
-Suporta: H200, H100, A100, RTX 4090, RTX 5090, etc.
+PyTorch 2.7.0 com suporte oficial a Blackwell (sm_120)
+
+Suporta:
+- Blackwell: RTX 5090, 5080, 5070, B100, B200 (sm_120) - CUDA 12.8
+- Hopper: H100, H200 (sm_90) - CUDA 12.8
+- Ada: RTX 4090, 4080, L40S (sm_89) - CUDA 12.4
+- Ampere: A100, RTX 3090 (sm_80/86) - CUDA 12.4
 
 Usage:
     python passo0_setup.py [--skip_test]
@@ -51,13 +57,15 @@ def get_gpu_info():
 
 def needs_cuda_128(gpu_name, compute_cap):
     """Check if GPU needs CUDA 12.8+ (Hopper/Blackwell architecture)."""
-    # H100, H200 = Hopper (sm_90), RTX 50xx = Blackwell (sm_100)
-    hopper_blackwell = ["H100", "H200", "RTX 50", "B100", "B200"]
+    # H100, H200 = Hopper (sm_90)
+    # RTX 50xx, B100, B200 = Blackwell (sm_120) - requires PyTorch 2.7+
+    hopper_blackwell = ["H100", "H200", "RTX 50", "5090", "5080", "5070", "B100", "B200"]
     if any(arch in gpu_name for arch in hopper_blackwell):
         return True
-    # Compute capability 9.0+ needs CUDA 12.8
+    # Compute capability 9.0+ (Hopper) or 12.0+ (Blackwell) needs CUDA 12.8
     try:
-        if float(compute_cap) >= 9.0:
+        cap = float(compute_cap)
+        if cap >= 9.0:  # sm_90 (Hopper) or sm_120 (Blackwell)
             return True
     except:
         pass
@@ -85,13 +93,16 @@ def main():
         run("apt-get update -qq && apt-get install -y -qq espeak-ng espeak libsndfile1 ffmpeg git wget curl build-essential sox libsox-fmt-all", check=False)
 
     # 2. PyTorch (version depends on GPU architecture)
-    log("\n[2/7] PyTorch...")
+    # PyTorch 2.7.0 has official Blackwell (sm_120) support
+    log("\n[2/7] PyTorch 2.7.0 (Blackwell compatible)...")
     if use_cuda_128:
         # Hopper (H100/H200) and Blackwell (RTX 50xx) need CUDA 12.8
-        run("pip install torch==2.8.0 torchaudio==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128 -q")
+        log("  Installing PyTorch 2.7.0 with CUDA 12.8 (Blackwell/Hopper support)")
+        run("pip install torch==2.7.0 torchaudio==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128 -q")
     else:
         # Ampere (A100, RTX 30xx) and Ada (RTX 40xx) can use CUDA 12.4
-        run("pip install torch==2.5.0 torchaudio==2.5.0 torchvision==0.20.0 --index-url https://download.pytorch.org/whl/cu124 -q")
+        log("  Installing PyTorch 2.7.0 with CUDA 12.4 (Ampere/Ada)")
+        run("pip install torch==2.7.0 torchaudio==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu124 -q")
 
     # 3. Core packages
     log("\n[3/7] Core packages...")
@@ -116,11 +127,24 @@ def main():
     # We use native Whisper from transformers (already installed above)
     # No whisperx/ctranslate2 needed - they don't support H200 (sm_90)
 
-    # 6. Soprano TTS + lmdeploy (fastest backend - 2000x realtime)
-    log("\n[6/8] Soprano TTS + lmdeploy...")
+    # 6. Soprano TTS 80M (ultra-lightweight, 2000x realtime)
+    # See: https://github.com/ekwek1/soprano
+    log("\n[6/8] Soprano TTS 80M...")
+
+    # Install soprano-tts with lmdeploy for fastest inference
+    # lmdeploy provides 2000x realtime speed vs ~10x for transformers backend
     run("pip install soprano-tts -q")
-    # lmdeploy is the fastest backend for Soprano TTS (2000x realtime vs ~10x for transformers)
-    run("pip install lmdeploy -q")
+
+    # lmdeploy doesn't support Blackwell (RTX 50xx) yet, so only install on supported GPUs
+    blackwell_gpus = ["RTX 50", "5090", "5080", "5070", "B100", "B200"]
+    is_blackwell = any(arch in gpu_name for arch in blackwell_gpus)
+
+    if is_blackwell:
+        log("  Blackwell GPU detected - skipping lmdeploy (not supported yet)")
+        log("  Soprano will use transformers backend (slower but compatible)")
+    else:
+        log("  Installing lmdeploy for 2000x realtime speed...")
+        run("pip install lmdeploy -q")
 
     # 7. Create directories
     log("\n[7/8] Creating directories...")

passo2_finetune_stage1.py

@@ -8,862 +8,213 @@ Based on IST-LM paper combined with LLaMA-Omni 2 staging:
 - LLM is completely frozen (adapter gets a "head start")
 
 Usage:
-    python finetune_stage1.py --data data.pt --epochs 2
+    python passo2_finetune_stage1.py --data data.pt --epochs 2
 
 Next: Stage 2 trains Adapter + LoRA together
 """
 
-import os
-import sys
 import argparse
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.data import Dataset, DataLoader, ConcatDataset
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from accelerate import Accelerator
-from accelerate.utils import set_seed
-from huggingface_hub import login
-from tqdm import tqdm
-import threading
-
-# ============================================================
-# Config
-# ============================================================
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
-torch.set_float32_matmul_precision('high')
-
-# SNAC token offsets for Orpheus
-SNAC_BASE_OFFSET = 128266
-EOS_TOKEN = 128009
-
-
-def log(msg):
-    print(msg)
-    sys.stdout.flush()
-
-
-def apply_snac_offset(token_idx, position):
-    """Apply position-based offset to SNAC token.
-    If token is already offset (>= SNAC_BASE_OFFSET), return as-is.
-    """
-    if int(token_idx) >= SNAC_BASE_OFFSET:
-        # Already has offset applied
-        return int(token_idx)
-    offset = SNAC_BASE_OFFSET + (position % 7) * 4096
-    return int(token_idx) + offset
+from typing import Dict, Any
+
+from training.config import TrainingConfig, GPUConfig
+from training.models import SpeechAdapter, ModelFactory
+from training.checkpoint import Stage1CheckpointManager, TrainingState
+from training.trainer import BaseTrainer
+from training.utils import log, should_enable_gradient_checkpointing
 
 
-def get_text_ratio(global_step, decay_steps=300, initial_ratio=0.9, min_ratio=0.0):
+class Stage1Trainer(BaseTrainer):
     """
-    IST-LM: Start with 90% text, decrease by 0.1 every 300 steps.
+    Stage 1 Trainer: Adapter only (LLM frozen).
+
+    Trains only the SpeechAdapter while keeping the LLM completely frozen.
+    This gives the adapter a "head start" in learning the audio→embedding mapping.
     """
-    num_decays = global_step // decay_steps
-    text_ratio = initial_ratio - (num_decays * 0.1)
-    return max(min_ratio, text_ratio)
-
-
-# ============================================================
-# Async Checkpoint Saving
-# ============================================================
-_save_threads = []
-
-def save_checkpoint_async(state_dict, path, is_main=True):
-    global _save_threads
-    _save_threads = [t for t in _save_threads if t.is_alive()]
-
-    def copy_to_cpu(obj):
-        if isinstance(obj, torch.Tensor):
-            return obj.detach().cpu().clone()
-        elif isinstance(obj, dict):
-            return {k: copy_to_cpu(v) for k, v in obj.items()}
-        return obj
-
-    state_copy = copy_to_cpu(state_dict)
-
-    def _save():
-        try:
-            torch.save(state_copy, path)
-            if is_main:
-                log(f"[ASYNC] Saved: {path}")
-        except Exception as e:
-            if is_main:
-                log(f"[ASYNC] Error: {e}")
-
-    thread = threading.Thread(target=_save, daemon=True)
-    thread.start()
-    _save_threads.append(thread)
-
-
-def wait_for_checkpoints():
-    global _save_threads
-    for t in _save_threads:
-        t.join()
-    _save_threads = []
-
-
-# ============================================================
-# GPU Auto-Detection (Multi-backend: CUDA, ROCm, MPS, XPU)
-# ============================================================
-def auto_detect_gpu_config():
-    """Detect GPU and return optimal batch size config."""
-    vram_gb = 0
-    gpu_name = "Unknown"
-
-    # Try CUDA (NVIDIA)
-    if torch.cuda.is_available():
-        try:
-            props = torch.cuda.get_device_properties(0)
-            vram_gb = props.total_memory // (1024**3)
-            gpu_name = props.name
-        except:
-            pass
-
-    # Try MPS (Apple Silicon)
-    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-        gpu_name = "Apple Silicon (MPS)"
-        # Apple Silicon unified memory - estimate based on system
-        try:
-            import subprocess
-            result = subprocess.run(['sysctl', '-n', 'hw.memsize'], capture_output=True, text=True)
-            total_mem = int(result.stdout.strip()) // (1024**3)
-            vram_gb = total_mem // 2  # Assume half for GPU
-        except:
-            vram_gb = 8  # Conservative default
-
-    # Try ROCm (AMD) - uses same API as CUDA
-    elif hasattr(torch, 'hip') or os.environ.get('ROCM_HOME'):
-        try:
-            if torch.cuda.is_available():  # ROCm exposes as CUDA
-                props = torch.cuda.get_device_properties(0)
-                vram_gb = props.total_memory // (1024**3)
-                gpu_name = f"AMD {props.name}"
-        except:
-            gpu_name = "AMD ROCm"
-            vram_gb = 16
-
-    # Fallback: try nvidia-smi
-    if vram_gb == 0:
-        try:
-            import subprocess
-            result = subprocess.run(
-                ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader,nounits'],
-                capture_output=True, text=True
-            )
-            lines = result.stdout.strip().split('\n')
-            gpu_name, vram_mb = lines[0].split(', ')
-            vram_gb = int(vram_mb) // 1024
-        except:
-            pass
-
-    # Determine batch size based on VRAM
-    if vram_gb >= 150:
-        return {"name": gpu_name, "batch_size": 8, "grad_accum": 4, "vram_gb": vram_gb}
-    elif vram_gb >= 80:
-        return {"name": gpu_name, "batch_size": 6, "grad_accum": 5, "vram_gb": vram_gb}
-    elif vram_gb >= 35:
-        return {"name": gpu_name, "batch_size": 4, "grad_accum": 8, "vram_gb": vram_gb}
-    elif vram_gb >= 16:
-        return {"name": gpu_name, "batch_size": 2, "grad_accum": 16, "vram_gb": vram_gb}
-    else:
-        return {"name": gpu_name, "batch_size": 1, "grad_accum": 32, "vram_gb": max(vram_gb, 8)}
-
-def get_ram_info():
-    """Get RAM info in GB."""
-    try:
-        import psutil
-        total = psutil.virtual_memory().total / 1024**3
-        available = psutil.virtual_memory().available / 1024**3
-        return total, available
-    except ImportError:
-        try:
-            import subprocess
-            result = subprocess.run(
-                ['free', '-g', '--output=SIZE,AVAILABLE'],
-                capture_output=True, text=True
+
+    def _get_stage_name(self) -> str:
+        return "STAGE 1: Adapter Only (LLM Frozen)"
+
+    def _setup_models(self):
+        """Setup adapter and frozen LLM."""
+        gpu_config = GPUConfig.auto_detect()
+        dtype = gpu_config.dtype
+
+        if self.is_main:
+            log(f"\nLoading LLM (FROZEN): {self.config.model_path}")
+
+        # Check if gradient checkpointing needed
+        use_checkpointing = self.config.gradient_checkpointing or \
+            should_enable_gradient_checkpointing(
+                gpu_config.vram_gb,
+                self.config.dynamic_decay
             )
-            lines = result.stdout.strip().split('\n')
-            if len(lines) >= 2:
-                total, available = map(float, lines[1].split())
-                return total, available
-        except:
-            pass
-    except:
-        pass
-    return 0, 0
-
-def limit_ram_usage(max_ram_gb):
-    """Limit RAM usage by setting resource limits."""
-    try:
-        import resource
-        max_bytes = int(max_ram_gb * 1024**3)
-        resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
-    except:
-        pass
-
-def log_memory_usage():
-    """Log current memory usage."""
-    msg = []
-    if torch.cuda.is_available():
-        used = torch.cuda.memory_allocated() / 1024**3
-        reserved = torch.cuda.memory_reserved() / 1024**3
-        msg.append(f"GPU: {used:.2f}GB / {reserved:.2f}GB")
-    try:
-        import psutil
-        ram_used = psutil.virtual_memory().used / 1024**3
-        ram_total = psutil.virtual_memory().total / 1024**3
-        msg.append(f"RAM: {ram_used:.1f}GB / {ram_total:.1f}GB")
-    except:
-        pass
-    return " | ".join(msg)
-
-
-# ============================================================
-# Speech Adapter (LLaMA-Omni 2 Style)
-# ============================================================
-class SpeechAdapter(nn.Module):
-    """
-    5× downsampling + FFN with intermediate dim 2048
-    """
-    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
-        super().__init__()
-        self.downsample = downsample
-        concat_dim = whisper_dim * downsample
-
-        self.ffn = nn.Sequential(
-            nn.Linear(concat_dim, intermediate_dim),
-            nn.GELU(),
-            nn.Linear(intermediate_dim, llm_dim),
-            nn.LayerNorm(llm_dim)
+
+        # Load frozen LLM
+        self.llm = ModelFactory.create_llm(
+            self.config.model_path,
+            dtype=dtype,
+            freeze=True,
+            gradient_checkpointing=use_checkpointing,
         )
 
-    def forward(self, x):
-        B, T, D = x.shape
-        T_new = (T // self.downsample) * self.downsample
-        x = x[:, :T_new]
-        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
-        return self.ffn(x)
+        if self.is_main and use_checkpointing:
+            log("[Memory] Gradient checkpointing enabled")
 
+        # Create adapter
+        self.adapter = SpeechAdapter().to(dtype=dtype)
 
-# ============================================================
-# Scheduled Interleaved Sequence Creation with Word Alignment
-# ============================================================
-def create_interleaved_sequence(text_tokens, snac_tokens, text_ratio=0.9, word_alignments=None, tokenizer=None, answer_text=None):
-    """
-    Create interleaved sequence based on text_ratio with word-level alignment.
-    - text_ratio=0.9 means 90% of words are replaced by text tokens
-    - text_ratio=0.0 means 100% audio (no text replacement)
+        if self.is_main:
+            params = self.adapter.get_num_params()
+            log(f"\nTrainable: Adapter only ({params:,} = {params/1e6:.1f}M params)")
+            log("LLM: FROZEN")
 
-    With word_alignments: replaces aligned audio spans with corresponding text tokens
-    Without word_alignments: falls back to positional interleaving
-    """
-    interleaved = []
-    is_audio_mask = []
-
-    if len(snac_tokens) == 0:
-        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
-
-    # Group SNAC into frames of 7
-    frames = []
-    for i in range(0, len(snac_tokens), 7):
-        frame = snac_tokens[i:i+7]
-        if len(frame) == 7:
-            frames.append(frame)
-
-    if len(frames) == 0:
-        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
-
-    total_frames = len(frames)
-
-    # If we have word alignments, use semantic interleaving
-    # Check if alignments have pre-computed tokens (preferred) or need tokenizer
-    has_precomputed = word_alignments and len(word_alignments) > 0 and 'tokens' in word_alignments[0] and word_alignments[0]['tokens']
-    can_interleave = word_alignments and text_ratio > 0 and (has_precomputed or (tokenizer and answer_text))
-
-    if can_interleave:
-        import random
-
-        # Decide which words to replace with text based on text_ratio
-        num_words = len(word_alignments)
-        num_text_words = int(num_words * text_ratio)
-
-        # Randomly select which word indices to replace with text
-        word_indices = list(range(num_words))
-        random.shuffle(word_indices)
-        text_word_indices = set(word_indices[:num_text_words])
-
-        # Build interleaved sequence frame by frame
-        frame_idx = 0
-        snac_position = 0
-
-        for word_idx, alignment in enumerate(word_alignments):
-            word = alignment['word']
-            start_frame = alignment['start_frame']
-            end_frame = min(alignment['end_frame'], total_frames)
-
-            if word_idx in text_word_indices:
-                # Replace this word's audio with text tokens
-                # Use pre-computed tokens if available, else tokenize on-the-fly
-                word_tokens = alignment.get('tokens', [])
-                if not word_tokens and tokenizer:
-                    word_tokens = tokenizer.encode(word, add_special_tokens=False)
-
-                for tok in word_tokens:
-                    interleaved.append(tok)
-                    is_audio_mask.append(False)
-                # Skip the audio frames for this word
-                snac_position = end_frame * 7
-            else:
-                # Keep audio for this word
-                for f_idx in range(start_frame, end_frame):
-                    if f_idx < total_frames:
-                        frame = frames[f_idx]
-                        for tok in frame:
-                            interleaved.append(apply_snac_offset(tok, snac_position))
-                            is_audio_mask.append(True)
-                            snac_position += 1
-
-            frame_idx = end_frame
-
-        # Add any remaining frames after the last word
-        while frame_idx < total_frames:
-            frame = frames[frame_idx]
-            for tok in frame:
-                interleaved.append(apply_snac_offset(tok, snac_position))
-                is_audio_mask.append(True)
-                snac_position += 1
-            frame_idx += 1
-
-    else:
-        # Fallback: positional interleaving (original behavior)
-        total_text = len(text_tokens)
-
-        # Determine interleaving pattern based on text_ratio
-        if text_ratio >= 0.9:
-            text_per_chunk, frames_per_chunk = 1, 3
-        elif text_ratio >= 0.7:
-            text_per_chunk, frames_per_chunk = 1, 5
-        elif text_ratio >= 0.5:
-            text_per_chunk, frames_per_chunk = 1, 7
-        elif text_ratio >= 0.3:
-            text_per_chunk, frames_per_chunk = 1, 10
-        else:
-            text_per_chunk, frames_per_chunk = 0, 1
-
-        text_idx = 0
-        frame_idx = 0
-        snac_position = 0
-
-        while frame_idx < total_frames:
-            if text_per_chunk > 0 and text_idx < total_text:
-                for _ in range(text_per_chunk):
-                    if text_idx < total_text:
-                        interleaved.append(text_tokens[text_idx])
-                        is_audio_mask.append(False)
-                        text_idx += 1
-
-            for _ in range(frames_per_chunk):
-                if frame_idx < total_frames:
-                    frame = frames[frame_idx]
-                    for tok in frame:
-                        interleaved.append(apply_snac_offset(tok, snac_position))
-                        is_audio_mask.append(True)
-                        snac_position += 1
-                    frame_idx += 1
-
-        while text_idx < total_text:
-            interleaved.append(text_tokens[text_idx])
-            is_audio_mask.append(False)
-            text_idx += 1
-
-    # Add EOS
-    interleaved.append(EOS_TOKEN)
-    is_audio_mask.append(False)
-
-    return interleaved, is_audio_mask
-
-
-# ============================================================
-# Dataset
-# ============================================================
-class InterleavedDataset(Dataset):
-    def __init__(self, data, tokenizer, max_audio_len=500, max_seq_len=2048):
-        self.data = data
-        self.tokenizer = tokenizer
-        self.max_audio = max_audio_len * 5
-        self.max_seq_len = max_seq_len
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        item = self.data[idx]
-
-        # Whisper features
-        whisper = item["whisper_features"][:self.max_audio]
-
-        # Text tokens - use pre-computed if available, otherwise tokenize
-        if "text_tokens" in item and len(item["text_tokens"]) > 0:
-            tt = item["text_tokens"]
-            text_tokens = tt.tolist() if hasattr(tt, 'tolist') else list(tt)
-        else:
-            text = item.get("answer", item.get("text", ""))
-            if isinstance(text, str) and len(text) > 0:
-                text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
-            else:
-                text_tokens = []
-
-        # SNAC tokens
-        snac = item["snac_tokens"]
-        snac_len = (len(snac) // 7) * 7
-        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
-        snac_list = snac.tolist() if hasattr(snac, 'tolist') else list(snac)
-
-        # Word alignments (if available)
-        word_alignments = item.get("word_alignments", None)
-        answer_text = item.get("answer", "")
-
-        return {
-            "whisper": whisper,
-            "text_tokens": text_tokens,
-            "snac_tokens": snac_list,
-            "word_alignments": word_alignments,
-            "answer_text": answer_text
-        }
-
-
-def collate_fn(batch, text_ratio=0.9, tokenizer=None):
-    """Collate with dynamic interleaving based on text_ratio and word alignments."""
-    max_w = max(b["whisper"].shape[0] for b in batch)
-    max_w = ((max_w + 4) // 5) * 5
-
-    whisper_batch = []
-    interleaved_batch = []
-    is_audio_batch = []
-
-    max_seq = 0
-    sequences = []
-
-    for b in batch:
-        interleaved, is_audio = create_interleaved_sequence(
-            b["text_tokens"],
-            b["snac_tokens"],
-            text_ratio,
-            word_alignments=b.get("word_alignments"),
-            tokenizer=tokenizer,
-            answer_text=b.get("answer_text")
+    def _setup_optimizer(self):
+        """Setup optimizer for adapter only."""
+        self.optimizer = torch.optim.AdamW(
+            self.adapter.parameters(),
+            lr=self.config.learning_rate,
+            weight_decay=self.config.weight_decay,
+        )
+
+    def _setup_checkpoint_manager(self):
+        """Setup Stage 1 checkpoint manager."""
+        self.checkpoint_manager = Stage1CheckpointManager(
+            self.config.output_dir,
+            verbose=self.is_main,
+        )
+
+    def _load_checkpoint(self, ckpt: Dict[str, Any]):
+        """Load checkpoint state."""
+        if "adapter" in ckpt:
+            self.adapter.load_state_dict(ckpt["adapter"])
+        if "optimizer" in ckpt:
+            self.optimizer.load_state_dict(ckpt["optimizer"])
+        if "step" in ckpt:
+            self.global_step = ckpt["step"]
+        if "epoch" in ckpt:
+            self.start_epoch = ckpt["epoch"]
+        if "text_ratio" in ckpt:
+            self.current_text_ratio = ckpt["text_ratio"]
+
+    def _get_trainable_params(self):
+        """Get trainable parameters (adapter only)."""
+        return self.adapter.parameters()
+
+    def _save_step_checkpoint(self, loss: float):
+        """Save step checkpoint."""
+        self.accelerator.wait_for_everyone()
+        unwrapped = self.accelerator.unwrap_model(self.adapter)
+
+        state = TrainingState(
+            step=self.global_step,
+            epoch=self.start_epoch,
+            loss=loss,
+            text_ratio=self.current_text_ratio,
+            best_loss=self.best_loss,
         )
-        sequences.append((interleaved, is_audio))
-        max_seq = max(max_seq, len(interleaved))
 
-    for i, b in enumerate(batch):
-        w = b["whisper"]
-        w_pad = F.pad(w, (0, 0, 0, max_w - w.shape[0]))
-        whisper_batch.append(w_pad)
+        self.checkpoint_manager.save_step(
+            unwrapped.state_dict(),
+            self.optimizer.state_dict(),
+            state,
+        )
 
-        interleaved, is_audio = sequences[i]
-        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
-        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)
+    def _finish_epoch(self, epoch: int, epoch_loss: float):
+        """Finish epoch and save checkpoints."""
+        avg_loss = epoch_loss / max(1, self.steps_per_epoch) / self.config.grad_accum
 
-        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
-        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)
+        if self.is_main:
+            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {self.current_text_ratio:.1f}")
 
-        interleaved_batch.append(seq_pad)
-        is_audio_batch.append(mask_pad)
+            self.accelerator.wait_for_everyone()
+            unwrapped = self.accelerator.unwrap_model(self.adapter)
 
-    return {
-        "whisper": torch.stack(whisper_batch),
-        "interleaved": torch.stack(interleaved_batch),
-        "is_audio_mask": torch.stack(is_audio_batch)
-    }
+            state = TrainingState(
+                step=self.global_step,
+                epoch=epoch + 1,
+                loss=avg_loss,
+                text_ratio=self.current_text_ratio,
+                best_loss=self.best_loss,
+            )
+
+            # Save epoch checkpoint
+            self.checkpoint_manager.save_epoch(
+                unwrapped.state_dict(),
+                self.optimizer.state_dict(),
+                state,
+            )
+
+            # Save best if improved
+            if avg_loss < self.best_loss:
+                self.best_loss = avg_loss
+                state.best_loss = self.best_loss
+                self.checkpoint_manager.save_best(unwrapped.state_dict(), state)
+                log(f"Best model saved! Loss: {self.best_loss:.4f}")
 
 
-# ============================================================
-# Arguments
-# ============================================================
 def parse_args():
-    parser = argparse.ArgumentParser(description="Stage 1: Adapter Only with Interleaved Output")
-    parser.add_argument("--data", type=str, required=True)
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Stage 1: Adapter Only with Interleaved Output"
+    )
+
+    # Data
+    parser.add_argument("--data", type=str, required=True,
+                        help="Dataset path(s), comma-separated")
     parser.add_argument("--output_dir", type=str, default="./checkpoints")
+
+    # Training
     parser.add_argument("--lr", type=float, default=5e-5)
-    parser.add_argument("--epochs", type=int, default=2, help="1-2 epochs for adapter warmup")
+    parser.add_argument("--epochs", type=int, default=2,
+                        help="1-2 epochs for adapter warmup")
     parser.add_argument("--batch_size", type=int, default=None)
     parser.add_argument("--grad_accum", type=int, default=None)
     parser.add_argument("--warmup_ratio", type=float, default=0.03)
     parser.add_argument("--max_grad_norm", type=float, default=1.0)
-    parser.add_argument("--save_steps", type=int, default=200)
     parser.add_argument("--label_smoothing", type=float, default=0.1)
+    parser.add_argument("--max_seq_len", type=int, default=2048)
+
     # Scheduled interleaving
     parser.add_argument("--initial_text_ratio", type=float, default=0.9)
     parser.add_argument("--decay_steps", type=int, default=300)
+    parser.add_argument("--dynamic_decay", action="store_true")
+    parser.add_argument("--no_decay", action="store_true", default=True,
+                        help="Keep text_ratio fixed (recommended for Stage 1)")
+    parser.add_argument("--use_decay", action="store_true",
+                        help="Enable decay (overrides --no_decay)")
+    parser.add_argument("--final_audio_portion", type=float, default=0.2)
+
     # Model
-    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
+    parser.add_argument("--model_path", type=str,
+                        default="canopylabs/3b-es_it-ft-research_release")
+
+    # Checkpointing
+    parser.add_argument("--save_steps", type=int, default=200)
     parser.add_argument("--resume", type=str, default=None)
-    # Memory limits
-    parser.add_argument("--vram_fraction", type=float, default=0.80, help="VRAM fraction to use (default 0.80)")
-    parser.add_argument("--ram_limit_gb", type=float, default=None, help="RAM limit in GB (auto if not specified)")
+
+    # Memory
+    parser.add_argument("--vram_fraction", type=float, default=0.80)
+    parser.add_argument("--ram_limit_gb", type=float, default=None)
+    parser.add_argument("--gradient_checkpointing", action="store_true")
+
     # Modes
     parser.add_argument("--demo", action="store_true")
     parser.add_argument("--test", action="store_true")
-    # Gradient checkpointing
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Enable gradient checkpointing (auto-enabled if VRAM < 20GB)")
+
     return parser.parse_args()
 
 
-# ============================================================
-# Main
-# ============================================================
 def main():
     args = parse_args()
-
-    # Determine mixed precision based on device
-    # MPS doesn't support bf16, CUDA/ROCm do
-    mixed_precision = None
-    if torch.cuda.is_available():
-        mixed_precision = "bf16"
-    # MPS and others use fp32 or fp16
-
-    # Initialize Accelerator (auto-detects CUDA, ROCm, MPS, XPU)
-    accelerator = Accelerator(
-        gradient_accumulation_steps=args.grad_accum if args.grad_accum else 1,
-        mixed_precision=mixed_precision,
-    )
-
-    device = accelerator.device
-    is_main = accelerator.is_main_process
-
-    # Set seed for reproducibility
-    set_seed(42)
-
-    # GPU config
-    gpu_config = auto_detect_gpu_config()
-    if args.batch_size is None:
-        args.batch_size = gpu_config["batch_size"]
-    if args.grad_accum is None:
-        args.grad_accum = gpu_config["grad_accum"]
-        # Update accelerator's gradient accumulation
-        accelerator.gradient_accumulation_steps = args.grad_accum
-
-    # Use bf16 for CUDA, fp32 for MPS/others
-    torch_dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32
-
-    # Get RAM info and set limits
-    ram_total, ram_available = get_ram_info()
-    if args.ram_limit_gb is None:
-        args.ram_limit_gb = ram_total * 0.80  # Default to 80% of RAM
-    limit_ram_usage(args.ram_limit_gb)
-
-    if is_main:
-        log("=" * 60)
-        log("STAGE 1: Adapter Only (LLM Frozen) + Interleaved Output")
-        log("=" * 60)
-        log(f"Device: {device} ({accelerator.device.type})")
-        log(f"GPU: {gpu_config['name']} ({gpu_config['vram_gb']}GB)")
-        log(f"Num processes: {accelerator.num_processes}")
-        log(f"RAM: {ram_total:.1f}GB total, {ram_available:.1f}GB available")
-        log(f"Batch: {args.batch_size}, Grad accum: {args.grad_accum}")
-        log(f"LR: {args.lr}, Epochs: {args.epochs}")
-        log(f"Initial text ratio: {args.initial_text_ratio}")
-        log(f"Decay steps: {args.decay_steps}")
-
-    # Apply memory limits for CUDA
-    if device.type == 'cuda':
-        torch.cuda.set_per_process_memory_fraction(args.vram_fraction)
-        torch.cuda.empty_cache()
-        torch.backends.cudnn.benchmark = True
-        torch.set_float32_matmul_precision('high')
-        if is_main:
-            log(f"[MEMORY] VRAM limited to {args.vram_fraction*100:.0f}%")
-            log(f"[MEMORY] RAM limited to {args.ram_limit_gb:.1f}GB")
-
-    # HuggingFace login
-    hf_token = os.environ.get("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    # Load datasets
-    data_paths = [p.strip() for p in args.data.split(",")]
-    all_datasets = []
-
-    if is_main:
-        log("\nLoading datasets...")
-
-    for path in data_paths:
-        if os.path.exists(path):
-            data = torch.load(path, weights_only=False, mmap=True)
-            dataset = InterleavedDataset(data, tokenizer)
-            all_datasets.append(dataset)
-            if is_main:
-                log(f"  {os.path.basename(path)}: {len(data):,} samples")
-
-    if len(all_datasets) == 0:
-        raise ValueError("No datasets loaded!")
-
-    combined_dataset = ConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0]
-
-    # Demo/Test mode
-    if args.test:
-        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(5, len(combined_dataset))))
-        args.batch_size = min(args.batch_size, len(combined_dataset))
-        args.grad_accum = 1
-    elif args.demo:
-        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(1000, len(combined_dataset))))
-        args.batch_size = min(4, args.batch_size)
-        args.grad_accum = max(8, args.grad_accum)
-
-    if is_main:
-        log(f"Total samples: {len(combined_dataset):,}")
-
-    # Load LLM (FROZEN)
-    if is_main:
-        log(f"\nLoading LLM (FROZEN): {args.model_path}")
-
-    llm = AutoModelForCausalLM.from_pretrained(
-        args.model_path,
-        torch_dtype=torch_dtype,
-        attn_implementation="sdpa",
-    )
-
-    # Freeze LLM completely
-    for p in llm.parameters():
-        p.requires_grad = False
-    llm.eval()
-
-    # Gradient checkpointing auto-detection
-    use_gradient_checkpointing = args.gradient_checkpointing
-    if not use_gradient_checkpointing and torch.cuda.is_available():
-        vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
-        if vram_gb < 20:
-            use_gradient_checkpointing = True
-            if is_main:
-                log(f"[AUTO] Enabling gradient checkpointing (VRAM={vram_gb:.1f}GB < 20GB)")
-        else:
-            if is_main:
-                log(f"[AUTO] Gradient checkpointing disabled for speed (VRAM={vram_gb:.1f}GB >= 20GB)")
-
-    if use_gradient_checkpointing:
-        llm.gradient_checkpointing_enable()
-        if is_main:
-            log("[MEMORY] Gradient checkpointing enabled")
-
-    # Create adapter (TRAINABLE)
-    adapter = SpeechAdapter(
-        whisper_dim=1280,
-        llm_dim=3072,
-        downsample=5,
-        intermediate_dim=2048
-    ).to(dtype=torch_dtype)
-
-    adapter_params = sum(p.numel() for p in adapter.parameters())
-    if is_main:
-        log(f"\nTrainable: Adapter only ({adapter_params:,} = {adapter_params/1e6:.1f}M params)")
-        log("LLM: FROZEN")
-
-    # Optimizer (only adapter)
-    optimizer = torch.optim.AdamW(adapter.parameters(), lr=args.lr, weight_decay=0.01)
-
-    # Training state
-    global_step = 0
-    start_epoch = 0
-    best_loss = float("inf")
-    current_text_ratio = args.initial_text_ratio
-
-    # Resume
-    if args.resume and os.path.exists(args.resume):
-        if is_main:
-            log(f"\nResuming from: {args.resume}")
-        ckpt = torch.load(args.resume, map_location="cpu", weights_only=False)
-        if "adapter" in ckpt:
-            adapter.load_state_dict(ckpt["adapter"])
-        if "optimizer" in ckpt:
-            optimizer.load_state_dict(ckpt["optimizer"])
-        if "step" in ckpt:
-            global_step = ckpt["step"]
-        if "epoch" in ckpt:
-            start_epoch = ckpt["epoch"]
-        if "text_ratio" in ckpt:
-            current_text_ratio = ckpt["text_ratio"]
-
-    # Training
-    os.makedirs(args.output_dir, exist_ok=True)
-
-    # Create dataloader with dynamic collate
-    def collate_with_ratio(batch):
-        current_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
-        return collate_fn(batch, current_ratio, tokenizer=tokenizer)
-
-    train_loader = DataLoader(
-        combined_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        collate_fn=collate_with_ratio,
-        num_workers=4,
-        pin_memory=True
-    )
-
-    # Prepare with accelerator (handles DDP, device placement, etc.)
-    adapter, llm, optimizer, train_loader = accelerator.prepare(
-        adapter, llm, optimizer, train_loader
-    )
-
-    steps_per_epoch = max(1, len(train_loader) // args.grad_accum)
-    total_steps = max(1, steps_per_epoch * args.epochs)
-    warmup_steps = int(total_steps * args.warmup_ratio)
-
-    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, total_steps - warmup_steps), eta_min=1e-6)
-
-    if is_main:
-        log(f"Steps per epoch: {steps_per_epoch}, Total: {total_steps}, Warmup: {warmup_steps}")
-
-    if is_main:
-        log("\n" + "=" * 60)
-        log("STARTING STAGE 1 TRAINING")
-        log("=" * 60)
-
-    for epoch in range(start_epoch, args.epochs):
-        # Update text ratio based on global step
-        current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
-
-        # Get unwrapped adapter for forward pass
-        unwrapped_adapter = accelerator.unwrap_model(adapter)
-        unwrapped_llm = accelerator.unwrap_model(llm)
-
-        adapter.train()
-        epoch_loss = 0
-        accum_loss = 0
-
-        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}", disable=not is_main)
-
-        for batch_idx, batch in enumerate(pbar):
-            # Update text ratio dynamically
-            current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
-
-            whisper = batch["whisper"]
-            interleaved = batch["interleaved"]
-
-            # Use gradient accumulation context
-            with accelerator.accumulate(adapter):
-                # Forward through adapter
-                audio_embeds = adapter(whisper)
-
-                # Get token embeddings for interleaved sequence (teacher forcing)
-                input_tokens = interleaved[:, :-1].clamp(min=0)
-                with torch.no_grad():
-                    token_embeds = unwrapped_llm.model.embed_tokens(input_tokens)
-
-                # Combine: [audio_embeds] + [token_embeds]
-                combined = torch.cat([audio_embeds, token_embeds], dim=1)
-
-                # Forward through frozen LLM
-                outputs = llm(inputs_embeds=combined, use_cache=False)
-                logits = outputs.logits
-
-                # Loss: predict interleaved tokens after audio prefix
-                audio_len = audio_embeds.shape[1]
-                seq_len = interleaved.shape[1]
-
-                seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]
-                targets = interleaved
-
-                loss = F.cross_entropy(
-                    seq_logits.reshape(-1, logits.size(-1)),
-                    targets.reshape(-1),
-                    ignore_index=-100,
-                    label_smoothing=args.label_smoothing
-                )
-
-                accelerator.backward(loss)
-                accum_loss += loss.item()
-
-                # Clip gradients and step
-                if accelerator.sync_gradients:
-                    accelerator.clip_grad_norm_(adapter.parameters(), args.max_grad_norm)
-
-                optimizer.step()
-                optimizer.zero_grad()
-
-            # Update after accumulation
-            if accelerator.sync_gradients:
-                if global_step < warmup_steps:
-                    lr_scale = (global_step + 1) / warmup_steps
-                    for pg in optimizer.param_groups:
-                        pg["lr"] = args.lr * lr_scale
-                else:
-                    scheduler.step()
-
-                global_step += 1
-                epoch_loss += accum_loss
-
-                pbar.set_postfix(
-                    loss=f"{accum_loss:.4f}",
-                    text_ratio=f"{current_text_ratio:.1f}",
-                    lr=f"{optimizer.param_groups[0]['lr']:.2e}"
-                )
-
-                # Save checkpoint
-                if global_step % args.save_steps == 0 and is_main:
-                    accelerator.wait_for_everyone()
-                    unwrapped_adapter = accelerator.unwrap_model(adapter)
-                    ckpt_path = os.path.join(args.output_dir, f"stage1_step{global_step}.pt")
-                    save_checkpoint_async({
-                        "adapter": unwrapped_adapter.state_dict(),
-                        "optimizer": optimizer.state_dict(),
-                        "step": global_step,
-                        "epoch": epoch,
-                        "loss": accum_loss,
-                        "text_ratio": current_text_ratio
-                    }, ckpt_path, is_main)
-
-                accum_loss = 0
-
-        # Epoch end - divide by grad_accum to get per-batch average
-        avg_loss = epoch_loss / max(1, steps_per_epoch) / args.grad_accum
-
-        if is_main:
-            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {current_text_ratio:.1f}")
-
-            accelerator.wait_for_everyone()
-            unwrapped_adapter = accelerator.unwrap_model(adapter)
-            ckpt_path = os.path.join(args.output_dir, f"stage1_epoch{epoch+1}.pt")
-            save_checkpoint_async({
-                "adapter": unwrapped_adapter.state_dict(),
-                "optimizer": optimizer.state_dict(),
-                "step": global_step,
-                "epoch": epoch + 1,
-                "loss": avg_loss,
-                "text_ratio": current_text_ratio
-            }, ckpt_path, is_main)
-
-            if avg_loss < best_loss:
-                best_loss = avg_loss
-                best_path = os.path.join(args.output_dir, "stage1_best.pt")
-                save_checkpoint_async({
-                    "adapter": unwrapped_adapter.state_dict(),
-                    "step": global_step,
-                    "epoch": epoch + 1,
-                    "loss": best_loss,
-                    "text_ratio": current_text_ratio
-                }, best_path, is_main)
-                log(f"Best model saved! Loss: {best_loss:.4f}")
-
-    # Finish
-    accelerator.wait_for_everyone()
-    if is_main:
-        wait_for_checkpoints()
-        log("\n" + "=" * 60)
-        log("STAGE 1 COMPLETE!")
-        log(f"Best loss: {best_loss:.4f}")
-        log(f"Final text_ratio: {current_text_ratio:.1f}")
-        log("Next: Stage 2 (Adapter + LoRA)")
-        log("=" * 60)
+    # --use_decay overrides --no_decay
+    if args.use_decay:
+        args.no_decay = False
+    config = TrainingConfig.from_args(args)
+
+    trainer = Stage1Trainer(config)
+    trainer.setup()
+    trainer.train()
 
 
 if __name__ == "__main__":

passo3_finetune_stage2.py

@@ -9,916 +9,289 @@ Continues from Stage 1 checkpoint:
 - Continues scheduled interleaving (90% text -> 0% text)
 
 Usage:
-    python finetune_stage2.py --data data.pt --stage1_ckpt checkpoints/stage1_best.pt --epochs 3
+    python passo3_finetune_stage2.py --data data.pt --stage1_ckpt checkpoints/stage1_best.pt --epochs 3
 
 Based on IST-LM paper + LLaMA-Omni 2 staging approach.
 """
 
-import os
-import sys
 import argparse
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.data import Dataset, DataLoader, ConcatDataset
-from torch.optim.lr_scheduler import CosineAnnealingLR
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import LoraConfig, get_peft_model, TaskType
-from accelerate import Accelerator
-from accelerate.utils import set_seed
-from huggingface_hub import login
-from tqdm import tqdm
-import threading
-
-# ============================================================
-# Config
-# ============================================================
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
-torch.set_float32_matmul_precision('high')
-
-# SNAC token offsets for Orpheus
-SNAC_BASE_OFFSET = 128266
-EOS_TOKEN = 128009
-
-
-def log(msg):
-    print(msg)
-    sys.stdout.flush()
-
-
-def apply_snac_offset(token_idx, position):
-    """Apply position-based offset to SNAC token.
-    If token is already offset (>= SNAC_BASE_OFFSET), return as-is.
-    """
-    if int(token_idx) >= SNAC_BASE_OFFSET:
-        # Already has offset applied
-        return int(token_idx)
-    offset = SNAC_BASE_OFFSET + (position % 7) * 4096
-    return int(token_idx) + offset
+from typing import Dict, Any
+from itertools import chain
 
-
-def get_text_ratio(global_step, decay_steps=300, initial_ratio=0.9, min_ratio=0.0):
-    """
-    IST-LM: Start with 90% text, decrease by 0.1 every 300 steps.
-    """
-    num_decays = global_step // decay_steps
-    text_ratio = initial_ratio - (num_decays * 0.1)
-    return max(min_ratio, text_ratio)
-
-
-# ============================================================
-# Async Checkpoint Saving
-# ============================================================
-_save_threads = []
-
-def save_checkpoint_async(state_dict, path, is_main=True):
-    global _save_threads
-    _save_threads = [t for t in _save_threads if t.is_alive()]
-
-    def copy_to_cpu(obj):
-        if isinstance(obj, torch.Tensor):
-            return obj.detach().cpu().clone()
-        elif isinstance(obj, dict):
-            return {k: copy_to_cpu(v) for k, v in obj.items()}
-        return obj
-
-    state_copy = copy_to_cpu(state_dict)
-
-    def _save():
-        try:
-            torch.save(state_copy, path)
-            if is_main:
-                log(f"[ASYNC] Saved: {path}")
-        except Exception as e:
-            if is_main:
-                log(f"[ASYNC] Error: {e}")
-
-    thread = threading.Thread(target=_save, daemon=True)
-    thread.start()
-    _save_threads.append(thread)
-
-
-def wait_for_checkpoints():
-    global _save_threads
-    for t in _save_threads:
-        t.join()
-    _save_threads = []
-
-
-# ============================================================
-# GPU Auto-Detection (Multi-backend: CUDA, ROCm, MPS, XPU)
-# ============================================================
-def auto_detect_gpu_config():
-    """Detect GPU and return optimal batch size config."""
-    vram_gb = 0
-    gpu_name = "Unknown"
-
-    # Try CUDA (NVIDIA)
-    if torch.cuda.is_available():
-        try:
-            props = torch.cuda.get_device_properties(0)
-            vram_gb = props.total_memory // (1024**3)
-            gpu_name = props.name
-        except:
-            pass
-
-    # Try MPS (Apple Silicon)
-    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-        gpu_name = "Apple Silicon (MPS)"
-        try:
-            import subprocess
-            result = subprocess.run(['sysctl', '-n', 'hw.memsize'], capture_output=True, text=True)
-            total_mem = int(result.stdout.strip()) // (1024**3)
-            vram_gb = total_mem // 2
-        except:
-            vram_gb = 8
-
-    # Try ROCm (AMD)
-    elif hasattr(torch, 'hip') or os.environ.get('ROCM_HOME'):
-        try:
-            if torch.cuda.is_available():
-                props = torch.cuda.get_device_properties(0)
-                vram_gb = props.total_memory // (1024**3)
-                gpu_name = f"AMD {props.name}"
-        except:
-            gpu_name = "AMD ROCm"
-            vram_gb = 16
-
-    # Fallback: try nvidia-smi
-    if vram_gb == 0:
-        try:
-            import subprocess
-            result = subprocess.run(
-                ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader,nounits'],
-                capture_output=True, text=True
-            )
-            lines = result.stdout.strip().split('\n')
-            gpu_name, vram_mb = lines[0].split(', ')
-            vram_gb = int(vram_mb) // 1024
-        except:
-            pass
-
-    # Determine batch size based on VRAM
-    if vram_gb >= 150:
-        return {"name": gpu_name, "batch_size": 8, "grad_accum": 4, "vram_gb": vram_gb}
-    elif vram_gb >= 80:
-        return {"name": gpu_name, "batch_size": 6, "grad_accum": 5, "vram_gb": vram_gb}
-    elif vram_gb >= 35:
-        return {"name": gpu_name, "batch_size": 4, "grad_accum": 8, "vram_gb": vram_gb}
-    elif vram_gb >= 16:
-        return {"name": gpu_name, "batch_size": 2, "grad_accum": 16, "vram_gb": vram_gb}
-    else:
-        return {"name": gpu_name, "batch_size": 1, "grad_accum": 32, "vram_gb": max(vram_gb, 8)}
-
-def get_ram_info():
-    """Get RAM info in GB."""
-    try:
-        import psutil
-        total = psutil.virtual_memory().total / 1024**3
-        available = psutil.virtual_memory().available / 1024**3
-        return total, available
-    except ImportError:
-        try:
-            import subprocess
-            result = subprocess.run(
-                ['free', '-g', '--output=SIZE,AVAILABLE'],
-                capture_output=True, text=True
-            )
-            lines = result.stdout.strip().split('\n')
-            if len(lines) >= 2:
-                total, available = map(float, lines[1].split())
-                return total, available
-        except:
-            pass
-    except:
-        pass
-    return 0, 0
-
-def limit_ram_usage(max_ram_gb):
-    """Limit RAM usage by setting resource limits."""
-    try:
-        import resource
-        max_bytes = int(max_ram_gb * 1024**3)
-        resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
-    except:
-        pass
-
-def log_memory_usage():
-    """Log current memory usage."""
-    msg = []
-    if torch.cuda.is_available():
-        used = torch.cuda.memory_allocated() / 1024**3
-        reserved = torch.cuda.memory_reserved() / 1024**3
-        msg.append(f"GPU: {used:.2f}GB / {reserved:.2f}GB")
-    try:
-        import psutil
-        ram_used = psutil.virtual_memory().used / 1024**3
-        ram_total = psutil.virtual_memory().total / 1024**3
-        msg.append(f"RAM: {ram_used:.1f}GB / {ram_total:.1f}GB")
-    except:
-        pass
-    return " | ".join(msg)
-
-
-# ============================================================
-# Speech Adapter (Same as Stage 1 - LLaMA-Omni 2 Style)
-# ============================================================
-class SpeechAdapter(nn.Module):
-    """
-    5x downsampling + FFN with intermediate dim 2048
-    MUST match Stage 1 architecture exactly for checkpoint loading.
-    """
-    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
-        super().__init__()
-        self.downsample = downsample
-        concat_dim = whisper_dim * downsample
-
-        self.ffn = nn.Sequential(
-            nn.Linear(concat_dim, intermediate_dim),
-            nn.GELU(),
-            nn.Linear(intermediate_dim, llm_dim),
-            nn.LayerNorm(llm_dim)
-        )
-
-    def forward(self, x):
-        B, T, D = x.shape
-        T_new = (T // self.downsample) * self.downsample
-        x = x[:, :T_new]
-        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
-        return self.ffn(x)
+from training.config import TrainingConfig, GPUConfig, LoRAConfig
+from training.models import SpeechAdapter, ModelFactory
+from training.checkpoint import Stage2CheckpointManager, TrainingState
+from training.trainer import BaseTrainer
+from training.utils import log, should_enable_gradient_checkpointing
 
 
-# ============================================================
-# Scheduled Interleaved Sequence Creation with Word Alignment
-# ============================================================
-def create_interleaved_sequence(text_tokens, snac_tokens, text_ratio=0.9, word_alignments=None, tokenizer=None, answer_text=None):
+class Stage2Trainer(BaseTrainer):
     """
-    Create interleaved sequence based on text_ratio with word-level alignment.
-    - text_ratio=0.9 means 90% of words are replaced by text tokens
-    - text_ratio=0.0 means 100% audio (no text replacement)
+    Stage 2 Trainer: Adapter + LoRA together.
 
-    With word_alignments: replaces aligned audio spans with corresponding text tokens
-    Without word_alignments: falls back to positional interleaving
+    Continues from Stage 1, adding LoRA to the LLM and training
+    both adapter and LoRA jointly.
     """
-    interleaved = []
-    is_audio_mask = []
-
-    if len(snac_tokens) == 0:
-        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
-
-    # Group SNAC into frames of 7
-    frames = []
-    for i in range(0, len(snac_tokens), 7):
-        frame = snac_tokens[i:i+7]
-        if len(frame) == 7:
-            frames.append(frame)
-
-    if len(frames) == 0:
-        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
-
-    total_frames = len(frames)
-
-    # If we have word alignments, use semantic interleaving
-    # Check if alignments have pre-computed tokens (preferred) or need tokenizer
-    has_precomputed = word_alignments and len(word_alignments) > 0 and 'tokens' in word_alignments[0] and word_alignments[0]['tokens']
-    can_interleave = word_alignments and text_ratio > 0 and (has_precomputed or (tokenizer and answer_text))
-
-    if can_interleave:
-        import random
-
-        # Decide which words to replace with text based on text_ratio
-        num_words = len(word_alignments)
-        num_text_words = int(num_words * text_ratio)
-
-        # Randomly select which word indices to replace with text
-        word_indices = list(range(num_words))
-        random.shuffle(word_indices)
-        text_word_indices = set(word_indices[:num_text_words])
-
-        # Build interleaved sequence frame by frame
-        frame_idx = 0
-        snac_position = 0
-
-        for word_idx, alignment in enumerate(word_alignments):
-            word = alignment['word']
-            start_frame = alignment['start_frame']
-            end_frame = min(alignment['end_frame'], total_frames)
-
-            if word_idx in text_word_indices:
-                # Replace this word's audio with text tokens
-                # Use pre-computed tokens if available, else tokenize on-the-fly
-                word_tokens = alignment.get('tokens', [])
-                if not word_tokens and tokenizer:
-                    word_tokens = tokenizer.encode(word, add_special_tokens=False)
-
-                for tok in word_tokens:
-                    interleaved.append(tok)
-                    is_audio_mask.append(False)
-                # Skip the audio frames for this word
-                snac_position = end_frame * 7
-            else:
-                # Keep audio for this word
-                for f_idx in range(start_frame, end_frame):
-                    if f_idx < total_frames:
-                        frame = frames[f_idx]
-                        for tok in frame:
-                            interleaved.append(apply_snac_offset(tok, snac_position))
-                            is_audio_mask.append(True)
-                            snac_position += 1
-
-            frame_idx = end_frame
-
-        # Add any remaining frames after the last word
-        while frame_idx < total_frames:
-            frame = frames[frame_idx]
-            for tok in frame:
-                interleaved.append(apply_snac_offset(tok, snac_position))
-                is_audio_mask.append(True)
-                snac_position += 1
-            frame_idx += 1
-
-    else:
-        # Fallback: positional interleaving (original behavior)
-        total_text = len(text_tokens)
-
-        # Determine interleaving pattern based on text_ratio
-        if text_ratio >= 0.9:
-            text_per_chunk, frames_per_chunk = 1, 3
-        elif text_ratio >= 0.7:
-            text_per_chunk, frames_per_chunk = 1, 5
-        elif text_ratio >= 0.5:
-            text_per_chunk, frames_per_chunk = 1, 7
-        elif text_ratio >= 0.3:
-            text_per_chunk, frames_per_chunk = 1, 10
-        else:
-            text_per_chunk, frames_per_chunk = 0, 1
-
-        text_idx = 0
-        frame_idx = 0
-        snac_position = 0
-
-        while frame_idx < total_frames:
-            if text_per_chunk > 0 and text_idx < total_text:
-                for _ in range(text_per_chunk):
-                    if text_idx < total_text:
-                        interleaved.append(text_tokens[text_idx])
-                        is_audio_mask.append(False)
-                        text_idx += 1
-
-            for _ in range(frames_per_chunk):
-                if frame_idx < total_frames:
-                    frame = frames[frame_idx]
-                    for tok in frame:
-                        interleaved.append(apply_snac_offset(tok, snac_position))
-                        is_audio_mask.append(True)
-                        snac_position += 1
-                    frame_idx += 1
-
-        while text_idx < total_text:
-            interleaved.append(text_tokens[text_idx])
-            is_audio_mask.append(False)
-            text_idx += 1
-
-    # Add EOS
-    interleaved.append(EOS_TOKEN)
-    is_audio_mask.append(False)
-
-    return interleaved, is_audio_mask
-
-
-# ============================================================
-# Dataset
-# ============================================================
-class InterleavedDataset(Dataset):
-    def __init__(self, data, tokenizer, max_audio_len=500, max_seq_len=2048):
-        self.data = data
-        self.tokenizer = tokenizer
-        self.max_audio = max_audio_len * 5
-        self.max_seq_len = max_seq_len
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        item = self.data[idx]
-
-        # Whisper features
-        whisper = item["whisper_features"][:self.max_audio]
-
-        # Text tokens - use pre-computed if available, otherwise tokenize
-        if "text_tokens" in item and len(item["text_tokens"]) > 0:
-            tt = item["text_tokens"]
-            text_tokens = tt.tolist() if hasattr(tt, 'tolist') else list(tt)
-        else:
-            text = item.get("answer", item.get("text", ""))
-            if isinstance(text, str) and len(text) > 0:
-                text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
-            else:
-                text_tokens = []
-
-        # SNAC tokens
-        snac = item["snac_tokens"]
-        snac_len = (len(snac) // 7) * 7
-        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
-        snac_list = snac.tolist() if hasattr(snac, 'tolist') else list(snac)
-
-        # Word alignments (if available)
-        word_alignments = item.get("word_alignments", None)
-        answer_text = item.get("answer", "")
 
-        return {
-            "whisper": whisper,
-            "text_tokens": text_tokens,
-            "snac_tokens": snac_list,
-            "word_alignments": word_alignments,
-            "answer_text": answer_text
-        }
+    def __init__(self, config: TrainingConfig, stage1_ckpt: str, lora_config: LoRAConfig = None):
+        super().__init__(config)
+        self.stage1_ckpt = stage1_ckpt
+        self.lora_config = lora_config or LoRAConfig()
 
+    def _get_stage_name(self) -> str:
+        return "STAGE 2: Adapter + LoRA (Joint Training)"
 
-def collate_fn(batch, text_ratio=0.9, tokenizer=None):
-    """Collate with dynamic interleaving based on text_ratio and word alignments."""
-    max_w = max(b["whisper"].shape[0] for b in batch)
-    max_w = ((max_w + 4) // 5) * 5
+    def _setup_models(self):
+        """Setup adapter and LLM with LoRA."""
+        gpu_config = GPUConfig.auto_detect()
+        dtype = gpu_config.dtype
 
-    whisper_batch = []
-    interleaved_batch = []
-    is_audio_batch = []
+        # Check gradient checkpointing
+        use_checkpointing = self.config.gradient_checkpointing or \
+            should_enable_gradient_checkpointing(
+                gpu_config.vram_gb,
+                self.config.dynamic_decay
+            )
 
-    max_seq = 0
-    sequences = []
+        if self.is_main:
+            log(f"\nLoading LLM: {self.config.model_path}")
 
-    for b in batch:
-        interleaved, is_audio = create_interleaved_sequence(
-            b["text_tokens"],
-            b["snac_tokens"],
-            text_ratio,
-            word_alignments=b.get("word_alignments"),
-            tokenizer=tokenizer,
-            answer_text=b.get("answer_text")
+        # Load LLM (not frozen, will apply LoRA)
+        self.llm = ModelFactory.create_llm(
+            self.config.model_path,
+            dtype=dtype,
+            freeze=False,
+            gradient_checkpointing=use_checkpointing,
         )
-        sequences.append((interleaved, is_audio))
-        max_seq = max(max_seq, len(interleaved))
-
-    for i, b in enumerate(batch):
-        w = b["whisper"]
-        w_pad = F.pad(w, (0, 0, 0, max_w - w.shape[0]))
-        whisper_batch.append(w_pad)
-
-        interleaved, is_audio = sequences[i]
-        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
-        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)
-
-        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
-        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)
-
-        interleaved_batch.append(seq_pad)
-        is_audio_batch.append(mask_pad)
 
-    return {
-        "whisper": torch.stack(whisper_batch),
-        "interleaved": torch.stack(interleaved_batch),
-        "is_audio_mask": torch.stack(is_audio_batch)
-    }
+        # Apply LoRA
+        if self.is_main:
+            log(f"Applying LoRA (r={self.lora_config.r}, alpha={self.lora_config.alpha})")
 
+        self.llm = ModelFactory.apply_lora(self.llm, self.lora_config)
 
-# ============================================================
-# Arguments
-# ============================================================
-def parse_args():
-    parser = argparse.ArgumentParser(description="Stage 2: Adapter + LoRA with Interleaved Output")
-    parser.add_argument("--data", type=str, required=True)
-    parser.add_argument("--output_dir", type=str, default="./checkpoints")
-    parser.add_argument("--stage1_ckpt", type=str, default=None,
-                        help="Path to Stage 1 adapter checkpoint (required for proper init)")
-    parser.add_argument("--lr", type=float, default=5e-5)
-    parser.add_argument("--epochs", type=int, default=3, help="3+ epochs recommended")
-    parser.add_argument("--batch_size", type=int, default=None)
-    parser.add_argument("--grad_accum", type=int, default=None)
-    parser.add_argument("--warmup_ratio", type=float, default=0.03)
-    parser.add_argument("--max_grad_norm", type=float, default=1.0)
-    parser.add_argument("--save_steps", type=int, default=200)
-    parser.add_argument("--label_smoothing", type=float, default=0.1)
-    # Scheduled interleaving
-    parser.add_argument("--initial_text_ratio", type=float, default=0.9)
-    parser.add_argument("--decay_steps", type=int, default=300)
-    # LoRA config
-    parser.add_argument("--lora_r", type=int, default=16)
-    parser.add_argument("--lora_alpha", type=int, default=32)
-    parser.add_argument("--lora_dropout", type=float, default=0.05)
-    # Model
-    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
-    parser.add_argument("--resume", type=str, default=None)
-    # Memory limits
-    parser.add_argument("--vram_fraction", type=float, default=0.80, help="VRAM fraction to use (default 0.80)")
-    parser.add_argument("--ram_limit_gb", type=float, default=None, help="RAM limit in GB (auto if not specified)")
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Enable gradient checkpointing to reduce VRAM")
-    # Modes
-    parser.add_argument("--demo", action="store_true")
-    parser.add_argument("--test", action="store_true")
-    return parser.parse_args()
+        if self.is_main and use_checkpointing:
+            log("[Memory] Gradient checkpointing enabled")
 
+        # Create and load adapter from Stage 1
+        self.adapter = SpeechAdapter().to(dtype=dtype)
 
-# ============================================================
-# Main
-# ============================================================
-def main():
-    args = parse_args()
+        if self.is_main:
+            log(f"\nLoading adapter from Stage 1: {self.stage1_ckpt}")
 
-    # Determine mixed precision based on device
-    # MPS doesn't support bf16, CUDA/ROCm do
-    mixed_precision = None
-    if torch.cuda.is_available():
-        mixed_precision = "bf16"
+        ckpt = torch.load(self.stage1_ckpt, map_location="cpu", weights_only=False)
+        if "adapter" in ckpt:
+            self.adapter.load_state_dict(ckpt["adapter"])
+        else:
+            self.adapter.load_state_dict(ckpt)
+
+        # Log parameters
+        if self.is_main:
+            adapter_params = self.adapter.get_num_params()
+            lora_params = sum(p.numel() for p in self.llm.parameters() if p.requires_grad)
+            total_params = adapter_params + lora_params
+
+            log(f"\nTrainable parameters:")
+            log(f"  Adapter: {adapter_params:,} ({adapter_params/1e6:.1f}M)")
+            log(f"  LoRA: {lora_params:,} ({lora_params/1e6:.1f}M)")
+            log(f"  Total: {total_params:,} ({total_params/1e6:.1f}M)")
+
+    def _setup_optimizer(self):
+        """Setup optimizer for adapter + LoRA."""
+        trainable_params = chain(
+            self.adapter.parameters(),
+            (p for p in self.llm.parameters() if p.requires_grad)
+        )
 
-    # Initialize Accelerator (auto-detects CUDA, ROCm, MPS, XPU)
-    accelerator = Accelerator(
-        gradient_accumulation_steps=args.grad_accum if args.grad_accum else 1,
-        mixed_precision=mixed_precision,
-    )
+        self.optimizer = torch.optim.AdamW(
+            trainable_params,
+            lr=self.config.learning_rate,
+            weight_decay=self.config.weight_decay,
+        )
 
-    device = accelerator.device
-    is_main = accelerator.is_main_process
-
-    # Set seed for reproducibility
-    set_seed(42)
-
-    # GPU config
-    gpu_config = auto_detect_gpu_config()
-    if args.batch_size is None:
-        args.batch_size = gpu_config["batch_size"]
-    if args.grad_accum is None:
-        args.grad_accum = gpu_config["grad_accum"]
-        accelerator.gradient_accumulation_steps = args.grad_accum
-
-    # Use bf16 for CUDA, fp32 for MPS/others
-    torch_dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32
-
-    # Get RAM info and set limits
-    ram_total, ram_available = get_ram_info()
-    if args.ram_limit_gb is None:
-        args.ram_limit_gb = ram_total * 0.80
-    limit_ram_usage(args.ram_limit_gb)
-
-    if is_main:
-        log("=" * 60)
-        log("STAGE 2: Adapter + LoRA (Both Trainable) + Interleaved Output")
-        log("=" * 60)
-        log(f"Device: {device} ({accelerator.device.type})")
-        log(f"GPU: {gpu_config['name']} ({gpu_config['vram_gb']}GB)")
-        log(f"Num processes: {accelerator.num_processes}")
-        log(f"RAM: {ram_total:.1f}GB total, {ram_available:.1f}GB available")
-        log(f"Batch: {args.batch_size}, Grad accum: {args.grad_accum}")
-        log(f"LR: {args.lr}, Epochs: {args.epochs}")
-        log(f"LoRA: r={args.lora_r}, alpha={args.lora_alpha}")
-        log(f"Initial text ratio: {args.initial_text_ratio}")
-        log(f"Decay steps: {args.decay_steps}")
-
-    # Apply memory limits for CUDA
-    if device.type == 'cuda':
-        torch.cuda.set_per_process_memory_fraction(args.vram_fraction)
-        torch.cuda.empty_cache()
-        torch.backends.cudnn.benchmark = True
-        torch.set_float32_matmul_precision('high')
-        if is_main:
-            log(f"[MEMORY] VRAM limited to {args.vram_fraction*100:.0f}%")
-            log(f"[MEMORY] RAM limited to {args.ram_limit_gb:.1f}GB")
-
-    # HuggingFace login
-    hf_token = os.environ.get("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    # Load datasets
-    data_paths = [p.strip() for p in args.data.split(",")]
-    all_datasets = []
-
-    if is_main:
-        log("\nLoading datasets...")
-
-    for path in data_paths:
-        if os.path.exists(path):
-            data = torch.load(path, weights_only=False, mmap=True)
-            dataset = InterleavedDataset(data, tokenizer)
-            all_datasets.append(dataset)
-            if is_main:
-                log(f"  {os.path.basename(path)}: {len(data):,} samples")
-
-    if len(all_datasets) == 0:
-        raise ValueError("No datasets loaded!")
-
-    combined_dataset = ConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0]
-
-    # Demo/Test mode
-    if args.test:
-        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(5, len(combined_dataset))))
-        args.batch_size = min(args.batch_size, len(combined_dataset))
-        args.grad_accum = 1
-    elif args.demo:
-        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(1000, len(combined_dataset))))
-        args.batch_size = min(4, args.batch_size)
-        args.grad_accum = max(8, args.grad_accum)
-
-    if is_main:
-        log(f"Total samples: {len(combined_dataset):,}")
-
-    # Load LLM with LoRA
-    if is_main:
-        log(f"\nLoading LLM with LoRA: {args.model_path}")
-
-    llm = AutoModelForCausalLM.from_pretrained(
-        args.model_path,
-        torch_dtype=torch_dtype,
-        attn_implementation="sdpa",
-    )
+    def _setup_checkpoint_manager(self):
+        """Setup Stage 2 checkpoint manager."""
+        self.checkpoint_manager = Stage2CheckpointManager(
+            self.config.output_dir,
+            verbose=self.is_main,
+        )
 
-    # Apply LoRA
-    lora_config = LoraConfig(
-        r=args.lora_r,
-        lora_alpha=args.lora_alpha,
-        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
-        lora_dropout=args.lora_dropout,
-        bias="none",
-        task_type=TaskType.CAUSAL_LM
-    )
-    llm = get_peft_model(llm, lora_config)
-
-    # Auto-detect gradient checkpointing based on VRAM
-    # Enable if VRAM < 20GB (needed for Stage 2 with LoRA)
-    use_gradient_checkpointing = args.gradient_checkpointing
-    if not use_gradient_checkpointing and torch.cuda.is_available():
-        vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
-        if vram_gb < 20:
-            use_gradient_checkpointing = True
-            if is_main:
-                log(f"[AUTO] Enabling gradient checkpointing (VRAM={vram_gb:.1f}GB < 20GB)")
-        else:
-            if is_main:
-                log(f"[AUTO] Gradient checkpointing disabled for speed (VRAM={vram_gb:.1f}GB >= 20GB)")
-
-    if use_gradient_checkpointing:
-        llm.gradient_checkpointing_enable()
-        if is_main:
-            log("[MEMORY] Gradient checkpointing enabled")
-
-    if is_main:
-        llm.print_trainable_parameters()
-
-    # Create adapter (TRAINABLE - same architecture as Stage 1)
-    adapter = SpeechAdapter(
-        whisper_dim=1280,
-        llm_dim=3072,
-        downsample=5,
-        intermediate_dim=2048
-    ).to(dtype=torch_dtype)
-
-    # Load Stage 1 adapter checkpoint
-    if args.stage1_ckpt and os.path.exists(args.stage1_ckpt):
-        if is_main:
-            log(f"\nLoading Stage 1 adapter from: {args.stage1_ckpt}")
-        ckpt = torch.load(args.stage1_ckpt, map_location="cpu", weights_only=False)
-        if "adapter" in ckpt:
-            adapter.load_state_dict(ckpt["adapter"])
-            if is_main:
-                log("  Adapter weights loaded successfully!")
-                if "loss" in ckpt:
-                    log(f"  Stage 1 final loss: {ckpt['loss']:.4f}")
-                if "text_ratio" in ckpt:
-                    log(f"  Stage 1 final text_ratio: {ckpt['text_ratio']:.1f}")
-        else:
-            if is_main:
-                log("  WARNING: No 'adapter' key in checkpoint, using random init")
-    else:
-        if is_main:
-            log("\nWARNING: No Stage 1 checkpoint provided, adapter starting from random init")
-            log("  Recommended: --stage1_ckpt checkpoints/stage1_best.pt")
-
-    adapter_params = sum(p.numel() for p in adapter.parameters())
-    lora_params = sum(p.numel() for p in llm.parameters() if p.requires_grad)
-
-    if is_main:
-        log(f"\nTrainable parameters:")
-        log(f"  Adapter: {adapter_params:,} ({adapter_params/1e6:.1f}M)")
-        log(f"  LoRA: {lora_params:,} ({lora_params/1e6:.1f}M)")
-        log(f"  Total: {adapter_params + lora_params:,} ({(adapter_params + lora_params)/1e6:.1f}M)")
-
-    # Optimizer (Adapter + LoRA together)
-    all_params = list(adapter.parameters()) + [p for p in llm.parameters() if p.requires_grad]
-    optimizer = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=0.01)
-
-    # Training state
-    global_step = 0
-    start_epoch = 0
-    best_loss = float("inf")
-    current_text_ratio = args.initial_text_ratio
-
-    # Resume from Stage 2 checkpoint
-    if args.resume and os.path.exists(args.resume):
-        if is_main:
-            log(f"\nResuming from: {args.resume}")
-        ckpt = torch.load(args.resume, map_location="cpu", weights_only=False)
+    def _load_checkpoint(self, ckpt: Dict[str, Any]):
+        """Load checkpoint state."""
         if "adapter" in ckpt:
-            adapter.load_state_dict(ckpt["adapter"])
+            self.adapter.load_state_dict(ckpt["adapter"])
         if "lora" in ckpt:
-            llm.load_state_dict(ckpt["lora"], strict=False)
+            # Load LoRA weights
+            self.llm.load_state_dict(ckpt["lora"], strict=False)
         if "optimizer" in ckpt:
-            optimizer.load_state_dict(ckpt["optimizer"])
+            self.optimizer.load_state_dict(ckpt["optimizer"])
         if "step" in ckpt:
-            global_step = ckpt["step"]
+            self.global_step = ckpt["step"]
         if "epoch" in ckpt:
-            start_epoch = ckpt["epoch"]
+            self.start_epoch = ckpt["epoch"]
         if "text_ratio" in ckpt:
-            current_text_ratio = ckpt["text_ratio"]
-        if "loss" in ckpt:
-            best_loss = ckpt["loss"]
+            self.current_text_ratio = ckpt["text_ratio"]
 
-    # Training
-    os.makedirs(args.output_dir, exist_ok=True)
-
-    # Create dataloader with dynamic collate
-    def collate_with_ratio(batch):
-        current_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
-        return collate_fn(batch, current_ratio, tokenizer=tokenizer)
-
-    train_loader = DataLoader(
-        combined_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        collate_fn=collate_with_ratio,
-        num_workers=4,
-        pin_memory=True
-    )
+    def _get_trainable_params(self):
+        """Get trainable parameters (adapter + LoRA)."""
+        return chain(
+            self.adapter.parameters(),
+            (p for p in self.llm.parameters() if p.requires_grad)
+        )
 
-    # Prepare with accelerator
-    adapter, llm, optimizer, train_loader = accelerator.prepare(
-        adapter, llm, optimizer, train_loader
-    )
+    def _get_lora_state_dict(self) -> Dict[str, Any]:
+        """Get LoRA state dict."""
+        unwrapped_llm = self.accelerator.unwrap_model(self.llm)
+        return {
+            k: v for k, v in unwrapped_llm.state_dict().items()
+            if "lora" in k.lower()
+        }
 
-    steps_per_epoch = max(1, len(train_loader) // args.grad_accum)
-    total_steps = max(1, steps_per_epoch * args.epochs)
-    warmup_steps = int(total_steps * args.warmup_ratio)
+    def _save_step_checkpoint(self, loss: float):
+        """Save step checkpoint."""
+        self.accelerator.wait_for_everyone()
+        unwrapped_adapter = self.accelerator.unwrap_model(self.adapter)
+
+        state = TrainingState(
+            step=self.global_step,
+            epoch=self.start_epoch,
+            loss=loss,
+            text_ratio=self.current_text_ratio,
+            best_loss=self.best_loss,
+        )
 
-    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, total_steps - warmup_steps), eta_min=1e-6)
+        # Get LoRA state
+        lora_state = self._get_lora_state_dict()
 
-    if is_main:
-        log(f"Steps per epoch: {steps_per_epoch}, Total: {total_steps}, Warmup: {warmup_steps}")
+        # Save checkpoint
+        state_dict = {
+            "adapter": unwrapped_adapter.state_dict(),
+            "lora": lora_state,
+            "optimizer": self.optimizer.state_dict(),
+            **state.to_dict()
+        }
 
-    if is_main:
-        log("\n" + "=" * 60)
-        log("STARTING STAGE 2 TRAINING")
-        log("=" * 60)
+        self.checkpoint_manager.save(
+            state_dict,
+            f"stage2_step{self.global_step}.pt",
+        )
 
-    for epoch in range(start_epoch, args.epochs):
-        # Update text ratio based on global step
-        current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+    def _finish_epoch(self, epoch: int, epoch_loss: float):
+        """Finish epoch and save checkpoints."""
+        avg_loss = epoch_loss / max(1, self.steps_per_epoch) / self.config.grad_accum
 
-        # Get unwrapped models for forward pass
-        unwrapped_adapter = accelerator.unwrap_model(adapter)
-        unwrapped_llm = accelerator.unwrap_model(llm)
+        if self.is_main:
+            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {self.current_text_ratio:.1f}")
 
-        adapter.train()
-        llm.train()
-        epoch_loss = 0
-        accum_loss = 0
+            self.accelerator.wait_for_everyone()
+            unwrapped_adapter = self.accelerator.unwrap_model(self.adapter)
+            lora_state = self._get_lora_state_dict()
 
-        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}", disable=not is_main)
+            state = TrainingState(
+                step=self.global_step,
+                epoch=epoch + 1,
+                loss=avg_loss,
+                text_ratio=self.current_text_ratio,
+                best_loss=self.best_loss,
+            )
+
+            # Save epoch checkpoint
+            state_dict = {
+                "adapter": unwrapped_adapter.state_dict(),
+                "lora": lora_state,
+                "optimizer": self.optimizer.state_dict(),
+                **state.to_dict()
+            }
+
+            self.checkpoint_manager.save(
+                state_dict,
+                f"stage2_epoch{epoch+1}.pt",
+            )
+
+            # Save best if improved
+            if avg_loss < self.best_loss:
+                self.best_loss = avg_loss
+                state.best_loss = self.best_loss
+
+                self.checkpoint_manager.save_best(
+                    unwrapped_adapter.state_dict(),
+                    state,
+                    lora_state=lora_state,
+                )
+                log(f"Best model saved! Loss: {self.best_loss:.4f}")
 
-        for batch_idx, batch in enumerate(pbar):
-            # Update text ratio dynamically
-            current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
 
-            whisper = batch["whisper"]
-            interleaved = batch["interleaved"]
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Stage 2: Adapter + LoRA with Interleaved Output"
+    )
 
-            # Use gradient accumulation context
-            with accelerator.accumulate(adapter, llm):
-                # Forward through adapter
-                audio_embeds = adapter(whisper)
+    # Data
+    parser.add_argument("--data", type=str, required=True,
+                        help="Dataset path(s), comma-separated")
+    parser.add_argument("--output_dir", type=str, default="./checkpoints")
 
-                # Get token embeddings for interleaved sequence (teacher forcing)
-                input_tokens = interleaved[:, :-1].clamp(min=0)
-                with torch.no_grad():
-                    # Access base model's embed_tokens through PEFT wrapper
-                    token_embeds = unwrapped_llm.get_base_model().model.embed_tokens(input_tokens)
+    # Stage 1 checkpoint (required)
+    parser.add_argument("--stage1_ckpt", type=str, required=True,
+                        help="Path to Stage 1 best checkpoint")
 
-                # Combine: [audio_embeds] + [token_embeds]
-                combined = torch.cat([audio_embeds, token_embeds], dim=1)
+    # Training
+    parser.add_argument("--lr", type=float, default=5e-5)
+    parser.add_argument("--epochs", type=int, default=3,
+                        help="3+ epochs for joint training")
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--grad_accum", type=int, default=None)
+    parser.add_argument("--warmup_ratio", type=float, default=0.03)
+    parser.add_argument("--max_grad_norm", type=float, default=1.0)
+    parser.add_argument("--label_smoothing", type=float, default=0.1)
+    parser.add_argument("--max_seq_len", type=int, default=2048)
 
-                # Forward through LLM with LoRA
-                outputs = llm(inputs_embeds=combined, use_cache=False)
-                logits = outputs.logits
+    # LoRA
+    parser.add_argument("--lora_r", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=int, default=32)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
 
-                # Loss: predict interleaved tokens after audio prefix
-                audio_len = audio_embeds.shape[1]
-                seq_len = interleaved.shape[1]
+    # Scheduled interleaving
+    parser.add_argument("--initial_text_ratio", type=float, default=0.9)
+    parser.add_argument("--decay_steps", type=int, default=300)
+    parser.add_argument("--dynamic_decay", action="store_true")
+    parser.add_argument("--final_audio_portion", type=float, default=0.2)
 
-                seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]
-                targets = interleaved
+    # Model
+    parser.add_argument("--model_path", type=str,
+                        default="canopylabs/3b-es_it-ft-research_release")
 
-                loss = F.cross_entropy(
-                    seq_logits.reshape(-1, logits.size(-1)),
-                    targets.reshape(-1),
-                    ignore_index=-100,
-                    label_smoothing=args.label_smoothing
-                )
+    # Checkpointing
+    parser.add_argument("--save_steps", type=int, default=200)
+    parser.add_argument("--resume", type=str, default=None)
 
-                accelerator.backward(loss)
-                accum_loss += loss.item()
+    # Memory
+    parser.add_argument("--vram_fraction", type=float, default=0.80)
+    parser.add_argument("--ram_limit_gb", type=float, default=None)
+    parser.add_argument("--gradient_checkpointing", action="store_true")
 
-                # Clip gradients and step
-                if accelerator.sync_gradients:
-                    accelerator.clip_grad_norm_(all_params, args.max_grad_norm)
+    # Modes
+    parser.add_argument("--demo", action="store_true")
+    parser.add_argument("--test", action="store_true")
 
-                optimizer.step()
-                optimizer.zero_grad()
+    return parser.parse_args()
 
-            # Update after accumulation
-            if accelerator.sync_gradients:
-                if global_step < warmup_steps:
-                    lr_scale = (global_step + 1) / warmup_steps
-                    for pg in optimizer.param_groups:
-                        pg["lr"] = args.lr * lr_scale
-                else:
-                    scheduler.step()
 
-                global_step += 1
-                epoch_loss += accum_loss
+def main():
+    args = parse_args()
+    config = TrainingConfig.from_args(args)
 
-                pbar.set_postfix(
-                    loss=f"{accum_loss:.4f}",
-                    text_ratio=f"{current_text_ratio:.1f}",
-                    lr=f"{optimizer.param_groups[0]['lr']:.2e}"
-                )
+    lora_config = LoRAConfig(
+        r=args.lora_r,
+        alpha=args.lora_alpha,
+        dropout=args.lora_dropout,
+    )
 
-                # Save checkpoint
-                if global_step % args.save_steps == 0 and is_main:
-                    accelerator.wait_for_everyone()
-                    unwrapped_adapter = accelerator.unwrap_model(adapter)
-                    unwrapped_llm = accelerator.unwrap_model(llm)
-                    ckpt_path = os.path.join(args.output_dir, f"stage2_step{global_step}.pt")
-                    save_checkpoint_async({
-                        "adapter": unwrapped_adapter.state_dict(),
-                        "lora": unwrapped_llm.state_dict(),
-                        "optimizer": optimizer.state_dict(),
-                        "step": global_step,
-                        "epoch": epoch,
-                        "loss": accum_loss,
-                        "text_ratio": current_text_ratio
-                    }, ckpt_path, is_main)
-
-                accum_loss = 0
-
-        # Epoch end - divide by grad_accum to get per-batch average
-        avg_loss = epoch_loss / max(1, steps_per_epoch) / args.grad_accum
-
-        if is_main:
-            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {current_text_ratio:.1f}")
-
-            accelerator.wait_for_everyone()
-            unwrapped_adapter = accelerator.unwrap_model(adapter)
-            unwrapped_llm = accelerator.unwrap_model(llm)
-            ckpt_path = os.path.join(args.output_dir, f"stage2_epoch{epoch+1}.pt")
-            save_checkpoint_async({
-                "adapter": unwrapped_adapter.state_dict(),
-                "lora": unwrapped_llm.state_dict(),
-                "optimizer": optimizer.state_dict(),
-                "step": global_step,
-                "epoch": epoch + 1,
-                "loss": avg_loss,
-                "text_ratio": current_text_ratio
-            }, ckpt_path, is_main)
-
-            if avg_loss < best_loss:
-                best_loss = avg_loss
-                best_path = os.path.join(args.output_dir, "stage2_best.pt")
-                save_checkpoint_async({
-                    "adapter": unwrapped_adapter.state_dict(),
-                    "lora": unwrapped_llm.state_dict(),
-                    "step": global_step,
-                    "epoch": epoch + 1,
-                    "loss": best_loss,
-                    "text_ratio": current_text_ratio
-                }, best_path, is_main)
-                log(f"Best model saved! Loss: {best_loss:.4f}")
-
-    # Finish
-    accelerator.wait_for_everyone()
-    if is_main:
-        wait_for_checkpoints()
-        log("\n" + "=" * 60)
-        log("STAGE 2 COMPLETE!")
-        log(f"Best loss: {best_loss:.4f}")
-        log(f"Final text_ratio: {current_text_ratio:.1f}")
-        log("=" * 60)
-        log("\nCheckpoints saved:")
-        log(f"  Best: {args.output_dir}/stage2_best.pt")
-        log(f"  Last: {args.output_dir}/stage2_epoch{args.epochs}.pt")
+    trainer = Stage2Trainer(config, args.stage1_ckpt, lora_config)
+    trainer.setup()
+    trainer.train()
 
 
 if __name__ == "__main__":

passo4_inference.py

@@ -14,12 +14,12 @@ from pathlib import Path
 
 # SNAC token configuration
 SNAC_BASE = 128266
+SNAC_MAX = SNAC_BASE + 7 * 4096  # 7 positions per frame, 4096 tokens each
 EOS_TOKEN = 128009
 
 def load_models(checkpoint_path: str, device: str = "cuda"):
     """Load all models for inference."""
     from transformers import WhisperModel, WhisperFeatureExtractor, AutoTokenizer
-    from peft import PeftModel
     import snac
 
     print("Loading models...")
@@ -34,8 +34,11 @@ def load_models(checkpoint_path: str, device: str = "cuda"):
     print(f"  Loading checkpoint: {checkpoint_path}")
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
 
+    # Debug: show checkpoint keys
+    print(f"  Checkpoint keys: {list(checkpoint.keys())}")
+
     # Load LLM with LoRA
-    print("  Loading LLM with LoRA...")
+    print("  Loading LLM...")
     from transformers import AutoModelForCausalLM
 
     llm = AutoModelForCausalLM.from_pretrained(
@@ -45,20 +48,41 @@ def load_models(checkpoint_path: str, device: str = "cuda"):
     )
 
     # Load LoRA weights if present
-    if 'lora' in checkpoint or 'lora_state_dict' in checkpoint:
-        print("  Loading LoRA weights...")
-        from peft import LoraConfig, get_peft_model
-        lora_config = LoraConfig(
-            r=16,
-            lora_alpha=32,
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
-            lora_dropout=0.05,
-            bias="none",
-            task_type="CAUSAL_LM"
-        )
-        llm = get_peft_model(llm, lora_config)
-        lora_key = 'lora' if 'lora' in checkpoint else 'lora_state_dict'
-        llm.load_state_dict(checkpoint[lora_key], strict=False)
+    lora_loaded = False
+    for lora_key in ['lora', 'lora_state_dict']:
+        if lora_key in checkpoint:
+            print(f"  Found LoRA weights with key '{lora_key}'")
+            from peft import LoraConfig, get_peft_model
+            lora_config = LoraConfig(
+                r=16,
+                lora_alpha=32,
+                target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+                lora_dropout=0.0,  # No dropout for inference
+                bias="none",
+                task_type="CAUSAL_LM"
+            )
+            llm = get_peft_model(llm, lora_config)
+
+            # Load the LoRA state dict
+            lora_state = checkpoint[lora_key]
+            print(f"  LoRA state dict has {len(lora_state)} keys")
+            if len(lora_state) > 0:
+                print(f"  Sample keys: {list(lora_state.keys())[:3]}")
+
+            # Try to load with strict=False to see what matches
+            result = llm.load_state_dict(lora_state, strict=False)
+            print(f"  LoRA load - Missing: {len(result.missing_keys)}, Unexpected: {len(result.unexpected_keys)}")
+            if result.missing_keys:
+                print(f"    Missing keys (first 3): {result.missing_keys[:3]}")
+            if result.unexpected_keys:
+                print(f"    Unexpected keys (first 3): {result.unexpected_keys[:3]}")
+
+            lora_loaded = True
+            break
+
+    if not lora_loaded:
+        print("  WARNING: No LoRA weights found in checkpoint! Model will use base weights only.")
+        print("  This might result in no SNAC token generation.")
 
     llm.eval()
 
@@ -71,9 +95,18 @@ def load_models(checkpoint_path: str, device: str = "cuda"):
         downsample=5
     ).to(device)
 
-    adapter_key = 'adapter' if 'adapter' in checkpoint else 'adapter_state_dict'
-    if adapter_key in checkpoint:
-        adapter.load_state_dict(checkpoint[adapter_key])
+    adapter_loaded = False
+    for adapter_key in ['adapter', 'adapter_state_dict']:
+        if adapter_key in checkpoint:
+            print(f"  Found adapter weights with key '{adapter_key}'")
+            result = adapter.load_state_dict(checkpoint[adapter_key], strict=False)
+            print(f"  Adapter load - Missing: {len(result.missing_keys)}, Unexpected: {len(result.unexpected_keys)}")
+            adapter_loaded = True
+            break
+
+    if not adapter_loaded:
+        print("  WARNING: No adapter weights found in checkpoint!")
+
     adapter.eval()
 
     # Load SNAC decoder
@@ -84,6 +117,12 @@ def load_models(checkpoint_path: str, device: str = "cuda"):
     # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained("canopylabs/3b-es_it-ft-research_release")
 
+    # Debug tokenizer info
+    print(f"  Tokenizer vocab size: {tokenizer.vocab_size}")
+    print(f"  BOS token: {tokenizer.bos_token} (id={tokenizer.bos_token_id})")
+    print(f"  EOS token: {tokenizer.eos_token} (id={tokenizer.eos_token_id})")
+    print(f"  SNAC token range: {SNAC_BASE} to {SNAC_BASE + 3*4096 - 1}")
+
     print("Models loaded!")
     return whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer
 
@@ -109,39 +148,60 @@ def encode_audio(audio_path: str, whisper_model, feature_extractor, adapter, dev
 
 
 def decode_snac_tokens(tokens: list, snac_model, device: str):
-    """Decode SNAC tokens to audio waveform."""
-    # Remove offsets and reshape
-    # SNAC has 3 layers with 1:2:4 ratio
+    """Decode SNAC tokens to audio waveform.
+
+    SNAC uses 3 hierarchical layers with 1:2:4 ratio.
+    Each "frame" has 7 tokens in order:
+      - 1 token from layer 0 (position 0)
+      - 2 tokens from layer 1 (positions 1, 2)
+      - 4 tokens from layer 2 (positions 3, 4, 5, 6)
+
+    Tokens are offset by: SNAC_BASE + (position % 7) * 4096
+    """
+    if len(tokens) == 0:
+        print("  Warning: No SNAC tokens to decode")
+        return np.zeros(24000, dtype=np.float32)  # 1 second of silence
+
     layer0_tokens = []
     layer1_tokens = []
     layer2_tokens = []
 
-    # Parse interleaved tokens: 1 from layer0, 2 from layer1, 4 from layer2 per frame
-    i = 0
-    while i < len(tokens):
-        if i < len(tokens):
-            layer0_tokens.append(tokens[i] - SNAC_BASE)
-            i += 1
-        if i < len(tokens):
-            layer1_tokens.append(tokens[i] - SNAC_BASE - 4096)
-            i += 1
-        if i < len(tokens):
-            layer1_tokens.append(tokens[i] - SNAC_BASE - 4096)
-            i += 1
-        for _ in range(4):
-            if i < len(tokens):
-                layer2_tokens.append(tokens[i] - SNAC_BASE - 8192)
-                i += 1
-
-    # Create tensors
-    min_len = min(len(layer0_tokens), len(layer1_tokens) // 2, len(layer2_tokens) // 4)
-    if min_len == 0:
-        return np.zeros(24000, dtype=np.float32)  # 1 second of silence
+    # Parse tokens by removing position-based offsets
+    for i, tok in enumerate(tokens):
+        pos = i % 7
+        # Remove the offset to get the original code
+        original = tok - SNAC_BASE - (pos * 4096)
+
+        # Ensure codes are valid (0-4095)
+        if original < 0 or original >= 4096:
+            print(f"  Warning: Invalid code {original} at position {i} (token={tok}, pos={pos})")
+            original = max(0, min(4095, original))
+
+        if pos == 0:
+            layer0_tokens.append(original)
+        elif pos in [1, 2]:
+            layer1_tokens.append(original)
+        else:  # pos in [3, 4, 5, 6]
+            layer2_tokens.append(original)
+
+    # Calculate how many complete frames we have
+    n_frames = len(tokens) // 7
+
+    if n_frames == 0:
+        print(f"  Warning: Not enough tokens for a complete frame ({len(tokens)} tokens)")
+        return np.zeros(24000, dtype=np.float32)
+
+    print(f"  Decoding {n_frames} frames ({len(layer0_tokens)} L0, {len(layer1_tokens)} L1, {len(layer2_tokens)} L2)")
+
+    # Truncate to complete frames
+    layer0_tokens = layer0_tokens[:n_frames]
+    layer1_tokens = layer1_tokens[:n_frames * 2]
+    layer2_tokens = layer2_tokens[:n_frames * 4]
 
     codes = [
-        torch.tensor([layer0_tokens[:min_len]], dtype=torch.long, device=device),
-        torch.tensor([layer1_tokens[:min_len * 2]], dtype=torch.long, device=device),
-        torch.tensor([layer2_tokens[:min_len * 4]], dtype=torch.long, device=device)
+        torch.tensor([layer0_tokens], dtype=torch.long, device=device),
+        torch.tensor([layer1_tokens], dtype=torch.long, device=device),
+        torch.tensor([layer2_tokens], dtype=torch.long, device=device)
     ]
 
     # Decode
@@ -156,45 +216,104 @@ def generate_response(
     whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
     device: str,
     max_new_tokens: int = 500,
-    temperature: float = 0.7
+    temperature: float = 0.7,
+    debug: bool = True,
+    use_prompt: bool = False
 ):
     """Generate speech response from audio input."""
 
     # Encode input audio
     print("Encoding input audio...")
     audio_embeddings = encode_audio(audio_input, whisper_model, feature_extractor, adapter, device)
+    print(f"  Audio embeddings shape: {audio_embeddings.shape}")
 
     # Generate with LLM
     print("Generating response...")
 
-    # Create input embeds by concatenating audio embeddings with start tokens
     with torch.no_grad():
         # Get embeddings layer
         embed_layer = llm.get_input_embeddings()
 
-        # Add start token
-        start_tokens = tokenizer.encode("<|start|>", add_special_tokens=False, return_tensors="pt").to(device)
-        start_embeds = embed_layer(start_tokens)
-
-        # Concatenate: audio + start (ensure same dtype)
-        input_embeds = torch.cat([audio_embeddings.to(torch.bfloat16), start_embeds.to(torch.bfloat16)], dim=1)
-
-        # Generate
-        outputs = llm.generate(
-            inputs_embeds=input_embeds,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            do_sample=True,
-            top_p=0.9,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=EOS_TOKEN
-        )
-
-    # Extract SNAC tokens from output
+        if use_prompt:
+            # Option A: Add BOS token after audio
+            bos_id = tokenizer.bos_token_id or 128000
+            prompt_tokens = torch.tensor([[bos_id]], dtype=torch.long, device=device)
+            prompt_embeds = embed_layer(prompt_tokens)
+            input_embeds = torch.cat([audio_embeddings.to(torch.bfloat16), prompt_embeds.to(torch.bfloat16)], dim=1)
+        else:
+            # Option B: Just audio embeddings (as trained)
+            # During training, target directly follows audio embeddings
+            input_embeds = audio_embeddings.to(torch.bfloat16)
+
+        print(f"  Input embeds shape: {input_embeds.shape}")
+
+        # Create attention mask
+        attention_mask = torch.ones(input_embeds.shape[:2], dtype=torch.long, device=device)
+
+        # Different generation strategies
+        # Strategy 1: Pure greedy (most stable)
+        if temperature <= 0.01:
+            outputs = llm.generate(
+                inputs_embeds=input_embeds,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=EOS_TOKEN,
+                use_cache=True
+            )
+        # Strategy 2: Sampling with nucleus
+        else:
+            outputs = llm.generate(
+                inputs_embeds=input_embeds,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                top_k=50,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=EOS_TOKEN,
+                use_cache=True
+            )
+
+    # Extract generated tokens
     generated_tokens = outputs[0].tolist()
-    snac_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_BASE + 3 * 4096]
 
-    print(f"Generated {len(snac_tokens)} SNAC tokens")
+    if debug:
+        print(f"\n=== DEBUG: Generated {len(generated_tokens)} total tokens ===")
+        # Show first 50 tokens
+        print(f"First 50 tokens: {generated_tokens[:50]}")
+        # Show last 20 tokens
+        print(f"Last 20 tokens: {generated_tokens[-20:]}")
+
+        # Show token ranges with correct SNAC range
+        text_tokens = [t for t in generated_tokens if t < SNAC_BASE]
+        snac_range_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_MAX]
+        other_high = [t for t in generated_tokens if t >= SNAC_MAX]
+
+        print(f"\nToken distribution:")
+        print(f"  Text tokens (<{SNAC_BASE}): {len(text_tokens)}")
+        print(f"  SNAC tokens ({SNAC_BASE}-{SNAC_MAX}): {len(snac_range_tokens)}")
+        print(f"  Other high tokens (>={SNAC_MAX}): {len(other_high)}")
+
+        # Try to decode text tokens
+        if text_tokens:
+            try:
+                decoded_text = tokenizer.decode(text_tokens, skip_special_tokens=False)
+                print(f"\nDecoded text: {decoded_text[:500]}")
+            except Exception as e:
+                print(f"Could not decode text: {e}")
+
+        # Show some SNAC tokens if any
+        if snac_range_tokens:
+            print(f"\nFirst 20 SNAC tokens: {snac_range_tokens[:20]}")
+
+    # Extract SNAC tokens from output (correct range)
+    snac_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_MAX]
+
+    print(f"\nGenerated {len(snac_tokens)} SNAC tokens")
 
     # Decode to audio
     print("Decoding to audio...")
@@ -208,8 +327,10 @@ def main():
     parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
     parser.add_argument("--input", type=str, required=True, help="Input audio file")
     parser.add_argument("--output", type=str, default="output.wav", help="Output audio file")
-    parser.add_argument("--max_tokens", type=int, default=500, help="Max tokens to generate")
-    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
+    parser.add_argument("--max_tokens", type=int, default=1000, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature (0 for greedy)")
+    parser.add_argument("--use_prompt", action="store_true", help="Add BOS token after audio embeddings")
+    parser.add_argument("--no_debug", action="store_true", help="Disable debug output")
     args = parser.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -226,7 +347,9 @@ def main():
         whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
         device,
         max_new_tokens=args.max_tokens,
-        temperature=args.temperature
+        temperature=args.temperature,
+        debug=not args.no_debug,
+        use_prompt=args.use_prompt
     )
 
     # Save output

training/__init__.py

@@ -0,0 +1,30 @@
+"""
+Training module for Speech-to-Speech finetuning.
+
+Follows SOLID principles:
+- Single Responsibility: Each module handles one concern
+- Open/Closed: Base classes allow extension without modification
+- Dependency Inversion: Components depend on abstractions
+"""
+
+from .config import TrainingConfig, GPUConfig
+from .data import load_sharded_dataset, InterleavedDataset, create_dataloader
+from .models import SpeechAdapter
+from .interleaving import create_interleaved_sequence, get_text_ratio
+from .checkpoint import CheckpointManager
+from .utils import setup_logging, log, get_device_info
+
+__all__ = [
+    'TrainingConfig',
+    'GPUConfig',
+    'load_sharded_dataset',
+    'InterleavedDataset',
+    'create_dataloader',
+    'SpeechAdapter',
+    'create_interleaved_sequence',
+    'get_text_ratio',
+    'CheckpointManager',
+    'setup_logging',
+    'log',
+    'get_device_info',
+]

training/checkpoint.py

@@ -0,0 +1,273 @@
+"""
+Checkpoint management for training.
+
+Single Responsibility: Only handles saving and loading checkpoints.
+"""
+
+import os
+import torch
+import threading
+from typing import Dict, Any, Optional, List
+from pathlib import Path
+from dataclasses import dataclass
+from .utils import log
+
+
+@dataclass
+class TrainingState:
+    """Immutable training state for checkpointing."""
+    step: int
+    epoch: int
+    loss: float
+    text_ratio: float
+    best_loss: float = float("inf")
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "step": self.step,
+            "epoch": self.epoch,
+            "loss": self.loss,
+            "text_ratio": self.text_ratio,
+            "best_loss": self.best_loss,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> 'TrainingState':
+        return cls(
+            step=d.get("step", 0),
+            epoch=d.get("epoch", 0),
+            loss=d.get("loss", 0.0),
+            text_ratio=d.get("text_ratio", 0.9),
+            best_loss=d.get("best_loss", float("inf")),
+        )
+
+
+class CheckpointManager:
+    """
+    Manages checkpoint saving and loading with async support.
+
+    Single Responsibility: Only handles checkpoint I/O.
+    Open/Closed: Can extend with new checkpoint formats without modification.
+    """
+
+    def __init__(
+        self,
+        output_dir: str,
+        prefix: str = "checkpoint",
+        verbose: bool = True,
+        max_checkpoints: Optional[int] = None
+    ):
+        """
+        Initialize checkpoint manager.
+
+        Args:
+            output_dir: Directory for saving checkpoints
+            prefix: Prefix for checkpoint filenames
+            verbose: Whether to log operations
+            max_checkpoints: Maximum checkpoints to keep (None = keep all)
+        """
+        self.output_dir = Path(output_dir)
+        self.prefix = prefix
+        self.verbose = verbose
+        self.max_checkpoints = max_checkpoints
+        self._save_threads: List[threading.Thread] = []
+
+        # Create output directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+    def save(
+        self,
+        state_dict: Dict[str, Any],
+        filename: str,
+        async_save: bool = True
+    ) -> str:
+        """
+        Save checkpoint.
+
+        Args:
+            state_dict: State dictionary to save
+            filename: Checkpoint filename
+            async_save: Whether to save asynchronously
+
+        Returns:
+            Path to saved checkpoint
+        """
+        path = self.output_dir / filename
+
+        if async_save:
+            self._save_async(state_dict, path)
+        else:
+            self._save_sync(state_dict, path)
+
+        return str(path)
+
+    def save_step(
+        self,
+        adapter_state: Dict[str, Any],
+        optimizer_state: Dict[str, Any],
+        training_state: TrainingState,
+        async_save: bool = True
+    ) -> str:
+        """Save step checkpoint."""
+        state_dict = {
+            "adapter": adapter_state,
+            "optimizer": optimizer_state,
+            **training_state.to_dict()
+        }
+        filename = f"{self.prefix}_step{training_state.step}.pt"
+        return self.save(state_dict, filename, async_save)
+
+    def save_epoch(
+        self,
+        adapter_state: Dict[str, Any],
+        optimizer_state: Dict[str, Any],
+        training_state: TrainingState,
+        async_save: bool = True
+    ) -> str:
+        """Save epoch checkpoint."""
+        state_dict = {
+            "adapter": adapter_state,
+            "optimizer": optimizer_state,
+            **training_state.to_dict()
+        }
+        filename = f"{self.prefix}_epoch{training_state.epoch}.pt"
+        return self.save(state_dict, filename, async_save)
+
+    def save_best(
+        self,
+        adapter_state: Dict[str, Any],
+        training_state: TrainingState,
+        lora_state: Optional[Dict[str, Any]] = None,
+        async_save: bool = True
+    ) -> str:
+        """Save best model checkpoint."""
+        state_dict = {
+            "adapter": adapter_state,
+            **training_state.to_dict()
+        }
+        if lora_state is not None:
+            state_dict["lora"] = lora_state
+
+        filename = f"{self.prefix}_best.pt"
+        return self.save(state_dict, filename, async_save)
+
+    def load(self, path: str) -> Dict[str, Any]:
+        """
+        Load checkpoint.
+
+        Args:
+            path: Path to checkpoint
+
+        Returns:
+            Loaded state dictionary
+        """
+        if self.verbose:
+            log(f"Loading checkpoint: {path}")
+        return torch.load(path, map_location="cpu", weights_only=False)
+
+    def load_latest(self) -> Optional[Dict[str, Any]]:
+        """Load the most recent checkpoint."""
+        checkpoints = self._get_checkpoints()
+        if not checkpoints:
+            return None
+        return self.load(str(checkpoints[-1]))
+
+    def wait_for_saves(self):
+        """Wait for all async saves to complete."""
+        for t in self._save_threads:
+            t.join()
+        self._save_threads = []
+
+    def _save_sync(self, state_dict: Dict[str, Any], path: Path):
+        """Synchronous save."""
+        # Copy tensors to CPU
+        state_copy = self._copy_to_cpu(state_dict)
+        torch.save(state_copy, path)
+        if self.verbose:
+            log(f"[Checkpoint] Saved: {path.name}")
+
+    def _save_async(self, state_dict: Dict[str, Any], path: Path):
+        """Asynchronous save."""
+        # Clean up completed threads
+        self._save_threads = [t for t in self._save_threads if t.is_alive()]
+
+        # Copy tensors to CPU
+        state_copy = self._copy_to_cpu(state_dict)
+
+        def _save():
+            try:
+                torch.save(state_copy, path)
+                if self.verbose:
+                    log(f"[Checkpoint] Saved: {path.name}")
+            except Exception as e:
+                if self.verbose:
+                    log(f"[Checkpoint] Error saving {path.name}: {e}")
+
+        thread = threading.Thread(target=_save, daemon=True)
+        thread.start()
+        self._save_threads.append(thread)
+
+        # Cleanup old checkpoints if needed
+        if self.max_checkpoints:
+            self._cleanup_old_checkpoints()
+
+    def _copy_to_cpu(self, obj: Any) -> Any:
+        """Recursively copy tensors to CPU."""
+        if isinstance(obj, torch.Tensor):
+            return obj.detach().cpu().clone()
+        elif isinstance(obj, dict):
+            return {k: self._copy_to_cpu(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._copy_to_cpu(v) for v in obj]
+        return obj
+
+    def _get_checkpoints(self) -> List[Path]:
+        """Get sorted list of checkpoint files."""
+        pattern = f"{self.prefix}_step*.pt"
+        checkpoints = sorted(
+            self.output_dir.glob(pattern),
+            key=lambda p: int(p.stem.split("step")[-1])
+        )
+        return checkpoints
+
+    def _cleanup_old_checkpoints(self):
+        """Remove old checkpoints beyond max_checkpoints."""
+        if not self.max_checkpoints:
+            return
+
+        checkpoints = self._get_checkpoints()
+        while len(checkpoints) > self.max_checkpoints:
+            oldest = checkpoints.pop(0)
+            try:
+                oldest.unlink()
+                if self.verbose:
+                    log(f"[Checkpoint] Removed old: {oldest.name}")
+            except Exception:
+                pass
+
+
+class Stage1CheckpointManager(CheckpointManager):
+    """Checkpoint manager for Stage 1 training."""
+
+    def __init__(self, output_dir: str, **kwargs):
+        super().__init__(output_dir, prefix="stage1", **kwargs)
+
+
+class Stage2CheckpointManager(CheckpointManager):
+    """Checkpoint manager for Stage 2 training."""
+
+    def __init__(self, output_dir: str, **kwargs):
+        super().__init__(output_dir, prefix="stage2", **kwargs)
+
+    def save_best(
+        self,
+        adapter_state: Dict[str, Any],
+        training_state: TrainingState,
+        lora_state: Dict[str, Any],
+        async_save: bool = True
+    ) -> str:
+        """Save best model with LoRA weights."""
+        return super().save_best(
+            adapter_state, training_state,
+            lora_state=lora_state, async_save=async_save
+        )

training/config.py

@@ -0,0 +1,231 @@
+"""
+Configuration management using dataclasses.
+
+Single Responsibility: Only handles configuration.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, List
+import os
+import torch
+
+
+# ============================================================
+# Constants (Single source of truth)
+# ============================================================
+SNAC_BASE_OFFSET = 128266
+SNAC_LAYERS_PER_FRAME = 7
+SNAC_LAYER_OFFSET = 4096
+EOS_TOKEN = 128009
+
+# Model defaults
+DEFAULT_WHISPER_DIM = 1280
+DEFAULT_LLM_DIM = 3072
+DEFAULT_DOWNSAMPLE = 5
+DEFAULT_INTERMEDIATE_DIM = 2048
+DEFAULT_MODEL_PATH = "canopylabs/3b-es_it-ft-research_release"
+
+# LoRA defaults
+DEFAULT_LORA_R = 16
+DEFAULT_LORA_ALPHA = 32
+DEFAULT_LORA_DROPOUT = 0.05
+DEFAULT_LORA_MODULES = [
+    "q_proj", "k_proj", "v_proj", "o_proj",
+    "gate_proj", "up_proj", "down_proj"
+]
+
+
+@dataclass
+class GPUConfig:
+    """GPU configuration detected at runtime."""
+    name: str = "Unknown"
+    vram_gb: int = 0
+    batch_size: int = 2
+    grad_accum: int = 16
+    dtype: torch.dtype = torch.float32
+    device_type: str = "cpu"
+
+    @classmethod
+    def auto_detect(cls) -> 'GPUConfig':
+        """Detect GPU and return optimal configuration."""
+        config = cls()
+
+        # Try CUDA (NVIDIA)
+        if torch.cuda.is_available():
+            try:
+                props = torch.cuda.get_device_properties(0)
+                config.vram_gb = props.total_memory // (1024**3)
+                config.name = props.name
+                config.device_type = "cuda"
+                config.dtype = torch.bfloat16
+            except Exception:
+                pass
+
+        # Try MPS (Apple Silicon)
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            config.name = "Apple Silicon (MPS)"
+            config.device_type = "mps"
+            config.dtype = torch.float32
+            try:
+                import subprocess
+                result = subprocess.run(
+                    ['sysctl', '-n', 'hw.memsize'],
+                    capture_output=True, text=True
+                )
+                total_mem = int(result.stdout.strip()) // (1024**3)
+                config.vram_gb = total_mem // 2
+            except Exception:
+                config.vram_gb = 8
+
+        # Try ROCm (AMD)
+        elif hasattr(torch, 'hip') or os.environ.get('ROCM_HOME'):
+            try:
+                if torch.cuda.is_available():
+                    props = torch.cuda.get_device_properties(0)
+                    config.vram_gb = props.total_memory // (1024**3)
+                    config.name = f"AMD {props.name}"
+                    config.device_type = "cuda"
+                    config.dtype = torch.bfloat16
+            except Exception:
+                config.name = "AMD ROCm"
+                config.vram_gb = 16
+
+        # Fallback: nvidia-smi
+        if config.vram_gb == 0:
+            try:
+                import subprocess
+                result = subprocess.run(
+                    ['nvidia-smi', '--query-gpu=name,memory.total',
+                     '--format=csv,noheader,nounits'],
+                    capture_output=True, text=True
+                )
+                lines = result.stdout.strip().split('\n')
+                config.name, vram_mb = lines[0].split(', ')
+                config.vram_gb = int(vram_mb) // 1024
+                config.device_type = "cuda"
+                config.dtype = torch.bfloat16
+            except Exception:
+                pass
+
+        # Set batch size based on VRAM
+        config.batch_size, config.grad_accum = cls._get_batch_config(config.vram_gb)
+        return config
+
+    @staticmethod
+    def _get_batch_config(vram_gb: int) -> tuple:
+        """Get optimal batch size and gradient accumulation based on VRAM."""
+        if vram_gb >= 140:  # H200 (141GB)
+            return 12, 3  # H200: batch=12, effective=36
+        elif vram_gb >= 80:
+            return 6, 5
+        elif vram_gb >= 35:
+            return 4, 8
+        elif vram_gb >= 16:
+            return 2, 16
+        else:
+            return 1, 32
+
+
+@dataclass
+class TrainingConfig:
+    """Training configuration with sensible defaults."""
+    # Data
+    data_paths: List[str] = field(default_factory=list)
+    output_dir: str = "./checkpoints"
+
+    # Training hyperparameters
+    learning_rate: float = 5e-5
+    epochs: int = 2
+    batch_size: Optional[int] = None
+    grad_accum: Optional[int] = None
+    warmup_ratio: float = 0.03
+    max_grad_norm: float = 1.0
+    label_smoothing: float = 0.1
+    weight_decay: float = 0.01
+
+    # Sequence limits
+    max_audio_len: int = 500
+    max_seq_len: int = 2048
+
+    # Scheduled interleaving (IST-LM)
+    initial_text_ratio: float = 0.9
+    decay_steps: int = 300
+    dynamic_decay: bool = False
+    no_decay: bool = False  # Stage 1: keep text_ratio fixed
+    final_audio_portion: float = 0.2
+
+    # Model
+    model_path: str = DEFAULT_MODEL_PATH
+
+    # Checkpointing
+    save_steps: int = 200
+    resume_from: Optional[str] = None
+
+    # Memory
+    vram_fraction: float = 0.80
+    ram_limit_gb: Optional[float] = None
+    gradient_checkpointing: bool = False
+
+    # Mode flags
+    demo_mode: bool = False
+    test_mode: bool = False
+
+    def __post_init__(self):
+        """Apply GPU auto-detection if batch_size not set."""
+        if self.batch_size is None or self.grad_accum is None:
+            gpu_config = GPUConfig.auto_detect()
+            if self.batch_size is None:
+                self.batch_size = gpu_config.batch_size
+            if self.grad_accum is None:
+                self.grad_accum = gpu_config.grad_accum
+
+    @classmethod
+    def from_args(cls, args) -> 'TrainingConfig':
+        """Create config from argparse namespace."""
+        return cls(
+            data_paths=[p.strip() for p in args.data.split(",")],
+            output_dir=args.output_dir,
+            learning_rate=args.lr,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            grad_accum=args.grad_accum,
+            warmup_ratio=args.warmup_ratio,
+            max_grad_norm=args.max_grad_norm,
+            label_smoothing=args.label_smoothing,
+            max_seq_len=args.max_seq_len,
+            initial_text_ratio=args.initial_text_ratio,
+            decay_steps=args.decay_steps,
+            dynamic_decay=getattr(args, 'dynamic_decay', False),
+            no_decay=getattr(args, 'no_decay', False),
+            final_audio_portion=getattr(args, 'final_audio_portion', 0.2),
+            model_path=args.model_path,
+            save_steps=args.save_steps,
+            resume_from=args.resume,
+            vram_fraction=args.vram_fraction,
+            ram_limit_gb=args.ram_limit_gb,
+            gradient_checkpointing=args.gradient_checkpointing,
+            demo_mode=args.demo,
+            test_mode=args.test,
+        )
+
+
+@dataclass
+class LoRAConfig:
+    """LoRA configuration for Stage 2."""
+    r: int = DEFAULT_LORA_R
+    alpha: int = DEFAULT_LORA_ALPHA
+    dropout: float = DEFAULT_LORA_DROPOUT
+    target_modules: List[str] = field(default_factory=lambda: DEFAULT_LORA_MODULES.copy())
+    bias: str = "none"
+
+    def to_peft_config(self):
+        """Convert to PEFT LoraConfig."""
+        from peft import LoraConfig as PeftLoraConfig, TaskType
+        return PeftLoraConfig(
+            r=self.r,
+            lora_alpha=self.alpha,
+            lora_dropout=self.dropout,
+            target_modules=self.target_modules,
+            bias=self.bias,
+            task_type=TaskType.CAUSAL_LM,
+        )

training/data.py

@@ -0,0 +1,734 @@
+"""
+Data loading and dataset management.
+
+Single Responsibility: Only handles data loading and dataset creation.
+Open/Closed: Can extend with new dataset formats without modifying existing code.
+
+Optimizations:
+- Lazy loading for memory efficiency
+- Sequence length bucketing for reduced padding overhead
+- LRU cache for batch files
+"""
+
+import gc
+import random
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, ConcatDataset, Sampler
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Callable, Tuple, Iterator
+from .utils import log
+
+
+# ============================================================
+# Sequence Length Bucketing (Reduces padding overhead)
+# ============================================================
+class BucketBatchSampler(Sampler[List[int]]):
+    """
+    Batch sampler that groups samples by sequence length into buckets.
+
+    This reduces padding overhead by batching similar-length sequences together.
+    Based on TensorFlow's bucket_by_sequence_length concept.
+
+    Benefits:
+    - Reduces wasted computation on padding tokens
+    - More consistent memory usage per batch
+    - Can improve training speed by 10-30%
+
+    Args:
+        lengths: List of sequence lengths for each sample
+        batch_size: Number of samples per batch
+        bucket_boundaries: Length boundaries for buckets (auto-computed if None)
+        shuffle: Whether to shuffle within buckets
+        drop_last: Whether to drop incomplete batches
+    """
+
+    def __init__(
+        self,
+        lengths: List[int],
+        batch_size: int,
+        bucket_boundaries: Optional[List[int]] = None,
+        shuffle: bool = True,
+        drop_last: bool = False,
+    ):
+        self.lengths = lengths
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+        # Auto-compute bucket boundaries if not provided
+        if bucket_boundaries is None:
+            # Create ~10 buckets based on length distribution
+            sorted_lens = sorted(lengths)
+            n = len(sorted_lens)
+            bucket_boundaries = [
+                sorted_lens[int(n * p)] for p in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+            ]
+            # Remove duplicates and sort
+            bucket_boundaries = sorted(set(bucket_boundaries))
+
+        self.bucket_boundaries = bucket_boundaries
+
+        # Assign samples to buckets
+        self.buckets: Dict[int, List[int]] = {i: [] for i in range(len(bucket_boundaries) + 1)}
+        for idx, length in enumerate(lengths):
+            bucket_id = self._get_bucket_id(length)
+            self.buckets[bucket_id].append(idx)
+
+    def _get_bucket_id(self, length: int) -> int:
+        """Find which bucket a length belongs to."""
+        for i, boundary in enumerate(self.bucket_boundaries):
+            if length <= boundary:
+                return i
+        return len(self.bucket_boundaries)
+
+    def __iter__(self) -> Iterator[List[int]]:
+        """Generate batches from buckets."""
+        # Collect all batches from all buckets
+        all_batches = []
+
+        for bucket_id, indices in self.buckets.items():
+            if not indices:
+                continue
+
+            # Shuffle within bucket
+            bucket_indices = indices.copy()
+            if self.shuffle:
+                random.shuffle(bucket_indices)
+
+            # Create batches
+            for i in range(0, len(bucket_indices), self.batch_size):
+                batch = bucket_indices[i:i + self.batch_size]
+                if len(batch) == self.batch_size or not self.drop_last:
+                    all_batches.append(batch)
+
+        # Shuffle batches across buckets
+        if self.shuffle:
+            random.shuffle(all_batches)
+
+        yield from all_batches
+
+    def __len__(self) -> int:
+        """Return total number of batches."""
+        total = 0
+        for indices in self.buckets.values():
+            n_batches = len(indices) // self.batch_size
+            if not self.drop_last and len(indices) % self.batch_size != 0:
+                n_batches += 1
+            total += n_batches
+        return total
+
+
+# ============================================================
+# Lazy Sharded Dataset (Memory Efficient)
+# ============================================================
+class LazyShardedDataset(Dataset):
+    """
+    Memory-efficient dataset that loads batch files on-demand.
+
+    Instead of loading all data into memory, maintains an index of
+    which sample is in which batch file, and loads batches as needed.
+
+    Also stores approximate sequence lengths for bucketing optimization.
+    """
+
+    def __init__(
+        self,
+        batch_files: List[Path],
+        tokenizer,
+        max_audio_len: int = 500,
+        max_seq_len: int = 2048,
+        cache_size: int = 3,  # Number of batches to keep in memory
+        verbose: bool = True
+    ):
+        self.batch_files = batch_files
+        self.tokenizer = tokenizer
+        self.max_audio = max_audio_len * 5
+        self.max_seq_len = max_seq_len
+        self.cache_size = cache_size
+        self.verbose = verbose
+
+        # Build index: sample_idx -> (batch_idx, local_idx)
+        self.index: List[Tuple[int, int]] = []
+        self.batch_sizes: List[int] = []
+        # Store sequence lengths for bucketing
+        self.sequence_lengths: List[int] = []
+
+        if verbose:
+            log(f"  Indexing {len(batch_files)} batch files...")
+
+        for batch_idx, bf in enumerate(batch_files):
+            # Quick load to get size and lengths
+            data = torch.load(bf, map_location="cpu", weights_only=False)
+            batch_size = len(data)
+            self.batch_sizes.append(batch_size)
+
+            for local_idx in range(batch_size):
+                self.index.append((batch_idx, local_idx))
+                # Estimate sequence length (SNAC tokens dominate)
+                item = data[local_idx]
+                snac_len = len(item.get("snac_tokens", [])) // 7 * 7
+                self.sequence_lengths.append(snac_len)
+
+            del data
+            if (batch_idx + 1) % 100 == 0:
+                gc.collect()
+                if verbose:
+                    log(f"    Indexed {batch_idx+1}/{len(batch_files)} batches ({len(self.index):,} samples)")
+
+        if verbose:
+            log(f"  Total indexed: {len(self.index):,} samples")
+
+        # LRU cache for loaded batches
+        self._cache: Dict[int, List[Dict]] = {}
+        self._cache_order: List[int] = []
+
+    def get_sequence_lengths(self) -> List[int]:
+        """Return sequence lengths for bucketing."""
+        return self.sequence_lengths
+
+    def __len__(self) -> int:
+        return len(self.index)
+
+    def _load_batch(self, batch_idx: int) -> List[Dict]:
+        """Load a batch file, using cache."""
+        if batch_idx in self._cache:
+            # Move to end of cache order (most recently used)
+            self._cache_order.remove(batch_idx)
+            self._cache_order.append(batch_idx)
+            return self._cache[batch_idx]
+
+        # Load from disk
+        data = torch.load(self.batch_files[batch_idx], map_location="cpu", weights_only=False)
+
+        # Add to cache
+        self._cache[batch_idx] = data
+        self._cache_order.append(batch_idx)
+
+        # Evict old batches if cache is full
+        while len(self._cache_order) > self.cache_size:
+            old_idx = self._cache_order.pop(0)
+            del self._cache[old_idx]
+            gc.collect()
+
+        return data
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        batch_idx, local_idx = self.index[idx]
+        batch_data = self._load_batch(batch_idx)
+        item = batch_data[local_idx]
+
+        # Process item (same as InterleavedDataset)
+        whisper = item["whisper_features"][:self.max_audio]
+        text_tokens = self._get_text_tokens(item)
+        snac_list = self._get_snac_tokens(item)
+        word_alignments = item.get("word_alignments", None)
+        answer_text = item.get("answer", "")
+
+        return {
+            "whisper": whisper,
+            "text_tokens": text_tokens,
+            "snac_tokens": snac_list,
+            "word_alignments": word_alignments,
+            "answer_text": answer_text
+        }
+
+    def _get_text_tokens(self, item: Dict[str, Any]) -> List[int]:
+        if "text_tokens" in item and len(item["text_tokens"]) > 0:
+            tt = item["text_tokens"]
+            return tt.tolist() if hasattr(tt, 'tolist') else list(tt)
+        text = item.get("answer", item.get("text", ""))
+        if isinstance(text, str) and len(text) > 0:
+            return self.tokenizer.encode(text, add_special_tokens=False)
+        return []
+
+    def _get_snac_tokens(self, item: Dict[str, Any]) -> List[int]:
+        snac = item["snac_tokens"]
+        snac_len = (len(snac) // 7) * 7
+        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
+        return snac.tolist() if hasattr(snac, 'tolist') else list(snac)
+
+
+# ============================================================
+# Sharded Dataset Loading (Industry Standard)
+# ============================================================
+class ShardedDatasetLoader:
+    """
+    Load datasets from single files or sharded batch directories.
+
+    Supports:
+    - Single .pt file
+    - Batch directory with batch_*.pt files
+    - Mixed (base file + batch files)
+
+    Dependency Inversion: Uses abstract path interface, not concrete file operations.
+    """
+
+    def __init__(self, verbose: bool = True):
+        self.verbose = verbose
+
+    def load(self, path: str) -> List[Dict[str, Any]]:
+        """Load dataset from path (file or directory)."""
+        path = Path(path)
+        samples = []
+
+        # Case 1: Explicit batches directory
+        if path.name.endswith('.batches') and path.is_dir():
+            samples = self._load_batches_dir(path)
+
+        # Case 2: Single file (possibly with batches)
+        elif path.exists() and path.is_file():
+            samples = self._load_file_with_batches(path)
+
+        # Case 3: Only batches directory exists
+        else:
+            batches_dir = Path(f"{path}.batches")
+            if batches_dir.exists() and batches_dir.is_dir():
+                samples = self._load_batches_dir(batches_dir)
+            else:
+                raise FileNotFoundError(f"No dataset found at {path}")
+
+        return samples
+
+    def _load_batches_dir(self, batches_dir: Path) -> List[Dict[str, Any]]:
+        """Load all batch files from a directory."""
+        batch_files = sorted(batches_dir.glob("batch_*.pt"))
+        if not batch_files:
+            raise FileNotFoundError(f"No batch files in {batches_dir}")
+
+        if self.verbose:
+            log(f"  Loading {len(batch_files)} batch files from {batches_dir.name}/")
+
+        samples = []
+        for i, bf in enumerate(batch_files):
+            items = torch.load(bf, map_location="cpu", weights_only=False)
+            samples.extend(items)
+            del items
+            if (i + 1) % 100 == 0:
+                gc.collect()
+                if self.verbose:
+                    log(f"    Loaded {i+1}/{len(batch_files)} batches ({len(samples):,} samples)")
+
+        return samples
+
+    def _load_file_with_batches(self, path: Path) -> List[Dict[str, Any]]:
+        """Load base file and any associated batch files."""
+        # Load base file
+        base = torch.load(path, map_location="cpu", weights_only=False, mmap=True)
+        samples = list(base)
+        del base
+        gc.collect()
+
+        if self.verbose:
+            log(f"  Base file: {len(samples):,} samples")
+
+        # Check for batch files
+        batches_dir = Path(f"{path}.batches")
+        if batches_dir.exists() and batches_dir.is_dir():
+            batch_files = sorted(batches_dir.glob("batch_*.pt"))
+            if batch_files:
+                if self.verbose:
+                    log(f"  Found {len(batch_files)} batch files")
+                for i, bf in enumerate(batch_files):
+                    items = torch.load(bf, map_location="cpu", weights_only=False)
+                    samples.extend(items)
+                    del items
+                    if (i + 1) % 100 == 0:
+                        gc.collect()
+                        if self.verbose:
+                            log(f"    Loaded {i+1}/{len(batch_files)} batches ({len(samples):,} total)")
+
+        return samples
+
+
+def load_sharded_dataset(path: str, verbose: bool = True) -> List[Dict[str, Any]]:
+    """Convenience function for loading sharded datasets."""
+    loader = ShardedDatasetLoader(verbose=verbose)
+    return loader.load(path)
+
+
+# ============================================================
+# Dataset Classes
+# ============================================================
+class InterleavedDataset(Dataset):
+    """
+    Dataset that prepares samples for interleaved training.
+
+    Single Responsibility: Only handles sample access and preprocessing.
+    """
+
+    def __init__(
+        self,
+        data: List[Dict[str, Any]],
+        tokenizer,
+        max_audio_len: int = 500,
+        max_seq_len: int = 2048
+    ):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_audio = max_audio_len * 5  # Account for downsampling
+        self.max_seq_len = max_seq_len
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        item = self.data[idx]
+
+        # Whisper features (truncate if needed)
+        whisper = item["whisper_features"][:self.max_audio]
+
+        # Text tokens - use pre-computed if available
+        text_tokens = self._get_text_tokens(item)
+
+        # SNAC tokens (ensure multiple of 7)
+        snac_list = self._get_snac_tokens(item)
+
+        # Optional fields
+        word_alignments = item.get("word_alignments", None)
+        answer_text = item.get("answer", "")
+
+        return {
+            "whisper": whisper,
+            "text_tokens": text_tokens,
+            "snac_tokens": snac_list,
+            "word_alignments": word_alignments,
+            "answer_text": answer_text
+        }
+
+    def _get_text_tokens(self, item: Dict[str, Any]) -> List[int]:
+        """Extract or generate text tokens from item."""
+        if "text_tokens" in item and len(item["text_tokens"]) > 0:
+            tt = item["text_tokens"]
+            return tt.tolist() if hasattr(tt, 'tolist') else list(tt)
+
+        text = item.get("answer", item.get("text", ""))
+        if isinstance(text, str) and len(text) > 0:
+            return self.tokenizer.encode(text, add_special_tokens=False)
+
+        return []
+
+    def _get_snac_tokens(self, item: Dict[str, Any]) -> List[int]:
+        """Extract SNAC tokens, ensuring multiple of 7."""
+        snac = item["snac_tokens"]
+        snac_len = (len(snac) // 7) * 7
+        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
+        return snac.tolist() if hasattr(snac, 'tolist') else list(snac)
+
+
+# ============================================================
+# Collate Functions
+# ============================================================
+def collate_simple(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Simple collate that pads whisper features.
+    Interleaving is done in training loop for correct text_ratio.
+    """
+    max_w = max(b["whisper"].shape[0] for b in batch)
+    max_w = ((max_w + 4) // 5) * 5  # Align to downsample factor
+
+    whisper_batch = []
+    raw_data = []
+
+    for b in batch:
+        w = b["whisper"]
+        w_pad = F.pad(w, (0, 0, 0, max_w - w.shape[0]))
+        whisper_batch.append(w_pad)
+        raw_data.append({
+            "text_tokens": b["text_tokens"],
+            "snac_tokens": b["snac_tokens"],
+            "word_alignments": b.get("word_alignments"),
+            "answer_text": b.get("answer_text", "")
+        })
+
+    return {
+        "whisper": torch.stack(whisper_batch),
+        "raw_data": raw_data
+    }
+
+
+# ============================================================
+# DataLoader Factory
+# ============================================================
+class DataLoaderFactory:
+    """
+    Factory for creating DataLoaders with optimal settings.
+
+    Single Responsibility: Only handles DataLoader creation.
+    Open/Closed: Can extend multiprocessing strategies without modification.
+
+    Optimizations:
+    - Sequence length bucketing for reduced padding
+    - Optimal worker count based on system resources
+    """
+
+    @staticmethod
+    def get_optimal_workers() -> int:
+        """Calculate optimal num_workers based on system resources."""
+        if not torch.cuda.is_available():
+            return 0
+
+        try:
+            import os
+            num_gpus = torch.cuda.device_count()
+            cpu_cores = os.cpu_count() or 4
+            max_workers = max(1, cpu_cores // 2)
+
+            # 2 workers per GPU, capped by CPU
+            ideal_workers = num_gpus * 2
+            num_workers = min(ideal_workers, max_workers)
+
+            # Check VRAM pressure
+            try:
+                free_bytes, _ = torch.cuda.mem_get_info(0)
+                if free_bytes / 1024**3 < 5:
+                    num_workers = min(num_workers, 1)
+            except Exception:
+                pass
+
+            return max(0, num_workers)
+        except Exception:
+            return 2
+
+    @classmethod
+    def create(
+        cls,
+        dataset: Dataset,
+        batch_size: int,
+        shuffle: bool = True,
+        collate_fn: Callable = collate_simple,
+        verbose: bool = True,
+        use_bucketing: bool = True,
+    ) -> DataLoader:
+        """
+        Create DataLoader with optimal settings.
+
+        Args:
+            dataset: The dataset to load from
+            batch_size: Batch size
+            shuffle: Whether to shuffle data
+            collate_fn: Function to collate samples
+            verbose: Whether to log details
+            use_bucketing: Use sequence length bucketing (reduces padding overhead)
+        """
+        optimal_workers = cls.get_optimal_workers()
+
+        # Try to use sequence length bucketing
+        batch_sampler = None
+        if use_bucketing and shuffle:
+            try:
+                # Get sequence lengths from dataset
+                lengths = cls._get_sequence_lengths(dataset)
+                if lengths and len(lengths) > batch_size * 10:  # Only bucket if enough samples
+                    batch_sampler = BucketBatchSampler(
+                        lengths=lengths,
+                        batch_size=batch_size,
+                        shuffle=True,
+                        drop_last=False,
+                    )
+                    if verbose:
+                        log(f"[DataLoader] Using sequence length bucketing ({len(batch_sampler)} batches)")
+            except Exception as e:
+                if verbose:
+                    log(f"[DataLoader] Bucketing failed: {e}, using standard batching")
+
+        # Try different multiprocessing methods
+        for mp_method in ['spawn', 'fork', None]:
+            try:
+                loader = cls._try_create_loader(
+                    dataset, batch_size, shuffle, collate_fn,
+                    optimal_workers, mp_method, batch_sampler
+                )
+                if loader is not None:
+                    if verbose and mp_method:
+                        log(f"[DataLoader] Using '{mp_method}' with {optimal_workers} workers")
+                    elif verbose:
+                        log("[DataLoader] Using single-process mode")
+                    return loader
+            except Exception as e:
+                if verbose:
+                    log(f"[DataLoader] {mp_method} failed: {str(e)[:50]}...")
+                continue
+
+        # Final fallback
+        if verbose:
+            log("[DataLoader] Fallback to num_workers=0")
+        return DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=collate_fn,
+            num_workers=0,
+            pin_memory=True
+        )
+
+    @staticmethod
+    def _get_sequence_lengths(dataset: Dataset) -> Optional[List[int]]:
+        """Extract sequence lengths from dataset for bucketing."""
+        # Check if dataset has get_sequence_lengths method
+        if hasattr(dataset, 'get_sequence_lengths'):
+            return dataset.get_sequence_lengths()
+
+        # For ConcatDataset, try to combine lengths from components
+        if isinstance(dataset, ConcatDataset):
+            lengths = []
+            for ds in dataset.datasets:
+                if hasattr(ds, 'get_sequence_lengths'):
+                    lengths.extend(ds.get_sequence_lengths())
+                else:
+                    return None  # Can't get lengths for all, skip bucketing
+            return lengths
+
+        return None
+
+    @staticmethod
+    def _try_create_loader(
+        dataset: Dataset,
+        batch_size: int,
+        shuffle: bool,
+        collate_fn: Callable,
+        num_workers: int,
+        mp_method: Optional[str],
+        batch_sampler: Optional[Sampler] = None,
+    ) -> Optional[DataLoader]:
+        """Try to create DataLoader with given settings."""
+        import multiprocessing
+
+        # When using batch_sampler, don't set batch_size, shuffle, or sampler
+        common_kwargs = {
+            'collate_fn': collate_fn,
+            'pin_memory': True,
+        }
+
+        if batch_sampler is not None:
+            common_kwargs['batch_sampler'] = batch_sampler
+        else:
+            common_kwargs['batch_size'] = batch_size
+            common_kwargs['shuffle'] = shuffle
+
+        if mp_method and num_workers > 0:
+            mp_context = multiprocessing.get_context(mp_method)
+            loader = DataLoader(
+                dataset,
+                num_workers=num_workers,
+                multiprocessing_context=mp_context,
+                persistent_workers=True,
+                **common_kwargs
+            )
+            # Test if it works
+            test_iter = iter(loader)
+            del test_iter
+            return loader
+        else:
+            return DataLoader(
+                dataset,
+                num_workers=0,
+                **common_kwargs
+            )
+
+
+def create_dataloader(
+    dataset: Dataset,
+    batch_size: int,
+    shuffle: bool = True,
+    verbose: bool = True,
+    use_bucketing: bool = True,
+) -> DataLoader:
+    """
+    Convenience function for creating DataLoaders.
+
+    Args:
+        dataset: Dataset to load from
+        batch_size: Batch size
+        shuffle: Whether to shuffle
+        verbose: Whether to log
+        use_bucketing: Use sequence length bucketing
+    """
+    return DataLoaderFactory.create(
+        dataset, batch_size, shuffle,
+        collate_fn=collate_simple, verbose=verbose,
+        use_bucketing=use_bucketing,
+    )
+
+
+# ============================================================
+# Dataset Loading Pipeline
+# ============================================================
+def load_datasets(
+    paths: List[str],
+    tokenizer,
+    max_audio_len: int = 500,
+    max_seq_len: int = 2048,
+    verbose: bool = True,
+    lazy_loading: bool = True,  # Use memory-efficient lazy loading
+) -> Dataset:
+    """
+    Load and combine multiple datasets.
+
+    Args:
+        paths: List of dataset paths (files or directories)
+        tokenizer: Tokenizer for text encoding
+        max_audio_len: Maximum audio length
+        max_seq_len: Maximum sequence length
+        verbose: Whether to log progress
+        lazy_loading: Use memory-efficient lazy loading for batch directories
+
+    Returns:
+        Combined dataset
+    """
+    if verbose:
+        log("\nLoading datasets (lazy loading enabled)..." if lazy_loading else "\nLoading datasets...")
+
+    all_datasets = []
+
+    for path in paths:
+        path = Path(path)
+        try:
+            # Check if this is a batch directory (use lazy loading)
+            batches_dir = None
+            if path.name.endswith('.batches') and path.is_dir():
+                batches_dir = path
+            elif Path(f"{path}.batches").exists():
+                batches_dir = Path(f"{path}.batches")
+            elif path.is_dir():
+                batch_files = list(path.glob("batch_*.pt"))
+                if batch_files:
+                    batches_dir = path
+
+            if batches_dir and lazy_loading:
+                # Use memory-efficient lazy loading
+                batch_files = sorted(batches_dir.glob("batch_*.pt"))
+                if batch_files:
+                    dataset = LazyShardedDataset(
+                        batch_files,
+                        tokenizer,
+                        max_audio_len=max_audio_len,
+                        max_seq_len=max_seq_len,
+                        cache_size=5,  # Keep 5 batches in memory
+                        verbose=verbose
+                    )
+                    all_datasets.append(dataset)
+                    if verbose:
+                        log(f"  {path.name}: {len(dataset):,} samples (lazy loading)")
+            else:
+                # Fall back to full loading for single files
+                loader = ShardedDatasetLoader(verbose=verbose)
+                data = loader.load(str(path))
+                if data:
+                    dataset = InterleavedDataset(
+                        data, tokenizer,
+                        max_audio_len=max_audio_len,
+                        max_seq_len=max_seq_len
+                    )
+                    all_datasets.append(dataset)
+                    if verbose:
+                        log(f"  {path.name}: {len(data):,} samples")
+
+        except FileNotFoundError as e:
+            if verbose:
+                log(f"  [WARN] {e}")
+
+    if not all_datasets:
+        raise ValueError("No datasets loaded!")
+
+    return ConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0]

training/interleaving.py

@@ -0,0 +1,420 @@
+"""
+Interleaved sequence creation for IST-LM training.
+
+Single Responsibility: Only handles interleaved sequence generation.
+"""
+
+import random
+import torch
+import torch.nn.functional as F
+from typing import List, Dict, Tuple, Optional, Any
+from .config import SNAC_BASE_OFFSET, SNAC_LAYERS_PER_FRAME, SNAC_LAYER_OFFSET, EOS_TOKEN
+
+
+def apply_snac_offset(token_idx: int, position: int) -> int:
+    """
+    Apply position-based offset to SNAC token.
+
+    SNAC uses 7 tokens per frame with position-based offsets:
+    offset = SNAC_BASE_OFFSET + (position % 7) * 4096
+
+    Args:
+        token_idx: Raw SNAC token index
+        position: Position in the sequence
+
+    Returns:
+        Offset-adjusted token index
+    """
+    if int(token_idx) >= SNAC_BASE_OFFSET:
+        # Already has offset applied
+        return int(token_idx)
+    offset = SNAC_BASE_OFFSET + (position % SNAC_LAYERS_PER_FRAME) * SNAC_LAYER_OFFSET
+    return int(token_idx) + offset
+
+
+def get_text_ratio(
+    global_step: int,
+    decay_steps: int = 300,
+    initial_ratio: float = 0.9,
+    min_ratio: float = 0.0
+) -> float:
+    """
+    Calculate text ratio based on training step (IST-LM schedule).
+
+    Schedule: Start at 90% text, decrease by 10% every decay_steps.
+    - Step 0-299: 0.9
+    - Step 300-599: 0.8
+    - Step 600-899: 0.7
+    - ...
+    - Step 2700+: 0.0 (pure audio)
+
+    Args:
+        global_step: Current training step
+        decay_steps: Steps between each 10% decay
+        initial_ratio: Starting text ratio
+        min_ratio: Minimum text ratio
+
+    Returns:
+        Current text ratio
+    """
+    num_decays = global_step // decay_steps
+    text_ratio = initial_ratio - (num_decays * 0.1)
+    return max(min_ratio, text_ratio)
+
+
+def calculate_dynamic_decay_steps(
+    total_steps: int,
+    steps_per_epoch: int = None,
+    initial_ratio: float = 0.9,
+    final_audio_portion: float = 0.2
+) -> int:
+    """
+    Calculate decay_steps for scheduled interleaving.
+
+    Two modes:
+    1. If steps_per_epoch provided: Complete decay in first epoch only,
+       remaining epochs are pure audio.
+    2. Otherwise: Use final_audio_portion to spread decay across training.
+
+    Args:
+        total_steps: Total training steps
+        steps_per_epoch: Steps per epoch (if provided, decay completes in epoch 1)
+        initial_ratio: Starting text ratio (default 0.9)
+        final_audio_portion: Portion of training with p=0 (only used if steps_per_epoch=None)
+
+    Returns:
+        Calculated decay_steps
+    """
+    # Number of decay stages: 0.9 -> 0.0 = 9 steps (not 10)
+    num_decay_stages = int(initial_ratio / 0.1)
+
+    if steps_per_epoch is not None:
+        # Complete decay in first epoch - remaining epochs are pure audio
+        return max(1, steps_per_epoch // num_decay_stages)
+    else:
+        # Original behavior: spread across training
+        steps_until_pure_audio = int(total_steps * (1 - final_audio_portion))
+        return max(1, steps_until_pure_audio // num_decay_stages)
+
+
+class InterleavingStrategy:
+    """
+    Base class for interleaving strategies.
+
+    Open/Closed: Can create new strategies without modifying this class.
+    """
+
+    def create_sequence(
+        self,
+        text_tokens: List[int],
+        snac_tokens: List[int],
+        text_ratio: float,
+        **kwargs
+    ) -> Tuple[List[int], List[bool]]:
+        """
+        Create interleaved sequence.
+
+        Args:
+            text_tokens: Text token IDs
+            snac_tokens: SNAC audio token IDs
+            text_ratio: Ratio of text vs audio (0.0 = pure audio)
+
+        Returns:
+            Tuple of (interleaved_tokens, is_audio_mask)
+        """
+        raise NotImplementedError
+
+
+class PositionalInterleaving(InterleavingStrategy):
+    """
+    Positional interleaving strategy (fallback when no alignments).
+
+    Interleaves text and audio tokens based on fixed patterns
+    determined by text_ratio.
+    """
+
+    # Pattern lookup: text_ratio -> (text_per_chunk, frames_per_chunk)
+    PATTERNS = {
+        0.9: (1, 3),   # 1 text token + 3 audio frames
+        0.7: (1, 5),   # 1 text token + 5 audio frames
+        0.5: (1, 7),   # 1 text token + 7 audio frames
+        0.3: (1, 10),  # 1 text token + 10 audio frames
+        0.0: (0, 1),   # Pure audio
+    }
+
+    def create_sequence(
+        self,
+        text_tokens: List[int],
+        snac_tokens: List[int],
+        text_ratio: float,
+        **kwargs
+    ) -> Tuple[List[int], List[bool]]:
+        interleaved = []
+        is_audio_mask = []
+
+        if len(snac_tokens) == 0:
+            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+        # Group SNAC into frames of 7
+        frames = self._group_into_frames(snac_tokens)
+        if len(frames) == 0:
+            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+        # Get interleaving pattern
+        text_per_chunk, frames_per_chunk = self._get_pattern(text_ratio)
+
+        total_text = len(text_tokens)
+        total_frames = len(frames)
+
+        text_idx = 0
+        frame_idx = 0
+        snac_position = 0
+
+        while frame_idx < total_frames:
+            # Add text tokens
+            if text_per_chunk > 0 and text_idx < total_text:
+                for _ in range(text_per_chunk):
+                    if text_idx < total_text:
+                        interleaved.append(text_tokens[text_idx])
+                        is_audio_mask.append(False)
+                        text_idx += 1
+
+            # Add audio frames
+            for _ in range(frames_per_chunk):
+                if frame_idx < total_frames:
+                    frame = frames[frame_idx]
+                    for tok in frame:
+                        interleaved.append(apply_snac_offset(tok, snac_position))
+                        is_audio_mask.append(True)
+                        snac_position += 1
+                    frame_idx += 1
+
+        # Add remaining text (only if not pure audio mode)
+        if text_per_chunk > 0:
+            while text_idx < total_text:
+                interleaved.append(text_tokens[text_idx])
+                is_audio_mask.append(False)
+                text_idx += 1
+
+        # Add EOS
+        interleaved.append(EOS_TOKEN)
+        is_audio_mask.append(False)
+
+        return interleaved, is_audio_mask
+
+    def _group_into_frames(self, snac_tokens: List[int]) -> List[List[int]]:
+        """Group SNAC tokens into frames of 7."""
+        frames = []
+        for i in range(0, len(snac_tokens), SNAC_LAYERS_PER_FRAME):
+            frame = snac_tokens[i:i + SNAC_LAYERS_PER_FRAME]
+            if len(frame) == SNAC_LAYERS_PER_FRAME:
+                frames.append(frame)
+        return frames
+
+    def _get_pattern(self, text_ratio: float) -> Tuple[int, int]:
+        """Get interleaving pattern for given text_ratio."""
+        if text_ratio >= 0.9:
+            return self.PATTERNS[0.9]
+        elif text_ratio >= 0.7:
+            return self.PATTERNS[0.7]
+        elif text_ratio >= 0.5:
+            return self.PATTERNS[0.5]
+        elif text_ratio >= 0.3:
+            return self.PATTERNS[0.3]
+        else:
+            return self.PATTERNS[0.0]
+
+
+class AlignedInterleaving(InterleavingStrategy):
+    """
+    Word-aligned interleaving strategy.
+
+    Uses word alignments to semantically replace audio spans
+    with corresponding text tokens.
+    """
+
+    def create_sequence(
+        self,
+        text_tokens: List[int],
+        snac_tokens: List[int],
+        text_ratio: float,
+        word_alignments: Optional[List[Dict]] = None,
+        tokenizer=None,
+        answer_text: str = "",
+        **kwargs
+    ) -> Tuple[List[int], List[bool]]:
+        # Fall back to positional if no alignments
+        if not word_alignments or text_ratio <= 0:
+            return PositionalInterleaving().create_sequence(
+                text_tokens, snac_tokens, text_ratio
+            )
+
+        # Check for pre-computed tokens
+        has_precomputed = (
+            len(word_alignments) > 0 and
+            'tokens' in word_alignments[0] and
+            word_alignments[0]['tokens']
+        )
+
+        if not has_precomputed and not tokenizer:
+            return PositionalInterleaving().create_sequence(
+                text_tokens, snac_tokens, text_ratio
+            )
+
+        interleaved = []
+        is_audio_mask = []
+
+        # Group SNAC into frames
+        frames = []
+        for i in range(0, len(snac_tokens), SNAC_LAYERS_PER_FRAME):
+            frame = snac_tokens[i:i + SNAC_LAYERS_PER_FRAME]
+            if len(frame) == SNAC_LAYERS_PER_FRAME:
+                frames.append(frame)
+
+        if len(frames) == 0:
+            return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+        total_frames = len(frames)
+
+        # Randomly select words to replace with text
+        num_words = len(word_alignments)
+        num_text_words = int(num_words * text_ratio)
+        word_indices = list(range(num_words))
+        random.shuffle(word_indices)
+        text_word_indices = set(word_indices[:num_text_words])
+
+        # Frame rate conversion
+        max_alignment_frame = max(wa['end_frame'] for wa in word_alignments)
+        frame_ratio = total_frames / max_alignment_frame if max_alignment_frame > total_frames * 2 else 1.0
+
+        snac_position = 0
+
+        for word_idx, alignment in enumerate(word_alignments):
+            word = alignment['word']
+            start_frame = int(alignment['start_frame'] * frame_ratio)
+            end_frame = min(int(alignment['end_frame'] * frame_ratio), total_frames)
+
+            if word_idx in text_word_indices:
+                # Replace audio with text
+                word_tokens = alignment.get('tokens', [])
+                if not word_tokens and tokenizer:
+                    word_tokens = tokenizer.encode(word, add_special_tokens=False)
+
+                for tok in word_tokens:
+                    interleaved.append(tok)
+                    is_audio_mask.append(False)
+                snac_position = end_frame * SNAC_LAYERS_PER_FRAME
+            else:
+                # Keep audio
+                for f_idx in range(start_frame, end_frame):
+                    if f_idx < total_frames:
+                        frame = frames[f_idx]
+                        for tok in frame:
+                            interleaved.append(apply_snac_offset(tok, snac_position))
+                            is_audio_mask.append(True)
+                            snac_position += 1
+
+        # Add remaining frames
+        remaining_start = max(wa['end_frame'] for wa in word_alignments)
+        remaining_start = min(int(remaining_start * frame_ratio), total_frames)
+        for f_idx in range(remaining_start, total_frames):
+            frame = frames[f_idx]
+            for tok in frame:
+                interleaved.append(apply_snac_offset(tok, snac_position))
+                is_audio_mask.append(True)
+                snac_position += 1
+
+        # Add EOS
+        interleaved.append(EOS_TOKEN)
+        is_audio_mask.append(False)
+
+        return interleaved, is_audio_mask
+
+
+def create_interleaved_sequence(
+    text_tokens: List[int],
+    snac_tokens: List[int],
+    text_ratio: float = 0.9,
+    word_alignments: Optional[List[Dict]] = None,
+    tokenizer=None,
+    answer_text: str = ""
+) -> Tuple[List[int], List[bool]]:
+    """
+    Create interleaved sequence (convenience function).
+
+    Automatically selects the best strategy based on available data.
+    """
+    if word_alignments and text_ratio > 0:
+        strategy = AlignedInterleaving()
+    else:
+        strategy = PositionalInterleaving()
+
+    return strategy.create_sequence(
+        text_tokens=text_tokens,
+        snac_tokens=snac_tokens,
+        text_ratio=text_ratio,
+        word_alignments=word_alignments,
+        tokenizer=tokenizer,
+        answer_text=answer_text,
+    )
+
+
+def apply_interleaving(
+    batch: Dict[str, Any],
+    text_ratio: float,
+    tokenizer=None,
+    max_seq_len: int = 2048
+) -> Dict[str, torch.Tensor]:
+    """
+    Apply interleaving to a batch of samples.
+
+    Args:
+        batch: Batch from DataLoader with 'whisper' and 'raw_data'
+        text_ratio: Current text ratio
+        tokenizer: Tokenizer for on-the-fly encoding
+        max_seq_len: Maximum sequence length
+
+    Returns:
+        Batch with 'whisper', 'interleaved', 'is_audio_mask'
+    """
+    raw_data = batch["raw_data"]
+    sequences = []
+    max_seq = 0
+
+    for item in raw_data:
+        interleaved, is_audio = create_interleaved_sequence(
+            item["text_tokens"],
+            item["snac_tokens"],
+            text_ratio,
+            word_alignments=item.get("word_alignments"),
+            tokenizer=tokenizer,
+            answer_text=item.get("answer_text", "")
+        )
+
+        # Truncate if needed
+        if len(interleaved) > max_seq_len:
+            interleaved = interleaved[:max_seq_len]
+            is_audio = is_audio[:max_seq_len]
+
+        sequences.append((interleaved, is_audio))
+        max_seq = max(max_seq, len(interleaved))
+
+    # Pad and stack
+    interleaved_batch = []
+    is_audio_batch = []
+
+    for interleaved, is_audio in sequences:
+        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
+        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)
+
+        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
+        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)
+
+        interleaved_batch.append(seq_pad)
+        is_audio_batch.append(mask_pad)
+
+    return {
+        "whisper": batch["whisper"],
+        "interleaved": torch.stack(interleaved_batch),
+        "is_audio_mask": torch.stack(is_audio_batch)
+    }

training/models.py

@@ -0,0 +1,247 @@
+"""
+Model components for Speech-to-Speech training.
+
+Single Responsibility: Only defines model architectures.
+Open/Closed: Can extend with new adapters without modifying existing code.
+
+Optimizations:
+- Flash Attention 2 for memory-efficient attention (10-20x savings on long sequences)
+- BFloat16 for better numerical stability than FP16
+- Gradient checkpointing for memory savings
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional
+from .config import (
+    DEFAULT_WHISPER_DIM,
+    DEFAULT_LLM_DIM,
+    DEFAULT_DOWNSAMPLE,
+    DEFAULT_INTERMEDIATE_DIM,
+)
+from .utils import log
+
+
+class SpeechAdapter(nn.Module):
+    """
+    Speech adapter that maps Whisper features to LLM embedding space.
+
+    Architecture: 5× downsampling + FFN with intermediate dim
+
+    Based on LLaMA-Omni 2 design:
+    - Concatenates 5 consecutive Whisper frames
+    - Projects through 2-layer FFN
+    - Applies LayerNorm for stability
+
+    Args:
+        whisper_dim: Dimension of Whisper features (default: 1280)
+        llm_dim: Dimension of LLM embeddings (default: 3072)
+        downsample: Downsampling factor (default: 5)
+        intermediate_dim: Hidden dimension of FFN (default: 2048)
+    """
+
+    def __init__(
+        self,
+        whisper_dim: int = DEFAULT_WHISPER_DIM,
+        llm_dim: int = DEFAULT_LLM_DIM,
+        downsample: int = DEFAULT_DOWNSAMPLE,
+        intermediate_dim: int = DEFAULT_INTERMEDIATE_DIM
+    ):
+        super().__init__()
+        self.whisper_dim = whisper_dim
+        self.llm_dim = llm_dim
+        self.downsample = downsample
+        self.intermediate_dim = intermediate_dim
+
+        concat_dim = whisper_dim * downsample
+
+        self.ffn = nn.Sequential(
+            nn.Linear(concat_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, llm_dim),
+            nn.LayerNorm(llm_dim)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Whisper features [B, T, D]
+
+        Returns:
+            LLM embeddings [B, T // downsample, llm_dim]
+        """
+        B, T, D = x.shape
+
+        # Ensure T is divisible by downsample
+        T_new = (T // self.downsample) * self.downsample
+        x = x[:, :T_new]
+
+        # Reshape: [B, T, D] -> [B, T // downsample, D * downsample]
+        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
+
+        return self.ffn(x)
+
+    def get_num_params(self) -> int:
+        """Return total number of parameters."""
+        return sum(p.numel() for p in self.parameters())
+
+    def get_config(self) -> dict:
+        """Return configuration dict for serialization."""
+        return {
+            "whisper_dim": self.whisper_dim,
+            "llm_dim": self.llm_dim,
+            "downsample": self.downsample,
+            "intermediate_dim": self.intermediate_dim,
+        }
+
+    @classmethod
+    def from_config(cls, config: dict) -> 'SpeechAdapter':
+        """Create adapter from configuration dict."""
+        return cls(**config)
+
+
+class ModelFactory:
+    """
+    Factory for creating models with consistent settings.
+
+    Single Responsibility: Only handles model instantiation.
+    Dependency Inversion: Depends on abstractions (config), not concretions.
+    """
+
+    @staticmethod
+    def create_adapter(
+        whisper_dim: int = DEFAULT_WHISPER_DIM,
+        llm_dim: int = DEFAULT_LLM_DIM,
+        dtype: torch.dtype = torch.float32,
+        checkpoint_path: Optional[str] = None
+    ) -> SpeechAdapter:
+        """
+        Create a SpeechAdapter, optionally loading from checkpoint.
+
+        Args:
+            whisper_dim: Whisper feature dimension
+            llm_dim: LLM embedding dimension
+            dtype: Tensor dtype
+            checkpoint_path: Optional path to checkpoint
+
+        Returns:
+            Initialized SpeechAdapter
+        """
+        adapter = SpeechAdapter(
+            whisper_dim=whisper_dim,
+            llm_dim=llm_dim,
+        ).to(dtype=dtype)
+
+        if checkpoint_path:
+            ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+            if "adapter" in ckpt:
+                adapter.load_state_dict(ckpt["adapter"])
+            elif "state_dict" in ckpt:
+                adapter.load_state_dict(ckpt["state_dict"])
+            else:
+                adapter.load_state_dict(ckpt)
+
+        return adapter
+
+    @staticmethod
+    def create_llm(
+        model_path: str,
+        dtype: torch.dtype = torch.bfloat16,
+        freeze: bool = True,
+        gradient_checkpointing: bool = False,
+        use_flash_attention: bool = True,
+        verbose: bool = True,
+    ):
+        """
+        Create and configure the LLM with memory optimizations.
+
+        Args:
+            model_path: HuggingFace model path
+            dtype: Tensor dtype (BF16 recommended for stability)
+            freeze: Whether to freeze all parameters
+            gradient_checkpointing: Enable gradient checkpointing
+            use_flash_attention: Try to use Flash Attention 2 (10-20x memory savings)
+            verbose: Log configuration details
+
+        Returns:
+            Configured LLM model
+        """
+        from transformers import AutoModelForCausalLM
+
+        # Determine attention implementation
+        # Flash Attention 2 provides 10-20x memory savings on long sequences
+        # Requires: Ampere/Ada/Hopper GPU (RTX 30xx, 40xx, A100, H100)
+        attn_impl = "sdpa"  # Default fallback
+
+        if use_flash_attention and torch.cuda.is_available():
+            try:
+                # Actually try to import flash_attn to verify it works
+                import flash_attn
+                # Check GPU capability (Flash Attention 2 requires SM 80+)
+                major, _ = torch.cuda.get_device_capability()
+                if major >= 8:  # Ampere or newer (RTX 30xx, 40xx, A100, H100)
+                    attn_impl = "flash_attention_2"
+                    if verbose:
+                        log(f"[LLM] Using Flash Attention 2 v{flash_attn.__version__} (10-20x memory savings)")
+                else:
+                    if verbose:
+                        log(f"[LLM] GPU SM {major}.x too old for Flash Attention 2, using SDPA")
+            except ImportError:
+                if verbose:
+                    log("[LLM] flash_attn not installed, using SDPA")
+            except Exception as e:
+                if verbose:
+                    log(f"[LLM] Flash Attention check failed: {e}, using SDPA")
+
+        # Load model with optimizations
+        try:
+            llm = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                torch_dtype=dtype,
+                attn_implementation=attn_impl,
+            )
+        except Exception as e:
+            # Fallback if flash_attention_2 fails
+            if attn_impl == "flash_attention_2":
+                if verbose:
+                    log(f"[LLM] Flash Attention failed ({e}), falling back to SDPA")
+                llm = AutoModelForCausalLM.from_pretrained(
+                    model_path,
+                    torch_dtype=dtype,
+                    attn_implementation="sdpa",
+                )
+            else:
+                raise
+
+        if freeze:
+            for p in llm.parameters():
+                p.requires_grad = False
+            llm.eval()
+
+        if gradient_checkpointing:
+            llm.gradient_checkpointing_enable()
+            if verbose:
+                log("[LLM] Gradient checkpointing enabled")
+
+        return llm
+
+    @staticmethod
+    def apply_lora(
+        llm,
+        lora_config: 'LoRAConfig'
+    ):
+        """
+        Apply LoRA to an LLM.
+
+        Args:
+            llm: The LLM model
+            lora_config: LoRA configuration
+
+        Returns:
+            LLM with LoRA applied
+        """
+        from peft import get_peft_model
+        peft_config = lora_config.to_peft_config()
+        return get_peft_model(llm, peft_config)

training/trainer.py

@@ -0,0 +1,565 @@
+"""
+Training loop abstraction.
+
+Single Responsibility: Handles the training loop logic.
+Open/Closed: Base class can be extended for different training stages.
+Liskov Substitution: Stage1Trainer and Stage2Trainer are interchangeable where Trainer is expected.
+
+Optimizations implemented:
+- OOM handling with proper recovery (PyTorch FAQ best practices)
+- Gradient NaN/Inf detection and skipping
+- Gradient norm monitoring for stability
+- Dynamic sequence length based on text ratio
+- CUDA memory fragmentation reduction
+"""
+
+import os
+import gc
+import math
+import torch
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from tqdm import tqdm
+from typing import Optional, Dict, Any, Tuple
+from abc import ABC, abstractmethod
+
+# Reduce CUDA memory fragmentation (PyTorch recommendation)
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+from .config import TrainingConfig, GPUConfig
+from .data import load_datasets, create_dataloader
+from .models import SpeechAdapter, ModelFactory
+from .interleaving import (
+    get_text_ratio,
+    calculate_dynamic_decay_steps,
+    apply_interleaving,
+)
+from .checkpoint import CheckpointManager, TrainingState
+from .utils import (
+    log,
+    setup_cuda_optimizations,
+    setup_hf_login,
+    load_tokenizer,
+    should_enable_gradient_checkpointing,
+    write_step,
+    get_device_info,
+    limit_ram_usage,
+    get_ram_info,
+)
+
+
+class BaseTrainer(ABC):
+    """
+    Abstract base class for trainers.
+
+    Implements Template Method pattern: defines training skeleton,
+    subclasses implement specific steps.
+    """
+
+    def __init__(self, config: TrainingConfig):
+        self.config = config
+        self.accelerator: Optional[Accelerator] = None
+        self.tokenizer = None
+        self.adapter: Optional[SpeechAdapter] = None
+        self.llm = None
+        self.optimizer = None
+        self.scheduler = None
+        self.train_loader = None
+        self.checkpoint_manager: Optional[CheckpointManager] = None
+
+        # Training state
+        self.global_step = 0
+        self.start_epoch = 0
+        self.best_loss = float("inf")
+        self.current_text_ratio = config.initial_text_ratio
+
+        # Stability monitoring
+        self.nan_count = 0
+        self.oom_count = 0
+        self.last_grad_norm = 0.0
+        self.max_grad_norm_seen = 0.0
+
+    @property
+    def is_main(self) -> bool:
+        """Check if this is the main process."""
+        return self.accelerator is None or self.accelerator.is_main_process
+
+    @property
+    def device(self):
+        """Get current device."""
+        return self.accelerator.device if self.accelerator else torch.device("cpu")
+
+    def setup(self):
+        """Setup training environment."""
+        self._setup_accelerator()
+        self._setup_memory()
+        self._setup_tokenizer()
+        self._setup_data()
+        self._setup_models()
+        self._setup_optimizer()
+        self._setup_checkpoint_manager()
+        self._resume_if_needed()
+        self._prepare_for_training()
+
+    def _setup_accelerator(self):
+        """Initialize Accelerator."""
+        mixed_precision = "bf16" if torch.cuda.is_available() else None
+
+        self.accelerator = Accelerator(
+            gradient_accumulation_steps=self.config.grad_accum,
+            mixed_precision=mixed_precision,
+        )
+
+        set_seed(42)
+
+        if self.is_main:
+            self._log_setup_info()
+
+    def _log_setup_info(self):
+        """Log setup information."""
+        device_info = get_device_info()
+        gpu_config = GPUConfig.auto_detect()
+
+        log("=" * 60)
+        log(self._get_stage_name())
+        log("=" * 60)
+        log(f"Device: {device_info}")
+        log(f"GPU: {gpu_config.name} ({gpu_config.vram_gb}GB)")
+        log(f"Num processes: {self.accelerator.num_processes}")
+        log(f"Batch: {self.config.batch_size}, Grad accum: {self.config.grad_accum}")
+        log(f"LR: {self.config.learning_rate}, Epochs: {self.config.epochs}")
+        if getattr(self.config, 'no_decay', False):
+            log(f"Fixed text ratio: {self.config.initial_text_ratio} (no decay)")
+        else:
+            log(f"Initial text ratio: {self.config.initial_text_ratio}")
+        log("[Optimizations] OOM recovery, NaN detection, grad monitoring, expandable_segments")
+
+    def _setup_memory(self):
+        """Configure memory settings."""
+        if self.device.type == 'cuda':
+            setup_cuda_optimizations(self.config.vram_fraction)
+
+        ram_total, _ = get_ram_info()
+        ram_limit = self.config.ram_limit_gb or (ram_total * 0.80)
+        limit_ram_usage(ram_limit)
+
+        if self.is_main:
+            log(f"[Memory] VRAM: {self.config.vram_fraction*100:.0f}%, RAM: {ram_limit:.1f}GB")
+
+    def _setup_tokenizer(self):
+        """Load tokenizer."""
+        setup_hf_login()
+        self.tokenizer = load_tokenizer(self.config.model_path)
+
+    def _setup_data(self):
+        """Load datasets and create dataloader."""
+        if self.is_main:
+            log("\nLoading datasets...")
+
+        dataset = load_datasets(
+            self.config.data_paths,
+            self.tokenizer,
+            max_audio_len=self.config.max_audio_len,
+            max_seq_len=self.config.max_seq_len,
+            verbose=self.is_main,
+        )
+
+        # Apply demo/test mode
+        if self.config.test_mode:
+            dataset = torch.utils.data.Subset(dataset, range(min(5, len(dataset))))
+            self.config.batch_size = min(self.config.batch_size, len(dataset))
+            self.config.grad_accum = 1
+        elif self.config.demo_mode:
+            dataset = torch.utils.data.Subset(dataset, range(min(1000, len(dataset))))
+            self.config.batch_size = min(4, self.config.batch_size)
+
+        if self.is_main:
+            log(f"Total samples: {len(dataset):,}")
+
+        self.train_loader = create_dataloader(
+            dataset,
+            self.config.batch_size,
+            shuffle=True,
+            verbose=self.is_main,
+        )
+
+    @abstractmethod
+    def _setup_models(self):
+        """Setup models (implemented by subclasses)."""
+        pass
+
+    @abstractmethod
+    def _setup_optimizer(self):
+        """Setup optimizer (implemented by subclasses)."""
+        pass
+
+    @abstractmethod
+    def _setup_checkpoint_manager(self):
+        """Setup checkpoint manager (implemented by subclasses)."""
+        pass
+
+    def _resume_if_needed(self):
+        """Resume from checkpoint if specified."""
+        if not self.config.resume_from or not os.path.exists(self.config.resume_from):
+            return
+
+        if self.is_main:
+            log(f"\nResuming from: {self.config.resume_from}")
+
+        ckpt = self.checkpoint_manager.load(self.config.resume_from)
+        self._load_checkpoint(ckpt)
+
+    @abstractmethod
+    def _load_checkpoint(self, ckpt: Dict[str, Any]):
+        """Load checkpoint state (implemented by subclasses)."""
+        pass
+
+    def _prepare_for_training(self):
+        """Prepare models and optimizer for training."""
+        # Prepare with accelerator
+        prepared = self.accelerator.prepare(
+            self.adapter, self.llm, self.optimizer, self.train_loader
+        )
+        self.adapter, self.llm, self.optimizer, self.train_loader = prepared
+
+        # Calculate steps
+        self.steps_per_epoch = max(1, len(self.train_loader) // self.config.grad_accum)
+        self.total_steps = max(1, self.steps_per_epoch * self.config.epochs)
+        self.warmup_steps = int(self.total_steps * self.config.warmup_ratio)
+
+        # Calculate decay steps
+        if self.config.dynamic_decay:
+            self.effective_decay_steps = calculate_dynamic_decay_steps(
+                self.total_steps,
+                steps_per_epoch=self.steps_per_epoch,  # Complete decay in epoch 1
+                initial_ratio=self.config.initial_text_ratio,
+            )
+            if self.is_main:
+                log(f"[Dynamic Decay] decay_steps={self.effective_decay_steps} (decay completes in epoch 1)")
+        else:
+            self.effective_decay_steps = self.config.decay_steps
+
+        # Setup scheduler
+        self.scheduler = CosineAnnealingLR(
+            self.optimizer,
+            T_max=max(1, self.total_steps - self.warmup_steps),
+            eta_min=1e-6,
+        )
+
+        if self.is_main:
+            log(f"Steps: {self.steps_per_epoch}/epoch, {self.total_steps} total, {self.warmup_steps} warmup")
+
+    @abstractmethod
+    def _get_stage_name(self) -> str:
+        """Get stage name for logging."""
+        pass
+
+    def train(self):
+        """Main training loop."""
+        if self.is_main:
+            log("\n" + "=" * 60)
+            log(f"STARTING {self._get_stage_name()}")
+            log("=" * 60)
+
+        for epoch in range(self.start_epoch, self.config.epochs):
+            self._train_epoch(epoch)
+
+        self._finish_training()
+
+    def _train_epoch(self, epoch: int):
+        """Train one epoch."""
+        self.adapter.train()
+        epoch_loss = 0
+        accum_loss = 0
+
+        pbar = tqdm(
+            self.train_loader,
+            desc=f"Epoch {epoch+1}/{self.config.epochs}",
+            disable=not self.is_main,
+        )
+
+        for batch_idx, raw_batch in enumerate(pbar):
+            loss = self._train_step(raw_batch, batch_idx)
+            accum_loss += loss
+
+            if self.accelerator.sync_gradients:
+                self._update_after_step(accum_loss, pbar)
+                epoch_loss += accum_loss
+                accum_loss = 0
+
+        self._finish_epoch(epoch, epoch_loss)
+
+    def _check_numerical_stability(
+        self,
+        loss: torch.Tensor,
+        params,
+    ) -> Tuple[bool, str]:
+        """
+        Check for NaN/Inf in loss and gradients.
+
+        Returns:
+            (is_stable, reason): Tuple of stability flag and reason if unstable
+        """
+        # Check loss
+        if torch.isnan(loss) or torch.isinf(loss):
+            return False, f"loss={loss.item()}"
+
+        # Check gradients
+        for name, param in params:
+            if param.grad is not None:
+                if torch.isnan(param.grad).any():
+                    return False, f"NaN gradient in {name}"
+                if torch.isinf(param.grad).any():
+                    return False, f"Inf gradient in {name}"
+
+        return True, ""
+
+    def _compute_grad_norm(self, params) -> float:
+        """Compute total gradient norm for monitoring."""
+        total_norm = 0.0
+        for _, param in params:
+            if param.grad is not None:
+                param_norm = param.grad.data.norm(2)
+                total_norm += param_norm.item() ** 2
+        return math.sqrt(total_norm)
+
+    def _train_step(self, raw_batch: Dict[str, Any], batch_idx: int) -> float:
+        """Single training step with OOM and NaN/Inf handling.
+
+        Stability features:
+        - OOM recovery outside except clause (PyTorch FAQ best practice)
+        - NaN/Inf detection in loss and gradients
+        - Gradient norm monitoring
+        - gc.collect() before empty_cache() for better cleanup
+
+        See: https://pytorch.org/docs/stable/notes/faq.html
+        """
+        # Update text ratio (skip if no_decay mode for Stage 1)
+        if not getattr(self.config, 'no_decay', False):
+            self.current_text_ratio = get_text_ratio(
+                self.global_step,
+                self.effective_decay_steps,
+                self.config.initial_text_ratio,
+            )
+        # else: keep current_text_ratio fixed at initial value
+
+        # Dynamic max_seq_len based on text ratio
+        dynamic_max_seq = self._get_dynamic_max_seq()
+
+        # Apply interleaving
+        batch = apply_interleaving(
+            raw_batch,
+            self.current_text_ratio,
+            tokenizer=self.tokenizer,
+            max_seq_len=dynamic_max_seq,
+        )
+
+        # Get adapter dtype for proper casting
+        adapter_dtype = next(self.adapter.parameters()).dtype
+        whisper = batch["whisper"].to(self.device, dtype=adapter_dtype)
+        interleaved = batch["interleaved"].to(self.device)
+
+        # Clear cache periodically
+        if batch_idx % 100 == 0 and self.device.type == 'cuda':
+            torch.cuda.empty_cache()
+
+        # Forward and backward with OOM and NaN handling
+        # Use flag pattern to move recovery outside except (PyTorch FAQ recommendation)
+        oom_error = False
+        nan_error = False
+        error_reason = ""
+        seq_len_for_log = interleaved.shape[1]
+        loss_value = 0.0
+
+        try:
+            with self.accelerator.accumulate(self.adapter):
+                loss = self._compute_loss(whisper, interleaved)
+
+                # Check for NaN/Inf in loss before backward
+                if torch.isnan(loss) or torch.isinf(loss):
+                    nan_error = True
+                    error_reason = f"loss={loss.item()}"
+                else:
+                    self.accelerator.backward(loss)
+
+                    if self.accelerator.sync_gradients:
+                        # Check gradient stability
+                        trainable_params = list(self._get_trainable_params_named())
+                        is_stable, reason = self._check_numerical_stability(loss, trainable_params)
+
+                        if not is_stable:
+                            nan_error = True
+                            error_reason = reason
+                        else:
+                            # Track gradient norm for monitoring
+                            self.last_grad_norm = self._compute_grad_norm(trainable_params)
+                            self.max_grad_norm_seen = max(self.max_grad_norm_seen, self.last_grad_norm)
+
+                            # Clip gradients
+                            self.accelerator.clip_grad_norm_(
+                                [p for _, p in trainable_params],
+                                self.config.max_grad_norm,
+                            )
+
+                    if not nan_error:
+                        self.optimizer.step()
+                        self.optimizer.zero_grad(set_to_none=True)
+                        loss_value = loss.item()
+
+        except torch.cuda.OutOfMemoryError:
+            oom_error = True
+
+        # Recovery outside except clause (prevents memory leak from exception stack frame)
+        if oom_error:
+            self.oom_count += 1
+            if self.is_main:
+                log(f"[OOM #{self.oom_count}] Skipping batch {batch_idx} (seq_len={seq_len_for_log})")
+            del whisper, interleaved, batch
+            gc.collect()
+            torch.cuda.empty_cache()
+            self.optimizer.zero_grad(set_to_none=True)
+            return 0.0
+
+        if nan_error:
+            self.nan_count += 1
+            if self.is_main:
+                log(f"[NaN #{self.nan_count}] Skipping batch {batch_idx}: {error_reason}")
+            self.optimizer.zero_grad(set_to_none=True)
+            return 0.0
+
+        return loss_value
+
+    def _get_trainable_params_named(self):
+        """Get named trainable parameters for gradient checking."""
+        # Default: adapter parameters
+        for name, param in self.adapter.named_parameters():
+            if param.requires_grad:
+                yield f"adapter.{name}", param
+
+    def _get_dynamic_max_seq(self) -> int:
+        """Get dynamic max sequence length based on text ratio.
+
+        More conservative limits to prevent CUDA OOM.
+        RTX 4090 (24GB) at 80% VRAM can handle ~1280-1536 max seq with LLM.
+        """
+        # Base limit more conservative for memory safety
+        base_limit = min(self.config.max_seq_len, 1536)
+
+        if self.current_text_ratio >= 0.7:
+            return base_limit
+        elif self.current_text_ratio >= 0.5:
+            return int(base_limit * 0.75)  # ~1152
+        elif self.current_text_ratio >= 0.3:
+            return int(base_limit * 0.6)   # ~922
+        else:
+            return int(base_limit * 0.5)   # ~768
+
+    def _compute_loss(
+        self,
+        whisper: torch.Tensor,
+        interleaved: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Compute training loss with numerical stability.
+
+        Numerical stability measures:
+        - Clamp logits to prevent extreme values
+        - Use BF16 which has better dynamic range than FP16
+        """
+        unwrapped_llm = self.accelerator.unwrap_model(self.llm)
+
+        # Forward through adapter
+        audio_embeds = self.adapter(whisper)
+
+        # Get token embeddings
+        input_tokens = interleaved[:, :-1].clamp(min=0)
+        with torch.no_grad():
+            token_embeds = unwrapped_llm.model.embed_tokens(input_tokens)
+
+        # Combine embeddings
+        combined = torch.cat([audio_embeds, token_embeds], dim=1)
+
+        # Forward through LLM
+        outputs = self.llm(inputs_embeds=combined, use_cache=False)
+        logits = outputs.logits
+
+        # Compute loss with numerical stability
+        audio_len = audio_embeds.shape[1]
+        seq_len = interleaved.shape[1]
+        seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]
+
+        # Clamp logits to prevent numerical issues in softmax
+        # Large logits can cause overflow in exp() during cross_entropy
+        seq_logits = seq_logits.clamp(min=-100, max=100)
+
+        return F.cross_entropy(
+            seq_logits.reshape(-1, logits.size(-1)),
+            interleaved.reshape(-1),
+            ignore_index=-100,
+            label_smoothing=self.config.label_smoothing,
+        )
+
+    @abstractmethod
+    def _get_trainable_params(self):
+        """Get trainable parameters for gradient clipping."""
+        pass
+
+    def _update_after_step(self, accum_loss: float, pbar):
+        """Update after gradient accumulation step."""
+        # Learning rate schedule
+        if self.global_step < self.warmup_steps:
+            lr_scale = (self.global_step + 1) / self.warmup_steps
+            for pg in self.optimizer.param_groups:
+                pg["lr"] = self.config.learning_rate * lr_scale
+        else:
+            self.scheduler.step()
+
+        self.global_step += 1
+        if self.is_main:
+            write_step(self.global_step)
+
+        # Update progress bar with gradient norm for stability monitoring
+        avg_loss = accum_loss / self.config.grad_accum
+        pbar.set_postfix(
+            loss=f"{avg_loss:.4f}",
+            lr=f"{self.optimizer.param_groups[0]['lr']:.2e}",
+            text_ratio=f"{self.current_text_ratio:.1f}",
+            grad=f"{self.last_grad_norm:.1f}",
+        )
+
+        # Log warning if gradient norm is very high (potential instability)
+        if self.last_grad_norm > 100 and self.global_step % 50 == 0 and self.is_main:
+            log(f"[WARN] High gradient norm: {self.last_grad_norm:.1f} at step {self.global_step}")
+
+        # Save checkpoint
+        if self.global_step % self.config.save_steps == 0 and self.is_main:
+            self._save_step_checkpoint(accum_loss)
+
+    @abstractmethod
+    def _save_step_checkpoint(self, loss: float):
+        """Save step checkpoint (implemented by subclasses)."""
+        pass
+
+    @abstractmethod
+    def _finish_epoch(self, epoch: int, epoch_loss: float):
+        """Finish epoch (implemented by subclasses)."""
+        pass
+
+    def _finish_training(self):
+        """Finish training."""
+        self.accelerator.wait_for_everyone()
+
+        if self.is_main:
+            self.checkpoint_manager.wait_for_saves()
+            log("\n" + "=" * 60)
+            log(f"{self._get_stage_name()} COMPLETE!")
+            log(f"Best loss: {self.best_loss:.4f}")
+            log(f"Final text_ratio: {self.current_text_ratio:.1f}")
+            log(f"Max gradient norm seen: {self.max_grad_norm_seen:.2f}")
+            if self.oom_count > 0:
+                log(f"OOM errors recovered: {self.oom_count}")
+            if self.nan_count > 0:
+                log(f"NaN/Inf errors recovered: {self.nan_count}")
+            log("=" * 60)

training/utils.py

@@ -0,0 +1,255 @@
+"""
+Utility functions for training.
+
+Single Responsibility: General utilities that don't fit elsewhere.
+"""
+
+import os
+import sys
+import torch
+from typing import Tuple, Optional
+from dataclasses import dataclass
+
+
+# ============================================================
+# Logging
+# ============================================================
+_verbose = True
+
+
+def setup_logging(verbose: bool = True):
+    """Configure logging verbosity."""
+    global _verbose
+    _verbose = verbose
+
+
+def log(msg: str):
+    """Log message to stdout with flush."""
+    if _verbose:
+        print(msg)
+        sys.stdout.flush()
+
+
+# ============================================================
+# Device Information
+# ============================================================
+@dataclass
+class DeviceInfo:
+    """Information about the compute device."""
+    device_type: str
+    device_name: str
+    vram_gb: float
+    ram_total_gb: float
+    ram_available_gb: float
+    num_gpus: int
+
+    def __str__(self) -> str:
+        parts = [f"Device: {self.device_type}"]
+        if self.device_name:
+            parts.append(f"({self.device_name})")
+        if self.vram_gb > 0:
+            parts.append(f"VRAM: {self.vram_gb:.0f}GB")
+        if self.num_gpus > 1:
+            parts.append(f"x{self.num_gpus} GPUs")
+        return " | ".join(parts)
+
+
+def get_device_info() -> DeviceInfo:
+    """Get information about the compute device."""
+    device_type = "cpu"
+    device_name = ""
+    vram_gb = 0.0
+    num_gpus = 0
+
+    # CUDA
+    if torch.cuda.is_available():
+        device_type = "cuda"
+        num_gpus = torch.cuda.device_count()
+        try:
+            props = torch.cuda.get_device_properties(0)
+            device_name = props.name
+            vram_gb = props.total_memory / (1024**3)
+        except Exception:
+            pass
+
+    # MPS
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        device_type = "mps"
+        device_name = "Apple Silicon"
+        num_gpus = 1
+
+    # RAM info
+    ram_total, ram_available = get_ram_info()
+
+    return DeviceInfo(
+        device_type=device_type,
+        device_name=device_name,
+        vram_gb=vram_gb,
+        ram_total_gb=ram_total,
+        ram_available_gb=ram_available,
+        num_gpus=num_gpus,
+    )
+
+
+def get_ram_info() -> Tuple[float, float]:
+    """Get RAM info in GB (total, available)."""
+    try:
+        import psutil
+        total = psutil.virtual_memory().total / 1024**3
+        available = psutil.virtual_memory().available / 1024**3
+        return total, available
+    except ImportError:
+        pass
+
+    try:
+        import subprocess
+        result = subprocess.run(
+            ['free', '-b'],
+            capture_output=True, text=True
+        )
+        lines = result.stdout.strip().split('\n')
+        if len(lines) >= 2:
+            parts = lines[1].split()
+            total = float(parts[1]) / 1024**3
+            available = float(parts[6]) / 1024**3 if len(parts) > 6 else float(parts[3]) / 1024**3
+            return total, available
+    except Exception:
+        pass
+
+    return 0.0, 0.0
+
+
+def log_memory_usage() -> str:
+    """Get current memory usage string."""
+    parts = []
+
+    if torch.cuda.is_available():
+        used = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+        parts.append(f"GPU: {used:.2f}GB / {reserved:.2f}GB")
+
+    try:
+        import psutil
+        ram_used = psutil.virtual_memory().used / 1024**3
+        ram_total = psutil.virtual_memory().total / 1024**3
+        parts.append(f"RAM: {ram_used:.1f}GB / {ram_total:.1f}GB")
+    except ImportError:
+        pass
+
+    return " | ".join(parts)
+
+
+# ============================================================
+# Memory Management
+# ============================================================
+def limit_ram_usage(max_ram_gb: float):
+    """Limit RAM usage via resource limits."""
+    try:
+        import resource
+        max_bytes = int(max_ram_gb * 1024**3)
+        resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
+    except Exception:
+        pass
+
+
+def setup_cuda_optimizations(vram_fraction: float = 0.80):
+    """Configure CUDA optimizations."""
+    if not torch.cuda.is_available():
+        return
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+    torch.set_float32_matmul_precision('high')
+
+    try:
+        torch.cuda.set_per_process_memory_fraction(vram_fraction)
+        torch.cuda.empty_cache()
+    except Exception:
+        pass
+
+
+def should_enable_gradient_checkpointing(
+    vram_gb: float,
+    dynamic_decay: bool = False,
+    threshold_fraction: float = 0.4
+) -> bool:
+    """
+    Determine if gradient checkpointing should be enabled.
+
+    Args:
+        vram_gb: Total VRAM in GB
+        dynamic_decay: Whether using dynamic decay (longer sequences over time)
+        threshold_fraction: Fraction of VRAM that should be free
+
+    Returns:
+        Whether to enable gradient checkpointing
+    """
+    if not torch.cuda.is_available():
+        return False
+
+    # With dynamic_decay, sequences get longer over time
+    if dynamic_decay and vram_gb <= 32:
+        return True
+
+    # Check available VRAM
+    try:
+        torch.cuda.empty_cache()
+        free_bytes, total_bytes = torch.cuda.mem_get_info(0)
+        free_gb = free_bytes / 1024**3
+        threshold_gb = vram_gb * threshold_fraction
+        return free_gb < threshold_gb
+    except Exception:
+        # Conservative: enable if VRAM < 20GB
+        return vram_gb < 20
+
+
+# ============================================================
+# Step Sharing (for DDP + DataLoader workers)
+# ============================================================
+STEP_FILE = "/tmp/training_step.txt"
+
+
+def write_step(step: int):
+    """Write current training step to file (main process only)."""
+    try:
+        with open(STEP_FILE, "w") as f:
+            f.write(str(step))
+    except Exception:
+        pass
+
+
+def read_step() -> int:
+    """Read current training step from file."""
+    try:
+        with open(STEP_FILE, "r") as f:
+            return int(f.read().strip())
+    except Exception:
+        return 0
+
+
+# ============================================================
+# HuggingFace Helpers
+# ============================================================
+def setup_hf_login():
+    """Setup HuggingFace login from environment."""
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        try:
+            from huggingface_hub import login
+            login(token=hf_token)
+            return True
+        except Exception:
+            pass
+    return False
+
+
+def load_tokenizer(model_path: str):
+    """Load tokenizer with proper padding token."""
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    return tokenizer