Alit/mamba by JRD971000 · Pull Request #9696 · NVIDIA-NeMo/NeMo

JRD971000 · 2024-07-11T15:14:03Z

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Collection: [Note which collection this PR will affect]

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: JRD971000 <JRD971000@users.noreply.github.com>

…/mamba

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

…/mamba

Signed-off-by: arendu <adithya.r@gmail.com>

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

…/mamba

JRD971000 · 2024-07-11T15:53:42Z

examples/nlp/language_modeling/mamba_change_num_partition.py

+
+        # print("Combined tp model: {}".format(combined_tp_model.keys()))
+
+        for ii, (key, original_tensor) in enumerate(combined_tp_model.items()):


examples/nlp/language_modeling/megatron_mamba_eval.py

+    # Have to turn off activations_checkpoint_method for inference
+    try:
+        model.model.language_model.encoder.activations_checkpoint_method = None
+    except AttributeError:


examples/nlp/language_modeling/mamba_change_num_partition.py

+            try:
+                layer_num = int(re.findall(r'\d+', key)[0])
+                new_key = key.replace(str(layer_num), str(layer_num), 1)
+            except:


examples/nlp/language_modeling/megatron_mamba_eval.py

+        with open(prompts_jsonl, 'rt') as fp:
+            try:
+                prompts += list(map(json.loads, map(str.rstrip, fp)))
+            except:


examples/nlp/language_modeling/megatron_mamba_eval.py

+try:
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True


examples/nlp/language_modeling/megatron_mamba_eval.py

+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False


scripts/checkpoint_converters/convert_mamba2_pyt_to_nemo.py

+            new_state_dict["model." + key] = value
+
+        # Tokenizer settings
+        tokenizer_library = 'megatron'


scripts/checkpoint_converters/convert_mamba2_pyt_to_nemo.py

+
+        # Tokenizer settings
+        tokenizer_library = 'megatron'
+        tokenizer_type = 'GPTSentencePieceTokenizer'


scripts/checkpoint_converters/convert_mamba2_pyt_to_nemo.py

+        # Tokenizer settings
+        tokenizer_library = 'megatron'
+        tokenizer_type = 'GPTSentencePieceTokenizer'
+        tokenizer_model = args.tokenizer_model_dir


Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

examples/nlp/language_modeling/mamba_change_num_partition.py

+            # copy model
+            combined_tp_model[key] = original_tensor
+
+        for _, (local_key, original_tensor) in enumerate(combined_tp_model.items()):


examples/nlp/language_modeling/mamba_change_num_partition.py

+            try:
+                layer_num = int(re.findall(r'\d+', local_key)[0])
+                new_key = local_key.replace(str(layer_num), str(layer_num), 1)
+            except:


Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

examples/nlp/language_modeling/mamba_change_num_partition.py

+            try:
+                layer_num = int(re.findall(r'\d+', local_key)[0])
+                new_key = local_key.replace(str(layer_num), str(layer_num), 1)
+            except:


…/mamba

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com> Signed-off-by: Tugrul Konuk <ertkonuk@gmail.com>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com> Signed-off-by: Malay Nagda <malayn@malayn-mlt.client.nvidia.com>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com> Signed-off-by: tonyjie <jl4257@cornell.edu>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com> Signed-off-by: Hainan Xu <hainanx@nvidia.com>

* adding mamba support * fix import mixins * rm convert jamba * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * more cleanups * use GPT text gen * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * fixing gbs in TP convetor * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * add reqs * add tutorial * minor fix to tutorial * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * moving finetuning files Signed-off-by: arendu <adithya.r@gmail.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * address comments * add mamba dependancies * add mcore tag * modify dockerfile ci * modify dockerfile ci * fix TP>1 to TP1 * add inference, update based on latest mcore commits * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * minor fix * Apply isort and black reformatting Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> * bug fix, tutorial update --------- Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com> Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: Ali Taghibakhshi <ataghibakhsh@login-eos01.eos.clusters.nvidia.com> Co-authored-by: JRD971000 <JRD971000@users.noreply.github.com> Co-authored-by: arendu <adithya.r@gmail.com>

Ali Taghibakhshi and others added 30 commits July 1, 2024 07:39

adding mamba support

8a26848

fix import mixins

73d7c4c

rm convert jamba

66886b5

Apply isort and black reformatting

f9e2066

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

more cleanups

96ab05c

Merge branch 'alit/mamba' of https://github.com/NVIDIA/NeMo into alit…

f24cd69

…/mamba

use GPT text gen

2e74b64

Apply isort and black reformatting

05c377a

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

fixing gbs in TP convetor

59f176a

resolve merge conflicts

74a30de

Apply isort and black reformatting

dfc24e2

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

add reqs

7edd5cc

Merge branch 'alit/mamba' of https://github.com/NVIDIA/NeMo into alit…

3eee1c7

…/mamba

add tutorial

c0afdc4

minor fix to tutorial

6097379

moving finetuning files

8e7aea0

Signed-off-by: arendu <adithya.r@gmail.com>

moving finetuning files

1db8269

Signed-off-by: arendu <adithya.r@gmail.com>

address comments

0f326d6

Apply isort and black reformatting

da7461a

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

address comments

022622e

Apply isort and black reformatting

7b67568

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

Merge branch 'main' into alit/mamba

7dce2bf

address comments

0d5cc37

add mamba dependancies

2cf9040

Merge branch 'main' into alit/mamba

0353eb9

add mcore tag

23a2d20

merge main

a9a24b7

Merge branch 'alit/mamba' of https://github.com/NVIDIA/NeMo into alit…

53792ab

…/mamba

Merge branch 'main' into alit/mamba

2b97d0b

modify dockerfile ci

0747052

JRD971000 added the Run CICD label Jul 11, 2024

github-advanced-security bot found potential problems Jul 11, 2024

View reviewed changes

minor fix

1294e3a

JRD971000 added Run CICD and removed Run CICD labels Jul 11, 2024

Apply isort and black reformatting

260bdfb

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

github-advanced-security bot found potential problems Jul 11, 2024

View reviewed changes

Ali Taghibakhshi and others added 4 commits July 11, 2024 09:57

minor fix

a1bb04f

resolve conflict:

73d02a4

Merge branch 'main' into alit/mamba

841b1e1

Apply isort and black reformatting

71cfb8b

Signed-off-by: JRD971000 <JRD971000@users.noreply.github.com>

JRD971000 added Run CICD and removed Run CICD labels Jul 11, 2024

github-advanced-security bot found potential problems Jul 11, 2024

View reviewed changes

yaoyu-33 previously approved these changes Jul 11, 2024

View reviewed changes

Ali Taghibakhshi added 2 commits July 11, 2024 10:34

bug fix, tutorial update

a583dad

Merge branch 'alit/mamba' of https://github.com/NVIDIA/NeMo into alit…

41f6229

…/mamba

JRD971000 dismissed yaoyu-33’s stale review via 41f6229 July 11, 2024 17:35

JRD971000 added Run CICD and removed Run CICD labels Jul 11, 2024

yaoyu-33 approved these changes Jul 11, 2024

View reviewed changes

JRD971000 merged commit 3482dc1 into main Jul 11, 2024

JRD971000 deleted the alit/mamba branch July 11, 2024 19:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Alit/mamba#9696

Alit/mamba#9696
JRD971000 merged 45 commits intomainfrom
alit/mamba

JRD971000 commented Jul 11, 2024

Uh oh!

Check notice

JRD971000 Jul 11, 2024

Uh oh!

Check notice

Check notice

Check notice

Check notice

Check notice

Check warning

Check warning

Check warning

Check notice

Check notice

Check notice

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		# print("Combined tp model: {}".format(combined_tp_model.keys()))

		for ii, (key, original_tensor) in enumerate(combined_tp_model.items()):


		except (ImportError, ModuleNotFoundError):

		HAVE_MEGATRON_CORE = False

Conversation

JRD971000 commented Jul 11, 2024

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

Check notice

JRD971000 Jul 11, 2024

Choose a reason for hiding this comment

Uh oh!

Check notice

Check notice

Check notice

Check notice

Check notice

Check warning

Check warning

Check warning

Check notice

Check notice

Check notice

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants