support auto hsdp (#804)

stanford-crfm · Nov 14, 2024 · a885f20 · a885f20
1 parent 55db7fe
commit a885f20
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py
@@ -783,6 +783,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