Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add mcore spec for full TE TransformerLayer #8316

Closed
wants to merge 17 commits into from
Closed

Add mcore spec for full TE TransformerLayer #8316

wants to merge 17 commits into from

Conversation

jbaczek
Copy link
Collaborator

@jbaczek jbaczek commented Feb 2, 2024

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

To run Jenkins, a NeMo User with write access must comment jenkins on the PR.

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)

@github-actions github-actions bot added core Changes to NeMo Core NLP CI labels Feb 2, 2024
jiemingz and others added 14 commits February 2, 2024 20:06
@jiemingz @jbaczek
add first_val_step for mcore schedules
5eaab1f 
Signed-off-by: jiemingz <jiemingz@nvidia.com>
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jiemingz @jbaczek
fix if non fp8
2e52799 
Signed-off-by: jiemingz <jiemingz@nvidia.com>
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@pre-commit-ci @jbaczek
[pre-commit.ci] auto fixes from pre-commit.com hooks
3f69086 
for more information, see https://pre-commit.ci

Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Revert "context parallelism (NVIDIA#7739)"
8f74c8c 
This reverts commit 58d6bce.

Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Add ModuleSpec changes
8d9a60a 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Add is_first_microbatch
f93b01e 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Mostly remove manual knobs
7f952c5 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Remove unused use_emha flag, remove debugging prints
4fc6089 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Copy over AutocastTransformerLayer, remove direct support for TE Tran…
7b80208 
…sformerLayer

Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Fix checkpoint loading
e0724f7 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@sanandaraj5597 @layalir @jbaczek
Fixed the tp overlap switch (NVIDIA#8195)
22ab1ed 
Signed-off-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: Selvaraj Anandaraj <selvaraja@login-eos01.eos.clusters.nvidia.com>
Co-authored-by: Layali R <31741533+layalir@users.noreply.github.com>
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
remove custom checkpoint code
3f25c6d 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Comment out mapping. We should add it together with config update
3b1cb2d 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@jbaczek
Add mcore spec test
d56297d 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@ShriyaPalsamudram
Copy link
Collaborator

jenkins

ShriyaPalsamudram
ShriyaPalsamudram reviewed Feb 2, 2024
View reviewed changes
nemo/collections/nlp/data/language_modeling/megatron/dataset_utils.py Outdated
@@ -1330,7 +1330,7 @@ def get_samples_mapping(
)
torch.distributed.barrier()
counts = torch.cuda.LongTensor([1])
torch.distributed.all_reduce(counts, group=parallel_state.get_data_parallel_group(with_context_parallel=True))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please separate context_parallel changes from the full TE spec changes to make review easier.

minitu
minitu reviewed Feb 2, 2024
View reviewed changes
nemo/collections/nlp/models/language_modeling/megatron/gpt_full_te_layer_autocast_spec.py Outdated
apply_residual_connection_post_layernorm=config.apply_residual_connection_post_layernorm,
autocast_dtype=precision,
#use_emha=False, # Use default 'False'
ub_tp_comm_overlap=config.tp_comm_overlap, # TODO: ub_tp_comm_overlap?
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
ub_tp_comm_overlap=config.tp_comm_overlap, # TODO: ub_tp_comm_overlap?
ub_tp_comm_overlap=config.tp_comm_overlap,

minitu
minitu reviewed Feb 2, 2024
View reviewed changes
nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py Outdated
@@ -225,6 +227,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):

self.mcore_gpt = cfg.get('mcore_gpt', False)
self.spec_name = cfg.get('name', '')
if cfg.get('fp8', False):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this is an artifact from Jimmy's memory fixes (the first two commits), which are now in main.
So those changes should go away as well.

@jbaczek
Revert contex parallel changes
73763cc 
Signed-off-by: Jan Baczek <jbaczek@nvidia.com>
@github-actions github-actions bot removed the core Changes to NeMo Core label Feb 5, 2024
github-advanced-security[bot]
github-advanced-security bot found potential problems Feb 5, 2024
View reviewed changes
nemo/collections/nlp/models/language_modeling/megatron/gpt_full_te_layer_autocast_spec.py
Comment on lines +19 to +23
from nemo.collections.nlp.modules.common.megatron.utils import (
ApexGuardDefaults,
init_method_normal,
scaled_init_method_normal,
)

Check notice

Code scanning / CodeQL

Unused import Note

Import of 'init_method_normal' is not used.
Import of 'scaled_init_method_normal' is not used.
nemo/collections/nlp/models/language_modeling/megatron/gpt_full_te_layer_autocast_spec.py

try:
from megatron.core import parallel_state, tensor_parallel
from megatron.core.dist_checkpointing.utils import apply_prefix_mapping

Check notice

Code scanning / CodeQL

Unused import Note

Import of 'apply_prefix_mapping' is not used.
nemo/collections/nlp/models/language_modeling/megatron/gpt_full_te_layer_autocast_spec.py
"TransformerEngine was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
)

return TransformerBlockSubmodules(layer_specs=ModuleSpec(module=TETransformerLayerAutocast))

Check failure

Code scanning / CodeQL

Wrong name for an argument in a class instantiation Error

Keyword argument 'module' is not a supported parameter name of
ApexGuardDefaults.__init__
Loading
.
@jbaczek
Copy link
Collaborator Author

jbaczek commented Feb 5, 2024

Closing this one in favour of #8328

@jbaczek jbaczek closed this Feb 5, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ShriyaPalsamudram ShriyaPalsamudram ShriyaPalsamudram left review comments

@minitu minitu minitu left review comments

Assignees
No one assigned
Labels
CI NLP
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@jbaczek @ShriyaPalsamudram @minitu @jiemingz @sanandaraj5597