github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 02:51:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

[gemini] APIs to set cpu memory capacity (#809)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-19 16:05:22 +08:00
committed by GitHub
parent f6dcd23fb9
commit 227d1cd4b3
3 changed files with 29 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/initialize.py
View File

@@ -275,9 +275,6 @@ def initialize(model: nn.Module,
optimizer_config=optimizer_config)
logger.info("Initializing ZeRO model and optimizer finished!", ranks=[0])
# FIXME() throw a warning if using zero with MP
if gpc.get_world_size(ParallelMode.MODEL) > 1:
logger.warning("ZeRO currently has not been tested with model parallelism.", ranks=[0])
else:
if isinstance(model, nn.Module):
# first sync model across dp ranks