Added a max reshard retry count

MaxText/elasticutils.py

@@ -83,7 +83,8 @@ def __init__(
       total_slice_count: int,
       save_period: Optional[int] = None,
       reshard_check_period: Optional[int] = None,
-      max_failures: Optional[int] = None,
+      max_failure_count: Optional[int] = None,
+      max_reshard_retry_count: Optional[int] = None,
   ):
     self.devices = devices
     self.total_slice_count = total_slice_count
@@ -96,25 +97,40 @@ def __init__(
       reshard_check_period = 1
     self.reshard_check_period = reshard_check_period
 
-    if max_failures is None:
-      max_failures = float("inf")
-    self.max_failures = max_failures
+    if max_failure_count is None:
+      max_failure_count = float("inf")
+    self.max_failure_count = max_failure_count
+
+    if max_reshard_retry_count is None:
+      max_reshard_retry_count = float("inf")
+    self.max_reshard_retry_count = max_reshard_retry_count
 
     self.failure_count = 0
+    self.reshard_retry_count = 0
     self.good_slice_indices = self.get_slice_availability()
     self.data = {}
 
-  def slice_down(self):
+  def slice_down(self, reshard_retry: bool = False):
     """Slice down."""
     logger.info("Slice down")
     self.good_slice_indices = self.get_slice_availability()
     self.failure_count += 1
+    if reshard_retry:
+      self.reshard_retry_count += 1
+    else:
+      self.reshard_retry_count = 0
+
+    logger.info(f"{self.failure_count=} {self.max_failure_count=}")
+    if self.failure_count >= self.max_failure_count:
+      logger.fatal(f"Max failure count reached {self.max_failure_count}")
 
     logger.info(
-        f"Failure count: {self.failure_count} with max {self.max_failures}"
+        f"{self.reshard_retry_count=} {self.max_reshard_retry_count=}"
     )
-    if self.failure_count >= self.max_failures:
-      logger.fatal(f"Max failures reached {self.max_failures}")
+    if self.reshard_retry_count > self.max_reshard_retry_count:
+      logger.fatal(
+          f"Max reshard retry count reached {self.max_reshard_retry_count}"
+      )
 
   @timeit
   def save(self, save_step: int, blocking: bool = True, **kwargs):
@@ -486,7 +502,6 @@ def put_array_device_put3(
         arr.shape, dst_sharding, arrays
     )
 
-
   def scale_by_good_slices(self, x: int | float) -> int | float:
     """Scale x by the number of good slices."""
     if isinstance(x, int):

MaxText/elasticutils_fake.py

@@ -28,7 +28,8 @@ def __init__(
       total_slice_count: int,
       save_period: Optional[int] = None,
       reshard_check_period: Optional[int] = None,
-      max_failures: Optional[int] = None,
+      max_failure_count: Optional[int] = None,
+      max_reshard_retry_count: Optional[int] = None,
   ):
     self.fake_good_slice_indices = set(d.slice_index for d in devices)
 
@@ -38,6 +39,7 @@ def __init__(
         save_period,
         reshard_check_period,
         max_failures,
+        max_reshard_retry_count,
     )
 
   def update_good_slice_indices(self, good_slice_indices: set[int]):

MaxText/train.py

@@ -863,7 +863,7 @@ def reshard(arr):
         max_logging.log("Unknown JaxRuntimeError during resharding!")
         raise
 
-      config.eu.slice_down()
+      config.eu.slice_down(reshard_retry=True)
 
   return (
       restore_step,