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

Debug Transformer Engine FP8 support with Megatron-core infrastructure #6791

Merged
merged 3 commits into from
Jun 4, 2023
Merged

Debug Transformer Engine FP8 support with Megatron-core infrastructure #6791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0e3b940
Construct FP8 amax reduction group
timmoon10 Jun 3, 2023
bb1b049
Update Megatron-core version in CI
timmoon10 Jun 3, 2023
bedad2c
Merge branch 'main' into cherry-pick-main-4bbb3c663b07045a88b9e095e89…
timmoon10 Jun 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions Jenkinsfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,15 @@ pipeline {
}
}

stage('Megatron Core installation') {
steps {
sh 'git clone https://github.com/NVIDIA/Megatron-LM.git && \
cd Megatron-LM && \
git checkout cd2537d444792b487b1ab5a6fa685e09c9957409 && \
pip install -e .'
}
}

stage('PyTorch Lightning version') {
steps {
sh 'python -c "import pytorch_lightning; print(pytorch_lightning.__version__)"'
Expand Down
3 changes: 2 additions & 1 deletion nemo/collections/nlp/parts/nlp_overrides.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ def init_model_parallel(self, global_rank: int, world_size: int) -> None:
pipeline_model_parallel_size=app_state.pipeline_model_parallel_size,
virtual_pipeline_model_parallel_size=app_state.virtual_pipeline_model_parallel_size,
pipeline_model_parallel_split_rank=app_state.pipeline_model_parallel_split_rank,
use_fp8=app_state.use_fp8,
)

# assert that fake tp and pp rank match after model parallel init
Expand Down Expand Up @@ -406,7 +407,7 @@ class PEFTSaveRestoreConnector(NLPSaveRestoreConnector):
peft_model_nemo_path: Used to provide the .nemo file corresponding to a PEFT model (which will only contain a small set of params)
peft_model_ckpt_path: Used to provide the path to .ckpt files of a PEFT model. This is required when no .nemo is available (yet) such as during resumed training.
peft_model_ckpt_name: The filename of the ckpt file inside the peft_model_ckpt_path folder
If both are provided the peft_model_ckpt_path takes precedence.
If both are provided the peft_model_ckpt_path takes precedence.
If neither are provided, PEFT params are initialized at random (not loaded from any external source).
"""

Expand Down