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: Execute forward & backward when utilizing pipeline parallel.
  Return loss or Huggingface style model outputs if needed.
  
  Warning: This function is tailored for the scenario of pipeline parallel.
  As a result, please don't do the forward/backward pass in the conventional way (model(input)/loss.backward())
  when doing pipeline parallel training with booster, which will cause unexpected errors.
  
  Args:
      data_iter(Iterator): The iterator for getting the next batch of data. Usually there are two ways to obtain this argument:
                           1. wrap the dataloader to iterator through: iter(dataloader)
                           2. get the next batch from dataloader, and wrap this batch to iterator: iter([batch])
      model (nn.Module): The model to execute forward/backward, it should be a model wrapped by a plugin that supports pipeline.
      criterion: (Callable[[Any, Any], torch.Tensor]): Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor.
                                                       'lambda y, x: loss_fn(y)' can turn a normal loss function into a valid two-argument criterion here.
      optimizer (Optimizer, optional): The optimizer for execution of backward. Can be None when only doing forward (i.e. evaluation). Defaults to None.
      return_loss (bool, optional): Whether to return loss in the dict returned by this method. Defaults to True.
      return_output (bool, optional): Whether to return Huggingface style model outputs in the dict returned by this method. Defaults to False.
  
  Returns:
      Dict[str, Any]: Output dict in the form of {'loss': ..., 'outputs': ...}.
                      ret_dict['loss'] is the loss of forward if return_loss is set to True, else None.
                      ret_dict['outputs'] is the Huggingface style model outputs during forward if return_output is set to True, else None.
docstr_tok: ["Execute","forward","&","backward","when","utilizing","pipeline","parallel",".","Return","loss","or","Huggingface","style","model","outputs","if","needed",".","Warning",":","This","function","is","tailored","for","the","scenario","of","pipeline","parallel",".","As","a","result",",","please","do","n't","do","the","forward\/backward","pass","in","the","conventional","way","(","model","(","input",")","\/loss.backward","(",")",")","when","doing","pipeline","parallel","training","with","booster",",","which","will","cause","unexpected","errors",".","Args",":","data_iter","(","Iterator",")",":","The","iterator","for","getting","the","next","batch","of","data",".","Usually","there","are","two","ways","to","obtain","this","argument",":","1.","wrap","the","dataloader","to","iterator","through",":","iter","(","dataloader",")","2.","get","the","next","batch","from","dataloader",",","and","wrap","this","batch","to","iterator",":","iter","(","[","batch","]",")","model","(","nn.Module",")",":","The","model","to","execute","forward\/backward",",","it","should","be","a","model","wrapped","by","a","plugin","that","supports","pipeline",".","criterion",":","(","Callable","[","[","Any",",","Any","]",",","torch.Tensor","]",")",":","Criterion","to","be","used",".","It","should","take","two","arguments",":","model","outputs","and","inputs",",","and","returns","loss","tensor",".","'lambda","y",",","x",":","loss_fn","(","y",")","'","can","turn","a","normal","loss","function","into","a","valid","two-argument","criterion","here",".","optimizer","(","Optimizer",",","optional",")",":","The","optimizer","for","execution","of","backward",".","Can","be","None","when","only","doing","forward","(","i.e",".","evaluation",")",".","Defaults","to","None",".","return_loss","(","bool",",","optional",")",":","Whether","to","return","loss","in","the","dict","returned","by","this","method",".","Defaults","to","True",".","return_output","(","bool",",","optional",")",":","Whether","to","return","Huggingface","style","model","outputs","in","the","dict","returned","by","this","method",".","Defaults","to","False",".","Returns",":","Dict","[","str",",","Any","]",":","Output","dict","in","the","form","of","{","'loss","'",":","...",",","'outputs","'",":","...","}",".","ret_dict","[","'loss","'","]","is","the","loss","of","forward","if","return_loss","is","set","to","True",",","else","None",".","ret_dict","[","'outputs","'","]","is","the","Huggingface","style","model","outputs","during","forward","if","return_output","is","set","to","True",",","else","None","."]
code: def execute_pipeline(
      self,
      data_iter: Iterator,
      model: nn.Module,
      criterion: Callable[[Any, Any], torch.Tensor],
      optimizer: Optional[Optimizer] = None,
      return_loss: bool = True,
      return_outputs: bool = False,
  ) -> Dict[str, Any]:
      """
      Execute forward & backward when utilizing pipeline parallel.
      Return loss or Huggingface style model outputs if needed.
  
      Warning: This function is tailored for the scenario of pipeline parallel.
      As a result, please don't do the forward/backward pass in the conventional way (model(input)/loss.backward())
      when doing pipeline parallel training with booster, which will cause unexpected errors.
  
      Args:
          data_iter(Iterator): The iterator for getting the next batch of data. Usually there are two ways to obtain this argument:
                               1. wrap the dataloader to iterator through: iter(dataloader)
                               2. get the next batch from dataloader, and wrap this batch to iterator: iter([batch])
          model (nn.Module): The model to execute forward/backward, it should be a model wrapped by a plugin that supports pipeline.
          criterion: (Callable[[Any, Any], torch.Tensor]): Criterion to be used. It should take two arguments: model outputs and inputs, and returns loss tensor.
                                                           'lambda y, x: loss_fn(y)' can turn a normal loss function into a valid two-argument criterion here.
          optimizer (Optimizer, optional): The optimizer for execution of backward. Can be None when only doing forward (i.e. evaluation). Defaults to None.
          return_loss (bool, optional): Whether to return loss in the dict returned by this method. Defaults to True.
          return_output (bool, optional): Whether to return Huggingface style model outputs in the dict returned by this method. Defaults to False.
  
      Returns:
          Dict[str, Any]: Output dict in the form of {'loss': ..., 'outputs': ...}.
                          ret_dict['loss'] is the loss of forward if return_loss is set to True, else None.
                          ret_dict['outputs'] is the Huggingface style model outputs during forward if return_output is set to True, else None.
      """
      assert isinstance(
          self.plugin, PipelinePluginBase
      ), f"The plugin {self.plugin.__class__.__name__} does not support pipeline."
      return self.plugin.execute_pipeline(
          data_iter,
          model,
          criterion,
          optimizer,
          return_loss,
          return_outputs,
      )
  
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 4 for calling the Booster obj's underlying member method code with example usage: obj.execute_pipeline(data_iter, model, criterion, optimizer, return_loss, return_outputs) and returns: self