PyTorch

https://docs.pytorch.org/docs/2.11/
Admin
PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.
Tokens:2,101,080
Snippets:17,408
Trust Score:10
Update:2 months ago
Show doc for...
Tokens:
Raw
### Initializing Parallel Models and Optimizers

Source: https://docs.pytorch.org/docs/2.11/distributed.checkpoint.html

Setup examples for FSDP and DDP models with corresponding Adam optimizers.

```python
>>> fsdp_model = FSDP(copy.deepcopy(model))
>>> fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)
>>> ddp_model = DDP(copy.deepcopy(model))
>>> ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)
```

--------------------------------

### Example of Guard Installation for Constant Input

Source: https://docs.pytorch.org/docs/2.11/_sources/user_guide/torch_compiler/torch.compiler_dynamo_deepdive.md.txt

This snippet demonstrates how a constant string input installs guards to ensure the input type and value match expectations. Running this with TORCH_LOGS=guards will show the generated checks.

```python
import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")
```

--------------------------------

### Setting up Copy Engine Collectives

Source: https://docs.pytorch.org/docs/2.11/_sources/symmetric_memory.md.txt

This example demonstrates the full setup for using Copy Engine (CE) Collectives, including initializing the NCCL process group with the zero-CTA policy, setting the symmetric memory backend, allocating and registering tensors, and performing a collective operation.

```python
import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem

# Initialize process group with zero-CTA policy for CE collectives
opts = dist.ProcessGroupNCCL.Options()
opts.config.cta_policy = dist.ProcessGroupNCCL.NCCL_CTA_POLICY_ZERO
device = torch.device("cuda", rank)
dist.init_process_group(backend="nccl", pg_options=opts, device_id=device)

# Set up symmetric memory with NCCL backend
symm_mem.set_backend("NCCL")
group_name = dist.group.WORLD.group_name

# Allocate tensors using symmetric memory
numel = 1024 * 1024
inp = symm_mem.empty(numel, device=device)
out = symm_mem.empty(numel * world_size, device=device)

# Register tensors for symmetric memory operations
symm_mem.rendezvous(inp, group=group_name)
symm_mem.rendezvous(out, group=group_name)

# Perform collective operation using copy engines
# This now runs on DMA engines instead of SMs
work = dist.all_gather_into_tensor(out, inp, async_op=True)
work.wait()
```

--------------------------------

### TorchDynamo compilation log output

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Example of the diagnostic output generated by TorchDynamo when starting the compilation process.

```text
ChromiumEventLogger initialized with id 523d3804-d49f-4b2a-8103-4c5529f23c3e
torchdynamo start compiling fn /tmp/ipykernel_289/2126697152.py:5, stack (elided 5 frames):
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel_launcher.py", line 18, in <module>
    app.launch_new_instance()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/traitlets/config/application.py", line 1075, in launch_instance
    app.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 758, in start
    self.io_loop.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/tornado/platform/asyncio.py", line 211, in start
    self.asyncio_loop.run_forever()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
    self._run_once()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
    handle._run()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/utils.py", line 71, in preserve_context
    return await f(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 621, in shell_main
    await self.dispatch_shell(msg, subshell_id=subshell_id)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 478, in dispatch_shell
    await result
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 372, in execute_request
    await super().execute_request(stream, ident, parent)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 834, in execute_request
    reply_content = await reply_content
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 464, in do_execute
    res = shell.run_cell(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
    return super().run_cell(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3006, in run_cell
    result = self._run_cell(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3061, in _run_cell
    result = runner(coro)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
    coro.send(None)
```

--------------------------------

### Start and stop an EtcdServer instance

Source: https://docs.pytorch.org/docs/2.11/elastic/rendezvous.html

Demonstrates the lifecycle of a local etcd server, including initialization, starting, client retrieval, and stopping.

```python
server = EtcdServer("/usr/bin/etcd", 2379, "/tmp/default.etcd")
server.start()
client = server.get_client()
# use client
server.stop()
```

--------------------------------

### Identify 32-bit Package Installation Failure

Source: https://docs.pytorch.org/docs/2.11/notes/windows.html

Example of the error output when attempting to install PyTorch on an unsupported 32-bit system.

```text
Solving environment: failed

PackagesNotFoundError: The following packages are not available from current channels:

- pytorch

Current channels:
- https://repo.continuum.io/pkgs/main/win-32
- https://repo.continuum.io/pkgs/main/noarch
- https://repo.continuum.io/pkgs/free/win-32
- https://repo.continuum.io/pkgs/free/noarch
- https://repo.continuum.io/pkgs/r/win-32
- https://repo.continuum.io/pkgs/r/noarch
- https://repo.continuum.io/pkgs/pro/win-32
- https://repo.continuum.io/pkgs/pro/noarch
- https://repo.continuum.io/pkgs/msys2/win-32
- https://repo.continuum.io/pkgs/msys2/noarch
```

--------------------------------

### Launching Multiple Process Replicas

Source: https://docs.pytorch.org/docs/2.11/elastic/multiprocessing.html

Examples demonstrating how to use start_processes with functions and binaries, including handling of arguments and environment variables.

```python
log_dir = "/tmp/test"

# ok; two copies of foo: foo("bar0"), foo("bar1")
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
)

# invalid; envs missing for local rank 1
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}},
   log_dir=log_dir
)

# ok; two copies of /usr/bin/touch: touch file1, touch file2
start_processes(
   name="trainer",
   entrypoint="/usr/bin/touch",
   args:{0:("file1",), 1:("file2",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )

# caution; arguments casted to string, runs:
# echo "1" "2" "3" and echo "[1, 2, 3]"
start_processes(
   name="trainer",
   entrypoint="/usr/bin/echo",
   args:{0:(1,2,3), 1:([1,2,3],),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )
```

--------------------------------

### Prepare a model for quantization using prepare_fx

Source: https://docs.pytorch.org/docs/2.11/generated/torch.ao.quantization.quantize_fx.prepare_fx.html

This example demonstrates defining a model, configuring a QConfigMapping, and using prepare_fx to insert observers for calibration.

```python
import torch
from torch.ao.quantization import get_default_qconfig_mapping
from torch.ao.quantization.quantize_fx import prepare_fx

class Submodule(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)
    def forward(self, x):
        x = self.linear(x)
        return x

class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)
        self.sub = Submodule()

    def forward(self, x):
        x = self.linear(x)
        x = self.sub(x) + x
        return x

# initialize a floating point model
float_model = M().eval()

# define calibration function
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# qconfig is the configuration for how we insert observers for a particular
# operator
# qconfig = get_default_qconfig("fbgemm")
# Example of customizing qconfig:
# qconfig = torch.ao.quantization.QConfig(
#    activation=MinMaxObserver.with_args(dtype=torch.qint8),
#    weight=MinMaxObserver.with_args(dtype=torch.qint8))
# `activation` and `weight` are constructors of observer module

# qconfig_mapping is a collection of quantization configurations, user can
# set the qconfig for each operator (torch op calls, functional calls, module calls)
# in the model through qconfig_mapping
# the following call will get the qconfig_mapping that works best for models
# that target "fbgemm" backend
qconfig_mapping = get_default_qconfig_mapping("fbgemm")

# We can customize qconfig_mapping in different ways.
# e.g. set the global qconfig, which means we will use the same qconfig for
# all operators in the model, this can be overwritten by other settings
# qconfig_mapping = QConfigMapping().set_global(qconfig)
# e.g. quantize the linear submodule with a specific qconfig
# qconfig_mapping = QConfigMapping().set_module_name("linear", qconfig)
# e.g. quantize all nn.Linear modules with a specific qconfig
# qconfig_mapping = QConfigMapping().set_object_type(torch.nn.Linear, qconfig)
# for a more complete list, please see the docstring for :class:`torch.ao.quantization.QConfigMapping`
# argument

# example_inputs is a tuple of inputs, that is used to infer the type of the
# outputs in the model
# currently it's not used, but please make sure model(*example_inputs) runs
example_inputs = (torch.randn(1, 3, 224, 224),)

# TODO: add backend_config after we split the backend_config for fbgemm and qnnpack
# e.g. backend_config = get_default_backend_config("fbgemm")
# `prepare_fx` inserts observers in the model based on qconfig_mapping and
# backend_config. If the configuration for an operator in qconfig_mapping
# is supported in the backend_config (meaning it's supported by the target
# hardware), we'll insert observer modules according to the qconfig_mapping
# otherwise the configuration in qconfig_mapping will be ignored
#
# Example:
# in qconfig_mapping, user sets linear module to be quantized with quint8 for
# activation and qint8 for weight:
# qconfig = torch.ao.quantization.QConfig(
#     observer=MinMaxObserver.with_args(dtype=torch.quint8),
#     weight=MinMaxObserver.with-args(dtype=torch.qint8))
# Note: current qconfig api does not support setting output observer, but
# we may extend this to support these more fine grained control in the
# future

# qconfig_mapping = QConfigMapping().set_object_type(torch.nn.Linear, qconfig)
# in backend config, linear module also supports in this configuration:
# weighted_int8_dtype_config = DTypeConfig(
#   input_dtype=torch.quint8,
#   output_dtype=torch.quint8,
#   weight_dtype=torch.qint8,
#   bias_type=torch.float)

# linear_pattern_config = BackendPatternConfig(torch.nn.Linear) \
#    .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \
#    .add_dtype_config(weighted_int8_dtype_config) \
#    ...

# backend_config = BackendConfig().set_backend_pattern_config(linear_pattern_config)
# `prepare_fx` will check that the setting requested by suer in qconfig_mapping
# is supported by the backend_config and insert observers and fake quant modules
# in the model
prepared_model = prepare_fx(float_model, qconfig_mapping, example_inputs)
# Run calibration
calibrate(prepared_model, sample_inference_data)
```

--------------------------------

### Prepare a model for QAT using prepare_qat_fx

Source: https://docs.pytorch.org/docs/2.11/generated/torch.ao.quantization.quantize_fx.prepare_qat_fx.html

This example demonstrates defining a model, configuring the quantization mapping, and preparing the model for training with fake quantization modules.

```python
import torch
from torch.ao.quantization import get_default_qat_qconfig_mapping
from torch.ao.quantization.quantize_fx import prepare_qat_fx


class Submodule(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)

    def forward(self, x):
        x = self.linear(x)
        return x


class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(5, 5)
        self.sub = Submodule()

    def forward(self, x):
        x = self.linear(x)
        x = self.sub(x) + x
        return x


# initialize a floating point model
float_model = M().train()
# (optional, but preferred) load the weights from pretrained model
# float_model.load_weights(...)


# define the training loop for quantization aware training
def train_loop(model, train_data):
    model.train()
    for image, target in data_loader:
        ...


# qconfig is the configuration for how we insert observers for a particular
# operator
# qconfig = get_default_qconfig("fbgemm")
# Example of customizing qconfig:
# qconfig = torch.ao.quantization.QConfig(
#    activation=FakeQuantize.with_args(observer=MinMaxObserver.with_args(dtype=torch.qint8)),
#    weight=FakeQuantize.with_args(observer=MinMaxObserver.with_args(dtype=torch.qint8)))
# `activation` and `weight` are constructors of observer module

# qconfig_mapping is a collection of quantization configurations, user can
# set the qconfig for each operator (torch op calls, functional calls, module calls)
# in the model through qconfig_mapping
# the following call will get the qconfig_mapping that works best for models
# that target "fbgemm" backend
qconfig_mapping = get_default_qat_qconfig_mapping("fbgemm")

# We can customize qconfig_mapping in different ways, please take a look at
# the docstring for :func:`~torch.ao.quantization.prepare_fx` for different ways
# to configure this

# example_inputs is a tuple of inputs, that is used to infer the type of the
# outputs in the model
# currently it's not used, but please make sure model(*example_inputs) runs
example_inputs = (torch.randn(1, 3, 224, 224),)

# TODO: add backend_config after we split the backend_config for fbgemm and qnnpack
# e.g. backend_config = get_default_backend_config("fbgemm")
# `prepare_qat_fx` inserts observers in the model based on qconfig_mapping and
# backend_config, if the configuration for an operator in qconfig_mapping
# is supported in the backend_config (meaning it's supported by the target
# hardware), we'll insert fake_quantize modules according to the qconfig_mapping
# otherwise the configuration in qconfig_mapping will be ignored
# see :func:`~torch.ao.quantization.prepare_fx` for a detailed explanation of
# how qconfig_mapping interacts with backend_config
prepared_model = prepare_qat_fx(float_model, qconfig_mapping, example_inputs)
# Run training
train_loop(prepared_model, train_loop)
```

--------------------------------

### Generate tensors with torch.arange

Source: https://docs.pytorch.org/docs/2.11/generated/torch.arange.html

Examples demonstrating the creation of tensors using different start, end, and step arguments.

```python
>>> torch.arange(5)
tensor([ 0,  1,  2,  3,  4])
>>> torch.arange(1, 4)
tensor([ 1,  2,  3])
>>> torch.arange(1, 2.5, 0.5)
tensor([ 1.0000,  1.5000,  2.0000])
```

--------------------------------

### Initialize FSDP with custom parameter initialization

Source: https://docs.pytorch.org/docs/2.11/fsdp.html

Demonstrates how to use a custom initialization function for meta-device modules and how to integrate with torchdistX for deferred initialization.

```python
>>> module = MyModule(device="meta")
>>> def my_init_fn(module: nn.Module):
>>>     # E.g. initialize depending on the module type
>>>     ...
>>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
>>> print(next(fsdp_model.parameters()).device) # current CUDA device
>>> # With torchdistX
>>> module = deferred_init.deferred_init(MyModule, device="cuda")
>>> # Will initialize via deferred_init.materialize_module().
>>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
```

--------------------------------

### Generating Logarithmically Spaced Tensors

Source: https://docs.pytorch.org/docs/2.11/generated/torch.logspace.html

Examples demonstrating how to use torch.logspace with different start, end, steps, and base parameters.

```python
>>> torch.logspace(start=-10, end=10, steps=5)
tensor([ 1.0000e-10,  1.0000e-05,  1.0000e+00,  1.0000e+05,  1.0000e+10])
>>> torch.logspace(start=0.1, end=1.0, steps=5)
tensor([  1.2589,   2.1135,   3.5481,   5.9566,  10.0000])
>>> torch.logspace(start=0.1, end=1.0, steps=1)
tensor([1.2589])
>>> torch.logspace(start=2, end=2, steps=1, base=2)
tensor([4.0])
```

--------------------------------

### Shared File-System Initialization Example

Source: https://docs.pytorch.org/docs/2.11/_sources/distributed.md.txt

Initialize using a shared file system. The URL should point to a non-existent file in an existing directory on a shared file system. This method requires manual rank specification and careful cleanup of the initialization file between runs to prevent deadlocks.

```python
import torch.distributed as dist

# rank should always be specified
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)
```

--------------------------------

### Graph Break with Context Managers

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Example demonstrating the use of contextlib.nullcontext and torch.no_grad with a manual graph break inside a compiled function.

```python
import contextlib
@torch.compile
def fn(x):
    with contextlib.nullcontext():
        with torch.no_grad():
            x = x + 1
            torch._dynamo.graph_break()
            return x + 1
fn(torch.randn(3))
```

--------------------------------

### TorchDynamo Trace Output

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

The internal trace log generated by TorchDynamo when executing the graph break example.

```text
torchdynamo start compiling fn /tmp/ipykernel_289/2546874632.py:1, stack (elided 5 frames):
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel_launcher.py", line 18, in <module>
    app.launch_new_instance()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/traitlets/config/application.py", line 1075, in launch_instance
    app.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 758, in start
    self.io_loop.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/tornado/platform/asyncio.py", line 211, in start
    self.asyncio_loop.run_forever()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
    self._run_once()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
    handle._run()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/utils.py", line 71, in preserve_context
    return await f(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 621, in shell_main
    await self.dispatch_shell(msg, subshell_id=subshell_id)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 478, in dispatch_shell
    await result
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 372, in execute_request
    await super().execute_request(stream, ident, parent)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 834, in execute_request
    reply_content = await reply_content
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 464, in do_execute
    res = shell.run_cell(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/zmqshell.py", line 663, in run_cell
    return super().run_cell(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3006, in run_cell
    result = self._run_cell(
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3061, in _run_cell
    result = runner(coro)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner
    coro.send(None)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3266, in run_cell_async
    has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3445, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3505, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_289/2546874632.py", line 9, in <module>
    fn(torch.randn(3))
Step 1: torchdynamo start tracing fn /tmp/ipykernel_289/2546874632.py:1
create_env
TRACE starts_line /tmp/ipykernel_289/2546874632.py:3 in fn
        try:
TRACE SETUP_FINALLY 30 []
TRACE starts_line /tmp/ipykernel_289/2546874632.py:4 in fn
            x = x + 1
TRACE LOAD_FAST x []
TRACE LOAD_CONST 1 [LazyVariableTracker(unrealized: <class 'torch.Tensor'>)]
TRACE BINARY_ADD None [LazyVariableTracker(unrealized: <class 'torch.Tensor'>), ConstantVariable(int: 1)]
wrap_to_fake L['x'] (3,) StatefulSymbolicContext(dynamic_sizes=[<DimDynamic.STATIC: 2>], dynamic_strides=[<DimDynamic.INFER_STRIDE: 4>], constraint_sizes=[None], constraint_strides=[None], specialize_on=[[]], view_base_context=None, shape_ids=None, tensor_source=LocalSource(local_name='x', is_input=True, dynamism=None, is_derefed_cell_contents=False), shape_env_to_source_to_symbol_cache={}) <class 'torch.Tensor'>
create_graph_input L_x_ L['x'] FakeTensor(..., size=(3,)) at debug_level 0 before=False
TRACE STORE_FAST x [TensorVariable()]
TRACE starts_line /tmp/ipykernel_289/2546874632.py:5 in fn
            torch._dynamo.graph_break()
TRACE LOAD_GLOBAL torch []
TRACE LOAD_ATTR _dynamo [LazyVariableTracker(unrealized: <class 'module'>)]
TRACE LOAD_ATTR graph_break [LazyVariableTracker(unrealized: <class 'module'>)]
```

--------------------------------

### Compiler Function Execution Trace

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Example of internal compiler logs showing the execution of an eager function and the subsequent guard tracking process.

```text
# File: /tmp/ipykernel_289/273153676.py:2 in inner1, code: return x + 1
        x: "f32[3][1]cpu" = l_x_ + 1;  l_x_ = None
        return (x,)
```

--------------------------------

### Basic DDP Example

Source: https://docs.pytorch.org/docs/2.11/_sources/notes/ddp.rst.txt

Illustrates setting up and running a single forward/backward pass with DDP. Requires environment variables for initialization.

```python
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
import os
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()
```

--------------------------------

### Generate evenly spaced tensors with torch.linspace

Source: https://docs.pytorch.org/docs/2.11/generated/torch.linspace.html

Examples demonstrating the creation of tensors with varying start, end, and step values.

```python
>>> torch.linspace(3, 10, steps=5)
tensor([  3.0000,   4.7500,   6.5000,   8.2500,  10.0000])
>>> torch.linspace(-10, 10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=5)
tensor([-10.,  -5.,   0.,   5.,  10.])
>>> torch.linspace(start=-10, end=10, steps=1)
tensor([-10.])
```

--------------------------------

### Graph Break in a Loop Example

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Demonstrates a function that triggers a graph break using torch._dynamo.graph_break() inside a loop.

```python
@torch.compile
def fn(x):
    for i in range(5):
        x = x + 1
        if i == 3:
            torch._dynamo.graph_break()
    return x
fn(torch.randn(3))
```

--------------------------------

### Compiler Intermediate Representation and Guard Logs

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Example of the internal compiler output showing function compilation, guard tracking, and the identification of skipped functions.

```text
# File: /tmp/ipykernel_289/2126697152.py:2 in inner1, code: return x + 1
        add: "f32[3][1]cpu" = l_x_ + 1;  l_x_ = None
        return (add,)
Step 2: calling compiler function eager
Step 2: done compiler function eager
produce_guards
track_symint L['x'].size()[0] 3 None
track_symint L['x'].stride()[0] 1 None
track_symint L['x'].storage_offset() 0 None
Skipping guard L['x'].size()[0] == 3
Skipping guard L['x'].stride()[0] == 1
Skipping guard L['x'].storage_offset() == 0
GUARDS:

TREE_GUARD_MANAGER:
+- RootGuardManager
| +- LAMBDA_GUARD: torch._functorch.aot_autograd.utils.top_saved_tensors_hooks ids == None  # _dynamo/output_graph.py:964 in init_ambient_guards
| +- GLOBAL_STATE: ___check_global_state() against {"allow_bf16_reduce":0,"allow_fp16_reduce":0,"allow_tf32":false,"autocast_state":{"cached_enabled":true,"dtype":[15,5,5,15,5,5,15,15,5,5],"enabled":[false,false,false,false,false,false,false,false,false,false]},"default_dtype":6,"deterministic_algorithms":false,"deterministic_algorithms_warn_only":false,"grad_mode":true,"num_threads":4,"torch_function":true,"torch_function_all_disabled":false}
| +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:951 in init_ambient_guards
| +- GuardManager: source=L['x'], accessed_by=FrameLocalsGuardAccessor(key='x', framelocals_idx=0), type=<class 'torch.Tensor'>, tag_safe=(True, False)
| | +- TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[3], stride=[1])  # return x + 1  # mp/ipykernel_289/2126697152.py:2 in inner1
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 1  # mp/ipykernel_289/2126697152.py:2 in inner1
Guard eval latency = 33.92 us
put_code_state: no cache key, skipping
run_gc_after_compile: running gc
skipping: inner (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_compile.py)
skipping: disable (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/decorators.py)
skipping: innermost_fn (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: __init__ (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: __init__ (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: nothing (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: innermost_backend (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: __call__ (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: _fn (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py)
skipping: skip_frame (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/decorators.py)
torchdynamo start compiling inner2 /tmp/ipykernel_289/2126697152.py:3, stack (elided 5 frames):
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/envs/py_3.10/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel_launcher.py", line 18, in <module>
    app.launch_new_instance()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/traitlets/config/application.py", line 1075, in launch_instance
    app.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelapp.py", line 758, in start
    self.io_loop.start()
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/tornado/platform/asyncio.py", line 211, in start
    self.asyncio_loop.run_forever()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
    self._run_once()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
    handle._run()
  File "/opt/conda/envs/py_3.10/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/utils.py", line 71, in preserve_context
    return await f(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 621, in shell_main
    await self.dispatch_shell(msg, subshell_id=subshell_id)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/kernelbase.py", line 478, in dispatch_shell
```

--------------------------------

### Wrapping Constants in Tensors to Prevent Recompilations

Source: https://docs.pytorch.org/docs/2.11/_sources/user_guide/torch_compiler/torch.compiler_troubleshooting.md.txt

Shows how to wrap float variables in tensors to avoid recompilations. This is applied to both the simple addition example and the optimizer/scheduler setup.

```python
# first example
for i in range(1, 10):
    fn(torch.ones(i), torch.tensor(0.5 + i))

# second example
opt = torch.optim.Adam(mod.parameters(), lr=torch.tensor(0.01))
sched = torch.optim.lr_scheduler.ExponentialLR(opt, torch.tensor(0.9))
```

--------------------------------

### setup_context(ctx, inputs, output)

Source: https://docs.pytorch.org/docs/2.11/generated/torch.autograd.function.NestedIOFunction.html

Defines the context setup for the backward pass when the forward method signature does not include the context object.

```APIDOC
### Method
static setup_context(ctx, inputs, output)

### Description
Used to set up the context for the backward pass when the forward method is defined without a ctx argument. This is an alternative to setting up the context inside the forward method.
```

--------------------------------

### Graph Break in Try Block Example

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

A function decorated with @torch.compile that triggers a graph break inside a try block.

```python
@torch.compile
def fn(x):
    try:
        x = x + 1
        torch._dynamo.graph_break()
        return x + 1
    except Exception as e:
        pass
fn(torch.randn(3))
```

--------------------------------

### Initialize Modules with Custom Settings

Source: https://docs.pytorch.org/docs/2.11/notes/modules.html

Demonstrates initializing modules directly on specific devices, using different dtypes, or skipping default initialization for custom logic.

```python
# Initialize module directly onto GPU.
m = nn.Linear(5, 3, device='cuda')

# Initialize module with 16-bit floating point parameters.
m = nn.Linear(5, 3, dtype=torch.half)

# Skip default parameter initialization and perform custom (e.g. orthogonal) initialization.
m = torch.nn.utils.skip_init(nn.Linear, 5, 3)
nn.init.orthogonal_(m.weight)
```

--------------------------------

### Debug Specific NCCL Subsystems

Source: https://docs.pytorch.org/docs/2.11/_sources/distributed.md.txt

Use NCCL_DEBUG_SUBSYS to get detailed logs for specific NCCL aspects. For example, 'COLL' for collective calls or 'GRAPH' for topology detection.

```bash
export NCCL_DEBUG_SUBSYS=COLL
```

```bash
export NCCL_DEBUG_SUBSYS=GRAPH
```

--------------------------------

### Initialize TCPStore and Wait for Key

Source: https://docs.pytorch.org/docs/2.11/distributed.html

Demonstrates initializing a TCPStore and using the wait method with a timeout.

```python
>>> import torch.distributed as dist
>>> from datetime import timedelta
>>> # Using TCPStore as an example, other store types can also be used
>>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30))
>>> # This will throw an exception after 10 seconds
>>> store.wait(["bad_key"], timedelta(seconds=10))
```

--------------------------------

### Basic torch.compile Profiling Example

Source: https://docs.pytorch.org/docs/2.11/_sources/user_guide/torch_compiler/torch.compiler_profiling_torch_compile.md.txt

Profile a compiled model's forward and backward pass. Includes warm-up, profiling context, and trace export. Ensure you have torchvision installed.

```python
import torch
from torchvision.models import resnet18

device = 'cuda'      # or 'cpu', 'xpu', etc.
model = resnet18().to(device)

inputs = [torch.randn((5, 3, 224, 224), device=device) for _ in range(10)]

model_c = torch.compile(model)

def fwd_bwd(inp):
    out = model_c(inp)
    out.sum().backward()

# warm up
fwd_bwd(inputs[0])

with torch.profiler.profile() as prof:
    for i in range(1, 4):
        fwd_bwd(inputs[i])
        prof.step()

prof.export_chrome_trace("trace.json")
```

--------------------------------

### Triggering a graph break with skip_frame

Source: https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/compile/programming_model.skipped_functions.html

Example of a function call that invokes torch._dynamo.skip_frame(), resulting in a graph break during tracing.

```python
        def problematic_code():
            torch._dynamo.skip_frame()
```

--------------------------------

### Distributed RPC Example

Source: https://docs.pytorch.org/docs/2.11/_sources/rpc/distributed_autograd.rst.txt

Demonstrates a simple distributed computation using RPC, involving remote function calls and local operations on results. Requires setup of distributed environment.

```python
import torch
import torch.distributed.rpc as rpc

def my_add(t1, t2):
  return torch.add(t1, t2)

# On worker 0:
t1 = torch.rand((3, 3), requires_grad=True)
t2 = torch.rand((3, 3), requires_grad=True)

# Perform some computation remotely.
t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))

# Perform some computation locally based on remote result.
t4 = torch.rand((3, 3), requires_grad=True)
t5 = torch.mul(t3, t4)

# Compute some loss.
loss = t5.sum()
```

--------------------------------

### Implement example_arguments

Source: https://docs.pytorch.org/docs/2.11/onnx_export.html

Define sample inputs for tracing and testing. Return a tuple of positional arguments and a dictionary of keyword arguments.

```python
def example_arguments(self):
    # For a model expecting a single tensor input
    return (torch.randn(1, 3, 224, 224),), None
```

```python
def example_arguments(self):
    # For a model with multiple inputs and keyword arguments
    return (torch.randn(1, 3, 224, 224), torch.randn(1, 512)), {
        "temperature": 1.0
    }
```

--------------------------------

### Configure QConfigMapping

Source: https://docs.pytorch.org/docs/2.11/generated/torch.ao.quantization.qconfig_mapping.QConfigMapping.html

Demonstrates how to chain configuration methods to define quantization settings for various model components.

```python
qconfig_mapping = QConfigMapping()
    .set_global(global_qconfig)
    .set_object_type(torch.nn.Linear, qconfig1)
    .set_object_type(torch.nn.ReLU, qconfig1)
    .set_module_name_regex("foo.*bar.*conv[0-9]+", qconfig1)
    .set_module_name_regex("foo.*", qconfig2)
    .set_module_name("module1", qconfig1)
    .set_module_name("module2", qconfig2)
    .set_module_name_object_type_order("foo.bar", torch.nn.functional.linear, 0, qconfig3)
```

--------------------------------

### Log PyTorch Model and Metrics with SummaryWriter

Source: https://docs.pytorch.org/docs/2.11/_sources/tensorboard.rst.txt

Initialize SummaryWriter to log data. This example logs a model graph and images. Ensure TensorBoard is installed and run 'tensorboard --logdir=runs' to visualize.

```python
import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(False)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
images, labels = next(iter(trainloader))

grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()
```

--------------------------------

### Autocast Behavior: Custom kwargs Example

Source: https://docs.pytorch.org/docs/2.11/_sources/user_guide/torch_compiler/torch.compiler_backward.md.txt

Configure the backward pass to run under a specific autocast context defined by provided kwargs. This is useful when the backward pass requires a different autocast setup than the forward pass.

```python
y = torch.compile(region)(x)
...
# Backward pass runs under special context manager
with torch.amp.autocast(**kwargs):
    z.backward()
```

--------------------------------

### Configuring prepare_custom_config_dict

Source: https://docs.pytorch.org/docs/2.11/generated/torch.ao.quantization.prepare.html

Example showing how to define a custom configuration dictionary to map float custom modules to their observed counterparts.

```python
# Example of prepare_custom_config_dict:
prepare_custom_config_dict = {
    # user will manually define the corresponding observed
    # module class which has a from_float class method that converts
    # float custom module to observed custom module
    "float_to_observed_custom_module_class": {CustomModule: ObservedCustomModule}
}
```

--------------------------------

### FAST Mode Distributed Autograd Example

Source: https://docs.pytorch.org/docs/2.11/_sources/rpc/distributed_autograd.rst.txt

This code snippet shows a basic setup for distributed autograd using the FAST mode algorithm. It defines a simple addition function to be used within the distributed context.

```python
import torch
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc

def my_add(t1, t2):
  return torch.add(t1, t2)
```
Tokens:
Raw