Support narrow() on batch dim for NJT #136444

jbschlosser · 2024-09-23T17:09:06Z

Stack from ghstack (oldest at bottom):

Requested in #136270

TODO:

I expect there will be torch.compile() problems from the use of the item() calls in the naive impl; fix these

cc @cpuhrsch @bhosmer @drisspg @soulitzer @davidberard98 @YuqingJ

[ghstack-poisoned]

pytorch-bot · 2024-09-23T17:09:10Z

🔗 Helpful Links

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

📄 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

As of commit 36ba8d8 with merge base 03ec250 ():

NEW FAILURES - The following jobs have failed:

Lint / lintrunner-noclang / linux-job (gh)
>>> Lint for torch/fx/passes/runtime_assert.py:
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_narrow_on_batch_dim_compiled_True_cuda_float32
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_narrow_on_batch_dim_compiled_True_cuda_float32
pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-focal-py3.11-clang10 / test (default, 3, 5, linux.4xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-focal-py3.11-clang10 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_False_cpu_float32
pull / linux-focal-py3.12-clang10 / test (default, 3, 5, linux.4xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-focal-py3.12-clang10 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_False_cpu_float32
pull / linux-focal-py3.9-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-focal-py3.9-clang10 / test (default, 3, 5, linux.4xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-focal-py3.9-clang10 / test (dynamo_wrapped, 3, 3, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_False_cpu_float32
pull / linux-jammy-py3.10-clang15-asan / test (default, 4, 6, linux.4xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32
pull / linux-jammy-py3.9-gcc11 / test (default, 3, 5, linux.2xlarge) (gh)
test_nestedtensor.py::TestNestedTensorSubclassCPU::test_narrow_on_batch_dim_compiled_True_cpu_float32

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

ghstack-source-id: 768b2b4 Pull Request resolved: #136444

jbschlosser · 2024-10-22T19:52:41Z

Update on this: I'm looking into but it's not easy to get this working graph-break-free with torch.compile. The heart of the issue is data-dependency: the output of narrow() on the batch dim depends on the actual data of the offsets. In particular, the returned NJT should have a values buffer with a packed sum(*) ragged dimension whose size depends on the offsets data.

Alternatively, we could return the same values buffer and just change the offsets data, avoiding the data-dependency for the shape of the returned values buffer. Now we have the following problems:

I believe there is an implicit assumption throughout the NJT logic that offsets start with zero. Investigation is required to identify and remove areas where this assumption is made.
Operations that operate on the underlying values buffer will do more computation than is needed, somewhat defeating the purpose of taking the narrow() view in the first place. For example, unary and binary pointwise operations simply redispatch on the underlying values buffer. It may be possible to implement these by composing some_custom_op_that_gets_the_shape(offsets, etc.) + the unary_op() / binary_op() itself.

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

jbschlosser · 2024-12-19T19:21:06Z

Closing in favor of #142063

ghstack-source-id: f763a01 Pull Request resolved: pytorch/pytorch#136444

Support narrow() on batch dim for NJT

b82ff20

[ghstack-poisoned]

jbschlosser requested review from albanD and soulitzer as code owners September 23, 2024 17:09

jbschlosser added a commit that referenced this pull request Sep 23, 2024

Support narrow() on batch dim for NJT

88b3355

ghstack-source-id: 768b2b4 Pull Request resolved: #136444

jbschlosser mentioned this pull request Sep 23, 2024

Training (backward) crashes when using torch.narrow, nested tensors, and scaled_dot_product_attention #136270

Open

jbschlosser added module: nestedtensor NestedTensor tag see issue #25032 topic: improvements topic category release notes: nested tensor Changes that have a direct impact on nested tensors labels Sep 23, 2024

albanD removed their request for review September 23, 2024 18:11

jbschlosser marked this pull request as draft October 29, 2024 14:58

Update on "Support narrow() on batch dim for NJT"

6926d78

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

jbschlosser mentioned this pull request Oct 29, 2024

Support view_as() on NJT; allow nested int swapping #139196

Closed

Update on "Support narrow() on batch dim for NJT"

dc019ff

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

jbschlosser mentioned this pull request Oct 30, 2024

Fix and test several NJT reductions #139317

Closed

jbschlosser added 5 commits October 30, 2024 17:23

Update on "Support narrow() on batch dim for NJT"

d52bb97

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

Update on "Support narrow() on batch dim for NJT"

29da2a3

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

Update on "Support narrow() on batch dim for NJT"

5037b1a

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

Update on "Support narrow() on batch dim for NJT"

980c52d

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

Update on "Support narrow() on batch dim for NJT"

36ba8d8

Requested in #136270 TODO: * I expect there will be `torch.compile()` problems from the use of the `item()` calls in the naive impl; fix these cc cpuhrsch bhosmer drisspg soulitzer davidberard98 YuqingJ [ghstack-poisoned]

jbschlosser mentioned this pull request Nov 8, 2024

Support attn_mask in jagged_scaled_dot_product_attention #138993

Open

jbschlosser closed this Dec 19, 2024

github-actions bot deleted the gh/jbschlosser/177/head branch January 19, 2025 02:06

desai0007 pushed a commit to desai0007/test-repo-pytorch that referenced this pull request Feb 26, 2025

Support narrow() on batch dim for NJT

ef75a8a

ghstack-source-id: f763a01 Pull Request resolved: pytorch/pytorch#136444

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Support narrow() on batch dim for NJT #136444

Support narrow() on batch dim for NJT #136444

Uh oh!

jbschlosser commented Sep 23, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Sep 23, 2024 •

edited

Loading

Uh oh!

jbschlosser commented Oct 22, 2024 •

edited

Loading

Uh oh!

jbschlosser commented Dec 19, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Support narrow() on batch dim for NJT #136444

Support narrow() on batch dim for NJT #136444

Uh oh!

Conversation

jbschlosser commented Sep 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 13 New Failures

Uh oh!

jbschlosser commented Oct 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jbschlosser commented Dec 19, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

jbschlosser commented Sep 23, 2024 •

edited

Loading

pytorch-bot bot commented Sep 23, 2024 •

edited

Loading

jbschlosser commented Oct 22, 2024 •

edited

Loading