Wanchao Liang
2ee6b97464
[dtensor] move DTensor to public namespace ( #133113 )
...
Moving DTensor to be in the public namespace, to formally add the
documentation page that includes all the public APIs. This includes:
* many path renames and path import fixes
* a dedicated doc page without too much content yet (adding in the next
PRs)
* To preserve the BC for users still using the `torch.distributed._tensor`,
I added a shim script to redirect old path calls to the new module
The BC preserving is evidented by the fact that all DTensor tests are still
working without changing the public imports. So it's safe to land the
changes
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133113
Approved by: https://github.com/XilunWu
ghstack dependencies: #133305 , #133306
2024-08-17 05:09:52 +00:00
xinyu-intel
5ae979ab10
[Dynamo] Support torch.autograd._is_checkpoint_valid ( #132611 )
...
Hi, we got `torch._dynamo.exc.Unsupported: torch.* op returned non-Tensor bool call_function <function _is_checkpoint_valid at 0x7f0b0d22e290>` while tracing activation [checkpointing function in deepspeed](324ee65cb0/deepspeed/runtime/activation_checkpointing/checkpointing.py (L630) ). Consider to add it to constant_folding list which is similar with https://github.com/pytorch/pytorch/pull/126196
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132611
Approved by: https://github.com/anijain2305 , https://github.com/williamwen42
2024-08-08 04:05:08 +00:00
Oguz Ulgen
6e79932543
Add basic mypy annotations to dynamo ( #132415 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132415
Approved by: https://github.com/XuehaiPan , https://github.com/jamesjwu
2024-08-04 18:43:36 +00:00
PyTorch MergeBot
3558a8cf4a
Revert "Add basic mypy annotations to dynamo ( #132415 )"
...
This reverts commit 71e22e0959 .
Reverted https://github.com/pytorch/pytorch/pull/132415 on behalf of https://github.com/ZainRizvi due to Sorry, this PR has entered a weird state in the diff train. Trying to revert it to skip it, and then we can try relanding it ([comment](https://github.com/pytorch/pytorch/pull/132415#issuecomment-2267631785 ))
2024-08-04 18:39:29 +00:00
Oguz Ulgen
71e22e0959
Add basic mypy annotations to dynamo ( #132415 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132415
Approved by: https://github.com/XuehaiPan , https://github.com/jamesjwu
2024-08-01 20:14:25 +00:00
Xuehai Pan
e74ba1b34a
[BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ ( #129767 )
...
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501 . Most changes are auto-generated by linter.
You can review these PRs via:
```bash
git diff --ignore-all-space --ignore-blank-lines HEAD~1
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/129767
Approved by: https://github.com/anijain2305
2024-07-31 21:18:11 +00:00
Oguz Ulgen
7a42470bcb
Annotate all InstructionTranslator ( #131509 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131509
Approved by: https://github.com/zou3519
2024-07-24 23:45:53 +00:00
PyTorch MergeBot
5db5865614
Revert "Annotate all InstructionTranslator ( #131509 )"
...
This reverts commit eafbd20f23 .
Reverted https://github.com/pytorch/pytorch/pull/131509 on behalf of https://github.com/clee2000 due to sorry need to revert this to revert something else, I think you only need to rebase and remerge ([comment](https://github.com/pytorch/pytorch/pull/131509#issuecomment-2249000843 ))
2024-07-24 22:29:49 +00:00
Oguz Ulgen
b56939dae1
Annotate more InstructionTranslator ( #131680 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131680
Approved by: https://github.com/zou3519
ghstack dependencies: #131676
2024-07-24 22:14:29 +00:00
Oguz Ulgen
eafbd20f23
Annotate all InstructionTranslator ( #131509 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131509
Approved by: https://github.com/zou3519
2024-07-24 05:31:01 +00:00
Animesh Jain
fa4e489d70
[dynamo][dynamic-shapes] Graph break if out shape changes on out= variants ( #130074 )
...
Fixes https://github.com/pytorch/pytorch/issues/130068
Pull Request resolved: https://github.com/pytorch/pytorch/pull/130074
Approved by: https://github.com/ezyang
ghstack dependencies: #129913 , #129914
2024-07-04 08:36:12 +00:00
William Wen
79aabaf626
[3.13, dynamo] codegen PUSH_NULL when callable is codegen'd ( #129172 )
...
Significant bytecode generation API change!
The new suggested convention to generating bytecode to call a function is now to wrap instructions that push a callable to the stack with `add_push_null`, then that callable is called with `create_call_function` with `push_null=False` (see diff for examples).
In Python 3.13, NULL is now expected to be pushed after the callable. In <=3.12, the NULL was pushed before the callable. This change abstracts away the exact placement of the NULL, but the developer must be aware that a NULL may be needed when codegen'ing a callable.
This abstraction also reduces the need for the `push_null=True` option in `create_call_function`, which removes the need to rotate a NULL to the right place on the stack with a sequence of `SWAP` instructions.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/129172
Approved by: https://github.com/jansel
2024-06-22 17:25:23 +00:00
Brian Hirsh
8c2542623b
[Traceable FSDP2] [Dynamo] Add tracing support for out-variant custom ops that return None ( #129078 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/129078
Approved by: https://github.com/yanboliang
2024-06-20 17:46:13 +00:00
Animesh Jain
bdffd9f0c6
[export] Graph break on nn.Parameter construction ( #128935 )
...
Fixes https://github.com/pytorch/pytorch/issues/126109
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128935
Approved by: https://github.com/angelayi
2024-06-18 18:37:44 +00:00
Animesh Jain
b0282071c4
[dynamo] override torch.nn.modules.activation._is_make_fx_tracing ( #128748 )
...
Discovered while inlining `MultiHeadAttention` nn Module.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128748
Approved by: https://github.com/jansel
ghstack dependencies: #128315
2024-06-17 08:49:29 +00:00
Will Feng
979edbbe12
[Traceable FSDP2] Dynamo support FSDP2 use_training_state context manager ( #127854 )
...
Improve Dynamo to support the FSDP2 `use_training_state()` context manager.
Test command:
`
pytest -rA test/distributed/_composable/fsdp/test_fully_shard_compile.py::TestFullyShardCompile::test_dynamo_trace_use_training_state
`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127854
Approved by: https://github.com/yanboliang
2024-06-16 08:48:52 +00:00
Animesh Jain
c0b87afcad
[RELAND2][dynamo][nn-modules] Trace through nn.Module dunder methods for UnspecializedNNModule ( #126578 )
...
Tracing through `__init__` is important because it initializes (calls STORE_ATTR) on members. By doing that, we kick in the mutation tracking for these objects. So, things like mutating `_modules` etc is tracked automatically.
Fixes https://github.com/pytorch/pytorch/issues/111837
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126578
Approved by: https://github.com/jansel
2024-06-12 04:09:23 +00:00
PyTorch MergeBot
adb699189b
Revert "[RELAND][dynamo][nn-modules] Trace through nn.Module dunder methods for UnspecializedNNModule ( #126578 )"
...
This reverts commit b2d602306a .
Reverted https://github.com/pytorch/pytorch/pull/126578 on behalf of https://github.com/clee2000 due to failed internal test D58394084. Author has forward fix but includes external changes so reverting is a bit easier to coordinate ([comment](https://github.com/pytorch/pytorch/pull/126578#issuecomment-2161481839 ))
2024-06-11 19:41:41 +00:00
Animesh Jain
b2d602306a
[RELAND][dynamo][nn-modules] Trace through nn.Module dunder methods for UnspecializedNNModule ( #126578 )
...
Tracing through `__init__` is important because it initializes (calls STORE_ATTR) on members. By doing that, we kick in the mutation tracking for these objects. So, things like mutating `_modules` etc is tracked automatically.
Fixes https://github.com/pytorch/pytorch/issues/111837
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126578
Approved by: https://github.com/jansel
ghstack dependencies: #128295
2024-06-10 23:11:04 +00:00
Aaron Orenstein
dcfa7702c3
Flip default value for mypy disallow_untyped_defs [1/11] ( #127838 )
...
See #127836 for details.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127838
Approved by: https://github.com/oulgen
2024-06-08 18:16:33 +00:00
PyTorch MergeBot
44371bd432
Revert "[dynamo][nn-modules] Trace through nn.Module dunder methods for UnspecializedNNModule ( #126578 )"
...
This reverts commit 7ede78f9f5 .
Reverted https://github.com/pytorch/pytorch/pull/126578 on behalf of https://github.com/anijain2305 due to pippy tests fail ([comment](https://github.com/pytorch/pytorch/pull/126578#issuecomment-2155836555 ))
2024-06-08 06:35:34 +00:00
Animesh Jain
7ede78f9f5
[dynamo][nn-modules] Trace through nn.Module dunder methods for UnspecializedNNModule ( #126578 )
...
Tracing through `__init__` is important because it initializes (calls STORE_ATTR) on members. By doing that, we kick in the mutation tracking for these objects. So, things like mutating `_modules` etc is tracked automatically.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126578
Approved by: https://github.com/jansel
ghstack dependencies: #128001
2024-06-06 23:05:49 +00:00
Animesh Jain
e5b3387166
[dynamo] Bugfix for nn parameter construction ( #128001 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128001
Approved by: https://github.com/jansel
2024-06-06 23:05:49 +00:00
PyTorch MergeBot
6dc0a291b9
Revert "[dynamo] Bugfix for nn parameter construction ( #127806 )"
...
This reverts commit f27c4dd862 .
Reverted https://github.com/pytorch/pytorch/pull/127806 on behalf of https://github.com/PaliC due to causing nn tests to fail ([comment](https://github.com/pytorch/pytorch/pull/127806#issuecomment-2148393903 ))
2024-06-04 20:51:41 +00:00
Animesh Jain
f27c4dd862
[dynamo] Bugfix for nn parameter construction ( #127806 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127806
Approved by: https://github.com/jansel
ghstack dependencies: #127785 , #127802
2024-06-04 18:25:46 +00:00
Yu, Guangye
e7a42702f9
generalize custom_fwd&custom_bwd to be device-agnostic ( #126531 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126531
Approved by: https://github.com/jgong5 , https://github.com/gujinghui , https://github.com/albanD , https://github.com/EikanWang
ghstack dependencies: #126527
2024-05-25 06:48:16 +00:00
Yanbo Liang
51ed4c46cf
[Dynamo] Supports torch._C._is_any_autocast_enabled ( #126196 )
...
Fixes #126026
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126196
Approved by: https://github.com/anijain2305
2024-05-15 03:16:13 +00:00
Xuehai Pan
7e1c98c171
[dynamo] support object.__setattr__(obj, name, value) ( #124068 )
...
Resolves #114964
Resolves #114966
- #114964
- #114966
Pull Request resolved: https://github.com/pytorch/pytorch/pull/124068
Approved by: https://github.com/jansel
2024-04-17 15:57:14 +00:00
Edward Z. Yang
b36b523c05
Fix guard_size_oblivious on non-symbolic expression ( #123743 )
...
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123743
Approved by: https://github.com/avikchaudhuri
2024-04-10 22:45:54 +00:00
Guilherme Leobas
4eaa000acc
Teach dynamo about torch.func.jvp ( #119926 )
...
List of changes:
- Replace JVP_NESTING by torch._C._functorch.maybe_current_level()
- Remove all increment nesting functions from wrap_fx_proxy_cls
- fwAD.make_dual receives the dual_level as keyword argument
- Add jvp_increment_nesting, set_fwd_grad_enabled and dual_level context managers to dynamo
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119926
Approved by: https://github.com/zou3519
2024-03-22 20:25:47 +00:00
Peter Bell
5790096059
[dynamo] Remove uses of raise unimplemented ( #122136 )
...
`unimplemented` is a function that raises an error, so
`raise unimplemented(...)` never reaches the `raise`.
Another related issue is that `raise unimplemented(...) from e`
doesn't attach the exception cause correctly. I fix this by adding
a `from_exc` argument to `unimplemented`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122136
Approved by: https://github.com/lezcano
2024-03-22 19:29:58 +00:00
PyTorch MergeBot
0696db8202
Revert "Teach dynamo about torch.func.jvp ( #119926 )"
...
This reverts commit 17489784b6 .
Reverted https://github.com/pytorch/pytorch/pull/119926 on behalf of https://github.com/peterbell10 due to broken mac jobs on main ([comment](https://github.com/pytorch/pytorch/pull/119926#issuecomment-2010327997 ))
2024-03-20 18:34:43 +00:00
Guilherme Leobas
17489784b6
Teach dynamo about torch.func.jvp ( #119926 )
...
List of changes:
- Replace JVP_NESTING by torch._C._functorch.maybe_current_level()
- Remove all increment nesting functions from wrap_fx_proxy_cls
- fwAD.make_dual receives the dual_level as keyword argument
- Add jvp_increment_nesting, set_fwd_grad_enabled and dual_level context managers to dynamo
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119926
Approved by: https://github.com/zou3519
2024-03-20 13:09:19 +00:00
PyTorch MergeBot
36e5c1dcab
Revert "Teach dynamo about torch.func.jvp ( #119926 )"
...
This reverts commit edd04b7c16 .
Reverted https://github.com/pytorch/pytorch/pull/119926 on behalf of https://github.com/jeanschmidt due to lots of breakages in pull jobs, checking if reverting this one will help ([comment](https://github.com/pytorch/pytorch/pull/119926#issuecomment-2007915919 ))
2024-03-19 18:59:46 +00:00
Guilherme Leobas
edd04b7c16
Teach dynamo about torch.func.jvp ( #119926 )
...
List of changes:
- Replace JVP_NESTING by torch._C._functorch.maybe_current_level()
- Remove all increment nesting functions from wrap_fx_proxy_cls
- fwAD.make_dual receives the dual_level as keyword argument
- Add jvp_increment_nesting, set_fwd_grad_enabled and dual_level context managers to dynamo
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119926
Approved by: https://github.com/zou3519
2024-03-19 13:06:42 +00:00
Jason Ansel
153a01833b
[dynamo] Optimize SourcelessBuilder ( #122063 )
...
Improves `benchmarks/dynamo/microbenchmarks/dynamo_microbenchmarks.py`
from 2.7s to 2.5s.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122063
Approved by: https://github.com/anijain2305
ghstack dependencies: #122039 , #122043 , #122055 , #122058 , #122060
2024-03-19 04:23:30 +00:00
Jason Ansel
2bec55c5f9
[dynamo] Remove VariableTracker.parents_tracker ( #122058 )
...
This is leftover from mutable variable tracker days and no longer needed.
Improves benchmarks/dynamo/microbenchmarks/dynamo_microbenchmarks.py
from 4.2s to 3.9s.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122058
Approved by: https://github.com/oulgen , https://github.com/anijain2305
ghstack dependencies: #122039 , #122043 , #122055
2024-03-19 04:23:24 +00:00
Jason Ansel
6ca0323615
[dynamo] Optimize VariableTracker.__post_init__ ( #122034 )
...
Improves `benchmarks/dynamo/microbenchmarks/dynamo_microbenchmarks.py`
from 8.6s to 7.3s.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122034
Approved by: https://github.com/Skylion007
ghstack dependencies: #122032 , #122033
2024-03-18 18:08:06 +00:00
Jason Ansel
4d92928fe2
[dynamo] Add tests for fake FSDP ( #121610 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121610
Approved by: https://github.com/yanboliang
ghstack dependencies: #121735 , #120965
2024-03-16 04:29:59 +00:00
Jason Ansel
0b7d9711d4
[dynamo] Add support for nn.Parameter constructor (part 2) ( #120965 )
...
This handles the case where the tensor isn't an input.
The changes to dynamo tests are cases where we would previously fall back to eager.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120965
Approved by: https://github.com/yanboliang
ghstack dependencies: #121735
2024-03-16 04:29:58 +00:00
Jason Ansel
559ca13b3f
[dynamo] Refactor TorchInGraphFunctionVariable for compile time ( #121616 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121616
Approved by: https://github.com/oulgen
ghstack dependencies: #121615
2024-03-13 14:21:21 +00:00
Jason Ansel
a13dd92d88
[dynamo] Minor compile time optimizations in torch.py ( #121615 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121615
Approved by: https://github.com/oulgen
2024-03-13 05:36:22 +00:00
Jason Ansel
7cc476ea16
[dynamo] Fix support for nn.Parameter constructor (part 1) ( #120163 )
...
This captures calls to `torch.nn.Parameter` by lifting them to graph inputs.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120163
Approved by: https://github.com/albanD , https://github.com/yanboliang
ghstack dependencies: #121086
2024-03-11 05:14:42 +00:00
Jason Ansel
32488b0664
[dynamo] Support _unsafe_set_version_counter ( #121086 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121086
Approved by: https://github.com/yanboliang
2024-03-11 05:14:42 +00:00
Ze Sheng
7a4e451184
[Dynamo] Fix function overrides ( #120885 )
...
To check existence of `__torch_function__`, the code intended to iterate each element but got `TupleVariable` when the ordinary `has_torch_function()` was being used. Needs further unpack in this case
Fixes #120653
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120885
Approved by: https://github.com/yanboliang
2024-03-11 02:18:43 +00:00
Jason Ansel
0e0a621e0c
[dynamo] Minor refactors ( #120966 )
...
These are changes I pulled out of the above PRs due to not being
related.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120966
Approved by: https://github.com/yanboliang
2024-03-03 02:20:48 +00:00
Jason Ansel
e3dbd194f4
[dynamo] Support module backwards hooks ( #120685 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120685
Approved by: https://github.com/yanboliang , https://github.com/xmfan
2024-03-01 02:24:26 +00:00
Oleg Khabinov
4b18ab869f
[torch.export] Support is_compiling() flag for non-strict mode ( #119602 )
...
Summary: In non-strict mode of torch.export() we didn't set those `is_compiling()` to `True` which is needed by some models.
Test Plan: Unit tests and manual testing.
Differential Revision: D53624452
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119602
Approved by: https://github.com/suo
2024-02-29 05:52:51 +00:00
Guilherme Leobas
491c2b4665
Let torch dynamo inline torch.func.grad ( #118407 )
...
When dynamo sees torch.func.grad, it tries to inline all frames related
to.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/118407
Approved by: https://github.com/zou3519
2024-02-28 20:05:00 +00:00
Animesh Jain
e3d64c4d5d
[dynamo] Desugar accumulate_grad, fix .grad handling ( #120590 )
...
Fixes https://github.com/pytorch/pytorch/issues/118435
Fixes https://github.com/pytorch/pytorch/issues/119906
Pull Request resolved: https://github.com/pytorch/pytorch/pull/120590
Approved by: https://github.com/ezyang , https://github.com/jansel
ghstack dependencies: #120520
2024-02-27 10:12:26 +00:00