Add SMP v2 notebook for accelerating training with FP8 on P5. #4578

viclzhu · 2024-03-01T00:49:53Z

Issue #, if available:

Description of changes:
Update scripts to latest PT-2.2-TSM-2.2 release.

Testing done:

Merge Checklist

Put an x in the boxes that apply. You can also fill these out after creating the PR. If you're unsure about any of them, don't hesitate to ask. We're here to help! This is simply a reminder of what we are going to look for before merging your pull request.

I have read the CONTRIBUTING doc and adhered to the example notebook best practices
I have updated any necessary documentation, including READMEs
I have tested my notebook(s) and ensured it runs end-to-end
I have linted my notebook(s) and code using black-nb -l 100 {path}/{notebook-name}.ipynb

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

review-notebook-app · 2024-03-01T00:49:57Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

sagemaker-bot · 2024-03-01T00:51:28Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: ef20644
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:51:45Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: ef20644
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:51:48Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: ef20644
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:53:11Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: ef20644
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:53:42Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 82224fe
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:53:56Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 82224fe
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:54:10Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 82224fe
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T00:54:53Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 82224fe
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/train_lib.py

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/memory_tracker.py

sagemaker-bot · 2024-03-01T19:06:23Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 74a54eb
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T19:07:23Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 74a54eb
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T19:07:28Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 74a54eb
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T19:07:45Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 74a54eb
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T19:18:22Z

We identified new issues on unchanged lines of code. Navigate to the Amazon CodeGuru Reviewer console to view the recommendations to fix them.

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/fsdp_utils.py

...ing/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-train-gpt-neox-fsdp-tp.ipynb

...istributed_training/pytorch/model_parallel_v2/gpt-neox/smp-finetuning-gpt-neox-fsdp-tp.ipynb

...ng/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-train-llama-fsdp-tp-fp8.ipynb

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/arguments.py

sagemaker-bot · 2024-03-01T23:19:40Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 0f81769
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:20:10Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 0f81769
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:20:20Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 0f81769
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:21:17Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 0f81769
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:21:24Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 5ac3e7f
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:21:52Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 5ac3e7f
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-01T23:22:40Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 5ac3e7f
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

Revert jupyter notebook python version in metadata. Set activation_offloading=False for FP8 notebook. Explicitly enable use_smp_implementation in all SMP v2 notebooks.

sagemaker-bot · 2024-03-04T22:25:59Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: d66fca8
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-04T22:27:48Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: d66fca8
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-04T22:28:28Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: d66fca8
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-04T22:28:35Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: d66fca8
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/utils.py

...ing/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-train-gpt-neox-fsdp-tp.ipynb

sagemaker-bot · 2024-03-05T03:14:33Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 4ec778d
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-05T03:14:46Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 4ec778d
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-05T03:15:15Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 4ec778d
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-05T03:15:37Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 4ec778d
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:39:24Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 22956c1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:39:31Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 22956c1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:40:01Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 22956c1
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:40:23Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 22956c1
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:55:25Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-amazon-sagemaker-examples-pr
Commit ID: 24a9d87
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:55:48Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-link-check
Commit ID: 24a9d87
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:56:20Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-code-formatting
Commit ID: 24a9d87
Result: SUCCEEDED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

sagemaker-bot · 2024-03-06T22:56:29Z

AWS CodeBuild CI Report

CodeBuild project: SageMakerNotebooksProd-USWEST2-sagemaker-examples-grammar
Commit ID: 24a9d87
Result: FAILED
Build Logs (available for 30 days)

Powered by github-codebuild-logs, available on the AWS Serverless Application Repository

viclzhu force-pushed the add-mp_v2-fp8 branch from ef20644 to 82224fe Compare March 1, 2024 00:52

sagemaker-bot reviewed Mar 1, 2024

View reviewed changes

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/fsdp_utils.py Show resolved Hide resolved

yongyanrao reviewed Mar 1, 2024

View reviewed changes

...ing/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-train-gpt-neox-fsdp-tp.ipynb Outdated Show resolved Hide resolved

yongyanrao reviewed Mar 1, 2024

View reviewed changes

...istributed_training/pytorch/model_parallel_v2/gpt-neox/smp-finetuning-gpt-neox-fsdp-tp.ipynb Outdated Show resolved Hide resolved

cavdard reviewed Mar 1, 2024

View reviewed changes

...ng/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-train-llama-fsdp-tp-fp8.ipynb Outdated Show resolved Hide resolved

cavdard reviewed Mar 1, 2024

View reviewed changes

...ng/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-train-llama-fsdp-tp-fp8.ipynb Outdated Show resolved Hide resolved

cavdard reviewed Mar 1, 2024

View reviewed changes

...ng/distributed_training/pytorch/model_parallel_v2/llama_v2/smp-train-llama-fsdp-tp-fp8.ipynb Show resolved Hide resolved

yongyanrao reviewed Mar 1, 2024

View reviewed changes

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/arguments.py Show resolved Hide resolved

viclzhu added 4 commits March 4, 2024 14:24

Add example notebook for accelerating Llama-v2 training with FP8 on P5.

476472b

Fix typo in version check

36a9da3

Update configurations.

c620d42

Revert jupyter notebook python version in metadata. Set activation_offloading=False for FP8 notebook. Explicitly enable use_smp_implementation in all SMP v2 notebooks.

Update FP8 notebook docs.

d66fca8

viclzhu force-pushed the add-mp_v2-fp8 branch from 7f66dd2 to d66fca8 Compare March 4, 2024 22:24

sagemaker-bot reviewed Mar 4, 2024

View reviewed changes

training/distributed_training/pytorch/model_parallel_v2/shared-scripts/utils.py Show resolved Hide resolved

mchoi8739 reviewed Mar 5, 2024

View reviewed changes

...ing/distributed_training/pytorch/model_parallel_v2/gpt-neox/smp-train-gpt-neox-fsdp-tp.ipynb Show resolved Hide resolved

Set zipped_data=0 for use_fsx=False FP8 notebook.

4ec778d

viclzhu marked this pull request as ready for review March 5, 2024 18:29

Update compute_tflops() script.

22956c1

viclzhu changed the title ~~[Draft] Add SMP v2 notebook for accelerating training with FP8 on P5.~~ Add SMP v2 notebook for accelerating training with FP8 on P5. Mar 6, 2024

Update minimum sagemaker pysdk version to 2.212.

24a9d87

knikure approved these changes Mar 7, 2024

View reviewed changes

knikure merged commit bf3cadc into aws:main Mar 7, 2024
4 of 5 checks passed

Add SMP v2 notebook for accelerating training with FP8 on P5. #4578

Add SMP v2 notebook for accelerating training with FP8 on P5. #4578

Conversation

viclzhu commented Mar 1, 2024 • edited Loading

Merge Checklist

review-notebook-app bot commented Mar 1, 2024

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 1, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 4, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 4, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 4, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 4, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 5, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 5, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 5, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 5, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

sagemaker-bot commented Mar 6, 2024

AWS CodeBuild CI Report

viclzhu commented Mar 1, 2024 •

edited

Loading