class_docstr:     Trainer to use while training reward model.
  
  Args:
      model (torch.nn.Module): the model to train
      strategy (Strategy): the strategy to use for training
      optim(Optimizer): the optimizer to use for training
      lr_scheduler(_LRScheduler): the lr scheduler to use for training
      max_epochs (int, defaults to 2): the number of epochs to train
      accumulation_steps (int, defaults to 8): the number of steps to accumulate gradients
class_docstr_tok: ["Trainer","to","use","while","training","reward","model",".","Args",":","model","(","torch.nn.Module",")",":","the","model","to","train","strategy","(","Strategy",")",":","the","strategy","to","use","for","training","optim","(","Optimizer",")",":","the","optimizer","to","use","for","training","lr_scheduler","(","_LRScheduler",")",":","the","lr","scheduler","to","use","for","training","max_epochs","(","int",",","defaults","to","2",")",":","the","number","of","epochs","to","train","accumulation_steps","(","int",",","defaults","to","8",")",":","the","number","of","steps","to","accumulate","gradients"]
code: def _train(self, epoch: int):
      self.model.train()
      step_bar = tqdm.trange(
          len(self.train_dataloader) // self.accumulation_steps,
          desc=f"Epoch {epoch + 1}/{self.max_epochs}",
          disable=not is_rank_0(),
      )
      for i, batch in enumerate(self.train_dataloader):
          batch = to_device(batch, torch.cuda.current_device())
          outputs = self.model(
              batch["input_ids"],
              attention_mask=batch["attention_mask"],
              labels=batch["labels"],
          )
          loss = outputs.loss / self.accumulation_steps
          self.total_loss += loss.item()
          self.strategy.backward(loss, self.model, self.optimizer)
          # gradient accumulation
          if (i + 1) % self.accumulation_steps == 0:
              self.strategy.optimizer_step(self.optimizer)
              self.optimizer.zero_grad()
              self.scheduler.step()
              if self.writer:
                  self.writer.add_scalar(
                      "train/loss", self.total_loss, self.num_train_step
                  )
                  self.writer.add_scalar(
                      "train/lr",
                      self.scheduler.get_last_lr()[0],
                      self.num_train_step,
                  )
                  self.num_train_step += 1
              self.total_loss = 0
              step_bar.update()
      step_bar.close()
  
file_path: ColossalAI/applications/Chat/coati/trainer/sft.py
imports: from typing import Optional
  import torch
  import torch.distributed
  import tqdm
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from colossalai.logging import DistributedLogger
  from .base import SLTrainer
  from .strategies import GeminiStrategy, Strategy
  from .utils import is_rank_0, to_device
label_desc: Use image node_id 2 for calling the SFTTrainer obj's underlying member method code with example usage: obj._train(epoch) without return types