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","."]
docstr: Build communication structures for tied modules. 
docstr_tok: ["Build","communication","structures","for","tied","modules","."]
code: def _index_tied_modules(self):
      """Build communication structures for tied modules."""
      tied_comms = {}
      if self._topo.get_dim("pipe") == 1:
          return tied_comms
  
      specs = self._layer_specs
      tie_keys = set(
          s.key for s in specs if isinstance(s, TiedLayerSpec)
      )
      for key in tie_keys:
          # Find the layers that the tied module appears in
          tied_layers = []
          for idx, layer in enumerate(specs):
              if isinstance(layer, TiedLayerSpec) and layer.key == key:
                  tied_layers.append(idx)
          # Find all stages with this tied module
          # TODO: Would be nice to remove the nested data/model parallelism loops and
          # TODO: instead generalize in some way, since we really just care about the
          # TODO: stage that owns the tied layer. Then loop over each (dp, mp, ...)
          # TODO: fiber to generate process groups.
          tied_stages = set(
              self.stage_owner(idx) for idx in tied_layers
          )
          for dp in range(self._grid.data_parallel_size):
              for mp in range(
                  self._grid.get_slice_parallel_world_size()
              ):
                  tied_ranks = []
                  for s in sorted(tied_stages):
                      if self._grid.get_slice_parallel_world_size() > 1:
                          tied_ranks.append(
                              self._grid.stage_to_global(
                                  stage_id=s, data=dp, model=mp
                              )
                          )
                      else:
                          tied_ranks.append(
                              self._grid.stage_to_global(
                                  stage_id=s, data=dp
                              )
                          )
                  group = dist.new_group(ranks=tied_ranks)
  
                  # Record this tied module if we own a local copy of it.
                  if self.global_rank in tied_ranks:
                      assert key in self.tied_modules
                      if key in self.tied_modules:
                          tied_comms[key] = {
                              "ranks": tied_ranks,
                              "group": group,
                              "weight_attr": self.tied_weight_attrs[
                                  key
                              ],
                              "module": self.tied_modules[key],
                          }
                          # Only count the tied module once in the eyes of the FP16 optimizer
                          if self.global_rank != tied_ranks[0]:
                              for p in self.tied_modules[
                                  key
                              ].parameters():
                                  p.ds_pipe_replicated = True
      """
          if len(tied_comms) > 0:
              print(f'RANK={self.global_rank} tied_comms={tied_comms}')
          """
  
      return tied_comms
  
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 10 for calling the PipelineModule obj's underlying member method code with example usage: obj._index_tied_modules() and returns: tied_comms, tied_comms