class_docstr: Modules to be parallelized with pipeline parallelism.
  
  The key constraint that enables pipeline parallelism is the
  representation of the forward pass as a sequence of layers
  and the enforcement of a simple interface between them. The
  forward pass is implicitly defined by the module ``layers``. The key
  assumption is that the output of each layer can be directly fed as
  input to the next, like a ``torch.nn.Sequence``. The forward pass is
  implicitly:
  
  .. code-block:: python
  
      def forward(self, inputs):
          x = inputs
          for layer in self.layers:
              x = layer(x)
          return x
  
  .. note::
      Pipeline parallelism is not compatible with ZeRO-2 and ZeRO-3.
  
  Args:
      layers (Iterable): A sequence of layers defining pipeline structure. Can be a ``torch.nn.Sequential`` module.
      num_stages (int, optional): The degree of pipeline parallelism. If not specified, ``topology`` must be provided.
      topology (``deepspeed.runtime.pipe.ProcessTopology``, optional): Defines the axes of parallelism axes for training. Must be provided if ``num_stages`` is ``None``.
      loss_fn (callable, optional): Loss is computed ``loss = loss_fn(outputs, label)``
      seed_layers(bool, optional): Use a different seed for each layer. Defaults to False.
      seed_fn(type, optional): The custom seed generating function. Defaults to random seed generator.
      base_seed (int, optional): The starting seed. Defaults to 1234.
      partition_method (str, optional): The method upon which the layers are partitioned. Defaults to 'parameters'.
      activation_checkpoint_interval (int, optional): The granularity activation checkpointing in terms of number of layers. 0 disables activation checkpointing.
      activation_checkpoint_func (callable, optional): The function to use for activation checkpointing. Defaults to ``deepspeed.checkpointing.checkpoint``.
      checkpointable_layers(list, optional): Checkpointable layers may not be checkpointed. Defaults to None which does not additional filtering.
class_docstr_tok: ["Modules","to","be","parallelized","with","pipeline","parallelism",".","The","key","constraint","that","enables","pipeline","parallelism","is","the","representation","of","the","forward","pass","as","a","sequence","of","layers","and","the","enforcement","of","a","simple","interface","between","them",".","The","forward","pass","is","implicitly","defined","by","the","module","``","layers","``",".","The","key","assumption","is","that","the","output","of","each","layer","can","be","directly","fed","as","input","to","the","next",",","like","a","``","torch.nn.Sequence","``",".","The","forward","pass","is","implicitly",":","..","code-block",":",":","python","def","forward","(","self",",","inputs",")",":","x","=","inputs","for","layer","in","self.layers",":","x","=","layer","(","x",")","return","x","..","note",":",":","Pipeline","parallelism","is","not","compatible","with","ZeRO-2","and","ZeRO-3",".","Args",":","layers","(","Iterable",")",":","A","sequence","of","layers","defining","pipeline","structure",".","Can","be","a","``","torch.nn.Sequential","``","module",".","num_stages","(","int",",","optional",")",":","The","degree","of","pipeline","parallelism",".","If","not","specified",",","``","topology","``","must","be","provided",".","topology","(","``","deepspeed.runtime.pipe.ProcessTopology","``",",","optional",")",":","Defines","the","axes","of","parallelism","axes","for","training",".","Must","be","provided","if","``","num_stages","``","is","``","None","``",".","loss_fn","(","callable",",","optional",")",":","Loss","is","computed","``","loss","=","loss_fn","(","outputs",",","label",")","``","seed_layers","(","bool",",","optional",")",":","Use","a","different","seed","for","each","layer",".","Defaults","to","False",".","seed_fn","(","type",",","optional",")",":","The","custom","seed","generating","function",".","Defaults","to","random","seed","generator",".","base_seed","(","int",",","optional",")",":","The","starting","seed",".","Defaults","to","1234.","partition_method","(","str",",","optional",")",":","The","method","upon","which","the","layers","are","partitioned",".","Defaults","to","'parameters","'",".","activation_checkpoint_interval","(","int",",","optional",")",":","The","granularity","activation","checkpointing","in","terms","of","number","of","layers",".","0","disables","activation","checkpointing",".","activation_checkpoint_func","(","callable",",","optional",")",":","The","function","to","use","for","activation","checkpointing",".","Defaults","to","``","deepspeed.checkpointing.checkpoint","``",".","checkpointable_layers","(","list",",","optional",")",":","Checkpointable","layers","may","not","be","checkpointed",".","Defaults","to","None","which","does","not","additional","filtering","."]
code: def __init__(
      self,
      layers,
      num_stages=None,
      topology=None,
      loss_fn=None,
      seed_layers=False,
      seed_fn=None,
      base_seed=1234,
      partition_method="parameters",
      activation_checkpoint_interval=0,
      activation_checkpoint_func=checkpointing.checkpoint,
      checkpointable_layers=None,
  ):
      super().__init__()
  
      if num_stages is None and topology is None:
          raise RuntimeError("must provide num_stages or topology")
  
      self.micro_offset = 0
  
      self.loss_fn = loss_fn
  
      self.checkpointable_layers = checkpointable_layers
      if checkpointable_layers is not None:
          assert isinstance(
              checkpointable_layers, list
          ), "param `checkpointable_layers` must be type of list."
  
      self.seed_layers = seed_layers
      self.seed_fn = seed_fn
      self.base_seed = base_seed
      if dist.get_rank() == 0:
          try:
              seed_str = self.seed_fn.__name__
          except AttributeError:
              seed_str = None
          print(
              f"SEED_LAYERS={self.seed_layers} BASE_SEED={self.base_seed} SEED_FN={seed_str}"
          )
  
      # Setup world info
      self.world_group = dist.new_group(
          ranks=range(dist.get_world_size())
      )
      self.global_rank = dist.get_rank(group=self.world_group)
      self.world_size = dist.get_world_size(group=self.world_group)
      self.local_rank = int(os.environ.get("LOCAL_RANK", None))
      assert self.local_rank != None
  
      if topology:
          self._topo = topology
          self.num_stages = self._topo.get_dim("pipe")
      else:
          self.num_stages = num_stages
          if topology is None:
              if self.world_size % self.num_stages != 0:
                  raise RuntimeError(
                      f"num_stages ({self.num_stages}) must divide distributed world size ({self.world_size})"
                  )
              dp = self.world_size // num_stages
              topology = PipeDataParallelTopology(
                  num_pp=num_stages, num_dp=dp
              )
              self._topo = topology
  
      # Construct communicators for pipeline topology
      self._grid = PipelineParallelGrid(
          process_group=self.world_group, topology=self._topo
      )
  
      self.stage_id = self._topo.get_coord(self.global_rank).pipe
  
      # Initialize partition information
      self._layer_specs = list(layers)
      self._num_layers = len(self._layer_specs)
      self._local_start = 0
      self._local_stop = None
      self._partition_layers(method=partition_method)
  
      self.forward_funcs = []
      self.fwd_map = {}
      self.tied_modules = nn.ModuleDict()
      self.tied_weight_attrs = {}
  
      # Offset the random seed by the stage ID.
      # newseed = torch.cuda.initial_seed() + self._grid.get_stage_id()
      # ds_utils.set_random_seed(newseed)
  
      # with torch.random.fork_rng(devices=[torch.cuda.current_device()]):
      self._build()
      self.to(f"cuda:{self.local_rank}")
  
      self.tied_comms = self._index_tied_modules()
      self._synchronize_tied_weights()
  
      self.activation_checkpoint_interval = (
          activation_checkpoint_interval
      )
      self.activation_checkpoint_func = activation_checkpoint_func
  
imports: import os
  import glob
  import re
  from functools import partial
  import torch
  import torch.nn
  from deepspeed import comm
  from deepspeed.utils import logger
  from ..None import utils
  from ..activation_checkpointing import checkpointing
  from .topology import PipeDataParallelTopology, PipelineParallelGrid
  from deepspeed.runtime.state_dict_factory import SDLoaderFactory
label_desc: Use image node_id 1 to create a new PipelineModule object from inherited base classes: nn with example: obj = PipelineModule(layers, num_stages, topology, loss_fn, seed_layers, seed_fn, base_seed, partition_method, activation_checkpoint_interval, activation_checkpoint_func, checkpointable_layers)