add check

mosaicml · Jan 2, 2024 · a41581c · a41581c
1 parent 1ac5c4b
commit a41581c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -120,6 +120,13 @@ def process_init_device(model_cfg: DictConfig, fsdp_config: Optional[Dict]):
             # Set defaults for mixed initialization
             fsdp_config.setdefault('use_orig_params', False)
             fsdp_config.setdefault('load_monolith_rank0_only', True)
+        # Always set `sync_module_states` to True when using hybrid sharding
+        if fsdp_config.get('sharding_strategy', 'FULL_SHARD') in ['HYBRID_SHARD', '_HYBRID_SHARD_ZERO2'] \
+                and not fsdp_config.get('sync_module_states', False):
+            warnings.warn((
+                'Setting `sync_module_states = True` for FSDP. This is required '
+                'when using hybrid sharding.'))
+            fsdp_config['sync_module_states'] = True
 
     # no mixed precision needed for weights when they're already 16 bits
     master_dtype = model_cfg.get('master_weights_dtype')