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

[shardformer] fix base policy (#4229)

pull/4445/head
Hongxin Liu 2023-07-14 09:51:53 +08:00
parent 208ac8f2ba
commit 7e4de520e1
1 changed files with 8 additions and 2 deletions
Show Stats 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 # append or create a new description
if target_key in policy: 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: else:
policy[target_key] = ModulePolicyDescription(sub_module_replacement=description) 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 target_key (Union[str, nn.Module]): the key of the policy to be updated
""" """
if target_key in policy: 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: else:
policy[target_key] = ModulePolicyDescription(method_replacement=description) policy[target_key] = ModulePolicyDescription(method_replacement=description)