github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
colossalai/legacy/zero/__init__.py
View File

@@ -11,8 +11,9 @@ from .sharded_model import ShardedModelV2
from .sharded_optim import ShardedOptimizerV2
def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model_config,
optimizer_config) -> Tuple[ShardedModelV2, ShardedOptimizerV2]:
def convert_to_zero_v2(
model: nn.Module, optimizer: torch.optim.Optimizer, model_config, optimizer_config
) -> Tuple[ShardedModelV2, ShardedOptimizerV2]:
"""
A helper function to integrate the model and optimizer with ZeRO optimizer and off-loading
@@ -25,12 +26,12 @@ def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model
:rtype: Tuple
"""
logger = get_dist_logger('convert_to_zero_v2')
logger = get_dist_logger("convert_to_zero_v2")
logger.info(f'optimizer_config is {optimizer_config}', ranks=[0])
logger.info(f"optimizer_config is {optimizer_config}", ranks=[0])
if optimizer_config is None:
optimizer_config = dict()
logger.info(f'model_config is {model_config}', ranks=[0])
logger.info(f"model_config is {model_config}", ranks=[0])
if model_config is None:
model_config = dict()
@@ -40,6 +41,12 @@ def convert_to_zero_v2(model: nn.Module, optimizer: torch.optim.Optimizer, model
__all__ = [
'convert_to_zero_v2', 'ShardedModelV2', 'ShardedOptimizerV2', 'ZeroInitContext', 'no_shard_zero_context',
'no_shard_zero_decrator', 'TensorShardStrategy', 'BucketTensorShardStrategy'
"convert_to_zero_v2",
"ShardedModelV2",
"ShardedOptimizerV2",
"ZeroInitContext",
"no_shard_zero_context",
"no_shard_zero_decrator",
"TensorShardStrategy",
"BucketTensorShardStrategy",
]