Merge branch 'main' into fsdp-tutorial-update

Author: surprisedPikachu007

.jenkins/build.sh

@@ -21,9 +21,9 @@ sudo apt-get install -y pandoc
 
 #Install PyTorch Nightly for test.
 # Nightly - pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cu102/torch_nightly.html
-# Install 2.2 for testing - uncomment to install nightly binaries (update the version as needed).
-# pip uninstall -y torch torchvision torchaudio torchtext torchdata
-# pip3 install torch==2.3.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu121
+# Install 2.4 to merge all 2.4 PRs - uncomment to install nightly binaries (update the version as needed).
+pip uninstall -y torch torchvision torchaudio torchtext torchdata
+pip3 install torch==2.4.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu124
 
 # Install two language tokenizers for Translation with TorchText tutorial
 python -m spacy download en_core_web_sm

.jenkins/validate_tutorials_built.py

@@ -29,7 +29,6 @@
     "intermediate_source/fx_conv_bn_fuser",
     "intermediate_source/_torch_export_nightly_tutorial",  # does not work on release
     "advanced_source/super_resolution_with_onnxruntime",
-    "advanced_source/python_custom_ops",  # https://github.com/pytorch/pytorch/issues/127443
     "advanced_source/usb_semisup_learn", # fails with CUDA OOM error, should try on a different worker
     "prototype_source/fx_graph_mode_ptq_dynamic",
     "prototype_source/vmap_recipe",
@@ -54,8 +53,6 @@
     "intermediate_source/flask_rest_api_tutorial",
     "intermediate_source/text_to_speech_with_torchaudio",
     "intermediate_source/tensorboard_profiler_tutorial", # reenable after 2.0 release.
-    "intermediate_source/inductor_debug_cpu", # reenable after 2942 
-    "beginner_source/onnx/onnx_registry_tutorial", # reenable after 2941 is fixed.
     "intermediate_source/torch_export_tutorial" # reenable after 2940 is fixed.
 ]
 

_static/img/distributed/tcpstore_barrier_time.png

@@ -417,4 +417,4 @@ Conclusion
 In this tutorial, we went over the recommended approach to integrating Custom C++
 and CUDA operators with PyTorch. The ``TORCH_LIBRARY/torch.library`` APIs are fairly
 low-level. For more information about how to use the API, see
-`The Custom Operators Manual <https://pytorch.org/docs/main/notes/custom_operators.html>`_.
+`The Custom Operators Manual <https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html#the-custom-operators-manual>`_.

_static/img/distributed/tcpstore_init_time.png

@@ -2,6 +2,10 @@ Custom C++ and CUDA Extensions
 ==============================
 **Author**: `Peter Goldsborough <https://www.goldsborough.me/>`_
 
+.. warning::
+
+    This tutorial is deprecated as of PyTorch 2.4. Please see :ref:`custom-ops-landing-page`
+    for the newest up-to-date guides on extending PyTorch with Custom C++/CUDA Extensions.
 
 PyTorch provides a plethora of operations related to neural networks, arbitrary
 tensor algebra, data wrangling and other purposes. However, you may still find
@@ -225,7 +229,7 @@ Instead of:
 Currently open issue for nvcc bug `here
 <https://github.com/pytorch/pytorch/issues/69460>`_.
 Complete workaround code example `here
-<https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48>`_. 
+<https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48>`_.
 
 Forward Pass
 ************

_static/img/onnx/custom_aten_add_function.png

@@ -1,7 +1,7 @@
 .. _custom-ops-landing-page:
 
-PyTorch Custom Operators Landing Page
-=====================================
+PyTorch Custom Operators
+===========================
 
 PyTorch offers a large library of operators that work on Tensors (e.g. ``torch.add``,
 ``torch.sum``, etc). However, you may wish to bring a new custom operation to PyTorch
@@ -10,26 +10,27 @@ In order to do so, you must register the custom operation with PyTorch via the P
 `torch.library docs <https://pytorch.org/docs/stable/library.html>`_ or C++ ``TORCH_LIBRARY``
 APIs.
 
-TL;DR
------
+
 
 Authoring a custom operator from Python
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 Please see :ref:`python-custom-ops-tutorial`.
 
 You may wish to author a custom operator from Python (as opposed to C++) if:
+
 - you have a Python function you want PyTorch to treat as an opaque callable, especially with
-respect to ``torch.compile`` and ``torch.export``.
+  respect to ``torch.compile`` and ``torch.export``.
 - you have some Python bindings to C++/CUDA kernels and want those to compose with PyTorch
-subsystems (like ``torch.compile`` or ``torch.autograd``)
+  subsystems (like ``torch.compile`` or ``torch.autograd``)
 
 Integrating custom C++ and/or CUDA code with PyTorch
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 Please see :ref:`cpp-custom-ops-tutorial`.
 
 You may wish to author a custom operator from C++ (as opposed to Python) if:
+
 - you have custom C++ and/or CUDA code.
 - you plan to use this code with ``AOTInductor`` to do Python-less inference.
 

_static/img/onnx/custom_aten_gelu_function.png

@@ -1,6 +1,11 @@
 Registering a Dispatched Operator in C++
 ========================================
 
+.. warning::
+
+    This tutorial is deprecated as of PyTorch 2.4. Please see :ref:`custom-ops-landing-page`
+    for the newest up-to-date guides on extending PyTorch with Custom Operators.
+
 The dispatcher is an internal component of PyTorch which is responsible for
 figuring out what code should actually get run when you call a function like
 ``torch::add``.  This can be nontrivial, because PyTorch operations need

_static/img/onnx/custom_aten_gelu_model.png

@@ -260,5 +260,5 @@ def f(x):
 # For more detailed information, see:
 #
 # - `the torch.library documentation <https://pytorch.org/docs/stable/library.html>`_
-# - `the Custom Operators Manual <https://pytorch.org/docs/main/notes/custom_operators.html>`_
+# - `the Custom Operators Manual <https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html#the-custom-operators-manual>`_
 #

advanced_source/cpp_custom_ops.rst

@@ -1,6 +1,11 @@
 Extending TorchScript with Custom C++ Operators
 ===============================================
 
+.. warning::
+
+    This tutorial is deprecated as of PyTorch 2.4. Please see :ref:`custom-ops-landing-page`
+    for the newest up-to-date guides on PyTorch Custom Operators.
+
 The PyTorch 1.0 release introduced a new programming model to PyTorch called
 `TorchScript <https://pytorch.org/docs/master/jit.html>`_. TorchScript is a
 subset of the Python programming language which can be parsed, compiled and

advanced_source/cpp_extension.rst

@@ -99,7 +99,6 @@ def forward(self, input_x, input_y):
 # NOTE: All attributes must be annotated with type hints.
 @onnxscript.script(custom_aten)
 def custom_aten_add(input_x, input_y, alpha: float = 1.0):
-    alpha = opset18.CastLike(alpha, input_y)
     input_y = opset18.Mul(input_y, alpha)
     return opset18.Add(input_x, input_y)
 
@@ -130,9 +129,9 @@ def custom_aten_add(input_x, input_y, alpha: float = 1.0):
 # graph node name is the function name
 assert onnx_program.model_proto.graph.node[0].op_type == "custom_aten_add"
 # function node domain is empty because we use standard ONNX operators
-assert onnx_program.model_proto.functions[0].node[3].domain == ""
+assert {node.domain for node in onnx_program.model_proto.functions[0].node} == {""}
 # function node name is the standard ONNX operator name
-assert onnx_program.model_proto.functions[0].node[3].op_type == "Add"
+assert {node.op_type for node in onnx_program.model_proto.functions[0].node} == {"Add", "Mul", "Constant"}
 
 
 ######################################################################
@@ -231,33 +230,25 @@ def custom_aten_gelu(input_x, approximate: str = "none"):
 
 
 ######################################################################
-# Let's inspect the model and verify the model uses :func:`custom_aten_gelu` instead of
-# :class:`aten::gelu`. Note the graph has one graph nodes for
-# ``custom_aten_gelu``, and inside ``custom_aten_gelu``, there is a function
-# node for ``Gelu`` with namespace ``com.microsoft``.
+# Let's inspect the model and verify the model uses op_type ``Gelu``
+# from namespace ``com.microsoft``.
+#
+# .. note::
+#     :func:`custom_aten_gelu` does not exist in the graph because
+#     functions with fewer than three operators are inlined automatically.
 #
 
 # graph node domain is the custom domain we registered
 assert onnx_program.model_proto.graph.node[0].domain == "com.microsoft"
 # graph node name is the function name
-assert onnx_program.model_proto.graph.node[0].op_type == "custom_aten_gelu"
-# function node domain is the custom domain we registered
-assert onnx_program.model_proto.functions[0].node[0].domain == "com.microsoft"
-# function node name is the node name used in the function
-assert onnx_program.model_proto.functions[0].node[0].op_type == "Gelu"
+assert onnx_program.model_proto.graph.node[0].op_type == "Gelu"
 
 
 ######################################################################
-# The following diagram shows ``custom_aten_gelu_model`` ONNX graph using Netron:
+# The following diagram shows ``custom_aten_gelu_model`` ONNX graph using Netron,
+# we can see the ``Gelu`` node from module ``com.microsoft`` used in the function:
 #
 # .. image:: /_static/img/onnx/custom_aten_gelu_model.png
-#    :width: 70%
-#    :align: center
-#
-# Inside the ``custom_aten_gelu`` function, we can see the ``Gelu`` node from module
-# ``com.microsoft`` used in the function:
-#
-# .. image:: /_static/img/onnx/custom_aten_gelu_function.png
 #
 # That is all we need to do. As an additional step, we can use ONNX Runtime to run the model,
 # and compare the results with PyTorch.

advanced_source/custom_ops_landing_page.rst

@@ -3,11 +3,11 @@ Welcome to PyTorch Tutorials
 
 **What's new in PyTorch tutorials?**
 
-* `Using User-Defined Triton Kernels with torch.compile <https://pytorch.org/tutorials/recipes/torch_compile_user_defined_triton_kernel_tutorial.html>`__
-* `Large Scale Transformer model training with Tensor Parallel (TP) <https://pytorch.org/tutorials/intermediate/TP_tutorial.html>`__
-* `Accelerating BERT with semi-structured (2:4) sparsity <https://pytorch.org/tutorials/advanced/semi_structured_sparse.html>`__
-* `torch.export Tutorial with torch.export.Dim <https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html>`__
-* `Extension points in nn.Module for load_state_dict and tensor subclasses <https://pytorch.org/tutorials/recipes/recipes/swap_tensors.html>`__
+* `Introduction to Distributed Pipeline Parallelism <https://pytorch.org/tutorials/intermediate/pipelining_tutorial.html>`__
+* `Introduction to Libuv TCPStore Backend <https://pytorch.org/tutorials/intermediate/TCPStore_libuv_backend.html>`__ 
+* `Asynchronous Saving with Distributed Checkpoint (DCP) <https://pytorch.org/tutorials/recipes/distributed_async_checkpoint_recipe.html>`__
+* `Python Custom Operators <https://pytorch.org/tutorials/advanced/python_custom_ops.html>`__
+* Updated `Getting Started with DeviceMesh <https://pytorch.org/tutorials/recipes/distributed_device_mesh.html>`__
 
 .. raw:: html
 
@@ -779,6 +779,13 @@ Welcome to PyTorch Tutorials
    :link: intermediate/FSDP_adavnced_tutorial.html
    :tags: Parallel-and-Distributed-Training
 
+.. customcarditem::
+   :header: Introduction to Libuv TCPStore Backend
+   :card_description: TCPStore now uses a new server backend for faster connection and better scalability.
+   :image: _static/img/thumbnails/cropped/Introduction-to-Libuv-Backend-TCPStore.png
+   :link: intermediate/TCPStore_libuv_backend.html
+   :tags: Parallel-and-Distributed-Training
+
 .. Edge
 
 .. customcarditem::
@@ -1134,6 +1141,7 @@ Additional Resources
    intermediate/dist_tuto
    intermediate/FSDP_tutorial
    intermediate/FSDP_adavnced_tutorial
+   intermediate/TCPStore_libuv_backend
    intermediate/TP_tutorial
    intermediate/pipelining_tutorial
    intermediate/process_group_cpp_extension_tutorial