Removed extraneous usage of fsdp_config.load_monolith_rank0_only since that's unreliable (#3901)

Author: rithwik-db

composer/distributed/shared_utils.py

@@ -181,9 +181,6 @@ def update_sync_module_states_if_needed(model: nn.Module, fsdp_config: FSDP2Conf
     dist.all_reduce(any_ranks_meta, reduce_operation='MAX')
     requires_sync = all_ranks_meta.item() == 0 and any_ranks_meta.item() == 1
 
-    if fsdp_config.load_monolith_rank0_only:
-        fsdp_config.sync_module_states = True
-
     if not fsdp_config.sync_module_states and requires_sync:
         fsdp_config.sync_module_states = True
 

composer/trainer/trainer.py

@@ -1701,12 +1701,6 @@ def __init__(
                 log.info('No previous autoresume checkpoint found')
         # Actually load the checkpoint from potentially updated arguments
         if load_path is not None:
-            # If we are using FSDP and load_monolith_rank0_only is True, then the state_dict must be `full`
-            # when we are loading a checkpoint
-            if self.state.fsdp_config and self.state.fsdp_config.load_monolith_rank0_only:  # type: ignore
-                err_msg = 'state_dict_type must be `full` when load_monolith_rank0_only is True when loading a checkpoint'
-                assert self.state.fsdp_config.state_dict_type == 'full', err_msg  # type: ignore
-
             log.info(f'Loading checkpoint from {load_path}')
             if load_object_store is None:
                 load_object_store = maybe_create_object_store_from_uri(load_path)