class_docstr: PipelineHelper assists with implementing pipelining in model parallelism.
  
  For a tutorial on model parallelism, as it's implemented in parts of ParlAI,
  see https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html.
  
  Usage:
  >>> my_model = PipelineHelper().make_parallel(my_model)
  
  Note that you will need to manually implement logic which handles the
  moved layers.
class_docstr_tok: ["PipelineHelper","assists","with","implementing","pipelining","in","model","parallelism",".","For","a","tutorial","on","model","parallelism",",","as","it","'s","implemented","in","parts","of","ParlAI",",","see","https",":","\/\/pytorch.org\/tutorials\/intermediate\/model_parallel_tutorial.html",".","Usage",":",">",">",">","my_model","=","PipelineHelper","(",")",".make_parallel","(","my_model",")","Note","that","you","will","need","to","manually","implement","logic","which","handles","the","moved","layers","."]
docstr: Allocate specific layers in a model to be ModelParallel.
  
  Limited to only ModuleLists within the model.  Uses some heuristics to
  attempt to evenly distribute layers across GPUs, in order to balance
  memory usage. They are:
  
  - Assume the 0th GPU will host the optimizer, word embeddings, etc.
  - Assume activation memory is linear with the number of parameters.
  - All layers are approximately equal in size.
docstr_tok: ["Allocate","specific","layers","in","a","model","to","be","ModelParallel",".","Limited","to","only","ModuleLists","within","the","model",".","Uses","some","heuristics","to","attempt","to","evenly","distribute","layers","across","GPUs",",","in","order","to","balance","memory","usage",".","They","are",":","-","Assume","the","0th","GPU","will","host","the","optimizer",",","word","embeddings",",","etc",".","-","Assume","activation","memory","is","linear","with","the","number","of","parameters",".","-","All","layers","are","approximately","equal","in","size","."]
code: def make_parallel(self, model: torch.nn.Module) -> torch.nn.Module:
      """
      Allocate specific layers in a model to be ModelParallel.
  
      Limited to only ModuleLists within the model.  Uses some heuristics to
      attempt to evenly distribute layers across GPUs, in order to balance
      memory usage. They are:
  
      - Assume the 0th GPU will host the optimizer, word embeddings, etc.
      - Assume activation memory is linear with the number of parameters.
      - All layers are approximately equal in size.
      """
  
      # first assume all layers will go onto gpu 0 as an optimizer. The
      # optimizer and embeddings are not quite as expensive as the
      # activations (which scale via batchsize), Empirically, I found this
      # heuristic works well enough. The weighting factor of 3 is more or
      # less made up.
      self.__device_allocations["cuda:0"] += (
          trainable_parameters(model) * 3
      )
  
      model.apply(self._place_modulelist)
      model._apply(self._move_rest_to_cuda0)  # type: ignore
      return model
  
imports: import os
  from typing import Union, Optional, Tuple, Any, List, Sized, TypeVar
  import itertools
  from collections import namedtuple
  import parlai.utils.logging
  import parlai.utils.io
  import torch.optim
label_desc: Use image node_id 3 for calling the PipelineHelper obj's underlying member method code with example usage: obj.make_parallel(model) and returns: model