Add effective_batch_size to auto-adjust gradient accumulation

ludwig/api.py

@@ -613,11 +613,7 @@ def on_epoch_end(self, trainer, progress_tracker, save_path):
                 random_seed=random_seed,
             ) as trainer:
                 # auto tune batch size
-                if (
-                    self.config_obj.trainer.to_dict().get(BATCH_SIZE, None) == AUTO
-                    or self.config_obj.trainer.to_dict().get(EVAL_BATCH_SIZE, None) == AUTO
-                ):
-                    self._tune_batch_size(trainer, training_set, random_seed=random_seed)
+                self._tune_batch_size(trainer, training_set, random_seed=random_seed)
 
                 # train model
                 if self.backend.is_coordinator():
@@ -795,35 +791,60 @@ def train_online(
                 config=self.config_obj.trainer, model=self.model, random_seed=random_seed
             )
 
-            if (
-                self.config_obj.trainer.to_dict().get(BATCH_SIZE, None) == AUTO
-                or self.config_obj.trainer.to_dict().get(EVAL_BATCH_SIZE, None) == AUTO
-            ):
-                self._tune_batch_size(self._online_trainer, dataset, random_seed=random_seed)
+            self._tune_batch_size(self._online_trainer, dataset, random_seed=random_seed)
 
         self.model = self._online_trainer.train_online(training_dataset)
 
     def _tune_batch_size(self, trainer, dataset, random_seed: int = default_random_seed):
+        if not self.config_obj.trainer.can_tune_batch_size():
+            # Models like GBMs don't have batch sizes to be tuned
+            return
+
+        # Render the batch size and gradient accumulation steps prior to batch size tuning. This is needed in the event
+        # the effective_batch_size and gradient_accumulation_steps are set explicitly, but batch_size is AUTO. In this
+        # case, we can infer the batch_size directly without tuning.
+        num_workers = self.backend.num_training_workers
+        self.config_obj.trainer.update_batch_size_grad_accum(num_workers)
+
         # TODO (ASN): add support for substitute_with_max parameter
         # TODO(travis): detect train and eval batch sizes separately (enable / disable gradients)
-        if self.backend.supports_batch_size_tuning():
-            tuned_batch_size = trainer.tune_batch_size(self.config_obj.to_dict(), dataset, random_seed=random_seed)
-        else:
-            logger.warning(
-                f"Backend {self.backend.BACKEND_TYPE} does not support batch size tuning, "
-                f"using fallback batch size {FALLBACK_BATCH_SIZE}."
-            )
-            tuned_batch_size = FALLBACK_BATCH_SIZE
-
-        # TODO(travis): pass these in as args to trainer when we call train,
-        #  to avoid setting state on possibly remote trainer
         if self.config_obj.trainer.batch_size == AUTO:
+            if self.backend.supports_batch_size_tuning():
+                tuned_batch_size = trainer.tune_batch_size(
+                    self.config_obj.to_dict(), dataset, random_seed=random_seed, tune_for_training=True
+                )
+            else:
+                logger.warning(
+                    f"Backend {self.backend.BACKEND_TYPE} does not support batch size tuning, "
+                    f"using fallback training batch size {FALLBACK_BATCH_SIZE}."
+                )
+                tuned_batch_size = FALLBACK_BATCH_SIZE
+
+            # TODO(travis): pass these in as args to trainer when we call train,
+            #  to avoid setting state on possibly remote trainer
             self.config_obj.trainer.batch_size = tuned_batch_size
-            trainer.batch_size = tuned_batch_size
+
+            # Re-render the gradient_accumulation_steps to account for the explicit batch size.
+            self.config_obj.trainer.update_batch_size_grad_accum(num_workers)
 
         if self.config_obj.trainer.eval_batch_size in {AUTO, None}:
+            if self.backend.supports_batch_size_tuning():
+                tuned_batch_size = trainer.tune_batch_size(
+                    self.config_obj.to_dict(), dataset, random_seed=random_seed, tune_for_training=False
+                )
+            else:
+                logger.warning(
+                    f"Backend {self.backend.BACKEND_TYPE} does not support batch size tuning, "
+                    f"using fallback eval batch size {FALLBACK_BATCH_SIZE}."
+                )
+                tuned_batch_size = FALLBACK_BATCH_SIZE
+
             self.config_obj.trainer.eval_batch_size = tuned_batch_size
-            trainer.eval_batch_size = tuned_batch_size
+
+        # Update trainer params separate to config params for backends with stateful trainers
+        trainer.batch_size = self.config_obj.trainer.batch_size
+        trainer.eval_batch_size = self.config_obj.trainer.eval_batch_size
+        trainer.gradient_accumulation_steps = self.config_obj.trainer.gradient_accumulation_steps
 
     def predict(
         self,

ludwig/backend/base.py

@@ -118,6 +118,11 @@ def read_binary_files(self, column: Series, map_fn: Optional[Callable] = None) -
     def num_nodes(self) -> int:
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def num_training_workers(self) -> int:
+        raise NotImplementedError()
+
     @abstractmethod
     def get_available_resources(self) -> Resources:
         raise NotImplementedError()
@@ -250,6 +255,10 @@ def __init__(self, **kwargs):
     def num_nodes(self) -> int:
         return 1
 
+    @property
+    def num_training_workers(self) -> int:
+        return 1
+
     def get_available_resources(self) -> Resources:
         return Resources(cpus=psutil.cpu_count(), gpus=torch.cuda.device_count())
 
@@ -308,6 +317,10 @@ def is_coordinator(self):
 
     @property
     def num_nodes(self) -> int:
+        return self._distributed.size() // self._distributed.local_size()
+
+    @property
+    def num_training_workers(self) -> int:
         return self._distributed.size()
 
     def get_available_resources(self) -> Resources:

ludwig/backend/ray.py

@@ -250,6 +250,7 @@ def tune_batch_size_fn(
     training_set_metadata: TrainingSetMetadataDict = None,
     features: Dict[str, Dict] = None,
     remote_trainer_cls: Callable[[], Trainer] = None,
+    tune_for_training: bool = True,
     **kwargs,
 ):
     # Pin GPU before loading the model to prevent memory leaking onto other devices
@@ -276,6 +277,7 @@ def on_best_batch_size_updated(best_batch_size: int, best_samples_per_sec: float
             train_shard,
             snapshot_weights=False,
             on_best_batch_size_updated=on_best_batch_size_updated,
+            tune_for_training=tune_for_training,
             **kwargs,
         )
         session.report(
@@ -539,6 +541,7 @@ def tune_batch_size(
         self,
         config: ModelConfigDict,
         training_set: RayDataset,
+        tune_for_training: bool = True,
         **kwargs,
     ) -> int:
         with create_runner(**self.trainer_kwargs) as runner:
@@ -552,6 +555,7 @@ def tune_batch_size(
                     ludwig_config=config,
                     training_set_metadata=training_set.training_set_metadata,
                     features=training_set.features,
+                    tune_for_training=tune_for_training,
                     **kwargs,
                 ),
                 exception_on_error=False,
@@ -592,6 +596,14 @@ def eval_batch_size(self) -> int:
     def eval_batch_size(self, value: int):
         self.config.eval_batch_size = value
 
+    @property
+    def gradient_accumulation_steps(self) -> int:
+        return self.config.gradient_accumulation_steps
+
+    @gradient_accumulation_steps.setter
+    def gradient_accumulation_steps(self, value: int):
+        self.config.gradient_accumulation_steps = value
+
     @property
     def resources_per_worker(self) -> Dict[str, Any]:
         trainer_kwargs = get_trainer_kwargs(**self.trainer_kwargs)
@@ -876,7 +888,7 @@ def __init__(
         super().__init__(dataset_manager=RayDatasetManager(self), **kwargs)
         self._preprocessor_kwargs = preprocessor_kwargs or {}
         self._df_engine = _get_df_engine(processor)
-        self._horovod_kwargs = trainer or {}
+        self._distributed_kwargs = trainer or {}
         self._pytorch_kwargs = {}
         self._data_loader_kwargs = loader or {}
         self._preprocessor_pg = None
@@ -943,7 +955,7 @@ def create_trainer(self, model: BaseModel, **kwargs) -> "BaseTrainer":  # noqa:
 
         all_kwargs = {
             "model": model,
-            "trainer_kwargs": self._horovod_kwargs,
+            "trainer_kwargs": self._distributed_kwargs,
             "data_loader_kwargs": self._data_loader_kwargs,
             "executable_kwargs": executable_kwargs,
         }
@@ -956,18 +968,18 @@ def create_predictor(self, model: BaseModel, **kwargs):
         return RayPredictor(
             model,
             self.df_engine,
-            self._horovod_kwargs,
+            self._distributed_kwargs,
             self._data_loader_kwargs,
             **executable_kwargs,
         )
 
     @property
     def distributed_kwargs(self):
-        return self._horovod_kwargs
+        return self._distributed_kwargs
 
     @distributed_kwargs.setter
     def distributed_kwargs(self, value):
-        self._horovod_kwargs = value
+        self._distributed_kwargs = value
 
     @property
     def df_engine(self):
@@ -1068,6 +1080,11 @@ def num_nodes(self) -> int:
             return 1
         return len(ray.nodes())
 
+    @property
+    def num_training_workers(self) -> int:
+        trainer_kwargs = get_trainer_kwargs(**self._distributed_kwargs)
+        return trainer_kwargs["num_workers"]
+
     def get_available_resources(self) -> Resources:
         resources = ray.cluster_resources()
         return Resources(cpus=resources.get("CPU", 0), gpus=resources.get("GPU", 0))
@@ -1122,7 +1139,7 @@ def batch_transform(self, df: DataFrame, batch_size: int, transform_fn: Callable
             return self.df_engine.from_ray_dataset(ds)
 
     def _get_transform_kwargs(self) -> Dict[str, Any]:
-        trainer_kwargs = get_trainer_kwargs(**self._horovod_kwargs)
+        trainer_kwargs = get_trainer_kwargs(**self._distributed_kwargs)
         resources_per_worker = trainer_kwargs.get("resources_per_worker", {})
         num_gpus = resources_per_worker.get("GPU", 0)
         num_cpus = resources_per_worker.get("CPU", (1 if num_gpus == 0 else 0))

ludwig/distributed/horovod.py

@@ -13,6 +13,7 @@
 from torch import nn
 from torch.optim import Optimizer
 
+from ludwig.constants import AUTO
 from ludwig.distributed.base import DistributedStrategy
 from ludwig.modules.optimization_modules import create_optimizer
 from ludwig.utils.horovod_utils import gather_all_tensors, is_distributed_available
@@ -35,10 +36,13 @@ def prepare(
         base_learning_rate: float,
     ) -> Tuple[nn.Module, Optimizer]:
         optimizer = create_optimizer(model, trainer_config.optimizer, base_learning_rate)
+        grad_accum_steps = (
+            trainer_config.gradient_accumulation_steps if trainer_config.gradient_accumulation_steps != AUTO else 1
+        )
         dist_optimizer = hvd.DistributedOptimizer(
             optimizer,
             named_parameters=model.named_parameters(),
-            backward_passes_per_step=trainer_config.gradient_accumulation_steps,
+            backward_passes_per_step=grad_accum_steps,
         )
         return model, dist_optimizer
 

ludwig/schema/metadata/configs/trainer.yaml

@@ -1,4 +1,11 @@
 ecd:
+    effective_batch_size:
+        commonly_used: true
+        expected_impact: 2
+        related_parameters:
+            - batch_size
+        suggested_values: auto
+        ui_display_name: Effective Batch Size
     batch_size:
         commonly_used: true
         default_value_reasoning: Not too big, not too small.