Hyena Operator #9264

guyjacob · 2024-05-21T14:12:06Z

What does this PR do ?

Add Hyena operator (https://arxiv.org/abs/2302.10866)

Collection: NLP

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Guy Jacob <guyj@nvidia.com>

* Add missing input LayerNorm to spec (in the default attention spec it's fused with the projection Linear layer, so not explicitly defined) * Shape conversion at start and end of Hyena forward Signed-off-by: Guy Jacob <guyj@nvidia.com>

Signed-off-by: Guy Jacob <guyj@nvidia.com>

See: HazyResearch/safari#26 (comment) Signed-off-by: Guy Jacob <guyj@nvidia.com>

(torch.fft doesn't support bf16) Signed-off-by: Guy Jacob <guyj@nvidia.com>

Signed-off-by: Guy Jacob <guyj@nvidia.com>

(made redundant by the default inmplementation in Megatron) Signed-off-by: Guy Jacob <guyj@nvidia.com>

Signed-off-by: Guy Jacob <guyj@nvidia.com>

* Refactor in spirit of implementation in MAD-Lab repo: https://github.com/athms/mad-lab/blob/main/mad/model/layers/hyena.py Signed-off-by: Guy Jacob <guyj@nvidia.com>

Signed-off-by: Guy Jacob <guyj@nvidia.com>

examples/nlp/language_modeling/conf/megatron_gpt_config_hyena.yaml

nemo/collections/nlp/modules/common/hyena/fftconv_wrapper.py

nemo/collections/nlp/modules/common/hyena/hyena.py

akoumpa

LGTM, just a few minor comments, if you have to address these that would be great, thanks!

* Clearer names + more documentation for config params * Clearer README * Check seq len < 8K with safari-fftconv * Avoid 0*bias op during forward Signed-off-by: Guy Jacob <guyj@nvidia.com>

tests/collections/nlp/test_hyena_operator.py

+from nemo.collections.nlp.parts.utils_funcs import torch_dtype_from_precision
+
+try:
+    import fftconv


tests/collections/nlp/test_hyena_operator.py

+    HAVE_FFTCONV = False
+
+try:
+    import flashfftconv


tests/collections/nlp/test_hyena_operator.py

+    HAVE_FLASHFFTCONV = False
+
+try:
+    import causal_conv1d


Signed-off-by: Guy Jacob <guyj@nvidia.com>

* Initial reference code commit, unchanged Signed-off-by: Guy Jacob <guyj@nvidia.com> * Hyena code changes for NeMO compatibility Signed-off-by: Guy Jacob <guyj@nvidia.com> * MCore spec override functionality + example config w. hyena Signed-off-by: Guy Jacob <guyj@nvidia.com> * Additional changes - now working on char-level TinyShakespeare * Add missing input LayerNorm to spec (in the default attention spec it's fused with the projection Linear layer, so not explicitly defined) * Shape conversion at start and end of Hyena forward Signed-off-by: Guy Jacob <guyj@nvidia.com> * Add fftconv cuda impl from safari Signed-off-by: Guy Jacob <guyj@nvidia.com> * Workaround for shape error in fftconv See: HazyResearch/safari#26 (comment) Signed-off-by: Guy Jacob <guyj@nvidia.com> * Explicitly convert kernel to FP32 (torch.fft doesn't support bf16) Signed-off-by: Guy Jacob <guyj@nvidia.com> * Working run configs Signed-off-by: Guy Jacob <guyj@nvidia.com> * Remove sharded_state_dict from HyenaOperator (made redundant by the default inmplementation in Megatron) Signed-off-by: Guy Jacob <guyj@nvidia.com> * Update configs Signed-off-by: Guy Jacob <guyj@nvidia.com> * Testing TE Linear classes in HyenaOperator Signed-off-by: Guy Jacob <guyj@nvidia.com> * Revert to FusedDense for in/out projections after merging with 24.01.01 Signed-off-by: Guy Jacob <guyj@nvidia.com> * Fix bug (use fused LNorm+Linear), bring back TE layers Signed-off-by: Guy Jacob <guyj@nvidia.com> * Configs rename + cleanup Signed-off-by: Guy Jacob <guyj@nvidia.com> * FlashFFTConv, Multi-head, some cleanup Signed-off-by: Guy Jacob <guyj@nvidia.com> * Bug fix - init FlashFFTConv with 2*seq_len Signed-off-by: Guy Jacob <guyj@nvidia.com> * ModuleSpec + replace nn.Conv1d with causal_conv1d Signed-off-by: Guy Jacob <guyj@nvidia.com> * Remove unneeded arguments Signed-off-by: Guy Jacob <guyj@nvidia.com> * More cleanup, remove fftconv ref functions Signed-off-by: Guy Jacob <guyj@nvidia.com> * Refactor HyenaFilter + more cleanup * Refactor in spirit of implementation in MAD-Lab repo: https://github.com/athms/mad-lab/blob/main/mad/model/layers/hyena.py Signed-off-by: Guy Jacob <guyj@nvidia.com> * Add missing attributions Signed-off-by: Guy Jacob <guyj@nvidia.com> * Remove fftconv sources Signed-off-by: Guy Jacob <guyj@nvidia.com> * Bug fixes Signed-off-by: Guy Jacob <guyj@nvidia.com> * Remove d_model from external API, take from TransformerConfig Signed-off-by: Guy Jacob <guyj@nvidia.com> * cleanup config Signed-off-by: Guy Jacob <guyj@nvidia.com> * Remove spec override logic (possibly push separately) Signed-off-by: Guy Jacob <guyj@nvidia.com> * Add tests Signed-off-by: Guy Jacob <guyj@nvidia.com> * Keep only megatron_gpt_config_hyena (w. 153m parameters) Signed-off-by: Guy Jacob <guyj@nvidia.com> * Black + isort formatting changes Signed-off-by: Guy Jacob <guyj@nvidia.com> * Fixes following PR review * Clearer names + more documentation for config params * Clearer README * Check seq len < 8K with safari-fftconv * Avoid 0*bias op during forward Signed-off-by: Guy Jacob <guyj@nvidia.com> * Fix tests following param name changes Signed-off-by: Guy Jacob <guyj@nvidia.com> --------- Signed-off-by: Guy Jacob <guyj@nvidia.com>

guyjacob added 28 commits May 16, 2024 18:32

Initial reference code commit, unchanged

5a64cb8

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Hyena code changes for NeMO compatibility

ee4c809

Signed-off-by: Guy Jacob <guyj@nvidia.com>

MCore spec override functionality + example config w. hyena

cf3aef9

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Add fftconv cuda impl from safari

104d454

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Workaround for shape error in fftconv

dcdfeba

See: HazyResearch/safari#26 (comment) Signed-off-by: Guy Jacob <guyj@nvidia.com>

Explicitly convert kernel to FP32

86eba0c

(torch.fft doesn't support bf16) Signed-off-by: Guy Jacob <guyj@nvidia.com>

Working run configs

eeceb77

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Remove sharded_state_dict from HyenaOperator

7c4e461

(made redundant by the default inmplementation in Megatron) Signed-off-by: Guy Jacob <guyj@nvidia.com>

Update configs

51aaed1

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Testing TE Linear classes in HyenaOperator

8a251b6

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Revert to FusedDense for in/out projections after merging with 24.01.01

4fefb04

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Fix bug (use fused LNorm+Linear), bring back TE layers

ab86a68

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Configs rename + cleanup

ae5fb19

Signed-off-by: Guy Jacob <guyj@nvidia.com>

FlashFFTConv, Multi-head, some cleanup

1d74c74

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Bug fix - init FlashFFTConv with 2*seq_len

46b1ddf

Signed-off-by: Guy Jacob <guyj@nvidia.com>

ModuleSpec + replace nn.Conv1d with causal_conv1d

129a0c1

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Remove unneeded arguments

2fbac28

Signed-off-by: Guy Jacob <guyj@nvidia.com>

More cleanup, remove fftconv ref functions

86d53d4

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Refactor HyenaFilter + more cleanup

56bfc03

* Refactor in spirit of implementation in MAD-Lab repo: https://github.com/athms/mad-lab/blob/main/mad/model/layers/hyena.py Signed-off-by: Guy Jacob <guyj@nvidia.com>

Add missing attributions

cb73eab

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Remove fftconv sources

392e934

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Bug fixes

9796463

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Remove d_model from external API, take from TransformerConfig

52aa6b9

Signed-off-by: Guy Jacob <guyj@nvidia.com>

cleanup config

68d821f

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Remove spec override logic (possibly push separately)

abbbbfe

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Add tests

19a51ff

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Keep only megatron_gpt_config_hyena (w. 153m parameters)

93ac1a7

Signed-off-by: Guy Jacob <guyj@nvidia.com>

github-actions bot added the NLP label May 21, 2024

Merge branch 'main' into guyj/hyena

e79167e

Black + isort formatting changes

a1bcad7

Signed-off-by: Guy Jacob <guyj@nvidia.com>

guyjacob force-pushed the guyj/hyena branch from f00e84c to a1bcad7 Compare May 23, 2024 14:33

ericharper requested a review from akoumpa May 28, 2024 17:00

akoumpa added the Run CICD label May 28, 2024

akoumpa reviewed May 28, 2024

View reviewed changes

examples/nlp/language_modeling/conf/megatron_gpt_config_hyena.yaml Outdated Show resolved Hide resolved

akoumpa reviewed May 28, 2024

View reviewed changes

examples/nlp/language_modeling/conf/megatron_gpt_config_hyena.yaml Show resolved Hide resolved

akoumpa reviewed May 28, 2024

View reviewed changes

nemo/collections/nlp/modules/common/hyena/fftconv_wrapper.py Outdated Show resolved Hide resolved

akoumpa reviewed May 28, 2024

View reviewed changes

nemo/collections/nlp/modules/common/hyena/hyena.py Outdated Show resolved Hide resolved

akoumpa reviewed May 28, 2024

View reviewed changes

nemo/collections/nlp/modules/common/hyena/hyena.py Outdated Show resolved Hide resolved

akoumpa previously approved these changes May 28, 2024

View reviewed changes

guyjacob dismissed akoumpa’s stale review via 60edb10 June 2, 2024 16:04

Fixes following PR review

bce73db

* Clearer names + more documentation for config params * Clearer README * Check seq len < 8K with safari-fftconv * Avoid 0*bias op during forward Signed-off-by: Guy Jacob <guyj@nvidia.com>

guyjacob force-pushed the guyj/hyena branch from 60edb10 to bce73db Compare June 2, 2024 16:05

akoumpa previously approved these changes Jun 9, 2024

View reviewed changes

github-advanced-security bot found potential problems Jun 9, 2024

View reviewed changes

guyjacob added 2 commits June 10, 2024 15:08

Merge branch 'main' into guyj/hyena

fb35209

Signed-off-by: Guy Jacob <guyj@nvidia.com>

Fix tests following param name changes

af99952

Signed-off-by: Guy Jacob <guyj@nvidia.com>

guyjacob dismissed akoumpa’s stale review via af99952 June 10, 2024 12:32

akoumpa approved these changes Jun 10, 2024

View reviewed changes

ericharper added Run CICD and removed Run CICD labels Jun 13, 2024

ericharper merged commit a6a0aee into NVIDIA:main Jun 13, 2024
109 checks passed

ko3n1g mentioned this pull request Jul 18, 2024

Release 2.0.0rc1 #9786

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Hyena Operator #9264

Hyena Operator #9264

guyjacob commented May 21, 2024

akoumpa left a comment

Hyena Operator #9264

Hyena Operator #9264

Conversation

guyjacob commented May 21, 2024

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

akoumpa left a comment

Choose a reason for hiding this comment