Streamable Conformer-Transducer ASR model for LibriSpeech

[asumagic]

Contribution in a nutshell

This PR implements a streaming ASR model for SpeechBrain.

Resolves: #1970

This PR will (probably) not be implementing the high-level interfaces (think interfaces.py) to perform inference as my work ends at the end of August, but I'm hoping to get everything in a functional state with a streaming inference proof-of-concept done. The high level interfaces will be implemented separately.

Scope

• Implement the Dynamic Chunk Training method to mask transformer keys according to chunks of time, and vary chunk size across batches.
• Implement AWS' Dynamic Chunk Convolution technique for the Conformer.
• Allow training a model with fully streaming compliant constraints and emulate streaming at inference through masking.
• Provide a tool to empirically analyze dependencies between output tokens and input tokens, in order to validate whether the intended time constraints are respected.
• Write a LibriSpeech Conformer-Transducer recipe.
• Provide a fully trained streaming model.
• Implementing a live microphone demo (using PyAudio)

Additionally, still TODO for me:

• Upload the model somewhere: https://drive.google.com/drive/folders/1vL0Wr15eQZWS8SHp6RJQ_acj__7iri6J?usp=sharing to be transferred to Dropbox
• Update README
• Solve essential TODOs
• Solve pre-commit issues
• Change the experiment seed
• Document new things heavily
• Ensure the code isn't breaking other things
• Probably figuring out what to do with custom_tgt_module, again -- probably needs undoing and patching the model same as the other conformer-transducer
• Implement a fully functional proof-of-concept of chunked streaming (WIP)
• Implementing transducer greedy search

Out-of-scope for this PR (will be done separately):

• Implementing high level interfaces (for HF demos etc.)

Notes for reviewing (optional)

See thread, there is an accuracy issue for streaming decoding ATM. Solved

Explanations and technical details

This should go quite in a bit of detail for the big picture. It could probably eventually be turned into a general documentation on streaming once SB's support matures out.

https://gist.github.com/asumagic/aaf875929d1c03a96fd08a897599aaa8

Modifying the existing recipe?

Streaming has required very little hparams changes so far. I figure that considering the only current difference between a non-streaming model and a streaming model is changing one hyperparameter, it made more sense to keep it at one file for now.

Additionally, as per the above technical doc, the model is also trained with a context extremely close to non-streaming for ~40% of the batches. The result is that the streaming model can operate in a non-streaming fashion, and for LibriSpeech, the non-streaming performance is very close between a streaming: False and a streaming: True model, at least in my testing.

Self-attention logic changes

The attention.py change which moves the scaling before in order to reduce the likelihood of float16 overflow should be reviewed thoroughly as it is extensively used. Namely, even though it looks mathematically correct, the rel_shift trick scares me a bit and someone with a better understanding than me on the subject should double-check it.

Performing the softmax in fp32 space did not seem to affect performance for me and it seems more correct.

Ultimately I trained the model using bf16. I am not sure if training in fp16 is fully stable yet, but I know that these changes reduced the amount of training failures significantly in fp16 mode.

Pre-review

• N/A (if applicable) add an extra_requirements.txt file
• N/A (if applicable) add database preparation scripts & use symlinks for nested folders (to the level of task READMEs)
• (if applicable) add a recipe test entry in the depending CSV file under: tests/recipes
• create a fresh testing environment (install SpeechBrain from cloned repo branch of this PR)
• (if applicable) run a recipe test for each yaml/your recipe dataset
• check function comments: are there docstrings w/ arguments & returns? If you're not the verbose type, put a comment every three lines of code (better: every line)
• use CI locally: pre-commit run -a to check linters; run pytest tests/consistency
• (optional) run tests/.run-doctests.sh & tests/.run-unittests.sh
• exhausted patience before clicking « Ready for review » in the merge box 🍄

[asumagic]

I wrote a relatively extensive description of the whole streaming approach.

This should hopefully clear up a lot of how and why the approach works and hopefully in a more understandable/intuitive way than some of the litterature.

Additionally, I got a simple greedy search prototype working yesterday and could successfully encode several minutes of spoken voice correctly.
Currently, some code is missing as part of a notebook as I needed some external code to glue everything up for streaming, which I'll clean up and share.
I even got a basic, local microphone demo working which got to harm my ego a bit considering how hard of a time it has with french accents, apparently. I blame LibriSpeech 😉

[asumagic]

I probably won't be able to spend much more time on this since my work ends today, and I'll share my notebook with the WIP streaming code today. The PR and code is generally cleaned up.

There is however an issue... The WER% when using the streaming code is bad (when it is fairly decent when emulating streaming by masking), so something is broken somewhere, and I am not sure what. I edited the technical document with a section on what I have attempted so far to debug the issue, and things to try.

So unless my priorities change, someone will have to pick up work from this. Hopefully a new set of eyes would spot the issues that I couldn't find in time.
With this issue solved (and I don't know how much effort that would be, hopefully it wasn't improper masking in training), there technically isn't a lot remaining to have streaming working, though: some cleanups and proper generic interfaces that could be reused for other models, and stuff to actually push it to HF.

I can still answer questions, clarify things and provide some help/guidance if needed of course.

So will all that said, I will switch the PR to "ready to review" once I upload the notebook somewhere, but it is up to you how to figure out what's the next best course of action.

[asumagic]

Archive containing the streaming test jupyter notebook. Should be relatively straightforward to run (checkout this branch in a new environment, edit the paths, get the model directory from gdrive and place it in the proper results subdirectory for the conformer-transducer recipe). Outdated

The WIP code to actually do streaming inference lives there (though this PR includes a bunch of model methods to help with it), alongside some demo and test code.

[TParcollet]

Hey @asumagic please ping me when it's ready to review again :)

[asumagic]

Trying to fix fp16 training NaN loss issues so that I can retrain a model after fixing a mismatch in the new convolution code, and rebasing to be up to date with the core changes.

[asumagic]

Trying to fix fp16 training NaN loss issues so that I can retrain a model after fixing a mismatch in the new convolution code, and rebasing to be up to date with the core changes.

Status update:

• I gave up on trying to fix fp16 for now, I'm training using bf16 again on 4x A100 (on narval).
• I had to merge Add warnings and fix numba #2271 and Fix checkpointing #2267 (in a currently not pushed version) to fix issues related to numba bf16 handling and checkpointing.

Not sure when the training will be done especially as narval will go under maintenance tomorrow, but at least I finally got the ball rolling there.

[TParcollet]

Did you try removing the loss computation from the mixed training + the last softmax after the joiner network?

One thing to check as well, is the RNNT loss a sum or an average? If it's a sum, dynamic batching might be slightly less stable.

[asumagic]

Did you try removing the loss computation from the mixed training + the last softmax after the joiner network?

One thing to check as well, is the RNNT loss a sum or an average? If it's a sum, dynamic batching might be slightly less stable.

I think I did try forcing the loss computation to fp32 to no avail. I didn't try touching the joiner network/softmax though. For now I'm tempted to keep training with bf16 as I started one (.. once compute canada goes out of maintenance) and to try to investigate the fp16 problem after that so that the PR is ready ASAP.

The RNN-T loss uses a sum here.

[TParcollet]

The reason could be the sum here then. Could you please try to initialise the grad scaler with a much smaller value? Like 10 instead of 65536? Or even 2.0? The rational behind this is that torch initialise this value to a crazy high number that is ok when you have averaged losses -- see here. You can certainly also try to play with growth_factor=2.0, backoff_factor=0.5). Another good thing to do to see exactly what is happening and adapt accordingly could be to print the value of the scaling factor and see how it behaves, it may either vanish or explode due to the fact that we use a sum... @Adel-Moumen might be interested by this as well. But I am pretty sure that it comes from here.

[asumagic]

cc @TParcollet @Adel-Moumen: nevermind, a mean is actually used in the transducer loss itself, I was looking at the wrong thing (the joiner network).

I've tried reducing the initial gradscaler scale very significantly, but it did not help. Printing its state_dict to monitor what's going on shows that the scale is vanishing, so it seems to be endlessly trying to reduce scale.

I'll be retrying to force different parts to float32 to diagnose the issue, since it's possible I had done it wrong the last time I tried. I think disabling autocast with enabled=False might not actually suffice as the inputs may need to be converted back to float explicitly, see: https://pytorch.org/docs/stable/amp.html#autocasting

Looking into this now.

[TParcollet]

@asumagic right ... what about the rest of the recipe, is it ready? If so, I'll review one last time and try it.

[asumagic]

@asumagic right ... what about the rest of the recipe, is it ready? If so, I'll review one last time and try it.

The recipe itself should™ be ready, other than that fp16 is still not fixed. The new training (with bf16) should end tomorrow or on Monday. Once done, I will test accuracy and everything once again.

Notes and reminders:

• This PR depends on Sync before and after deleting #2268 for training with DDP. Rebased against unstable-v0.6, so this is now in.
• In TransformerASR, I still have to launch training with the custom_tgt_module line commented out. Because that field is not used when only encoding, DDP raises an error during backpropagation. To work around it, when the model is done training, I manually add it back to the model to avoid missing key errors. This issue was mentioned back with the other Conformer PR but we postponed a fix.
• The interfaces to conveniently do chunked streaming are not currently part of the PR. This PR only includes the recipe changes and the model changes. I'll update and reupload my test notebook when it's done.

The high-level interfaces shouldn't take too long to implement (comparatively) but I kind of wanted to make sure everything works before trying to design something that may also end up being reused for other streaming models (or for adapting non-streaming models to streaming contexts like whispercpp etc. do IIRC).

[TParcollet]

@asumagic as soon as you are confident in the results, ping me and I'll review. For the three notes: 1. Ok. 2. Who decided to postpone this issue, it sounds like something that would make this recipe unusable for non-expert people; 3. Ok, this is totally fine and I actually think that the interface should be another PR. Could you ping to an issue showing 2? I will fix this.

[asumagic]

The custom_tgt_module issue was brought up in the original PR starting with this comment: #1782 (comment)

[TParcollet]

wouldn't --find_unused_parameters solve the issue without commenting out the code? "The problem is that simply adding this check will not cut it as it will break every existing trained model, since trained models will include keys that will not exist (as the custom_tgt_module will not have been initialized)." -- right. This is a problem. Any idea from our checkpoint experts @Gastron and @pplantinga?

[asumagic]

wouldn't --find_unused_parameters solve the issue without commenting out the code?

No, all this does is to print which parameters were unused, but you still get an error. If you don't specify it, it does not tell you which parameter was the problem.

(I imagine the default just checks a counter, whereas this option also tracks and sends the list of parameter names over DDP, which is probably why it's not default).

[TParcollet]

@asumagic also, it looks like when streaming = True, the training speed slowdown seriously - on 8 GPUs it goes from 10min per epoch to 15-20. Is this expected? I guess it's due to the processing per chunk with all the for loops? Should we try to jit this? GPU util is also seems to be lower, once we have merged that, maybe we might want to have a look at optimising the code, if this is of interest to you.

[asumagic]

I'll try to profile it. It's possible that some of the list comprehensions inside of the convolution module are what harm performance, especially with longer utterances and smaller chunk sizes.
I might have underestimated their effect on performance, but I think most of that could be mitigated if it's that.

Should this be solved as part of this PR or another? I imagine the change would be not too extensive, but might take a bit to properly implement and test. nvm, didn't read it through, didn't realize you suggested doing this after.

[asumagic]

On a side note, I did try poking with torch.compile a few days ago and impressively enough, it succeeds compiling the entire compute_forward. However, the dynamic shape support is still too lacking as compilation crashes hard with it enabled (and it's kind of a requirement to avoid extremely slow constant recompiles), and the BS fails to compile with it (but not the GS). So right now it's a dead end, but it's promising.

[asumagic]

Yet again updated the notebook due to the renaming refactor:
streaming-tests.zip

[asumagic]

@TParcollet Ok, I think it should be all done now. There are some conversations I didn't resolve which were mostly answers to some remarks/questions.

With the very latest code, I have checked that:

• training runs for a few iterations without errors
• inference still works with the expected accuracy with the real streaming code
• the integration test recipe is actually able to output tokens after a much longer number of epochs than is defined in the code
• verifying that masking and streaming encoder code paths produce almost exactly the same results when chunking after feature extraction (the feature extraction may behave slightly differently in streaming, which does not cause issues)

[TParcollet]

LGTM. Awesome work @asumagic !