Feature Normalization in the ASR preprocessor is too slow.

[galv]

Another issue I discovered while working on #8673 .

Right now, when running Parakeet CTC 1.1B at batch size 32, the preprocessor takes about 10milliseconds (of course, this depends on the lengths of the inputs, but this was with librispeech test other). 90% of the time is taken by feature normalization, and the GPU is hardly active during that time.

You can reproduce via running this:

nsys profile -c cudaProfilerApi $(which python) examples/asr/speech_to_text_eval.py  pretrained_name=nvidia/parakeet-ctc-1.1b dataset_manifest=/home/dgalvez/scratch/data/test_other_sorted_downward.json  \
batch_size=$batch_size  output_filename=test_other_decoded.jsonl  amp=$amp  amp_dtype=bfloat16  use_cer=false num_workers=1

The culprit is this code that sequentially does a for loop over the batch:

NeMo/nemo/collections/asr/parts/preprocessing/features.py

Lines 65 to 73 in 468d5b6

	for i in range(x.shape[0]):
	if x[i, :, : seq_len[i]].shape[1] == 1:
	raise ValueError(
	"normalize_batch with `per_feature` normalize_type received a tensor of length 1. This will result "
	"in torch.std() returning nan. Make sure your audio length has enough samples for a single "
	"feature (ex. at least `hop_length` for Mel Spectrograms)."
	)
	x_mean[i, :] = x[i, :, : seq_len[i]].mean(dim=1)
	x_std[i, :] = x[i, :, : seq_len[i]].std(dim=1)

First of all, it's actually copying a single scalar value of seq_len to CPU three times per loop. If you have to do a GPU to CPU copy, please do it just once via seq_len = seq_len.cpu(). But the actual kernels being launched are likely to go faster if we just do a single batched operation. Note that the slowness of this sequential for loop operation only increases as you run at larger batch sizes. An extra 10ms of latency during streaming inference is not really acceptable. The person who wrote the code initially probably wasn't sure how to deal with a ragged tensor for doing mean and std reductions, but we can do it straightforwardly with torch.where() and a mask built from seq_len, without ever communicating with the CPU (which is important for cuda stream capture to a graph).

I will either fix this as part of #8673 or make a separate PR, if the change is larger and riskier than I initially think.

You can download the .nsys-rep file I took a screenshot of here:
report6.nsys-rep.gz

[titu1994]

Could you put it to a separate pr ? Let's keep prs small / medium sized and focussed so we can debug them if they cause issues later.

[titu1994]

Also, don't make the assumption that a GPU will always be there (or that cuda graphs are supported). We can have fast paths for those cases when GPU is available or cudapython 12.3 and above is installed via conda / pip but they aren't guaranteed.

[galv]

[don't make the assumption] that cuda graphs are supported

It is good practice to avoid synchronizing with the CPU in general, with GPU code. If we do that, we can get cuda graphs "for free". No part of that is building in cuda graphs as a requirement to do execution.