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

7db3ccc7 · Jiarui Fang · GitHub · 600e769a · 7db3ccc7
Unverified Commit 7db3ccc7 authored Apr 12, 2022 by Jiarui Fang Committed by GitHub Apr 12, 2022
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 5 deletions

colossalai/zero/sharded_model/sharded_model_v2.py colossalai/zero/sharded_model/sharded_model_v2.py +3 -5

No files found.
--- a/colossalai/zero/sharded_model/sharded_model_v2.py
+++ b/colossalai/zero/sharded_model/sharded_model_v2.py
@@ -106,9 +106,7 @@ 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 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)
            self._start_collect_memstats = disposable(self._memstats_collector.start_collection)