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: Iterate through chunks and layers that should be pipelined.
  
  Each iteration of this generator yields the following properties:
  
      - layer_nos: a list of indices of layers for you to forward through
      - chunk_idx: the index of the chunk we are manipulating. Use this
        if you need to update chunk representations.
      - next_device: where the chunk should be moved to AFTER the layer
        computation is done.
docstr_tok: ["Iterate","through","chunks","and","layers","that","should","be","pipelined",".","Each","iteration","of","this","generator","yields","the","following","properties",":","-","layer_nos",":","a","list","of","indices","of","layers","for","you","to","forward","through","-","chunk_idx",":","the","index","of","the","chunk","we","are","manipulating",".","Use","this","if","you","need","to","update","chunk","representations",".","-","next_device",":","where","the","chunk","should","be","moved","to","AFTER","the","layer","computation","is","done","."]
code: def schedule_work_items(
      layers: torch.nn.ModuleList, chunks: List[Chunk]
  ):
      """
      Iterate through chunks and layers that should be pipelined.
  
      Each iteration of this generator yields the following properties:
  
          - layer_nos: a list of indices of layers for you to forward through
          - chunk_idx: the index of the chunk we are manipulating. Use this
            if you need to update chunk representations.
          - next_device: where the chunk should be moved to AFTER the layer
            computation is done.
      """
      # We want to pipeline our computations so that each GPU is working on
      # chunks of the problem at the same of the time. The load of the will
      # look like this, assuming there are 5 chunks (A, B, C, D, E) and 4
      # GPUs. Each slot fill means that gpu is working on that chunk.
      #
      #         +-----------------+
      #         |       Time      |
      #         | 1 2 3 4 5 6 7 8 |
      # +-------+-----------------+
      # |  G  0 | A B C D E       |
      # |  P  1 |   A B C D E     |
      # |  U  2 |     A B C D E   |
      # |     3 |       A B C D E |
      # +-------+-----------------+
      #
      # Note that some GPUs will be idle much of the time. In reality, we
      # will use 2 * num_gpus as the number of chunks, to minimize idle
      # time.
      num_chunks = len(chunks)
      for l in layers:
          if not hasattr(l, "_mp_gpu"):
              raise RuntimeError(
                  "You must run PipelineHelper.make_parallel on the ModuleList "
                  "before you can use iterate_layers_chunks."
              )
  
      # devices maps device_idx -> (device, [layer_idx, layer_idx, ...])
      # for example, if devices is 2 and there are 4 layers, we might have
      #   devices = {
      #     0: ('cuda:0', [0]),
      #     1: ('cuda:1', [1, 2, 3]),
      #   }
      # This means layers 0 is on cuda:0, but layers 1-3 are on cuda:1.
      devices = {
          device_idx: (dev, list(grp))
          for device_idx, (dev, grp) in enumerate(
              itertools.groupby(
                  range(len(layers)), lambda x: layers[x]._mp_gpu
              )
          )
      }
      num_timesteps = len(devices) + num_chunks
      for timestep in range(num_timesteps):
          for chunk_idx in range(num_chunks):
              device_idx = timestep - chunk_idx
              if device_idx >= 0 and device_idx < len(devices):
                  dev, layers_nos = devices[device_idx]
                  next_device, _ = devices[
                      (device_idx + 1) % len(devices)
                  ]
                  assert device_idx in devices
                  yield PipelineWorkItem(
                      chunk_idx=chunk_idx,
                      layer_nos=layers_nos,
                      next_device=next_device,
                  )
  
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 10 for calling the PipelineHelper obj's underlying member method code with example usage: obj.schedule_work_items(layers, chunks) without return types