Trainer parameter limit-train-batches was meant to be per-worker

[wingkitlee0]

📚 Documentation

When using lightning with RayDDPStrategy, we found that Trainer's limit_train_batches parameter was actually meant for each worker instead of "global".

I wanna confirm with the developer if

1. this is the case with other parallel training strategies
2. this applies to other limit_*_batches parameters of the Trainer.

It will be great to update the doc to mention this, e.g.,
https://lightning.ai/docs/pytorch/stable/common/trainer.html#limit-train-batches

Thanks

cc @lantiga @Borda @justusschock

[SkafteNicki]

Hi @wingkitlee0, the short answer is yes and yes. The limit applies to all parallel training strategies and it is the same for all limit_*_batches parameters. Just to give a bit of code:

This is the relevant inner code that is executed per device when training:

pytorch-lightning/src/lightning/pytorch/loops/training_epoch_loop.py

Lines 147 to 157 in 6a09f27

	def run(self, data_fetcher: _DataFetcher) -> None:
	self.reset()
	self.on_run_start(data_fetcher)
	while not self.done:
	try:
	self.advance(data_fetcher)
	self.on_advance_end(data_fetcher)
	except StopIteration:
	break
	finally:
	self.on_iteration_done()

where in the self.done method is ending up calling this method, that checks if num_training_batches is reached

pytorch-lightning/src/lightning/pytorch/loops/training_epoch_loop.py

Lines 434 to 437 in 6a09f27

	def _num_ready_batches_reached(self) -> bool:
	"""Checks if we are in the last batch or if there are more batches to follow."""
	epoch_finished_on_ready = self.batch_progress.current.ready == self.trainer.num_training_batches
	return epoch_finished_on_ready or self.batch_progress.is_last_batch

num_training_batches are essentially calculated by this function

pytorch-lightning/src/lightning/pytorch/trainer/connectors/data_connector.py

Lines 447 to 474 in 6a09f27

	def _parse_num_batches(
	stage: RunningStage, length: Union[int, float], limit_batches: Union[int, float]
	) -> Union[int, float]:
	if length == 0:
	return int(length)
	num_batches = length
	# limit num batches either as a percent or num steps
	if isinstance(limit_batches, int):
	num_batches = min(length, limit_batches)
	elif isinstance(limit_batches, float) and length != float("inf"):
	num_batches = int(length * limit_batches)
	elif limit_batches != 1.0:
	raise MisconfigurationException(
	f"When using an `IterableDataset`, `Trainer(limit_{stage.dataloader_prefix}_batches)` must be"
	f" `1.0` or an int. An int specifies `num_{stage.dataloader_prefix}_batches` to use."
	)
	if num_batches == 0 and limit_batches > 0.0 and isinstance(limit_batches, float) and length != float("inf"):
	min_percentage = 1.0 / length
	raise MisconfigurationException(
	f"You requested to check {limit_batches} of the `{stage.dataloader_prefix}_dataloader` but"
	f" {limit_batches} * {length} < 1. Please increase the"
	f" `limit_{stage.dataloader_prefix}_batches` argument. Try at least"
	f" `limit_{stage.dataloader_prefix}_batches={min_percentage}`"
	)
	return num_batches

this function does not account in anyway for the number of devices uses by the training, so that is another way of seeing that the limit is applied per device and not globally.