def train_step(self, samples, raise_oom=False):
        """Do forward, backward and parameter update."""
        self._set_seed()
        self.model.train()
        self.criterion.train()
        self.zero_grad()

        metrics.log_start_time("train_wall", priority=800, round=2)

        grad_norm = torch.tensor(0.0).cuda()

        # let fairseq handle grad accumlation scaling
        self.model.scale_wrt_gas = False

        # forward and backward pass
        logging_outputs, sample_size, ooms = [], 0, 0
        # print(f'** samples size: {len(samples)}')
        sample_count = len(samples)
        for i, sample in enumerate(samples):  # delayed update loop
            sample, is_dummy_batch = self._prepare_sample(sample)

            if self.cfg.common.fp16:
                self.model.optimizer.override_loss_scale(self.scaler.loss_scale)

            self.model.set_gradient_accumulation_boundary(is_boundary=False)
            try:
                # forward and backward
                # print(f'i={i}, rank={dist.get_rank()}, pre task.train_step')
                dist.barrier()
                loss, sample_size_i, logging_output = self.task.train_step(
                    sample=sample,
                    model=self.model,
                    criterion=self.criterion,
                    optimizer=self.optimizer,
                    update_num=self.get_num_updates(),
                    ignore_grad=is_dummy_batch,
                )
                self.train_step_count += 1
                
                # increment deepspeed micro step on non-final train step since optimizer.step will increment it for us
                if (i + 1) != sample_count:
                    self.model.micro_steps += 1

                self.model.set_gradient_accumulation_boundary(is_boundary=True)
                if self.zero_stage <= 2:
                    self.model.allreduce_gradients()

                # print(f'grads[0]={list([p.grad for p in self.model.optimizer.fp16_groups[0]])}')
                # print(f'train_step={self.train_step_count}, loss_scale={self.model.optimizer.cur_scale}')
                # print(f'i={i}, rank={dist.get_rank()}, loss={loss}')
                if torch.distributed.get_rank() == 0:
                    _loss_scale = self.model.optimizer.external_loss_scale if self.cfg.common.fp16 else 0
                    print(f"[{torch.distributed.get_rank()}], " \
                            f"micro_step={self.model.micro_steps}, " \
                            f"gas_boundary={self.model.is_gradient_accumulation_boundary()}, " \
                            f"train_step={self.train_step_count}, " \
                            f"lr={self.get_lr()}, " \
                            f"loss_scale={_loss_scale}, " \
                            f"loss={loss}")

                if self.cfg.common.exit_interval and self.train_step_count % self.cfg.common.exit_interval == 0:
                    if torch.distributed.get_rank() == 0:
                        print("exiting early...")
                    sys.exit()

                logging_outputs.append(logging_output)
                sample_size += sample_size_i

                # emptying the CUDA cache after the first step can
                # reduce the chance of OOM
                if self.cuda and self.get_num_updates() == 0:
                    torch.cuda.empty_cache()
            except RuntimeError as e:
                if "out of memory" in str(e):
                    self._log_oom(e)
                    if raise_oom:
                        raise e
                    logger.warning(
                        "attempting to recover from OOM in forward/backward pass"
                    )
                    ooms += 1
                    self.zero_grad()
                    if self.cuda:
                        torch.cuda.empty_cache()
                    if self.cfg.distributed_training.distributed_world_size == 1:
                        return None
                else:
                    raise e

        if is_dummy_batch:
            if torch.is_tensor(sample_size):
                sample_size.zero_()
            else:
                sample_size *= 0.0

        if torch.is_tensor(sample_size):
            sample_size = sample_size.float()
        else:
            sample_size = float(sample_size)

        # gather logging outputs from all replicas
        if self.data_parallel_world_size > 1:
            train_time = self._local_cumulative_training_time()
            logging_outputs, (
                sample_size,
                ooms,
                total_train_time,
            ) = self._aggregate_logging_outputs(
                logging_outputs, sample_size, ooms, train_time, ignore=is_dummy_batch
            )
            self._cumulative_training_time = (
                total_train_time / self.data_parallel_world_size
            )

        # grad_norms = []
        # for param in self.optimizer.fp16_params:
        #     if param.grad is not None:
        #         grad_norms.append(param.grad.norm())
        # print(grad_norms)

        # self.model.optimizer.dump_grad_norms()

        ######################
        ### multiply grads ###
        ######################
        # numer = (
        #     self.data_parallel_world_size
        #     if not self.cfg.optimization.use_bmuf or self._sync_stats()
        #     else 1
        # )
        # self.optimizer.multiply_grads(numer / (sample_size or 1.0))
        # grad_norm = self.optimizer.clip_grad_norm(self.cfg.optimization.clip_norm)

        overflow = False
        try:
            with torch.autograd.profiler.record_function("optimizer"):
                # take an optimization step
                self.task.optimizer_step(
                    self.optimizer, model=self.model, update_num=self.get_num_updates()
                )
                # pass overflow flag from ds to fairseq
                if self.cfg.common.fp16:
                    overflow = self.model.optimizer.overflow
                    self.scaler.check_overflow(overflow=overflow)
                    self.scaler.update()

                _grad_norm = self.model.get_global_grad_norm()
                if _grad_norm is not None:
                    grad_norm = torch.tensor(_grad_norm).to(self.device)
        except FloatingPointError:
            raise
        except OverflowError as e:
            logger.info(f"NOTE: gradient overflow detected, ignoring gradient, {str(e)}")
            overflow = True
        
        # import sys; sys.exit()
        # except RuntimeError as e:
        #     raise e

        # except FloatingPointError:
        #     # re-run the forward and backward pass with hooks attached to print
        #     # out where it fails
        #     self.zero_grad()
        #     with NanDetector(self.get_model()):
        #         for _, sample in enumerate(samples):
        #             sample, _ = self._prepare_sample(sample)
        #             self.task.train_step(
        #                 sample,
        #                 self.model,
        #                 self.criterion,
        #                 self.optimizer,
        #                 self.get_num_updates(),
        #                 ignore_grad=False,
        #             )
        #     raise
        # except OverflowError as e:
        #     overflow = True
        #     logger.info(f"NOTE: gradient overflow detected, ignoring gradient, {str(e)}")
        #     grad_norm = torch.tensor(0.0).cuda()
        #     self.zero_grad()
        # except RuntimeError as e:
        #     if "out of memory" in str(e):
        #         self._log_oom(e)
        #         logger.error("OOM during optimization, irrecoverable")
        #     raise e

        # Some distributed wrappers (e.g., SlowMo) need access to the optimizer
        # after the step
        # if hasattr(self.model, "perform_additional_optimizer_actions"):
        #     if hasattr(self.optimizer, "fp32_params"):
        #         self.model.perform_additional_optimizer_actions(
        #             self.optimizer.optimizer, self.optimizer.fp32_params
        #         )
        #     else:
        #         self.model.perform_additional_optimizer_actions(
        #             self.optimizer.optimizer
        #         )

        logging_output = None
        if not overflow or self.cfg.distributed_training.ddp_backend == "slow_mo":
            self.set_num_updates(self.get_num_updates() + 1)

            if self.cuda and self.cuda_env is not None:
                # log minimum free memory over the iteration
                gb_used = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024
                torch.cuda.reset_peak_memory_stats()
                gb_free = self.cuda_env.total_memory_in_GB - gb_used
                metrics.log_scalar(
                    "gb_free", gb_free, priority=1500, round=1, weight=0
                )

            # log stats
            logging_output = self._reduce_and_log_stats(
                logging_outputs, sample_size, grad_norm
            )

            # clear CUDA cache to reduce memory fragmentation
            # if (
            #     self.cuda
            #     and self.cfg.common.empty_cache_freq > 0
            #     and (
            #         (self.get_num_updates() + self.cfg.common.empty_cache_freq - 1)
            #         % self.cfg.common.empty_cache_freq
            #     )
            #     == 0
            # ):
            #     torch.cuda.empty_cache()

        if self.cfg.common.fp16:
            metrics.log_scalar(
                "loss_scale",
                # self.optimizer.loss_scaler.loss_scaler,
                self.optimizer.cur_scale,
                priority=700,
                round=4,
                weight=0,
            )

        metrics.log_stop_time("train_wall")
        return logging_output
Do forward, backward and parameter update.
kosmos-2/fairseq/fairseq/ds_trainer.py