Merge branch 'main' into patch-1

Author: svekars

.jenkins/build.sh

@@ -24,10 +24,10 @@ pip install --progress-bar off -r $DIR/../requirements.txt
 
 #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
-pip uninstall -y torch torchvision torchaudio torchtext torchdata
-pip3 install torch==2.2.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu121
-pip3 install torchdata torchtext --index-url https://download.pytorch.org/whl/test/cpu
+# 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.2.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu121
+# pip3 install torchdata torchtext --index-url https://download.pytorch.org/whl/test/cpu
 
 # Install two language tokenizers for Translation with TorchText tutorial
 python -m spacy download en_core_web_sm

advanced_source/coding_ddpg.py

@@ -63,16 +63,25 @@
 #      %%bash
 #      pip3 install torchrl mujoco glfw
 
-import torchrl
-import torch
-import tqdm
-from typing import Tuple
-
 # sphinx_gallery_start_ignore
 import warnings
 warnings.filterwarnings("ignore")
+import multiprocessing
+# TorchRL prefers spawn method, that restricts creation of  ``~torchrl.envs.ParallelEnv`` inside
+# `__main__` method call, but for the easy of reading the code switch to fork
+# which is also a default spawn method in Google's Colaboratory
+try:
+    multiprocessing.set_start_method("fork")
+except RuntimeError:
+    assert multiprocessing.get_start_method() == "fork"
 # sphinx_gallery_end_ignore
 
+
+import torchrl
+import torch
+import tqdm
+from typing import Tuple
+
 ###############################################################################
 # We will execute the policy on CUDA if available
 device =  torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -1219,6 +1228,6 @@ def ceil_div(x, y):
 # 
 # To iterate further on this loss module we might consider:
 # 
-# - Using `@dispatch` (see `[Feature] Distpatch IQL loss module <https://github.com/pytorch/rl/pull/1230>`_.
+# - Using `@dispatch` (see `[Feature] Distpatch IQL loss module <https://github.com/pytorch/rl/pull/1230>`_.)
 # - Allowing flexible TensorDict keys.
 # 

index.rst

@@ -3,15 +3,11 @@ Welcome to PyTorch Tutorials
 
 What's new in PyTorch tutorials?
 
-* `Getting Started with Distributed Checkpoint (DCP) <https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html>`__
-* `torch.export Tutorial <https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html>`__
-* `Facilitating New Backend Integration by PrivateUse1 <https://pytorch.org/tutorials/advanced/privateuseone.html>`__
-* `(prototype) Accelerating BERT with semi-structured (2:4) sparsity <https://pytorch.org/tutorials/prototype/semi_structured_sparse.html>`__
-* `(prototype) PyTorch 2 Export Quantization-Aware Training (QAT) <https://pytorch.org/tutorials/prototype/pt2e_quant_qat.html>`__
-* `(prototype) PyTorch 2 Export Post Training Quantization with X86 Backend through Inductor <https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_x86_inductor.html>`__
-* `(prototype) Inductor C++ Wrapper Tutorial <https://pytorch.org/tutorials/prototype/inductor_cpp_wrapper_tutorial.html>`__
-* `How to save memory by fusing the optimizer step into the backward pass <https://pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html>`__
-* `Tips for Loading an nn.Module from a Checkpoint <https://pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html>`__
+* `PyTorch Inference Performance Tuning on AWS Graviton Processors <https://pytorch.org/tutorials/recipes/inference_tuning_on_aws_graviton.html>`__
+* `Using TORCH_LOGS python API with torch.compile <https://pytorch.org/tutorials/recipes/torch_logs.html>`__
+* `PyTorch 2 Export Quantization with X86 Backend through Inductor <https://pytorch.org/tutorials/prototype/pt2e_quant_x86_inductor.html>`__
+* `Getting Started with DeviceMesh <https://pytorch.org/tutorials/recipes/distributed_device_mesh.html>`__
+* `Compiling the optimizer with torch.compile <https://pytorch.org/tutorials/recipes/compiling_optimizer.html>`__
 
 
 .. raw:: html

recipes_source/compiling_optimizer.rst

@@ -86,3 +86,9 @@ Sample Results:
 
 * Eager runtime: 747.2437149845064us
 * Compiled runtime: 392.07384741178us
+
+See Also
+~~~~~~~~~
+
+* For an in-depth technical overview, see
+`Compiling the optimizer with PT2 <https://dev-discuss.pytorch.org/t/compiling-the-optimizer-with-pt2/1669>`__

recipes_source/torch_logs.py

@@ -8,7 +8,7 @@
 
 ######################################################################
 #
-# This tutorial introduces the ``TORCH_LOGS`` environment variable, as well ass the Python API, and
+# This tutorial introduces the ``TORCH_LOGS`` environment variable, as well as the Python API, and
 # demonstrates how to apply it to observe the phases  of ``torch.compile``.
 #
 # .. note::
@@ -34,53 +34,49 @@
 
 # exit cleanly if we are on a device that doesn't support torch.compile
 if torch.cuda.get_device_capability() < (7, 0):
-    print("Exiting because torch.compile is not supported on this device.")
-    import sys
+    print("Skipping because torch.compile is not supported on this device.")
+else:
+    @torch.compile()
+    def fn(x, y):
+        z = x + y
+        return z + 2
 
-    sys.exit(0)
 
-
-@torch.compile()
-def fn(x, y):
-    z = x + y
-    return z + 2
-
-
-inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
+    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
 
 
 # print separator and reset dynamo
 # between each example
-def separator(name):
-    print(f"==================={name}=========================")
-    torch._dynamo.reset()
+    def separator(name):
+        print(f"==================={name}=========================")
+        torch._dynamo.reset()
 
 
-separator("Dynamo Tracing")
+    separator("Dynamo Tracing")
 # View dynamo tracing
 # TORCH_LOGS="+dynamo"
-torch._logging.set_logs(dynamo=logging.DEBUG)
-fn(*inputs)
+    torch._logging.set_logs(dynamo=logging.DEBUG)
+    fn(*inputs)
 
-separator("Traced Graph")
+    separator("Traced Graph")
 # View traced graph
 # TORCH_LOGS="graph"
-torch._logging.set_logs(graph=True)
-fn(*inputs)
+    torch._logging.set_logs(graph=True)
+    fn(*inputs)
 
-separator("Fusion Decisions")
+    separator("Fusion Decisions")
 # View fusion decisions
 # TORCH_LOGS="fusion"
-torch._logging.set_logs(fusion=True)
-fn(*inputs)
+    torch._logging.set_logs(fusion=True)
+    fn(*inputs)
 
-separator("Output Code")
+    separator("Output Code")
 # View output code generated by inductor
 # TORCH_LOGS="output_code"
-torch._logging.set_logs(output_code=True)
-fn(*inputs)
+    torch._logging.set_logs(output_code=True)
+    fn(*inputs)
 
-separator("")
+    separator("")
 
 ######################################################################
 # Conclusion