[BUG] fuse_loss_wer=true deadlocks multi-GPU validation when validation_ds.batch_size > fused_batch_size

[jaeminh]

Describe the bug

While following the tutorial for fine-tuning Nemotron 3.5 ASR asr-finetune-nemotron-3.5-asr-streaming-prompt.ipynb, I encountered a deadlock during multi-gpu training.

This issue occurs during the validation phase with fuse_loss_wer=true and validation_ds.batch_size > joint.fused_batch_size.

RNNTJoint.forward splits each validation batch into multiple sub-batches and calls self.wer.compute() for every sub-batch.

Speech/nemo/collections/asr/modules/rnnt.py

Line 1527 in 1c82990

for batch_idx in range(0, batch_size, self._fused_batch_size):

Since the number of sub-batches depends on the local batch size, different ranks can execute a different number of wer.compute() calls. Because wer.compute() performs a synchronized collective (all_gather), the mismatched number of synchronization calls across ranks causes a deadlock.

With 2 GPUs, 20 validation samples, validation_ds.batch_size=16, and joint.fused_batch_size=2, the validation loader produces two batches (16 and 4 samples). One rank receives the batch of 16 samples and executes 8 sub-batches, while the other rank receives the batch of 4 samples and executes 2 sub-batches. The first two wer.compute() calls match across ranks, but after that the rank processing the larger batch continues issuing collectives while the other rank has already exited the loop, causing the remaining collectives to block indefinitely until the NCCL watchdog times out.

The training path already disables sync before calling self.wer.compute(), but this is only applied when self.training is True.

Speech/nemo/collections/asr/modules/rnnt.py

Lines 1598 to 1601 in 1c82990

	# Update WER on each process without syncing
	if self.training:
	original_sync = self.wer._to_sync
	self.wer._to_sync = False

Steps/Code to reproduce bug

Follow the tutorial in asr-finetune-nemotron-3.5-asr-streaming-prompt.ipynb .
The AN4 test set contains130 samples. Running validation with 2 GPUs and overriding the validation batch size to 24 consistently reproduces the deadlock.
Set to ++trainer.devices=2 and add ++model.validation_ds.batch_size=24.

N=130 W=2 B=24 f=2 -> rank0 [24,24,24]=36 vs rank1 [24,24,10]=29 collectives -> hang

python /{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.optim.sched.d_model=1024 \
  ++trainer.devices=2 \
  ++trainer.max_epochs=1 \
  ++trainer.limit_train_batches=60 \
  ++trainer.precision=bf16 \
  ++model.train_ds.batch_duration=200 \
  ++model.optim.name="adamw" \
  ++model.optim.lr=0.1 \
  ++model.optim.weight_decay=0.001 \
  ++model.optim.sched.warmup_steps=100 \
  ++exp_manager.version=test \
  ++exp_manager.use_datetime_version=False \
  ++exp_manager.exp_dir=$DATA_DIR/checkpoints \
  ++model.validation_ds.batch_size=24

Expected behavior

Validation completes on all ranks. The number of collectives a rank issues should not depend on how
many samples that rank was handed.

Environment overview (please complete the following information)

• Environment location: Bare-metal
• Method of NeMo install: git clone https://github.com/NVIDIA-NeMo/Speech && uv sync --extra all --extra cu13

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• Ubuntu 24.04.3 LTS
• 2.12.0+cu132
• 3.12.3

Additional context

GPU model: NVIDIA A100-SXM4-80GB

[Sahith59]

Hi @jaeminh, thanks for the detailed write-up! I'd like to work on a fix for this...!

My plan:

--> During validation sub-batch loops in RNNTJoint.forward, disable wer.sync before the loop (the same way training does at L1598-1601)
--> Re-enable and call wer.compute() once after the loop completes
--> Add a unit test that verifies 2-rank validation doesn't hang with mismatched batch sizes

Does this approach align with what the maintainers would want? Happy to submit a draft PR.