PyTorch

https://docs.pytorch.org/docs/2.10/
Admin
PyTorch is an optimized tensor library for deep learning using GPUs and CPUs, providing...
Tokens:1,297,329
Snippets:11,513
Trust Score:10
Update:2 months ago
Show doc for...
Tokens:
Raw
### Initializing Distributed Models and Optimizers

Source: https://docs.pytorch.org/docs/2.10/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)
```

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

### Initialize ZeroRedundancyOptimizer

Source: https://docs.pytorch.org/docs/2.10/distributed.optim.html

Example showing the setup of ZeroRedundancyOptimizer with a DistributedDataParallel model.

```python
>>> import torch.nn as nn
>>> from torch.distributed.optim import ZeroRedundancyOptimizer
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
>>> ddp = DDP(model, device_ids=[rank])
>>> opt = ZeroRedundancyOptimizer(
>>>     ddp.parameters(),
>>>     optimizer_class=torch.optim.Adam,
>>>     lr=0.01
>>> )
>>> ddp(inputs).sum().backward()
>>> opt.step()
```

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

### Install and run TensorBoard

Source: https://docs.pytorch.org/docs/2.10/tensorboard.html

Commands to install the TensorBoard package and launch the server to view logs.

```bash
pip install tensorboard
tensorboard --logdir=runs
```

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

### Example Output

Source: https://docs.pytorch.org/docs/2.10/user_guide/torch_compiler/compile/programming_model.compiler_disable.html

The expected output when running the provided code example.

```text
tensor([8., 8., 8.])
```

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

### Start and stop an EtcdServer

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

Demonstrates the lifecycle of an EtcdServer instance, including starting the server, retrieving a client, and stopping the server.

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

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

### TorchDynamo compilation log output

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

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

```text
ChromiumEventLogger initialized with id 65691516-4673-4093-b6a8-b2840eac3a8c
torchdynamo start compiling fn /tmp/ipykernel_533/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 614, 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 471, in dispatch_shell
    await result
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 366, 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 827, in execute_request
    reply_content = await reply_content
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 458, 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
```

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

### Tensor Output Example

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

Example output of a tensor operation.

```python
tensor([0.6105, 0.0799, 2.6129])
```

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

### Compute matrix inverse examples

Source: https://docs.pytorch.org/docs/2.10/generated/torch.linalg.inv.html

Examples demonstrating the inversion of standard, batched, and complex matrices.

```python
>>> A = torch.randn(4, 4)
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(1.1921e-07)

>>> A = torch.randn(2, 3, 4, 4)  # Batch of matrices
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(1.9073e-06)

>>> A = torch.randn(4, 4, dtype=torch.complex128)  # Complex matrix
>>> Ainv = torch.linalg.inv(A)
>>> torch.dist(A @ Ainv, torch.eye(4))
tensor(7.5107e-16, dtype=torch.float64)
```

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

### Minifier Output Example

Source: https://docs.pytorch.org/docs/2.10/user_guide/torch_compiler/torch.compiler_aot_inductor_minifier.html

Example console output generated when running the minifier with the command='minify' option.

```text
...
W1031 16:48:08.938000 3598491 torch/_dynamo/repro/aoti.py:89] Writing checkpoint with 3 nodes to /data/users/shangdiy/pytorch/torch_compile_debug/run_2024_10_31_16_48_02_720863-pid_3598491/minifier/checkpoints/3.py
W1031 16:48:08.975000 3598491 torch/_dynamo/repro/aoti.py:101] Copying repro file for convenience to /data/users/shangdiy/pytorch/repro.py
Wrote minimal repro out to repro.py
```

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

### Install PyTorch for Intel GPU via Binaries

Source: https://docs.pytorch.org/docs/2.10/notes/get_start_xpu.html

Use these commands to install PyTorch, torchvision, and torchaudio with Intel GPU support from official wheels.

```bash
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu
```

```bash
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu
```

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

### Example Output of Graph Break

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

Example of the internal trace output generated during a graph break.

```text
TRACE RETURN_VALUE None [ConstantVariable(NoneType: None)]
DONE INLINING <code object __exit__ at 0x7fd3ce1c5dc0, file "/tmp/ipykernel_533/4148913404.py", line 4>
TRACE POP_TOP None [TensorVariable(), ConstantVariable(NoneType: None)]
TRACE RETURN_VALUE None [TensorVariable()]
Step 1: torchdynamo done tracing torch_dynamo_resume_in_fn_at_5 (RETURN_VALUE)
return triggered compile
COMPILING GRAPH due to GraphCompileReason(reason='return_value', user_stack=[<FrameSummary file /tmp/ipykernel_533/2124425154.py, line 6 in torch_dynamo_resume_in_fn_at_5>], graph_break=False)
TRACED GRAPH
 ===== __compiled_fn_23_63f4b076_8bfa_4361_ab6a_96eb31fb2676 =====
 /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[3][1]cpu"):
        l_x_ = L_x_
        
        # File: /tmp/ipykernel_533/2124425154.py:7 in torch_dynamo_resume_in_fn_at_5, 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:807 in init_ambient_guards
| +- GLOBAL_STATE: ___check_global_state()
| +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:794 in init_ambient_guards
| +- GuardManager: source=L['x'], accessed_by=FrameLocalsGuardAccessor(key='x', framelocals_idx=3), 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_533/2124425154.py:7 in torch_dynamo_resume_in_fn_at_5
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 1  # mp/ipykernel_533/2124425154.py:7 in torch_dynamo_resume_in_fn_at_5
| +- GuardManager: source=G, accessed_by=GlobalsGuardAccessor, type=<class 'dict'>, tag_safe=(True, False)
| | +- GuardManager: source=G['CustomCtxManager'], accessed_by=DictGetItemGuardAccessor('CustomCtxManager'), type=<class 'type'>, tag_safe=(True, False)
| | | +- ID_MATCH: ___check_obj_id(G['CustomCtxManager'], 94362076592320), type=<class '__main__.CustomCtxManager'>  # with CustomCtxManager():  # mp/ipykernel_533/2124425154.py:6 in torch_dynamo_resume_in_fn_at_5
Guard eval latency = 154.77 us
put_code_state: no cache key, skipping
run_gc_after_compile: running gc
```

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

### Execute C++ binary

Source: https://docs.pytorch.org/docs/2.10/user_guide/torch_compiler/torch.compiler_aot_inductor.html

Example output after running the compiled C++ binary.

```bash
$ ./aoti_example
Result from the first inference:
0.4866
0.5184
0.4462
0.4611
0.4744
0.4811
0.4938
0.4193
[ CUDAFloatType{8,1} ]
Result from the second inference:
0.4883
0.4703
[ CUDAFloatType{2,1} ]
```

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

### Launching processes with start_processes

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

Examples demonstrating valid and invalid configurations for launching function and binary entrypoints, including argument handling and environment variable requirements.

```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
 )
```

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

### View PackagesNotFoundError output

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

Example of the error message displayed when attempting to install PyTorch on a 32-bit Windows environment.

```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
```

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

### Initialize and draw from SobolEngine

Source: https://docs.pytorch.org/docs/2.10/generated/torch.quasirandom.SobolEngine.html

Demonstrates initializing a SobolEngine with a specified dimension and drawing a sequence of points.

```python
>>> soboleng = torch.quasirandom.SobolEngine(dimension=5)
>>> soboleng.draw(3)
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
        [0.7500, 0.2500, 0.2500, 0.2500, 0.7500]])
```

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

### Initialize FileStore

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

Demonstrates using a file-based store to share data between processes.

```python
>>> import torch.distributed as dist
>>> store1 = dist.FileStore("/tmp/filestore", 2)
>>> store2 = dist.FileStore("/tmp/filestore", 2)
>>> # Use any of the store methods from either the client or server after initialization
>>> store1.set("first_key", "first_value")
>>> store2.get("first_key")
```

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

### Generating Logarithmically Spaced Tensors

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

Examples demonstrating the creation of tensors with various 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])
```

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

### Nested Graph Break Tracing Example

Source: https://docs.pytorch.org/docs/2.10/user_guide/torch_compiler/compile/programming_model.nested_graph_breaks.html

Shows the initial setup where a graph break in a nested function triggers compilation.

```python
def inner1(x):
    x = x + 1
    torch._dynamo.graph_break()  # stop tracing due to graph break
    return x + 2

def inner2(x):
    x = x + 4
    x = inner1(x)
    x = x + 8

@torch.compile
def f(x):
    # start tracing from here
    x = x + 16
    x = inner2(x)
    x = x + 32

f(torch.randn(3))
```

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

### Initialize Checkpoint Loading Components

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

Setup the model, optimizer, and file system storage reader required for loading a checkpoint.

```python
>>> my_model = MyModule()
>>> optimizer = Adagrad(my_model.parameters())
>>> model_state_dict = my_model.state_dict()
>>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
...     "/checkpoint/1"
... )
```

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

### Compilation Log Output

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

Log output showing the start of the compilation process.

```text
torchdynamo start compiling fn /tmp/ipykernel_533/3152636365.py:2, 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)
```

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

### View launch utility help

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

Display available optional arguments for the launch module.

```bash
python -m torch.distributed.launch --help
```

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

### Generating evenly spaced tensors with torch.linspace

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

Examples demonstrating how to create tensors with a specific number of steps between start and end 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.])
```

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

### Example of Traced Graph and Guard Logs

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

Logs showing the tracing process, graph compilation, and guard evaluation for a function.

```text
Step 1: torchdynamo done tracing inner2 (RETURN_VALUE)
return triggered compile
COMPILING GRAPH due to GraphCompileReason(reason='return_value', user_stack=[<FrameSummary file /tmp/ipykernel_533/2126697152.py, line 4 in inner2>], graph_break=False)
TRACED GRAPH
 ===== __compiled_fn_4_348a4625_356f_4891_aef2_52f2fc37ef72 =====
 /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[3][1]cpu"):
        l_x_ = L_x_
        
        # File: /tmp/ipykernel_533/2126697152.py:4 in inner2, code: return x + 2
        add: "f32[3][1]cpu" = l_x_ + 2;  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:807 in init_ambient_guards
| +- GLOBAL_STATE: ___check_global_state()
| +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:794 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 + 2  # mp/ipykernel_533/2126697152.py:4 in inner2
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # return x + 2  # mp/ipykernel_533/2126697152.py:4 in inner2
Guard eval latency = 14.18 us
put_code_state: no cache key, skipping
run_gc_after_compile: running gc
skipping: remove (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/weak.py)
skipping: __hash__ (reason: in skipfiles, file: /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/weak.py)
```

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

### Prepare a model for QAT using prepare_qat_fx

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

Demonstrates the workflow for preparing a floating-point model for quantization aware training by defining a qconfig_mapping and applying prepare_qat_fx.

```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)
```

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

### Configure Copy Engine Collectives

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

Example demonstrating the initialization of a process group with zero-CTA policy, symmetric memory setup, and collective operation execution using copy engines.

```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
symm_mem.enable_symm_mem_for_group(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 Guard and Compilation Logs

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

Example logs showing the tracking of symbolic integers, skipped guards, and the tree guard manager state during compilation.

```text
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:807 in init_ambient_guards
| +- GLOBAL_STATE: ___check_global_state()
| +- TORCH_FUNCTION_MODE_STACK: ___check_torch_function_mode_stack()
| +- DEFAULT_DEVICE: utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:794 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])  # x = inner1(x)  # mp/ipykernel_533/273153676.py:7 in fn
| | +- NO_HASATTR: hasattr(L['x'], '_dynamo_dynamic_indices') == False           # x = inner1(x)  # mp/ipykernel_533/273153676.py:7 in fn
| +- GuardManager: source=G, accessed_by=GlobalsGuardAccessor, type=<class 'dict'>, tag_safe=(True, False)
| | +- GuardManager: source=G['inner1'], accessed_by=DictGetItemGuardAccessor('inner1'), type=<class 'function'>, tag_safe=(True, False)
| | | +- GuardManager: source=G['inner1'].__code__, accessed_by=CodeGuardAccessor, type=<code'>, tag_safe=(True, False)
| | | | +- ID_MATCH: ___check_obj_id(G['inner1'].__code__, 140546383531328), type=<code object inner1 at 0x7fd38143cd40, file "/tmp/ipykernel_533/273153676.py", line 1>  # x = inner1(x)  # mp/ipykernel_533/273153676.py:7 in fn
Guard eval latency = 286.53 us
put_code_state: no cache key, skipping
run_gc_after_compile: running gc
torchdynamo start compiling problematic_code /tmp/ipykernel_533/273153676.py:8, 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 614, 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 471, in dispatch_shell
    await result
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 366, 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 827, in execute_request
    reply_content = await reply_content
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/ipykernel/ipkernel.py", line 458, 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,
```

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

### torch.narrow(input, dim, start, length)

Source: https://docs.pytorch.org/docs/2.10/generated/torch.narrow.html

Returns a new tensor that is a narrowed version of the input tensor. The dimension dim is input from start to start + length.

```APIDOC
## torch.narrow(input, dim, start, length)

### Description
Returns a new tensor that is a narrowed version of the input tensor. The dimension `dim` is input from `start` to `start + length`. The returned tensor and `input` tensor share the same underlying storage.

### Parameters
- **input** (Tensor) - The tensor to narrow.
- **dim** (int) - The dimension along which to narrow.
- **start** (int or Tensor) - Index of the element to start the narrowed dimension from. Can be negative, which means indexing from the end of dim. If Tensor, it must be an 0-dim integral Tensor.
- **length** (int) - Length of the narrowed dimension, must be weakly positive.

### Example
```python
>>> x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
>>> torch.narrow(x, 0, 0, 2)
tensor([[ 1,  2,  3],
        [ 4,  5,  6]])
>>> torch.narrow(x, 1, 1, 2)
tensor([[ 2,  3],
        [ 5,  6],
        [ 8,  9]])
>>> torch.narrow(x, -1, torch.tensor(-1), 1)
tensor([[3],
        [6],
        [9]])
```
```

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

### Prepare a model for quantization using prepare_fx

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

Demonstrates the end-to-end process of defining a model, configuring quantization settings, 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)
```

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

### Graph Break in a Loop Example

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

Demonstrates a function that triggers a graph break inside a loop, which prevents torch.compile from resuming.

```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))
```

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

### Configure PrepareCustomConfig

Source: https://docs.pytorch.org/docs/2.10/generated/torch.ao.quantization.fx.custom_config.PrepareCustomConfig.html

Demonstrates the fluent API usage to set various configuration options for quantization preparation.

```python
prepare_custom_config = PrepareCustomConfig()             .set_standalone_module_name("module1", qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_standalone_module_class(MyStandaloneModule, qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_float_to_observed_mapping(FloatCustomModule, ObservedCustomModule)             .set_non_traceable_module_names(["module2", "module3"])             .set_non_traceable_module_classes([NonTraceableModule1, NonTraceableModule2])             .set_input_quantized_indexes([0])             .set_output_quantized_indexes([0])             .set_preserved_attributes(["attr1", "attr2"])
```

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

### Graph Break in a Try Block Example

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

Demonstrates a function containing a graph break inside a try block, which prevents resumption.

```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))
```

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

### FutureWarning Example

Source: https://docs.pytorch.org/docs/2.10/user_guide/torch_compiler/export.html

Example of a warning message encountered during the export process.

```text
/opt/conda/envs/py_3.10/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
```

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

### Run Training with Multiple Processes

Source: https://docs.pytorch.org/docs/2.10/notes/multiprocessing.html

Command-line example for executing a training script with a specified number of processes.

```bash
python main.py --num-processes 4
```

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

### Manual Graph Break Trigger

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

Example of a manual graph break insertion using the torch._dynamo.graph_break() function.

```python
                torch._dynamo.graph_break()
```
Tokens:
Raw