Amazon EC2上でのNVIDIA Nemotron Speech ASRのドメイン適応のためのファインチューニング

*This post is a collaboration between AWS, NVIDIA and Heidi. *

Automatic speech recognition (ASR), often called speech-to-text (STT) is becoming increasingly critical across industries like healthcare, customer service, and media production. While pre-trained models offer strong capabilities for general speech, fine-tuning for specific domains and use cases can enhance accuracy and performance.

In this post, we explore how to fine-tune a leaderboard-topping, NVIDIA Nemotron Speech Automatic Speech Recognition (ASR) model; Parakeet TDT 0.6B V2. Using synthetic speech data to achieve superior transcription results for specialised applications, we’ll walk through an end-to-end workflow that combines AWS infrastructure with the following popular open-source frameworks:

• Amazon Elastic Compute Cloud (Amazon EC2) GPU instances (p4d.24xlarge with NVIDIA A100 GPUs) for distributed training at scale

• NVIDIA NeMo framework for ASR model fine-tuning and optimization

• DeepSpeed for memory-efficient distributed training across multiple nodes

• MLflow and TensorBoard for comprehensive experiment tracking

• Amazon Elastic Kubernetes Services (Amazon EKS) for scalable model serving

• Amazon FSx for Lustre for high-performance model weight storage

• AI Gateway and Langfuse for production-grade API management and observability

• Docker for consistent, reproducible environments across training and inference

This architecture demonstrates how to use the managed services of AWS alongside best-in-class open-source AI tools to build production-ready, domain-adapted ASR systems that deliver measurable business value—from initial fine-tuning through to elastic, observable deployment.

Solution Overview: Heidi’s AI Care Partner

Heidi is an AI Care Partner that removes the work around care—handling documentation, clinical evidence, and patient communications so clinicians can stay focused on patients. The platform supports over 2.4 million consultations per week in 110 languages across 190 countries. Used across emergency departments, general practice, and specialist clinics, Heidi helps clinicians reclaim hours each day while maintaining the accuracy and integrity of the clinical record.

Out-of-the-box ASR models struggle with medical terminology, regional accents, and code-switching between clinical and conversational language. These limitations lead to transcription errors, lost context, and increased cognitive load which forces clinicians to spend time correcting what should have saved them time. For clinicians, accurate documentation isn’t just convenience. It’s clinical safety, liability protection, and trust in the tool. A note with errors undermines all three.

To address this, Heidi collaborated with AWS Generative AI Innovation Center (GenAIIC) to fine-tune and adapt the model to the unique linguistic, acoustic, and contextual nuances of real-world clinical environments—enabling accurate and reliable performance at scale. Using recent advancements in text-to-speech (TTS) models, Heidi generated high-quality, multilingual synthetic speech interleaved with real-world noises from conversations emulated with large language models (LLMs). This approach allowed the team to scale training across a wide range of accents, and medical contexts without compromising patient privacy. Using synthetic data also enabled targeted augmentation with focus on low-resource languages and rare medical terms that are underrepresented in open datasets.

Fine-tuning was conducted using Amazon EC2 GPU instances optimized for deep learning workloads. By usings pre-configured AWS Deep Learning AMIs, the team was able to accelerate experimentation and model iteration while maintaining control over performance and security. The combination of scalable compute and tightly integrated AWS services enabled fast, cost-effective development within a highly regulated environment.

*Figure 1: This figure shows the diagram for fine-tuning and deploying NVIDIA Parakeet TDT 0.6B V2 model on AWS services.*

Synthesizing domain-specific data

To help improve NVIDIA Parakeet TDT 0.6B V2‘s performance on medically specialized terminology, we developed a targeted synthetic data generation pipeline combining large language models (LLMs), neural text-to-speech (TTS) synthesis, and noise augmentation. The process began by compiling a lexicon of medical terms, primarily drug names, anatomical entities, and procedural phrases that exhibited low recall in prior evaluation runs. These terms were used as conditioning inputs for a domain-adapted LLM, which generated semantically coherent and contextually diverse transcripts resembling real-world clinical dictations. The prompts were designed to elicit varied linguistic structures and natural code-switching patterns that are frequently observed in medical speech. For example, intermixing abbreviations, Latin-origin words, and colloquial phrasing.

The generated transcripts were later transformed into speech using a neural TTS system, configured for specific accents and natural prosody. To increase data diversity and realism, we employed a multi-stage audio augmentation pipeline. This involved overlaying hospital and clinic ambient noise recordings such as background conversations, machine alarms, and hospital ambient noises sampled from open acoustic datasets. We also applied controlled perturbations, including random gain adjustments, reverberation filters, and additive Gaussian noise to simulate real-world microphone and environment conditions. Each synthetic utterance was assigned a randomized signal-to-noise ratio (SNR) between 10–25 dB to balance clarity and authenticity.

Introduction to NVIDIA Parakeet TDT 0.6B V2 model

The NVIDIA Parakeet TDT 0.6B V2 is a 600-million parameter automatic speech recognition (ASR) model designed for high-quality English transcription. Built on NVIDIA NeMo’s FastConformer architecture with a Token-and-Duration Transducer (TDT) decoder, this model helps provide exceptional speech recognition capabilities along with valuable features like:

• Automatic punctuation and capitalization

• Word-level timestamp predictions

• Robust performance on spoken numbers and song lyrics

• Support for audio segments up to 24 minutes in a single pass

The open model achieves an impressive 6.05% average Word Error Rate (WER) across various benchmark datasets on the Hugging Face Open-ASR Leaderboard, demonstrating its effectiveness for general English speech recognition.

Why fine-tune the model?

Despite its strong out-of-the-box performance, fine-tuning the NVIDIA Parakeet TDT 0.6B V2 model for specific domains offers several compelling advantages:

• Domain-specific terminology – Enhanced recognition of specialized vocabulary and jargon that can be rare in general training datasets

• Accent and dialect adaptation – Improved performance for specific regional speech patterns or linguistic variations

• Noise resilience – Better handling of domain-specific background noises and acoustic environments

• Cost efficiency – Reduced need for expensive third-party API calls with your own optimized model

• Customized optimization – Balance between accuracy and inference speed tailored to your specific application requirements

Model architecture

The NVIDIA Parakeet TDT 0.6B V2 model uses a Token-and-Duration Transducer (TDT) architecture with these key components:

Encoder architecture

The encoder transforms input audio features into contextual representations through multiple processing stages:

• Conformer Encoder: 24 conformer layers with 1024 hidden dimensions

• Convolutional Subsampling: Initial feature extraction with four convolutional layers

3×3 convolutions with stride 2 for progressive downsampling

• Channel dimension of 256 throughout the convolutional stack

• Positional Encoding: Relative positional encoding with dropout (0.1)

Conformer layers

Each conformer layer includes:

• Feed Forward Modules: Two FFN blocks with 4096 inner dimension

• Convolutional Module: 9×1 depthwise convolution with causal padding, batch normalization, and Swish activation

• Multi-Head Attention: Self-attention mechanism with relative positional encoding

• Normalization: Layer normalization between components

• Dropout: 0.1 dropout rate for regularization

Decoder & Joint Network

The decoder and joint network combine encoder outputs with predicted tokens to generate final transcriptions:

RNN Decoder:

• Embedding layer (1025 vocabulary size, 640 embedding dimension)

• 2-layer LSTM with 640 hidden units and 0.2 dropout

Joint Network:

• Linear transformations for encoder and decoder

• ReLU activation with 0.2 dropout

• Output projection to 1030 dimensions

Audio processing and loss function

The model employs specialized preprocessing and training objectives to help optimize transcription accuracy:

• Mel Spectrogram Preprocessing: Converts raw audio to mel-scale spectrograms

• Spectral Augmentation: Time and frequency masking to help enhance robustness

• TDT Loss: Token-and-Duration Transducer loss for simultaneous prediction of tokens and their durations

• Word Error Rate (WER): Primary evaluation metric for transcription quality

Setting up your environment for fine-tuning

Our fine-tuning approach leverages distributed training on Amazon EC2 instances, with a path to deployment on Amazon SageMaker AI. The environment is encapsulated in a Docker container to facilitate consistency and reproducibility.

Docker-based environment setup

The Docker container includes the necessary dependencies for fine-tuning as shown in the following code snippet. You can access the complete file in the associated GitHub repository here:

FROM nvcr.io/nvidia/pytorch:24.01-py3

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=UTC

# Set the working directory
WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    libsox-fmt-mp3 \
    gnupg \
    && rm -rf /var/lib/apt/lists/*
    
# Install Cython (needed for NeMo)
RUN pip install Cython

# Clone the specified branch of the pytorch-lightning repository and install it
RUN git clone -b bug_fix https://github.com/athitten/pytorch-lightning.git && \
    cd pytorch-lightning && \
    PACKAGE_NAME=pytorch pip install -e .

# Install TransformerEngine for optimization
RUN git clone https://github.com/NVIDIA/TransformerEngine.git && \
    cd TransformerEngine && \
    git fetch origin 8c9abbb80dba196f086b8b602a7cf1bce0040a6a && \
    git checkout FETCH_HEAD && \
    git submodule init && git submodule update && \

This container provides:

• NVIDIA’s optimized PyTorch container as the base

• NeMo framework for ASR model handling

• DeepSpeed for efficient distributed training

• MLflow and TensorBoard for experiment tracking

Resource requirements

For efficient fine-tuning, we recommend:

• EC2 instance type: p4d.24xlarge (8 NVIDIA A100 GPUs per instance) – The A100 GPUs provide 80GB of high-bandwidth memory per GPU, essential for handling the Parakeet TDT 0.6B V2 model’s 600 million parameters with large batch sizes. A100 tensor cores accelerate mixed-precision training, reducing training time compared to previous-generation GPUs. Note that there are newer, more powerful EC2 instances such as P5 that are also available to cater for more demanding requirements.

• Cluster size: 8 nodes (64 GPUs total) for full-scale training – Distributed training across multiple nodes enables larger batch sizes and faster convergence through parallelization, supporting efficient training on large-scale audio datasets. With 64 GPUs, you can train on extensive datasets (100+ hours of audio) in hours rather than days, enabling rapid experimentation for production timelines.

• Storage: At least 500 GB per node for model checkpoints and data. This capacity accommodates frequent checkpoint saving, intermediate training states, and local caching of preprocessed audio features. Adequate storage prevents I/O bottlenecks that could idle expensive GPU resources during training.

For smaller datasets or budget constraints, the training can scale down to fewer nodes or GPU instances like g6e.2xlarge.

Data preparation for fine-tuning

The Parakeet TDT 0.6B V2 model requires audio data and corresponding transcriptions in NeMo’s JSONL manifest format. Each line in these files points to the synthesized audio and the corresponding transcript that it was synthesized from.

{   "audio_filepath": "/path/to/audio.wav", "duration": 5.2, 
    "text": "The transcription with punctuation and capitalization."
}

Creating training manifests

Our fine-tuning approach uses three separate manifest files:

• Training manifest: Contains the bulk of your data for model training

• Validation manifest: Used to evaluate model performance during training

• Test manifest: Used for final evaluation of the fine-tuned model

The dataset should ideally include diverse speakers, accents, and acoustic conditions relevant to your domain for improved performance.

Fine-Tuning configuration deep dive

Our fine-tuning configuration is comprehensive and carefully tailored for the Parakeet TDT 0.6B V2 model:

Model configuration

The configuration specifies the Parakeet TDT 0.6B V2 architecture with 17 conformer blocks:

model:
  sample_rate: 16000
  normalize_text: true
  symbols_to_keep: ["'"]
  
  encoder:
    _target_: nemo.collections.asr.modules.ConformerEncoder
    feat_in: \${model.preprocessor.features}
    n_layers: 17
    d_model: 512
    subsampling: dw_striding
    subsampling_factor: 8
    subsampling_conv_channels: 256

Token-and-Duration Transducer (TDT) settings

The TDT architecture is configured with specific duration values to predict token emission timing:

model:
  model_defaults:
    tdt_durations: [0, 1, 2, 3, 4]
    num_tdt_durations: 5
    
  loss:
    loss_name: "tdt"
    tdt_kwargs:
      durations: \${model.max_duration}
      sigma: 0.02
      omega: 0.1

Audio preprocessing

Audio preprocessing is configured for optimal feature extraction:

preprocessor:
  _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
  sample_rate: \${model.sample_rate}
  normalize: "per_feature"
  window_size: 0.025
  window_stride: 0.01
  features: 128
  n_fft: 512
  dither: 0.00001

Data augmentation

To help improve model robustness, multiple augmentation techniques such as Spectrogram augmentation are applied:

train_ds:
  augmentor:
    speed:
      prob: 0.4
      min_speed_rate: 0.9
      max_speed_rate: 1.1
      
spec_augment:
  _target_: nemo.collections.asr.modules.SpectrogramAugmentation
  freq_masks: 2
  time_masks: 10
  freq_width: 27
  time_width: 0.05

Additional augmentation strategies like noise augmentation and shift perturbation are available as commented options in the configuration file.

Distributed training strategy

The training is configured for efficient distribution across multiple GPUs and nodes:

trainer:
  devices: 8
  num_nodes: 8
  strategy:
    _target_: "lightning.pytorch.strategies.DeepSpeedStrategy"
    stage: 2
    offload_optimizer: true
    partition_activations: true
    gradient_as_bucket_view: true
    cpu_checkpointing: true
    contiguous_gradients: true
    overlap_comm: true

This DeepSpeed configuration improves memory usage and communication between GPUs for maximum training efficiency.

Implementing the fine-tuning process

Our implementation uses a modular approach with a dedicated ASRTrainer class:

Model initialization and unfreezing

def get_base_model(self, trainer):

"""Get the base model to start training from based on config settings."""

pretrained_name = self.config.init_from_pretrained_model

Handle multi-GPU download efficiently

num_ranks = trainer.num_devices * trainer.num_nodes

if num_ranks > 1 and is_global_rank_zero():

logging.info(f"Downloading pretrained model '{pretrained_name}' on main process")

asr_model = ASRModel.from_pretrained(model_name=pretrained_name)

else:

Wait for model download to complete on main process

wait_time = 1 if is_global_rank_zero() else 60

logging.info(f"Waiting {wait_time}s for model download")

time.sleep(wait_time)

asr_model = ASRModel.from_pretrained(model_name=pretrained_name)

asr_model.to(f"cuda:{int(os.environ.get('LOCAL_RANK', 0))}")