mirrors
/
mmengine
mirror of https://github.com/open-mmlab/mmengine.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Fix typos and documents of colossalai (#1315)

pull/1321/head
Mashiro 2023-08-22 16:13:55 +08:00 committed by GitHub
parent db32234241
commit 19ab172b2d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 14 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/en/api/strategy.rst
View File

@ -38,4 +38,4 @@ mmengine._strategy
:template: classtemplate.rst
CollosalAIModelWrapper
ColossalAIOpitmWrapper
ColossalAIOptimWrapper

6
docs/en/common_usage/large_model_training.md
View File

@ -205,8 +205,8 @@ pip install colossalai
Once ColossalAI is installed, configure the `strategy` and `optim_wrapper` parameters for FlexibleRunner:
- `strategy`: Specify `type='ColossalAIStrategy'` and configure the parameters. Detailed parameter descriptions can be found in [ColossalAIStrategy](mmengine._strategy.ColossalAI).
- `optim_wrapper`: Default to no `type` parameter or specify `type=ColossalAIOpitmWrapper`. It is recommended to choose `HybridAdam` as the optimizer type. Other configurable types are listed in [ColossalAIOptimWrapper](mmengine._strategy.ColossalAIOptimWrapper).
- `strategy`: Specify `type='ColossalAIStrategy'` and configure the parameters. Detailed parameter descriptions can be found in [ColossalAIStrategy](mmengine._strategy.ColossalAIStrategy).
- `optim_wrapper`: Default to no `type` parameter or specify `type=ColossalAIOptimWrapper`. It is recommended to choose `HybridAdam` as the optimizer type. Other configurable types are listed in [ColossalAIOptimWrapper](mmengine._strategy.ColossalAIOptimWrapper).
Here's the configuration related to ColossalAI:
@ -243,7 +243,7 @@ torchrun --nproc-per-node 2 examples/distributed_training_with_flexible_runner.p
<summary>Training Logs</summary>
```
08/18 11:56:34 - mmengine - INFO - Epoch(train) [1][ 10/196] lr: 3.3333e-04 eta: 0:10:31 time: 0.3238 data_time: 0.0344 memory: 597 loss: 3.8766
08/18 11:56:34 - mmengine - INFO - Epoch(train) [1][ 10/196] lr: 3.3333e-04 eta: 0:10:31 time: 0.3238 data_time: 0.0344 memory: 597 loss: 3.8766
08/18 11:56:35 - mmengine - INFO - Epoch(train) [1][ 20/196] lr: 3.3333e-04 eta: 0:06:56 time: 0.1057 data_time: 0.0338 memory: 597 loss: 2.3797
08/18 11:56:36 - mmengine - INFO - Epoch(train) [1][ 30/196] lr: 3.3333e-04 eta: 0:05:45 time: 0.1068 data_time: 0.0342 memory: 597 loss: 2.3219
08/18 11:56:37 - mmengine - INFO - Epoch(train) [1][ 40/196] lr: 3.3333e-04 eta: 0:05:08 time: 0.1059 data_time: 0.0337 memory: 597 loss: 2.2641

2
docs/zh_cn/api/strategy.rst
View File

@ -38,4 +38,4 @@ mmengine._strategy
:template: classtemplate.rst
CollosalAIModelWrapper
ColossalAIOpitmWrapper
ColossalAIOptimWrapper

4
docs/zh_cn/common_usage/large_model_training.md
View File

@ -204,8 +204,8 @@ pip install colossalai
安装好 ColossalAI 后，需配置 FlexibleRunner 的 strategy 和 optim_wrapper 参数：
- strategy：指定 `type='ColossalAIStrategy'` 并配置参数。参数的详细介绍可阅读 [ColossalAIStrategy](mmengine._strategy.ColossalAI)。
- optim_wrapper：缺省 `type` 参数，或指定 `type=ColossalAIOpitmWrapper`，优化器类型时建议选择 `HybridAdam`。其他可配置类型可阅读 [ColossalAIOptimWrapper](mmengine._strategy.ColossalAIOptimWrapper)。
- strategy：指定 `type='ColossalAIStrategy'` 并配置参数。参数的详细介绍可阅读 [ColossalAIStrategy](mmengine._strategy.ColossalAIStrategy)。
- optim_wrapper：缺省 `type` 参数，或指定 `type=ColossalAIOptimWrapper`，优化器类型建议选择 `HybridAdam`。其他可配置类型可阅读 [ColossalAIOptimWrapper](mmengine._strategy.ColossalAIOptimWrapper)。
下面是 ColossalAI 相关的配置：

14
mmengine/_strategy/colossalai.py
View File

@ -68,7 +68,7 @@ def register_mixed_precisions():
@OPTIM_WRAPPERS.register_module()
class ColossalAIOpitmWrapper(OptimWrapper):
class ColossalAIOptimWrapper(OptimWrapper):
"""OptimWrapper for ColossalAI.
The available optimizers are:
@ -83,7 +83,7 @@ class ColossalAIOpitmWrapper(OptimWrapper):
You can find more details in the `colossalai tutorial`_
Args:
optimizer (dict or collossal.booster.Booster): The optimizer to be
optimizer (dict or torch.optim.Optimizer): The optimizer to be
wrapped.
accumulative_counts (int): The number of iterations to accumulate
gradients. The parameters will be updated per
@ -128,7 +128,7 @@ class CollosalAIModelWrapper:
def train_step(
self,
data: Union[dict, tuple, list],
optim_wrapper: ColossalAIOpitmWrapper,
optim_wrapper: ColossalAIOptimWrapper,
) -> Dict[str, torch.Tensor]:
data = self.model.data_preprocessor(data, training=True)
with optim_wrapper.optim_context(self.model):
@ -234,7 +234,7 @@ class ColossalAIStrategy(BaseStrategy):
MODEL_DIR = 'model' # directory to save model
SCHEDULER_DIR = 'scheduler' # directory to save scheduelrs
model: CollosalAIModelWrapper # type: ignore
optim_wrapper: ColossalAIOpitmWrapper # type: ignore
optim_wrapper: ColossalAIOptimWrapper # type: ignore
def __init__(
self,
@ -304,7 +304,7 @@ class ColossalAIStrategy(BaseStrategy):
# optim_wrapper is required by booster
if optim_wrapper is not None and isinstance(optim_wrapper, dict):
optim_wrapper.setdefault('type', 'ColossalAIOpitmWrapper')
optim_wrapper.setdefault('type', 'ColossalAIOptimWrapper')
optim_wrapper.setdefault('booster', self.booster)
optim_wrapper_type = OPTIM_WRAPPERS.get(optim_wrapper['type'])
if optim_wrapper_type is None:
@ -312,7 +312,7 @@ class ColossalAIStrategy(BaseStrategy):
'`OPTIM_WRAPPERS`.')
if 'clip_grad' in optim_wrapper:
raise ValueError('`Please configure `clip_grad` in `plugin`')
if not issubclass(optim_wrapper_type, ColossalAIOpitmWrapper):
if not issubclass(optim_wrapper_type, ColossalAIOptimWrapper):
raise ValueError(
'The type of `optim_wrapper` must be '
'`ColossalAIOptimWrapper` (or subclass), but got '
@ -503,7 +503,7 @@ class ColossalAIStrategy(BaseStrategy):
self,
model: nn.Module,
optim_wrapper: Optional[OptimWrapper] = None,
) -> Union[Tuple[CollosalAIModelWrapper, ColossalAIOpitmWrapper],
) -> Union[Tuple[CollosalAIModelWrapper, ColossalAIOptimWrapper],
CollosalAIModelWrapper]: # type: ignore
"""Wrap model with :class:`ModelWrapper`."""
if self.model_wrapper is None: