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 __init__(
      self,
      model,
      strategy: Strategy,
      optim: Optimizer,
      lr_scheduler: _LRScheduler,
      max_epochs: int = 2,
      accumulation_steps: int = 8,
  ) -> None:
      if accumulation_steps > 1:
          assert not isinstance(
              strategy, GeminiStrategy
          ), "Accumulation steps are not supported in stage 3 of ColossalAI"
  
      super().__init__(strategy, max_epochs, model, optim)
  
      self.accumulation_steps = accumulation_steps
      self.scheduler = lr_scheduler
  
      self.num_train_step = 0
      self.num_eval_step = 0
  
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 1 to create a new SFTTrainer object from inherited base classes: SLTrainer with example: obj = SFTTrainer(model, strategy, optim, lr_scheduler, max_epochs, accumulation_steps)