MCPcopy
hub / github.com/hpcaitech/ColossalAI / _get_dfs_module_list

Function _get_dfs_module_list

colossalai/zero/gemini/utils.py:32–45  ·  view source on GitHub ↗

Get a dfs module list of the given module. Its order is same as the order of creations of modules.

(module: nn.Module, memo: Optional[Set[nn.Module]] = None, prefix: str = "")

Source from the content-addressed store, hash-verified

30
31
32def _get_dfs_module_list(module: nn.Module, memo: Optional[Set[nn.Module]] = None, prefix: str = ""):
33 """Get a dfs module list of the given module. Its order is same as the order of creations of modules."""
34 if memo is None:
35 memo = set()
36 if module not in memo:
37 for name, submodule in module._modules.items():
38 if submodule is None:
39 continue
40 submodule_prefix = prefix + ("." if prefix else "") + name
41 for m in _get_dfs_module_list(submodule, memo, submodule_prefix):
42 yield m
43
44 memo.add(module)
45 yield prefix, module
46
47
48def _get_shallow_copy_model(model: nn.Module):

Callers 2

_get_shallow_copy_modelFunction · 0.85
get_static_torch_modelFunction · 0.85

Calls 1

addMethod · 0.45

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…