2.3/fsdp.html



<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta name="robots" content="noindex">
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>FullyShardedDataParallel &mdash; PyTorch 2.3 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/fsdp.html"/>
  

  <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="_static/copybutton.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.13.11/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="_static/katex-math.css" type="text/css" />
  <link rel="stylesheet" href="_static/sphinx-dropdown.css" type="text/css" />
  <link rel="stylesheet" href="_static/panels-bootstrap.min.css" type="text/css" />
  <link rel="stylesheet" href="_static/css/jit.css" type="text/css" />
    <link rel="index" title="Index" href="genindex.html" />
    <link rel="search" title="Search" href="search.html" />
    <link rel="next" title="Distributed Optimizers" href="distributed.optim.html" />
    <link rel="prev" title="TorchElastic Kubernetes" href="elastic/kubernetes.html" />


  <!-- Google Tag Manager -->
    <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start':
    new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0],
    j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src=
    'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f);
    })(window,document,'script','dataLayer','GTM-T8XT4PS');</script>
    <!-- End Google Tag Manager -->
  

  <script src="_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
  <link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.15.2/css/all.css" integrity="sha384-vSIIfh2YWi9wW0r9iZe7RJPrKwp6bG+s9QZMoITbCckVJqGCCRhc+ccxNcdpHuYu" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>

          <li class="main-menu-item">
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Learn
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/get-started">
                  <span class=dropdown-title>Get Started</span>
                  <p>Run PyTorch locally or get started quickly with one of the supported cloud platforms</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials">
                  <span class="dropdown-title">Tutorials</span>
                  <p>Whats new in PyTorch tutorials</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/beginner/basics/intro.html">
                  <span class="dropdown-title">Learn the Basics</span>
                  <p>Familiarize yourself with PyTorch concepts and modules</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/recipes/recipes_index.html">
                  <span class="dropdown-title">PyTorch Recipes</span>
                  <p>Bite-size, ready-to-deploy PyTorch code examples</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/beginner/introyt.html">
                  <span class="dropdown-title">Intro to PyTorch - YouTube Series</span>
                  <p>Master PyTorch basics with our engaging YouTube tutorial series</p>
                </a>
              </div>
            </div>
          </li>

          <li>
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Ecosystem
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem">
                  <span class="dropdown-title">Tools</span>
                  <p>Learn about the tools and frameworks in the PyTorch Ecosystem</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/#community-module">
                  <span class=dropdown-title>Community</span>
                  <p>Join the PyTorch developer community to contribute, learn, and get your questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://discuss.pytorch.org/" target="_blank">
                  <span class=dropdown-title>Forums</span>
                  <p>A place to discuss PyTorch code, issues, install, research</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/resources">
                  <span class=dropdown-title>Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem/contributor-awards-2023">
                  <span class="dropdown-title">Contributor Awards - 2023</span>
                  <p>Award winners announced at this year's PyTorch Conference</p>
                </a>
              </div>
            </div>
          </li>

          <li>
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Edge
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/edge">
                  <span class="dropdown-title">About PyTorch Edge</span>
                  <p>Build innovative and privacy-aware AI experiences for edge devices</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/executorch-overview">
                  <span class="dropdown-title">ExecuTorch</span>
                  <p>End-to-end solution for enabling on-device inference capabilities across mobile and edge devices</p>
                </a>
              </div>
            </div>  
          </li>

          <li class="main-menu-item">
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Docs
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/docs/stable/index.html">
                  <span class="dropdown-title">PyTorch</span>
                  <p>Explore the documentation for comprehensive guidance on how to use PyTorch</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/pytorch-domains">
                  <span class="dropdown-title">PyTorch Domains</span>
                  <p>Read the PyTorch Domains documentation to learn more about domain-specific libraries</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Blogs & News 
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/blog/">
                  <span class="dropdown-title">PyTorch Blog</span>
                  <p>Catch up on the latest technical news and happenings</p>
                </a>
                 <a class="nav-dropdown-item" href="https://pytorch.org/community-blog">
                  <span class="dropdown-title">Community Blog</span>
                  <p>Stories from the PyTorch ecosystem</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/videos">
                  <span class="dropdown-title">Videos</span>
                  <p>Learn about the latest PyTorch tutorials, new, and more </p>
                <a class="nav-dropdown-item" href="https://pytorch.org/community-stories">
                  <span class="dropdown-title">Community Stories</span>
                  <p>Learn how our community solves real, everyday machine learning problems with PyTorch</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/events">
                  <span class="dropdown-title">Events</span>
                  <p>Find events, webinars, and podcasts</p>
                </a>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                About
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/foundation">
                  <span class="dropdown-title">PyTorch Foundation</span>
                  <p>Learn more about the PyTorch Foundation</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/governing-board">
                  <span class="dropdown-title">Governing Board</span>
                  <p></p>
                </a>
              </div>
            </div>
          </li>

          <li class="main-menu-item">
            <div class="no-dropdown">
              <a href="https://pytorch.org/join" data-cta="join">
                Become a Member
              </a>
            </div>
          </li>
          <li>
           <div class="main-menu-item">
             <a href="https://github.com/pytorch/pytorch" class="github-icon">
             </a>
           </div>
          </li>
          <!--- TODO: This block adds the search icon to the nav bar. We will enable it later. 
          <li>
            <div class="main-menu-item">
             <a href="https://github.com/pytorch/pytorch" class="search-icon">
             </a>
            </div>
          </li>
          --->
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>
  </div>
</div>

<body class="pytorch-body">

   
    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            
    <div class="version">
      <a href='https://pytorch.org/docs/versions.html'>2.3 &#x25BC</a>
    </div>
    

<div role="search">
  <form id="rtd-search-form" class="wy-form" action="search.html" method="get">
    <input type="text" name="q" placeholder="Search Docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          </div>

          
              <p class="caption" role="heading"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="community/build_ci_governance.html">PyTorch Governance | Build + CI</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/design.html">PyTorch Design Philosophy</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/governance.html">PyTorch Governance | Mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/persons_of_interest.html">PyTorch Governance | Maintainers</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Developer Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="notes/amp_examples.html">CUDA Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/extending.func.html">Extending torch.func with autograd.Function</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/fsdp.html">FSDP Notes</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/gradcheck.html">Gradcheck mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/hip.html">HIP (ROCm) semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/modules.html">Modules</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/mps.html">MPS backend</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/numerical_accuracy.html">Numerical accuracy</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="cpp_index.html">C++</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/javadoc/">Javadoc</a></li>
<li class="toctree-l1"><a class="reference internal" href="deploy.html">torch::deploy</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Python API</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="amp.html">torch.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="library.html">torch.library</a></li>
<li class="toctree-l1"><a class="reference internal" href="cpu.html">torch.cpu</a></li>
<li class="toctree-l1"><a class="reference internal" href="cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch_cuda_memory.html">Understanding CUDA Memory Usage</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch_cuda_memory.html#generating-a-snapshot">Generating a Snapshot</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch_cuda_memory.html#using-the-visualizer">Using the visualizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch_cuda_memory.html#snapshot-api-reference">Snapshot API Reference</a></li>
<li class="toctree-l1"><a class="reference internal" href="mps.html">torch.mps</a></li>
<li class="toctree-l1"><a class="reference internal" href="xpu.html">torch.xpu</a></li>
<li class="toctree-l1"><a class="reference internal" href="meta.html">Meta device</a></li>
<li class="toctree-l1"><a class="reference internal" href="backends.html">torch.backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="export.html">torch.export</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.algorithms.join.html">torch.distributed.algorithms.join</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.elastic.html">torch.distributed.elastic</a></li>
<li class="toctree-l1 current"><a class="current reference internal" href="#">torch.distributed.fsdp</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.optim.html">torch.distributed.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.tensor.parallel.html">torch.distributed.tensor.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.checkpoint.html">torch.distributed.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler.html">torch.compiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="fft.html">torch.fft</a></li>
<li class="toctree-l1"><a class="reference internal" href="func.html">torch.func</a></li>
<li class="toctree-l1"><a class="reference internal" href="futures.html">torch.futures</a></li>
<li class="toctree-l1"><a class="reference internal" href="fx.html">torch.fx</a></li>
<li class="toctree-l1"><a class="reference internal" href="fx.experimental.html">torch.fx.experimental</a></li>
<li class="toctree-l1"><a class="reference internal" href="hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="linalg.html">torch.linalg</a></li>
<li class="toctree-l1"><a class="reference internal" href="monitor.html">torch.monitor</a></li>
<li class="toctree-l1"><a class="reference internal" href="signal.html">torch.signal</a></li>
<li class="toctree-l1"><a class="reference internal" href="special.html">torch.special</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.overrides.html">torch.overrides</a></li>
<li class="toctree-l1"><a class="reference internal" href="package.html">torch.package</a></li>
<li class="toctree-l1"><a class="reference internal" href="profiler.html">torch.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.attention.html">torch.nn.attention</a></li>
<li class="toctree-l1"><a class="reference internal" href="onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="complex_numbers.html">Complex Numbers</a></li>
<li class="toctree-l1"><a class="reference internal" href="ddp_comm_hooks.html">DDP Communication Hooks</a></li>
<li class="toctree-l1"><a class="reference internal" href="pipeline.html">Pipeline Parallelism</a></li>
<li class="toctree-l1"><a class="reference internal" href="quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="rpc.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="masked.html">torch.masked</a></li>
<li class="toctree-l1"><a class="reference internal" href="nested.html">torch.nested</a></li>
<li class="toctree-l1"><a class="reference internal" href="sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="testing.html">torch.testing</a></li>
<li class="toctree-l1"><a class="reference internal" href="utils.html">torch.utils</a></li>
<li class="toctree-l1"><a class="reference internal" href="benchmark_utils.html">torch.utils.benchmark</a></li>
<li class="toctree-l1"><a class="reference internal" href="bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="deterministic.html">torch.utils.deterministic</a></li>
<li class="toctree-l1"><a class="reference internal" href="jit_utils.html">torch.utils.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="mobile_optimizer.html">torch.utils.mobile_optimizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="config_mod.html">torch.__config__</a></li>
<li class="toctree-l1"><a class="reference internal" href="future_mod.html">torch.__future__</a></li>
<li class="toctree-l1"><a class="reference internal" href="logging.html">torch._logging</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch_environment_variables.html">Torch Environment Variables</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio/stable">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/data">TorchData</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/torchrec">TorchRec</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text/stable">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/vision/stable">torchvision</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>

            
        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          

<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
      <li>FullyShardedDataParallel</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
            
            <a href="_sources/fsdp.rst.txt" rel="nofollow"><img src="_static/images/view-page-source-icon.svg"></a>
          
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          <!-- Google Tag Manager (noscript) -->
          <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-T8XT4PS"
          height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>
          <!-- End Google Tag Manager (noscript) -->
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <div class="section" id="module-torch.distributed.fsdp">
<span id="fullyshardeddataparallel"></span><h1>FullyShardedDataParallel<a class="headerlink" href="#module-torch.distributed.fsdp" title="Permalink to this heading">¶</a></h1>
<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">FullyShardedDataParallel</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">process_group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">sharding_strategy</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cpu_offload</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">auto_wrap_policy</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backward_prefetch</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">BackwardPrefetch.BACKWARD_PRE</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mixed_precision</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">ignored_modules</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">param_init_fn</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">device_id</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">sync_module_states</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_prefetch</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">limit_all_gathers</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">use_orig_params</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">ignored_states</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">device_mesh</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="Permalink to this definition">¶</a></dt>
<dd><p>A wrapper for sharding module parameters across data parallel workers.</p>
<p>This is inspired by <a class="reference external" href="https://arxiv.org/abs/2004.13336">Xu et al.</a> as well as the ZeRO Stage 3 from <a class="reference external" href="https://www.deepspeed.ai/">DeepSpeed</a>.
FullyShardedDataParallel is commonly shortened to FSDP.</p>
<p>For advanced notes please refer to <a class="reference internal" href="notes/fsdp.html#fsdp-notes"><span class="std std-ref">FSDP Notes</span></a>.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">torch</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">set_device</span><span class="p">(</span><span class="n">device_id</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">sharded_module</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span><span class="n">my_module</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">optim</span><span class="o">.</span><span class="n">Adam</span><span class="p">(</span><span class="n">sharded_module</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">lr</span><span class="o">=</span><span class="mf">0.0001</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">x</span> <span class="o">=</span> <span class="n">sharded_module</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span> <span class="n">z</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">([</span><span class="mi">1</span><span class="p">]))</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">loss</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
</pre></div>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>The optimizer must be initialized <em>after</em> the module has been wrapped
with FSDP since FSDP will shard and transform the module’s parameters
in a way that may not preserve the original parameter variables. Thus,
the previously initialized optimizer may have stale references to the
parameters.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>If the destination CUDA device has ID <code class="docutils literal notranslate"><span class="pre">dev_id</span></code>, either (1)
<code class="docutils literal notranslate"><span class="pre">module</span></code> should already be placed on that device, (2) the device
should be set using <code class="docutils literal notranslate"><span class="pre">torch.cuda.set_device(dev_id)</span></code>, or (3)
<code class="docutils literal notranslate"><span class="pre">dev_id</span></code> should be passed into the <code class="docutils literal notranslate"><span class="pre">device_id</span></code> 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 <code class="docutils literal notranslate"><span class="pre">module</span></code> ‘s current
device, which may be CPU.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>FSDP currently does not support gradient accumulation outside
<code class="docutils literal notranslate"><span class="pre">no_sync()</span></code> 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.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Changing the original parameter variable names after construction will
lead to undefined behavior.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Passing in the <code class="docutils literal notranslate"><span class="pre">sync_module_states=True</span></code> flag requires <code class="docutils literal notranslate"><span class="pre">module</span></code> to
be on GPU or to use the <code class="docutils literal notranslate"><span class="pre">device_id</span></code> argument to specify a CUDA device
that FSDP will move <code class="docutils literal notranslate"><span class="pre">module</span></code> to in the FSDP constructor. This is
because <code class="docutils literal notranslate"><span class="pre">sync_module_states=True</span></code> requires GPU communication.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>As of PyTorch 1.12, FSDP only offers limited support for shared parameters
(for example, setting one <code class="docutils literal notranslate"><span class="pre">Linear</span></code> 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 <a class="reference external" href="https://github.com/pytorch/pytorch/issues/77724">https://github.com/pytorch/pytorch/issues/77724</a></p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>FSDP has some constraints on freezing parameters (i.e. setting
<code class="docutils literal notranslate"><span class="pre">param.requires_grad=False</span></code>). For <code class="docutils literal notranslate"><span class="pre">use_orig_params=False</span></code>, each
FSDP instance must manage parameters that are all frozen or all
non-frozen. For <code class="docutils literal notranslate"><span class="pre">use_orig_params=True</span></code>, FSDP supports mixing frozen
and non-frozen, but we recommend not doing so since then the gradient
memory usage will be higher than expected (namely, equivalent to not
freezing those parameters). This means that ideally, frozen parameters
should be isolated into their own <code class="docutils literal notranslate"><span class="pre">nn.Module</span></code> s and wrapped
separately with FSDP.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Attempting to run the forward pass of a submodule that is contained in an
FSDP instance is not supported and will result in errors. This is because the
submodule’s parameters will be sharded, but it itself is not an FSDP instance,
so its forward pass will not all-gather the full parameters appropriately.
This could potentially happen when attempting to run only the encoder of a
encoder-decoder model, and the encoder is not wrapped in its own FSDP instance. To
resolve this, please wrap the submodule in its own FSDP unit.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>FSDP moves input tensors to the <code class="docutils literal notranslate"><span class="pre">forward</span></code> method to the GPU compute
device, so the user does not need to manually move them from CPU.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>The user should not modify the parameters between forward and backward
without using the <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params" title="torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params"><code class="xref py py-meth docutils literal notranslate"><span class="pre">summon_full_params()</span></code></a> context since the
modifications may not persist. Moreover, for <code class="docutils literal notranslate"><span class="pre">use_orig_params=False</span></code>,
accessing the original parameters between forward and backward may
raise an illegal memory access.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>For <code class="docutils literal notranslate"><span class="pre">use_orig_params=True</span></code>, <code class="docutils literal notranslate"><span class="pre">ShardingStrategy.SHARD_GRAD_OP</span></code>
exposes the unsharded parameters, not the sharded parameters, after
forward since it does not free the unsharded ones, unlike
<code class="docutils literal notranslate"><span class="pre">ShardingStrategy.FULL_SHARD</span></code>. One caveat is that, since gradients
are always sharded or <code class="docutils literal notranslate"><span class="pre">None</span></code>, <code class="docutils literal notranslate"><span class="pre">ShardingStrategy.SHARD_GRAD_OP</span></code> will
not expose the sharded gradients with the unsharded parameters after
forward. If you want to inspect the gradients, try
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params" title="torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params"><code class="xref py py-meth docutils literal notranslate"><span class="pre">summon_full_params()</span></code></a> with <code class="docutils literal notranslate"><span class="pre">with_grads=True</span></code>.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>FSDP replaces managed modules’ parameters with <code class="docutils literal notranslate"><span class="pre">torch.Tensor</span></code> views
during forward and backward computation for autograd-related reasons.
If your module’s forward relies on saved references to the parameters
instead of reacquiring the references each iteration, then it will not
see FSDP’s newly created views, and autograd will not work correctly.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>With <code class="docutils literal notranslate"><span class="pre">limit_all_gathers=True</span></code>, you may see a gap in the FSDP
pre-forward where the CPU thread is not issuing any kernels. This is
intentional and shows the rate limiter in effect. Synchronizing the CPU
thread in that way prevents over-allocating memory for subsequent
all-gathers, and it should not actually delay GPU kernel execution.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>When using <code class="docutils literal notranslate"><span class="pre">sharding_strategy=ShardingStrategy.HYBRID_SHARD</span></code> with the
sharding process group being intra-node and the replication process
group being inter-node, setting <code class="docutils literal notranslate"><span class="pre">NCCL_CROSS_NIC=1</span></code> can help improve
the all-reduce times over the replication process group for some
cluster setups.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>FSDP does not work with double backwards due to how it registers
backward hooks.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>module</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>nn.Module</em></a>) – This is the module to be wrapped with FSDP.</p></li>
<li><p><strong>process_group</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><em>ProcessGroup</em><em>, </em><em>Tuple</em><em>[</em><em>ProcessGroup</em><em>, </em><em>ProcessGroup</em><em>]</em><em>]</em><em>]</em>) – This is the process group over which the model is sharded and thus
the one used for FSDP’s all-gather and reduce-scatter collective
communications. If <code class="docutils literal notranslate"><span class="pre">None</span></code>, then FSDP uses the default process
group. For hybrid sharding strategies such as
<code class="docutils literal notranslate"><span class="pre">ShardingStrategy.HYBRID_SHARD</span></code>, users can pass in a tuple of
process groups, representing the groups over which to shard and
replicate, respectively. If <code class="docutils literal notranslate"><span class="pre">None</span></code>, then FSDP constructs process
groups for the user to shard intra-node and replicate inter-node.
(Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>sharding_strategy</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.ShardingStrategy" title="torch.distributed.fsdp.ShardingStrategy"><em>ShardingStrategy</em></a><em>]</em>) – This configures the sharding strategy, which may trade off memory
saving and communication overhead. See <a class="reference internal" href="#torch.distributed.fsdp.ShardingStrategy" title="torch.distributed.fsdp.ShardingStrategy"><code class="xref py py-class docutils literal notranslate"><span class="pre">ShardingStrategy</span></code></a>
for details. (Default: <code class="docutils literal notranslate"><span class="pre">FULL_SHARD</span></code>)</p></li>
<li><p><strong>cpu_offload</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.CPUOffload" title="torch.distributed.fsdp.CPUOffload"><em>CPUOffload</em></a><em>]</em>) – This configures CPU offloading. If this is set to <code class="docutils literal notranslate"><span class="pre">None</span></code>, then
no CPU offloading happens. See <a class="reference internal" href="#torch.distributed.fsdp.CPUOffload" title="torch.distributed.fsdp.CPUOffload"><code class="xref py py-class docutils literal notranslate"><span class="pre">CPUOffload</span></code></a> for details.
(Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>auto_wrap_policy</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><em>Callable</em><em>[</em><em>[</em><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>nn.Module</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em>]</em><em>, </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>]</em><em>, </em><em>ModuleWrapPolicy</em><em>, </em><em>CustomPolicy</em><em>]</em><em>]</em>) – <p>This specifies a policy to apply FSDP to submodules of <code class="docutils literal notranslate"><span class="pre">module</span></code>,
which is needed for communication and computation overlap and thus
affects performance. If <code class="docutils literal notranslate"><span class="pre">None</span></code>, then FSDP only applies to
<code class="docutils literal notranslate"><span class="pre">module</span></code>, and users should manually apply FSDP to parent modules
themselves (proceeding bottom-up). For convenience, this accepts
<code class="docutils literal notranslate"><span class="pre">ModuleWrapPolicy</span></code> directly, which allows users to specify the
module classes to wrap (e.g. the transformer block). Otherwise,
this should be a callable that takes in three arguments
<code class="docutils literal notranslate"><span class="pre">module:</span> <span class="pre">nn.Module</span></code>, <code class="docutils literal notranslate"><span class="pre">recurse:</span> <span class="pre">bool</span></code>, and
<code class="docutils literal notranslate"><span class="pre">nonwrapped_numel:</span> <span class="pre">int</span></code> and should return a <code class="docutils literal notranslate"><span class="pre">bool</span></code> specifying
whether the passed-in <code class="docutils literal notranslate"><span class="pre">module</span></code> should have FSDP applied if
<code class="docutils literal notranslate"><span class="pre">recurse=False</span></code> or if the traversal should continue into the
module’s subtree if <code class="docutils literal notranslate"><span class="pre">recurse=True</span></code>. Users may add additional
arguments to the callable. The <code class="docutils literal notranslate"><span class="pre">size_based_auto_wrap_policy</span></code> in
<code class="docutils literal notranslate"><span class="pre">torch.distributed.fsdp.wrap.py</span></code> gives an example callable that
applies FSDP to a module if the parameters in its subtree exceed
100M numel. We recommend printing the model after applying FSDP
and adjusting as needed.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="k">def</span> <span class="nf">custom_auto_wrap_policy</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">recurse</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">nonwrapped_numel</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="c1"># Additional custom arguments</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">min_num_params</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="mf">1e8</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="k">return</span> <span class="n">nonwrapped_numel</span> <span class="o">&gt;=</span> <span class="n">min_num_params</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Configure a custom `min_num_params`</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">my_auto_wrap_policy</span> <span class="o">=</span> <span class="n">functools</span><span class="o">.</span><span class="n">partial</span><span class="p">(</span><span class="n">custom_auto_wrap_policy</span><span class="p">,</span> <span class="n">min_num_params</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="mf">1e5</span><span class="p">))</span>
</pre></div>
</div>
</p></li>
<li><p><strong>backward_prefetch</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.BackwardPrefetch" title="torch.distributed.fsdp.BackwardPrefetch"><em>BackwardPrefetch</em></a><em>]</em>) – This configures explicit backward prefetching of all-gathers. If
<code class="docutils literal notranslate"><span class="pre">None</span></code>, then FSDP does not backward prefetch, and there is no
communication and computation overlap in the backward pass. See
<a class="reference internal" href="#torch.distributed.fsdp.BackwardPrefetch" title="torch.distributed.fsdp.BackwardPrefetch"><code class="xref py py-class docutils literal notranslate"><span class="pre">BackwardPrefetch</span></code></a> for details. (Default: <code class="docutils literal notranslate"><span class="pre">BACKWARD_PRE</span></code>)</p></li>
<li><p><strong>mixed_precision</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.MixedPrecision" title="torch.distributed.fsdp.MixedPrecision"><em>MixedPrecision</em></a><em>]</em>) – This configures native mixed precision for FSDP. If this is set to
<code class="docutils literal notranslate"><span class="pre">None</span></code>, then no mixed precision is used. Otherwise, parameter,
buffer, and gradient reduction dtypes can be set. See
<a class="reference internal" href="#torch.distributed.fsdp.MixedPrecision" title="torch.distributed.fsdp.MixedPrecision"><code class="xref py py-class docutils literal notranslate"><span class="pre">MixedPrecision</span></code></a> for details. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>ignored_modules</strong> (<em>Optional</em><em>[</em><em>Iterable</em><em>[</em><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a><em>]</em><em>]</em>) – Modules whose
own parameters and child modules’ parameters and buffers are
ignored by this instance. None of the modules directly in
<code class="docutils literal notranslate"><span class="pre">ignored_modules</span></code> should be <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a>
instances, and any child modules that are already-constructed
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> 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
<code class="docutils literal notranslate"><span class="pre">auto_wrap_policy</span></code> or if parameters’ sharding is not managed by
FSDP. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>param_init_fn</strong> (<em>Optional</em><em>[</em><em>Callable</em><em>[</em><em>[</em><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>nn.Module</em></a><em>]</em><em>, </em><em>None</em><em>]</em><em>]</em>) – <p>A <code class="docutils literal notranslate"><span class="pre">Callable[torch.nn.Module]</span> <span class="pre">-&gt;</span> <span class="pre">None</span></code> that
specifies how modules that are currently on the meta device should
be initialized onto an actual device. As of v1.12, FSDP detects
modules with parameters or buffers on meta device via <code class="docutils literal notranslate"><span class="pre">is_meta</span></code>
and either applies <code class="docutils literal notranslate"><span class="pre">param_init_fn</span></code> if specified or calls
<code class="docutils literal notranslate"><span class="pre">nn.Module.reset_parameters()</span></code> otherwise. For both cases, the
implementation should <em>only</em> initialize the parameters/buffers of
the module, not those of its submodules. This is to avoid
re-initialization. In addition, FSDP also supports deferred
initialization via torchdistX’s (<a class="reference external" href="https://github.com/pytorch/torchdistX">https://github.com/pytorch/torchdistX</a>)
<code class="docutils literal notranslate"><span class="pre">deferred_init()</span></code> API, where the deferred modules are initialized
by calling <code class="docutils literal notranslate"><span class="pre">param_init_fn</span></code> if specified or torchdistX’s default
<code class="docutils literal notranslate"><span class="pre">materialize_module()</span></code> otherwise. If <code class="docutils literal notranslate"><span class="pre">param_init_fn</span></code> is
specified, then it is applied to all meta-device modules, meaning
that it should probably case on the module type. FSDP calls the
initialization function before parameter flattening and sharding.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">module</span> <span class="o">=</span> <span class="n">MyModule</span><span class="p">(</span><span class="n">device</span><span class="o">=</span><span class="s2">&quot;meta&quot;</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">def</span> <span class="nf">my_init_fn</span><span class="p">(</span><span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="c1"># E.g. initialize depending on the module type</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">fsdp_model</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">param_init_fn</span><span class="o">=</span><span class="n">my_init_fn</span><span class="p">,</span> <span class="n">auto_wrap_policy</span><span class="o">=</span><span class="n">size_based_auto_wrap_policy</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="nb">print</span><span class="p">(</span><span class="nb">next</span><span class="p">(</span><span class="n">fsdp_model</span><span class="o">.</span><span class="n">parameters</span><span class="p">())</span><span class="o">.</span><span class="n">device</span><span class="p">)</span> <span class="c1"># current CUDA device</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># With torchdistX</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">module</span> <span class="o">=</span> <span class="n">deferred_init</span><span class="o">.</span><span class="n">deferred_init</span><span class="p">(</span><span class="n">MyModule</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Will initialize via deferred_init.materialize_module().</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">fsdp_model</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">auto_wrap_policy</span><span class="o">=</span><span class="n">size_based_auto_wrap_policy</span><span class="p">)</span>
</pre></div>
</div>
</p></li>
<li><p><strong>device_id</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a><em>, </em><a class="reference internal" href="tensor_attributes.html#torch.device" title="torch.device"><em>torch.device</em></a><em>]</em><em>]</em>) – An <code class="docutils literal notranslate"><span class="pre">int</span></code> or
<code class="docutils literal notranslate"><span class="pre">torch.device</span></code> giving the CUDA device on which FSDP
initialization takes place, including the module initialization
if needed and the parameter sharding. This should be specified to
improve initialization speed if <code class="docutils literal notranslate"><span class="pre">module</span></code> is on CPU. If the
default CUDA device was set (e.g. via <code class="docutils literal notranslate"><span class="pre">torch.cuda.set_device</span></code>),
then the user may pass <code class="docutils literal notranslate"><span class="pre">torch.cuda.current_device</span></code> to this.
(Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>sync_module_states</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then each FSDP module will
broadcast module parameters and buffers from rank 0 to ensure that
they are replicated across ranks (adding communication overhead to
this constructor). This can help load <code class="docutils literal notranslate"><span class="pre">state_dict</span></code> checkpoints
via <code class="docutils literal notranslate"><span class="pre">load_state_dict</span></code> in a memory efficient way. See
<a class="reference internal" href="#torch.distributed.fsdp.FullStateDictConfig" title="torch.distributed.fsdp.FullStateDictConfig"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullStateDictConfig</span></code></a> for an example of this. (Default:
<code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>forward_prefetch</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP <em>explicitly</em> prefetches
the next forward-pass all-gather before the current forward
computation. This is only useful for CPU-bound workloads, in which
case issuing the next all-gather earlier may improve overlap. This
should only be used for static-graph models since the prefetching
follows the first iteration’s execution order. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>limit_all_gathers</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP explicitly
synchronizes the CPU thread to ensure GPU memory usage from only
<em>two</em> consecutive FSDP instances (the current instance running
computation and the next instance whose all-gather is prefetched).
If <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP allows the CPU thread to issue all-gathers
without any extra synchronization. (Default: <code class="docutils literal notranslate"><span class="pre">True</span></code>) We often
refer to this feature as the “rate limiter”. This flag should only
be set to <code class="docutils literal notranslate"><span class="pre">False</span></code> for specific CPU-bound workloads with low
memory pressure in which case the CPU thread can aggressively issue
all kernels without concern for the GPU memory usage.</p></li>
<li><p><strong>use_orig_params</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – Setting this to <code class="docutils literal notranslate"><span class="pre">True</span></code> has FSDP use
<code class="docutils literal notranslate"><span class="pre">module</span></code> ‘s original parameters. FSDP exposes those original
parameters to the user via <code class="xref py py-meth docutils literal notranslate"><span class="pre">nn.Module.named_parameters()</span></code>
instead of FSDP’s internal <code class="xref py py-class docutils literal notranslate"><span class="pre">FlatParameter</span></code> s. This means
that the optimizer step runs on the original parameters, enabling
per-original-parameter hyperparameters. FSDP preserves the original
parameter variables and manipulates their data between unsharded
and sharded forms, where they are always views into the underlying
unsharded or sharded <code class="xref py py-class docutils literal notranslate"><span class="pre">FlatParameter</span></code>, respectively. With the
current algorithm, the sharded form is always 1D, losing the
original tensor structure. An original parameter may have all,
some, or none of its data present for a given rank. In the none
case, its data will be like a size-0 empty tensor. Users should not
author programs relying on what data is present for a given
original parameter in its sharded form. <code class="docutils literal notranslate"><span class="pre">True</span></code> is required to
use <code class="docutils literal notranslate"><span class="pre">torch.compile()</span></code>. Setting this to <code class="docutils literal notranslate"><span class="pre">False</span></code> exposes FSDP’s
internal <code class="xref py py-class docutils literal notranslate"><span class="pre">FlatParameter</span></code> s to the user via
<code class="xref py py-meth docutils literal notranslate"><span class="pre">nn.Module.named_parameters()</span></code>. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>ignored_states</strong> (<em>Optional</em><em>[</em><em>Iterable</em><em>[</em><em>torch.nn.Parameter</em><em>]</em><em>]</em><em>, </em><em>Optional</em><em>[</em><em>Iterable</em><em>[</em><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a><em>]</em><em>]</em>) – Ignored parameters or modules that will not be managed by this FSDP
instance, meaning that the parameters are not sharded and their
gradients are not reduced across ranks. This argument unifies with
the existing <code class="docutils literal notranslate"><span class="pre">ignored_modules</span></code> argument, and we may deprecate
<code class="docutils literal notranslate"><span class="pre">ignored_modules</span></code> soon. For backward compatibility, we keep both
<code class="docutils literal notranslate"><span class="pre">ignored_states</span></code> and <cite>ignored_modules`</cite>, but FSDP only allows one
of them to be specified as not <code class="docutils literal notranslate"><span class="pre">None</span></code>.</p></li>
</ul>
</dd>
</dl>
<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.apply">
<span class="sig-name descname"><span class="pre">apply</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fn</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.apply"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.apply" title="Permalink to this definition">¶</a></dt>
<dd><p>Apply <code class="docutils literal notranslate"><span class="pre">fn</span></code> recursively to every submodule (as returned by <code class="docutils literal notranslate"><span class="pre">.children()</span></code>) as well as self.</p>
<p>Typical use includes initializing the parameters of a model (see also <a class="reference internal" href="nn.init.html#nn-init-doc"><span class="std std-ref">torch.nn.init</span></a>).</p>
<p>Compared to <code class="docutils literal notranslate"><span class="pre">torch.nn.Module.apply</span></code>, this version additionally gathers
the full parameters before applying <code class="docutils literal notranslate"><span class="pre">fn</span></code>. It should not be called from
within another <code class="docutils literal notranslate"><span class="pre">summon_full_params</span></code> context.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><p><strong>fn</strong> (<code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code> -&gt; None) – function to be applied to each submodule</p>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>self</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><p id="torch.nn.Module"/><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module">Module</a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.check_is_root">
<span class="sig-name descname"><span class="pre">check_is_root</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.check_is_root"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.check_is_root" title="Permalink to this definition">¶</a></dt>
<dd><p>Check if this instance is a root FSDP module.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)">bool</a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.clip_grad_norm_">
<span class="sig-name descname"><span class="pre">clip_grad_norm_</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">max_norm</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">norm_type</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">2.0</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.clip_grad_norm_"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.clip_grad_norm_" title="Permalink to this definition">¶</a></dt>
<dd><p>Clip the gradient norm of all parameters.</p>
<p>The norm is computed over all parameters’ gradients as viewed as a single vector, and the
gradients are modified in-place.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>max_norm</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#float" title="(in Python v3.12)"><em>float</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a>) – max norm of the gradients</p></li>
<li><p><strong>norm_type</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#float" title="(in Python v3.12)"><em>float</em></a><em> or </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.12)"><em>int</em></a>) – type of the used p-norm. Can be <code class="docutils literal notranslate"><span class="pre">'inf'</span></code>
for infinity norm.</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>Total norm of the parameters (viewed as a single vector).</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><em>Tensor</em></a></p>
</dd>
</dl>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>If every FSDP instance uses <code class="docutils literal notranslate"><span class="pre">NO_SHARD</span></code>, meaning that no
gradients are sharded across ranks, then you may directly use
<a class="reference internal" href="generated/torch.nn.utils.clip_grad_norm_.html#torch.nn.utils.clip_grad_norm_" title="torch.nn.utils.clip_grad_norm_"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.utils.clip_grad_norm_()</span></code></a>.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>If at least some FSDP instance uses a sharded strategy (i.e.
one other than <code class="docutils literal notranslate"><span class="pre">NO_SHARD</span></code>), then you should use this method
instead of <a class="reference internal" href="generated/torch.nn.utils.clip_grad_norm_.html#torch.nn.utils.clip_grad_norm_" title="torch.nn.utils.clip_grad_norm_"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.utils.clip_grad_norm_()</span></code></a> since this method
handles the fact that gradients are sharded across ranks.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>The total norm returned will have the “largest” dtype across
all parameters/gradients as defined by PyTorch’s type promotion
semantics. For example, if <em>all</em> 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.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This needs to be called on all ranks since it uses
collective communications.</p>
</div>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.flatten_sharded_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">flatten_sharded_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">sharded_optim_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.flatten_sharded_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.flatten_sharded_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Flatten a sharded optimizer state-dict.</p>
<p>The API is similar to <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a>. The only
difference is that the input <code class="docutils literal notranslate"><span class="pre">sharded_optim_state_dict</span></code> should be
returned from <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.sharded_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.sharded_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">sharded_optim_state_dict()</span></code></a>. Therefore, there will
be all-gather calls on each rank to gather <code class="docutils literal notranslate"><span class="pre">ShardedTensor</span></code> s.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>sharded_optim_state_dict</strong> (<em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em>) – Optimizer state dict
corresponding to the unflattened parameters and holding the
sharded optimizer state.</p></li>
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Refer to <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a>.</p></li>
<li><p><strong>optim</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer for <code class="docutils literal notranslate"><span class="pre">model</span></code> ‘s
parameters.</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>Refer to <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a>.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.12)"><em>Dict</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.12)"><em>Any</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.forward">
<span class="sig-name descname"><span class="pre">forward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">**</span></span><span class="n"><span class="pre">kwargs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.forward"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.forward" title="Permalink to this definition">¶</a></dt>
<dd><p>Run the forward pass for the wrapped module, inserting FSDP-specific pre- and post-forward sharding logic.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.12)"><em>Any</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.fsdp_modules">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">fsdp_modules</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">root_only</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.fsdp_modules"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.fsdp_modules" title="Permalink to this definition">¶</a></dt>
<dd><p>Return all nested FSDP instances.</p>
<p>This possibly includes <code class="docutils literal notranslate"><span class="pre">module</span></code> itself and only includes FSDP root modules if <code class="docutils literal notranslate"><span class="pre">root_only=True</span></code>.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>module</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module, which may or may not be an
<code class="docutils literal notranslate"><span class="pre">FSDP</span></code> module.</p></li>
<li><p><strong>root_only</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – Whether to return only FSDP root modules.
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>FSDP modules that are nested in
the input <code class="docutils literal notranslate"><span class="pre">module</span></code>.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p>List[<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel">FullyShardedDataParallel</a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">full_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_input</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">rank0_only</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.full_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Return the full optimizer state-dict.</p>
<p>Consolidates the full optimizer state on rank 0 and returns it
as a <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a> following the convention of
<a class="reference internal" href="generated/torch.optim.Optimizer.state_dict.html#torch.optim.Optimizer.state_dict" title="torch.optim.Optimizer.state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">torch.optim.Optimizer.state_dict()</span></code></a>, i.e. with keys <code class="docutils literal notranslate"><span class="pre">&quot;state&quot;</span></code>
and <code class="docutils literal notranslate"><span class="pre">&quot;param_groups&quot;</span></code>. The flattened parameters in <code class="docutils literal notranslate"><span class="pre">FSDP</span></code> modules
contained in <code class="docutils literal notranslate"><span class="pre">model</span></code> are mapped back to their unflattened parameters.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This needs to be called on all ranks since it uses
collective communications. However, if <code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code>, then
the state dict is only populated on rank 0, and all other ranks
return an empty <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a>.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Unlike <code class="docutils literal notranslate"><span class="pre">torch.optim.Optimizer.state_dict()</span></code>, this method
uses full parameter names as keys instead of parameter IDs.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Like in <a class="reference internal" href="generated/torch.optim.Optimizer.state_dict.html#torch.optim.Optimizer.state_dict" title="torch.optim.Optimizer.state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">torch.optim.Optimizer.state_dict()</span></code></a>, 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
<code class="docutils literal notranslate"><span class="pre">torch.save()</span></code>.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module (which may or may not be a
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> instance) whose parameters
were passed into the optimizer <code class="docutils literal notranslate"><span class="pre">optim</span></code>.</p></li>
<li><p><strong>optim</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer for <code class="docutils literal notranslate"><span class="pre">model</span></code> ‘s
parameters.</p></li>
<li><p><strong>optim_input</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><em>List</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em><em>, </em><em>Iterable</em><em>[</em><em>torch.nn.Parameter</em><em>]</em><em>]</em><em>]</em>) – Input passed into the optimizer <code class="docutils literal notranslate"><span class="pre">optim</span></code> representing either a
<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">list</span></code></a> of parameter groups or an iterable of parameters;
if <code class="docutils literal notranslate"><span class="pre">None</span></code>, then this method assumes the input was
<code class="docutils literal notranslate"><span class="pre">model.parameters()</span></code>. This argument is deprecated, and there
is no need to pass it in anymore. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>rank0_only</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, saves the populated <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a>
only on rank 0; if <code class="docutils literal notranslate"><span class="pre">False</span></code>, saves it on all ranks. (Default:
<code class="docutils literal notranslate"><span class="pre">True</span></code>)</p></li>
<li><p><strong>group</strong> (<em>dist.ProcessGroup</em>) – Model’s process group or <code class="docutils literal notranslate"><span class="pre">None</span></code> if using
the default process group. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>A <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a> containing the optimizer state for
<code class="docutils literal notranslate"><span class="pre">model</span></code> ‘s original unflattened parameters and including keys
“state” and “param_groups” following the convention of
<a class="reference internal" href="generated/torch.optim.Optimizer.state_dict.html#torch.optim.Optimizer.state_dict" title="torch.optim.Optimizer.state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">torch.optim.Optimizer.state_dict()</span></code></a>. If <code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code>,
then nonzero ranks return an empty <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a>.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p>Dict[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, Any]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.get_state_dict_type">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">get_state_dict_type</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.get_state_dict_type"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.get_state_dict_type" title="Permalink to this definition">¶</a></dt>
<dd><p>Get the state_dict_type and the corresponding configurations for the FSDP modules rooted at <code class="docutils literal notranslate"><span class="pre">module</span></code>.</p>
<p>The target module does not have to be an FSDP module.</p>
<dl class="field-list simple">
<dt class="field-odd">Returns</dt>
<dd class="field-odd"><p>A <code class="docutils literal notranslate"><span class="pre">StateDictSettings</span></code> containing the state_dict_type and
state_dict / optim_state_dict configs that are currently set.</p>
</dd>
<dt class="field-even">Raises</dt>
<dd class="field-even"><ul class="simple">
<li><p><strong>AssertionError` if the StateDictSettings for differen</strong> – </p></li>
<li><p><strong>FSDP submodules differ.</strong> – </p></li>
</ul>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.distributed.fsdp.StateDictSettings" title="torch.distributed.fsdp.api.StateDictSettings"><em>StateDictSettings</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py property">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.module">
<em class="property"><span class="pre">property</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">module</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.modules.module.Module"><span class="pre">Module</span></a></em><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.module" title="Permalink to this definition">¶</a></dt>
<dd><p>Return the wrapped module.</p>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.named_buffers">
<span class="sig-name descname"><span class="pre">named_buffers</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">**</span></span><span class="n"><span class="pre">kwargs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.named_buffers"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.named_buffers" title="Permalink to this definition">¶</a></dt>
<dd><p>Return an iterator over module buffers, yielding both the name of the buffer and the buffer itself.</p>
<p>Intercepts buffer names and removes all occurrences of the FSDP-specific flattened buffer prefix
when inside the <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params" title="torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params"><code class="xref py py-meth docutils literal notranslate"><span class="pre">summon_full_params()</span></code></a> context manager.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.12)"><em>Iterator</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.12)"><em>Tuple</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, <a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><em>Tensor</em></a>]]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.named_parameters">
<span class="sig-name descname"><span class="pre">named_parameters</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">**</span></span><span class="n"><span class="pre">kwargs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.named_parameters"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.named_parameters" title="Permalink to this definition">¶</a></dt>
<dd><p>Return an iterator over module parameters, yielding both the name of the parameter and the parameter itself.</p>
<p>Intercepts parameter names and removes all occurrences of the FSDP-specific flattened parameter prefix
when inside the <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params" title="torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params"><code class="xref py py-meth docutils literal notranslate"><span class="pre">summon_full_params()</span></code></a> context manager.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.12)"><em>Iterator</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.12)"><em>Tuple</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, <a class="reference internal" href="generated/torch.nn.parameter.Parameter.html#torch.nn.parameter.Parameter" title="torch.nn.parameter.Parameter"><em>Parameter</em></a>]]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.no_sync">
<span class="sig-name descname"><span class="pre">no_sync</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.no_sync"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.no_sync" title="Permalink to this definition">¶</a></dt>
<dd><p>Disable gradient synchronizations across FSDP instances.</p>
<p>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.</p>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This likely results in higher memory usage because FSDP will
accumulate the full model gradients (instead of gradient shards)
until the eventual sync.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>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.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Generator" title="(in Python v3.12)"><em>Generator</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_dict</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Transform the state-dict of an optimizer corresponding to a sharded model.</p>
<p>The given state-dict can be transformed to one of three types:
1) full optimizer state_dict, 2) sharded optimizer state_dict, 3) local optimizer state_dict.</p>
<p>For full optimizer state_dict, all states are unflattened and not sharded.
Rank0 only and CPU only can be specified via <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type" title="torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type"><code class="xref py py-meth docutils literal notranslate"><span class="pre">state_dict_type()</span></code></a> to
avoid OOM.</p>
<p>For sharded optimizer state_dict, all states are unflattened but sharded.
CPU only can be specified via <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type" title="torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type"><code class="xref py py-meth docutils literal notranslate"><span class="pre">state_dict_type()</span></code></a> to further save
memory.</p>
<p>For local state_dict, no transformation will be performed. But a state
will be converted from nn.Tensor to ShardedTensor to represent its sharding
nature (this is not supported yet).</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">StateDictType</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullStateDictConfig</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullOptimStateDictConfig</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Save a checkpoint</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">FSDP</span><span class="o">.</span><span class="n">set_state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">FULL_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullOptimStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">state_dict</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">optim_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">save_a_checkpoint</span><span class="p">(</span><span class="n">state_dict</span><span class="p">,</span> <span class="n">optim_state_dict</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Load a checkpoint</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">state_dict</span><span class="p">,</span> <span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">load_a_checkpoint</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">FSDP</span><span class="o">.</span><span class="n">set_state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">FULL_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullOptimStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">optim_state_dict_to_load</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">,</span> <span class="n">optim_state_dict</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">optim_state_dict</span><span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module (which may or may not be a
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> instance) whose parameters
were passed into the optimizer <code class="docutils literal notranslate"><span class="pre">optim</span></code>.</p></li>
<li><p><strong>optim</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer for <code class="docutils literal notranslate"><span class="pre">model</span></code> ‘s
parameters.</p></li>
<li><p><strong>optim_state_dict</strong> (<em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em>) – the target optimizer state_dict to
transform. If the value is None, optim.state_dict() will be used. (
Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>group</strong> (<em>dist.ProcessGroup</em>) – Model’s process group across which parameters
are sharded or <code class="docutils literal notranslate"><span class="pre">None</span></code> if using the default process group. (
Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>A <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a> containing the optimizer state for
<code class="docutils literal notranslate"><span class="pre">model</span></code>. The sharding of the optimizer state is based on
<code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code>.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p>Dict[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, Any]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict_to_load">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">optim_state_dict_to_load</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">is_named_optimizer</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">load_directly</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.optim_state_dict_to_load"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict_to_load" title="Permalink to this definition">¶</a></dt>
<dd><p>Convert an optimizer state-dict so that it can be loaded into the optimizer associated with the FSDP model.</p>
<p>Given a <code class="docutils literal notranslate"><span class="pre">optim_state_dict</span></code> that is transformed through
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">optim_state_dict()</span></code></a>, it gets converted to the flattened optimizer
state_dict that can be loaded to <code class="docutils literal notranslate"><span class="pre">optim</span></code> which is the optimizer for
<code class="docutils literal notranslate"><span class="pre">model</span></code>. <code class="docutils literal notranslate"><span class="pre">model</span></code> must be sharded by FullyShardedDataParallel.</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">StateDictType</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullStateDictConfig</span>
<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullOptimStateDictConfig</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Save a checkpoint</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">FSDP</span><span class="o">.</span><span class="n">set_state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">FULL_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullOptimStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">state_dict</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">original_osd</span> <span class="o">=</span> <span class="n">optim</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">optim_state_dict</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">optim</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">optim_state_dict</span><span class="o">=</span><span class="n">original_osd</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">save_a_checkpoint</span><span class="p">(</span><span class="n">state_dict</span><span class="p">,</span> <span class="n">optim_state_dict</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Load a checkpoint</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">state_dict</span><span class="p">,</span> <span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">load_a_checkpoint</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">FSDP</span><span class="o">.</span><span class="n">set_state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">FULL_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">FullOptimStateDictConfig</span><span class="p">(</span><span class="n">rank0_only</span><span class="o">=</span><span class="kc">False</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">optim_state_dict_to_load</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">,</span> <span class="n">optim_state_dict</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">optim_state_dict</span><span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module (which may or may not be a
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> instance) whose parameters
were passed into the optimizer <code class="docutils literal notranslate"><span class="pre">optim</span></code>.</p></li>
<li><p><strong>optim</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer for <code class="docutils literal notranslate"><span class="pre">model</span></code> ‘s
parameters.</p></li>
<li><p><strong>optim_state_dict</strong> (<em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em>) – The optimizer states to be loaded.</p></li>
<li><p><strong>is_named_optimizer</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – Is this optimizer a NamedOptimizer or
KeyedOptimizer. Only set to True if <code class="docutils literal notranslate"><span class="pre">optim</span></code> is TorchRec’s
KeyedOptimizer or torch.distributed’s NamedOptimizer.</p></li>
<li><p><strong>load_directly</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If this is set to True, this API will also
call optim.load_state_dict(result) before returning the result.
Otherwise, users are responsible to call <code class="docutils literal notranslate"><span class="pre">optim.load_state_dict()</span></code>
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>group</strong> (<em>dist.ProcessGroup</em>) – Model’s process group across which parameters
are sharded or <code class="docutils literal notranslate"><span class="pre">None</span></code> if using the default process group. (
Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Return type</dt>
<dd class="field-even"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.12)"><em>Dict</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.12)"><em>Any</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.register_comm_hook">
<span class="sig-name descname"><span class="pre">register_comm_hook</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">state</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">hook</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.register_comm_hook"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.register_comm_hook" title="Permalink to this definition">¶</a></dt>
<dd><p>Register a communication hook.</p>
<p>This 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
<a class="reference external" href="https://arxiv.org/abs/1803.05880">GossipGrad</a> and gradient compression
which involve different communication strategies for
parameter syncs while training with <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a>.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>FSDP communication hook should be registered before running an initial forward pass
and only once.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>state</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#object" title="(in Python v3.12)"><em>object</em></a>) – <p>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 <a class="reference external" href="https://arxiv.org/abs/1803.05880">GossipGrad</a>, etc.
It is locally stored by each worker
and shared by all the gradient tensors on the worker.</p>
</p></li>
<li><p><strong>hook</strong> (<em>Callable</em>) – Callable, which has one of the following signatures:
1) <code class="docutils literal notranslate"><span class="pre">hook:</span> <span class="pre">Callable[torch.Tensor]</span> <span class="pre">-&gt;</span> <span class="pre">None</span></code>:
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 <code class="docutils literal notranslate"><span class="pre">None</span></code>;
2) <code class="docutils literal notranslate"><span class="pre">hook:</span> <span class="pre">Callable[torch.Tensor,</span> <span class="pre">torch.Tensor]</span> <span class="pre">-&gt;</span> <span class="pre">None</span></code>:
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 <code class="docutils literal notranslate"><span class="pre">None</span></code>.
Callables with signature 1 are expected to handle gradient communication for a <cite>NO_SHARD</cite> case.
Callables with signature 2 are expected to handle gradient communication for sharded cases.</p></li>
</ul>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.rekey_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">rekey_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">optim_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_key_type</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_input</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.rekey_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.rekey_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Re-keys the optimizer state dict <code class="docutils literal notranslate"><span class="pre">optim_state_dict</span></code> to use the key type <code class="docutils literal notranslate"><span class="pre">optim_state_key_type</span></code>.</p>
<p>This can be used to achieve compatibility between optimizer state dicts from models with FSDP
instances and ones without.</p>
<p>To re-key an FSDP full optimizer state dict (i.e. from
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">full_optim_state_dict()</span></code></a>) to use parameter IDs and be loadable to
a non-wrapped model:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">wrapped_model</span><span class="p">,</span> <span class="n">wrapped_optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">full_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">full_optim_state_dict</span><span class="p">(</span><span class="n">wrapped_model</span><span class="p">,</span> <span class="n">wrapped_optim</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">nonwrapped_model</span><span class="p">,</span> <span class="n">nonwrapped_optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">rekeyed_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">rekey_optim_state_dict</span><span class="p">(</span><span class="n">full_osd</span><span class="p">,</span> <span class="n">OptimStateKeyType</span><span class="o">.</span><span class="n">PARAM_ID</span><span class="p">,</span> <span class="n">nonwrapped_model</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">nonwrapped_optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">rekeyed_osd</span><span class="p">)</span>
</pre></div>
</div>
<p>To re-key a normal optimizer state dict from a non-wrapped model to be
loadable to a wrapped model:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">nonwrapped_model</span><span class="p">,</span> <span class="n">nonwrapped_optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">osd</span> <span class="o">=</span> <span class="n">nonwrapped_optim</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">rekeyed_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">rekey_optim_state_dict</span><span class="p">(</span><span class="n">osd</span><span class="p">,</span> <span class="n">OptimStateKeyType</span><span class="o">.</span><span class="n">PARAM_NAME</span><span class="p">,</span> <span class="n">nonwrapped_model</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">wrapped_model</span><span class="p">,</span> <span class="n">wrapped_optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">sharded_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">shard_full_optim_state_dict</span><span class="p">(</span><span class="n">rekeyed_osd</span><span class="p">,</span> <span class="n">wrapped_model</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">wrapped_optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">sharded_osd</span><span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Returns</dt>
<dd class="field-odd"><p>The optimizer state dict re-keyed using the
parameter keys specified by <code class="docutils literal notranslate"><span class="pre">optim_state_key_type</span></code>.</p>
</dd>
<dt class="field-even">Return type</dt>
<dd class="field-even"><p>Dict[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, Any]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">scatter_full_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">full_optim_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_input</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.scatter_full_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Scatter the full optimizer state dict from rank 0 to all other ranks.</p>
<p>Returns the sharded optimizer state dict on each rank.
The return value is the same as <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a>, and on rank
0, the first argument should be the return value of
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">full_optim_state_dict()</span></code></a>.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">full_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">full_optim_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">)</span>  <span class="c1"># only non-empty on rank 0</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Define new model with possibly different world size</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">new_model</span><span class="p">,</span> <span class="n">new_optim</span><span class="p">,</span> <span class="n">new_group</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">sharded_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">scatter_full_optim_state_dict</span><span class="p">(</span><span class="n">full_osd</span><span class="p">,</span> <span class="n">new_model</span><span class="p">,</span> <span class="n">group</span><span class="o">=</span><span class="n">new_group</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">new_optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">sharded_osd</span><span class="p">)</span>
</pre></div>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Both <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a> and
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">scatter_full_optim_state_dict()</span></code></a> 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.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>full_optim_state_dict</strong> (<em>Optional</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em>) – 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.</p></li>
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module (which may or may not be a
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> instance) whose parameters
correspond to the optimizer state in <code class="docutils literal notranslate"><span class="pre">full_optim_state_dict</span></code>.</p></li>
<li><p><strong>optim_input</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><em>List</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em><em>, </em><em>Iterable</em><em>[</em><em>torch.nn.Parameter</em><em>]</em><em>]</em><em>]</em>) – Input passed into the optimizer representing either a
<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">list</span></code></a> of parameter groups or an iterable of parameters;
if <code class="docutils literal notranslate"><span class="pre">None</span></code>, then this method assumes the input was
<code class="docutils literal notranslate"><span class="pre">model.parameters()</span></code>. This argument is deprecated, and there
is no need to pass it in anymore. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>optim</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a><em>]</em>) – Optimizer that will load
the state dict returned by this method. This is the preferred
argument to use over <code class="docutils literal notranslate"><span class="pre">optim_input</span></code>. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>group</strong> (<em>dist.ProcessGroup</em>) – Model’s process group or <code class="docutils literal notranslate"><span class="pre">None</span></code> if
using the default process group. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>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.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p>Dict[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, Any]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">set_state_dict_type</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict_type</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict_config</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_dict_config</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.set_state_dict_type"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type" title="Permalink to this definition">¶</a></dt>
<dd><p>Set the <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code> of all the descendant FSDP modules of the target module.</p>
<p>Also takes (optional) configuration for the model’s and optimizer’s state dict.
The target module does not have to be a FSDP module. If the target
module is a FSDP module, its <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code> will also be changed.</p>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This API should be called for only the top-level (root)
module.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This API enables users to transparently use the conventional
<code class="docutils literal notranslate"><span class="pre">state_dict</span></code> API to take model checkpoints in cases where the
root FSDP module is wrapped by another <code class="docutils literal notranslate"><span class="pre">nn.Module</span></code>. For example,
the following will ensure <code class="docutils literal notranslate"><span class="pre">state_dict</span></code> is called on all non-FSDP
instances, while dispatching into <cite>sharded_state_dict</cite> implementation
for FSDP:</p>
</div>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">model</span> <span class="o">=</span> <span class="n">DDP</span><span class="p">(</span><span class="n">FSDP</span><span class="p">(</span><span class="o">...</span><span class="p">))</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">FSDP</span><span class="o">.</span><span class="n">set_state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">SHARDED_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">state_dict_config</span> <span class="o">=</span> <span class="n">ShardedStateDictConfig</span><span class="p">(</span><span class="n">offload_to_cpu</span><span class="o">=</span><span class="kc">True</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">optim_state_dict_config</span> <span class="o">=</span> <span class="n">OptimStateDictConfig</span><span class="p">(</span><span class="n">offload_to_cpu</span><span class="o">=</span><span class="kc">True</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">param_state_dict</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">optim_state_dict</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">optim_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>module</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module.</p></li>
<li><p><strong>state_dict_type</strong> (<em>StateDictType</em>) – the desired <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code> to set.</p></li>
<li><p><strong>state_dict_config</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.StateDictConfig" title="torch.distributed.fsdp.StateDictConfig"><em>StateDictConfig</em></a><em>]</em>) – the configuration for the
target <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code>.</p></li>
<li><p><strong>optim_state_dict_config</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.OptimStateDictConfig" title="torch.distributed.fsdp.OptimStateDictConfig"><em>OptimStateDictConfig</em></a><em>]</em>) – the configuration
for the optimizer state dict.</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>A StateDictSettings that include the previous state_dict type and
configuration for the module.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.distributed.fsdp.StateDictSettings" title="torch.distributed.fsdp.api.StateDictSettings"><em>StateDictSettings</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">shard_full_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">full_optim_state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_input</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.shard_full_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Shard a full optimizer state-dict.</p>
<p>Remaps the state in <code class="docutils literal notranslate"><span class="pre">full_optim_state_dict</span></code> to flattened parameters instead of unflattened
parameters and restricts to only this rank’s part of the optimizer state.
The first argument should be the return value of <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">full_optim_state_dict()</span></code></a>.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">full_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">full_optim_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optim</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">torch</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">full_osd</span><span class="p">,</span> <span class="n">PATH</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Define new model with possibly different world size</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">new_model</span><span class="p">,</span> <span class="n">new_optim</span> <span class="o">=</span> <span class="o">...</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">full_osd</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">PATH</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">sharded_osd</span> <span class="o">=</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">shard_full_optim_state_dict</span><span class="p">(</span><span class="n">full_osd</span><span class="p">,</span> <span class="n">new_model</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">new_optim</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">sharded_osd</span><span class="p">)</span>
</pre></div>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Both <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">shard_full_optim_state_dict()</span></code></a> and
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">scatter_full_optim_state_dict()</span></code></a> 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.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>full_optim_state_dict</strong> (<em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em>) – Optimizer state dict
corresponding to the unflattened parameters and holding the
full non-sharded optimizer state.</p></li>
<li><p><strong>model</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module (which may or may not be a
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a> instance) whose parameters
correspond to the optimizer state in <code class="docutils literal notranslate"><span class="pre">full_optim_state_dict</span></code>.</p></li>
<li><p><strong>optim_input</strong> (<em>Optional</em><em>[</em><em>Union</em><em>[</em><em>List</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em><em>, </em><em>Iterable</em><em>[</em><em>torch.nn.Parameter</em><em>]</em><em>]</em><em>]</em>) – Input passed into the optimizer representing either a
<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#list" title="(in Python v3.12)"><code class="xref py py-class docutils literal notranslate"><span class="pre">list</span></code></a> of parameter groups or an iterable of parameters;
if <code class="docutils literal notranslate"><span class="pre">None</span></code>, then this method assumes the input was
<code class="docutils literal notranslate"><span class="pre">model.parameters()</span></code>. This argument is deprecated, and there
is no need to pass it in anymore. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>optim</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a><em>]</em>) – Optimizer that will load
the state dict returned by this method. This is the preferred
argument to use over <code class="docutils literal notranslate"><span class="pre">optim_input</span></code>. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>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.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p>Dict[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, Any]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.sharded_optim_state_dict">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">sharded_optim_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">model</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">group</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.sharded_optim_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.sharded_optim_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Return the optimizer state-dict in its sharded form.</p>
<p>The API is similar to <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">full_optim_state_dict()</span></code></a> but this API chunks
all non-zero-dimension states to <code class="xref py py-class docutils literal notranslate"><span class="pre">ShardedTensor</span></code> to save memory.
This API should only be used when the model <code class="docutils literal notranslate"><span class="pre">state_dict</span></code> is derived
with the context manager <code class="docutils literal notranslate"><span class="pre">with</span> <span class="pre">state_dict_type(SHARDED_STATE_DICT):</span></code>.</p>
<p>For the detailed usage, refer to <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict" title="torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">full_optim_state_dict()</span></code></a>.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>The returned state dict contains <code class="docutils literal notranslate"><span class="pre">ShardedTensor</span></code> and
cannot be directly used by the regular <code class="docutils literal notranslate"><span class="pre">optim.load_state_dict</span></code>.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.12)"><em>Dict</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.12)">str</a>, <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.12)"><em>Any</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">state_dict_type</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict_type</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict_config</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_dict_config</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.state_dict_type"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type" title="Permalink to this definition">¶</a></dt>
<dd><p>Set the <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code> of all the descendant FSDP modules of the target module.</p>
<p>This context manager has the same functions as <a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type" title="torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type"><code class="xref py py-meth docutils literal notranslate"><span class="pre">set_state_dict_type()</span></code></a>. Read the document of
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type" title="torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type"><code class="xref py py-meth docutils literal notranslate"><span class="pre">set_state_dict_type()</span></code></a> for the detail.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">model</span> <span class="o">=</span> <span class="n">DDP</span><span class="p">(</span><span class="n">FSDP</span><span class="p">(</span><span class="o">...</span><span class="p">))</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">with</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">state_dict_type</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">StateDictType</span><span class="o">.</span><span class="n">SHARDED_STATE_DICT</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">):</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">checkpoint</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>module</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><em>torch.nn.Module</em></a>) – Root module.</p></li>
<li><p><strong>state_dict_type</strong> (<em>StateDictType</em>) – the desired <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code> to set.</p></li>
<li><p><strong>state_dict_config</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.StateDictConfig" title="torch.distributed.fsdp.StateDictConfig"><em>StateDictConfig</em></a><em>]</em>) – the model <code class="docutils literal notranslate"><span class="pre">state_dict</span></code>
configuration for the target <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code>.</p></li>
<li><p><strong>optim_state_dict_config</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.distributed.fsdp.OptimStateDictConfig" title="torch.distributed.fsdp.OptimStateDictConfig"><em>OptimStateDictConfig</em></a><em>]</em>) – the optimizer
<code class="docutils literal notranslate"><span class="pre">state_dict</span></code> configuration for the target <code class="docutils literal notranslate"><span class="pre">state_dict_type</span></code>.</p></li>
</ul>
</dd>
<dt class="field-even">Return type</dt>
<dd class="field-even"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Generator" title="(in Python v3.12)"><em>Generator</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params">
<em class="property"><span class="pre">static</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">summon_full_params</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">recurse</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">writeback</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">rank0_only</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">with_grads</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/fully_sharded_data_parallel.html#FullyShardedDataParallel.summon_full_params"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params" title="Permalink to this definition">¶</a></dt>
<dd><p>Expose full params for FSDP instances with this context manager.</p>
<p>Can be useful <em>after</em> 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 <code class="docutils literal notranslate"><span class="pre">recurse</span></code> argument.</p>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This can be used on inner FSDPs.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This can <em>not</em> be used within a forward or backward pass. Nor
can forward and backward be started from within this context.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Parameters will revert to their local shards after the context
manager exits, storage behavior is the same as forward.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>The full parameters can be modified, but only the portion
corresponding to the local param shard will persist after the
context manager exits (unless <code class="docutils literal notranslate"><span class="pre">writeback=False</span></code>, in which case
changes will be discarded). In the case where FSDP does not shard
the parameters, currently only when <code class="docutils literal notranslate"><span class="pre">world_size</span> <span class="pre">==</span> <span class="pre">1</span></code>, or <code class="docutils literal notranslate"><span class="pre">NO_SHARD</span></code>
config, the modification is persisted regardless of <code class="docutils literal notranslate"><span class="pre">writeback</span></code>.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>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.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Note that <code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code> in conjunction with
<code class="docutils literal notranslate"><span class="pre">writeback=True</span></code> 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.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Note that <code class="docutils literal notranslate"><span class="pre">offload_to_cpu</span></code> and <code class="docutils literal notranslate"><span class="pre">rank0_only=False</span></code> 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 <code class="docutils literal notranslate"><span class="pre">offload_to_cpu</span></code> with
<code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code>.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>recurse</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><em>Optional</em>) – recursively summon all params for nested
FSDP instances (default: True).</p></li>
<li><p><strong>writeback</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><em>Optional</em>) – if <code class="docutils literal notranslate"><span class="pre">False</span></code>, modifications to params are
discarded after the context manager exits;
disabling this can be slightly more efficient (default: True)</p></li>
<li><p><strong>rank0_only</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><em>Optional</em>) – if <code class="docutils literal notranslate"><span class="pre">True</span></code>, 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
<code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code> with <code class="docutils literal notranslate"><span class="pre">writeback=True</span></code> 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.</p></li>
<li><p><strong>offload_to_cpu</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><em>Optional</em>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, 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 <code class="docutils literal notranslate"><span class="pre">NO_SHARD</span></code> config). It is recommended
to use <code class="docutils literal notranslate"><span class="pre">offload_to_cpu</span></code> with <code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code> to avoid
redundant copies of model parameters being offloaded to the same CPU memory.</p></li>
<li><p><strong>with_grads</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a><em>, </em><em>Optional</em>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, gradients are also
unsharded with the parameters. Currently, this is only
supported when passing <code class="docutils literal notranslate"><span class="pre">use_orig_params=True</span></code> to the FSDP
constructor and <code class="docutils literal notranslate"><span class="pre">offload_to_cpu=False</span></code> to this method.
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
</ul>
</dd>
<dt class="field-even">Return type</dt>
<dd class="field-even"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Generator" title="(in Python v3.12)"><em>Generator</em></a></p>
</dd>
</dl>
</dd></dl>

</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.BackwardPrefetch">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">BackwardPrefetch</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#BackwardPrefetch"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.BackwardPrefetch" title="Permalink to this definition">¶</a></dt>
<dd><p>This configures explicit backward prefetching, which improves throughput by
enabling communication and computation overlap in the backward pass at the
cost of slightly increased memory usage.</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">BACKWARD_PRE</span></code>: This enables the most overlap but increases memory
usage the most. This prefetches the next set of parameters <em>before</em> the
current set of parameters’ gradient computation. This overlaps the <em>next
all-gather</em> and the <em>current gradient computation</em>, and at the peak, it
holds the current set of parameters, next set of parameters, and current
set of gradients in memory.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">BACKWARD_POST</span></code>: This enables less overlap but requires less memory
usage. This prefetches the next set of parameters <em>after</em> the current
set of parameters’ gradient computation. This overlaps the <em>current
reduce-scatter</em> and the <em>next gradient computation</em>, and it frees the
current set of parameters before allocating memory for the next set of
parameters, only holding the next set of parameters and current set of
gradients in memory at the peak.</p></li>
<li><p>FSDP’s <code class="docutils literal notranslate"><span class="pre">backward_prefetch</span></code> argument accepts <code class="docutils literal notranslate"><span class="pre">None</span></code>, which disables
the backward prefetching altogether. This has no overlap and does not
increase memory usage. In general, we do not recommend this setting since
it may degrade throughput significantly.</p></li>
</ul>
<p>For more technical context: For a single process group using NCCL backend,
any collectives, even if issued from different streams, contend for the
same per-device NCCL stream, which implies that the relative order in which
the collectives are issued matters for overlapping. The two backward
prefetching values correspond to different issue orders.</p>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.ShardingStrategy">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">ShardingStrategy</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#ShardingStrategy"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.ShardingStrategy" title="Permalink to this definition">¶</a></dt>
<dd><p>This specifies the sharding strategy to be used for distributed training by
<a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel" title="torch.distributed.fsdp.FullyShardedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a>.</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">FULL_SHARD</span></code>: 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.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">SHARD_GRAD_OP</span></code>: 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 <code class="docutils literal notranslate"><span class="pre">no_sync()</span></code>, the parameters are not resharded
after the backward computation.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">NO_SHARD</span></code>: Parameters, gradients, and optimizer states are not sharded
but instead replicated across ranks similar to PyTorch’s
<code class="xref py py-class docutils literal notranslate"><span class="pre">DistributedDataParallel</span></code> API. For gradients, this strategy
synchronizes them (via all-reduce) after the backward computation. The
unsharded optimizer states are updated locally per rank.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">HYBRID_SHARD</span></code>: Apply <code class="docutils literal notranslate"><span class="pre">FULL_SHARD</span></code> 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.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">_HYBRID_SHARD_ZERO2</span></code>: Apply <code class="docutils literal notranslate"><span class="pre">SHARD_GRAD_OP</span></code> within a node, and replicate parameters across
nodes. This is like <code class="docutils literal notranslate"><span class="pre">HYBRID_SHARD</span></code>, 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.</p></li>
</ul>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.MixedPrecision">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">MixedPrecision</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">param_dtype=None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">reduce_dtype=None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">buffer_dtype=None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">keep_low_precision_grads=False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cast_forward_inputs=False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cast_root_forward_inputs=True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">_module_classes_to_ignore=(&lt;class</span> <span class="pre">'torch.nn.modules.batchnorm._BatchNorm'&gt;</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">)</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#MixedPrecision"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.MixedPrecision" title="Permalink to this definition">¶</a></dt>
<dd><p>This configures FSDP-native mixed precision training.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>param_dtype</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="tensor_attributes.html#torch.dtype" title="torch.dtype"><em>torch.dtype</em></a><em>]</em>) – This specifies the dtype for model
parameters during forward and backward and thus the dtype for
forward and backward computation. Outside forward and backward, the
<em>sharded</em> parameters are kept in full precision (e.g. for the
optimizer step), and for model checkpointing, the parameters are
always saved in full precision. (Default: <code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>reduce_dtype</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="tensor_attributes.html#torch.dtype" title="torch.dtype"><em>torch.dtype</em></a><em>]</em>) – This specifies the dtype for
gradient reduction (i.e. reduce-scatter or all-reduce). If this is
<code class="docutils literal notranslate"><span class="pre">None</span></code> but <code class="docutils literal notranslate"><span class="pre">param_dtype</span></code> is not <code class="docutils literal notranslate"><span class="pre">None</span></code>, then this takes on
the <code class="docutils literal notranslate"><span class="pre">param_dtype</span></code> value, still running gradient reduction in low
precision. This is permitted to differ from <code class="docutils literal notranslate"><span class="pre">param_dtype</span></code>, e.g.
to force gradient reduction to run in full precision. (Default:
<code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>buffer_dtype</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="tensor_attributes.html#torch.dtype" title="torch.dtype"><em>torch.dtype</em></a><em>]</em>) – This specifies the dtype for
buffers. FSDP does not shard buffers. Rather, FSDP casts them to
<code class="docutils literal notranslate"><span class="pre">buffer_dtype</span></code> in the first forward pass and keeps them in that
dtype thereafter. For model checkpointing, the buffers are saved
in full precision except for <code class="docutils literal notranslate"><span class="pre">LOCAL_STATE_DICT</span></code>. (Default:
<code class="docutils literal notranslate"><span class="pre">None</span></code>)</p></li>
<li><p><strong>keep_low_precision_grads</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP upcasts
gradients to full precision after the backward pass in preparation
for the optimizer step. If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP keeps the gradients
in the dtype used for gradient reduction, which can save memory if
using a custom optimizer that supports running in low precision.
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>cast_forward_inputs</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then this FSDP module casts
its forward args and kwargs to <code class="docutils literal notranslate"><span class="pre">param_dtype</span></code>. This is to ensure
that parameter and input dtypes match for forward computation, as
required by many ops. This may need to be set to <code class="docutils literal notranslate"><span class="pre">True</span></code> when only
applying mixed precision to some but not all FSDP modules, in which
case a mixed-precision FSDP submodule needs to recast its inputs.
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p></li>
<li><p><strong>cast_root_forward_inputs</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then the root FSDP module
casts its forward args and kwargs to <code class="docutils literal notranslate"><span class="pre">param_dtype</span></code>, overriding
the value of <code class="docutils literal notranslate"><span class="pre">cast_forward_inputs</span></code>. For non-root FSDP modules,
this does not do anything. (Default: <code class="docutils literal notranslate"><span class="pre">True</span></code>)</p></li>
<li><p><strong>_module_classes_to_ignore</strong> (<em>Sequence</em><em>[</em><em>Type</em><em>[</em><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.modules.module.Module"><em>torch.nn.modules.module.Module</em></a><em>]</em><em>]</em>) – (Sequence[Type[nn.Module]]): This specifies
module classes to ignore for mixed precision when using an
<code class="docutils literal notranslate"><span class="pre">auto_wrap_policy</span></code>: Modules of these classes will have FSDP
applied to them separately with mixed precision disabled (meaning
that the final FSDP construction would deviate from the specified
policy). If <code class="docutils literal notranslate"><span class="pre">auto_wrap_policy</span></code> is not specified, then this does
not do anything. This API is experimental and subject to change.
(Default: <code class="docutils literal notranslate"><span class="pre">(_BatchNorm,)</span></code>)</p></li>
</ul>
</dd>
</dl>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>This API is experimental and subject to change.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Only floating point tensors are cast to their specified dtypes.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>In <code class="docutils literal notranslate"><span class="pre">summon_full_params</span></code>, parameters are forced to full
precision, but buffers are not.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>Layer norm and batch norm accumulate in <code class="docutils literal notranslate"><span class="pre">float32</span></code> even when
their inputs are in a low precision like <code class="docutils literal notranslate"><span class="pre">float16</span></code> or <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>.
Disabling FSDP’s mixed precision for those norm modules only means that
the affine parameters are kept in <code class="docutils literal notranslate"><span class="pre">float32</span></code>. However, this incurs
separate all-gathers and reduce-scatters for those norm modules, which
may be inefficient, so if the workload permits, the user should prefer
to still apply mixed precision to those modules.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>By default, if the user passes a model with any <code class="docutils literal notranslate"><span class="pre">_BatchNorm</span></code>
modules and specifies an <code class="docutils literal notranslate"><span class="pre">auto_wrap_policy</span></code>, then the batch norm
modules will have FSDP applied to them separately with mixed precision
disabled. See the <code class="docutils literal notranslate"><span class="pre">_module_classes_to_ignore</span></code> argument.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p><code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code> has <code class="docutils literal notranslate"><span class="pre">cast_root_forward_inputs=True</span></code> and
<code class="docutils literal notranslate"><span class="pre">cast_forward_inputs=False</span></code> by default. For the root FSDP instance,
its <code class="docutils literal notranslate"><span class="pre">cast_root_forward_inputs</span></code> takes precedence over its
<code class="docutils literal notranslate"><span class="pre">cast_forward_inputs</span></code>. For non-root FSDP instances, their
<code class="docutils literal notranslate"><span class="pre">cast_root_forward_inputs</span></code> values are ignored. The default setting is
sufficient for the typical case where each FSDP instance has the same
<code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code> configuration and only needs to cast inputs to the
<code class="docutils literal notranslate"><span class="pre">param_dtype</span></code> at the beginning of the model’s forward pass.</p>
</div>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>For nested FSDP instances with different <code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code>
configurations, we recommend setting individual <code class="docutils literal notranslate"><span class="pre">cast_forward_inputs</span></code>
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 <code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code> configuration.</p>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">model</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">))</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">mixed_precision</span><span class="o">=</span><span class="n">MixedPrecision</span><span class="p">(</span><span class="n">param_dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">cast_forward_inputs</span><span class="o">=</span><span class="kc">True</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="p">,</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">mixed_precision</span><span class="o">=</span><span class="n">MixedPrecision</span><span class="p">(</span><span class="n">param_dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">,</span> <span class="n">cast_forward_inputs</span><span class="o">=</span><span class="kc">True</span><span class="p">),</span>
<span class="gp">&gt;&gt;&gt; </span><span class="p">)</span>
</pre></div>
</div>
<p>The above shows a working example. On the other hand, if <code class="docutils literal notranslate"><span class="pre">model[1]</span></code>
were replaced with <code class="docutils literal notranslate"><span class="pre">model[0]</span></code>, meaning that the submodule using
different <code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code> ran its forward first, then <code class="docutils literal notranslate"><span class="pre">model[1]</span></code>
would incorrectly see <code class="docutils literal notranslate"><span class="pre">float16</span></code> activations instead of <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>
ones.</p>
</div>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.CPUOffload">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">CPUOffload</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_params</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#CPUOffload"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.CPUOffload" title="Permalink to this definition">¶</a></dt>
<dd><p>This configures CPU offloading.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>offload_params</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – This specifies whether to offload parameters to
CPU when not involved in computation. If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then this
offloads gradients to CPU as well, meaning that the optimizer step
runs on CPU.</p>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.StateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">StateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#StateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.StateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><p><code class="docutils literal notranslate"><span class="pre">StateDictConfig</span></code> is the base class for all <code class="docutils literal notranslate"><span class="pre">state_dict</span></code> configuration
classes. Users should instantiate a child class (e.g.
<code class="docutils literal notranslate"><span class="pre">FullStateDictConfig</span></code>) in order to configure settings for the
corresponding <code class="docutils literal notranslate"><span class="pre">state_dict</span></code> type supported by FSDP.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>offload_to_cpu</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP offloads the state dict
values to CPU, and if <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP keeps them on GPU.
(Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">FullStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">rank0_only</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#FullStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><p><code class="docutils literal notranslate"><span class="pre">FullStateDictConfig</span></code> is a config class meant to be used with
<code class="docutils literal notranslate"><span class="pre">StateDictType.FULL_STATE_DICT</span></code>. We recommend enabling both
<code class="docutils literal notranslate"><span class="pre">offload_to_cpu=True</span></code> and <code class="docutils literal notranslate"><span class="pre">rank0_only=True</span></code> when saving full state
dicts to save GPU memory and CPU memory, respectively. This config class
is meant to be used via the <code class="xref py py-func docutils literal notranslate"><span class="pre">state_dict_type()</span></code> context manager as
follows:</p>
<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">torch.distributed.fsdp</span> <span class="kn">import</span> <span class="n">FullyShardedDataParallel</span> <span class="k">as</span> <span class="n">FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">fsdp</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">auto_wrap_policy</span><span class="o">=...</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">cfg</span> <span class="o">=</span> <span class="n">FullStateDictConfig</span><span class="p">(</span><span class="n">offload_to_cpu</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">rank0_only</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">with</span> <span class="n">FSDP</span><span class="o">.</span><span class="n">state_dict_type</span><span class="p">(</span><span class="n">fsdp</span><span class="p">,</span> <span class="n">StateDictType</span><span class="o">.</span><span class="n">FULL_STATE_DICT</span><span class="p">,</span> <span class="n">cfg</span><span class="p">):</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">state</span> <span class="o">=</span> <span class="n">fsdp</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="c1"># `state` will be empty on non rank 0 and contain CPU tensors on rank 0.</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># To reload checkpoint for inference, finetuning, transfer learning, etc:</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">model</span> <span class="o">=</span> <span class="n">model_fn</span><span class="p">()</span> <span class="c1"># Initialize model in preparation for wrapping with FSDP</span>
<span class="gp">&gt;&gt;&gt; </span><span class="k">if</span> <span class="n">dist</span><span class="o">.</span><span class="n">get_rank</span><span class="p">()</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="c1"># Load checkpoint only on rank 0 to avoid memory redundancy</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s2">&quot;my_checkpoint.pt&quot;</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span>    <span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># All ranks initialize FSDP module as usual. `sync_module_states` argument</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># communicates loaded checkpoint states from rank 0 to rest of the world.</span>
<span class="gp">&gt;&gt;&gt; </span><span class="n">fsdp</span> <span class="o">=</span> <span class="n">FSDP</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">device_id</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">current_device</span><span class="p">(),</span> <span class="n">auto_wrap_policy</span><span class="o">=...</span><span class="p">,</span> <span class="n">sync_module_states</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="gp">&gt;&gt;&gt; </span><span class="c1"># After this point, all ranks have FSDP model with loaded checkpoint.</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>rank0_only</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then only rank 0 saves the full state
dict, and nonzero ranks save an empty dict. If <code class="docutils literal notranslate"><span class="pre">False</span></code>, then all
ranks save the full state dict. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.ShardedStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">ShardedStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">_use_dtensor</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#ShardedStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.ShardedStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><p><code class="docutils literal notranslate"><span class="pre">ShardedStateDictConfig</span></code> is a config class meant to be used with
<code class="docutils literal notranslate"><span class="pre">StateDictType.SHARDED_STATE_DICT</span></code>.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>_use_dtensor</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP saves the state dict values
as <code class="docutils literal notranslate"><span class="pre">DTensor</span></code>, and if <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP saves them as
<code class="docutils literal notranslate"><span class="pre">ShardedTensor</span></code>. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p>
</dd>
</dl>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><code class="docutils literal notranslate"><span class="pre">_use_dtensor</span></code> is a private field of <a class="reference internal" href="#torch.distributed.fsdp.ShardedStateDictConfig" title="torch.distributed.fsdp.ShardedStateDictConfig"><code class="xref py py-class docutils literal notranslate"><span class="pre">ShardedStateDictConfig</span></code></a>
and it is used by FSDP to determine the type of state dict values. Users should not
manually modify <code class="docutils literal notranslate"><span class="pre">_use_dtensor</span></code>.</p>
</div>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.LocalStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">LocalStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><span class="pre">bool</span></a></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#LocalStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.LocalStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.OptimStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">OptimStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#OptimStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.OptimStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><p><code class="docutils literal notranslate"><span class="pre">OptimStateDictConfig</span></code> is the base class for all <code class="docutils literal notranslate"><span class="pre">optim_state_dict</span></code>
configuration classes.  Users should instantiate a child class (e.g.
<code class="docutils literal notranslate"><span class="pre">FullOptimStateDictConfig</span></code>) in order to configure settings for the
corresponding <code class="docutils literal notranslate"><span class="pre">optim_state_dict</span></code> type supported by FSDP.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>offload_to_cpu</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP offloads the state dict’s
tensor values to CPU, and if <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP keeps them on the
original device (which is GPU unless parameter CPU offloading is
enabled). (Default: <code class="docutils literal notranslate"><span class="pre">True</span></code>)</p>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.FullOptimStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">FullOptimStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">rank0_only</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#FullOptimStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.FullOptimStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>rank0_only</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then only rank 0 saves the full state
dict, and nonzero ranks save an empty dict. If <code class="docutils literal notranslate"><span class="pre">False</span></code>, then all
ranks save the full state dict. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.ShardedOptimStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">ShardedOptimStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">_use_dtensor</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#ShardedOptimStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.ShardedOptimStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><p><code class="docutils literal notranslate"><span class="pre">ShardedOptimStateDictConfig</span></code> is a config class meant to be used with
<code class="docutils literal notranslate"><span class="pre">StateDictType.SHARDED_STATE_DICT</span></code>.</p>
<dl class="field-list simple">
<dt class="field-odd">Variables</dt>
<dd class="field-odd"><p><strong>_use_dtensor</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><em>bool</em></a>) – If <code class="docutils literal notranslate"><span class="pre">True</span></code>, then FSDP saves the state dict values
as <code class="docutils literal notranslate"><span class="pre">DTensor</span></code>, and if <code class="docutils literal notranslate"><span class="pre">False</span></code>, then FSDP saves them as
<code class="docutils literal notranslate"><span class="pre">ShardedTensor</span></code>. (Default: <code class="docutils literal notranslate"><span class="pre">False</span></code>)</p>
</dd>
</dl>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><code class="docutils literal notranslate"><span class="pre">_use_dtensor</span></code> is a private field of <a class="reference internal" href="#torch.distributed.fsdp.ShardedOptimStateDictConfig" title="torch.distributed.fsdp.ShardedOptimStateDictConfig"><code class="xref py py-class docutils literal notranslate"><span class="pre">ShardedOptimStateDictConfig</span></code></a>
and it is used by FSDP to determine the type of state dict values. Users should not
manually modify <code class="docutils literal notranslate"><span class="pre">_use_dtensor</span></code>.</p>
</div>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.LocalOptimStateDictConfig">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">LocalOptimStateDictConfig</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">offload_to_cpu</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.12)"><span class="pre">bool</span></a></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#LocalOptimStateDictConfig"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.LocalOptimStateDictConfig" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.distributed.fsdp.StateDictSettings">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.distributed.fsdp.</span></span><span class="sig-name descname"><span class="pre">StateDictSettings</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">state_dict_type</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">torch.distributed.fsdp.api.StateDictType</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict_config</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="#torch.distributed.fsdp.StateDictConfig" title="torch.distributed.fsdp.api.StateDictConfig"><span class="pre">torch.distributed.fsdp.api.StateDictConfig</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">optim_state_dict_config</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="#torch.distributed.fsdp.OptimStateDictConfig" title="torch.distributed.fsdp.api.OptimStateDictConfig"><span class="pre">torch.distributed.fsdp.api.OptimStateDictConfig</span></a></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/distributed/fsdp/api.html#StateDictSettings"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.distributed.fsdp.StateDictSettings" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

</div>


             </article>
             
            </div>
            <footer>
  
    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="distributed.optim.html" class="btn btn-neutral float-right" title="Distributed Optimizers" accesskey="n" rel="next">Next <img src="_static/images/chevron-right-orange.svg" class="next-page"></a>
      
      
        <a href="elastic/kubernetes.html" class="btn btn-neutral" title="TorchElastic Kubernetes" accesskey="p" rel="prev"><img src="_static/images/chevron-right-orange.svg" class="previous-page"> Previous</a>
      
    </div>
  

    <hr>

  
  <div role="contentinfo">
    <p>
        &copy; Copyright 2023, PyTorch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
<script>

var match = window.location.href.match(/\/_[a-zA-Z0-9_]*.html|_dynamo/gi);
var url = window.location.href.lastIndexOf(match[match.length-1]);

if (url)
  {
    var div = '<div class="admonition note"><p class="admonition-title">Note</p><p><i class="fa fa-exclamation-circle" aria-hidden="true">&nbsp</i> This page describes an internal API which is not intended to be used outside of the PyTorch codebase and can be modified or removed without notice.</p></div>'
    document.getElementById("pytorch-article").insertAdjacentHTML('afterBegin', div)
  }
</script>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              <ul>
<li><a class="reference internal" href="#">FullyShardedDataParallel</a><ul>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel</span></code></a><ul>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.apply"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.apply()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.check_is_root"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.check_is_root()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.clip_grad_norm_"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.clip_grad_norm_()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.flatten_sharded_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.flatten_sharded_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.forward"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.forward()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.fsdp_modules"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.fsdp_modules()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.full_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.get_state_dict_type"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.get_state_dict_type()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.module"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.module</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.named_buffers"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.named_buffers()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.named_parameters"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.named_parameters()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.no_sync"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.no_sync()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.optim_state_dict_to_load"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.optim_state_dict_to_load()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.register_comm_hook"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.register_comm_hook()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.rekey_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.rekey_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.scatter_full_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.scatter_full_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.set_state_dict_type"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.set_state_dict_type()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.shard_full_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.shard_full_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.sharded_optim_state_dict"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.sharded_optim_state_dict()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.state_dict_type"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.state_dict_type()</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullyShardedDataParallel.summon_full_params"><code class="docutils literal notranslate"><span class="pre">FullyShardedDataParallel.summon_full_params()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.distributed.fsdp.BackwardPrefetch"><code class="docutils literal notranslate"><span class="pre">BackwardPrefetch</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.ShardingStrategy"><code class="docutils literal notranslate"><span class="pre">ShardingStrategy</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.MixedPrecision"><code class="docutils literal notranslate"><span class="pre">MixedPrecision</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.CPUOffload"><code class="docutils literal notranslate"><span class="pre">CPUOffload</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.StateDictConfig"><code class="docutils literal notranslate"><span class="pre">StateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullStateDictConfig"><code class="docutils literal notranslate"><span class="pre">FullStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.ShardedStateDictConfig"><code class="docutils literal notranslate"><span class="pre">ShardedStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.LocalStateDictConfig"><code class="docutils literal notranslate"><span class="pre">LocalStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.OptimStateDictConfig"><code class="docutils literal notranslate"><span class="pre">OptimStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.FullOptimStateDictConfig"><code class="docutils literal notranslate"><span class="pre">FullOptimStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.ShardedOptimStateDictConfig"><code class="docutils literal notranslate"><span class="pre">ShardedOptimStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.LocalOptimStateDictConfig"><code class="docutils literal notranslate"><span class="pre">LocalOptimStateDictConfig</span></code></a></li>
<li><a class="reference internal" href="#torch.distributed.fsdp.StateDictSettings"><code class="docutils literal notranslate"><span class="pre">StateDictSettings</span></code></a></li>
</ul>
</li>
</ul>

            </div>
          </div>
        </div>
      </section>
    </div>

  
       <script type="text/javascript" id="documentation_options" data-url_root="./" src="_static/documentation_options.js"></script>
         <script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
         <script src="_static/jquery.js"></script>
         <script src="_static/underscore.js"></script>
         <script src="_static/_sphinx_javascript_frameworks_compat.js"></script>
         <script src="_static/doctools.js"></script>
         <script src="_static/sphinx_highlight.js"></script>
         <script src="_static/clipboard.min.js"></script>
         <script src="_static/copybutton.js"></script>
     

  <script type="text/javascript" src="_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script script type="text/javascript">
  var collapsedSections = ['Developer Notes', 'Language Bindings', 'Libraries', 'Community'];
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">Stay up to date</li>
            <li><a href="https://www.facebook.com/pytorch" target="_blank">Facebook</a></li>
            <li><a href="https://twitter.com/pytorch" target="_blank">Twitter</a></li>
            <li><a href="https://www.youtube.com/pytorch" target="_blank">YouTube</a></li>
            <li><a href="https://www.linkedin.com/company/pytorch" target="_blank">LinkedIn</a></li>
          </ul>  
          </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">PyTorch Podcasts</li>
            <li><a href="https://open.spotify.com/show/6UzHKeiy368jKfQMKKvJY5" target="_blank">Spotify</a></li>
            <li><a href="https://podcasts.apple.com/us/podcast/pytorch-developer-podcast/id1566080008" target="_blank">Apple</a></li>
            <li><a href="https://www.google.com/podcasts?feed=aHR0cHM6Ly9mZWVkcy5zaW1wbGVjYXN0LmNvbS9PQjVGa0lsOA%3D%3D" target="_blank">Google</a></li>
            <li><a href="https://music.amazon.com/podcasts/7a4e6f0e-26c2-49e9-a478-41bd244197d0/PyTorch-Developer-Podcast?" target="_blank">Amazon</a></li>
          </ul>
         </div>
        </div>
        
        <div class="privacy-policy">
          <ul>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/terms/" target="_blank">Terms</a></li>
            <li class="privacy-policy-links">|</li>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/privacy-policy/" target="_blank">Privacy</a></li>
          </ul>
        </div>
        <div class="copyright">
        <p>© Copyright The Linux Foundation. The PyTorch Foundation is a project of The Linux Foundation.
          For web site terms of use, trademark policy and other policies applicable to The PyTorch Foundation please see
          <a href="https://www.linuxfoundation.org/policies/">www.linuxfoundation.org/policies/</a>. The PyTorch Foundation supports the PyTorch open source
          project, which has been established as PyTorch Project a Series of LF Projects, LLC. For policies applicable to the PyTorch Project a Series of LF Projects, LLC,
          please see <a href="https://www.lfprojects.org/policies/">www.lfprojects.org/policies/</a>.</p>
      </div>
     </div>

  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
           <li class="resources-mobile-menu-title">
             <a>Learn</a>
           </li>
           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/get-started">Get Started</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials">Tutorials</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/beginner/basics/intro.html">Learn the Basics</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/recipes/recipes_index.html">PyTorch Recipes</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/beginner/introyt.html">Introduction to PyTorch - YouTube Series</a>
             </li>
           </ul>
           <li class="resources-mobile-menu-title">
             <a>Ecosystem</a>
           </li>
           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/ecosystem">Tools</a>
             </li>
             <li>
               <a href="https://pytorch.org/#community-module">Community</a>
             </li>
             <li>
               <a href="https://discuss.pytorch.org/">Forums</a>
             </li>
             <li>
               <a href="https://pytorch.org/resources">Developer Resources</a>
             </li>
             <li>
               <a href="https://pytorch.org/ecosystem/contributor-awards-2023">Contributor Awards - 2023</a>
             </li>
           </ul>

           <li class="resources-mobile-menu-title">
             <a>Edge</a>
           </li>

           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/edge">About PyTorch Edge</a>
             </li>
             
             <li>
               <a href="https://pytorch.org/executorch-overview">ExecuTorch</a>
             </li>
           </ul>

           <li class="resources-mobile-menu-title">
             <a>Docs</a>
           </li>

           <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/docs/stable/index.html">PyTorch</a>
            </li>

            <li>
              <a href="https://pytorch.org/pytorch-domains">PyTorch Domains</a>
            </li>
          </ul>

          <li class="resources-mobile-menu-title">
            <a>Blog & News</a>
          </li>
            
           <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/blog/">PyTorch Blog</a>
            </li>
            <li>
              <a href="https://pytorch.org/community-blog">Community Blog</a>
            </li>

            <li>
              <a href="https://pytorch.org/videos">Videos</a>
            </li>

            <li>
              <a href="https://pytorch.org/community-stories">Community Stories</a>
            </li>
            <li>
              <a href="https://pytorch.org/events">Events</a>
            </li>
          </ul>
          
          <li class="resources-mobile-menu-title">
            <a>About</a>
          </li>

          <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/foundation">PyTorch Foundation</a>
            </li>
            <li>
              <a href="https://pytorch.org/governing-board">Governing Board</a>
            </li>
          </ul>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>