Commit a442244d authored Nov 19, 2018 by Halil Akin Committed by Facebook Github Bot Nov 19, 2018

Protect against failures in case of OOMs

Summary: Fixing some distributed failures that happen when OOMs are observed.

Reviewed By: myleott

Differential Revision: D13121054

fbshipit-source-id: f71a0a695332acbaa1797e89887b8b7c7ddaa727

parent 693894b6

fairseq/optim/fairseq_optimizer.py

+2 −1

Original line number	Diff line number	Diff line
		@@ -75,6 +75,7 @@ class FairseqOptimizer(object):
		def multiply_grads(self, c):
		"""Multiplies grads by a constant ``c``."""
		for p in self.params:
		if p.grad is not None:
		p.grad.data.mul_(c)

		def clip_grad_norm(self, max_norm):
		@@ -82,7 +83,7 @@ class FairseqOptimizer(object):
		if max_norm > 0:
		return torch.nn.utils.clip_grad_norm_(self.params, max_norm)
		else:
		return math.sqrt(sum(p.grad.data.norm()**2 for p in self.params))
		return math.sqrt(sum(p.grad.data.norm()**2 for p in self.params if p.grad is not None))

		def step(self, closure=None):
		"""Performs a single optimization step."""

fairseq/trainer.py

+1 −1

Original line number	Diff line number	Diff line
		@@ -210,7 +210,7 @@ class Trainer(object):
		sample_sizes = list(chain.from_iterable(sample_sizes))
		ooms = sum(ooms)

		if ooms == self.args.distributed_world_size:
		if ooms == self.args.distributed_world_size * len(samples):
		print('\| WARNING: OOM in all workers, skipping update')
		self.zero_grad()
		return None