Uh oh!
There was an error while loading. Please reload this page.
Store FP8 checkpointing data in CPU - #351
Merged
Merged
Conversation
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ksivaman
commented
Aug 2, 2023
MemberAuthor
/te-ci |
Comment on lines
+90
to
+96
| buffer = {} | ||
| # Map all tensors to CPU. | ||
| for k, v in _global_fp8_buffer.items(): | ||
| buffer[k] = [tensor.cpu() for tensor in v] | ||
| return buffer |
Member
There was a problem hiding this comment.
This doesn't seem like the right place for checkpoint-specific logic. How about we change
state["global_fp8_buffer"] = {
key: [tensor.cpu() fortensorintensor_list]
forkey, tensor_listinget_global_fp8_buffer().items()
}MemberAuthor
There was a problem hiding this comment.
get_global_fp8_buffer is a checkpoint specific function. Its only use is to return the current global buffer for saving the checkpoint.
Member
There was a problem hiding this comment.
I see this implementation of get_global_fp8_buffer mirrors the implementation of set_global_fp8_buffer:
TransformerEngine/transformer_engine/pytorch/fp8.py
Lines 93 to 101 in 95ec156
That said, it's unintuitive. The function names and documentation don't make me think of checkpointing.
timmoon10
self-requested a review
August 2, 2023 22:50
timmoon10
approved these changes
Aug 2, 2023
mingxu1067 pushed a commit
to mingxu1067/TransformerEngine
that referenced
this pull request
Aug 3, 2023
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Ming Huang <mingh@nvidia.com>
timmoon10 added a commit
that referenced
this pull request
Aug 3, 2023
* Cast Flax collections to FrozenDict as WAR to adapt Flax 0.7.1. Signed-off-by: Ming Huang <mingh@nvidia.com> * Adding min version of flax to requirements.txt in examples Signed-off-by: Ming Huang <mingh@nvidia.com> * Fix praxis tests and rename compare_frozen_dict to compare_dict Signed-off-by: Ming Huang <mingh@nvidia.com> * [Paddle] Refactor FP8 state (#350) Refactor fp8 state Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Ming Huang <mingh@nvidia.com> * Store FP8 checkpointing data in CPU (#351) Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Ming Huang <mingh@nvidia.com> * Make test_layer be able to run on both Flax >=0.7.1 and <=0.7.0 Signed-off-by: Ming Huang <mingh@nvidia.com> * Update Flax version Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Ming Huang <mingh@nvidia.com> Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tian Zheng <tizheng@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See comment for the need for this.
Note: This is backwards compatible with current checkpoints including all parallelism modes (DP/SP/TP/PP).