Method update_step

funasr/train_utils/trainer_ds.py:713–752

Update step. Args: model: Model instance or model name. optim: TODO. scheduler: TODO. scaler: TODO. loss_dict: TODO.

(self, model, optim, scheduler, scaler, loss_dict=None)

                loss.backward()

    def update_step(self, model, optim, scheduler, scaler, loss_dict=None):
        """Update step.
        
            Args:
                model: Model instance or model name.
                optim: TODO.
                scheduler: TODO.
                scaler: TODO.
                loss_dict: TODO.
            """
        batch_idx = loss_dict["batch_idx"]
        if self.use_deepspeed:
            model.step()
        else:
            if (batch_idx + 1) % self.accum_grad == 0:
                # Perform gradient clipping if it is set
                if self.grad_clip > 0:
                    grad_norm = torch.nn.utils.clip_grad_norm_(
                        model.parameters(),
                        max_norm=self.grad_clip,
                        norm_type=self.grad_clip_type,
                    )
                    if not torch.isfinite(grad_norm):
                        logging.warning(
                            f"The grad norm is {grad_norm}. Skipping updating the model."
                        )
                        optim.zero_grad()  # Reset gradients
                        return

                # Execute an optimization step (update model parameters)
                if self.use_ddp or self.use_fsdp:
                    dist.barrier()
                if scaler:
                    scaler.step(optim)
                    scaler.update()
                else:
                    optim.step()
                scheduler.step()
                # Clear gradients for the next accumulation stage
                optim.zero_grad(set_to_none=True)

    def validate_epoch(
        self,