Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Skips DDP parameter sync (#4301) 

* ddp no-sync

* Update pytorch_lightning/trainer/training_loop.py

Co-authored-by: ananthsub <ananth.subramaniam@gmail.com>

* Update training_loop.py

* factor __enter__ and __exit__ out to separate context manager

* delete _updated_model_last_step

Co-authored-by: justusschock <justusschock@pc125.lfb.rwth-aachen.de>
Co-authored-by: Teddy Koker <teddy.koker@gmail.com>
Co-authored-by: ananthsub <ananth.subramaniam@gmail.com>
Co-authored-by: chaton <thomas@grid.ai>
Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>
This commit is contained in:
Justus Schock 2020-10-29 18:31:37 +01:00 committed by GitHub
parent b459fd26ac
commit bbd81dfd55
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
pytorch_lightning/trainer/training_loop.py
View File

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import subprocess from contextlib import contextmanager
from copy import copy, deepcopy from copy import copy, deepcopy
import numpy as np import numpy as np
@ -655,6 +655,7 @@ class TrainLoop:
# checks if backward or backward + optimizer step (via closure) # checks if backward or backward + optimizer step (via closure)
accumulation_done = self._accumulated_batches_reached() accumulation_done = self._accumulated_batches_reached()
is_final_batch = self._num_training_batches_reached() is_final_batch = self._num_training_batches_reached()
should_accumulate = not (accumulation_done or is_final_batch)
# lightning module hook # lightning module hook
splits = self.tbptt_split_batch(batch) splits = self.tbptt_split_batch(batch)
@ -675,13 +676,17 @@ class TrainLoop:
model = self.trainer.get_model() model = self.trainer.get_model()
model.toggle_optimizer(optimizer, opt_idx) model.toggle_optimizer(optimizer, opt_idx)
if not (accumulation_done or is_final_batch): if should_accumulate:
# For gradient accumulation # For gradient accumulation
# ------------------- # -------------------
# calculate loss (train step + train step end) # calculate loss (train step + train step end)
# ------------------- # -------------------
self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)
# perform dpp sync only when performing optimizer_step
with self.block_ddp_sync_behaviour():
self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)
batch_outputs = self._process_closure_result( batch_outputs = self._process_closure_result(
batch_callback_metrics=batch_callback_metrics, batch_callback_metrics=batch_callback_metrics,
batch_log_metrics=batch_log_metrics, batch_log_metrics=batch_log_metrics,
@ -695,7 +700,6 @@ class TrainLoop:
# gradient update with accumulated gradients # gradient update with accumulated gradients
else: else:
if self.automatic_optimization: if self.automatic_optimization:
def train_step_and_backward_closure(): def train_step_and_backward_closure():
@ -760,6 +764,13 @@ class TrainLoop:
) )
return result return result
@contextmanager
def block_ddp_sync_behaviour(self):
if isinstance(self.trainer.model, torch.nn.parallel.DistributedDataParallel):
yield from self.trainer.model.no_sync()
else:
yield
def _process_closure_result( def _process_closure_result(
self, batch_callback_metrics: list, batch_log_metrics: list, batch_outputs: list, opt_idx: int self, batch_callback_metrics: list, batch_log_metrics: list, batch_outputs: list, opt_idx: int
) -> list: ) -> list: