reversed logic to ensure checkpointing is unsharded for single accele…

…rator
allenai · Dec 21, 2024 · 539f64a · 539f64a
1 parent 311286c
commit 539f64a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/train.py b/olmo/train.py
@@ -1251,7 +1251,7 @@ def on_trace_ready(p):
                             stop_at = min(stop_at, self.global_step + extra_steps)
 
                     # Maybe save sharded checkpoint.
-                    if self.cfg.distributed_strategy != DistributedStrategy.ddp:
+                    if self.cfg.distributed_strategy == DistributedStrategy.fsdp:
                         if save_checkpoints and (
                             cancel_initiated
                             or (