shrinath-suresh/blogs-docs-splitted huggingface dataset:
the full dict in CPU memory, where rank 0 moves each shard to
        GPU memory (for NCCL) and communicates it to ranks
        appropriately. Hence, the former has higher aggregate CPU
        memory cost, while the latter has higher communication cost.
  Parameters:
     * **full_optim_state_dict** (*Optional**[**Dict**[**str**,
       **Any**]**]*) -- Optimizer state dict corresponding to the
       unflattened parameters and holding the full non-sharded
       optimizer state if on rank 0; the argument is ignored on
       nonzero ranks.

     * **model** (*torch.nn.Module*) -- Root module (which may or
       may not be a "FullyShardedDataParallel" instance) whose
       parameters correspond to the optimizer state in
       "full_optim_state_dict".

     * **optim_input**
       (*Optional**[**Union**[**List**[**Dict**[**str**,
       **Any**]**]**, **Iterable**[**torch.nn.Parameter**]**]**]*)
-- Input passed into the optimizer representing either a
           "list" of parameter groups or an iterable of parameters; if
           "None", then this method assumes the input was
           "model.parameters()". This argument is deprecated, and
           there is no need to pass it in anymore. (Default: "None")
     * **optim** (*Optional**[**torch.optim.Optimizer**]*) --
       Optimizer that will load the state dict returned by this
       method. This is the preferred argument to use over
       "optim_input". (Default: "None")

     * **group** (*dist.ProcessGroup*) -- Model's process group or
       "None" if using the default process group. (Default:
       "None")

  Returns:
     The full optimizer state dict now remapped to flattened
     parameters instead of unflattened parameters and restricted
     to only include this rank's part of the optimizer state.

  Return type:
     Dict[str, Any]
Return type:
         Dict[str, Any]
static set_state_dict_type(module, state_dict_type, state_dict_config=None)
  Set the "state_dict_type" and the corresponding (optional)
  configurations of all the descendant FSDP modules of the target
  module. The target module does not have to be a FSDP module. If
  the target module is a FSDP module, its "state_dict_type" will
  also be changed.

  Note:

    This API should be called for only the top-level (root)
    module.

  Note:

    This API enables users to transparently use the conventional
    "state_dict" API to take model checkpoints in cases where the
    root FSDP module is wrapped by another "nn.Module". For
    example, the following will ensure "state_dict" is called on
    all non-FSDP instances, while dispatching into
    *sharded_state_dict* implementation for FSDP:

  Example:

     >>> model = DDP(FSDP(...))
     >>> FSDP.set_state_dict_type(
FSDP.set_state_dict_type(
         >>>     model,
         >>>     StateDictType.SHARDED_STATE_DICT,
         >>>     ShardedStateDictConfig(offload_to_cpu=True),
         >>> )
         >>> checkpoint = model.state_dict()



  Parameters:
     * **module** (*torch.nn.Module*) -- Root module.

     * **state_dict_type** (*StateDictType*) -- the desired
       "state_dict_type" to set.

     * **state_dict_config** (*Optional**[**StateDictConfig**]*)
       -- the configuration for the target "state_dict_type".

  Return type:
     *Tuple*[*StateDictType*, *StateDictConfig*]

static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None)
  Shards the full optimizer state dict "full_optim_state_dict" by
  remapping the state to flattened parameters instead of
  unflattened parameters and restricting to only this rank's part
  of the optimizer state. The first argument should be the return
value of "full_optim_state_dict()".
  Example:

     >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
     >>> model, optim = ...
     >>> full_osd = FSDP.full_optim_state_dict(model, optim)
     >>> torch.save(full_osd, PATH)
     >>> # Define new model with possibly different world size
     >>> new_model, new_optim = ...
     >>> full_osd = torch.load(PATH)
     >>> sharded_osd = FSDP.shard_full_optim_state_dict(full_osd, new_model)
     >>> new_optim.load_state_dict(sharded_osd)

  Note:

    Both "shard_full_optim_state_dict()" and
    "scatter_full_optim_state_dict()" may be used to get the
    sharded optimizer state dict to load. Assuming that the full
    optimizer state dict resides in CPU memory, the former
    requires each rank to have the full dict in CPU memory, where
    each rank individually shards the dict without any
    communication, while the latter requires only rank 0 to have
the full dict in CPU memory, where rank 0 moves each shard to
        GPU memory (for NCCL) and communicates it to ranks
        appropriately. Hence, the former has higher aggregate CPU
        memory cost, while the latter has higher communication cost.
  Parameters:
     * **full_optim_state_dict** (*Dict**[**str**, **Any**]*) --
       Optimizer state dict corresponding to the unflattened
       parameters and holding the full non-sharded optimizer
       state.

     * **model** (*torch.nn.Module*) -- Root module (which may or
       may not be a "FullyShardedDataParallel" instance) whose
       parameters correspond to the optimizer state in
       "full_optim_state_dict".

     * **optim_input**
       (*Optional**[**Union**[**List**[**Dict**[**str**,
       **Any**]**]**, **Iterable**[**torch.nn.Parameter**]**]**]*)
       -- Input passed into the optimizer representing either a
"list" of parameter groups or an iterable of parameters; if
           "None", then this method assumes the input was
           "model.parameters()". This argument is deprecated, and
           there is no need to pass it in anymore. (Default: "None")
     * **optim** (*Optional**[**torch.optim.Optimizer**]*) --
       Optimizer that will load the state dict returned by this
       method. This is the preferred argument to use over
       "optim_input". (Default: "None")

  Returns:
     The full optimizer state dict now remapped to flattened
     parameters instead of unflattened parameters and restricted
     to only include this rank's part of the optimizer state.

  Return type:
     Dict[str, Any]

static sharded_optim_state_dict(model, optim, group=None)
  The API is similar to "full_optim_state_dict()" but this API
  chunks all non-zero-dimension states to "ShardedTensor" to save
memory. This API should only be used when the model "state_dict"
      is derived with the context manager "with
      state_dict_type(SHARDED_STATE_DICT):".
  For the detailed usage, refer to "full_optim_state_dict()".

  Warning:

    The returned state dict contains "ShardedTensor" and cannot be
    directly used by the regular "optim.load_state_dict".

  Return type:
     *Dict*[str, *Any*]

static state_dict_type(module, state_dict_type, state_dict_config=None)
  A context manager to set the "state_dict_type" of all the
  descendant FSDP modules of the target module. This context
  manager has the same functions as "set_state_dict_type()". Read
  the document of "set_state_dict_type()" for the detail.

  Example:

     >>> model = DDP(FSDP(...))
     >>> with FSDP.state_dict_type(
     >>>     model,
     >>>     StateDictType.SHARDED_STATE_DICT,
     >>> ):
     >>>     checkpoint = model.state_dict()

  Parameters:
Parameters:
         * module (torch.nn.Module) -- Root module.
     * **state_dict_type** (*StateDictType*) -- the desired
       "state_dict_type" to set.

     * **state_dict_config** (*Optional**[**StateDictConfig**]*)
       -- the configuration for the target "state_dict_type".

  Return type:
     *Generator*

static summon_full_params(module, recurse=True, writeback=True, rank0_only=False, offload_to_cpu=False, with_grads=False)
  A context manager to expose full params for FSDP instances. Can
  be useful *after* forward/backward for a model to get the params
  for additional processing or checking. It can take a non-FSDP
  module and will summon full params for all contained FSDP
  modules as well as their children, depending on the "recurse"
  argument.

  Note:

    This can be used on inner FSDPs.

  Note:

    This can *not* be used within a forward or backward pass. Nor
can forward and backward be started from within this context.
  Note:

    Parameters will revert to their local shards after the context
    manager exits, storage behavior is the same as forward.

  Note:

    The full parameters can be modified, but only the portion
    corresponding to the local param shard will persist after the
    context manager exits (unless "writeback=False", in which case
    changes will be discarded). In the case where FSDP does not
    shard the parameters, currently only when "world_size == 1",
    or "NO_SHARD" config, the modification is persisted regardless
    of "writeback".

  Note:

    This method works on modules which are not FSDP themselves but
    may contain multiple independent FSDP units. In that case, the
    given arguments will apply to all contained FSDP units.

  Warning:

    Note that "rank0_only=True" in conjunction with
"writeback=True" is not currently supported and will raise an
        error. This is because model parameter shapes would be
        different across ranks within the context, and writing to them
        can lead to inconsistency across ranks when the context is
        exited.
  Warning:

    Note that "offload_to_cpu" and "rank0_only=False" will result
    in full parameters being redundantly copied to CPU memory for
    GPUs that reside on the same machine, which may incur the risk
    of CPU OOM. It is recommended to use "offload_to_cpu" with
    "rank0_only=True".

  Parameters:
     * **recurse** (*bool**, **Optional*) -- recursively summon
       all params for nested FSDP instances (default: True).

     * **writeback** (*bool**, **Optional*) -- if "False",
       modifications to params are discarded after the context
       manager exits; disabling this can be slightly more
       efficient (default: True)
efficient (default: True)
     * **rank0_only** (*bool**, **Optional*) -- if "True", full
       parameters are materialized on only global rank 0. This
       means that within the context, only rank 0 will have full
       parameters and the other ranks will have sharded
       parameters. Note that setting "rank0_only=True" with
       "writeback=True" is not supported, as model parameter
       shapes will be different across ranks within the context,
       and writing to them can lead to inconsistency across ranks
       when the context is exited.

     * **offload_to_cpu** (*bool**, **Optional*) -- If "True",
       full parameters are offloaded to CPU. Note that this
       offloading currently only occurs if the parameter is
       sharded (which is only not the case for world_size = 1 or
       "NO_SHARD" config). It is recommended to use
       "offload_to_cpu" with "rank0_only=True" to avoid redundant
copies of model parameters being offloaded to the same CPU
           memory.
     * **with_grads** (*bool**, **Optional*) -- If "True",
       gradients are also unsharded with the parameters.
       Currently, this is only supported when passing
       "use_orig_params=True" to the FSDP constructor and
       "offload_to_cpu=False" to this method. (Default: "False")

  Return type:
     *Generator*

class torch.distributed.fsdp.BackwardPrefetch(value)
This configures explicit backward prefetching, which can improve
   throughput but may slightly increase peak memory usage.
For NCCL backend, any collectives, even if issued in different
   streams, contend for the same per-device NCCL stream, which is why
   the relative order in which the collectives are issued matters for
   overlapping. The different backward prefetching settings correspond
   to different orderings.

"BACKWARD_PRE": This prefetches the next set of parameters before
the current set of parameter's gradient computation. This
     improves backward pass throughput by overlapping communication
     (next all-gather) and computation (current gradient computation).

"BACKWARD_POST": This prefetches the next set of parameters after
     the current set of parameter's gradient computation. This may
     improve backward pass throughput by overlapping communication
     (current reduce-scatter) and computation (next gradient
     computation). Specifically, the next all-gather is reordered to
     be before the current reduce-scatter.

Note:
 If the increase in peak memory usage from prefetching is an
 issue, you may consider passing "limit_all_gathers=True" to the
 FSDP constructor, which may help reduce peak memory usage in some
 cases.

class torch.distributed.fsdp.ShardingStrategy(value)
This specifies the sharding strategy to be used for distributed
   training by "FullyShardedDataParallel".
training by "FullyShardedDataParallel".


"FULL_SHARD": Parameters, gradients, and optimizer states are
     sharded. For the parameters, this strategy unshards (via all-
     gather) before the forward, reshards after the forward, unshards
     before the backward computation, and reshards after the backward
     computation. For gradients, it synchronizes and shards them (via
     reduce-scatter) after the backward computation. The sharded
     optimizer states are updated locally per rank.


"SHARD_GRAD_OP": Gradients and optimizer states are sharded
     during computation, and additionally, parameters are sharded
     outside computation. For the parameters, this strategy unshards
     before the forward, does not reshard them after the forward, and
     only reshards them after the backward computation. The sharded
     optimizer states are updated locally per rank. Inside
     "no_sync()", the parameters are not resharded after the backward
     computation.
computation.


"NO_SHARD": Parameters, gradients, and optimizer states are not
     sharded but instead replicated across ranks similar to PyTorch's
     "DistributedDataParallel" API. For gradients, this strategy
     synchronizes them (via all-reduce) after the backward
     computation. The unsharded optimizer states are updated locally
     per rank.


"HYBRID_SHARD": Apply "FULL_SHARD" within a node, and replicate
     parameters across
        nodes. This results in reduced communication volume as
        expensive all-gathers and reduce-scatters are only done within
        a node, which can be more performant for medium -sized models.


"_HYBRID_SHARD_ZERO2": Apply "SHARD_GRAD_OP" within a node, and
     replicate parameters across
        nodes. This is like "HYBRID_SHARD", except this may provide
        even higher throughput since the unsharded parameters are not
        freed after the forward pass, saving the all-gathers in the
        pre-backward.
pre-backward.
class torch.distributed.fsdp.MixedPrecision(param_dtype=None, reduce_dtype=None, buffer_dtype=None, keep_low_precision_grads=False, cast_forward_inputs=False, cast_root_forward_inputs=True)
This configures FSDP-native mixed precision training.
Variables:
      * param_dtype (torch.dtype) -- This specifies the dtype
        for model parameters, inputs (when "cast_forward_inputs" or
        "cast_root_forward_inputsis set toTrue"), and therefore
        the dtype for computation. However, outside the forward and
        backward passes, parameters are in full precision. Model
        checkpointing always happens in full precision.
  * **reduce_dtype** (*torch.dtype*) -- This specifies the dtype
    for gradient reduction, which is permitted to differ from
    "param_dtype".

  * **buffer_dtype** (*torch.dtype*) -- This specifies the dtype
    for buffers. FSDP does not shard buffers, casts them to
"buffer_dtype" in the first forward pass, and keeps them in
        that dtype thereafter. Model checkpointing always happens in
        full precision.
  * **keep_low_precision_grads** (*bool*) -- This specifies
    whether to upcast gradients back to the full parameter
    precision after the backward pass. This may be set to "False"
    to save memory if using custom optimizers that can perform the
    optimizer step in "reduce_dtype". (Default: "False")

  * **cast_forward_inputs** (*bool*) -- Cast floating point
    tensors in the forward arguments and keyword arguments to
    "param_dtype". (Default: "False")

  * **cast_root_forward_inputs** (*bool*) -- Cast floating point
    tensors in the forward arguments and keyword arguments to
    "param_dtype" for the root FSDP instance. It takes precedence
    over "cast_forward_inputs" for the root FSDP instance.
    (Default: "True")

Note:
 This API is experimental and subject to change.
Note:
 Only floating point tensors are cast to their specified dtypes.

Note:
 In "summon_full_params", parameters are forced to full precision,
 but buffers are not.

Note:
 "state_dict" checkpoints parameters and buffers in full
 precision. For buffers, this is only supported for
 "StateDictType.FULL_STATE_DICT".

Note:
 Each low precision dtype must be specified explicitly. For
 example, "MixedPrecision(reduce_dtype=torch.float16)" only
 specifies the reduction dtype to be low precision, and FSDP will
 not cast parameters or buffers.

Note:
 If a "reduce_dtype" is not specified, then gradient reduction
 happens in "param_dtype" if specified or the original parameter
 dtype otherwise.

Note:
 If the user passes a model with "BatchNorm" modules and an
 "auto_wrap_policy" to the FSDP constructor, then FSDP will
 disable mixed precision for "BatchNorm" modules by wrapping them
separately in their own FSDP instance with mixed precision
     disabled. This is due to some missing low precision "BatchNorm"
     kernels. If the user does not use an "auto_wrap_policy", then the
     user must take care to not use mixed precision for FSDP instances
     containing "BatchNorm" modules.
Note:
 "MixedPrecision" has "cast_root_forward_inputs=True" and
 "cast_forward_inputs=False" by default. For the root FSDP
 instance, its "cast_root_forward_inputs" takes precedence over
 its "cast_forward_inputs". For non-root FSDP instances, their
 "cast_root_forward_inputs" values are ignored. The default
 setting is sufficient for the typical case where each FSDP
 instance has the same "MixedPrecision" configuration and only
 needs to cast inputs to the "param_dtype" at the beginning of the
 model's forward pass.

Note:
 For nested FSDP instances with different "MixedPrecision"
 configurations, we recommend setting individual
"cast_forward_inputs" values to configure casting inputs or not
     before each instance's forward. In such a case, since the casts
     happen before each FSDP instance's forward, a parent FSDP
     instance should have its non-FSDP submodules run before its FSDP
     submodules to avoid the activation dtype being changed due to a
     different "MixedPrecision" configuration.Example:
    >>> model = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3))
    >>> model[1] = FSDP(
    >>>     model[1],
    >>>     mixed_precision=MixedPrecision(param_dtype=torch.float16, cast_forward_inputs=True),
    >>> )
    >>> model = FSDP(
    >>>     model,
    >>>     mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, cast_forward_inputs=True),
    >>> )

 The above shows a working example. On the other hand, if
 "model[1]" were replaced with "model[0]", meaning that the
 submodule using different "MixedPrecision" ran its forward first,
then "model[1]" would incorrectly see "float16" activations
     instead of "bfloat16" ones.
class torch.distributed.fsdp.CPUOffload(offload_params=False)
This configures CPU offloading.
Variables:
      offload_params (bool) -- This specifies whether to offload
      parameters to CPU when not involved in computation. If enabled,
      this implicitly offloads gradients to CPU as well. This is to
      support the optimizer step, which requires parameters and
      gradients to be on the same device.
torch.utils.cpp_extension
torch.utils.cpp_extension.CppExtension(name, sources, args, *kwargs)
Creates a "setuptools.Extension" for C++.
Convenience method that creates a "setuptools.Extension" with the
   bare minimum (but often sufficient) arguments to build a C++
   extension.
All arguments are forwarded to the "setuptools.Extension"
   constructor.
-[ Example ]-



from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
   ...     name='extension',
   ...     ext_modules=[
   ...         CppExtension(
   ...             name='extension',
   ...             sources=['extension.cpp'],
   ...             extra_compile_args=['-g']),
   ...     ],
   ...     cmdclass={
   ...         'build_ext': BuildExtension
   ...     })



torch.utils.cpp_extension.CUDAExtension(name, sources, args, *kwargs)
Creates a "setuptools.Extension" for CUDA/C++.
Creates a "setuptools.Extension" for CUDA/C++.
Convenience method that creates a "setuptools.Extension" with the
   bare minimum (but often sufficient) arguments to build a CUDA/C++
   extension. This includes the CUDA include path, library path and
   runtime library.
All arguments are forwarded to the "setuptools.Extension"
   constructor.
-[ Example ]-



from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
   ...     name='cuda_extension',
   ...     ext_modules=[
   ...         CUDAExtension(
   ...                 name='cuda_extension',
   ...                 sources=['extension.cpp', 'extension_kernel.cu'],
   ...                 extra_compile_args={'cxx': ['-g'],
   ...                                     'nvcc': ['-O2']})
   ...     ],
   ...     cmdclass={
   ...         'build_ext': BuildExtension
   ...     })



Compute capabilities:
By default the extension will be compiled to run on all archs of
the cards visible during the building process of the extension,
   plus PTX. If down the road a new card is installed the extension
   may need to be recompiled. If a visible card has a compute
   capability (CC) that's newer than the newest version for which your
   nvcc can build fully-compiled binaries, Pytorch will make nvcc fall
   back to building kernels with the newest version of PTX your nvcc
   does support (see below for details on PTX).
You can override the default behavior using TORCH_CUDA_ARCH_LIST
   to explicitly specify which CCs you want the extension to support:
TORCH_CUDA_ARCH_LIST="6.1 8.6" python build_my_extension.py
   TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" python
   build_my_extension.py
The +PTX option causes extension kernel binaries to include PTX
   instructions for the specified CC. PTX is an intermediate
   representation that allows kernels to runtime-compile for any CC >=
   the specified CC (for example, 8.6+PTX generates PTX that can
runtime-compile for any GPU with CC >= 8.6). This improves your
   binary's forward compatibility. However, relying on older PTX to
   provide forward compat by runtime-compiling for newer CCs can
   modestly reduce performance on those newer CCs. If you know exact
   CC(s) of the GPUs you want to target, you're always better off
   specifying them individually. For example, if you want your
   extension to run on 8.0 and 8.6, "8.0+PTX" would work functionally
   because it includes PTX that can runtime-compile for 8.6, but "8.0
   8.6" would be better.
Note that while it's possible to include all supported archs, the
   more archs get included the slower the building process will be, as
   it will build a separate kernel image for each arch.
Note that CUDA-11.5 nvcc will hit internal compiler error while
   parsing torch/extension.h on Windows. To workaround the issue, move
   python binding logic to pure C++ file.
Example use:
      #include  at::Tensor