[feature] new zero implementation #1623
Merged
Conversation
This file contains 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
ver217
reviewed
Sep 22, 2022
| self.fp16_param_to_fp32_param[p] = fp32_p | ||
| chunk_16 = self.chunk_manager.get_chunk(p) | ||
| chunk_32 = self.chunk_manager.get_chunk(fp32_p) | ||
| chunk_32.init_pair(chunk_16) |
There was a problem hiding this comment.
It's a problem for a legacy reason. I'll move it to ZeroDDP.
ver217
reviewed
Sep 22, 2022
| else: | ||
| assert self.cuda_shard is not None # only check in CUDA | ||
| valid_tensor = self.cuda_shard[:self.valid_end] |
There was a problem hiding this comment.
Where do you copy fp32 param to fp16 param when both of them are on CPU?
There was a problem hiding this comment.
See shard_move function in the Chunk class.
feifeibear
approved these changes
Sep 24, 2022
feifeibear
added a commit
that referenced
this pull request
Sep 26, 2022
This reverts commit 5be118f.
feifeibear
added a commit
that referenced
this pull request
Sep 26, 2022
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
A New ZeRO Implementation
Backgrounds
In the current version, our ZeRO has a performance issue. The reason is that our asymmetric distribution of chunks makes one process hinder other processes when reading the content of the chunk located in the CPU memory. It really prolongs the time of the data transmission and undermines the efficiency of ZeRO.
Implementation
In order to solve this problem, I refactored the class
Chunk. The new chunk is distributed evenly to all processes. All processes can move the data from the CPU memory to the CUDA memory in the same time. Furthermore, I provide an option to enable the pin memory for chunks. Now all chunks can have a copy in the pinned CPU memory. The above optimizations prominently improved the efficiency of data movements between CPU and CUDA.Another Advantage
The new ZeRO supports the true hybrid parallelism. It creates different chunk groups for parameters which have different DP communication groups. This brings huge flexibility to our up coming automatic configuration of parallelism.