{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a Speech LLM with NeMo Automodel and MoE\n",
    "\n",
    "This tutorial walks through the full **SALMAutomodel** pipeline: data preparation,\n",
    "training, checkpoint conversion, and evaluation.\n",
    "\n",
    "## What is SALMAutomodel?\n",
    "\n",
    "SALMAutomodel is a Speech-Augmented Language Model that uses\n",
    "[NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel) as the LLM backend.\n",
    "It connects a pretrained ASR encoder (e.g., Canary) to a pretrained LLM via a\n",
    "learnable linear projection layer.\n",
    "\n",
    "## Why MoE?\n",
    "\n",
    "We use **NVIDIA Nemotron Nano V3** as the LLM backbone — a Mixture-of-Experts\n",
    "model with 30B total parameters but only 3B active per token. NeMo Automodel\n",
    "provides two key MoE optimizations:\n",
    "\n",
    "- **Grouped GEMM**: Fuses expert computations into a single batched matrix multiply.\n",
    "- **DeepEP**: Efficient all-to-all expert routing across GPUs.\n",
    "\n",
    "## EP and FSDP2: Same Axis\n",
    "\n",
    "A key point: **Expert Parallelism (EP) reuses the FSDP2 data-parallel axis\n",
    "(`dp_size`)**. Dense layers are sharded via FSDP2, while MoE expert layers\n",
    "use EP for all-to-all routing — both on the same set of GPUs. Setting\n",
    "`ep_size` does *not* add a separate dimension.\n",
    "\n",
    "## What We Cover\n",
    "\n",
    "1. Download Mini LibriSpeech with Lhotse\n",
    "2. Train SALMAutomodel (only `perception.proj` is trained; LLM and ASR encoder are frozen)\n",
    "3. Convert the distributed checkpoint to HuggingFace format\n",
    "4. Evaluate with distributed inference and compute WER\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "- The tutorial was tested on **2x RTX 6000 Pro (Blackwell) GPUs**; for a smaller setup, you might need to use a smaller LLM backbone\n",
    "- `pip install nemo-toolkit[speechlm2]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install nemo-toolkit[speechlm2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Download Mini LibriSpeech\n",
    "\n",
    "We use [Lhotse](https://github.com/lhotse-speech/lhotse) to download and prepare\n",
    "Mini LibriSpeech — a small subset of LibriSpeech with two splits:\n",
    "- `train-clean-5` (~5 hours)\n",
    "- `dev-clean-2` (~2 hours)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "from lhotse.recipes import download_librispeech, prepare_librispeech\n",
    "from lhotse import CutSet\n",
    "\n",
    "def find_nemo_root(start: Path) -> Path:\n",
    "    start = start.resolve()\n",
    "    for candidate in (start, *start.parents):\n",
    "        if (candidate / \"examples\" / \"speechlm2\").is_dir() and (candidate / \"tutorials\" / \"speechlm2\").is_dir():\n",
    "            return candidate\n",
    "    raise RuntimeError(\n",
    "        f\"Could not locate the NeMo source tree from {start}. Open this notebook from inside the repo checkout.\"\n",
    "    )\n",
    "\n",
    "NEMO_ROOT = find_nemo_root(Path.cwd())\n",
    "TUTORIAL_ROOT = NEMO_ROOT / \"tutorials\" / \"speechlm2\"\n",
    "DATA_ROOT = TUTORIAL_ROOT / \"data\"\n",
    "SPEECHLM2_EXAMPLES_ROOT = NEMO_ROOT / \"examples\" / \"speechlm2\"\n",
    "RESULTS_ROOT = TUTORIAL_ROOT / \"salm_tutorial_results\"\n",
    "HF_OUTPUT_DIR = TUTORIAL_ROOT / \"salm_tutorial_hf\"\n",
    "GENERATIONS_PATH = TUTORIAL_ROOT / \"tutorial_generations.jsonl\"\n",
    "\n",
    "existing_pythonpath = os.environ.get(\"PYTHONPATH\")\n",
    "REPO_PYTHONPATH = str(NEMO_ROOT)\n",
    "if existing_pythonpath:\n",
    "    REPO_PYTHONPATH = os.pathsep.join([REPO_PYTHONPATH, existing_pythonpath])\n",
    "SCRIPT_ENV = dict(os.environ, PYTHONPATH=REPO_PYTHONPATH)\n",
    "\n",
    "print(f\"Using NeMo root: {NEMO_ROOT}\")\n",
    "print(f\"Using SpeechLM2 examples: {SPEECHLM2_EXAMPLES_ROOT}\")\n",
    "\n",
    "CORPUS_DIR = download_librispeech(DATA_ROOT, dataset_parts=\"mini_librispeech\")\n",
    "manifests = prepare_librispeech(\n",
    "    CORPUS_DIR, dataset_parts=\"mini_librispeech\", output_dir=DATA_ROOT / \"manifests\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create Lhotse CutSets\n",
    "\n",
    "Convert the raw manifests into CutSets and save as `.jsonl` files. LibriSpeech has all uppercase transcripts which we normalize: the first letter is capital, and each sentence ends with a full stop, to match LLM training data more closely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for split, parts in manifests.items():\n",
    "    cuts = CutSet.from_manifests(\n",
    "        recordings=parts[\"recordings\"], supervisions=parts[\"supervisions\"]\n",
    "    ).with_recording_path_prefix(DATA_ROOT.absolute())\n",
    "    path = DATA_ROOT / f\"cuts_{split}.jsonl\"\n",
    "    cuts = cuts.transform_text(lambda txt: txt.capitalize() + \".\")\n",
    "    cuts.to_file(path)\n",
    "    print(\n",
    "        f\"{split}: {len(cuts)} cuts, \"\n",
    "        f\"total duration: {sum(c.duration for c in cuts):.1f}s\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Inspect the Data\n",
    "\n",
    "Each `Cut` holds a pointer to an audio recording, its duration, and one or more\n",
    "supervisions (transcripts with timing information)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuts = CutSet.from_file(DATA_ROOT / \"cuts_train-clean-5.jsonl\")\n",
    "cut = cuts[0]\n",
    "print(f\"ID: {cut.id}\")\n",
    "print(f\"Duration: {cut.duration:.2f}s\")\n",
    "print(f\"Text: {cut.supervisions[0].text}\")\n",
    "cut.play_audio()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Write Training Config\n",
    "\n",
    "We write a YAML config for SALMAutomodel with:\n",
    "\n",
    "- **Nemotron Nano V3** as the LLM backbone\n",
    "- **EP=2** on 2 GPUs (dense layers → FSDP2, MoE layers → EP, same axis)\n",
    "- Only `perception.proj` (Linear 1024→4096) is trainable; LLM, ASR encoder,\n",
    "  preprocessor, and modality adapter are frozen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_cuts = str((DATA_ROOT / \"cuts_train-clean-5.jsonl\").resolve())\n",
    "val_cuts = str((DATA_ROOT / \"cuts_dev-clean-2.jsonl\").resolve())\n",
    "\n",
    "results_dir = str(RESULTS_ROOT.resolve())\n",
    "\n",
    "config_yaml = f\"\"\"\\\n",
    "model:\n",
    "  pretrained_llm: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16\n",
    "  pretrained_asr: nvidia/canary-1b-flash\n",
    "  pretrained_weights: true\n",
    "  use_nemo_automodel: true\n",
    "  trust_remote_code: true  # needed for Nemotron-3-Nano\n",
    "  automodel_backend:\n",
    "    dispatcher: torch # set to \"torch\" if your GPUs don't have NVLINK/NVSHMEM\n",
    "\n",
    "  prompt_format: nemotron-nano-v3\n",
    "  audio_locator_tag: \"<|audio|>\"\n",
    "\n",
    "  freeze_params:\n",
    "    - \"^llm\\\\\\\\..+$\"\n",
    "    - \"^perception\\\\\\\\.preprocessor\\\\\\\\..+$\"\n",
    "    - \"^perception\\\\\\\\.encoder\\\\\\\\..+$\"\n",
    "  prevent_freeze_params: []\n",
    "\n",
    "  # Uncomment to enable LoRA on the LLM:\n",
    "  # lora:\n",
    "  #   dim: 128\n",
    "  #   alpha: 256\n",
    "  #   dropout: 0.01\n",
    "  #   target_modules: [\"q_proj\", \"v_proj\"]\n",
    "\n",
    "  perception:\n",
    "    target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule\n",
    "    output_dim: 4096\n",
    "    modality_adapter:\n",
    "      _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector\n",
    "      d_model: 1024\n",
    "    spec_augment:\n",
    "      _target_: nemo.collections.asr.modules.SpectrogramAugmentation\n",
    "      freq_masks: 2 # set to zero to disable it\n",
    "      time_masks: 10 # set to zero to disable it\n",
    "      freq_width: 27\n",
    "      time_width: 5  # 5 frames = 50ms\n",
    "\n",
    "  optimizer:\n",
    "    _target_: torch.optim.AdamW\n",
    "    lr: 3e-4\n",
    "    betas: [0.9, 0.98]\n",
    "    weight_decay: 1e-3\n",
    "    foreach: true\n",
    "\n",
    "  lr_scheduler:\n",
    "    _target_: nemo.core.optim.lr_scheduler.CosineAnnealing\n",
    "    warmup_steps: 50\n",
    "    min_lr: 1e-6\n",
    "    max_steps: ${{trainer.max_steps}}\n",
    "\n",
    "trainer:\n",
    "  devices: 2\n",
    "  accelerator: gpu\n",
    "  num_nodes: 1\n",
    "  precision: bf16-true\n",
    "  logger: false\n",
    "  enable_checkpointing: false\n",
    "  use_distributed_sampler: false\n",
    "  max_steps: 500\n",
    "  val_check_interval: 1.0\n",
    "  limit_val_batches: 5\n",
    "  log_every_n_steps: 10\n",
    "  num_sanity_val_steps: 0\n",
    "  gradient_clip_val: 1.0\n",
    "  accumulate_grad_batches: 1\n",
    "  strategy:\n",
    "    _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy\n",
    "    dp_size: null\n",
    "    dp_replicate_size: 1\n",
    "    tp_size: 1\n",
    "    pp_size: 1\n",
    "    cp_size: 1\n",
    "    ep_size: 2\n",
    "\n",
    "data:\n",
    "  train_ds:\n",
    "    sample_rate: 16000\n",
    "    prompt_format: ${{model.prompt_format}}\n",
    "    token_equivalent_duration: 0.08\n",
    "    input_cfg:\n",
    "      - type: lhotse_as_conversation\n",
    "        cuts_path: {train_cuts}\n",
    "        audio_locator_tag: ${{model.audio_locator_tag}}\n",
    "        tags:\n",
    "          context: \"\"  # a fixed prompt can be added here since, or example-specific prompt can be attached in the manfiest directly\n",
    "    seed: 42\n",
    "    shuffle: true\n",
    "    shard_seed: randomized\n",
    "    num_workers: 1\n",
    "    batch_size: 16\n",
    "\n",
    "  validation_ds:\n",
    "    prompt_format: ${{model.prompt_format}}\n",
    "    token_equivalent_duration: 0.08\n",
    "    datasets:\n",
    "      dev_clean_2:\n",
    "        input_cfg:\n",
    "          - type: lhotse_as_conversation\n",
    "            cuts_path: {val_cuts}\n",
    "            audio_locator_tag: ${{model.audio_locator_tag}}\n",
    "            tags:\n",
    "              context: \"\"  # a fixed prompt can be added here since, or example-specific prompt can be attached in the manfiest directly\n",
    "    sample_rate: 16000\n",
    "    batch_size: 1\n",
    "    seed: 42\n",
    "    shard_seed: randomized\n",
    "\n",
    "exp_manager:\n",
    "  exp_dir: null\n",
    "  explicit_log_dir: {results_dir}\n",
    "  name: salm\n",
    "  create_tensorboard_logger: false\n",
    "  create_checkpoint_callback: true\n",
    "  use_datetime_version: true\n",
    "  resume_if_exists: true\n",
    "  resume_ignore_no_checkpoint: true\n",
    "  create_wandb_logger: false\n",
    "  checkpoint_callback_params:\n",
    "    filename: \"{{step}}\"\n",
    "    monitor: val_loss\n",
    "    mode: min\n",
    "    every_n_epochs: 1\n",
    "    always_save_nemo: false\n",
    "    save_top_k: 1\n",
    "    save_nemo_on_train_end: false\n",
    "\"\"\"\n",
    "\n",
    "config_path = DATA_ROOT / \"salm_automodel_tutorial.yaml\"\n",
    "config_path.write_text(config_yaml)\n",
    "print(f\"Config written to {config_path.resolve()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Config Explained\n",
    "\n",
    "Key settings in the config above:\n",
    "\n",
    "| Setting | Meaning |\n",
    "|---|---|\n",
    "| `use_nemo_automodel: true` | Selects `SALMAutomodel` (NeMo Automodel backend) |\n",
    "| `ep_size: 2` | Expert Parallelism on the FSDP data-parallel axis. Dense layers are sharded via FSDP2, MoE expert layers use EP — both on the same 2 GPUs |\n",
    "| `perception.output_dim: 4096` | Nemotron Nano V3 has `hidden_size=4096` |\n",
    "| `freeze_params` | Freezes the LLM, ASR encoder, preprocessor, and modality adapter. Only `perception.proj` (Linear 1024→4096) is trainable |\n",
    "| `lora` (commented) | Can be enabled for parameter-efficient LLM fine-tuning using Automodel-native LoRA |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Launch Training\n",
    "\n",
    "We use `torchrun` with 2 GPUs. The training script (`salm_train.py`) reads\n",
    "`use_nemo_automodel: true` and instantiates `SALMAutomodel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "TRAIN_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"salm_train.py\").resolve())\n",
    "\n",
    "cmd = [\n",
    "    \"torchrun\", \"--nproc_per_node=2\",\n",
    "    TRAIN_SCRIPT,\n",
    "    f\"--config-path={str(DATA_ROOT.resolve())}\",\n",
    "    \"--config-name=salm_automodel_tutorial\",\n",
    "]\n",
    "print(\"Running:\", \" \".join(cmd))\n",
    "subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Locate Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_paths = sorted((RESULTS_ROOT / \"checkpoints\").glob(\"*\"))\n",
    "ckpt_paths = [str(path.resolve()) for path in ckpt_paths]\n",
    "print(\"Checkpoints found:\", ckpt_paths)\n",
    "CKPT_PATH = ckpt_paths[-1]\n",
    "CONFIG_PATH = str((RESULTS_ROOT / \"exp_config.yaml\").resolve())\n",
    "print(f\"Using checkpoint: {CKPT_PATH}\")\n",
    "print(f\"Using config: {CONFIG_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Convert Checkpoint to HuggingFace Format\n",
    "\n",
    "Training with EP=2 produces **distributed checkpoints** (a directory with\n",
    "per-rank shards). The `to_hf.py` script consolidates DTensors into regular\n",
    "tensors and saves them as `config.json` + `model.safetensors`.\n",
    "\n",
    "We launch with `torchrun --nproc_per_node=2` to match the training topology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "TO_HF_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"to_hf.py\").resolve())\n",
    "\n",
    "cmd = [\n",
    "    \"torchrun\", \"--nproc_per_node=2\",\n",
    "    TO_HF_SCRIPT,\n",
    "    \"class_path=nemo.collections.speechlm2.models.SALMAutomodel\",\n",
    "    f\"ckpt_path='{CKPT_PATH}'\",\n",
    "    f\"ckpt_config='{CONFIG_PATH}'\",\n",
    "    f\"output_dir='{HF_OUTPUT_DIR.resolve()}'\",\n",
    "]\n",
    "print(\"Running:\", \" \".join(cmd))\n",
    "subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Evaluate with WER\n",
    "\n",
    "Run distributed inference on the dev set and compute Word Error Rate.\n",
    "The `salm_eval.py` script prints per-batch and overall WER."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "DEV_CUTS = str((DATA_ROOT / \"cuts_dev-clean-2.jsonl\").resolve())\n",
    "EVAL_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"salm_eval.py\").resolve())\n",
    "\n",
    "cmd = [\n",
    "    \"torchrun\", \"--nproc_per_node=2\",\n",
    "    EVAL_SCRIPT,\n",
    "    f\"pretrained_name='{HF_OUTPUT_DIR.resolve()}'\",\n",
    "    f\"inputs={DEV_CUTS}\",\n",
    "    # f\"inputs={train_cuts}\",\n",
    "    \"batch_size=128\",\n",
    "    \"max_new_tokens=128\",\n",
    "    \"ep_size=2\",\n",
    "    f\"output_manifest='{GENERATIONS_PATH.resolve()}'\",\n",
    "    \"user_prompt=Transcribe the following:\",\n",
    "    \"enable_thinking=False\",\n",
    "]\n",
    "print(\"Running:\", \" \".join(cmd))\n",
    "subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Inspect Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(GENERATIONS_PATH) as f:\n",
    "    results = [json.loads(line) for line in f]\n",
    "\n",
    "for r in results[:5]:\n",
    "    print(f\"REF: {r['text']}\")\n",
    "    print(f\"HYP: {r['pred_text']}\")\n",
    "    print()\n",
    "\n",
    "print(f\"Total examples: {len(results)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial we:\n",
    "\n",
    "1. Downloaded Mini LibriSpeech with Lhotse\n",
    "2. Trained SALMAutomodel with Nemotron Nano V3 MoE backbone (EP=2 on the FSDP data-parallel axis)\n",
    "3. Converted the distributed checkpoint to HuggingFace format\n",
    "4. Evaluated with distributed inference and computed WER\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "- **Scale up**: more GPUs, larger datasets\n",
    "- **Enable LoRA**: uncomment the `lora:` block for parameter-efficient LLM fine-tuning\n",
    "- **Try other parallelism combos**: TP + EP, HSDP (`dp_replicate_size > 1`)\n",
    "- See `docs/source/speechlm2/` for full documentation"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}