Always create ShapeEnv, always apply unspec logic #103302

ezyang · 2023-06-09T03:25:49Z

Stack from ghstack (oldest at bottom):

Originally, my goal for this PR was to remove the dynamic_shapes tests in torch/_dynamo/variables/builder.py. However, one thing lead to another, and it turns out that it was easiest to do all of the following in one go:

Unconditionally allocate a ShapeEnv, no matter if dynamic_shapes is enabled or not (torch/_dynamo/output_graph.py). There is a small adjustment to export torch/_dynamo/eval_frame.py to account for the fact that a ShapeEnv always exists, even if you're not doing symbolic export.
Remove dynamic_shapes test from unspec logic (torch/_dynamo/variables/builder.py), the original goal
Specialize strides and storage offset if all sizes are dynamic (torch/fx/experimental/symbolic_shapes.py). This is required to deal with unconditional ShapeEnv: if a ShapeEnv exist, fake tensor-ification may choose to allocate symbols. The idea is that with automatic_dynamic_shapes == False, Dynamo should never request dynamic sizes, but this invariant was not upheld for nontrivial strides/offset.

The rest are just auxiliary fixups from the above:

Workaround bug in FakeTensorProp where sometimes it doesn't return a FakeTensor (torch/fx/passes/fake_tensor_prop.py), see [FAILING] Tighten FakeTensorProp assert to require only fake tensor returns. #103395 for follow up
Make ShapeProp correctly handle int inputs (torch/fx/passes/shape_prop.py)
Disable indexing strength reduction if assume_static_by_default is False (torch/_inductor/codegen/triton.py)
Fix hf_T5_generate to NOT toggle assume_static_by_default if dynamic shapes is not enabled (benchmarks/dynamo/common.py); technically this is not necessary anymore but it's in for safety.

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @aakhundov @anijain2305

Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

pytorch-bot · 2023-06-09T03:25:51Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/103302

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ca69264:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 7080719c9b5976719a9c6fa3fa4744655cb115d5 Pull Request resolved: #103302

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 846160c0e503d7f49823ad1521c3478cbc48cc0c Pull Request resolved: #103302

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 aakhundov anijain2305 [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 924899944d9c2c61c4e28edb5f8df14f824f215e Pull Request resolved: #103302

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 aakhundov anijain2305 [ghstack-poisoned]

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 50472e03d68393c3084bb55b2b4114b60e3483d5 Pull Request resolved: #103302

ezyang · 2023-06-11T17:42:29Z

torch/_dynamo/variables/builder.py

-            automatic_dynamic = config.automatic_dynamic_shapes and (
-                frame_state_entry.size is None or frame_state_entry.size[i] is None
-            )
+    dynamic_dims = []


This is reindentation only, as shape_env is always non-None.

ezyang · 2023-06-12T12:44:43Z

@pytorchbot merge

pytorchmergebot · 2023-06-12T12:48:23Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Fix #16355. The root cause change in PyTorch is [#103302](pytorch/pytorch#103302), which seem blocking calling make_fx inside a dynamo backend. Changes: 1. Move decomposition to `register_backend.py`, so we don't have to call `make_fx` inside DORT, which triggers a bunch of new exceptions. 2. Remove shape inference based on FakeTensorProp since the FX graph received from dynamo contains all shapes now. 3. Fix a macro bug so that DORT can build without CUDA. Before (3), ``` #if defined(USE_CUDA) || defined(USE_ROCM) virtual PhiloxGenerator& PhiloxGenerator__Default() = 0; #ifdef ENABLE_TRAINING_TORCH_INTEROP ... #endif #endif ``` After (3), ``` #if defined(USE_CUDA) || defined(USE_ROCM) virtual PhiloxGenerator& PhiloxGenerator__Default() = 0; #endif #ifdef ENABLE_TRAINING_TORCH_INTEROP ... #endif ``` The later one looks better since the `ENABLE_TRAINING_TORCH_INTEROP` is for Python bridge code, not for random-number-generating kernels `PhiloxGenerator`.

Fix microsoft#16355. The root cause change in PyTorch is [#103302](pytorch/pytorch#103302), which seem blocking calling make_fx inside a dynamo backend. Changes: 1. Move decomposition to `register_backend.py`, so we don't have to call `make_fx` inside DORT, which triggers a bunch of new exceptions. 2. Remove shape inference based on FakeTensorProp since the FX graph received from dynamo contains all shapes now. 3. Fix a macro bug so that DORT can build without CUDA. Before (3), ``` #if defined(USE_CUDA) || defined(USE_ROCM) virtual PhiloxGenerator& PhiloxGenerator__Default() = 0; #ifdef ENABLE_TRAINING_TORCH_INTEROP ... #endif #endif ``` After (3), ``` #if defined(USE_CUDA) || defined(USE_ROCM) virtual PhiloxGenerator& PhiloxGenerator__Default() = 0; #endif #ifdef ENABLE_TRAINING_TORCH_INTEROP ... #endif ``` The later one looks better since the `ENABLE_TRAINING_TORCH_INTEROP` is for Python bridge code, not for random-number-generating kernels `PhiloxGenerator`.

Always apply unspec logic

80b302a

Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 9, 2023

Always apply unspec logic

fe69589

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 7080719c9b5976719a9c6fa3fa4744655cb115d5 Pull Request resolved: #103302

github-actions bot requested review from albanD, antoniojkim, bdhirsh, jbschlosser, miladm, SherlockNoMad, voznesenskym and wconstab June 9, 2023 03:26

github-actions bot added ciflow/inductor module: dynamo labels Jun 9, 2023

Update on "Always apply unspec logic"

952c2ec

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

ezyang mentioned this pull request Jun 9, 2023

Turn off automatic_dynamic_shapes in prep for dynamic-by-default #103320

Closed

Update on "Always apply unspec logic"

208960d

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

Update on "Always apply unspec logic"

ec23c3d

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 9, 2023

Always apply unspec logic

0bd1a7a

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 846160c0e503d7f49823ad1521c3478cbc48cc0c Pull Request resolved: #103302

ezyang mentioned this pull request Jun 9, 2023

Remove dynamic_shapes condition from ShapeEnv creation logic #103328

Closed

Update on "Always apply unspec logic"

ceb051f

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy aakhundov [ghstack-poisoned]

github-actions bot added the module: inductor label Jun 10, 2023

pytorch-bot bot added the release notes: fx release notes category label Jun 10, 2023

Update on "Always apply unspec logic"

02a4b32

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 aakhundov anijain2305 [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 11, 2023

Always apply unspec logic

090166d

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 924899944d9c2c61c4e28edb5f8df14f824f215e Pull Request resolved: #103302

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 11, 2023

ezyang mentioned this pull request Jun 11, 2023

Unconditionally create shape_env in Dynamo #103390

Closed

Update on "Always apply unspec logic"

ca69264

Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy ngimel yf225 aakhundov anijain2305 [ghstack-poisoned]

ezyang added a commit that referenced this pull request Jun 11, 2023

Always apply unspec logic

89ffdb0

Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 50472e03d68393c3084bb55b2b4114b60e3483d5 Pull Request resolved: #103302

ezyang changed the title ~~Always apply unspec logic~~ Always create ShapeEnv, always apply unspec logic Jun 11, 2023

ezyang commented Jun 11, 2023

View reviewed changes

ezyang requested a review from anijain2305 June 12, 2023 02:12

ezyang mentioned this pull request Jun 9, 2023

Refactor wrap_fx_proxy_cls, no symbolic_shapes test anymore #103303

Closed

voznesenskym approved these changes Jun 12, 2023

View reviewed changes

pytorchmergebot added the merging label Jun 12, 2023

pytorchmergebot added Merged and removed merging labels Jun 12, 2023

pytorchmergebot closed this in c3fdfca Jun 12, 2023

This was referenced Jun 12, 2023

Use assertEqual() instead of assertTrue(same()) #103429

Closed

Rewrite size/stride/numel TensorVariable handling #103438

Closed

facebook-github-bot deleted the gh/ezyang/2151/head branch June 15, 2023 14:16

wschin mentioned this pull request Jun 19, 2023

Update DORT to follow PyTorch changes microsoft/onnxruntime#16394

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Always create ShapeEnv, always apply unspec logic #103302

Always create ShapeEnv, always apply unspec logic #103302

ezyang commented Jun 9, 2023 •

edited

Loading

pytorch-bot bot commented Jun 9, 2023 •

edited

Loading

ezyang Jun 11, 2023

ezyang commented Jun 12, 2023

pytorchmergebot commented Jun 12, 2023

Always create ShapeEnv, always apply unspec logic #103302

Always create ShapeEnv, always apply unspec logic #103302

Conversation

ezyang commented Jun 9, 2023 • edited Loading

pytorch-bot bot commented Jun 9, 2023 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/103302

✅ No Failures

ezyang Jun 11, 2023

Choose a reason for hiding this comment

ezyang commented Jun 12, 2023

pytorchmergebot commented Jun 12, 2023

Merge started

ezyang commented Jun 9, 2023 •

edited

Loading

pytorch-bot bot commented Jun 9, 2023 •

edited

Loading