Short-term fix to preserve NJT metadata cache in torch.compile #122836

jbschlosser · 2024-03-27T21:48:02Z

Stack from ghstack (oldest at bottom):

Idea: close over min / max sequence length in the main NJT view func (_nested_view_from_jagged) so that view replay during fake-ification propagates these correctly in torch.compile.

For dynamic shapes support for min / max sequence length, this PR uses a hack that stores the values in (val, 0) shaped tensors.

NB: This PR changes SDPA to operate on real views instead of using buffer_from_jagged() / ViewNestedFromBuffer, which may impact the internal FIRST model. That is, it undoes the partial revert from #123215 alongside a fix to the problem that required the partial revert. We need to verify that there are no regressions there before landing.

Differential Revision: D55448636

[ghstack-poisoned]

pytorch-bot · 2024-03-27T21:48:05Z

🔗 Helpful Links

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

📄 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.

❌ 13 New Failures, 1 Pending

As of commit 06e068c with merge base edb45dc ():

NEW FAILURES - The following jobs have failed:

Lint / lintrunner-noclang / linux-job (gh)
>>> Lint for test/test_nestedtensor.py:
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-focal-py3.11-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_sdpa_cpu_float16
pull / linux-focal-py3.12-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_sdpa_cpu_float16
pull / linux-focal-py3.8-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-focal-py3.8-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_sdpa_cpu_float16
pull / linux-jammy-py3.10-clang15-asan / test (default, 6, 6, linux.4xlarge) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
pull / linux-jammy-py3.8-gcc11 / test (default, 1, 3, linux.2xlarge) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic
trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh)
dynamo/test_subclasses.py::TestNestedTensor::test_inputs_to_compiled_fn_are_views_nt_view_name_base_is_nt_False_basic

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

davidberard98 · 2024-03-27T22:17:11Z

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

tools/autograd/derivatives.yaml

torch/nested/_internal/ops.py

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

jbschlosser · 2024-03-29T21:24:15Z

aten/src/ATen/native/native_functions.yaml

@@ -6200,6 +6200,16 @@
  device_check: NoCheck
  dispatch: {}

+- func: _nested_get_min_seqlen(Tensor self) -> int


I think these need to be symbolic to work with dynamic shapes. TBD exploring this

This needs more thought / design work to fully cover dynamic shapes. @YuqingJ mentioned to me that a static max length is being used for now, so maybe we can get a patch to work for FIRST without this, but it's iffy.

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

clee2000 · 2024-04-02T21:56:55Z

@clee2000 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

jbschlosser · 2024-04-03T15:14:07Z

@jbschlosser has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ghstack-poisoned]

torch/nested/_internal/ops.py

soulitzer · 2024-05-02T16:27:34Z

test/test_nestedtensor.py

+            return torch.ones_like(nt) * nt.max_seqlen()
+
+        for dynamic in [False, True]:
+            self.assertFalse(_recompiles_for_inputs(f, (nt,), (nt2,), dynamic=dynamic))


jfyi there's with unittest.mock.patch("torch._dynamo.config.error_on_recompile", True): though _recompiles_for_inputs seems generally useful if you also want to make sure cases where recompiles=True

I lifted _recompiles_for_inputs() from some other test file, which isn't great :p

test/test_nestedtensor.py

torch/nested/_internal/nested_tensor.py

davidberard98 · 2024-05-02T16:41:43Z

test/test_nestedtensor.py

@@ -4339,6 +4353,115 @@ def forward(self, query, value, offsets):
        self.assertTrue(torch.allclose(attn_output_eager, attn_output))
        self.assertTrue(torch.allclose(value_grad, value.grad))

+    @dtypes(torch.float32)


do you know if this will work with a max_seqlen manually assigned in the graph? (and if so, can you add a test / or add as a followup)

good Q! this works with manual assignment as done in FIRST (i.e. via convert_jagged_to_nested_tensor() -> ViewNestedFromBuffer with the metadata cache info). Also see test_dummy_mha_with_nt which tries to match as closely as possible to what the FIRST model does.

We should arguably support and test manual max seq len assignment for the new API. I'll address this or add a TODO

Update: I think I'll boot handling manual assignment outside of the FIRST use case until we can have public @property getters / setters for min / max sequence length that work with PT2 (this is in progress)

[ghstack-poisoned]

…pile" Idea: close over min / max sequence length in the main NJT view func (`_nested_view_from_jagged`) so that view replay during fake-ification propagates these correctly in torch.compile. For dynamic shapes support for min / max sequence length, this PR uses a hack that stores the values in `(val, 0)` shaped tensors. **NB: This PR changes SDPA to operate on real views instead of using `buffer_from_jagged()` / `ViewNestedFromBuffer`, which may impact the internal FIRST model. That is, it undoes the partial revert from #123215 alongside a fix to the problem that required the partial revert. We need to verify that there are no regressions there before landing.** Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

davidberard98 · 2024-05-23T18:52:12Z

aten/src/ATen/native/native_functions.yaml

@@ -6185,12 +6185,12 @@
    CompositeExplicitAutogradNonFunctional: _nested_view_from_buffer_copy
  autogen: _nested_view_from_buffer_copy.out

- func: _nested_view_from_jagged(Tensor(a) self, Tensor offsets, Tensor dummy, Tensor? lengths=None, int ragged_idx=1) -> Tensor(a)
+- func: _nested_view_from_jagged(Tensor(a) self, Tensor offsets, Tensor dummy, Tensor? lengths=None, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None) -> Tensor(a)


quick question, do we want to eventually make max/min_seqlen available from the torch.nested.nested_tensor_from_jagged API?

We definitely need a way for the user to provide these manually. Adding them to the nested_tensor_from_jagged() API is one way to do this. One additional way / alternative is to provide public setter @properties for min / max sequence length. Some better Dynamo tracing support for @properties landed recently so this may be feasible now.

I was originally thinking we'd only allow manual setting through setter properties, but I could probably be convinced that the nested_tensor_from_jagged() API should accept these too. Opinions on this?

…pile" Idea: close over min / max sequence length in the main NJT view func (`_nested_view_from_jagged`) so that view replay during fake-ification propagates these correctly in torch.compile. For dynamic shapes support for min / max sequence length, this PR uses a hack that stores the values in `(val, 0)` shaped tensors. **NB: This PR changes SDPA to operate on real views instead of using `buffer_from_jagged()` / `ViewNestedFromBuffer`, which may impact the internal FIRST model. That is, it undoes the partial revert from #123215 alongside a fix to the problem that required the partial revert. We need to verify that there are no regressions there before landing.** Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

Short-term fix to preserve NJT metadata cache in torch.compile

fcb0820

[ghstack-poisoned]

jbschlosser requested review from albanD and soulitzer as code owners March 27, 2024 21:48

jbschlosser mentioned this pull request Mar 27, 2024

(WIP) to_padded_tensor() triton kernel for NJT #121947

Draft

jbschlosser requested a review from YuqingJ March 27, 2024 21:48

albanD removed their request for review March 27, 2024 22:00

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

ab7e444

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

soulitzer reviewed Mar 28, 2024

View reviewed changes

tools/autograd/derivatives.yaml Outdated Show resolved Hide resolved

YuqingJ reviewed Mar 28, 2024

View reviewed changes

torch/nested/_internal/ops.py Outdated Show resolved Hide resolved

jbschlosser added 2 commits March 29, 2024 10:51

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

b9606a7

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

ef02e6a

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

jbschlosser commented Mar 29, 2024

View reviewed changes

jbschlosser added 3 commits April 2, 2024 16:22

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

1bd5588

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

abf28a8

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

957d277

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

Update on "Short-term fix to preserve NJT metadata cache in torch.com…

78efc28

…pile" Differential Revision: [D55448636](https://our.internmc.facebook.com/intern/diff/D55448636) [ghstack-poisoned]

pytorch-bot bot added the ciflow/inductor label Apr 3, 2024

Update

ece048c

[ghstack-poisoned]

jbschlosser added the topic: not user facing topic category label May 1, 2024

Update

b6d1abe

[ghstack-poisoned]

soulitzer reviewed May 2, 2024

View reviewed changes

torch/nested/_internal/ops.py Show resolved Hide resolved

soulitzer reviewed May 2, 2024

View reviewed changes

test/test_nestedtensor.py Outdated Show resolved Hide resolved

soulitzer reviewed May 2, 2024

View reviewed changes

torch/nested/_internal/nested_tensor.py Show resolved Hide resolved

davidberard98 reviewed May 2, 2024

View reviewed changes

Update

6d75373

[ghstack-poisoned]

This was referenced May 17, 2024

Use return_and_correct_aliasing() for NJT + compatible storage setting #126552

Open

[Tracker] Move nested tensors to beta #112398

Open

jbschlosser added 11 commits May 17, 2024 14:15

davidberard98 reviewed May 23, 2024

View reviewed changes

jbschlosser mentioned this pull request May 23, 2024

Naive CPU kernels for jagged <-> padded dense conversions #127007

Closed

jbschlosser added 7 commits May 23, 2024 16:32

jbschlosser mentioned this pull request Jun 5, 2024

Skip debug asserts for mixed dense, subclass views in autograd_not_implemented_fallback #128057

Closed

jbschlosser added 6 commits June 7, 2024 11:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Short-term fix to preserve NJT metadata cache in torch.compile #122836

Short-term fix to preserve NJT metadata cache in torch.compile #122836

jbschlosser commented Mar 27, 2024 •

edited

pytorch-bot bot commented Mar 27, 2024 •

edited

davidberard98 commented Mar 27, 2024

jbschlosser Mar 29, 2024

jbschlosser Apr 2, 2024 •

edited

clee2000 commented Apr 2, 2024

jbschlosser commented Apr 3, 2024

soulitzer May 2, 2024

jbschlosser May 7, 2024

davidberard98 May 2, 2024

jbschlosser May 7, 2024

jbschlosser May 8, 2024

davidberard98 May 23, 2024

jbschlosser May 23, 2024

Short-term fix to preserve NJT metadata cache in torch.compile #122836

Are you sure you want to change the base?

Short-term fix to preserve NJT metadata cache in torch.compile #122836

Conversation

jbschlosser commented Mar 27, 2024 • edited

pytorch-bot bot commented Mar 27, 2024 • edited

🔗 Helpful Links

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

❌ 13 New Failures, 1 Pending

davidberard98 commented Mar 27, 2024

Choose a reason for hiding this comment

jbschlosser Apr 2, 2024 • edited

Choose a reason for hiding this comment

clee2000 commented Apr 2, 2024

jbschlosser commented Apr 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jbschlosser commented Mar 27, 2024 •

edited

pytorch-bot bot commented Mar 27, 2024 •

edited

jbschlosser Apr 2, 2024 •

edited