class_docstr: Wraps a :class:`PipelineModulesGraph` model to train on using synchronous pipeline
  parallelism. If the model requires lots of memory and doesn't fit on a single GPU,
  pipeline parallelism is a useful technique to employ for training.
  
  The implementation is based on the torchgpipe_ paper.
  
  .. _torchgpipe: https://arxiv.org/abs/2004.09910
  
  PipelineModulesGraph combines pipeline parallelism with checkpointing to reduce peak
  memory required to train while minimizing device under-utilization.
  
  You should place all the modules on the appropriate rpc workers and devices and wrap
  them into an :class:`PipelineModulesGraph` module defining the connection between them.
  
  Args:
      module (:class:`PipelineModulesGraph`):
      model to be parallelized using pipelining. Each module
          in the graph has to have all of its parameters on a single
          device.
      chunks (int):
          number of micro-batches (default: ``1``)
      checkpoint (str):
          when to enable checkpointing, one of ``'always'``,
          ``'except_last'``, or ``'never'`` (default: ``'except_last'``).
          ``'never'`` disables checkpointing completely, ``'except_last'``
          enables checkpointing for all micro-batches except the last one
          and ``'always'`` enables checkpointing for all micro-batches.
class_docstr_tok: ["Wraps","a",":","class",":","`","PipelineModulesGraph","`","model","to","train","on","using","synchronous","pipeline","parallelism",".","If","the","model","requires","lots","of","memory","and","does","n't","fit","on","a","single","GPU",",","pipeline","parallelism","is","a","useful","technique","to","employ","for","training",".","The","implementation","is","based","on","the","torchgpipe_","paper",".","..","_torchgpipe",":","https",":","\/\/arxiv.org\/abs\/2004.09910","PipelineModulesGraph","combines","pipeline","parallelism","with","checkpointing","to","reduce","peak","memory","required","to","train","while","minimizing","device","under-utilization",".","You","should","place","all","the","modules","on","the","appropriate","rpc","workers","and","devices","and","wrap","them","into","an",":","class",":","`","PipelineModulesGraph","`","module","defining","the","connection","between","them",".","Args",":","module","(",":","class",":","`","PipelineModulesGraph","`",")",":","model","to","be","parallelized","using","pipelining",".","Each","module","in","the","graph","has","to","have","all","of","its","parameters","on","a","single","device",".","chunks","(","int",")",":","number","of","micro-batches","(","default",":","``","1","``",")","checkpoint","(","str",")",":","when","to","enable","checkpointing",",","one","of","``","'always","'","``",",","``","'except_last","'","``",",","or","``","'never","'","``","(","default",":","``","'except_last","'","``",")",".","``","'never","'","``","disables","checkpointing","completely",",","``","'except_last","'","``","enables","checkpointing","for","all","micro-batches","except","the","last","one","and","``","'always","'","``","enables","checkpointing","for","all","micro-batches","."]
code: def __init__(
      self,
      graph: PipelineModulesGraph,
      chunks: int = 1,
      checkpoint: str = "except_last",
  ) -> None:
      super().__init__()
  
      check_pytorch_version()
  
      chunks = int(chunks)
      checkpoint = str(checkpoint)
  
      if chunks <= 0:
          raise ValueError("number of chunks must be positive integer")
      if checkpoint not in ["always", "except_last", "never"]:
          raise ValueError(
              "checkpoint is not one of 'always', 'except_last', or 'never'"
          )
  
      self.chunks = chunks
      # The micro-batch index where the checkpointing stops.
      checkpoint_stop = {
          "always": self.chunks,
          "except_last": self.chunks - 1,
          "never": 0,
      }[checkpoint]
  
      self.partitions = [
          self.Partition(
              nodes,
              rpc.remote(
                  handler.owner(),
                  PartitionHandler,
                  args=(
                      handler,
                      nodes[0].module.device,
                      len(nodes[0].inputs),
                      nodes[-1].num_outputs,
                      i,
                      self.chunks,
                      checkpoint_stop,
                  ),
              ),
          )
          for i, (nodes, handler) in enumerate(graph.partition_graph())
      ]
      self.input_consumers = [
          next(
              self.DataConsumer(
                  partition,
                  input_consumer.consumer_input_idx,
                  input_consumer.output_idx,
              )
              for partition in self.partitions
              if partition.nodes[0] is input_consumer.consumer
          )
          for input_consumer in graph.model_input_consumers
      ]
  
      self.graph = graph
  
file_path: fairscale/fairscale/experimental/nn/distributed_pipeline/pipeline.py
imports: from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Union
  import torch
  from torch import Tensor, nn
  from torch.distributed import rpc
  from fairscale.internal import torch_version
  from fairscale.nn.pipe import microbatch
  from .data import DataConsumer
  from .graph import Node, PipelineModulesGraph
  from .partition_handler import DistributedPipelineRecord, PartitionHandler
label_desc: Use image node_id 1 to create a new DistributedPipeline object from inherited base classes: nn with example: obj = DistributedPipeline(graph, chunks, checkpoint)