Simulated big batches (2d27ae08) · Commits · Simon Will / fairseq

fairseq/options.py

+2 −0

Changes for fairseq/options.py: 2 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -176,6 +176,8 @@ def add_optimization_args(parser):
		' dataset')
		group.add_argument('--curriculum', default=0, type=int, metavar='N',
		help='sort batches by source length for first N epochs')
		group.add_argument('--update-freq', default=1, type=int, metavar='N',
		help='update parameters every N batches')
		return group

fairseq/trainer.py

+54 −23

Changes for fairseq/trainer.py: 54 added lines, 23 removed lines.
Original line number	Diff line number	Diff line
		@@ -9,7 +9,8 @@
		Train a network on multiple GPUs.
		"""

		from collections import OrderedDict
		from collections import defaultdict, OrderedDict
		from itertools import chain
		import math
		import torch

		@@ -55,6 +56,7 @@ class Trainer(object):
		self.meters['clip'] = AverageMeter() # % of updates clipped
		self.meters['oom'] = AverageMeter() # out of memory

		self._buffered_stats = defaultdict(lambda: [])
		self._max_bsz_seen = 0
		self._num_updates = 0
		self._optim_history = None
		@@ -86,22 +88,46 @@ class Trainer(object):

		return extra_state

		def train_step(self, sample):
		def train_step(self, sample, update_params=True):
		"""Do forward, backward and parameter update."""

		sample = self._prepare_sample(sample, volatile=False)

		# forward pass
		loss, sample_sizes, logging_outputs, ooms_fwd = self._forward(sample)
		# forward and backward pass
		loss, sample_size, logging_output, oom_fwd = self._forward(sample)
		oom_bwd = self._backward(loss)

		# buffer stats and logging outputs
		self._buffered_stats['sample_sizes'].append(sample_size)
		self._buffered_stats['logging_outputs'].append(logging_output)
		self._buffered_stats['ooms_fwd'].append(oom_fwd)
		self._buffered_stats['ooms_bwd'].append(oom_bwd)

		# update parameters
		if update_params:
		# gather logging outputs from all GPUs
		sample_sizes = self._buffered_stats['sample_sizes']
		logging_outputs = self._buffered_stats['logging_outputs']
		ooms_fwd = self._buffered_stats['ooms_fwd']
		ooms_bwd = self._buffered_stats['ooms_bwd']
		if self.args.distributed_world_size > 1:
		sample_sizes, logging_outputs, ooms_fwd, ooms_bwd = map(
		lambda l: list(chain.from_iterable(l)),
		zip(*distributed_utils.all_gather_list(
		(sample_sizes, logging_outputs, ooms_fwd, ooms_bwd)
		))
		)
		ooms_fwd = sum(ooms_fwd)
		ooms_bwd = sum(ooms_bwd)

		# aggregate stats and logging outputs
		ntokens = sum(log.get('ntokens', 0) for log in logging_outputs)
		nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)
		grad_denom = self.criterion.__class__.grad_denom(sample_sizes)
		agg_logging_output = self.criterion.__class__.aggregate_logging_outputs(logging_outputs)

		# backward pass, all-reduce gradients and take an optimization step
		grad_norm, ooms_bwd = self._backward_and_opt(loss, grad_denom)
		# all-reduce gradients and take an optimization step
		grad_denom = self.criterion.__class__.grad_denom(sample_sizes)
		grad_norm = self._opt(grad_denom)

		# update meters
		self.meters['wps'].update(ntokens)
		@@ -119,7 +145,11 @@ class Trainer(object):
		if 'nll_loss' in agg_logging_output:
		self.meters['train_nll_loss'].update(agg_logging_output['nll_loss'], ntokens)

		self._buffered_stats.clear()

		return agg_logging_output
		else:
		return None # buffering updates

		def _forward(self, sample, eval=False):
		# prepare model and optimizer
		@@ -127,7 +157,6 @@ class Trainer(object):
		self.model.eval()
		else:
		self.model.train()
		self.optimizer.zero_grad()

		loss = None
		sample_size = 0
		@@ -152,19 +181,9 @@ class Trainer(object):
		else:
		raise e

		# synchronize logging outputs for multi-GPU training
		if self.args.distributed_world_size > 1:
		sample_sizes, logging_outputs, ooms = zip(*list(
		distributed_utils.all_gather_list((sample_size, logging_output, oom))))
		ooms = sum(ooms)
		else:
		sample_sizes = [sample_size]
		logging_outputs = [logging_output]
		ooms = oom

		return loss, sample_sizes, logging_outputs, ooms
		return loss, sample_size, logging_output, oom

		def _backward_and_opt(self, loss, grad_denom):
		def _backward(self, loss):
		oom = 0
		if loss is not None:
		try:
		@@ -179,7 +198,9 @@ class Trainer(object):
		self.optimizer.zero_grad()
		else:
		raise e
		return oom

		def _opt(self, grad_denom):
		# all-reduce grads and rescale by grad_denom
		if self.args.distributed_world_size > 1:
		grads = [p.grad.data for p in self.model.parameters() if p.requires_grad]
		@@ -197,12 +218,13 @@ class Trainer(object):

		# take an optimization step
		self.optimizer.step()
		self.optimizer.zero_grad()
		self._num_updates += 1

		# update learning rate
		self.lr_scheduler.step_update(self._num_updates)

		return grad_norm, oom
		return grad_norm

		def valid_step(self, sample):
		"""Do forward pass in evaluation mode."""
		@@ -210,8 +232,17 @@ class Trainer(object):
		sample = self._prepare_sample(sample, volatile=True)

		# forward pass
		loss, sample_sizes, logging_outputs, ooms_fwd = self._forward(sample, eval=True)
		assert not ooms_fwd, 'Ran out of memory during validation'
		_loss, sample_size, logging_output, oom_fwd = self._forward(sample, eval=True)
		assert not oom_fwd, 'Ran out of memory during validation'

		# gather logging outputs from all GPUs
		if self.args.distributed_world_size > 1:
		sample_sizes, logging_outputs = zip(*distributed_utils.all_gather_list(
		(sample_size, logging_output)
		))
		else:
		sample_sizes = [sample_size]
		logging_outputs = [logging_output]

		# aggregate stats and logging outputs
		ntokens = sum(log.get('ntokens', 0) for log in logging_outputs)

singleprocess_train.py

+8 −3

Changes for singleprocess_train.py: 8 added lines, 3 removed lines.
Original line number	Diff line number	Diff line
		@@ -132,6 +132,7 @@ def train(args, trainer, dataset, epoch, batch_offset):
		num_shards=args.distributed_world_size,
		)
		progress = progress_bar.build_progress_bar(args, itr, epoch, no_progress_bar='simple')
		epoch_size = len(itr)
		itr = itertools.islice(progress, batch_offset, None)

		# reset training meters
		@@ -143,7 +144,12 @@ def train(args, trainer, dataset, epoch, batch_offset):
		extra_meters = collections.defaultdict(lambda: AverageMeter())
		max_update = args.max_update or math.inf
		for i, sample in enumerate(itr, start=batch_offset):
		log_output = trainer.train_step(sample)
		if i < epoch_size - 1 and (i + 1) % args.update_freq > 0:
		# buffer updates according to --update-freq
		trainer.train_step(sample, update_params=False)
		continue
		else:
		log_output = trainer.train_step(sample, update_params=True)

		# log mid-epoch stats
		stats = get_training_stats(trainer)
		@@ -157,9 +163,8 @@ def train(args, trainer, dataset, epoch, batch_offset):
		stats[k] = extra_meters[k].avg
		progress.log(stats)

		# save mid-epoch checkpoints
		if i == batch_offset:
		# ignore the first mini-batch in words-per-second calculation
		if i == batch_offset:
		trainer.get_meter('wps').reset()

		# save mid-epoch checkpoints