github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix] remove duplicated param register to stateful tensor manager (#728)

pull/726/head^2
Jiarui Fang 2022-04-12 13:55:25 +08:00 committed by GitHub
parent 600e769a42
commit 7db3ccc79b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/zero/sharded_model/sharded_model_v2.py
View File

@ -106,11 +106,9 @@ class ShardedModelV2(nn.Module):
GLOBAL_MODEL_DATA_TRACER.register_model(self)
self._memstats_collector = MemStatsCollector()
self._stateful_tensor_mgr = StatefulTensorMgr(self._memstats_collector)
# for param in module.parameters():
for submodule in module.modules():
for param in submodule.parameters(recurse=False):
if hasattr(param, 'colo_attr'):
self._stateful_tensor_mgr.register_stateful_param(param.colo_attr)
for param in module.parameters():
if hasattr(param, 'colo_attr'):
self._stateful_tensor_mgr.register_stateful_param(param.colo_attr)
self._start_collect_memstats = disposable(self._memstats_collector.start_collection)
self._finish_collect_memstats = disposable(self._memstats_collector.finish_collection)
else: