add magic filename for sharded state dicts (#1001)

* add magic filename for sharded state dicts * Update scripts/train/train.py Co-authored-by: Daniel King <[email protected]> * oops forgot to push this * no shard if no fsdp * default to full on foundry --------- Co-authored-by: Daniel King <[email protected]>
mosaicml · Mar 5, 2024 · 96c8218 · 96c8218
1 parent 3880d04
commit 96c8218
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -262,10 +262,14 @@ def main(cfg: DictConfig) -> Trainer:
                                             'save_folder',
                                             must_exist=False,
                                             default_value=None)
-    save_latest_filename: str = pop_config(cfg,
-                                           'save_latest_filename',
-                                           must_exist=False,
-                                           default_value='latest-rank{rank}.pt')
+    is_state_dict_sharded: bool = (fsdp_config.get('state_dict_type', 'full')
+                                   == 'sharded') if fsdp_config else False
+    save_latest_filename: str = pop_config(
+        cfg,
+        'save_latest_filename',
+        must_exist=False,
+        default_value='latest-sharded-rank{rank}'
+        if is_state_dict_sharded else 'latest-rank{rank}.pt')
     save_overwrite: bool = pop_config(cfg,
                                       'save_overwrite',
                                       must_exist=False,