[shardformer] fix base policy (#4229)

colossalai/shardformer/policies/base_policy.py

@@ -156,6 +156,9 @@ class Policy(ABC):
 
         # append or create a new description
         if target_key in policy:
+            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,6 +177,9 @@ class Policy(ABC):
             target_key (Union[str, nn.Module]): the key of the policy to be updated
         """
         if target_key in policy:
+            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)