check str or list dtype

mosaicml · Dec 1, 2023 · eddccc2 · eddccc2
1 parent 3d6506b
commit eddccc2
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -739,9 +739,12 @@ def fsdp_wrap_fn(self, module: nn.Module) -> bool:
     def activation_checkpointing_fn(self, module: nn.Module) -> bool:
         act_ckpt_list = getattr(self.config, 'activation_checkpointing_target',
                                 None) or ['MPTBlock']
-        if not isinstance(act_ckpt_list, list):
-            # `activation_checkpointing_target` is a single value
+        if isinstance(act_ckpt_list, str):
             act_ckpt_list = [act_ckpt_list]
+        elif not isinstance(act_ckpt_list, list):
+            raise ValueError(
+                f'activation_checkpointing_target must be either a single string or a list, but got {type(act_ckpt_list)}'
+            )
 
         if 'MPTBlock' in act_ckpt_list or 'mptblock' in act_ckpt_list:
             if len(act_ckpt_list) > 1: