Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
6e553748
Unverified
Commit
6e553748
authored
Apr 14, 2022
by
ver217
Committed by
GitHub
Apr 14, 2022
Browse files
polish sharded optim docstr and warning (#770)
parent
80e37eec
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
4 deletions
+3
-4
colossalai/zero/sharded_optim/sharded_optim_v2.py
colossalai/zero/sharded_optim/sharded_optim_v2.py
+3
-4
No files found.
colossalai/zero/sharded_optim/sharded_optim_v2.py
View file @
6e553748
...
@@ -50,7 +50,7 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
...
@@ -50,7 +50,7 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
You must use ``ShardedOptimizerV2`` with ``ShardedModelV2``.
You must use ``ShardedOptimizerV2`` with ``ShardedModelV2``.
Note:
Note:
Make sure you
enable ``use_memory_tracer
`` in ``ShardedModelV2``,
Make sure you
set ``tensor_placement_policy
`` in ``ShardedModelV2``
to `"auto"`
,
if you set ``gpu_margin_mem_ratio > 0``.
if you set ``gpu_margin_mem_ratio > 0``.
Args:
Args:
...
@@ -59,7 +59,6 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
...
@@ -59,7 +59,6 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
optimizer (Optimizer): An Optimizer instance.
optimizer (Optimizer): An Optimizer instance.
gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward)
gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward)
which will be used when using hybrid CPU optimizer.
which will be used when using hybrid CPU optimizer.
Make sure `reuse_fp16_shard` is enabled in `ShardedModelV2`, if `gpu_margin_mem_ratio` > `0.0`.
This argument is meaningless when `tensor_placement_policy` of `ShardedModelV2` is not "auto".
This argument is meaningless when `tensor_placement_policy` of `ShardedModelV2` is not "auto".
Defaults to 0.0.
Defaults to 0.0.
initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
...
@@ -119,7 +118,7 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
...
@@ -119,7 +118,7 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
# Store fp32 param shards
# Store fp32 param shards
self
.
_register_master_weight
()
self
.
_register_master_weight
()
if
self
.
gpu_margin_mem_ratio
!=
0.0
and
isinstance
(
sharded_model
.
_tensor_placement_policy
,
if
self
.
gpu_margin_mem_ratio
!=
0.0
and
not
isinstance
(
sharded_model
.
_tensor_placement_policy
,
AutoTensorPlacementPolicy
):
AutoTensorPlacementPolicy
):
self
.
_logger
.
warning
(
f
'gpu_margin_mem_ratio is meaningless when tensor_placement_policy is not "auto"'
)
self
.
_logger
.
warning
(
f
'gpu_margin_mem_ratio is meaningless when tensor_placement_policy is not "auto"'
)
self
.
_logger
.
debug
(
f
"After init ShardedOptimizerV2 consumes
{
self
.
get_memory_usage
()[
0
]
/
1e6
}
MB CUDA Memory!"
,
self
.
_logger
.
debug
(
f
"After init ShardedOptimizerV2 consumes
{
self
.
get_memory_usage
()[
0
]
/
1e6
}
MB CUDA Memory!"
,
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment