support auto hsdp

stanford-crfm · Nov 14, 2024 · 3bcac0a · 3bcac0a
1 parent 433490e
commit 3bcac0a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py
@@ -766,6 +766,10 @@ def _validate_and_set_defaults(self):
         if self.per_device_eval_parallelism == -1:
             self.per_device_eval_parallelism = self.per_device_parallelism
 
+        if self.replica_dcn_axis_size == -1:
+            self.replica_dcn_axis_size = self.num_slices
+            logger.info(f"Setting replica_dcn_axis_size to {self.replica_dcn_axis_size}")
+
 
 class AllConfig(Protocol):
     trainer: TrainerConfig