class_docstr: Booster is a high-level API for training neural networks. It provides a unified interface for
  training with different precision, accelerator, and plugin.
  
  
  ```python
  # Following is pseudocode
  
  colossalai.launch(...)
  plugin = GeminiPlugin(...)
  booster = Booster(precision='fp16', plugin=plugin)
  
  model = GPT2()
  optimizer = HybridAdam(model.parameters())
  dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
  lr_scheduler = LinearWarmupScheduler()
  criterion = GPTLMLoss()
  
  model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model, optimizer, criterion, dataloader, lr_scheduler)
  
  for epoch in range(max_epochs):
      for input_ids, attention_mask in dataloader:
          outputs = model(input_ids.cuda(), attention_mask.cuda())
          loss = criterion(outputs.logits, input_ids)
          booster.backward(loss, optimizer)
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
  ```
  
  Args:
      device (str or torch.device): The device to run the training. Default: None.
                                    If plugin is not used or plugin doesn't control the device,
                                    this argument will be set as training device ('cuda' will be used if argument is None).
      mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.
                              If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.
                              'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.
      plugin (Plugin): The plugin to run the training. Default: None.
class_docstr_tok: ["Booster","is","a","high-level","API","for","training","neural","networks",".","It","provides","a","unified","interface","for","training","with","different","precision",",","accelerator",",","and","plugin",".","``","`","python","#","Following","is","pseudocode","colossalai.launch","(","...",")","plugin","=","GeminiPlugin","(","...",")","booster","=","Booster","(","precision='fp16","'",",","plugin=plugin",")","model","=","GPT2","(",")","optimizer","=","HybridAdam","(","model.parameters","(",")",")","dataloader","=","plugin.prepare_dataloader","(","train_dataset",",","batch_size=8",")","lr_scheduler","=","LinearWarmupScheduler","(",")","criterion","=","GPTLMLoss","(",")","model",",","optimizer",",","criterion",",","dataloader",",","lr_scheduler","=","booster.boost","(","model",",","optimizer",",","criterion",",","dataloader",",","lr_scheduler",")","for","epoch","in","range","(","max_epochs",")",":","for","input_ids",",","attention_mask","in","dataloader",":","outputs","=","model","(","input_ids.cuda","(",")",",","attention_mask.cuda","(",")",")","loss","=","criterion","(","outputs.logits",",","input_ids",")","booster.backward","(","loss",",","optimizer",")","optimizer.step","(",")","lr_scheduler.step","(",")","optimizer.zero_grad","(",")","``","`","Args",":","device","(","str","or","torch.device",")",":","The","device","to","run","the","training",".","Default",":","None",".","If","plugin","is","not","used","or","plugin","does","n't","control","the","device",",","this","argument","will","be","set","as","training","device","(","'cuda","'","will","be","used","if","argument","is","None",")",".","mixed_precision","(","str","or","MixedPrecision",")",":","The","mixed","precision","to","run","the","training",".","Default",":","None",".","If","the","argument","is","a","string",",","it","can","be","'fp16","'",",","'fp16_apex","'",",","'bf16","'",",","or","'fp8","'",".","'fp16","'","would","use","PyTorch","AMP","while","`","fp16_apex","`","would","use","Nvidia","Apex",".","plugin","(","Plugin",")",":","The","plugin","to","run","the","training",".","Default",":","None","."]
docstr: Wrap and inject features to the passed in model, optimizer, criterion, lr_scheduler, and dataloader.
  
  Args:
      model (nn.Module): Convert model into a wrapped model for distributive training.
                         The model might be decorated or partitioned by plugin's strategy after execution of this method.
      optimizer (Optimizer, optional): Convert optimizer into a wrapped optimizer for distributive training.
                                       The optimizer's param groups or states might be decorated or partitioned by plugin's strategy after execution of this method. Defaults to None.
      criterion (Callable, optional): The function that calculates loss. Defaults to None.
      dataloader (DataLoader, optional): The prepared dataloader for training. Defaults to None.
      lr_scheduler (LRScheduler, optional): The learning scheduler for training. Defaults to None.
  
  Returns:
      List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]: The list of boosted input arguments.
docstr_tok: ["Wrap","and","inject","features","to","the","passed","in","model",",","optimizer",",","criterion",",","lr_scheduler",",","and","dataloader",".","Args",":","model","(","nn.Module",")",":","Convert","model","into","a","wrapped","model","for","distributive","training",".","The","model","might","be","decorated","or","partitioned","by","plugin","'s","strategy","after","execution","of","this","method",".","optimizer","(","Optimizer",",","optional",")",":","Convert","optimizer","into","a","wrapped","optimizer","for","distributive","training",".","The","optimizer","'s","param","groups","or","states","might","be","decorated","or","partitioned","by","plugin","'s","strategy","after","execution","of","this","method",".","Defaults","to","None",".","criterion","(","Callable",",","optional",")",":","The","function","that","calculates","loss",".","Defaults","to","None",".","dataloader","(","DataLoader",",","optional",")",":","The","prepared","dataloader","for","training",".","Defaults","to","None",".","lr_scheduler","(","LRScheduler",",","optional",")",":","The","learning","scheduler","for","training",".","Defaults","to","None",".","Returns",":","List","[","Union","[","nn.Module",",","Optimizer",",","LRScheduler",",","DataLoader","]","]",":","The","list","of","boosted","input","arguments","."]
code: def boost(
      self,
      model: nn.Module,
      optimizer: Optional[Optimizer] = None,
      criterion: Optional[Callable] = None,
      dataloader: Optional[DataLoader] = None,
      lr_scheduler: Optional[LRScheduler] = None,
  ) -> List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]:
      """
      Wrap and inject features to the passed in model, optimizer, criterion, lr_scheduler, and dataloader.
  
      Args:
          model (nn.Module): Convert model into a wrapped model for distributive training.
                             The model might be decorated or partitioned by plugin's strategy after execution of this method.
          optimizer (Optimizer, optional): Convert optimizer into a wrapped optimizer for distributive training.
                                           The optimizer's param groups or states might be decorated or partitioned by plugin's strategy after execution of this method. Defaults to None.
          criterion (Callable, optional): The function that calculates loss. Defaults to None.
          dataloader (DataLoader, optional): The prepared dataloader for training. Defaults to None.
          lr_scheduler (LRScheduler, optional): The learning scheduler for training. Defaults to None.
  
      Returns:
          List[Union[nn.Module, Optimizer, LRScheduler, DataLoader]]: The list of boosted input arguments.
      """
      # TODO(FrankLeeeee): consider multi-model and multi-optimizer case
      # TODO(FrankLeeeee): consider multi-dataloader case
      pretrained_path = pretrained_utils.get_pretrained_path(model)
      # transform model for mixed precision
      if self.plugin:
          (
              model,
              optimizer,
              criterion,
              dataloader,
              lr_scheduler,
          ) = self.plugin.configure(
              model, optimizer, criterion, dataloader, lr_scheduler
          )
  
      if self.plugin and not self.plugin.control_device():
          # transform model for accelerator
          model = self.accelerator.configure_model(model)
  
      if self.mixed_precision and (
          self.plugin is None
          or self.plugin
          and not self.plugin.control_precision()
      ):
          # transform model for mixed precision
          # when mixed_precision is specified and the plugin is not given or does not control the precision
          model, optimizer, criterion = self.mixed_precision.configure(
              model, optimizer, criterion
          )
  
      if pretrained_path:
          self.load_model(model, pretrained_path)
          # clear pretrained path attr
          orig_model = (
              model.unwrap()
              if isinstance(model, ModelWrapper)
              else model
          )
          pretrained_utils.set_pretrained_path(orig_model, None)
  
      return model, optimizer, criterion, dataloader, lr_scheduler
  
imports: import warnings
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, Iterator, List, Optional, Union
  import torch
  import torch.nn
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  import colossalai.interface.pretrained
  from colossalai.checkpoint_io import GeneralCheckpointIO
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from .accelerator import Accelerator
  from .mixed_precision import MixedPrecision, mixed_precision_factory
  from .plugin import Plugin
  from .plugin.pp_plugin_base import PipelinePluginBase
label_desc: Use image node_id 2 for calling the Booster obj's underlying member method code with example usage: obj.boost(model, optimizer, criterion, dataloader, lr_scheduler) and returns: model,  optimizer,  criterion,  dataloader,  lr_scheduler