[NeMo-UX] Adding megatron strategy by marcromeyn · Pull Request #8995 · NVIDIA-NeMo/NeMo

marcromeyn · 2024-04-22T11:52:46Z

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

Jenkins CI

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

There's no need to comment jenkins on the PR to trigger Jenkins CI.
The GitHub Actions CI will run automatically when the PR is opened.
To run CI on an untrusted fork, a NeMo user with write access must 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: Marc Romeyn <marcromeyn@gmail.com>

for more information, see https://pre-commit.ci Signed-off-by: Marc Romeyn <marcromeyn@gmail.com>

nemo/lightning/_strategy_lib.py

+        except Exception:
+            NVIDIA_TORCH_MAJOR = 0
+        try:
+            NVIDIA_TORCH_MINOR = int(nvidia_torch_version.split(".")[1])


nemo/lightning/_strategy_lib.py

+        try:
+            NVIDIA_TORCH_MINOR = int(nvidia_torch_version.split(".")[1])
+        except Exception:
+            NVIDIA_TORCH_MINOR = 0


nemo/lightning/pytorch/strategies.py

+        # if self.trainer.state.fn == TrainerFn.FITTING:
+        #     if self.lightning_module.optimizers(use_pl_optimizer=False):
+        #         sharded_state_dict["optimizer_states"] = [self.optimizer_sharded_state_dict()]


nemo/lightning/_strategy_lib.py

+
+class SharedStateDictProtocol(Protocol):
+    def sharded_state_dict(self, prefix=""):
+        ...


Signed-off-by: Chen Cui <chcui@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Chen Cui <chcui@nvidia.com>

for more information, see https://pre-commit.ci

tests/lightning/test_strategy_lib.py

@@ -0,0 +1,211 @@
+from unittest.mock import ANY, MagicMock, patch


* Adding MegatronParallel Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * Move over _strategy_liMegatronCheckpointIO Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * add strategy lib test Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused test and add comment Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> Signed-off-by: Chen Cui <chcui@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com>

* Adding MegatronParallel Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * Move over _strategy_liMegatronCheckpointIO Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * add strategy lib test Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused test and add comment Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> Signed-off-by: Chen Cui <chcui@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com> Signed-off-by: Ao Tang <aot@nvidia.com>

* Adding MegatronParallel Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * Move over _strategy_liMegatronCheckpointIO Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> * add strategy lib test Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused test and add comment Signed-off-by: Chen Cui <chcui@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Marc Romeyn <marcromeyn@gmail.com> Signed-off-by: Chen Cui <chcui@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com>

marcromeyn requested a review from cuichenx April 22, 2024 11:52

marcromeyn and others added 3 commits April 22, 2024 13:56

Adding MegatronParallel

f4f885b

Signed-off-by: Marc Romeyn <marcromeyn@gmail.com>

Move over _strategy_liMegatronCheckpointIO

23f302e

Signed-off-by: Marc Romeyn <marcromeyn@gmail.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

ec5b8c9

for more information, see https://pre-commit.ci Signed-off-by: Marc Romeyn <marcromeyn@gmail.com>

marcromeyn force-pushed the nemo-ux/megatron-strategy branch from d111eb4 to ec5b8c9 Compare April 22, 2024 11:56

github-advanced-security bot found potential problems Apr 22, 2024

View reviewed changes

cuichenx and others added 5 commits April 22, 2024 13:45

Merge branch 'main' into nemo-ux/megatron-strategy

1e0bbe2

add strategy lib test

1403874

Signed-off-by: Chen Cui <chcui@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

2da8a83

for more information, see https://pre-commit.ci

remove unused test and add comment

107e870

Signed-off-by: Chen Cui <chcui@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

56e7707

for more information, see https://pre-commit.ci

github-advanced-security bot found potential problems Apr 22, 2024

View reviewed changes

tests/lightning/test_strategy_lib.py

@@ -0,0 +1,211 @@

from unittest.mock import ANY, MagicMock, patch

Check notice

Code scanning / CodeQL

Unused import

Import of 'ANY' is not used.

cuichenx approved these changes Apr 23, 2024

View reviewed changes

cuichenx merged commit a6424e9 into main Apr 23, 2024

cuichenx deleted the nemo-ux/megatron-strategy branch April 23, 2024 06:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[NeMo-UX] Adding megatron strategy#8995

[NeMo-UX] Adding megatron strategy#8995
cuichenx merged 8 commits intomainfrom
nemo-ux/megatron-strategy

marcromeyn commented Apr 22, 2024

Uh oh!

Check notice

Check notice

Check notice

Check notice

Check notice

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,211 @@
		from unittest.mock import ANY, MagicMock, patch

Conversation

marcromeyn commented Apr 22, 2024

What does this PR do ?

Changelog

Usage

Jenkins CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

Check notice

Check notice

Check notice

Check notice

Check notice

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants