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

MemStatsCollectorStatic (#1765)

Browse Source
This commit is contained in:
Zihao
2022-11-07 16:49:03 +08:00
committed by GitHub
parent 327d07c44a
commit 20e255d4e8
4 changed files with 142 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/nn/parallel/data_parallel.py
View File

@@ -267,7 +267,7 @@ class ZeroDDP(ColoDDP):
def forward(self, *args, **kwargs):
args, kwargs = _cast_float(args, torch.half), _cast_float(kwargs, torch.half)
self.module.zero_grad(set_to_none=True)
self.gemini_manager.pre_iter()
self.gemini_manager.pre_iter(*args)
with ParamOpHookManager.use_hooks(self.param_op_hook):
outputs = self.module(*args, **kwargs)
if self.force_outputs_fp32: