shrinath-suresh/blogs-docs-splitted huggingface dataset:
FullyShardedDataParallel
class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, sharding_strategy=None, cpu_offload=None, auto_wrap_policy=None, backward_prefetch=BackwardPrefetch.BACKWARD_PRE, mixed_precision=None, ignored_modules=None, param_init_fn=None, device_id=None, sync_module_states=False, forward_prefetch=False, limit_all_gathers=False, use_orig_params=False, ignored_parameters=None)
A wrapper for sharding Module parameters across data parallel
   workers. This is inspired by Xu et al. as well as the ZeRO Stage 3
   from DeepSpeed. FullyShardedDataParallel is commonly shortened to
   FSDP.
Example:
  >>> import torch
  >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
  >>> torch.cuda.set_device(device_id)
  >>> sharded_module = FSDP(my_module)
  >>> optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
  >>> x = sharded_module(x, y=3, z=torch.Tensor([1]))
  >>> loss = x.sum()
loss = x.sum()
      >>> loss.backward()
      >>> optim.step()



Warning:
 The optimizer must be initialized *after* the module has been
 wrapped, since FSDP will shard parameters in-place and this will
 break any previously initialized optimizers.

Warning:
 If the destination CUDA device has ID "dev_id", either (1)
 "module" should already be placed on that device, (2) the device
 should be set using "torch.cuda.set_device(dev_id)", or (3)
 "dev_id" should be passed into the "device_id" constructor
 argument. This FSDP instance's compute device will be that
 destination device. For (1) and (3), the FSDP initialization
 always occurs on GPU. For (2), the FSDP initialization happens on
 "module" 's current device, which may be CPU.

Warning:
 FSDP currently does not support gradient accumulation outside
 "no_sync()" when using CPU offloading. Trying to do so yields
 incorrect results since FSDP will use the newly-reduced gradient
instead of accumulating with any existing gradient.
Warning:
 Changing the original parameter variable names after construction
 will lead to undefined behavior.

Warning:
 Passing in *sync_module_states=True* flag requires module to be
 put on GPU, or to use "device_id" argument to specify a CUDA
 device that FSDP will move module to. This is because
 "sync_module_states=True" requires GPU communication.

Warning:
 As of PyTorch 1.12, FSDP only offers limited support for shared
 parameters (for example, setting one "Linear" layer's weight to
 another's). In particular, modules that share parameters must be
 wrapped as part of the same FSDP unit. If enhanced shared
 parameter support is needed for your use case, please ping
 https://github.com/pytorch/pytorch/issues/77724

Note:
 Inputs into FSDP "forward" function will be moved to compute
 device (same device FSDP module is on) before running "forward",
so user does not have to manually move inputs from CPU -> GPU.
Parameters:
      * module (nn.Module) -- This is the module to be wrapped
        with FSDP.
  * **process_group** (*Optional**[**Union**[**ProcessGroup**,
    **Tuple**[**ProcessGroup**, **ProcessGroup**]**]**]*) --
    Optional[Union[ProcessGroup, Tuple[ProcessGroup,
    ProcessGroup]]] This is the process group used for collective
    communications and the one over which the model is sharded.
    For hybrid sharding strategies such as
    "ShardingStrategy.HYBRID_SHARD" users can pass in a tuple of
    process groups representing the groups to shard and replicate
    across, respectively.

  * **sharding_strategy** (*Optional**[**ShardingStrategy**]*) --
    This configures the sharding strategy used by FSDP, which may
    trade off memory saving and communication overhead. See
    "ShardingStrategy" for details. (Default: "FULL_SHARD")
cpu_offload (Optional[CPUOffload]) -- This
        configures CPU offloading. If this is set to "None", then no
        CPU offloading happens. See "CPUOffload" for details.
        (Default: "None")


auto_wrap_policy
    (Optional[Union[Callable[[nn.Module,
    bool, int], bool], _FSDPPolicy]]) --
This is either "None", an "_FSDPPolicy", or a callable of a
fixed signature. If it is "None", then "module" is wrapped
with only a top-level FSDP instance without any nested
wrapping. If it is an "_FSDPPolicy", then the wrapping follows
the given policy. "ModuleWrapPolicy" in
"torch.distributed.fsdp.wrap.py" is an example. If it is a
callable, then it should take in three arguments "module:
nn.Module", "recurse: bool", and "nonwrapped_numel: int" and
should return a "bool" specifying whether the passed-in
"module" should be wrapped if "recurse=False" or if the
        traversal should continue down the subtree if "recurse=True".
        Additional custom arguments may be added to the callable. The
        "size_based_auto_wrap_policy" in
        "torch.distributed.fsdp.wrap.py" gives an example callable
        that wraps a module if the parameters in its subtree exceed
        100M numel. A good practice is to print the model after
        wrapping and adjust as needed.
    Example:

       >>> def custom_auto_wrap_policy(
       >>>     module: nn.Module,
       >>>     recurse: bool,
       >>>     nonwrapped_numel: int,
       >>>     # Additional custom arguments
       >>>     min_num_params: int = int(1e8),
       >>> ) -> bool:
       >>>     return nonwrapped_numel >= min_num_params
       >>> # Configure a custom `min_num_params`
       >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
backward_prefetch (Optional[BackwardPrefetch]) --
        This configures explicit backward prefetching of all-gathers.
        See "BackwardPrefetch" for details. (Default: "BACKWARD_PRE")


mixed_precision (Optional[MixedPrecision]) -- This
    configures native mixed precision for FSDP. If this is set to
    "None", then no mixed precision is used. Otherwise, parameter,
    buffer, and gradient reduction dtypes can be set. See
    "MixedPrecision" for details. (Default: "None")


ignored_modules
    (Optional[Iterable[torch.nn.Module]]) -- Modules
    whose own parameters and child modules' parameters and buffers
    are ignored by this instance. None of the modules directly in
    "ignored_modules" should be "FullyShardedDataParallel"
    instances, and any child modules that are already-constructed
    "FullyShardedDataParallel" instances will not be ignored if
they are nested under this instance. This argument may be used
        to avoid sharding specific parameters at module granularity
        when using an "auto_wrap_policy" or if parameters' sharding is
        not managed by FSDP. (Default: "None")
  * **param_init_fn**
    (*Optional**[**Callable**[**[**nn.Module**]**, **None**]**]*)
    --

    A "Callable[torch.nn.Module] -> None" that specifies how
    modules that are currently on the meta device should be
    initialized onto an actual device. Note that as of v1.12, we
    detect modules on the meta device via "is_meta" check and
    apply a default initialization that calls "reset_parameters"
    method on the passed in "nn.Module" if "param_init_fn" is not
    specified, otherwise we run "param_init_fn" to initialize the
    passed in "nn.Module". In particular, this means that if
    "is_meta=True" for any module parameters for modules that will
be wrapped with FSDP and "param_init_fn" is not specified, we
        assume your module properly implements a "reset_parameters()"
        and will throw errors if not. Note that additionally, we offer
        support for modules initialized with torchdistX's
        (https://github.com/pytorch/torchdistX) "deferred_init" API.
        In this case, deferred modules would be initialized by a
        default initialization function that calls torchdistX's
        "materialize_module", or the passed in "param_init_fn", if it
        is not "None". The same "Callable" is applied to initialize
        all meta modules. Note that this initialization function is
        applied before doing any FSDP sharding logic.
    Example:

       >>> module = MyModule(device="meta")
       >>> def my_init_fn(module):
       >>>     # responsible for initializing a module, such as with reset_parameters
       >>>     ...
...
       >>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
       >>> print(next(fsdp_model.parameters()).device) # current CUDA device
       >>> # With torchdistX
       >>> module = deferred_init.deferred_init(MyModule, device="cuda")
       >>> # Will initialize via deferred_init.materialize_module().
       >>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)




  * **device_id** (*Optional**[**Union**[**int**,
    **torch.device**]**]*) -- An "int" or "torch.device"
    describing the CUDA device the FSDP module should be moved to
    determining where initialization such as sharding takes place.
    If this argument is not specified and "module" is on CPU, we
    issue a warning mentioning that this argument can be specified
    for faster initialization. If specified, resulting FSDP
    instances will reside on this device, including moving ignored
modules' parameters if needed. Note that if "device_id" is
        specified but "module" is already on a different CUDA device,
        an error will be thrown. (Default: "None")
  * **sync_module_states** (*bool*) -- If "True", each
    individually wrapped FSDP unit will broadcast module
    parameters from rank 0 to ensure they are the same across all
    ranks after initialization. This helps ensure model parameters
    are the same across ranks before starting training, but adds
    communication overhead to "__init__", as at least one
    broadcast is triggered per individually wrapped FSDP unit.
    This can also help load checkpoints taken by "state_dict" and
    to be loaded by "load_state_dict" in a memory efficient way.
    See documentation for "FullStateDictConfig" for an example of
    this. (Default: "False")

  * **forward_prefetch** (*bool*) -- If "True", then FSDP
    *explicitly* prefetches the next upcoming all-gather while
executing in the forward pass. This may improve communication
        and computation overlap for CPU bound workloads. This should
        only be used for static graph models since the forward order
        is fixed based on the first iteration's execution. (Default:
        "False")
  * **limit_all_gathers** (*bool*) -- If "False", then FSDP allows
    the CPU thread to schedule all-gathers without any extra
    synchronization. If "True", then FSDP explicitly synchronizes
    the CPU thread to prevent too many in-flight all-gathers. This
    "bool" only affects the sharded strategies that schedule all-
    gathers. Enabling this can help lower the number of CUDA
    malloc retries.

  * **ignored_parameters**
    (*Optional**[**Iterable**[**torch.nn.Parameter**]**]*) --
    Ignored parameters will not be managed by this FSDP instance,
    that means these parameters will not be flattened and sharded
by FSDP, their gradients will not be synchronized as well.
        With this newly added argument, "ignored_modules" could be
        deprecated soon. For backward compatibility, both
        "ignored_parameters" and "ignored_modules" are kept for now,
        but FSDP only allows one of them to be specified as not
        "None".
apply(fn)
  Applies "fn" recursively to every submodule (as returned by
  ".children()") as well as self. Typical use includes
  initializing the parameters of a model (see also torch.nn.init).

  Compared to "torch.nn.Module.apply", this version additionally
  gathers the full parameters before applying "fn". It should not
  be called from within another "summon_full_params" context.

  Parameters:
     **fn** ("Module" -> None) -- function to be applied to each
     submodule

  Returns:
     self

  Return type:
     Module

clip_grad_norm_(max_norm, norm_type=2.0)
clip_grad_norm_(max_norm, norm_type=2.0)
  Clips the gradient norm of all parameters. The norm is computed
  over all parameters' gradients as viewed as a single vector, and
  the gradients are modified in-place.

  Parameters:
     * **max_norm** (*float** or **int*) -- max norm of the
       gradients

     * **norm_type** (*float** or **int*) -- type of the used
       p-norm. Can be "'inf'" for infinity norm.

  Returns:
     Total norm of the parameters (viewed as a single vector).

  Return type:
     *Tensor*

  Note:

    If every FSDP instance uses "NO_SHARD", meaning that no
    gradients are sharded across ranks, then you may directly use
    "torch.nn.utils.clip_grad_norm_()".

  Note:

    If at least some FSDP instance uses a sharded strategy (i.e.
    one other than "NO_SHARD"), then you should use this method
    instead of "torch.nn.utils.clip_grad_norm_()" since this
method handles the fact that gradients are sharded across
        ranks.
  Note:

    The total norm returned will have the "largest" dtype across
    all parameters/gradients as defined by PyTorch's type
    promotion semantics. For example, if *all*
    parameters/gradients use a low precision dtype, then the
    returned norm's dtype will be that low precision dtype, but if
    there exists at least one parameter/ gradient using FP32, then
    the returned norm's dtype will be FP32.

  Warning:

    This needs to be called on all ranks since it uses collective
    communications.

static flatten_sharded_optim_state_dict(sharded_optim_state_dict, model, optim)
  The API is similar to "shard_full_optim_state_dict()". The only
  difference is that the input "sharded_optim_state_dict" should
  be returned from "sharded_optim_state_dict()". Therefore, there
  will be all-gather calls on each rank to gather "ShardedTensor"
  s.
s.
  Parameters:
     * **sharded_optim_state_dict** (*Dict**[**str**, **Any**]*)
       -- Optimizer state dict corresponding to the unflattened
       parameters and holding the sharded optimizer state.

     * **model** (*torch.nn.Module*) -- Refer to
       :meth:"shard_full_optim_state_dict".

     * **optim** (*torch.optim.Optimizer*) -- Optimizer for
       "model" 's

     * **parameters.** --

  Returns:
     Refer to "shard_full_optim_state_dict()".

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

forward(args, *kwargs)
  Runs the forward pass for the wrapped module, inserting FSDP-
  specific pre- and post-forward sharding logic.

  Return type:
     *Any*

static fsdp_modules(module, root_only=False)
  Returns all nested FSDP instances, possibly including "module"
  itself and only including FSDP root modules if "root_only=True".

  Parameters:
Parameters:
         * module (torch.nn.Module) -- Root module, which may or
           may not be an "FSDP" module.
     * **root_only** (*bool*) -- Whether to return only FSDP root
       modules. (Default: "False")

  Returns:
     FSDP modules that are nested in the input "module".

  Return type:
     List[FullyShardedDataParallel]

static full_optim_state_dict(model, optim, optim_input=None, rank0_only=True, group=None)
  Consolidates the full optimizer state on rank 0 and returns it
  as a "dict" following the convention of
  "torch.optim.Optimizer.state_dict()", i.e. with keys ""state""
  and ""param_groups"". The flattened parameters in "FSDP" modules
  contained in "model" are mapped back to their unflattened
  parameters.

  Warning:

    This needs to be called on all ranks since it uses collective
    communications. However, if "rank0_only=True", then the state
dict is only populated on rank 0, and all other ranks return
        an empty "dict".
  Warning:

    Unlike "torch.optim.Optimizer.state_dict()", this method uses
    full parameter names as keys instead of parameter IDs.

  Note:

    Like in "torch.optim.Optimizer.state_dict()", the tensors
    contained in the optimizer state dict are not cloned, so there
    may be aliasing surprises. For best practices, consider saving
    the returned optimizer state dict immediately, e.g. using
    "torch.save()".

  Parameters:
     * **model** (*torch.nn.Module*) -- Root module (which may or
       may not be a "FullyShardedDataParallel" instance) whose
       parameters were passed into the optimizer "optim".

     * **optim** (*torch.optim.Optimizer*) -- Optimizer for
       "model" 's parameters.

     * **optim_input**
       (*Optional**[**Union**[**List**[**Dict**[**str**,
Any]], Iterable[torch.nn.Parameter]]]*)
           -- Input passed into the optimizer "optim" 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")
     * **rank0_only** (*bool*) -- If "True", saves the populated
       "dict" only on rank 0; if "False", saves it on all ranks.
       (Default: "True")

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

  Returns:
     A "dict" containing the optimizer state for "model" 's
     original unflattened parameters and including keys "state"
     and "param_groups" following the convention of
     "torch.optim.Optimizer.state_dict()". If "rank0_only=True",
then nonzero ranks return an empty "dict".
  Return type:
     Dict[str, Any]

property module: Module
  Returns the wrapped module (like "DistributedDataParallel").

named_buffers(args, *kwargs)
  Overrides "named_buffers()" to intercept buffer names and remove
  all occurrences of the FSDP-specific flattened buffer prefix
  when inside the "summon_full_params()" context manager.

  Return type:
     *Iterator*[*Tuple*[str, *Tensor*]]

named_parameters(args, *kwargs)
  Overrides "named_parameters()" to intercept parameter names and
  remove all occurrences of the FSDP-specific flattened parameter
  prefix when inside the "summon_full_params()" context manager.

  Return type:
     *Iterator*[*Tuple*[str, *Parameter*]]

no_sync()
  A context manager to disable gradient synchronizations across
  FSDP instances. Within this context, gradients will be
  accumulated in module variables, which will later be
synchronized in the first forward-backward pass after exiting
      the context. This should only be used on the root FSDP instance
      and will recursively apply to all children FSDP instances.
  Note:

    This likely results in higher memory usage because FSDP will
    accumulate the full model gradients (instead of gradient
    shards) until the eventual sync.

  Note:

    When used with CPU offloading, the gradients will not be
    offloaded to CPU when inside the context manager. Instead,
    they will only be offloaded right after the eventual sync.

  Return type:
     *Generator*

register_comm_hook(state, hook)
  Registers a communication hook which is an enhancement that
  provides a flexible hook to users where they can specify how
  FSDP aggregates gradients across multiple workers. This hook can
  be used to implement several algorithms like GossipGrad and
  gradient compression which involve different communication
strategies for parameter syncs while training with
      "FullyShardedDataParallel".
  Warning:

    FSDP communication hook should be registered before running an
    initial forward pass and only once.

  Parameters:
     * **state** (*object*) --

       Passed to the hook to maintain any state information during
       the training process. Examples include error feedback in
       gradient compression, peers to communicate with next in
       GossipGrad, etc. It is locally stored by each worker and
       shared by all the gradient tensors on the worker.

     * **hook** (*Callable*) -- Callable, which has one of the
       following signatures: 1) "hook: Callable[torch.Tensor] ->
       None": This function takes in a Python tensor, which
       represents the full, flattened, unsharded gradient with
       respect to all variables corresponding to the model this
       FSDP unit is wrapping (that are not wrapped by other FSDP
sub-units). It then performs all necessary processing and
           returns "None"; 2) "hook: Callable[torch.Tensor,
           torch.Tensor] -> None": This function takes in two Python
           tensors, the first one represents the full, flattened,
           unsharded gradient with respect to all variables
           corresponding to the model this FSDP unit is wrapping (that
           are not wrapped by other FSDP sub-units). The latter
           represents a pre-sized tensor to store a chunk of a sharded
           gradient after reduction. In both cases, callable performs
           all necessary processing and returns "None". Callables with
           signature 1 are expected to handle gradient communication
           for a NO_SHARD case. Callables with signature 2 are
           expected to handle gradient communication for sharded
           cases.
static rekey_optim_state_dict(optim_state_dict, optim_state_key_type, model, optim_input=None, optim=None)
Re-keys the optimizer state dict "optim_state_dict" to use the
      key type "optim_state_key_type". This can be used to achieve
      compatibility between optimizer state dicts from models with
      FSDP instances and ones without.
  To re-key an FSDP full optimizer state dict (i.e. from
  "full_optim_state_dict()") to use parameter IDs and be loadable
  to a non-wrapped model:

     >>> wrapped_model, wrapped_optim = ...
     >>> full_osd = FSDP.full_optim_state_dict(wrapped_model, wrapped_optim)
     >>> nonwrapped_model, nonwrapped_optim = ...
     >>> rekeyed_osd = FSDP.rekey_optim_state_dict(full_osd, OptimStateKeyType.PARAM_ID, nonwrapped_model)
     >>> nonwrapped_optim.load_state_dict(rekeyed_osd)

  To re-key a normal optimizer state dict from a non-wrapped model
  to be loadable to a wrapped model:

     >>> nonwrapped_model, nonwrapped_optim = ...
     >>> osd = nonwrapped_optim.state_dict()
osd = nonwrapped_optim.state_dict()
         >>> rekeyed_osd = FSDP.rekey_optim_state_dict(osd, OptimStateKeyType.PARAM_NAME, nonwrapped_model)
         >>> wrapped_model, wrapped_optim = ...
         >>> sharded_osd = FSDP.shard_full_optim_state_dict(rekeyed_osd, wrapped_model)
         >>> wrapped_optim.load_state_dict(sharded_osd)



  Returns:
     The optimizer state dict re-keyed using the parameter keys
     specified by "optim_state_key_type".

  Return type:
     Dict[str, Any]

static scatter_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None, group=None)
  Scatters the full optimizer state dict from rank 0 to all other
  ranks, returning the sharded optimizer state dict on each rank.
  The return value is the same as "shard_full_optim_state_dict()",
  and on rank 0, the first argument should be the return value of
  "full_optim_state_dict()".

  Example:
"full_optim_state_dict()".
  Example:

     >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
     >>> model, optim = ...
     >>> full_osd = FSDP.full_optim_state_dict(model, optim)  # only non-empty on rank 0
     >>> # Define new model with possibly different world size
     >>> new_model, new_optim, new_group = ...
     >>> sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, new_model, group=new_group)
     >>> 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