Add Numba CPU+CUDA support for RNNT Loss #1854

titu1994 · 2021-03-06T03:18:40Z

Changelog

Add Numba CUDA + CPU kernels for RNNT Loss ported from https://github.com/HawkAaron/warp-transducer
Add tests for CUDA + CPU kernels + reference numpy implementation.

Todo:

Add documentation for all kernels and wrapped modules
Add Typing information to RNNTLossNumba (Done by resolver)
Rename this module to avoid clash with original warprnnt_pytorch loss
Add WarpRNNT LICENSE header in all files that implement functionality as ported.
Fix imports inside __init__.py
Enable loss switching for RNNT / RNNTBPE models via config (done in next PR)

lgtm-com · 2021-03-06T03:37:13Z

This pull request introduces 1 alert when merging 9de7bc5 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T04:04:17Z

This pull request introduces 1 alert when merging 8f8a275 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T04:19:25Z

This pull request introduces 1 alert when merging e454c4c into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T05:12:59Z

This pull request introduces 1 alert when merging 4946314 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T06:25:05Z

This pull request introduces 1 alert when merging ea51c3b into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T06:42:28Z

This pull request introduces 1 alert when merging aea2b72 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T06:58:18Z

This pull request introduces 1 alert when merging bcb683a into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T07:57:29Z

This pull request introduces 1 alert when merging 5addd50 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T08:18:49Z

This pull request introduces 1 alert when merging af028e9 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-06T17:31:40Z

This pull request introduces 1 alert when merging 8edd7f1 into 74a1ed8 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-09T08:34:24Z

This pull request introduces 1 alert when merging 5bf8f3d into e8912ca - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-23T00:07:07Z

This pull request introduces 1 alert when merging b3f93ae into aaab317 - view on LGTM.com

new alerts:

1 for Unused local variable

lgtm-com · 2021-03-25T22:53:46Z

This pull request introduces 1 alert when merging 9cdc539 into d623109 - view on LGTM.com

new alerts:

1 for Unused local variable

nemo/collections/asr/losses/rnnt.py

nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py

nemo/collections/asr/parts/numba/__init__.py

VahidooX · 2021-03-25T23:31:57Z

nemo/collections/asr/parts/numba/rnnt_loss/__init__.py

+#
+# Copyright 2018-2019, Mingkun Huang
+#
+# Licensed under the Apache License, Version 2.0 (the "License");


is it necessary to have this license here?

Yes, here its needed since the 2 functions and the class is ported from that codebase.

But this is the init file which just contains the imports.

nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py

reinstall.sh

tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py

nemo/collections/asr/parts/numba/__init__.py

Signed-off-by: smajumdar <titu1994@gmail.com>

lgtm-com · 2021-03-26T00:47:48Z

This pull request introduces 1 alert when merging 19215a1 into 480bbc3 - view on LGTM.com

new alerts:

1 for Unused local variable

Signed-off-by: smajumdar <titu1994@gmail.com>

* Initial code merge Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Update reduce kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add tests and remove log messages Signed-off-by: smajumdar <titu1994@gmail.com> * Improve documentation Signed-off-by: smajumdar <titu1994@gmail.com> * Add alpha beta kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Document warp reduce.py Signed-off-by: smajumdar <titu1994@gmail.com> * Style fixes Signed-off-by: smajumdar <titu1994@gmail.com> * Add NeMo licenses Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Force update Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel gradient Signed-off-by: smajumdar <titu1994@gmail.com> * Update docs for cpu and gpu rnnt kernels + helpers Signed-off-by: smajumdar <titu1994@gmail.com> * Correct backward for numpy rnnt Signed-off-by: smajumdar <titu1994@gmail.com> * Update name of RNNT loss Signed-off-by: smajumdar <titu1994@gmail.com> * Formatting Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss License Signed-off-by: smajumdar <titu1994@gmail.com> * Add __init__.py correctly and improve rnnt loss resolution Signed-off-by: smajumdar <titu1994@gmail.com> * Address reviewer comments Signed-off-by: smajumdar <titu1994@gmail.com> * Address comments Signed-off-by: smajumdar <titu1994@gmail.com> Signed-off-by: Virginia Adams <vadams@nvidia.com>

* Initial code merge Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Update reduce kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add tests and remove log messages Signed-off-by: smajumdar <titu1994@gmail.com> * Improve documentation Signed-off-by: smajumdar <titu1994@gmail.com> * Add alpha beta kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Document warp reduce.py Signed-off-by: smajumdar <titu1994@gmail.com> * Style fixes Signed-off-by: smajumdar <titu1994@gmail.com> * Add NeMo licenses Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Force update Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel gradient Signed-off-by: smajumdar <titu1994@gmail.com> * Update docs for cpu and gpu rnnt kernels + helpers Signed-off-by: smajumdar <titu1994@gmail.com> * Correct backward for numpy rnnt Signed-off-by: smajumdar <titu1994@gmail.com> * Update name of RNNT loss Signed-off-by: smajumdar <titu1994@gmail.com> * Formatting Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss License Signed-off-by: smajumdar <titu1994@gmail.com> * Add __init__.py correctly and improve rnnt loss resolution Signed-off-by: smajumdar <titu1994@gmail.com> * Address reviewer comments Signed-off-by: smajumdar <titu1994@gmail.com> * Address comments Signed-off-by: smajumdar <titu1994@gmail.com> Signed-off-by: Micha Livne <mlivne@nvidia.com>

* Initial code merge Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss tests Signed-off-by: smajumdar <titu1994@gmail.com> * Update reduce kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Add tests and remove log messages Signed-off-by: smajumdar <titu1994@gmail.com> * Improve documentation Signed-off-by: smajumdar <titu1994@gmail.com> * Add alpha beta kernel tests Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Remove unused imports Signed-off-by: smajumdar <titu1994@gmail.com> * Document warp reduce.py Signed-off-by: smajumdar <titu1994@gmail.com> * Style fixes Signed-off-by: smajumdar <titu1994@gmail.com> * Add NeMo licenses Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Add Numba to dependencies Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Use safetynet api in case numba version is too old Signed-off-by: smajumdar <titu1994@gmail.com> * Force update Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel Signed-off-by: smajumdar <titu1994@gmail.com> * Add docs for gpu rnnt kernel gradient Signed-off-by: smajumdar <titu1994@gmail.com> * Update docs for cpu and gpu rnnt kernels + helpers Signed-off-by: smajumdar <titu1994@gmail.com> * Correct backward for numpy rnnt Signed-off-by: smajumdar <titu1994@gmail.com> * Update name of RNNT loss Signed-off-by: smajumdar <titu1994@gmail.com> * Formatting Signed-off-by: smajumdar <titu1994@gmail.com> * Add RNNT Loss License Signed-off-by: smajumdar <titu1994@gmail.com> * Add __init__.py correctly and improve rnnt loss resolution Signed-off-by: smajumdar <titu1994@gmail.com> * Address reviewer comments Signed-off-by: smajumdar <titu1994@gmail.com> * Address comments Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 marked this pull request as draft March 6, 2021 03:19

titu1994 force-pushed the add_rnnt_numba_loss branch from 487e12f to 6d28761 Compare March 9, 2021 08:17

titu1994 force-pushed the add_rnnt_numba_loss branch from 5bf8f3d to b3f93ae Compare March 22, 2021 23:56

titu1994 force-pushed the add_rnnt_numba_loss branch from b3f93ae to 9cdc539 Compare March 25, 2021 22:43

titu1994 marked this pull request as ready for review March 25, 2021 22:53

titu1994 requested a review from VahidooX March 25, 2021 23:11

VahidooX requested changes Mar 25, 2021

View reviewed changes

titu1994 force-pushed the add_rnnt_numba_loss branch from 9cdc539 to 19215a1 Compare March 26, 2021 00:35

titu1994 added 9 commits March 25, 2021 17:35

Initial code merge

ab0206f

Signed-off-by: smajumdar <titu1994@gmail.com>

Add RNNT Loss tests

2ba4b58

Signed-off-by: smajumdar <titu1994@gmail.com>

Add RNNT Loss tests

4325c6a

Signed-off-by: smajumdar <titu1994@gmail.com>

Update reduce kernel tests

60b029b

Signed-off-by: smajumdar <titu1994@gmail.com>

Add tests and remove log messages

3a5e7bd

Signed-off-by: smajumdar <titu1994@gmail.com>

Improve documentation

622cba1

Signed-off-by: smajumdar <titu1994@gmail.com>

Add alpha beta kernel tests

25c8a94

Signed-off-by: smajumdar <titu1994@gmail.com>

Remove unused imports

d42766b

Signed-off-by: smajumdar <titu1994@gmail.com>

Remove unused imports

2a0eb3a

Signed-off-by: smajumdar <titu1994@gmail.com>

titu1994 added 19 commits March 25, 2021 17:35

Style fixes

c184f6e

Signed-off-by: smajumdar <titu1994@gmail.com>

Add NeMo licenses

8f67e18

Signed-off-by: smajumdar <titu1994@gmail.com>

Add Numba to dependencies

a67d131

Signed-off-by: smajumdar <titu1994@gmail.com>

Add Numba to dependencies

8289fb3

Signed-off-by: smajumdar <titu1994@gmail.com>

Add Numba to dependencies

f8b5f51

Signed-off-by: smajumdar <titu1994@gmail.com>

Use safetynet api in case numba version is too old

0497751

Signed-off-by: smajumdar <titu1994@gmail.com>

Use safetynet api in case numba version is too old

e57bb99

Signed-off-by: smajumdar <titu1994@gmail.com>

Use safetynet api in case numba version is too old

79e1a87

Signed-off-by: smajumdar <titu1994@gmail.com>

Use safetynet api in case numba version is too old

2679681

Signed-off-by: smajumdar <titu1994@gmail.com>

Force update

14dbc99

Signed-off-by: smajumdar <titu1994@gmail.com>

Add docs for gpu rnnt kernel

69507c3

Signed-off-by: smajumdar <titu1994@gmail.com>

Add docs for gpu rnnt kernel gradient

9601eb2

Signed-off-by: smajumdar <titu1994@gmail.com>

Update docs for cpu and gpu rnnt kernels + helpers

22aa34e

Signed-off-by: smajumdar <titu1994@gmail.com>

Correct backward for numpy rnnt

47fad42

Signed-off-by: smajumdar <titu1994@gmail.com>

Update name of RNNT loss

c525dc5

Signed-off-by: smajumdar <titu1994@gmail.com>

Formatting

d37685f

Signed-off-by: smajumdar <titu1994@gmail.com>

Add RNNT Loss License

b62048a

Signed-off-by: smajumdar <titu1994@gmail.com>

Add __init__.py correctly and improve rnnt loss resolution

0115bbc

Signed-off-by: smajumdar <titu1994@gmail.com>

Address reviewer comments

19215a1

Signed-off-by: smajumdar <titu1994@gmail.com>

VahidooX approved these changes Mar 26, 2021

View reviewed changes

titu1994 added 3 commits March 25, 2021 23:25

Address comments

a3e0fff

Signed-off-by: smajumdar <titu1994@gmail.com>

Merge branch 'main' into add_rnnt_numba_loss

f9d4767

Merge branch 'main' into add_rnnt_numba_loss

e4e77b9

titu1994 merged commit bb01dbb into NVIDIA:main Mar 29, 2021

titu1994 deleted the add_rnnt_numba_loss branch March 29, 2021 05:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Numba CPU+CUDA support for RNNT Loss #1854

Add Numba CPU+CUDA support for RNNT Loss #1854

titu1994 commented Mar 6, 2021 •

edited

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 9, 2021

lgtm-com bot commented Mar 23, 2021

lgtm-com bot commented Mar 25, 2021

VahidooX Mar 25, 2021

titu1994 Mar 26, 2021

VahidooX Mar 26, 2021

titu1994 Mar 26, 2021

lgtm-com bot commented Mar 26, 2021

Add Numba CPU+CUDA support for RNNT Loss #1854

Add Numba CPU+CUDA support for RNNT Loss #1854

Conversation

titu1994 commented Mar 6, 2021 • edited

Changelog

Todo:

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 6, 2021

lgtm-com bot commented Mar 9, 2021

lgtm-com bot commented Mar 23, 2021

lgtm-com bot commented Mar 25, 2021

VahidooX Mar 25, 2021

Choose a reason for hiding this comment

titu1994 Mar 26, 2021

Choose a reason for hiding this comment

VahidooX Mar 26, 2021

Choose a reason for hiding this comment

titu1994 Mar 26, 2021

Choose a reason for hiding this comment

lgtm-com bot commented Mar 26, 2021

titu1994 commented Mar 6, 2021 •

edited