SubashSK777/Visual-Question-Answering
0
1# Copyright (c) Facebook, Inc. and its affiliates.2#3# This source code is licensed under the MIT license found in the4# LICENSE file in the root directory of this source tree.5 6"""7Train a network across multiple GPUs.8"""9 10import contextlib11import logging12import sys13import time14from argparse import Namespace15from itertools import chain16from typing import Any, Dict, List17 18import torch19from fairseq import models, optim, utils20from fairseq.dataclass.configs import FairseqConfig21from fairseq.dataclass.utils import convert_namespace_to_omegaconf22from fairseq.distributed import utils as distributed_utils23from fairseq.file_io import PathManager24from fairseq.logging import meters, metrics25from fairseq.models.ema import build_ema26from fairseq.nan_detector import NanDetector27from fairseq.optim import lr_scheduler28from omegaconf import OmegaConf29 30from utils import checkpoint_utils31 32logger = logging.getLogger(__name__)33 34 35class Trainer(object):36 """Main class for data parallel training.37 38 This class supports synchronous distributed data parallel training,39 where multiple workers each have a full model replica and gradients40 are accumulated across workers before each update. We use41 :class:`~torch.nn.parallel.DistributedDataParallel` to handle42 communication of the gradients across workers.43 """44 45 def __init__(self, cfg: FairseqConfig, task, model, criterion, quantizer=None):46 47 if isinstance(cfg, Namespace):48 logger.warning(49 "argparse.Namespace configuration is deprecated! Automatically converting to OmegaConf"50 )51 cfg = convert_namespace_to_omegaconf(cfg)52 53 self.cfg = cfg54 self.task = task55 56 # catalog shared parameters57 shared_params = _catalog_shared_params(model)58 self.tpu = cfg.common.tpu59 self.cuda = torch.cuda.is_available() and not cfg.common.cpu and not self.tpu60 if self.cuda:61 self.device = torch.device("cuda")62 elif self.tpu:63 self.device = utils.get_tpu_device()64 else:65 self.device = torch.device("cpu")66 67 if self.is_fsdp:68 import fairscale69 if self.cfg.common.bf16:70 raise ValueError(71 "FullyShardedDataParallel is not compatible with --bf16 or "72 "--memory-efficient-bf16"73 )74 if self.cfg.distributed_training.zero_sharding != "none":75 raise ValueError(76 "FullyShardedDataParallel is not compatible with --zero-sharding "77 "option (it's already built in)"78 )79 if max(self.cfg.optimization.update_freq) > 1 and fairscale.__version__ < "0.4.0":80 raise RuntimeError(81 "Please update to fairscale 0.4.0 or newer when combining "82 "--update-freq with FullyShardedDataParallel"83 )84 else:85 if (86 hasattr(self.cfg.distributed_training, "cpu_offload")87 and self.cfg.distributed_training.cpu_offload88 ):89 raise ValueError("--cpu-offload requires --ddp-backend=fully_sharded")90 91 # copy model and criterion to current device/dtype92 self._criterion = criterion93 self._model = model94 if not self.is_fsdp:95 if cfg.common.fp16:96 assert not cfg.common.amp, "Cannot use fp16 and AMP together"97 self._criterion = self._criterion.half()98 self._model = self._model.half()99 elif cfg.common.bf16:100 self._criterion = self._criterion.to(dtype=torch.bfloat16)101 self._model = self._model.to(dtype=torch.bfloat16)102 elif cfg.common.amp:103 self._amp_retries = 0104 if (105 not cfg.distributed_training.pipeline_model_parallel106 # the DistributedFairseqModel wrapper will handle moving to device,107 # so only handle cases which don't use the wrapper108 and not self.use_distributed_wrapper109 ):110 self._criterion = self._criterion.to(device=self.device)111 self._model = self._model.to(device=self.device)112 self.pipeline_model_parallel = cfg.distributed_training.pipeline_model_parallel113 self.last_device = None114 if self.cuda and self.pipeline_model_parallel:115 self.last_device = torch.device(116 cfg.distributed_training.pipeline_devices[-1]117 )118 119 # check that shared parameters are preserved after device transfer120 for shared_param in shared_params:121 ref = _get_module_by_path(self._model, shared_param[0])122 for path in shared_param[1:]:123 logger.info(124 "detected shared parameter: {} <- {}".format(shared_param[0], path)125 )126 _set_module_by_path(self._model, path, ref)127 128 self._dummy_batch = None # indicates we don't have a dummy batch at first129 self._lr_scheduler = None130 self._num_updates = 0131 self._num_xla_compiles = 0 # for TPUs132 self._optim_history = None133 self._optimizer = None134 self._warn_once = set()135 self._wrapped_criterion = None136 self._wrapped_model = None137 self._ema = None138 139 # TODO(myleott): support tpu140 if self.cuda and self.data_parallel_world_size > 1:141 self._grad_norm_buf = torch.cuda.DoubleTensor(self.data_parallel_world_size)142 else:143 self._grad_norm_buf = None144 145 self.quantizer = quantizer146 if self.quantizer is not None:147 self.quantizer.set_trainer(self)148 149 # get detailed cuda environment150 if self.cuda:151 self.cuda_env = utils.CudaEnvironment()152 if self.data_parallel_world_size > 1:153 self.cuda_env_arr = distributed_utils.all_gather_list(154 self.cuda_env, group=distributed_utils.get_global_group()155 )156 else:157 self.cuda_env_arr = [self.cuda_env]158 if self.data_parallel_rank == 0:159 utils.CudaEnvironment.pretty_print_cuda_env_list(self.cuda_env_arr)160 else:161 self.cuda_env = None162 self.cuda_env_arr = None163 164 metrics.log_start_time("wall", priority=790, round=0)165 166 self._start_time = time.time()167 self._previous_training_time = 0168 self._cumulative_training_time = None169 170 def reinitialize(self):171 """Reinitialize the Trainer, typically after model params change."""172 self._lr_scheduler = None173 self._optimizer = None174 self._wrapped_criterion = None175 self._wrapped_model = None176 177 @property178 def data_parallel_world_size(self):179 if self.cfg.distributed_training.distributed_world_size == 1:180 return 1181 return distributed_utils.get_data_parallel_world_size()182 183 @property184 def data_parallel_process_group(self):185 return distributed_utils.get_data_parallel_group()186 187 @property188 def data_parallel_rank(self):189 if self.cfg.distributed_training.distributed_world_size == 1:190 return 0191 return distributed_utils.get_data_parallel_rank()192 193 @property194 def is_data_parallel_master(self):195 # NOTE: this returns true for all model parallel replicas with data196 # parallel rank 0197 return self.data_parallel_rank == 0198 199 @property200 def use_distributed_wrapper(self) -> bool:201 return (202 self.data_parallel_world_size > 1 and not self.cfg.optimization.use_bmuf203 ) or (204 self.is_fsdp and self.cfg.distributed_training.cpu_offload205 )206 207 @property208 def should_save_checkpoint_on_current_rank(self) -> bool:209 """Indicates whether to save checkpoints on the current DDP rank."""210 if (211 self.is_fsdp and self.cfg.distributed_training.use_sharded_state212 ) or getattr(self.cfg.model, "base_layers", 0) > 0:213 return True214 else:215 return self.is_data_parallel_master216 217 @property218 def always_call_state_dict_during_save_checkpoint(self) -> bool:219 if self.is_fsdp and not self.cfg.distributed_training.use_sharded_state:220 # FSDP calls communication collective when consolidating checkpoints221 return True222 else:223 return False224 225 @property226 def checkpoint_suffix(self) -> str:227 """Suffix to add to the checkpoint file name."""228 if self.is_fsdp and self.cfg.distributed_training.use_sharded_state:229 return self.cfg.checkpoint.checkpoint_suffix + "-shard{0}".format(230 self.data_parallel_rank231 )232 else:233 return self.cfg.checkpoint.checkpoint_suffix or ""234 235 @property236 def criterion(self):237 if self._wrapped_criterion is None:238 if utils.has_parameters(self._criterion) and self.use_distributed_wrapper:239 self._wrapped_criterion = models.DistributedFairseqModel(240 self.cfg.distributed_training,241 self._criterion,242 process_group=self.data_parallel_process_group,243 device=self.device,244 )245 else:246 self._wrapped_criterion = self._criterion247 return self._wrapped_criterion248 249 @property250 def model(self):251 if self._wrapped_model is None:252 if self.use_distributed_wrapper:253 self._wrapped_model = models.DistributedFairseqModel(254 self.cfg.distributed_training,255 self._model,256 process_group=self.data_parallel_process_group,257 device=self.device,258 )259 else:260 self._wrapped_model = self._model261 return self._wrapped_model262 263 @property264 def ema(self):265 if self._ema is None:266 self._build_ema()267 return self._ema268 269 def _build_ema(self):270 if self.cfg.ema.store_ema:271 self._ema = build_ema(self._model, self.cfg.ema, self.device)272 logger.info(273 "Exponential Moving Average Shadow Model is initialized."274 )275 276 @property277 def optimizer(self):278 if self._optimizer is None:279 self._build_optimizer()280 return self._optimizer281 282 @property283 def lr_scheduler(self):284 if self._lr_scheduler is None:285 self._build_optimizer() # this will initialize self._lr_scheduler286 return self._lr_scheduler287 288 def _build_optimizer(self):289 params = list(290 filter(291 lambda p: p.requires_grad,292 chain(self.model.parameters(), self.criterion.parameters()),293 )294 )295 296 if self.is_fsdp and self.cfg.common.fp16:297 # FullyShardedDataParallel always uses MemoryEfficientFP16 wrapper,298 # mostly for the grad scaling. But if we don't have the299 # --memory-efficient-fp16 flag set, then we're effectively doing300 # regular --fp16 and can allow the use of optimizers that would301 # otherwise be unsupported by MemoryEfficientFP16Optimizer.302 allow_unsupported = not self.cfg.common.memory_efficient_fp16303 self._optimizer = optim.MemoryEfficientFP16Optimizer.build_optimizer(304 self.cfg, params, allow_unsupported=allow_unsupported305 )306 elif self.cfg.common.fp16 or self.cfg.common.bf16 or self.cfg.common.amp:307 if self.cuda and torch.cuda.get_device_capability(0)[0] < 7:308 logger.info(309 "NOTE: your device does NOT support faster training with --fp16 or --amp, "310 "please switch to FP32 which is likely to be faster"311 )312 if (313 self.cfg.common.memory_efficient_fp16314 or self.cfg.common.memory_efficient_bf16315 ):316 self._optimizer = optim.MemoryEfficientFP16Optimizer.build_optimizer(317 self.cfg, params318 )319 elif self.cfg.common.amp:320 self._optimizer = optim.AMPOptimizer.build_optimizer(self.cfg, params)321 else:322 self._optimizer = optim.FP16Optimizer.build_optimizer(self.cfg, params)323 else:324 if self.cuda and torch.cuda.get_device_capability(0)[0] >= 7:325 logger.info("NOTE: your device may support faster training with --fp16 or --amp")326 self._optimizer = optim.build_optimizer(self.cfg.optimizer, params)327 328 if self.is_fsdp:329 assert (330 not self.cfg.optimization.use_bmuf331 ), "--ddp-backend=fully_sharded is not compatible with BMUF"332 assert self._optimizer.supports_flat_params, (333 "--ddp-backend=fully_sharded is only compatible with pointwise "334 "optimizers (e.g., Adam, AdamW, Adadelta, Adamax, SGD, etc.). "335 "However, the sharding will result in slightly different results when "336 "using non-pointwise optimizers (e.g., Adagrad, Adafactor, LAMB)"337 )338 339 if self.cfg.optimization.use_bmuf:340 self._optimizer = optim.FairseqBMUF(341 self.cfg.bmuf,342 self._optimizer,343 )344 345 if self.cfg.distributed_training.zero_sharding == "os":346 if (347 self.cfg.common.fp16348 and not self.cfg.common.memory_efficient_fp16349 and not self.cfg.common.memory_efficient_bf16350 ) and not self.cfg.common.fp16_no_flatten_grads:351 raise ValueError(352 "ZeRO is incomptabile with fp16 and flattened grads. "353 "Please use --fp16-no-flatten-grads"354 )355 else:356 optim.shard_(self._optimizer, self.data_parallel_process_group)357 358 # We should initialize the learning rate scheduler immediately after359 # building the optimizer, so that the initial learning rate is set.360 self._lr_scheduler = lr_scheduler.build_lr_scheduler(361 self.cfg.lr_scheduler,362 self.optimizer,363 )364 self._lr_scheduler.step_update(0)365 366 @property367 def is_fsdp(self):368 return self.cfg.distributed_training.ddp_backend == "fully_sharded"369 370 def consolidate_optimizer(self):371 """For OSS, we need to consolidate the state dict."""372 if self.cfg.checkpoint.no_save_optimizer_state:373 return374 self._gathered_optim_state = None375 if hasattr(self.optimizer.optimizer, "consolidate_state_dict"):376 self.optimizer.optimizer.consolidate_state_dict()377 elif self.is_fsdp and not self.model.use_sharded_state:378 st = self.model.gather_full_optim_state_dict(379 self.optimizer380 ) # only returns on rank 0381 self._gathered_optim_state = st382 383 def state_dict(self):384 state_dict = {385 "args": None, # legacy386 "cfg": (387 OmegaConf.to_container(self.cfg, resolve=True, enum_to_str=True)388 if OmegaConf.is_config(self.cfg)389 else self.cfg390 ),391 "model": self.model.state_dict(),392 "criterion": (393 self.criterion.state_dict()394 if utils.has_parameters(self.criterion)395 else None396 ),397 "optimizer_history": (self._optim_history or [])398 + [399 {400 "criterion_name": self.get_criterion().__class__.__name__,401 "optimizer_name": self.optimizer.__class__.__name__,402 "lr_scheduler_state": self.lr_scheduler.state_dict(),403 "num_updates": self.get_num_updates(),404 }405 ],406 "task_state": self.task.state_dict() if self.task is not None else {},407 "extra_state": {408 "metrics": metrics.state_dict(),409 "previous_training_time": self.cumulative_training_time(),410 },411 }412 if self.cfg.ema.store_ema:413 # Save EMA model state as extra state414 state_dict["extra_state"]["ema"] = self.ema.get_model().state_dict()415 if self.cfg.ema.ema_fp32:416 # Save EMA params in fp32417 state_dict["extra_state"]["ema_fp32_params"] = self.ema.fp32_params418 if not self.cfg.checkpoint.no_save_optimizer_state:419 if self._gathered_optim_state is not None:420 state_dict["last_optimizer_state"] = self._gathered_optim_state421 self._gathered_optim_state = None422 else:423 state_dict["last_optimizer_state"] = self.optimizer.state_dict()424 if self.is_fsdp:425 # save meta data for recombining checkpoint upon loading426 state_dict["fsdp_metadata"] = self.model.local_metadata_dict()427 return state_dict428 429 def save_checkpoint(self, filename, extra_state):430 """Save all training state in a checkpoint file."""431 logger.info(f"Saving checkpoint to {filename}")432 # call state_dict on all ranks in case it needs internal communication433 state_dict = utils.move_to_cpu(self.state_dict())434 state_dict["extra_state"].update(extra_state)435 if self.should_save_checkpoint_on_current_rank:436 checkpoint_utils.torch_persistent_save(437 state_dict,438 filename,439 async_write=self.cfg.checkpoint.write_checkpoints_asynchronously,440 )441 logger.info(f"Finished saving checkpoint to {filename}")442 443 def load_checkpoint(444 self,445 filename,446 reset_optimizer=False,447 reset_lr_scheduler=False,448 optimizer_overrides=None,449 reset_meters=False,450 ):451 """452 Load all training state from a checkpoint file.453 rank = 0 will load the checkpoint, and then broadcast it to all454 other ranks.455 """456 extra_state, self._optim_history, last_optim_state = None, [], None457 458 logger.info(f"Preparing to load checkpoint {filename}")459 is_distributed = self.data_parallel_world_size > 1460 bexists = PathManager.isfile(filename)461 if bexists:462 load_on_all_ranks = (463 self.cfg.checkpoint.load_checkpoint_on_all_dp_ranks464 # TPUs don't support broadcast yet, so load checkpoints465 # on every worker for now466 or self.tpu467 # FSDP requires loading checkpoint shards on all ranks468 or (self.is_fsdp and self.cfg.distributed_training.use_sharded_state)469 or getattr(self.cfg.model, "base_layers", 0) > 0470 )471 472 if load_on_all_ranks or self.data_parallel_rank == 0:473 state = checkpoint_utils.load_checkpoint_to_cpu(474 filename, load_on_all_ranks=load_on_all_ranks475 )476 last_optim_state = state.get("last_optimizer_state", None)477 478 # If doing zero_sharding, do not broadcast global optimizer479 # state. Later we will broadcast sharded states to each rank480 # to avoid memory from exploding.481 if (482 not load_on_all_ranks483 and self.cfg.distributed_training.zero_sharding == "os"484 and "last_optimizer_state" in state485 and is_distributed486 ):487 state["last_optimizer_state"] = "SHARDED"488 else:489 last_optim_state = None490 state = None491 492 if is_distributed and not load_on_all_ranks:493 state = distributed_utils.broadcast_object(494 state,495 src_rank=0,496 group=self.data_parallel_process_group,497 dist_device=self.device,498 )499 if self.data_parallel_rank > 0:500 last_optim_state = state.get("last_optimizer_state", None)501 502 # load model parameters503 try:504 if self.cfg.checkpoint.use_ema_weights_to_init_param and "extra_state" in state and "ema" in state["extra_state"]:505 logger.info("use_ema_weights_to_init_param = True, will use EMA weights in the ckpt to init the model param...")506 ema_state_dict = state["extra_state"]["ema_fp32_params"] if "ema_fp32_params" in state["extra_state"] else state["extra_state"]["ema"]507 self.model.load_state_dict(508 ema_state_dict, strict=True, model_cfg=self.cfg.model509 )510 else:511 self.model.load_state_dict(512 state["model"], strict=True, model_cfg=self.cfg.model513 )514 # save memory for later steps515 if not (self.cfg.ema.store_ema and (self.cfg.checkpoint.use_latest_weights_to_init_ema or not ("extra_state" in state and "ema" in state["extra_state"]))):516 del state["model"]517 if utils.has_parameters(self.get_criterion()):518 self.get_criterion().load_state_dict(519 state["criterion"], strict=True520 )521 del state["criterion"]522 523 except Exception:524 raise Exception(525 "Cannot load model parameters from checkpoint {}; "526 "please ensure that the architectures match.".format(filename)527 )528 extra_state = state["extra_state"]529 self._optim_history = state["optimizer_history"]530 531 if last_optim_state is not None and not reset_optimizer:532 # rebuild optimizer after loading model, since params may have changed533 self._build_optimizer()534 535 # only reload optimizer and lr_scheduler if they match536 last_optim = self._optim_history[-1]537 assert (538 last_optim["criterion_name"] == self.get_criterion().__class__.__name__539 ), f"Criterion does not match; please reset the optimizer (--reset-optimizer). {last_optim['criterion_name']} vs {self.get_criterion().__class__.__name__}"540 assert (541 last_optim["optimizer_name"] == self.optimizer.__class__.__name__542 ), f"Optimizer does not match; please reset the optimizer (--reset-optimizer). {last_optim['optimizer_name']} vs {self.optimizer.__class__.__name__}"543 544 if not reset_lr_scheduler:545 self.lr_scheduler.load_state_dict(last_optim["lr_scheduler_state"])546 547 if self.is_fsdp and not self.model.use_sharded_state:548 # if use_sharded_state, the last_optim_state is already sharded, skip this549 last_optim_state = self.model.get_shard_from_optim_state_dict(550 last_optim_state551 )552 elif not load_on_all_ranks and is_distributed:553 last_optim_state = self.optimizer.broadcast_global_state_dict(554 last_optim_state555 )556 557 self.optimizer.load_state_dict(last_optim_state, optimizer_overrides)558 559 self.set_num_updates(last_optim["num_updates"])560 561 if extra_state is not None:562 itr_state = extra_state["train_iterator"]563 epoch = itr_state["epoch"]564 565 if "previous_training_time" in extra_state:566 self._previous_training_time = extra_state["previous_training_time"]567 self._start_time = time.time()568 569 self.lr_step(epoch)570 571 if (572 itr_state.get("version", 1) >= 2573 and itr_state["iterations_in_epoch"] == 0574 ):575 # reset meters at start of epoch576 reset_meters = True577 578 if "metrics" in extra_state and not reset_meters:579 metrics.load_state_dict(extra_state["metrics"])580 581 # reset TimeMeters, since their start times don't make sense anymore582 for meter in metrics.get_meters("default"):583 if isinstance(meter, meters.TimeMeter):584 meter.reset()585 586 if self.cfg.ema.store_ema:587 if self.cfg.checkpoint.use_latest_weights_to_init_ema or "ema" not in extra_state:588 if "ema" not in extra_state:589 logger.warn(590 "EMA not found in checkpoint. But store_ema is True. "591 "EMA is re-initialized from checkpoint."592 )593 elif self.cfg.checkpoint.use_latest_weights_to_init_ema:594 logger.info(595 "use_latest_weights_to_init_ema = True. EMA is re-initialized from checkpoint."596 )597 self.ema.restore(state["model"], build_fp32_params=self.cfg.ema.ema_fp32)598 del state["model"]599 else:600 logger.info(601 "Loading EMA from checkpoint"602 )603 self.ema.restore(extra_state["ema"], build_fp32_params=False)604 605 if self.cfg.ema.ema_fp32:606 if "ema_fp32_params" in extra_state:607 logger.info(608 "Loading EMA fp32 params from checkpoint"609 )610 self.ema.build_fp32_params(extra_state["ema_fp32_params"])611 else:612 logger.info(613 "Building EMA fp32 params from EMA model in checkpoint"614 )615 self.ema.build_fp32_params()616 617 logger.info(618 "Loaded checkpoint {} (epoch {} @ {} updates)".format(619 filename, epoch, self.get_num_updates()620 )621 )622 623 else:624 logger.info("No existing checkpoint found {}".format(filename))625 626 return extra_state627 628 def get_train_iterator(629 self,630 epoch,631 combine=True,632 load_dataset=True,633 data_selector=None,634 shard_batch_itr=True,635 disable_iterator_cache=False,636 ):637 """Return an EpochBatchIterator over the training set for a given epoch."""638 if load_dataset:639 logger.info("loading train data for epoch {}".format(epoch))640 self.task.load_dataset(641 self.cfg.dataset.train_subset,642 epoch=epoch,643 combine=combine,644 data_selector=data_selector,645 tpu=self.tpu,646 )647 batch_iterator = self.task.get_batch_iterator(648 dataset=self.task.dataset(self.cfg.dataset.train_subset),649 max_tokens=self.cfg.dataset.max_tokens,650 max_sentences=self.cfg.dataset.batch_size,651 max_positions=utils.resolve_max_positions(652 self.task.max_positions(),653 self.model.max_positions(),654 self.cfg.dataset.max_tokens,655 ),656 ignore_invalid_inputs=True,657 required_batch_size_multiple=self.cfg.dataset.required_batch_size_multiple,658 seed=self.cfg.common.seed,659 num_shards=self.data_parallel_world_size if shard_batch_itr else 1,660 shard_id=self.data_parallel_rank if shard_batch_itr else 0,661 num_workers=self.cfg.dataset.num_workers,662 epoch=epoch,663 data_buffer_size=self.cfg.dataset.data_buffer_size,664 disable_iterator_cache=disable_iterator_cache,665 )666 self.reset_dummy_batch(batch_iterator.first_batch)667 batch_iterator.dataset.dataset._seek()668 return batch_iterator669 670 def get_valid_iterator(671 self,672 subset,673 disable_iterator_cache=False,674 ):675 """Return an EpochBatchIterator over given validation subset for a given epoch."""676 self.task.dataset(subset).dataset._seek()677 batch_iterator = self.task.get_batch_iterator(678 dataset=self.task.dataset(subset),679 max_tokens=self.cfg.dataset.max_tokens_valid,680 max_sentences=self.cfg.dataset.batch_size_valid,681 max_positions=utils.resolve_max_positions(682 self.task.max_positions(),683 self.model.max_positions(),684 ),685 ignore_invalid_inputs=self.cfg.dataset.skip_invalid_size_inputs_valid_test,686 required_batch_size_multiple=self.cfg.dataset.required_batch_size_multiple,687 seed=self.cfg.common.seed,688 num_shards=self.data_parallel_world_size,689 shard_id=self.data_parallel_rank,690 num_workers=self.cfg.dataset.num_workers,691 # always pass a fixed "epoch" to keep validation data consistent692 # across training epochs693 epoch=1,694 data_buffer_size=self.cfg.dataset.data_buffer_size,695 disable_iterator_cache=disable_iterator_cache,696 )697 self.reset_dummy_batch(batch_iterator.first_batch)698 batch_iterator.dataset.dataset._seek()699 return batch_iterator700 701 def begin_epoch(self, epoch):702 """Called at the beginning of each epoch."""703 logger.info("begin training epoch {}".format(epoch))704 705 self.lr_step_begin_epoch(epoch)706 707 if self.quantizer is not None:708 self.quantizer.begin_epoch(epoch)709 710 # task specific setup per epoch711 self.task.begin_epoch(epoch, self.get_model())712 713 if self.tpu:714 import torch_xla.core.xla_model as xm715 716 xm.rendezvous("begin_epoch") # wait for all workers717 xm.mark_step()718 719 def begin_valid_epoch(self, epoch):720 """Called at the beginning of each validation epoch."""721 722 # task specific setup per validation epoch723 self.task.begin_valid_epoch(epoch, self.get_model())724 725 def reset_dummy_batch(self, batch):726 self._dummy_batch = batch727 728 @metrics.aggregate("train")729 def train_step(self, samples, raise_oom=False):730 """Do forward, backward and parameter update."""731 self._set_seed()732 self.model.train()733 self.criterion.train()734 self.zero_grad()735 736 metrics.log_start_time("train_wall", priority=800, round=0)737 738 # If EMA is enabled through store_ema=True739 # and task.uses_ema is True, pass the EMA model as a keyword740 # argument to the task.741 extra_kwargs = {}742 if self.cfg.ema.store_ema and getattr(self.task, "uses_ema", False):743 extra_kwargs["ema_model"] = self.ema.get_model()744 745 # forward and backward pass746 logging_outputs, sample_size, ooms = [], 0, 0747 for i, sample in enumerate(samples): # delayed update loop748 sample, is_dummy_batch = self._prepare_sample(sample)749 750 def maybe_no_sync():751 """752 Whenever *samples* contains more than one mini-batch, we753 want to accumulate gradients locally and only call754 all-reduce in the last backwards pass.755 """756 if (757 self.data_parallel_world_size > 1758 and hasattr(self.model, "no_sync")759 and i < len(samples) - 1760 # The no_sync context manager results in increased memory761 # usage with FSDP, since full-size gradients will be762 # accumulated on each GPU. It's typically a better tradeoff763 # to do the extra communication with FSDP.764 and not self.is_fsdp765 ):766 return self.model.no_sync()767 else:768 return contextlib.ExitStack() # dummy contextmanager769 770 try:771 with maybe_no_sync():772 # forward and backward773 loss, sample_size_i, logging_output = self.task.train_step(774 sample=sample,775 model=self.model,776 criterion=self.criterion,777 optimizer=self.optimizer,778 update_num=self.get_num_updates(),779 ignore_grad=is_dummy_batch,780 **extra_kwargs,781 )782 del loss783 784 logging_outputs.append(logging_output)785 sample_size += sample_size_i786 787 # emptying the CUDA cache after the first step can788 # reduce the chance of OOM789 if self.cuda and self.get_num_updates() == 0:790 torch.cuda.empty_cache()791 except RuntimeError as e:792 if "out of memory" in str(e):793 self._log_oom(e)794 if raise_oom:795 raise e796 logger.warning(797 "attempting to recover from OOM in forward/backward pass"798 )799 ooms += 1800 self.zero_grad()801 if self.cuda:802 torch.cuda.empty_cache()803 if self.cfg.distributed_training.distributed_world_size == 1:804 return None805 else:806 raise e807 808 if self.tpu and i < len(samples) - 1:809 # tpu-comment: every XLA operation before marking step is810 # appended to the IR graph, and processing too many batches811 # before marking step can lead to OOM errors.812 # To handle gradient accumulation use case, we explicitly813 # mark step here for every forward pass without a backward pass814 self._xla_markstep_and_send_to_cpu()815 816 if is_dummy_batch:817 if torch.is_tensor(sample_size):818 sample_size.zero_()819 else:820 sample_size *= 0.0821 822 if torch.is_tensor(sample_size):823 sample_size = sample_size.float()824 else:825 sample_size = float(sample_size)826 827 # gather logging outputs from all replicas828 if self._sync_stats():829 train_time = self._local_cumulative_training_time()830 logging_outputs, (831 sample_size,832 ooms,833 total_train_time,834 ) = self._aggregate_logging_outputs(835 logging_outputs, sample_size, ooms, train_time, ignore=is_dummy_batch836 )837 self._cumulative_training_time = (838 total_train_time / self.data_parallel_world_size839 )840 841 overflow = False842 try:843 with torch.autograd.profiler.record_function("reduce-grads"):844 # reduce gradients across workers845 self.optimizer.all_reduce_grads(self.model)846 if utils.has_parameters(self.criterion):847 self.optimizer.all_reduce_grads(self.criterion)848 849 with torch.autograd.profiler.record_function("multiply-grads"):850 # multiply gradients by (data_parallel_size / sample_size) since851 # DDP normalizes by the number of data parallel workers for852 # improved fp16 precision.853 # Thus we get (sum_of_gradients / sample_size) at the end.854 # In case of fp16, this step also undoes loss scaling.855 # (Debugging note: Some optimizers perform this scaling on the856 # fly, so inspecting model.parameters() or optimizer.params may857 # still show the original, unscaled gradients.)858 numer = (859 self.data_parallel_world_size860 if not self.cfg.optimization.use_bmuf or self._sync_stats()861 else 1862 )863 self.optimizer.multiply_grads(numer / (sample_size or 1.0))864 # Note: (sample_size or 1.0) handles the case of a zero gradient, in a865 # way that avoids CPU/device transfers in case sample_size is a GPU or866 # TPU object. The assumption is that the gradient itself is also 0.867 868 with torch.autograd.profiler.record_function("clip-grads"):869 # clip grads870 grad_norm = self.clip_grad_norm(self.cfg.optimization.clip_norm)871 872 # check that grad norms are consistent across workers873 # on tpu check tensor is slow874 if not self.tpu:875 if (876 not self.cfg.optimization.use_bmuf877 and self.cfg.distributed_training.ddp_backend != "slow_mo"878 ):879 self._check_grad_norms(grad_norm)880 if not torch.isfinite(grad_norm).all():881 # in case of AMP, if gradients are Nan/Inf then882 # optimizer step is still required883 if self.cfg.common.amp:884 overflow = True885 else:886 # check local gradnorm single GPU case, trigger NanDetector887 raise FloatingPointError("gradients are Nan/Inf")888 889 with torch.autograd.profiler.record_function("optimizer"):890 # take an optimization step891 self.task.optimizer_step(892 self.optimizer, model=self.model, update_num=self.get_num_updates()893 )894 if self.cfg.common.amp and overflow:895 if self._amp_retries == self.cfg.common.amp_batch_retries:896 logger.info("AMP: skipping this batch.")897 self._amp_retries = 0898 else:899 self._amp_retries += 1900 return self.train_step(samples, raise_oom) # recursion to feed in same batch901 902 except FloatingPointError:903 # re-run the forward and backward pass with hooks attached to print904 # out where it fails905 self.zero_grad()906 with NanDetector(self.get_model()):907 for _, sample in enumerate(samples):908 sample, _ = self._prepare_sample(sample)909 self.task.train_step(910 sample,911 self.model,912 self.criterion,913 self.optimizer,914 self.get_num_updates(),915 ignore_grad=False,916 **extra_kwargs,917 )918 raise919 except OverflowError as e:920 overflow = True921 logger.info(922 f"NOTE: gradient overflow detected, ignoring gradient, {str(e)}"923 )924 grad_norm = torch.tensor(0.0).cuda()925 self.zero_grad()926 except RuntimeError as e:927 if "out of memory" in str(e):928 self._log_oom(e)929 logger.error("OOM during optimization, irrecoverable")930 raise e931 932 # Some distributed wrappers (e.g., SlowMo) need access to the optimizer933 # after the step934 if hasattr(self.model, "perform_additional_optimizer_actions"):935 if hasattr(self.optimizer, "fp32_params"):936 self.model.perform_additional_optimizer_actions(937 self.optimizer.optimizer, self.optimizer.fp32_params938 )939 else:940 self.model.perform_additional_optimizer_actions(941 self.optimizer.optimizer942 )943 944 logging_output = None945 if not overflow or self.cfg.distributed_training.ddp_backend == "slow_mo":946 self.set_num_updates(self.get_num_updates() + 1)947 948 if self.cfg.ema.store_ema:949 # Step EMA forward with new model.950 self.ema.step(951 self.get_model(),952 self.get_num_updates(),953 )954 metrics.log_scalar(955 "ema_decay",956 self.ema.get_decay(),957 priority=10000,958 round=5,959 weight=0,960 )961 962 if self.tpu:963 import torch_xla.core.xla_model as xm964 965 # mark step on TPUs966 self._xla_markstep_and_send_to_cpu()967 968 # only log stats every log_interval steps969 # this causes wps to be misreported when log_interval > 1970 logging_output = {}971 if self.get_num_updates() % self.cfg.common.log_interval == 0:972 # log memory usage973 mem_info = xm.get_memory_info(self.device)974 gb_free = mem_info["kb_free"] / 1024 / 1024975 gb_total = mem_info["kb_total"] / 1024 / 1024976 metrics.log_scalar(977 "gb_free", gb_free, priority=1500, round=1, weight=0978 )979 metrics.log_scalar(980 "gb_total", gb_total, priority=1600, round=1, weight=0981 )982 logging_outputs = self._xla_markstep_and_send_to_cpu(983 logging_outputs984 )985 logging_output = self._reduce_and_log_stats(986 logging_outputs, sample_size, grad_norm987 )988 989 # log whenever there's an XLA compilation, since these990 # slow down training and may indicate opportunities for991 # optimization992 self._check_xla_compilation()993 else:994 if self.cuda and self.cuda_env is not None:995 # log minimum free memory over the iteration996 gb_used = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024997 torch.cuda.reset_peak_memory_stats()998 gb_free = self.cuda_env.total_memory_in_GB - gb_used999 metrics.log_scalar(1000 "gb_free", gb_free, priority=1500, round=1, weight=01001 )1002 1003 # log stats1004 logging_output = self._reduce_and_log_stats(1005 logging_outputs, sample_size, grad_norm1006 )1007 1008 # clear CUDA cache to reduce memory fragmentation1009 if (1010 self.cuda1011 and self.cfg.common.empty_cache_freq > 01012 and (1013 (self.get_num_updates() + self.cfg.common.empty_cache_freq - 1)1014 % self.cfg.common.empty_cache_freq1015 )1016 == 01017 ):1018 torch.cuda.empty_cache()1019 1020 if self.cfg.common.fp16 or self.cfg.common.amp:1021 metrics.log_scalar(1022 "loss_scale",1023 (1024 self.optimizer.scaler.loss_scale1025 if self.cfg.common.fp161026 else self.optimizer.scaler.get_scale()1027 ),1028 priority=700,1029 round=4,1030 weight=0,1031 )1032 1033 metrics.log_stop_time("train_wall")1034 return logging_output1035 1036 @metrics.aggregate("valid")1037 def valid_step(self, sample, raise_oom=False):1038 """Do forward pass in evaluation mode."""1039 if self.tpu:1040 import torch_xla.core.xla_model as xm1041 1042 xm.rendezvous("valid_step") # wait for all workers1043 1044 # If EMA is enabled through store_ema=True1045 # and task.uses_ema is True, pass the EMA model as a keyword1046 # argument to the task.1047 extra_kwargs = {}1048 if self.cfg.ema.store_ema and getattr(self.task, "uses_ema", False):1049 extra_kwargs["ema_model"] = self.ema.get_model()1050 1051 with torch.no_grad():1052 self.model.eval()1053 self.criterion.eval()1054 1055 sample, is_dummy_batch = self._prepare_sample(sample)1056 1057 try:1058 _loss, sample_size, logging_output = self.task.valid_step(1059 sample, self.model, self.criterion, **extra_kwargs1060 )1061 except RuntimeError as e:1062 if "out of memory" in str(e):1063 self._log_oom(e)1064 if not raise_oom:1065 logger.warning(1066 "ran out of memory in validation step, retrying batch"1067 )1068 for p in self.model.parameters():1069 if p.grad is not None:1070 p.grad = None # free some memory1071 if self.cuda:1072 torch.cuda.empty_cache()1073 return self.valid_step(sample, raise_oom=True)1074 raise e1075 1076 logging_outputs = [logging_output]1077 if is_dummy_batch:1078 if torch.is_tensor(sample_size):1079 sample_size.zero_()1080 else:1081 sample_size *= 0.01082 1083 # gather logging outputs from all replicas1084 if self.data_parallel_world_size > 1:1085 logging_outputs, (sample_size,) = self._aggregate_logging_outputs(1086 logging_outputs,1087 sample_size,1088 ignore=is_dummy_batch,1089 )1090 1091 # log validation stats1092 if self.tpu:1093 logging_outputs = self._xla_markstep_and_send_to_cpu(logging_outputs)1094 logging_output = self._reduce_and_log_stats(logging_outputs, sample_size)1095 1096 return logging_output1097 1098 def zero_grad(self):1099 self.optimizer.zero_grad()1100 1101 def lr_step_begin_epoch(self, epoch):1102 """Adjust the learning rate at the beginning of the epoch."""1103 self.lr_scheduler.step_begin_epoch(epoch)1104 # prefer updating the LR based on the number of steps1105 return self.lr_step_update()1106 1107 def lr_reinit(self, total_updates, num_updates):1108 self.lr_scheduler.reinit(total_updates, num_updates)1109 1110 def lr_step(self, epoch, val_loss=None):1111 """Adjust the learning rate at the end of the epoch."""1112 self.lr_scheduler.step(epoch, val_loss)1113 # prefer updating the LR based on the number of steps1114 return self.lr_step_update()1115 1116 def lr_step_update(self):1117 """Update the learning rate after each update."""1118 new_lr = self.lr_scheduler.step_update(self.get_num_updates())1119 if isinstance(new_lr, dict):1120 for k, v in new_lr.items():1121 metrics.log_scalar(f"lr_{k}", v, weight=0, priority=300)1122 new_lr = new_lr.get("default", next(iter(new_lr.values())))1123 else:1124 metrics.log_scalar("lr", new_lr, weight=0, priority=300)1125 return new_lr1126 1127 def get_lr(self):1128 """Get the current learning rate."""1129 return self.optimizer.get_lr()1130 1131 def get_model(self):1132 """Get the (non-wrapped) model instance."""1133 return self._model1134 1135 def get_criterion(self):1136 """Get the (non-wrapped) criterion instance."""1137 return self._criterion1138 1139 def get_meter(self, name):1140 """[deprecated] Get a specific meter by name."""1141 from fairseq import meters1142 1143 if "get_meter" not in self._warn_once:1144 self._warn_once.add("get_meter")1145 utils.deprecation_warning(1146 "Trainer.get_meter is deprecated. Please use fairseq.metrics instead."1147 )1148 1149 train_meters = metrics.get_meters("train")1150 if train_meters is None:1151 train_meters = {}1152 1153 if name == "train_loss" and "loss" in train_meters:1154 return train_meters["loss"]1155 elif name == "train_nll_loss":1156 # support for legacy train.py, which assumed this meter is1157 # always initialized1158 m = train_meters.get("nll_loss", None)1159 return m or meters.AverageMeter()1160 elif name == "wall":1161 # support for legacy train.py, which assumed this meter is1162 # always initialized1163 m = metrics.get_meter("default", "wall")1164 return m or meters.TimeMeter()1165 elif name == "wps":1166 m = metrics.get_meter("train", "wps")1167 return m or meters.TimeMeter()1168 elif name in {"valid_loss", "valid_nll_loss"}:1169 # support for legacy train.py, which assumed these meters1170 # are always initialized1171 k = name[len("valid_") :]1172 m = metrics.get_meter("valid", k)1173 return m or meters.AverageMeter()1174 elif name == "oom":1175 return meters.AverageMeter()1176 elif name in train_meters:1177 return train_meters[name]1178 return None1179 1180 def get_num_updates(self):1181 """Get the number of parameters updates."""1182 return self._num_updates1183 1184 def set_num_updates(self, num_updates):1185 """Set the number of parameters updates."""1186 self._num_updates = num_updates1187 self.lr_step_update()1188 if self.quantizer:1189 self.quantizer.step_update(self._num_updates)1190 metrics.log_scalar("num_updates", self._num_updates, weight=0, priority=200)1191 1192 def clip_grad_norm(self, clip_norm):1193 def agg_norm_fn(total_norm):1194 total_norm = total_norm.cuda().float() ** 21195 total_norm = distributed_utils.all_reduce(1196 total_norm, group=self.data_parallel_process_group1197 )1198 return total_norm ** 0.51199 1200 should_agg_norm = (