Merge branch 'site' into patch-1

Author: vincentqb

_posts/2024-02-06-pytorch-2-paper-tutorial.md

@@ -11,12 +11,12 @@ During the ASPLOS conference, we'll be conducting a tutorial on Saturday, April
 
 A preview of the paper is attached below:
 
-Title: **PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation.** [**Full Paper PDF**](/assets/pytorch_2.pdf)
+Title: **PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation.** [**Full Paper PDF**](/assets/pytorch2-2.pdf)
 
 ### Abstract
 This paper introduces two extensions to the popular PyTorch machine learning framework, TorchDynamo and TorchInductor, which implement the torch.compile feature released in PyTorch 2. TorchDynamo is a Python-level just-in-time (JIT) compiler that enables graph compilation in PyTorch programs without sacrificing the flexibility of Python. It achieves this by dynamically modifying Python bytecode before execution and extracting sequences of PyTorch operations into an FX graph, which is then JIT compiled using one of many extensible backends. TorchInductor is the default compiler backend for TorchDynamo, which translates PyTorch programs into OpenAI's Triton for GPUs and C++ for CPUs. Results show that TorchDynamo is able to capture graphs more robustly than prior approaches while adding minimal overhead, and TorchInductor is able to provide a 2.27x inference and 1.41x training geometric mean speedup on an NVIDIA A100 GPU across 180+ real-world models, which outperforms six other compilers. These extensions provide a new way to apply optimizations through compilers in eager mode frameworks like PyTorch.
 
 
 ### Authors
 
-Jason Ansel (Meta); Edward Yang (Meta); Horace He (Meta); Natalia Gimelshein (OpenAI); Animesh Jain (Meta); Michael Voznesensky (Meta); Bin Bao (Meta); David Berard (Meta); Geeta Chauhan (Meta); Anjali Chourdia (Meta); Will Constable (Meta); Alban Desmaison (Meta); Zachary DeVito (Meta); Elias Ellison (Meta); Will Feng (Meta); Jiong Gong (Intel); Michael Gschwind (Meta); Brian Hirsh (Meta); Sherlock Huang (Meta); Laurent Kirsch (Meta); Michael Lazos (Meta); Yanbo Liang (Meta); Jason Liang (Meta); Yinghai Lu (Meta); CK Luk (Meta); Bert Maher (Meta); Yunjie Pan (University of Michigan); Christian Puhrsch (Meta); Matthias Reso (Meta); Mark Saroufim (Meta); Helen Suk (Meta); Michael Suo (Meta); Phil Tillet (OpenAI); Eikan Wang (Intel); Xiaodong Wang (Meta); William Wen (Meta); Shunting Zhang (Meta); Xu Zhao (Meta); Keren Zhou (OpenAI & George Mason University); Richard Zou (Meta); Ajit Mathews (Meta); Gregory Chanan (Meta); Peng Wu (Meta); Soumith Chintala (Meta)
+Jason Ansel (Meta); Edward Yang (Meta); Horace He (Meta); Natalia Gimelshein (OpenAI); Animesh Jain (Meta); Michael Voznesensky (Meta); Bin Bao (Meta); Peter Bell (Quansight); David Berard (Meta); Evgeni Burovski Quansight; Geeta Chauhan (Meta); Anjali Chourdia (Meta); Will Constable (Meta); Alban Desmaison (Meta); Zachary DeVito (Meta); Elias Ellison (Meta); Will Feng (Meta); Jiong Gong (Intel); Michael Gschwind (Meta); Brian Hirsh (Meta); Sherlock Huang (Meta); Kshiteej Kalambarkar (Quansight); Laurent Kirsch (Meta); Michael Lazos (Meta); Mario Lezcano (Quansight); Yanbo Liang (Meta); Jason Liang (Meta); Yinghai Lu (Meta); CK Luk (Meta); Bert Maher (Meta); Yunjie Pan (University of Michigan); Christian Puhrsch (Meta); Matthias Reso (Meta); Mark Saroufim (Meta); Marcos Yukio Siraichi (Quansight); Helen Suk (Meta); Michael Suo (Meta); Phil Tillet (OpenAI); Eikan Wang (Intel); Xiaodong Wang (Meta); William Wen (Meta); Shunting Zhang (Meta); Xu Zhao (Meta); Keren Zhou (OpenAI & George Mason University); Richard Zou (Meta); Ajit Mathews (Meta); Gregory Chanan (Meta); Peng Wu (Meta); Soumith Chintala (Meta)

assets/pytorch_2.pdf renamed to assets/pytorch2-2.pdf

@@ -1126,7 +1126,13 @@ <h1>Source code for torch</h1><div class="highlight"><pre>
 
 
 <div class="viewcode-block" id="set_default_tensor_type"><a class="viewcode-back" href="../generated/torch.set_default_tensor_type.html#torch.set_default_tensor_type">[docs]</a><span class="k">def</span> <span class="nf">set_default_tensor_type</span><span class="p">(</span><span class="n">t</span><span class="p">):</span>
-<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Sets the default ``torch.Tensor`` type to floating point tensor type</span>
+<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    .. warning::</span>
+
+<span class="sd">        This function is deprecated as of PyTorch 2.1, please use :func:`torch.set_default_dtype()` and</span>
+<span class="sd">        :func:`torch.set_default_device()` as alternatives.</span>
+
+<span class="sd">    Sets the default ``torch.Tensor`` type to floating point tensor type</span>
 <span class="sd">    ``t``. This type will also be used as default floating point type for</span>
 <span class="sd">    type inference in :func:`torch.tensor`.</span>
 

assets/quick-start-module.js

@@ -621,7 +621,7 @@ <h1>Source code for torch.distributed.checkpoint.state_dict</h1><div class="high
             <span class="k">if</span> <span class="ow">not</span> <span class="n">skip_ddp_prefix</span><span class="p">:</span>
                 <span class="n">fqn_obj_names</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">curr_obj_name</span><span class="p">)</span>
         <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">curr_obj</span><span class="p">,</span> <span class="n">FSDP</span><span class="p">):</span>
-            <span class="k">if</span> <span class="n">obj_names</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="n">FLAT_PARAM</span><span class="p">:</span>
+            <span class="k">if</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">obj_names</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">obj_names</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="n">FLAT_PARAM</span><span class="p">:</span>
                 <span class="n">prefix</span> <span class="o">=</span> <span class="s2">&quot;.&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">fqn_obj_names</span><span class="p">)</span>
                 <span class="n">flat_param</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">curr_obj</span><span class="p">,</span> <span class="n">FLAT_PARAM</span><span class="p">)</span>
                 <span class="k">if</span> <span class="n">prefix</span><span class="p">:</span>
@@ -660,7 +660,7 @@ <h1>Source code for torch.distributed.checkpoint.state_dict</h1><div class="high
         <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">Union</span><span class="p">[</span><span class="n">Set</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span>
     <span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>
     <span class="n">all_fqns</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
-    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">():</span>
+    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">chain</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">(),</span> <span class="n">model</span><span class="o">.</span><span class="n">named_buffers</span><span class="p">()):</span>
         <span class="n">fqns</span> <span class="o">=</span> <span class="n">_get_fqns</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>
         <span class="n">fqn_param_mapping</span><span class="p">[</span><span class="n">param</span><span class="p">]</span> <span class="o">=</span> <span class="n">fqns</span>
         <span class="k">for</span> <span class="n">fqn</span> <span class="ow">in</span> <span class="n">fqns</span><span class="p">:</span>
@@ -859,7 +859,7 @@ <h1>Source code for torch.distributed.checkpoint.state_dict</h1><div class="high
     <span class="k">if</span> <span class="ow">not</span> <span class="n">info</span><span class="o">.</span><span class="n">handle_model</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">state_dict</span><span class="p">:</span>
         <span class="k">return</span> <span class="n">_IncompatibleKeys</span><span class="p">({},</span> <span class="p">{})</span>
 
-    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">():</span>
+    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">chain</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">(),</span> <span class="n">model</span><span class="o">.</span><span class="n">named_buffers</span><span class="p">()):</span>
         <span class="n">fqns</span> <span class="o">=</span> <span class="n">_get_fqns</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span>
         <span class="n">fqns_with_ddp_prefix</span> <span class="o">=</span> <span class="n">_get_fqns</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">skip_ddp_prefix</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
         <span class="k">for</span> <span class="n">fqn</span><span class="p">,</span> <span class="n">fqn_with_ddp_prefix</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">fqns</span><span class="p">,</span> <span class="n">fqns_with_ddp_prefix</span><span class="p">):</span>
@@ -1142,25 +1142,25 @@ <h1>Source code for torch.distributed.checkpoint.state_dict</h1><div class="high
 <span class="sd">    optimizer parameter IDs to the canonical FQNs.</span>
 
 <span class="sd">    Example:</span>
+<span class="sd">        &gt;&gt;&gt; # xdoctest: +SKIP</span>
+<span class="sd">        &gt;&gt;&gt; import torch</span>
+<span class="sd">        &gt;&gt;&gt; from torch.distributed.fsdp import FullyShardedDataParallel as FSDP</span>
+<span class="sd">        &gt;&gt;&gt; from torch.nn.parallel import DistributedDataParallel as DDP</span>
+<span class="sd">        &gt;&gt;&gt; from torch.distributed.checkpoint.state_dict import get_state_dict</span>
 
-<span class="sd">        import torch</span>
-<span class="sd">        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP</span>
-<span class="sd">        from torch.nn.parallel import DistributedDataParallel as DDP</span>
-<span class="sd">        from torch.distributed.checkpoint.state_dict import get_state_dict</span>
-
-<span class="sd">        fsdp_model = FSDP(copy.deepcopy(model))</span>
-<span class="sd">        fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)</span>
-<span class="sd">        ddp_model = DDP(copy.deepcopy(model))</span>
-<span class="sd">        ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)</span>
+<span class="sd">        &gt;&gt;&gt; fsdp_model = FSDP(copy.deepcopy(model))</span>
+<span class="sd">        &gt;&gt;&gt; fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)</span>
+<span class="sd">        &gt;&gt;&gt; ddp_model = DDP(copy.deepcopy(model))</span>
+<span class="sd">        &gt;&gt;&gt; ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)</span>
 
 
-<span class="sd">        ddp_state_dict, ddp_optim_state_dict = get_state_dict(ddp_model, ddp_optim)</span>
-<span class="sd">        fsdp_state_dict, fsdp_optim_state_dict = get_state_dict(fsdp_model, fsdp_optim)</span>
+<span class="sd">        &gt;&gt;&gt; ddp_state_dict, ddp_optim_state_dict = get_state_dict(ddp_model, ddp_optim)</span>
+<span class="sd">        &gt;&gt;&gt; fsdp_state_dict, fsdp_optim_state_dict = get_state_dict(fsdp_model, fsdp_optim)</span>
 
-<span class="sd">        # if we simply call ddp_model.state_dict() and fsdp_model.state_dict(),</span>
-<span class="sd">        # the asserts will fail.</span>
-<span class="sd">        assert ddp_state_dict == fsdp_state_dict</span>
-<span class="sd">        assert ddp_optim_state == fsdp_optim_state_dict</span>
+<span class="sd">        &gt;&gt;&gt; # if we simply call ddp_model.state_dict() and fsdp_model.state_dict(),</span>
+<span class="sd">        &gt;&gt;&gt; # the asserts will fail.</span>
+<span class="sd">        &gt;&gt;&gt; assert ddp_state_dict == fsdp_state_dict</span>
+<span class="sd">        &gt;&gt;&gt; assert ddp_optim_state == fsdp_optim_state_dict</span>
 
 
 <span class="sd">    Args:</span>
@@ -1175,6 +1175,8 @@ <h1>Source code for torch.distributed.checkpoint.state_dict</h1><div class="high
 
 <span class="sd">    Returns:</span>
 <span class="sd">        ``Tuple`` that contain model state_dict and optimizer state_dict.</span>
+
+<span class="sd">    :rtype: typing.Tuple[typing.Dict[str, ValueType], OptimizerStateType]</span>
 <span class="sd">    &quot;&quot;&quot;</span>
 
     <span class="k">with</span> <span class="n">gc_context</span><span class="p">():</span>

docs/2.2/_images/RReLU.png

@@ -932,7 +932,7 @@ <h1>Source code for torch.distributed.fsdp.fully_sharded_data_parallel</h1><div
                 <span class="s2">&quot;ignored_states&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">_ignored_params</span><span class="p">,</span>
                 <span class="s2">&quot;device_mesh&quot;</span><span class="p">:</span> <span class="n">device_mesh</span><span class="p">,</span>
             <span class="p">}</span>
-            <span class="k">if</span> <span class="n">sharding_strategy</span> <span class="ow">in</span> <span class="n">HYBRID_SHARDING_STRATEGIES</span><span class="p">:</span>
+            <span class="k">if</span> <span class="n">sharding_strategy</span> <span class="ow">in</span> <span class="n">HYBRID_SHARDING_STRATEGIES</span> <span class="ow">and</span> <span class="n">device_mesh</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                 <span class="c1"># Share root process groups with children to maintain</span>
                 <span class="c1"># the invariant that all FSDP modules will have the same</span>
                 <span class="c1"># process groups.</span>

docs/2.2/_modules/torch.html

@@ -495,7 +495,7 @@ <h1>Source code for torch.distributed.tensor.parallel.style</h1><div class="high
 
 <div class="viewcode-block" id="ColwiseParallel"><a class="viewcode-back" href="../../../../../distributed.tensor.parallel.html#torch.distributed.tensor.parallel.ColwiseParallel">[docs]</a><span class="k">class</span> <span class="nc">ColwiseParallel</span><span class="p">(</span><span class="n">ParallelStyle</span><span class="p">):</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
-<span class="sd">    Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.</span>
+<span class="sd">    Partition a compatible nn.Module in a column-wise fashion. Currently supports nn.Linear and nn.Embedding.</span>
 <span class="sd">    Users can compose it together with RowwiseParallel to achieve the sharding of more complicated modules.</span>
 <span class="sd">    (i.e. MLP, Attention)</span>