Check FSDP module has _device_mesh before getting it (#3091)

mosaicml · Mar 29, 2024 · 3d4d506 · 3d4d506
1 parent d05c40d
commit 3d4d506
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/composer/core/state.py b/composer/core/state.py
@@ -762,7 +762,7 @@ def fsdp_elastic_sharded_enabled(self):
     @property
     def fsdp_device_mesh(self):
         if self.fsdp_enabled:
-            if not hasattr(self.model, 'model'):
+            if not hasattr(self.model, 'model') or not hasattr(self.model.model, '_device_mesh'):
                 return None
             return self.model.model._device_mesh
         else: