[fx] Add pofo solver #1608
Merged
[fx] Add pofo solver #1608
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
Merge ColossalAI
Daily merge
FrankLeeeee
reviewed
Sep 20, 2022
Comment on lines
80
to
81
| class PofoSolver: | ||
|
|
There was a problem hiding this comment.
You should include a docstring here. You can move your PR description here so that it will appear on our API documentation.
FrankLeeeee
reviewed
Sep 20, 2022
| from colossalai.fx.profiler import parameter_size | ||
| import math | ||
| from .linearize import linearize | ||
| from .utils import * |
FrankLeeeee
reviewed
Sep 20, 2022
| def _discretize(mem_unit, values): | ||
| return [math.ceil(value / mem_unit) for value in values] | ||
|
|
||
|
|
||
| class Chain: |
There was a problem hiding this comment.
These classes in the utils file are not really utils, they are more like to be in a file like operation.py or action.py.
added 3 commits
September 20, 2022 10:48
…/ColossalAI into feature/add_ckpt_offload_solver
FrankLeeeee
approved these changes
Sep 20, 2022
This was referenced Sep 20, 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.
What's New?
The new pofo solver is based on paper Efficient Combination of Rematerialization and Offloading for Training DNNs and it's code given in the supplemental. Currently we doesn't use the whole set up in the original paper and reuse rotor solver for the backward sequence as suggested in supplemental. The solver now is able to find strategy with offload. I will redesign the codegen later and add graph annotation to this solver when the codegen is compatible.