diff --git a/examples/nemotron-asr-nemo/.python-version b/examples/nemotron-asr-nemo/.python-version new file mode 100644 index 0000000..92536a9 --- /dev/null +++ b/examples/nemotron-asr-nemo/.python-version @@ -0,0 +1 @@ +3.12.0 diff --git a/examples/nemotron-asr-nemo/README.md b/examples/nemotron-asr-nemo/README.md new file mode 100644 index 0000000..893ce5a --- /dev/null +++ b/examples/nemotron-asr-nemo/README.md @@ -0,0 +1,79 @@ +# Nemotron 3.5 ASR Fine-Tuning with NeMo + +This example fine-tunes NVIDIA's [Nemotron 3.5 ASR streaming](https://huggingface.co/nvidia/nemotron-3.5-asr-streaming-0.6b) speech recognition model using the NVIDIA NeMo framework on Baseten. + +Nemotron 3.5 ASR is a 600M-parameter, multilingual (40 language-locales), real-time streaming model built on a Cache-Aware FastConformer-RNNT architecture with prompt-based language conditioning. This recipe does a full fine-tune from the base `.nemo` checkpoint using NeMo's `speech_to_text_finetune.py`, following NVIDIA's [official fine-tuning recipe](https://huggingface.co/blog/nvidia/fine-tuning-nemotron-35-asr). It uses the small public AN4 corpus as a stand-in so it runs end-to-end quickly. + +The job runs on the `nvcr.io/nvidia/nemo:26.02` container — the last consolidated NeMo Framework image that bundles the ASR/TTS collection with a pre-tested CUDA/numba/lightning/kaldialign stack (the newer `26.06` image is Megatron/LLM-only, and the collection now lives in [`NVIDIA-NeMo/Speech`](https://github.com/NVIDIA-NeMo/Speech)). `run.sh` uses the container's bundled NeMo when it has the streaming-prompt recipe, and otherwise clones it and prepends it to `PYTHONPATH` so the recipe code runs against the container's already-installed, correctly CUDA-linked deps. The base checkpoint is delivered read-only via the [Baseten Delivery Network (BDN)](https://docs.baseten.co/training/concepts/storage) — declared as a `WeightsSource` in `config.py` — so it's on local disk before training starts and never costs billed GPU time to download. + +**Resources:** 1 node, 1x H100 GPU + +## Prerequisites + +1. [Create a Baseten account](https://baseten.co/signup) if you don't already have one. +2. Add a Hugging Face access token as the Baseten secret `hf_access_token` (used to download the gated base checkpoint). See [secrets](https://app.baseten.co/settings/secrets). +3. Install the Truss CLI: + ```bash + # pip + pip install -U truss + # or uv + uv add truss + ``` + +## Getting Started + +Initialize the example, navigate into the directory, and push the training job: + +```bash +truss train init --examples nemotron-asr-nemo +cd nemotron-asr-nemo +truss train push config.py +``` + +## Using Your Own Data + +The example trains on AN4 purely as a smoke test. To fine-tune on your own language, domain, or accent, replace `prepare_data.py` with your own manifest builder that emits NeMo JSON-lines manifests: + +```json +{"audio_filepath": "/abs/path/clip.wav", "duration": 4.27, "text": "Reference transcript.", "lang": "en-US", "target_lang": "en-US"} +``` + +Two details matter most: + +- **Every clip needs a `target_lang` tag** matching a locale the model recognizes (e.g. `en-US`, `es-ES`, `el-GR`, `bg-BG`). This drives the model's prompt-based language conditioning and is unforgiving of mismatched labels. +- **Match the base model's text style** — punctuated, properly-cased transcripts, since that's what the model produces. + +For larger datasets, point the trainer at tarred NeMo/Lhotse shards and drive training with a fixed step budget (`trainer.max_steps`) rather than epochs. When specializing on a few languages in the multilingual model, blend in a slice of the other languages ("replay") to avoid eroding them. + +For real, **frozen** datasets (e.g. tarred shards on Hugging Face, S3, or GCS), mount them through BDN the same way as the base checkpoint rather than downloading them in `run.sh` — add another `WeightsSource` in `config.py` and read from the mount path: + +```python +weights=[ + WeightsSource(source=f"hf://{INIT_MODEL}", mount_location=INIT_MODEL_MOUNT, auth_secret_name="hf_access_token"), + WeightsSource(source="s3://my-bucket/asr-shards", mount_location="/app/data"), +] +``` + +The AN4 smoke test stays as an in-container download because it's tiny and needs on-the-fly preprocessing (`.sph`→`.wav` + manifest building), which a read-only BDN mount can't do in place. + +## Checkpoints + +NeMo saves a `.nemo` tarball rather than the HF `config.json` + safetensors layout Baseten's checkpoint detector looks for, so `run.sh` moves the final `.nemo` into a `checkpoint-/` directory under `$BT_CHECKPOINT_DIR` and writes a minimal `config.json` marker. This registers the artifact so it's backed up and appears in `truss train checkpoints list`. It's a NeMo checkpoint — download it and load it with NeMo (see below), not the standard vLLM/HF deploy path. + +## Evaluate + +Evaluate at your deployment latency on held-out data using NeMo's streaming inference script. The lowest-latency setting (`att_context_size=[56,0]`, 80ms chunk, 0ms look-ahead) is the most demanding, honest condition: + +```bash +python ${NEMO_DIR}/examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py \ + model_path= \ + dataset_manifest= \ + target_lang=auto \ + att_context_size="[56,3]" \ + decoder_type=rnnt \ + pad_and_drop_preencoded=true \ + batch_size=8 \ + strip_lang_tags=false +``` + +The same checkpoint covers the whole latency/accuracy spectrum — pick the operating point at inference time via `att_context_size`, no retraining required. diff --git a/examples/nemotron-asr-nemo/config.py b/examples/nemotron-asr-nemo/config.py new file mode 100644 index 0000000..211d744 --- /dev/null +++ b/examples/nemotron-asr-nemo/config.py @@ -0,0 +1,59 @@ +from truss_train import definitions, WeightsSource +from truss.base import truss_config + +# Last consolidated NeMo Framework container that bundles the ASR/TTS collection +# with a pre-tested CUDA/numba/lightning/kaldialign stack (the newer 26.06 image +# is Megatron/LLM-only). NVIDIA-NeMo/Speech recommends 26.02 as the current +# stable runtime for speech models, so it sidesteps the source-build + numba-JIT +# issues you hit on a plain PyTorch image. +BASE_IMAGE = "nvcr.io/nvidia/nemo:26.02" + +# Base checkpoint. Delivered read-only via BDN before start commands run, so we +# never pay for the ~2.4GB download on billed GPU time (BDN mirrors + caches it +# across jobs). See https://docs.baseten.co/training/concepts/storage +INIT_MODEL = "nvidia/nemotron-3.5-asr-streaming-0.6b" +INIT_MODEL_MOUNT = f"/app/models/{INIT_MODEL}" + +training_runtime = definitions.Runtime( + start_commands=["/bin/sh -c 'chmod +x ./run.sh && ./run.sh'"], + environment_variables={ + # Gated on Hugging Face - used by BDN auth and any in-container downloads. + "HF_TOKEN": definitions.SecretReference(name="hf_access_token"), + "HF_HUB_ENABLE_HF_TRANSFER": "true", + # run.sh reads this to locate the BDN-mounted checkpoint. + "INIT_MODEL_MOUNT": INIT_MODEL_MOUNT, + }, + cache_config=definitions.CacheConfig( + enabled=True, + ), + checkpointing_config=definitions.CheckpointingConfig( + enabled=True, + ), +) + +training_compute = definitions.Compute( + node_count=1, + accelerator=truss_config.AcceleratorSpec( + accelerator=truss_config.Accelerator.H100, + count=1, + ), +) + +training_job = definitions.TrainingJob( + image=definitions.Image(base_image=BASE_IMAGE), + compute=training_compute, + runtime=training_runtime, + # Mount the base checkpoint via BDN instead of downloading it in run.sh. + # `auth_secret_name` authenticates against the gated HF repo. + weights=[ + WeightsSource( + source=f"hf://{INIT_MODEL}", + mount_location=INIT_MODEL_MOUNT, + auth_secret_name="hf_access_token", + ), + ], +) + +training_project = definitions.TrainingProject( + name="Nemotron-3.5-ASR-Streaming Finetuned", job=training_job +) diff --git a/examples/nemotron-asr-nemo/prepare_data.py b/examples/nemotron-asr-nemo/prepare_data.py new file mode 100644 index 0000000..a059ab7 --- /dev/null +++ b/examples/nemotron-asr-nemo/prepare_data.py @@ -0,0 +1,121 @@ +"""Prepare a NeMo ASR manifest for fine-tuning Nemotron 3.5 ASR. + +This uses the small, public AN4 corpus as a stand-in for your own data so the +example runs end-to-end quickly. Swap in your own audio + transcripts by +producing JSON-lines manifests with the same schema: + + {"audio_filepath": "/abs/path/clip.wav", "duration": 4.27, + "text": "reference transcript", "lang": "en-US", "target_lang": "en-US"} + +Two details matter most for Nemotron 3.5 ASR: + * Every clip carries a ``target_lang`` tag - this drives the model's + prompt-based language conditioning. Use a locale the model recognizes + (e.g. en-US, es-ES, el-GR, bg-BG). + * Transcripts should match the base model's text style: punctuated and + properly cased. AN4 is lowercase/unpunctuated, so expect it only as a + smoke test, not a quality benchmark. +""" + +import argparse +import glob +import json +import os +import subprocess +import tarfile +import urllib.request + +AN4_URL = "https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz" + + +def download_and_extract_an4(data_dir: str) -> str: + source_data_dir = os.path.join(data_dir, "an4") + if os.path.exists(source_data_dir): + print(f"AN4 already present at {source_data_dir}, skipping download.") + return source_data_dir + + tar_path = os.path.join(data_dir, "an4_sphere.tar.gz") + print(f"Downloading AN4 from {AN4_URL} ...") + urllib.request.urlretrieve(AN4_URL, tar_path) + print("Extracting AN4 ...") + with tarfile.open(tar_path) as tar: + tar.extractall(path=data_dir) + return source_data_dir + + +def build_manifest(transcripts_path, manifest_path, target_wavs_dir, target_lang): + import librosa + + with open(transcripts_path, "r") as fin, open(manifest_path, "w") as fout: + for line in fin: + # Lines look like: transcript (fileID) + transcript = line[: line.find("(") - 1].lower() + transcript = transcript.replace("", "").replace("", "").strip() + + file_id = line[line.find("(") + 1 : -2] + audio_path = os.path.join(target_wavs_dir, file_id + ".wav") + duration = librosa.core.get_duration(path=audio_path) + + metadata = { + "audio_filepath": audio_path, + "duration": duration, + "text": transcript, + "lang": target_lang, + "target_lang": target_lang, + } + json.dump(metadata, fout) + fout.write("\n") + print(f"Wrote manifest: {manifest_path}") + + +def main(args): + data_dir = args.data_dir + os.makedirs(data_dir, exist_ok=True) + + source_data_dir = download_and_extract_an4(data_dir) + target_data_dir = os.path.join(data_dir, "an4_converted") + target_wavs_dir = os.path.join(target_data_dir, "wavs") + os.makedirs(target_wavs_dir, exist_ok=True) + + # Convert the .sph source files to mono 16kHz .wav files. + sph_list = glob.glob(os.path.join(source_data_dir, "**/*.sph"), recursive=True) + print(f"Converting {len(sph_list)} .sph files to .wav ...") + for sph_path in sph_list: + wav_path = os.path.join( + target_wavs_dir, + os.path.splitext(os.path.basename(sph_path))[0] + ".wav", + ) + subprocess.run(["sox", sph_path, "-r", "16000", "-c", "1", wav_path], check=True) + + build_manifest( + os.path.join(source_data_dir, "etc/an4_train.transcription"), + os.path.join(target_data_dir, "train_manifest.json"), + target_wavs_dir, + args.target_lang, + ) + build_manifest( + os.path.join(source_data_dir, "etc/an4_test.transcription"), + os.path.join(target_data_dir, "test_manifest.json"), + target_wavs_dir, + args.target_lang, + ) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare AN4 manifests for NeMo ASR.") + parser.add_argument( + "--data_dir", + type=str, + default=os.environ.get("DATA_DIR", "./data"), + help="Directory to download and process the dataset into.", + ) + parser.add_argument( + "--target_lang", + type=str, + default="en-US", + help="Language tag applied to every clip (drives prompt conditioning).", + ) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/nemotron-asr-nemo/requirements.txt b/examples/nemotron-asr-nemo/requirements.txt new file mode 100644 index 0000000..d1c9ab3 --- /dev/null +++ b/examples/nemotron-asr-nemo/requirements.txt @@ -0,0 +1,4 @@ +# The NeMo container already provides nemo_toolkit[asr] and its full ASR stack +# (lightning, numba, kaldialign, librosa, ...). Only add the small extra used to +# speed up the Hugging Face checkpoint download in the local-dev fallback. +hf_transfer diff --git a/examples/nemotron-asr-nemo/run.sh b/examples/nemotron-asr-nemo/run.sh new file mode 100755 index 0000000..b743510 --- /dev/null +++ b/examples/nemotron-asr-nemo/run.sh @@ -0,0 +1,138 @@ +#!/bin/bash +set -eux + +# ------------------------------------------------------------------ +# Config +# ------------------------------------------------------------------ +export MODEL_ID="nvidia/nemotron-3.5-asr-streaming-0.6b" +export TARGET_LANG="en-US" + +# BDN mounts the base checkpoint here (set in config.py). Falls back to a local +# path for dev runs outside Baseten. +INIT_MODEL_MOUNT="${INIT_MODEL_MOUNT:-./assets/$MODEL_ID}" +export DATA_DIR="./data" +NEMO_REPO="https://github.com/NVIDIA-NeMo/NeMo" +# Pin to a release tag/commit for reproducibility once one is published for this +# recipe (the streaming-prompt config currently only lives on main). +NEMO_BRANCH="main" +NEMO_DIR="/opt/NeMo" + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "WARNING: HF_TOKEN is not set. The base checkpoint is gated on Hugging Face." + echo "Configure Baseten secret 'hf_access_token' and map it to HF_TOKEN in config.py." +fi + +# ------------------------------------------------------------------ +# 1. Python dependencies +# ------------------------------------------------------------------ +# The NeMo container already ships nemo_toolkit[asr] and all system audio deps +# (sox, ffmpeg, libsndfile) with a pre-tested CUDA/numba/lightning stack, so we +# only add the small extras used for downloads. +pip install -r requirements.txt + +# ------------------------------------------------------------------ +# 2. Locate the streaming-prompt fine-tune script + config +# ------------------------------------------------------------------ +# Prefer the NeMo source bundled in the container. If this container predates the +# streaming-prompt recipe, clone it and prepend to PYTHONPATH so the recipe code +# runs against the container's already-installed (and correctly CUDA-linked) deps +# - we do NOT reinstall nemo_toolkit, which would undo the container's tested stack. +CONFIG_REL="examples/asr/conf/fastconformer/cache_aware_streaming/fastconformer_transducer_bpe_streaming_prompt.yaml" +NEMO_DIR="" +for cand in /opt/NeMo /workspace/NeMo /opt/nemo; do + if [[ -f "$cand/$CONFIG_REL" ]]; then + NEMO_DIR="$cand" + break + fi +done +if [[ -z "$NEMO_DIR" ]]; then + echo "Streaming-prompt recipe not bundled in the container; cloning NeMo." + NEMO_DIR="/opt/NeMo-src" + [[ -d "$NEMO_DIR" ]] || git clone -b "$NEMO_BRANCH" "$NEMO_REPO" "$NEMO_DIR" + export PYTHONPATH="$NEMO_DIR:${PYTHONPATH:-}" +fi +echo "Using NeMo source at: $NEMO_DIR" + +# ------------------------------------------------------------------ +# 3. Locate the base checkpoint (.nemo) +# ------------------------------------------------------------------ +# On Baseten it's already on local disk via BDN (see config.py `weights`), so +# there's no download on billed GPU time. Outside Baseten, pull it once. +if [[ ! -d "$INIT_MODEL_MOUNT" ]]; then + echo "Checkpoint mount not found; downloading $MODEL_ID (local dev fallback)." + huggingface-cli download "$MODEL_ID" --local-dir "$INIT_MODEL_MOUNT" +fi +HF_CKPT="$(find "$INIT_MODEL_MOUNT" -name '*.nemo' | head -n 1)" +echo "Using base checkpoint: $HF_CKPT" + +# ------------------------------------------------------------------ +# 4. Prepare data (swap prepare_data.py for your own manifest builder) +# ------------------------------------------------------------------ +python3 prepare_data.py --data_dir "$DATA_DIR" --target_lang "$TARGET_LANG" + +# ------------------------------------------------------------------ +# 5. Fine-tune from the base checkpoint +# ------------------------------------------------------------------ +# Full fine-tune of the Cache-Aware FastConformer-RNNT streaming model. +# - Prefer a fixed step budget over epochs for streaming/iterable data; AN4 is +# tiny so we cap with limit_train_batches for a quick smoke run. +# - Increase trainer.max_epochs / point at more data for a real run. +# - Reduce ++model.train_ds.batch_duration if you hit out-of-memory errors. +python3 "$NEMO_DIR/examples/asr/speech_to_text_finetune.py" \ + --config-path="../asr/conf/fastconformer/cache_aware_streaming" \ + --config-name=fastconformer_transducer_bpe_streaming_prompt.yaml \ + +init_from_nemo_model="$HF_CKPT" \ + ++model.train_ds.manifest_filepath="$DATA_DIR/an4_converted/train_manifest.json" \ + ++model.validation_ds.manifest_filepath="$DATA_DIR/an4_converted/test_manifest.json" \ + ++model.train_ds.batch_duration=200 \ + ++model.optim.name="adamw" \ + ++model.optim.lr=0.1 \ + ++model.optim.weight_decay=0.001 \ + ++model.optim.sched.d_model=1024 \ + ++model.optim.sched.warmup_steps=100 \ + ++trainer.devices="$BT_NUM_GPUS" \ + ++trainer.max_epochs=1 \ + ++trainer.limit_train_batches=60 \ + ++trainer.precision=bf16 \ + ++exp_manager.exp_dir="$BT_CHECKPOINT_DIR" \ + ++exp_manager.use_datetime_version=False \ + ++exp_manager.version=finetune + +# ------------------------------------------------------------------ +# 6. Register the checkpoint with Baseten +# ------------------------------------------------------------------ +# Baseten's checkpoint detector registers an artifact when it finds a +# `checkpoint-/config.json` layout under $BT_CHECKPOINT_DIR (rank-0 root). +# NeMo emits a `.nemo` tarball rather than HF safetensors + config.json, so we +# move the final .nemo into a checkpoint-/ dir and drop a minimal config.json +# marker so the checkpoint is registered, backed up, and shows up in +# `truss train checkpoints list`. It's a NeMo checkpoint - load it with NeMo +# (see README "Evaluate"), not the vLLM/HF deploy path. +FINAL_NEMO="$(find "$BT_CHECKPOINT_DIR" -name '*.nemo' -printf '%T@ %p\n' 2>/dev/null \ + | sort -nr | head -n 1 | cut -d' ' -f2-)" + +if [[ -n "${FINAL_NEMO:-}" && -f "$FINAL_NEMO" ]]; then + # Derive N from the checkpoint name if it carries a step, else default to 1. + STEP="$(basename "$FINAL_NEMO" .nemo | grep -oE '[0-9]+' | tail -n 1)" + STEP="${STEP:-1}" + CKPT_OUT="$BT_CHECKPOINT_DIR/checkpoint-$STEP" + mkdir -p "$CKPT_OUT" + mv "$FINAL_NEMO" "$CKPT_OUT/" + NEMO_FILE="$(basename "$FINAL_NEMO")" + + # The detector registers on config.json with an `architectures` or `model_type` + # field (no safetensors needed) and reads the base model from `_name_or_path`. + cat > "$CKPT_OUT/config.json" <