Variable batch size in Multi-GPU trainings

Ryan_Xu: <p>I am trying to run a multi-gpu training where batch size could be not fixed across each GPU. To make a fixed batch size, I upsample the given batch to make sure it contains a fixed number of datapoints in the collate function. However, during training, I am seeing a tons of the following message in my log:</p>
<blockquote>
<p><code>[trainer:0:4]:[rank4]:[I728 20:47:12.342890349 ProcessGroupWrapper.cpp:587] [Rank 4] Running collective: CollectiveFingerPrint(SequenceNumber=10883, OpType=ALLREDUCE, TensorShape=[7077904], TensorDtypes=Float, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)))</code></p>
</blockquote>
<p>The training still runs but I am a little unsure why this message pops up in the first place, and if I am actually missing something..?</p>

ptrblck: <p>I guess you have enabled logging as this log should come from <a href="https://github.com/pytorch/pytorch/blob/fd2c64e286b7787080cd9028cc69561c769ffae4/torch/csrc/distributed/c10d/ProcessGroupWrapper.cpp#L586-L587">here</a>. Do you see any real errors in your use case?</p>

Ryan_Xu: <p>The training ran okay for a few thousands steps until it crashed with error messages like the following. I am not entirely sure what went wrong. In the end I move the logic that expands and upsamples a given batch to a fixed size from the collate function to inside the training loop (basically where model.predict happens), and it looks like the issue is no longer there.</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “/root/w/helm/car_demo/hcore/hcore/gcp/batch/worker_entrypoint.py”, line 101, in run_pmap<br>
function(**kwargs)<br>
File “/root/w/helm/planning/deterministic_planner/training_sweep.py”, line 177, in execute_training_task<br>
init_and_launch(<br>
File “/root/w/helm/planning/deterministic_planner/training_sweep.py”, line 64, in init_and_launch<br>
dataclasses.replace(reconstructed, **changes).launch_training()<br>
File “/root/w/helm/planning/common/training_job_config.py”, line 367, in launch_training<br>
training_task.execute()<br>
File “/root/w/helm/percept_training/task.py”, line 138, in execute<br>
self._run(trainer)<br>
File “/root/w/helm/percept_training/task.py”, line 289, in _run<br>
trainer.fit(<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py”, line 544, in fit<br>
call._call_and_handle_interrupt(<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py”, line 43, in _call_and_handle_interrupt<br>
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py”, line 102, in launch<br>
return function(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py”, line 580, in _fit_impl<br>
self._run(model, ckpt_path=ckpt_path)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py”, line 989, in _run<br>
results = self._run_stage()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py”, line 1035, in _run_stage<br>
self.fit_loop.run()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py”, line 202, in run<br>
self.advance()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py”, line 359, in advance<br>
self.epoch_loop.run(self._data_fetcher)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py”, line 136, in run<br>
self.advance(data_fetcher)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py”, line 240, in advance<br>
batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py”, line 187, in run<br>
self._optimizer_step(batch_idx, closure)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py”, line 265, in _optimizer_step<br>
call._call_lightning_module_hook(<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py”, line 157, in _call_lightning_module_hook<br>
output = fn(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/core/module.py”, line 1282, in optimizer_step<br>
optimizer.step(closure=optimizer_closure)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/core/optimizer.py”, line 151, in step<br>
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/strategies/ddp.py”, line 264, in optimizer_step<br>
optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py”, line 230, in optimizer_step<br>
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/amp.py”, line 74, in optimizer_step<br>
return super().optimizer_step(optimizer, model=model, closure=closure, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision.py”, line 117, in optimizer_step<br>
return optimizer.step(closure=closure, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/optim/lr_scheduler.py”, line 137, in wrapper<br>
return func.<strong>get</strong>(opt, opt.<strong>class</strong>)(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/optim/optimizer.py”, line 487, in wrapper<br>
out = func(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/adabelief_pytorch/AdaBelief.py”, line 136, in step<br>
loss = closure()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision.py”, line 104, in _wrap_closure<br>
closure_result = closure()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py”, line 140, in <strong>call</strong><br>
self._result = self.closure(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/utils/_contextlib.py”, line 116, in decorate_context<br>
return func(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py”, line 126, in closure<br>
step_output = self._step_fn()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py”, line 315, in _training_step<br>
training_step_output = call._call_strategy_hook(trainer, “training_step”, *kwargs.values())<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py”, line 309, in _call_strategy_hook<br>
output = fn(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py”, line 381, in training_step<br>
return self._forward_redirection(self.model, self.lightning_module, “training_step”, *args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py”, line 633, in <strong>call</strong><br>
wrapper_output = wrapper_module(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1736, in _wrapped_call_impl<br>
return self._call_impl(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1747, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 1639, in forward<br>
inputs, kwargs = self._pre_forward(*inputs, **kwargs)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 1535, in _pre_forward<br>
self._sync_buffers()<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 2172, in _sync_buffers<br>
self._sync_module_buffers(authoritative_rank)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 2176, in _sync_module_buffers<br>
self._default_broadcast_coalesced(authoritative_rank=authoritative_rank)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 2198, in _default_broadcast_coalesced<br>
self._distributed_broadcast_coalesced(bufs, bucket_size, authoritative_rank)<br>
File “/opt/pyenv/versions/3.10.6/lib/python3.10/site-packages/torch/nn/parallel/distributed.py”, line 2113, in _distributed_broadcast_coalesced<br>
dist._broadcast_coalesced(<br>
RuntimeError: [Rank 0]: Ranks 1 failed to pass monitoredBarrier in 900000 ms</p>
</blockquote>

ptrblck: <aside class="quote no-group" data-username="Ryan_Xu" data-post="3" data-topic="221940">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ryan_xu/48/77825_2.png" class="avatar"> Ryan_Xu:</div>
<blockquote>
<p>In the end I move the logic that expands and upsamples a given batch to a fixed size from the collate function to inside the training loop (basically where model.predict happens), and it looks like the issue is no longer there.</p>
</blockquote>
</aside>
<p>It’s good to hear you are not seeing the timeout anymore. The issue sounds like a hang caused by an imbalanced number of batches between the ranks which would cause such issues. We thus recommend using e.g. the <code>DistributedSampler</code> which will make sure all ranks receive the same number of batches. Since you already fixed it yourself you could of course also stay with your approach.</p>