Fixing error condition check when device microbatch size times seq pa…

…rallelism dim is not 1 due to floating point precision (mosaicml#3200) * .. * .. * lint
mvpatel2000 · Apr 19, 2024 · 7b25d90 · 7b25d90
1 parent 960654d
commit 7b25d90
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
@@ -409,7 +409,9 @@ def _validate_evaluator(evaluator: Evaluator, device: Device):
     if hasattr(
         evaluator.dataloader,
         'seq_parallel_world_size',
-    ) and evaluator.dataloader.seq_parallel_world_size > 1 and evaluator.dataloader.device_eval_batch_size * evaluator.dataloader.seq_parallel_world_size != 1:  # type: ignore
+    ) and evaluator.dataloader.seq_parallel_world_size > 1 and abs(  # type: ignore
+        evaluator.dataloader.device_eval_batch_size * evaluator.dataloader.seq_parallel_world_size - 1,  # type: ignore
+    ) > 1e-4:
         raise ValueError(
             'Sequence parallelism requires a microbatch size of 1 distributed over the sequence parallel group.',
         )
@@ -1126,7 +1128,9 @@ def __init__(
         if train_dataloader is not None and hasattr(
             train_dataloader,
             'seq_parallel_world_size',
-        ) and train_dataloader.seq_parallel_world_size > 1 and device_train_microbatch_size * train_dataloader.seq_parallel_world_size != 1:  # type: ignore
+        ) and train_dataloader.seq_parallel_world_size > 1 and abs( # type: ignore
+            device_train_microbatch_size * train_dataloader.seq_parallel_world_size - 1, # type: ignore
+        ) > 1e-4:
             raise ValueError(
                 '`Sequence parallelism requires a microbatch size of 1 distributed over the sequence parallel group.',
             )
@@ -2181,7 +2185,9 @@ def fit(
             if train_dataloader is not None and hasattr(
                 train_dataloader,
                 'seq_parallel_world_size',
-            ) and train_dataloader.seq_parallel_world_size > 1 and device_train_microbatch_size * train_dataloader.seq_parallel_world_size != 1:  # type: ignore
+            ) and train_dataloader.seq_parallel_world_size > 1 and abs(  # type: ignore
+                device_train_microbatch_size * train_dataloader.seq_parallel_world_size - 1, # type: ignore
+            ) > 1e-4:
                 raise ValueError(
                     '`Sequence parallelism requires a microbatch size of 1 distributed over the sequence parallel group.',
                 )