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

[shardformer] fix base policy (#4229)

Browse Source
This commit is contained in:
Hongxin Liu
2023-07-14 09:51:53 +08:00
parent 208ac8f2ba
commit 7e4de520e1
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
colossalai/shardformer/policies/base_policy.py
View File

@@ -156,7 +156,10 @@ class Policy(ABC):
# append or create a new description
if target_key in policy:
policy[target_key].sub_module_replacement.extend(description)
if policy[target_key].sub_module_replacement is None:
policy[target_key].sub_module_replacement = description
else:
policy[target_key].sub_module_replacement.extend(description)
else:
policy[target_key] = ModulePolicyDescription(sub_module_replacement=description)
@@ -174,7 +177,10 @@ class Policy(ABC):
target_key (Union[str, nn.Module]): the key of the policy to be updated
"""
if target_key in policy:
policy[target_key].method_replacement.update(description)
if policy[target_key].method_replacement is None:
policy[target_key].method_replacement = description
else:
policy[target_key].method_replacement.update(description)
else:
policy[target_key] = ModulePolicyDescription(method_replacement=description)