There are two motivating use cases for this change:
1) export (when we trace pytree calls into a graph, we don't want to accidentally trace the side effect bytecode which will pollute the initial state) -> We want to warn about side effects and don't want to actually apply them
2) VLLM -> They want to detect side effects and error out.
We implement this with two configs where one config controls whether we want to apply side effects (by default yes) and the warning level for side effects (warning for export and error for VLLM). We intentionally ignore input side effects, because they are captured in the graph and export would never trace the actual dynamo graph module when tracing the pytree calls).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167239
Approved by: https://github.com/williamwen42, https://github.com/anijain2305
This PR introduces inductor lite mode for opt-in optimizations and numeric correctness guarantees.
Different from default mode that applies all possible fusions, lite mode gives the control back to user and provides guarantee on numeric correctness. Specifically, this mode:
- **Fallback by Default**: Fallback for ALL nodes by default, unless users explicitly mark node for inductor fusion.
- **Selective Decomposition**: Skip decomposition for all nodes except for user marked nodes.
- **Regional inductor compile**
- Skip dead code elimination
- Skip buffer reues
- Skip reorder passes, such as reorder for peak memory, reorder for compute comm overlap, and reorder_for_reducing_graph_partitions.
- Skip all pre-grad, joint-graph, and post-grad passes.
## Example: Flex Attention
```python
import torch
import torch.fx.traceback as fx_traceback
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
def _squared(score, b, h, m, n):
return score * score
def mask_mod(b, h, q, k):
return q >= 0
a, b = 12, 64
block_mask = create_block_mask(mask_mod, None, None, a * b, a * b, device="cuda")
def fn(x):
x = torch.sin(x)
with fx_traceback.annotate({"compile_with_inductor": 0}):
x = flex_attention(x, x, x, block_mask=block_mask, score_mod=_squared)
return torch.cos(x)
x = torch.randn(1, 1, a * b, b, dtype=torch.bfloat16, device="cuda", requires_grad=True)
opt_fn = torch.compile(fn, mode="lite", fullgraph=True,)
opt_fn(x)
```
[code diff](https://www.internalfb.com/intern/diffing/?paste_number=2027441476)
[default mode tlp](https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpYAzDxX/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000) vs [lite mode tlp](https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpnnuh1W/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000)
## Numerics
Inductor lite mode provides bitwise equivalence with `aot_eager` backend on torchtitan llama3-8b and DeepSeek v3. https://github.com/pytorch/torchtitan/pull/2005close: #167012
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167115
Approved by: https://github.com/ezyang
This adds a workflow to run full set of UTs on default and distributed configs on ROCm MI3xx CI runners, to _eventually_ assess if the CI capacity can handle the PR-based workload for trunk.yml. The plan was to keep this workflow in unstable as we test out this new CI capacity, so it wouldn't impact PR merges. However, since upstream maintainers have indicated that, as of today, even unstable workflows will block PR merges, we are going with branch push-based triggers to at least pipeclean this workflow on the new CI capacity.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167587
Approved by: https://github.com/jeffdaily
The seq_nr doesn't always increment for gradient accumulation nodes, and they might be copying annotation from forward nodes.
I'm just going to skip copying the custom meta for any gradient accumulation nodes and give them a special tag e.g. node.meta["is_gradient_acc"]=True
Example repro for deepseek torchtitan (without using DTensor): https://gist.github.com/yushangdi/aae13ea382732f31d0fdfb3ffeda12c8
(side note: if you want some more hints on these gradient acc node: 1) they have torch.ops.aten.add.Tensor op, not add.default. 2) they have the highest seq_nr(s) )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167572
Approved by: https://github.com/mlazos
Summary:
Added a new flag called "use_legacy_dynamo_graph_capture" which defaults to True and only False with the updated test_strict_export_v2.py
In addiotion to this flag, we also use legacy tracer when the following features are used:
1. dynamic shape
2. preserve module call signature
3. retracing.
4. draft mode.
Test Plan:
test_strict_export_v2.py
Fixes #ISSUE_NUMBER
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167236
Approved by: https://github.com/tugsbayasgalan
# Summary
See title ;)
## Design
Currently once you install there is no going back in the same python process, this need not be the case, cc @mikaylagawarecki's work on being able to grab original impl. I'll leave for follow up.
Okay I added an open reg, but I really want the backends to be found so some weird typing but we get
<img width="523" height="197" alt="Screenshot 2025-11-07 at 3 30 32 PM" src="https://github.com/user-attachments/assets/586de943-bbed-40cf-abd1-131f747a4cf1" />
## Overheads:
<img width="799" height="735" alt="Screenshot 2025-11-07 at 2 35 04 PM" src="https://github.com/user-attachments/assets/f9217f31-3e42-4816-8fb3-29ea8b49d735" />
First call to forward -> majority of time is spent in jit for FA
First call to backward, 3sec interestingly it doesn't appear that with_stack gets events in the backwards loop @albanD is this expected?
<img width="948" height="385" alt="Screenshot 2025-11-07 at 2 35 50 PM" src="https://github.com/user-attachments/assets/a40bacd0-3fb0-4bd8-b33e-bec8fb3f36c0" />
Getting form Pt op to impl is about 43 us which is dwarfed by other cpu overheads
<img width="1227" height="649" alt="Screenshot 2025-11-07 at 2 37 41 PM" src="https://github.com/user-attachments/assets/51da0615-facd-41e1-a6e2-fb7778079ab6" />
Just invoking the jit object from cutesl is 100s of us
<img width="545" height="414" alt="Screenshot 2025-11-07 at 2 38 19 PM" src="https://github.com/user-attachments/assets/d20345a0-6c47-4dcb-892f-9ef9894a1cf5" />
### Example usage
```Py
#!/usr/bin/env python3
"""Minimal FA4 smoke test for scaled dot product attention."""
from __future__ import annotations
import sys
from jsonargparse import CLI
import torch
import torch.nn.functional as F
from torch.nn.attention import (
install_flash_attention_impl,
sdpa_kernel,
SDPBackend,
)
def _map_dtype(kind: str) -> torch.dtype:
return torch.bfloat16 if kind == "bf16" else torch.float16
# To infinity and beyond
install_flash_attention_impl("FA4")
@sdpa_kernel([SDPBackend.FLASH_ATTENTION])
def main(
module_path: str = "flash_attn.cute.interface",
batch: int = 4,
seq: int = 81292,
heads: int = 16,
head_dim: int = 128,
device: int = 0,
dtype: str = "bf16"
) -> None:
if not torch.cuda.is_available():
sys.exit("CUDA is required for FA4 smoke testing")
torch.cuda.set_device(device)
dtype = _map_dtype(dtype)
generator = torch.Generator(device="cuda").manual_seed(0)
q = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
k = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
v = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
from transformer_nuggets.utils.benchmark import profiler
with profiler("sdpa_FA4", with_stack=False):
for _ in range(3):
out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
loss = out.real.sum()
loss.backward()
print("Scaled dot product attention output norm:", out.norm().item())
print("dq norm:", q.grad.norm().item())
if __name__ == "__main__":
CLI(main)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167348
Approved by: https://github.com/albanD, https://github.com/malfet
This was intersting to debug:
(1) Python 3.14 ships with a lazy way for retrieving annotations. The annotations field can be a callable that lazily evaluates it
(2) On the dynamo side, `SET_FUNCTION_ATTRIBUTE` needs to handle an extra flag value (0x10)
(3) The decorator `functools.wraps` used extensively in the codebase (e.g. `make_fx`, `functionalize`) doesn't copy the `__annotations__` attribute by default. To correctly retrieve an annotatin, we need to walk on the chain of `__wrapped__` objects and retrieve the attribute from the first function. Fortunately, there are functions on the stdlib to do this.
Fixes:
```
'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_fx_out_op_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_fx_transpose_simple_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_optional_tensorlist2_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_fx_multi_out_op_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_fx_reapply_views_simple_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_fx_simple_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_nonfunctional_output_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_opt_tensor_list_cpu', 'test/functorch/test_eager_transforms.py::TestFunctionalizeCPU::test_functionalize_optional_tensorlist1_cpu'
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167573
Approved by: https://github.com/williamwen42
As titled, this PR adds manual bucketing pass to SimpleFSDP. Users will need to parse FQNs they wanted to bucket together using `module_bucket_plans`. Then, `_manual_bucket_collectives` will get the node of the subgraphs correspond to each `bucket_module`, and bucket bucketable (FSDP-style) AG/RS together. `_manual_reorder_graph` reorders them for overlapping.
For detailed performance, see this torchtitan PR: https://github.com/pytorch/torchtitan/pull/1881.
There are a few todo items isted in torchtitan PR. Let's start with this PR that implements FSDP+TP+llama3 manual bucketing. I will fix/add the rest in follow up PRs.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165487
Approved by: https://github.com/ezyang
# Summary
See title ;)
## Design
Currently once you install there is no going back in the same python process, this need not be the case, cc @mikaylagawarecki's work on being able to grab original impl. I'll leave for follow up.
Okay I added an open reg, but I really want the backends to be found so some weird typing but we get
<img width="523" height="197" alt="Screenshot 2025-11-07 at 3 30 32 PM" src="https://github.com/user-attachments/assets/586de943-bbed-40cf-abd1-131f747a4cf1" />
## Overheads:
<img width="799" height="735" alt="Screenshot 2025-11-07 at 2 35 04 PM" src="https://github.com/user-attachments/assets/f9217f31-3e42-4816-8fb3-29ea8b49d735" />
First call to forward -> majority of time is spent in jit for FA
First call to backward, 3sec interestingly it doesn't appear that with_stack gets events in the backwards loop @albanD is this expected?
<img width="948" height="385" alt="Screenshot 2025-11-07 at 2 35 50 PM" src="https://github.com/user-attachments/assets/a40bacd0-3fb0-4bd8-b33e-bec8fb3f36c0" />
Getting form Pt op to impl is about 43 us which is dwarfed by other cpu overheads
<img width="1227" height="649" alt="Screenshot 2025-11-07 at 2 37 41 PM" src="https://github.com/user-attachments/assets/51da0615-facd-41e1-a6e2-fb7778079ab6" />
Just invoking the jit object from cutesl is 100s of us
<img width="545" height="414" alt="Screenshot 2025-11-07 at 2 38 19 PM" src="https://github.com/user-attachments/assets/d20345a0-6c47-4dcb-892f-9ef9894a1cf5" />
### Example usage
```Py
#!/usr/bin/env python3
"""Minimal FA4 smoke test for scaled dot product attention."""
from __future__ import annotations
import sys
from jsonargparse import CLI
import torch
import torch.nn.functional as F
from torch.nn.attention import (
install_flash_attention_impl,
sdpa_kernel,
SDPBackend,
)
def _map_dtype(kind: str) -> torch.dtype:
return torch.bfloat16 if kind == "bf16" else torch.float16
# To infinity and beyond
install_flash_attention_impl("FA4")
@sdpa_kernel([SDPBackend.FLASH_ATTENTION])
def main(
module_path: str = "flash_attn.cute.interface",
batch: int = 4,
seq: int = 81292,
heads: int = 16,
head_dim: int = 128,
device: int = 0,
dtype: str = "bf16"
) -> None:
if not torch.cuda.is_available():
sys.exit("CUDA is required for FA4 smoke testing")
torch.cuda.set_device(device)
dtype = _map_dtype(dtype)
generator = torch.Generator(device="cuda").manual_seed(0)
q = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
k = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
v = torch.randn(
batch,
heads,
seq,
head_dim,
device="cuda",
dtype=dtype,
requires_grad=True,
generator=generator,
)
from transformer_nuggets.utils.benchmark import profiler
with profiler("sdpa_FA4", with_stack=False):
for _ in range(3):
out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
loss = out.real.sum()
loss.backward()
print("Scaled dot product attention output norm:", out.norm().item())
print("dq norm:", q.grad.norm().item())
if __name__ == "__main__":
CLI(main)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167348
Approved by: https://github.com/albanD
Fixes#164663
## Issue
The torch model with multiple layers that is wrapped with fsdp2 registers pre and post forward hooks in a group using `_MultiHandler`. This becomes an issue during the context manager of the tracker where the hooks are reset and replaced. The hooks are all using the same fsdp state pointer so one reset will reset all. So when the output layer was modified with a new pre and post forward hook it would delete the previous layer's initialization causing `KeyError` for the Norm layer as it is nonexistent.
## The Fix
Check to see if there are multiple `_MultiHandler` objects and `RemoveHandler` objects and only execute the remove hook once.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165662
Approved by: https://github.com/sanketpurandare
In pip 25.3 and newer, use of --no-use-pep517 has been removed (https://pip.pypa.io/en/stable/news/). In builds with pip 25.2, a warning message notes:
> DEPRECATION: Building 'torchvision' using the legacy setup.py bdist_wheel mechanism, which will be removed in a future version. pip 25.3 will enforce this behaviour change. A possible replacement is to use the standardized build interface by setting the `--use-pep517` option, (possibly combined with `--no-build-isolation`), or adding a `pyproject.toml` file to the source tree of 'torchvision'. Discussion can be found at https://github.com/pypa/pip/issues/6334
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167096
Approved by: https://github.com/atalman
* Trigger this workflow on every completed run of `docker-builds.yml`
* Uses `ubuntu-latest` for downloading artifacts from `docker-build` workflow run
* Uses `linux.rocm.gfx942.docker-cache` to cache docker images as tarballs for MI3xx CI
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167554
Approved by: https://github.com/jeffdaily
Resolves issues running the dynamo cpython math.fma tests.
Though math.fma is enabled to perform a multiply add in dynamo, torch.addcmul is currently used which doesn't guarantee the user request for fma. It was decided to not use inductor fma prim as it would break the contract of using aten/core ir in dynamo output - otherwise export=True may have issues.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167217
Approved by: https://github.com/guilhermeleobas
Logging propagate_op_sharding_non_cached is a compromise between
- logging in DTensor.__init__ to catch ALL DTensor usage
- sparing the overhead in a latency-senstitive region like
DTensor.__init__
- and 'real' DTensor usage should incur at least one call to sharding
propagation
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167374
Approved by: https://github.com/zpcore
## Problems
When exporting a compiled model with nested input like below
```python
import torch
from torch.export import export, Dim
def test_export_compiled_model_with_nested_dynamic_shapes():
"""Test exporting a compiled model with nested dict inputs and dynamic shapes."""
print("Running test_export_compiled_model_with_nested_dynamic_shapes...")
class M(torch.nn.Module):
def forward(self, data_batch):
return data_batch["a1"] + data_batch["a2"]
m = M()
compiled_m = torch.compile(m)
example_args = ({
"a1": torch.ones(3, 3),
"a2": torch.ones(3, 3),
},)
dynamic_shapes = ({
"a1": {0: Dim.DYNAMIC},
"a2": {0: Dim.DYNAMIC},
},)
try:
ep = export(compiled_m, example_args, dynamic_shapes=dynamic_shapes, strict=True)
gm = ep.module()
result_exported = gm(*example_args)
result_compiled = compiled_m(*example_args)
assert torch.allclose(result_exported, result_compiled), "Results don't match!"
print("✓ test_export_compiled_model_with_nested_dynamic_shapes PASSED")
return True
except Exception as e:
print(f"✗ test_export_compiled_model_with_nested_dynamic_shapes FAILED")
print(f"Error: {e}")
import traceback
traceback.print_exc()
return False
def test_export_compiled_model_with_kwargs_dynamic_shapes():
"""Test exporting a compiled model with kwargs and dynamic shapes."""
print("\nRunning test_export_compiled_model_with_kwargs_dynamic_shapes...")
class M(torch.nn.Module):
def forward(self, a1, a2):
return a1 + a2
m = M()
compiled_m = torch.compile(m)
example_args = ()
example_kwargs = {
"a1": torch.ones(3, 3),
"a2": torch.ones(3, 3),
}
dynamic_shapes = {
"a1": {0: Dim.DYNAMIC},
"a2": {0: Dim.DYNAMIC},
}
try:
ep = export(compiled_m, example_args, kwargs=example_kwargs, dynamic_shapes=dynamic_shapes, strict=True)
gm = ep.module()
result_exported = gm(**example_kwargs)
result_compiled = compiled_m(**example_kwargs)
assert torch.allclose(result_exported, result_compiled), "Results don't match!"
print("✓ test_export_compiled_model_with_kwargs_dynamic_shapes PASSED")
return True
except Exception as e:
print(f"✗ test_export_compiled_model_with_kwargs_dynamic_shapes FAILED")
print(f"Error: {e}")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
print("Testing export of compiled models with dynamic shapes\n")
print("=" * 70)
results = []
results.append(test_export_compiled_model_with_nested_dynamic_shapes())
results.append(test_export_compiled_model_with_kwargs_dynamic_shapes())
print("\n" + "=" * 70)
print(f"\nResults: {sum(results)}/{len(results)} tests passed")
if all(results):
print("✓ All tests passed!")
else:
print("✗ Some tests failed")
exit(1)
```
It will report
```
======================================================================
Running test_export_compiled_model_with_nested_dynamic_shapes...
✗ test_export_compiled_model_with_nested_dynamic_shapes FAILED
Error: Detected mismatch between the structure of `inputs` and `dynamic_shapes`: `inputs[0]` is a <class 'tuple'>, but `dynamic_shapes[0]` is a <class 'dict'>
For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#dynamic-shapes-validation
The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
Traceback (most recent call last):
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 614, in _tree_map_with_path
return tree_map_with_path(f, tree, *dynamic_shapes, is_leaf=is_leaf)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/utils/_pytree.py", line 2055, in tree_map_with_path
all_keypath_leaves = keypath_leaves + [treespec.flatten_up_to(r) for r in rests]
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/utils/_pytree.py", line 2055, in <listcomp>
all_keypath_leaves = keypath_leaves + [treespec.flatten_up_to(r) for r in rests]
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/utils/_pytree.py", line 1188, in flatten_up_to
helper(self, tree, subtrees)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/utils/_pytree.py", line 1185, in helper
helper(subspec, subtree, subtrees)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/utils/_pytree.py", line 1141, in helper
raise ValueError(
ValueError: Node type mismatch; expected <class 'tuple'>, but got <class 'dict'>.
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/home/chzhu/infinitrain/test_exprot.py", line 25, in test_export_compiled_model_with_nested_dynamic_shapes
ep = export(compiled_m, example_args, dynamic_shapes=dynamic_shapes, strict=True)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/__init__.py", line 311, in export
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/__init__.py", line 277, in export
return _export(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1163, in wrapper
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1129, in wrapper
ep = fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
return fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 2255, in _export
ep = _export_for_training(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1163, in wrapper
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1129, in wrapper
ep = fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
return fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 2071, in _export_for_training
export_artifact = export_func(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1415, in _strict_export
gm_torch_level = _export_to_torch_ir(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 785, in _export_to_torch_ir
_check_dynamic_shapes(combined_args, dynamic_shapes)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 1031, in _check_dynamic_shapes
_tree_map_with_path(check_shape, combined_args, dynamic_shapes, tree_name="inputs")
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 686, in _tree_map_with_path
_compare(tree_spec, other_tree_spec, [])
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 677, in _compare
_compare(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 652, in _compare
raise_mismatch_error(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 634, in raise_mismatch_error
raise UserError(
torch._dynamo.exc.UserError: Detected mismatch between the structure of `inputs` and `dynamic_shapes`: `inputs[0]` is a <class 'tuple'>, but `dynamic_shapes[0]` is a <class 'dict'>
For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#dynamic-shapes-validation
The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
Running test_export_compiled_model_with_kwargs_dynamic_shapes...
✗ test_export_compiled_model_with_kwargs_dynamic_shapes FAILED
Error: When `dynamic_shapes` is specified as a dict, its top-level keys must be the arg names ['kwargs'] of `inputs`, but here they are ['a1', 'a2']. Since here `inputs` is a list/tuple enclosing a single dict, maybe you just forgot to enclose `dynamic_shapes` in a list/tuple?
For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#dynamic-shapes-validation
The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
Traceback (most recent call last):
File "/home/chzhu/infinitrain/test_exprot.py", line 62, in test_export_compiled_model_with_kwargs_dynamic_shapes
ep = export(compiled_m, example_args, kwargs=example_kwargs, dynamic_shapes=dynamic_shapes, strict=True)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/__init__.py", line 311, in export
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/__init__.py", line 277, in export
return _export(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1163, in wrapper
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1129, in wrapper
ep = fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
return fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 2255, in _export
ep = _export_for_training(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1163, in wrapper
raise e
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1129, in wrapper
ep = fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/exported_program.py", line 124, in wrapper
return fn(*args, **kwargs)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 2071, in _export_for_training
export_artifact = export_func(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1415, in _strict_export
gm_torch_level = _export_to_torch_ir(
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/_trace.py", line 785, in _export_to_torch_ir
_check_dynamic_shapes(combined_args, dynamic_shapes)
File "/home/chzhu/infinitrain/build/infinitrain/environments/development-venv/lib/python3.10/site-packages/torch/export/dynamic_shapes.py", line 1007, in _check_dynamic_shapes
raise UserError(
torch._dynamo.exc.UserError: When `dynamic_shapes` is specified as a dict, its top-level keys must be the arg names ['kwargs'] of `inputs`, but here they are ['a1', 'a2']. Since here `inputs` is a list/tuple enclosing a single dict, maybe you just forgot to enclose `dynamic_shapes` in a list/tuple?
For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#dynamic-shapes-validation
The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can replace your `export()` call with `draft_export()`.
======================================================================
```
## Torch Version
(reproducible nightly version)
## Other Behavior
The model can export regularly when we test without compiling the model
```python
import torch
from torch.export import export, Dim
def test_export_compiled_model_with_nested_dynamic_shapes():
"""Test exporting a compiled model with nested dict inputs and dynamic shapes."""
print("Running test_export_compiled_model_with_nested_dynamic_shapes...")
class M(torch.nn.Module):
def forward(self, data_batch):
return data_batch["a1"] + data_batch["a2"]
m = M()
example_args = ({
"a1": torch.ones(3, 3),
"a2": torch.ones(3, 3),
},)
dynamic_shapes = ({
"a1": {0: Dim.DYNAMIC},
"a2": {0: Dim.DYNAMIC},
},)
try:
ep = export(m, example_args, dynamic_shapes=dynamic_shapes, strict=True)
gm = ep.module()
result_exported = gm(*example_args)
result_compiled = m(*example_args)
assert torch.allclose(result_exported, result_compiled), "Results don't match!"
print("✓ test_export_compiled_model_with_nested_dynamic_shapes PASSED")
return True
except Exception as e:
print(f"✗ test_export_compiled_model_with_nested_dynamic_shapes FAILED")
print(f"Error: {e}")
import traceback
traceback.print_exc()
return False
def test_export_compiled_model_with_kwargs_dynamic_shapes():
"""Test exporting a compiled model with kwargs and dynamic shapes."""
print("\nRunning test_export_compiled_model_with_kwargs_dynamic_shapes...")
class M(torch.nn.Module):
def forward(self, a1, a2):
return a1 + a2
m = M()
example_args = ()
example_kwargs = {
"a1": torch.ones(3, 3),
"a2": torch.ones(3, 3),
}
dynamic_shapes = {
"a1": {0: Dim.DYNAMIC},
"a2": {0: Dim.DYNAMIC},
}
try:
ep = export(m, example_args, kwargs=example_kwargs, dynamic_shapes=dynamic_shapes, strict=True)
gm = ep.module()
result_exported = gm(**example_kwargs)
result_compiled = m(**example_kwargs)
assert torch.allclose(result_exported, result_compiled), "Results don't match!"
print("✓ test_export_compiled_model_with_kwargs_dynamic_shapes PASSED")
return True
except Exception as e:
print(f"✗ test_export_compiled_model_with_kwargs_dynamic_shapes FAILED")
print(f"Error: {e}")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
print("Testing export of compiled models with dynamic shapes\n")
print("=" * 70)
results = []
results.append(test_export_compiled_model_with_nested_dynamic_shapes())
results.append(test_export_compiled_model_with_kwargs_dynamic_shapes())
print("\n" + "=" * 70)
print(f"\nResults: {sum(results)}/{len(results)} tests passed")
if all(results):
print("✓ All tests passed!")
else:
print("✗ Some tests failed")
exit(1)
```
## Root Cause
This is because of a side effect of torch.compile(model). When the model is being compiled, the input signature will become (*args, **kwargs) automatically. In the above example, the `data_batch` will be added into `args` in combined_args [here](dc011d3203/torch/export/dynamic_shapes.py (L720)), and it will look like
```
{'args': ({'a1': tensor([[1., 1., 1.]... 1., 1.]]), 'a2': tensor([[1., 1., 1.]... 1., 1.]])},)}
```
Without the compiling, the combined args will look like
```
{'data_batch': {'a1': tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]), 'a2': tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]])}}
```
Thus causing the mismatch when we use treemap to match the dynamic shape with the input argos
The error is also reproducible when we setup kwargs as example argos (see the 2nd test above)
## Fix
Proposed fix: In [_combine_args](dc011d3203/torch/export/dynamic_shapes.py (L720)) we explicitly flatten out the kwargs and args into combined args.
## Side Effects
There are 2 existing tests that assume this behavior and
1. add `args` explicitly to dynamic shapes
2. wrap args into nested format in dynamic_shape
I have modified those test to make args and dynamic_shapes to be in consistent format.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166358
Approved by: https://github.com/angelayi
Summary: When cudagraphs-parittion is enabled, we have seen an increase of cold compile time in the vllm benchmark (see https://github.com/vllm-project/vllm/issues/27080). After some profiling, we found Triton compilation time increased the most. Further investigation reveals it was caused by duplicated Triton kernels not being shared among different partitions. This PR fixes the issue by reusing the Trition kernel source code cache at the top-level PythonWrapperCodegen.
In theory we could further reduce the compilation time by completely skipping compiling duplicated partitions. That can come as a furture improvement.
Some vllm benchmarking data,
```
VLLM_USE_STANDALONE_COMPILE=0 VLLM_DISABLE_COMPILE_CACHE=1 vllm bench latency -O.cudagraph_mode=PIECEWISE -O.use_inductor_graph_partition=True --model meta-llama/Meta-Llama-3.1-8
```
Before:
```
torch.compile takes 69.18 s in total
```
After:
```
torch.compile takes 26.81 s in total
```
As a refrence, this is the compile time when turning off inductor graph partition. Looks like we still have some gap to close.
```
VLLM_USE_STANDALONE_COMPILE=0 VLLM_DISABLE_COMPILE_CACHE=1 vllm bench latency -O.cudagraph_mode=PIECEWISE -O.use_inductor_graph_partition=False --model meta-llama/Meta-Llama-3.1-8B
torch.compile takes 19.41 s in total
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167132
Approved by: https://github.com/eellison
ghstack dependencies: #167131
Summary: When cudagraphs partition and autotune_at_compile_time are enabled, currently each subgraph will generate its own auto-tuning code block and run them once by one. This PR improves it by only generating one auto-tuning code block at the main graph level and execute it once time to auto-tune all the kernels.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167131
Approved by: https://github.com/eellison
Avoids data-dependent errors for out-of-bounds & redundant slice checks.
The sharding logic that immediately depends on this only checks for redundant slices, and is saying: "it's safe to reuse the input placements if a) the slicing dimension isn't sharded, or b) the slice is redundant, so just pretend this op didn't happen".
This has a slight effect on output placements, when a slice is performed on a shared dim, and dynamic shapes are involved (size/start/end/step). Now if the slice isn't obviously redundant, we won't immediately consider the input placements valid (even if they could be for very particular runtime shapes), and select strategies valid for the general case - in this case I guess unsharding the slicing dim.
For backed symbols, we could choose to recompile when the redundant case is hit, by switching to `guard_or_false`, but it's not obvious how desirable this is.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166990
Approved by: https://github.com/laithsakka
This PR enables XPU devices in test_analysis.py.
For performance reason, it skips some slow tests, so a full scope should be enabled by using:
```
export PYTORCH_TEST_WTH_SLOW=1
```
**PR Stack:**
- https://github.com/pytorch/pytorch/pull/166840 : This PR enables the tests, ignores the tests that failed
- https://github.com/pytorch/pytorch/pull/166839 : This fixed the bug and enable the full tests for xpu
**Some skipped test time:**
```
test_augment_trace_against_flop_counter_maxat0_xpu_float16 [49.0863s]
test_augment_trace_against_flop_counter_maxat0_xpu_float32 [18.2268s]
test_augment_trace_against_flop_counter_maxat1_xpu_float16 [85.6549s]
test_augment_trace_against_flop_counter_maxat1_xpu_float32 [329.0832s]
test_augment_trace_against_flop_counter_maxat2_xpu_float16 [24.4825s]
test_augment_trace_against_flop_counter_maxat2_xpu_float32 [19.0688s]
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166840
Approved by: https://github.com/guangyey, https://github.com/jansel
Fixed a command injection vulnerability in PreCompiled Header (PCH) compilation where extra_cflags were passed to subprocess with shell=True, allowing arbitrary command execution through malicious compiler flags.
Changed subprocess.check_output(pch_cmd, shell=True) to use shlex.split() to safely parse the command without shell interpretation. This prevents shell metacharacters (;, |, &, etc.) in extra_cflags from being executed as shell commands.
Added test case test_pch_command_injection that verifies:
1. PCH compilation attempts with malicious payloads in extra_cflags
2. Shell commands embedded in flags are not executed
3. Exploit file is not created, proving no shell execution occurred
Note: On RHEL/Fedora and other systems with versioned GCC compilers, the test depends on https://github.com/pytorch/pytorch/pull/167501 being merged first, otherwise the test will be skipped due to GCC detection issues.
Fixes#167480
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167502
Approved by: https://github.com/malfet
Allocate ErrorMessages buffer associated with MPSStream and introduce `c10:🤘:report_error` method(and `TORCH_REPORT_ERROR` macro), that can be used to preserve up to `c10:🤘:error_message_count` messages
Add test that detects those
As results attempt to run something like
```python
import torch
x=torch.rand(10, 1, 10, device='mps')
y=x[:, [1]]
torch.mps.synchonize()
```
will raise `torch.AcceleratorError: index 1 is out of bounds for dimension 0 with size 1`
Fixes https://github.com/pytorch/pytorch/issues/111669
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166615
Approved by: https://github.com/manuelcandales, https://github.com/dcci
ghstack dependencies: #167444, #167445
The function was only returning True for compilers named 'c++', but failed to detect g++, gcc, g++-13, g++-14, etc. This fixes the detection to work for any GCC variant by checking for both COLLECT_GCC and 'gcc version' in the compiler output.
The previous implementation used os.path.basename() on the resolved compiler path and only checked if it exactly matched 'c++'. This caused false negatives for versioned GCC installations and direct g++ usage.
The fix simplifies the logic: if both COLLECT_GCC is present in the output (indicating GCC toolchain) and 'gcc version' appears in the version string, it's GCC.
Fixes#167499
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167501
Approved by: https://github.com/ezyang, https://github.com/malfet
Summary:
When PT2 sees an operation like `torch.empty_strided(8, device="cuda")` it returns a fake tensor on a specific device index, such as `"cuda:0"`. This is implemented via a list of multi-index devices in the fake tensor constructor. If a device supports indices but is not in this list, we can hit fake tensor propagation errors as `"cuda"` is not considered to be the same device as `"cuda:0"`.
This PR adds the `"mtia"` device to this list, to resolve some fake tensor propagation errors we're seeing with the full Inductor backend. (Internal task: T243176905)
Test Plan: Tests are stacked on the internal diff D86605248.
Reviewed By: StellarrZ
Differential Revision: D86605248
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167457
Approved by: https://github.com/eellison, https://github.com/mlazos
Summary: fix some type errors + instead of manually creating a filelock when dumping dcache's imc to file we simply use an odc (since this is the intended behavior of odc, anyways)
Test Plan:
```
buck test fbcode//mode/opt caffe2/test/inductor:caching
```
Differential Revision: D86345594
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167136
Approved by: https://github.com/aorenste
Prior to this PR, the IValue <-> StableIValue conversion for `DeviceObjType` (aka c10::Device) was to pack it into the leading bits of the StableIValue (which is a uint64_t)
After this PR, the IValue <-> StableIValue conversion for `DeviceObjType` expects DeviceType to be packed into the upper 32 bits of StableIValue and DeviceIndex to be packed into the lower 32 bits
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166579
Approved by: https://github.com/janeyx99
Make the PyObject preservation scheme thread-safe with free threaded (nogil) Python. The general idea is:
* Python Tensor and Storage objects always hold a strong reference to their underlying c10 object
* c10 objects hold a strong reference to their Python objects if there's at least one other reference to the c10 object
This is implemented in `intrusive_ptr`:
* The top most bit (`kHasPyObject`) from the weakref count is now used to indicate if the `intrusive_ptr_target` has an associated PyObject. So `kHasPyObject` is one bit, the weakref count is now 31 bits and the strong refcount remains 32 bits.
* When the reference count increases from one to two and `kHasPyObject` is set, we incref the associated Python object to ensure that it's kept alive.
* When the reference count decreases from two to one (i.e., there are no C++ reference to the `intrusive_ptr_target` other than from the Python object), we decre the associated Python object to break the cycle.
Other benefits:
* We can delete a lot of the copypasta from Python internal `subtype_dealloc`
* This fixes the weakref and GC bugs we had in the previous scheme. Python weakrefs on Tensors and Storages should just work as expected now.
Risks:
* Extra branch for reference count operations on `intrusive_ptr<TensorImpl>`, `intrusive_ptr<StorageImpl>`, and the generic `intrusive_ptr<intrusive_ptr_target>` even when we're not using Python.
* It's a big change
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166342
Approved by: https://github.com/albanD
Preferred solution to #166380
Changes:
- Moved summary print to bottom of CMakeLists.txt
- Fix the problem 'add_compile_options' should be called before targets defined, so opted for `append_cxx_flag_if_supported` and `append_c_flag_if_supported` ( new ).
- Added extra verbosity so it can be seen when linker script added.
( unfortunately linker script has to be added per-target rather than globally due to ninja/cmake depdendency tracking ).
Also move summary print to bottom of CMakeLists.txt and improve logging
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166407
Approved by: https://github.com/Aidyn-A, https://github.com/atalman
Summary:
https://github.com/pytorch/pytorch/pull/166609 updates `node.is_impure` to consider a submodule as impure if submodule contains impure node. This in turn changes `graph.eliminate_dead_code()` function behavior, which does not eliminate nodes with side effects, see [pytorch documentation](https://docs.pytorch.org/docs/stable/fx.html#torch.fx.Graph.eliminate_dead_code)
> Remove all dead code from the graph, based on each node’s number of users, and whether the nodes have any side effects.
While this is correct that a submodule containing side-effectful ops is side-effectful and should not be dead code eliminated, some customers rely on the dead code elimination to eliminate submodules that contain impure ops which is the behavior before #166609 fix.
Due to production environment constraints, we have to revert https://github.com/pytorch/pytorch/pull/166609 and move the side-effectful submodule check logic to `const_fold.py`, which will correctly **not** const-fold a submodule that contains impure ops.
NOTE other call sites that use `node.is_impure()` to make decisions are still incorrectly eliminating side-effectful submodules, but we can't safely change that today.
## This pr
- move `_subgraph_has_impure_op` into `fx/experimental/const_fold.py`, check and prevent const-folding an impure submodule
- added a note in `node.is_impure` to highlight the incorrect behavior and context in case people go looking in the future.
Test Plan: run test_fx_const_fold and all tests pass
Differential Revision: D86641994
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167443
Approved by: https://github.com/jfix71
We plan to use `StridedShard` to express `shard_order`. This PR adds the function to support the conversion between `StridedShard` and `shard_order`.
I moved some test related function into torch/testing/_internal/common_utils.py. We may only care about **_dtensor_spec.py** and **test_utils.py** in this PR for the review.
### How to convert shard order to StridedShard:
Considering the example:
- placements = $[x_0, x_1, x_2, x_3, x_4]$, all $x_?$ are shard on the same tensor dim.
Let's see how the shard order will impact the split_factor (sf). We loop from right to left in the placements to construct the split_factor by assuming different shard order. Starting from $x_4$, this should be a normal shard.
Then $x_3$. There are two possibilities, $x_3$'s order can be before $x_4$. If so, $x_3$'s sf=1, because $x_3$ is before $x_4$ in the placements. Else $x_3$'s order is after $x_4$, then the $x_3$'s sf should be the mesh dim size of $x_4$, which is $T(x_4)$:
<img width="820" height="431" alt="image" src="https://github.com/user-attachments/assets/f53b4b24-2523-42cc-ad6f-41f3c280db70" />
We can use this method to decide on the split factor for $x_2$, $x_1$ and so on.
### How to convert StridedShard to shard order:
This follows the same method above. We check all possible paths and use the real split_factor to see which path matchs the split_factor. If no such matches, the StridedShard is unable to be converted to shard order.
---
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166740
Approved by: https://github.com/ezyang
Pass `dim_map` to `_requires_data_exchange` and return False if both spatial and channels dimensions are replicated
Modify `test_conv1d` and `test_conv3d` to check values rather than just shape, and replicate `conv3d` across batch dimension
In general, feels like current Convolution implementation was written to work only if tensor is sharded across last dimention
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167402
Approved by: https://github.com/ezyang
- The logger name in test_fully_shard_logging.py was wrong so the logs didn't happen.
- The `device` variable in test_fully_shard_logging is expected to be a string, so quote it
- `unittest.skipIf` is used so importing `unittest` instead of `unittest.mock` is required
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167312
Approved by: https://github.com/Skylion007, https://github.com/cyyever
## Summary
- Fixes#163548 by replacing the quadratic chunk-overlap scan in `_validate_global_plan` with a sweep-line pass that sorts chunk intervals and keeps an active set via `bisect_right`, giving O(n log n) behavior for metadata validation.
- Add focused tests in `TestValidateGlobalPlan` covering overlapping and non-overlapping shard layouts to lock in the faster path.
## Testing
- python test/distributed/checkpoint/test_planner.py -k ValidateGlobalPlan
## Benchmarks
| chunks | old runtime | new runtime |
|--------|-------------|-------------|
| 1 024 | 0.121 s | 0.0014 s |
| 2 048 | 0.486 s | 0.0027 s |
| 4 096 | 2.474 s | 0.0058 s |
| 8 192 | 8.014 s | 0.0126 s |
| 16 384 | 32.740 s | 0.026 s |
@ezyang
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166820
Approved by: https://github.com/LucasLLC, https://github.com/Skylion007
Summary:
It seems
D80948073
has caused some issue on a lowering pkg built on trunk: https://fburl.com/mlhub/o6p60pno
error log: P2001933683
which we were able to lower successfully in older ien pkg: https://fburl.com/mlhub/1ro094zo
D85172267 fixed this issue for the if conditional, but issue still exists for the else conditional. Logic is moved right before if-else to cover both cases
Test Plan:
checkout D85605372
buck2 run -c fbcode.enable_gpu_sections=true -c fbcode.platform=platform010 -c fbcode.nvcc_arch=a100,h100 -c fbcode.split-dwarf=true -c fbcode.dwp=true -c fbcode.enable_distributed_thinlto=true -c fbcode.use_link_groups=true fbcode//inference_enablement/model_processing/infra/components/lowering/re:re_cinder -- -r "$(cat ./fbcode/minimal_viable_ai/umia_v1/ig/ss_omni_exp/re_lower_aoti.json)"
with the diff, no issue was encountered.
Reviewed By: tissue3
Differential Revision: D86474796
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167279
Approved by: https://github.com/pianpwk
This adds zero-bubble / DualPipeV support for (S)AC
Before:
- AC will always retrigger recompute upon every distinct backward.
After:
- Any checkpointed regions encountered by backward under the same instance of this context manager will only trigger recompute at most once, even if there are multiple calls to backward.
- Backward calls under the same instance of this context manager must execute over non-overlapping regions of the backward graph even if retain_graph=True.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166536
Approved by: https://github.com/albanD
**Summary:** While the implementation is correct, these checks are just a subset of the Partial placement checks that are done in https://github.com/pytorch/pytorch/pull/165962. This means for ops aten.linalg_vector_norm.default and aten._foreach_norm.Scalar, we're unnecessarily checking for Partial placements twice.
**Test Cases**
1. pytest test/distributed/tensor/test_math_ops.py -k test_vector_norm_partial
2. pytest test/distributed/tensor/test_math_ops.py -k test_foreach_norm_partial
3. pytest test/distributed/tensor/test_math_ops.py -k test_partial_reduction_ops
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167247
Approved by: https://github.com/XilunWu
# Problem
The FX backend currently allocates tensors on an exact device index, such as `"cuda:0"`. In contrast, the Python backend allocates on a device type, such as `"cuda"`. This avoids edge cases where fake tensor propagation can fail due to mismatched devices.
# Fix
Allocate tensors on `device.type` instead of the device.
# Test plan
Added a CI test passing in sample inputs on an indexed device, and checking that the output device in the generated FX graph is not indexed.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167358
Approved by: https://github.com/mlazos, https://github.com/nandesuka, https://github.com/eellison
Not sure why, but running some tests (for example `test_weights_only_safe_globals_build`) with `pytest` in 3.14 makes global name `test_serialization.ClassThatUsesBuildInstruction` instead of expected `__main__.ClassThatUsesBuildInstruction`
Also, change expected exception type from `AttributeError` to `PicklingError`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167333
Approved by: https://github.com/atalman
Some important notes:
a) Just like IValues steal the ownership of ArrayRefs and any std::vectors in order to convert the inner elements into IValues, we do the same thing with StableIValue. This O(N) traverse is ineluctable.
b) As a result, since StableIValues are owning and our contract is that to<T>(StableIValue) transfers ownership, you cannot ever convert from StableIValue to a nonowning HeaderOnlyArrayRef<V>.
We handle memory similar to AtenTensorHandle, but we have a StableListHandle!
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165953
Approved by: https://github.com/malfet
ghstack dependencies: #164991, #165152, #165153
Summary:
If quiesce ends up called twice (which is likely not possible with the timer based implementation, but possible with either manual calls, or with the context manager implementation), this assertion fires.
Instead make this assertion tolerant to rentrant calling of quiesce
Test Plan: Added a explicit test which calls quiesce twice.
Differential Revision: D86251534
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167025
Approved by: https://github.com/masnesral
Provides type coverage to torch/_dynamo/variables/dicts.py
Coverage report:
`mypy torch/_dynamo/variables/functions.py --linecount-report /tmp/coverage_log`
Compare before to after - we go from 0 lines and 0 funcs covered to 2698 lines and 166 funcs covered
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167103
Approved by: https://github.com/mlazos, https://github.com/fxdawnn
This fixes 2 issues with the DTensor data-dependent test case:
1) ShapeEnv not found when doing shard prop on data-dependent ops - fix was to detect the outer tracing fake mode. Maybe ShardingPropagator should just own a FakeMode & ShapeEnv for these purposes? The previous behavior was to initialize a new fake mode on every call.
2) Pending unbacked symbols not found. This happens because DTensor dispatch runs fake prop twice, once while figuring out the output sharding: 2bba37309b/torch/distributed/tensor/_sharding_prop.py (L175) and again to actually get the resulting local tensor: 2bba37309b/torch/distributed/tensor/_dispatch.py (L254-L255) With data-dependent ops, both calls will produce an unbacked symbol, but symbols in the first invocation are never surfaced, producing this error, so we ignore pending symbols from this site.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166989
Approved by: https://github.com/ezyang
TORCH_CHECK is non-fatal by design, but TORCH_CHECK_{COND} macros are fatal. this is confusing, and we should limit fatality to the set of debug macros.
Differential Revision: D86168955
Pull Request resolved: https://github.com/pytorch/pytorch/pull/167004
Approved by: https://github.com/malfet
Fixes#163971
Problem:
PyTorch's inductor compiler crashed with IndexError: list index out of range when compiling code that uses 0-dimensional tensors with operations like torch.softmax(scalar_tensor, dim=0).
A 0-dim tensor has shape = torch.Size([]) (empty shape)
```
ndim = 0 (zero dimensions)
len(shape) = 0 (no indices to access)
# Line 972: Pad other_shape to match inp dimensions
other_shape = [1] * (inp_ndim - len(other_shape)) + list(other_shape)
# For scalar tensors:
# inp_ndim = 0 # as input is scalar
# other_shape = []
# Result: [1] * (0 - 0) + [] = [] (still empty!)
dim = match.kwargs["dim"] # dim = 0
if isinstance(dim, int):
dim = (dim,)
# crash is happening here!
return all(statically_known_true(other_shape[d] == 1) for d in dim)
# ^^^^^^^^^^^^^^^^
# Tries other_shape[0] but other_shape = [] (empty!)
# → IndexError: list index out of range
```
The function _other_is_broadcasted_in_dim() is an optimization check for a softmax fusion pattern. It verifies whether it's safe to rewrite:
```
# From
scaled = inp * other
result = scaled - scaled.amax(dim, keepdim=True)
# To this more stable form:
result = (inp - inp.amax(dim, keepdim=True)) * other
```
The optimization is only valid if other is constant across the reduction dimension (i.e., broadcasted to size 1 in that dimension). Otherwise, scaling changes which element is the maximum.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166547
Approved by: https://github.com/jansel
@ -18,6 +18,8 @@ Please report security issues using https://github.com/pytorch/pytorch/security/
All reports submitted through the security advisories mechanism would **either be made public or dismissed by the team within 90 days of the submission**. If advisory has been closed on the grounds that it is not a security issue, please do not hesitate to create an [new issue](https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml) as it is still likely a valid issue within the framework.
**Note on crashes and out of bounds access**: PyTorch is a computational framework that performs operations on behalf of the caller. Like many low-level libraries, PyTorch generally does not validate all inputs to every function—the responsibility for providing valid arguments lies with the calling code. While crashes and out of bounds memory access should be reported as bugs, they are generally not considered security vulnerabilities in PyTorch's threat model.
Please refer to the following page for our responsible disclosure policy, reward guidelines, and those things that should not be reported:
@ -52,19 +52,18 @@ def test_sparse_coo_and_csr(m, n, k, nnz, test_count):
start.record()
coo.matmul(mat)
stop.record()
times.append(start.elapsed_time(stop))
coo_mean_time=sum(times)/len(times)
coo_mean_time=sum(times)/len(times)
times=[]
for_inrange(test_count):
start.record()
csr.matmul(mat)
stop.record()
times.append(start.elapsed_time(stop))
times=[]
for_inrange(test_count):
start.record()
csr.matmul(mat)
stop.record()
times.append(start.elapsed_time(stop))
csr_mean_time=sum(times)/len(times)
csr_mean_time=sum(times)/len(times)
returncoo_mean_time,csr_mean_time
@ -84,10 +83,13 @@ if __name__ == "__main__":
ifargs.outfile=="stdout":
outfile=sys.stdout
need_close=False
elifargs.outfile=="stderr":
outfile=sys.stderr
need_close=False
else:
outfile=open(args.outfile,"a")
need_close=True
test_count=args.test_count
m=args.m
@ -148,3 +150,5 @@ if __name__ == "__main__":
time,
file=outfile,
)
ifneed_close:
outfile.close()
Some files were not shown because too many files have changed in this diff
Show More
Reference in New Issue
Block a user
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.
