[AutoTuner]Add memory model #147
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
Caozhou1995
force-pushed
the
memory_model
branch
from
844c126
to
0a33cba
Compare
aoyulong
reviewed
Jun 17, 2024
flagscale/auto_tuner/memory_model.py
Outdated
There was a problem hiding this comment.
Please try to reuse this impl
There was a problem hiding this comment.
The impl has been reused and the activation section has been refined.
flagscale/auto_tuner/prune/memory.py
Outdated
|
|
||
|
|
||
| def prune_by_memory_model_util(config, strategy, history=[]): | ||
| if "modeling_memory" in strategy: |
There was a problem hiding this comment.
It's better to rename "modeling_memory" to "memory_model" as the other places.
flagscale/auto_tuner/tuner.py
Outdated
Comment on lines
80
to
97
| if os.environ.get("AIRS_ACCELERATOR_COUNT", None): | ||
| self.config.experiment.auto_tuner.nproc_per_node = int( | ||
| os.environ["AIRS_ACCELERATOR_COUNT"] | ||
| # Set config | ||
| self.config.experiment.auto_tuner.nproc_per_node = ( | ||
| int(os.environ["AIRS_ACCELERATOR_COUNT"]) * 2 | ||
| if "luvatar_BI" in os.environ["AIRS_ACCELERATOR_MODEL"] | ||
| else int(os.environ["AIRS_ACCELERATOR_COUNT"]) | ||
| ) | ||
| # Set original config | ||
| self.orig_config.experiment.runner.nproc_per_node = int( | ||
| os.environ["AIRS_ACCELERATOR_COUNT"] | ||
| self.orig_config.experiment.runner.nproc_per_node = ( | ||
| int(os.environ["AIRS_ACCELERATOR_COUNT"]) * 2 | ||
| if "luvatar_BI" in os.environ["AIRS_ACCELERATOR_MODEL"] | ||
| else int(os.environ["AIRS_ACCELERATOR_COUNT"]) | ||
| ) | ||
| # Set config | ||
| self.config.experiment.runner.nproc_per_node = int( | ||
| os.environ["AIRS_ACCELERATOR_COUNT"] | ||
| self.config.experiment.runner.nproc_per_node = ( | ||
| int(os.environ["AIRS_ACCELERATOR_COUNT"]) * 2 | ||
| if "luvatar_BI" in os.environ["AIRS_ACCELERATOR_MODEL"] | ||
| else int(os.environ["AIRS_ACCELERATOR_COUNT"]) |
There was a problem hiding this comment.
Is there any way to move these platform related code into a standalone place? We may support differnt cloud platform.
There was a problem hiding this comment.
The platform code has been removed to platform.py and different platforms code will be in this file.
Caozhou1995
force-pushed
the
memory_model
branch
5 times, most recently
from
d052639
to
69510a0
Compare
Caozhou1995
force-pushed
the
memory_model
branch
from
69510a0
to
e56b149
Compare
aoyulong
previously approved these changes
Jun 25, 2024
| @@ -0,0 +1,351 @@ | |||
| """ | |||
| Computes theoretical memory footprint for model training referring to megatron. | |||
| Activation memory is optimized with adding block recompute formula. | |||
There was a problem hiding this comment.
Please add reference to megatron original impl
| @@ -161,3 +167,114 @@ def compare_by_recompute(strategy1, strategy2): | |||
| result = True | |||
|
|
|||
| return result | |||
|
|
|||
|
|
|||
| def convert_config_to_megatron_args(config, strategy): | |||
There was a problem hiding this comment.
Is there any simpler way to impl this?
Caozhou1995
force-pushed
the
memory_model
branch
from
9d49a2a
to
9f6ad90
Compare
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.
This PR adds memory model, which can be used to speed up pruning and filter out OOM memory strategies and strategies with low memory usage.