Skip to content

[PyTorch] Experimental FP8 tensor class - #452

Merged
ksivaman merged 29 commits into
NVIDIA:mainfrom
timmoon10:float8tensor_experiments
Oct 31, 2023
Merged

[PyTorch] Experimental FP8 tensor class#452
ksivaman merged 29 commits into
NVIDIA:mainfrom
timmoon10:float8tensor_experiments

Conversation

@timmoon10

@timmoon10timmoon10 commented Sep 29, 2023

Copy link
Copy Markdown
Member

This FP8 tensor class is based on the implementation at https://github.com/facebookexperimental/protoquant/tree/fp8_poc and is primarily oriented toward enabling efficient FP8 support in Apex's DistributedFusedAdam. See NVIDIA-NeMo/Speech#7469 and NVIDIA-NeMo/Speech#7565.

CC @sudhakarsingh27@ksivaman

@timmoon10timmoon10 added the enhancement New feature or request label Sep 29, 2023
ptrendx
ptrendx reviewed Sep 29, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py Outdated
@ptrendx

ptrendx commented Sep 30, 2023

Copy link
Copy Markdown
Member

/te-ci

1 similar comment
@timmoon10

timmoon10 commented Oct 5, 2023

Copy link
Copy Markdown
MemberAuthor

/te-ci

@sudhakarsingh27

sudhakarsingh27 commented Oct 16, 2023

Copy link
Copy Markdown
Member

/te-ci pytorch

@ksivaman
ksivamanforce-pushed the float8tensor_experiments branch from de20156 to 4315115CompareOctober 16, 2023 21:08
@ksivaman@timmoon10
@sudhakarsingh27@ptrendx
Experimental FP8 tensor
b6bfddb
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivamanforce-pushed the float8tensor_experiments branch from 67f7cd3 to b6bfddbCompareOctober 19, 2023 23:45
@ksivaman
Add fp8 tensor to ci test
36093a5
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman marked this pull request as ready for review October 20, 2023 04:47
@ksivaman

ksivaman commented Oct 20, 2023

Copy link
Copy Markdown
Member

/te-ci pytorch

@ksivaman

ksivaman commented Oct 20, 2023

Copy link
Copy Markdown
Member

/te-ci pytorch

1 similar comment
@ksivaman

ksivaman commented Oct 20, 2023

Copy link
Copy Markdown
Member

/te-ci pytorch

ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
handled outside this class. If a tensor is initialized with an FP8
metadata object, it extracts the information it needs so it isn't
affected by later changes in the FP8 metadata (although its design
does cause us to leak some subtle side-effects into FP8 metadata).

@ptrendxptrendxOct 23, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This doc is not really correct since we are holding a view to the meta, right?

@timmoon10timmoon10Oct 24, 2023

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

Ops using the tensor class's __torch_dispatch__ are insensitive to external changes in the meta since we cache scale_inv. However, all bets are off when we extract _data and pass it to external ops like tex.fp8_gemm.

ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py Outdated
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/module/base.py
ptrendx
ptrendx reviewed Oct 23, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/module/layernorm_linear.py
timmoon10
timmoon10 commented Oct 24, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py Outdated
Comment threadtransformer_engine/pytorch/module/base.py Outdated
Comment threadtransformer_engine/pytorch/module/base.py
Comment threadtransformer_engine/pytorch/module/layernorm_linear.py
Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
handled outside this class. If a tensor is initialized with an FP8
metadata object, it extracts the information it needs so it isn't
affected by later changes in the FP8 metadata (although its design
does cause us to leak some subtle side-effects into FP8 metadata).

@timmoon10timmoon10Oct 24, 2023

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

Ops using the tensor class's __torch_dispatch__ are insensitive to external changes in the meta since we cache scale_inv. However, all bets are off when we extract _data and pass it to external ops like tex.fp8_gemm.

Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
Comment threadtransformer_engine/pytorch/float8_tensor.py
Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
sudhakarsingh27
sudhakarsingh27 reviewed Oct 24, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py
sudhakarsingh27
sudhakarsingh27 reviewed Oct 24, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py Outdated
@ksivaman
Merge branch 'main' into float8tensor_experiments
dfcbcf1
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
review comments and tests
8814ead
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
timmoon10and others added 7 commits October 25, 2023 14:27
@timmoon10
Rename FP8GlobalStateManager.with_fp8_parameters
202afcb
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@ksivaman
remove Float8Tensor from import API
1d0b1fe
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@timmoon10
Avoid caching FP8 transposes if not required
39add1a
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Fix import error in FP8 tensor tests
b845d32
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@ksivaman
Fix tranpose caching and checkpointing bug
a5351b3
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'main' into float8tensor_experiments
79932f1
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Improve caching and fix distopt case
7d95a91
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
timmoon10
timmoon10 commented Oct 27, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py Outdated
timmoon10and others added 2 commits October 26, 2023 17:34
@timmoon10
Update transformer_engine/pytorch/float8_tensor.py
20fc9a9
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
@ksivaman
Remove recursive logic
9f08be7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
timmoon10
timmoon10 commented Oct 27, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/float8_tensor.py
ksivamanand others added 4 commits October 27, 2023 22:20
@ksivaman
Fix cache reset bug
00b9c31
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@timmoon10
Store FP8 attributes in dict
4cf27a1
Easier for multiple tensors to share, e.g. detached tensors.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Make sure scale_inv is 1D tensor
718d284
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
Make sure scale_inv is 1D tensor
94848da
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
timmoon10force-pushed the float8tensor_experiments branch from 718d284 to 94848daCompareOctober 31, 2023 00:34
@ksivaman
Fixes and detach recipe
ac192d8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Merge branch 'float8tensor_experiments' ofhttps://github.com/timmoon…
918bda3
…10/TransformerEngine into float8tensor_experiments
@ksivaman
Set default fp8 data type
7b3f5cd
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx approved these changes Oct 31, 2023
View reviewed changes

@ptrendxptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Approving as experimental. We will iterate upon this in the next release.

@ksivaman
ksivaman merged commit b1820c4 into NVIDIA:mainOct 31, 2023
ksivaman added a commit that referenced this pull request Oct 31, 2023
@timmoon10@ksivaman
@sudhakarsingh27@ptrendx
[PyTorch] Experimental FP8 tensor class (#452)
d58c08c
* Experimental FP8 tensor
Co-authored-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add fp8 tensor to ci test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* review comments and tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Minor changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Default to FP8 usage
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix docs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Naming changes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* minor fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix transpose caching
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Debug transpose caching
Handle case where transpose cache is updated externally.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Rename FP8GlobalStateManager.with_fp8_parameters
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* remove Float8Tensor from import API
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Avoid caching FP8 transposes if not required
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix import error in FP8 tensor tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fix tranpose caching and checkpointing bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Improve caching and fix distopt case
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Update transformer_engine/pytorch/float8_tensor.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* Remove recursive logic
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix cache reset bug
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Store FP8 attributes in dict
Easier for multiple tensors to share, e.g. detached tensors.
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Make sure scale_inv is 1D tensor
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Fixes and detach recipe
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Set default fp8 data type
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
This was referenced Nov 14, 2023
Closed
Merged
carmocca
carmocca reviewed Nov 29, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/fp8.py
Comment on lines +507 to +508
* full model training using optimizer with master weights, where the high
precision copies of weights are already present in the optimizer.

@carmoccacarmoccaNov 29, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

How does this look in practice? If the model will be initialized directly with fp8 weights, how does the optimizer get high-precision copies?

@timmoon10timmoon10 mentioned this pull request Feb 13, 2024
Closed
@timmoon10timmoon10 mentioned this pull request Apr 27, 2024
Merged
@hungryGeek16hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sudhakarsingh27sudhakarsingh27sudhakarsingh27 left review comments
@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes
+1 more reviewer
@carmoccacarmoccacarmocca left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

1.0.0enhancementNew feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@timmoon10@ptrendx@sudhakarsingh27@ksivaman@carmocca