class_docstr: Wrap an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>` and shards its states across ranks in the group.
  
  The sharing is done as described by ZeRO_.
  
  The local optimizer instance in each rank is only
  responsible for updating approximately ``1 / world_size`` parameters and
  hence only needs to keep ``1 / world_size`` optimizer states. After
  parameters are updated locally, each rank will broadcast its parameters to
  all other peers to keep all model replicas in the same state.
  ``ZeroRedundancyOptimizer`` can be used in conjunction with
  :class:`torch.nn.parallel.DistributedDataParallel` to reduce per-rank peak
  memory consumption.
  
  ``ZeroRedundancyOptimizer`` uses a sorted-greedy algorithm to pack a number
  of parameters at each rank. Each parameter belongs to a single rank and is
  not divided among ranks. The partition is arbitrary and might not match the
  the parameter registration or usage order.
  
  Arguments:
      params (``Iterable``): an ``Iterable`` of :class:`torch.Tensor` s
          or :class:`dict` s giving all parameters, which will be sharded
          across ranks.
  
  Keyword Args:
      optimizer_class (:class:`torch.nn.Optimizer`): the class of the local
          optimizer.
      process_group (``ProcessGroup``, optional): ``torch.distributed``
          ``ProcessGroup`` (default: ``dist.group.WORLD`` initialized by
          :meth:`torch.distributed.init_process_group`).
      parameters_as_bucket_view (bool, optional): if ``True``, parameters are
          packed into buckets to speed up communication, and ``param.data``
          fields point to bucket views at different offsets; if ``False``,
          each individual parameter is communicated separately, and each
          ``params.data`` stays intact (default: ``False``).
      overlap_with_ddp (bool, optional): if ``True``, :meth:`step` is
          overlapped with :class:`DistributedDataParallel` 's gradient
          synchronization; this requires (1) either a functional optimizer
          for the ``optimizer_class`` argument or one with a functional
          equivalent and (2) registering a DDP communication hook
          constructed from one of the functions in ``ddp_zero_hook.py``;
          parameters are packed into buckets matching those in
          :class:`DistributedDataParallel`, meaning that the
          ``parameters_as_bucket_view`` argument is ignored.
          If ``False``, :meth:`step` runs disjointly after the backward pass
          (per normal).
          (default: ``False``)
      **defaults: any trailing arguments, which are forwarded to the local
          optimizer.
  
  Example::
  
      >>> # xdoctest: +SKIP
      >>> import torch.nn as nn
      >>> from torch.distributed.optim import ZeroRedundancyOptimizer
      >>> from torch.nn.parallel import DistributedDataParallel as DDP
      >>> model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
      >>> ddp = DDP(model, device_ids=[rank])
      >>> opt = ZeroRedundancyOptimizer(
      >>>     ddp.parameters(),
      >>>     optimizer_class=torch.optim.Adam,
      >>>     lr=0.01
      >>> )
      >>> ddp(inputs).sum().backward()
      >>> opt.step()
  
  .. warning::
      Currently, ``ZeroRedundancyOptimizer`` requires that all of the
      passed-in parameters are the same dense type.
  
  .. warning::
      If you pass ``overlap_with_ddp=True``, be wary of the following: Given
      the way that overlapping :class:`DistributedDataParallel` with
      :class:`ZeroRedundancyOptimizer` is currently implemented, the first
      two or three training iterations do not perform parameter updates in
      the optimizer step, depending on if ``static_graph=False`` or
      ``static_graph=True``, respectively. This is because it needs
      information about the gradient bucketing strategy used by
      :class:`DistributedDataParallel`, which is not finalized until the
      second forward pass if ``static_graph=False`` or until the third
      forward pass if ``static_graph=True``. To adjust for this, one option
      is to prepend dummy inputs.
  
  .. warning:: ZeroRedundancyOptimizer is experimental and subject to change.
  
  .. _ZeRO: https://arxiv.org/abs/1910.02054
class_docstr_tok: ["Wrap","an","arbitrary",":","class",":","`","optim.Optimizer","<","torch.optim.Optimizer",">","`","and","shards","its","states","across","ranks","in","the","group",".","The","sharing","is","done","as","described","by","ZeRO_",".","The","local","optimizer","instance","in","each","rank","is","only","responsible","for","updating","approximately","``","1","\/","world_size","``","parameters","and","hence","only","needs","to","keep","``","1","\/","world_size","``","optimizer","states",".","After","parameters","are","updated","locally",",","each","rank","will","broadcast","its","parameters","to","all","other","peers","to","keep","all","model","replicas","in","the","same","state",".","``","ZeroRedundancyOptimizer","``","can","be","used","in","conjunction","with",":","class",":","`","torch.nn.parallel.DistributedDataParallel","`","to","reduce","per-rank","peak","memory","consumption",".","``","ZeroRedundancyOptimizer","``","uses","a","sorted-greedy","algorithm","to","pack","a","number","of","parameters","at","each","rank",".","Each","parameter","belongs","to","a","single","rank","and","is","not","divided","among","ranks",".","The","partition","is","arbitrary","and","might","not","match","the","the","parameter","registration","or","usage","order",".","Arguments",":","params","(","``","Iterable","``",")",":","an","``","Iterable","``","of",":","class",":","`","torch.Tensor","`","s","or",":","class",":","`","dict","`","s","giving","all","parameters",",","which","will","be","sharded","across","ranks",".","Keyword","Args",":","optimizer_class","(",":","class",":","`","torch.nn.Optimizer","`",")",":","the","class","of","the","local","optimizer",".","process_group","(","``","ProcessGroup","``",",","optional",")",":","``","torch.distributed","``","``","ProcessGroup","``","(","default",":","``","dist.group.WORLD","``","initialized","by",":","meth",":","`","torch.distributed.init_process_group","`",")",".","parameters_as_bucket_view","(","bool",",","optional",")",":","if","``","True","``",",","parameters","are","packed","into","buckets","to","speed","up","communication",",","and","``","param.data","``","fields","point","to","bucket","views","at","different","offsets",";","if","``","False","``",",","each","individual","parameter","is","communicated","separately",",","and","each","``","params.data","``","stays","intact","(","default",":","``","False","``",")",".","overlap_with_ddp","(","bool",",","optional",")",":","if","``","True","``",",",":","meth",":","`","step","`","is","overlapped","with",":","class",":","`","DistributedDataParallel","`","'s","gradient","synchronization",";","this","requires","(","1",")","either","a","functional","optimizer","for","the","``","optimizer_class","``","argument","or","one","with","a","functional","equivalent","and","(","2",")","registering","a","DDP","communication","hook","constructed","from","one","of","the","functions","in","``","ddp_zero_hook.py","``",";","parameters","are","packed","into","buckets","matching","those","in",":","class",":","`","DistributedDataParallel","`",",","meaning","that","the","``","parameters_as_bucket_view","``","argument","is","ignored",".","If","``","False","``",",",":","meth",":","`","step","`","runs","disjointly","after","the","backward","pass","(","per","normal",")",".","(","default",":","``","False","``",")","*","*","defaults",":","any","trailing","arguments",",","which","are","forwarded","to","the","local","optimizer",".","Example",":",":",">",">",">","#","xdoctest",":","+SKIP",">",">",">","import","torch.nn","as","nn",">",">",">","from","torch.distributed.optim","import","ZeroRedundancyOptimizer",">",">",">","from","torch.nn.parallel","import","DistributedDataParallel","as","DDP",">",">",">","model","=","nn.Sequential","(","*","[","nn.Linear","(","2000",",","2000",")",".to","(","rank",")","for","_","in","range","(","20",")","]",")",">",">",">","ddp","=","DDP","(","model",",","device_ids=","[","rank","]",")",">",">",">","opt","=","ZeroRedundancyOptimizer","(",">",">",">","ddp.parameters","(",")",",",">",">",">","optimizer_class=torch.optim.Adam",",",">",">",">","lr=0.01",">",">",">",")",">",">",">","ddp","(","inputs",")",".sum","(",")",".backward","(",")",">",">",">","opt.step","(",")","..","warning",":",":","Currently",",","``","ZeroRedundancyOptimizer","``","requires","that","all","of","the","passed-in","parameters","are","the","same","dense","type",".","..","warning",":",":","If","you","pass","``","overlap_with_ddp=True","``",",","be","wary","of","the","following",":","Given","the","way","that","overlapping",":","class",":","`","DistributedDataParallel","`","with",":","class",":","`","ZeroRedundancyOptimizer","`","is","currently","implemented",",","the","first","two","or","three","training","iterations","do","not","perform","parameter","updates","in","the","optimizer","step",",","depending","on","if","``","static_graph=False","``","or","``","static_graph=True","``",",","respectively",".","This","is","because","it","needs","information","about","the","gradient","bucketing","strategy","used","by",":","class",":","`","DistributedDataParallel","`",",","which","is","not","finalized","until","the","second","forward","pass","if","``","static_graph=False","``","or","until","the","third","forward","pass","if","``","static_graph=True","``",".","To","adjust","for","this",",","one","option","is","to","prepend","dummy","inputs",".","..","warning",":",":","ZeroRedundancyOptimizer","is","experimental","and","subject","to","change",".","..","_ZeRO",":","https",":","\/\/arxiv.org\/abs\/1910.02054"]
docstr: Return the optimizer constructor using validation and transformation depending on ``overlap_with_ddp``.
  
  Returns:
      - ``optimizer_class`` if ``overlap_with_ddp=False`` and
          ``optimizer_class`` is not a functional optimizer.
      - ``optimizer_class`` if ``overlap_with_ddp=True`` and
          ``optimizer_class`` is already a functional optimizer.
      - The functional equivalent of ``optimizer_class`` if
          ``overlap_with_ddp=True`` and ``optimizer_class`` is not
          already a functional optimizer (assuming the equivalent
          exists).
  
  Raises:
      ValueError:
  
          - if ``overlap_with_ddp=True`` but ``optimizer_class`` is
              neither a functional optimizer nor translatable to a
              functional optimizer.
          - if ``overlap_with_ddp=False`` and ``optimizer_class`` is a
              functional optimizer.
docstr_tok: ["Return","the","optimizer","constructor","using","validation","and","transformation","depending","on","``","overlap_with_ddp","``",".","Returns",":","-","``","optimizer_class","``","if","``","overlap_with_ddp=False","``","and","``","optimizer_class","``","is","not","a","functional","optimizer",".","-","``","optimizer_class","``","if","``","overlap_with_ddp=True","``","and","``","optimizer_class","``","is","already","a","functional","optimizer",".","-","The","functional","equivalent","of","``","optimizer_class","``","if","``","overlap_with_ddp=True","``","and","``","optimizer_class","``","is","not","already","a","functional","optimizer","(","assuming","the","equivalent","exists",")",".","Raises",":","ValueError",":","-","if","``","overlap_with_ddp=True","``","but","``","optimizer_class","``","is","neither","a","functional","optimizer","nor","translatable","to","a","functional","optimizer",".","-","if","``","overlap_with_ddp=False","``","and","``","optimizer_class","``","is","a","functional","optimizer","."]
code: def _get_optimizer_constructor(self, optimizer_class: Any) -> Any:
      r"""
      Return the optimizer constructor using validation and transformation depending on ``overlap_with_ddp``.
  
      Returns:
          - ``optimizer_class`` if ``overlap_with_ddp=False`` and
              ``optimizer_class`` is not a functional optimizer.
          - ``optimizer_class`` if ``overlap_with_ddp=True`` and
              ``optimizer_class`` is already a functional optimizer.
          - The functional equivalent of ``optimizer_class`` if
              ``overlap_with_ddp=True`` and ``optimizer_class`` is not
              already a functional optimizer (assuming the equivalent
              exists).
  
      Raises:
          ValueError:
  
              - if ``overlap_with_ddp=True`` but ``optimizer_class`` is
                  neither a functional optimizer nor translatable to a
                  functional optimizer.
              - if ``overlap_with_ddp=False`` and ``optimizer_class`` is a
                  functional optimizer.
      """
      functional_optims = functional_optim_map.values()
      if not self._overlap_with_ddp:
          if optimizer_class in functional_optims:
              # Using a functional optimizer is only supported when
              # `overlap_with_ddp=True`
              raise ValueError(
                  f"Passing in a functional optimizer {optimizer_class} "
                  "when `overlap_with_ddp=False`"
              )
          else:
              return optimizer_class
      else:
          if optimizer_class in functional_optims:
              # Already a functional optimizer
              return optimizer_class
          elif optimizer_class in functional_optim_map:
              # Translate the passed-in optimizer class to its functional
              # equivalent if `overlap_with_ddp=True`
              optim_constructor = functional_optim_map[optimizer_class]
              logger.info(
                  "Using the functional optimizer %s "
                  "instead of %s since "
                  "`overlap_with_ddp=True`",
                  optim_constructor,
                  optimizer_class,
              )
              return optim_constructor
          else:
              raise ValueError(
                  "Using `ddp_with_overlap=True` requires using a "
                  "functional optimizer, but there is no supported functional "
                  f"optimizer equivalent for {optimizer_class}"
              )
  
imports: import collections
  import copy
  import enum
  import inspect
  import io
  import logging
  from itertools import chain
  from typing import Any, Callable, Dict, List, Optional, Set, Type, Union
  import torch
  import torch.distributed
  from torch.distributed.algorithms.join import Join, Joinable, JoinHook
  from torch.distributed.optim.utils import functional_optim_map
  from torch.optim import Optimizer
label_desc: Use image node_id 34 for calling the ZeroRedundancyOptimizer obj's underlying member method code with example usage: obj._get_optimizer_constructor(optimizer_class) and returns: optimizer_class, optimizer_class, optim_constructor