Getting Kaldi’s nnet3 TDNN to Train on GPU: A Practical Walk‑through

Problem: Slow CPU‑only training limits iteration speed

When experimenting with speech recognition models, waiting hours for each training epoch on CPU slows down research and development. Kaldi’s nnet3 framework provides time‑delay neural network (TDNN) acoustic models that can replace GMM‑HMMs, but to reap their speed benefits you need to run the training on a GPU.

Thesis: A correctly built Kaldi with CUDA lets nnet3 TDNNs train quickly on GPU, delivering state‑of‑the‑art accuracy while keeping the workflow familiar.

1. Building Kaldi with GPU support

First, obtain the Kaldi source and configure it for your CUDA toolkit. Assume you have root or sudo access to install dependencies.

1. Install CUDA (≥11.0) and cuDNN (≥8.0) according to NVIDIA’s instructions.
2. Clone Kaldi: git clone https://github.com/kaldi-asr/kaldi.git
3. Enter the tools directory and check prerequisites: cd kaldi/tools; ./extras/check_dependencies.sh
4. Compile OpenFST and other dependencies: make -j$(nproc)
5. Return to the root and configure with CUDA: cd ..; ./configure --cuda-dir=/usr/local/cuda
6. Build Kaldi: make -j$(nproc)

After the build, verify that the GPU binaries are present: ls src/nnet3bin/nnet3-train-* should show executables without a "_cpu" suffix.

2. Preparing a LibriSpeech subset and alignments

We will use the popular LibriSpeech "train-clean-100" subset as an example. Adjust paths as needed.

1. Run the data preparation script (provided in the egosuite): local/prepare_data.sh --subset train-clean-100 /path/to/LibriSpeech
2. Generate MFCC features: steps/make_mfcc.sh --nj 8 --cmd "run.pl" data/train_clean_100 exp/make_mfcc/train_clean_100 mfcc
3. Compute cepstral mean and variance normalization (CMVN) stats: steps/compute_cmvn_stats.sh data/train_clean_100 exp/make_mfcc/train_clean_100 mfcc
4. Train a monophone model and generate alignments needed for TDNN training: steps/align_si.sh --nj 8 --cmd "run.pl" data/train_clean_100 data/lang exp/mono exp/mono_ali_train_clean_100

At this point you have data/train_clean_100/feats.scp and exp/mono_ali_train_clean_100/ali.*.gz ready for the nnet3 pipeline.

3. Training a TDNN with nnet3 on GPU

Kaldi ships a ready‑made script for TDNN training in egs/librispeech/s5/local/tdnn/run_tdnn.sh. We will call it with GPU‑aware options.

cd egs/librispeech/s5
# Export the CUDA visible devices (adjust if you have multiple GPUs)
export CUDA_VISIBLE_DEVICES=0
# Launch the training script; it will automatically pick up the GPU binaries
local/tdnn/run_tdnn.sh --stage 0 --train-stage -1 \
    --num-jobs-initial 2 --num-jobs-final 4 \
    --num-epochs 4 --hidden-layer-dim 650 \
    --egs-dir exp/tdnn_sp/egs \
    --dir exp/tdnn_sp

While the script runs, you can open another terminal and watch GPU utilization:

nvidia-smi -l 1

You should see the Kaldi training process (e.g., nnet3-train-lstm-streams) consuming GPU memory and showing non‑zero GPU‑utilization percentages.

To confirm that the binary is indeed the CUDA version, inspect the log file:

grep -i cuda exp/tdnn_sp/log/train.*.log

You ought to find a line similar to Using GPU device 0 or CUDA version: 11.2. If you see only CPU‑related messages, the build likely fell back to CPU; revisit the configuration step.

4. Trade‑offs and practical verification

• GPU memory vs. model size: Larger TDNNs (more hidden layers or wider layers) increase accuracy but also GPU memory consumption. If nvidia-smi shows memory usage near the total capacity, reduce the mini‑batch size via the --minibatch-size option or lower the hidden dimension.
• Version compatibility: Mismatched CUDA/cuDNN versions can cause silent fallbacks. The safest route is to match the versions listed in Kaldi’s INSTALL file (e.g., CUDA 11.8 with cuDNN 8.9). After building, run the sanity check:

cd src
./kaldi_gpu_test

The test should report "All tests passed" and list the detected GPU architecture.

Online decoding latency: Once training finishes, you can test real‑time performance with the online2 decoder:

steps/online/nnet3/prepare_online_decoding.sh --mfcc-config conf/mfcc.conf \
    data/lang exp/tdnn_sp exp/tdnn_sp_online
online2-wav-nnet3-latency --frame-subsampling-factor=3 \
    exp/tdnn_sp_online/final.mdl exp/tdnn_sp_online/online_conf \
    "ark:echo utterance-id1 utterance-id2|" \
    "scp:path/to/test/wav.scp" "ark,t:path/to/output/text"

Measure latency with time or by checking the timestamps in the decoder’s log; typical latencies for a medium‑sized TDNN on a modern GPU are under 100 ms per chunk.

Actionable next steps

1. Re‑run the build with the exact CUDA/cuDNN versions your hardware supports. 2. Start with a modest TDNN (2‑3 hidden layers, 650‑dim) to verify GPU utilization before scaling up. 3. Monitor nvidia-smi and the training logs for any fallback to CPU. 4. If you hit memory limits, halve the mini‑batch size or enable mixed‑precision by setting --use-gpu=true --gpu-wait=false and recompiling with --cudnn-half if your Kaldi version supports it. 5. Once training converges, decode a short test set with the online2 pipeline to ensure low latency meets your application’s needs.

By following these steps, you turn Kaldi’s nnet3 TDNN from a CPU‑bound experiment into a GPU‑accelerated workflow that shortens iteration cycles and brings you closer to production‑ready speech recognition.