[NJT] can only chunk if the 2nd dimension is ragged #153238

mdeff · 2025-05-09T13:19:27Z

🐛 Describe the bug

The following works as expected:

import torch
x = torch.nested.nested_tensor(
    [
        torch.randn((8, 3, 9)),
        torch.randn((9, 3, 9)),
    ],
    layout=torch.jagged,
)
assert x.is_contiguous()
print(x)
torch.chunk(x, 3, dim=-1)

While the following fails:

import torch
x = torch.nested.nested_tensor(
    [
        torch.randn((9, 3, 9)),
        torch.randn((9, 4, 9)),
    ],
    layout=torch.jagged,
)
assert x.is_contiguous()
print(x)
torch.chunk(x, 3, dim=-1)

With the unhelpful exception NestedTensor chunk_default(self: jt, chunks: any, dim: any?): expected self to be a contiguous jagged layout NestedTensor.

Not a problem for the strided layout.

In real-world use, that'd allow to simplify

qkv = self.in_proj(input)

# (N, L, embed_dim*3) → 3 * (N, num_heads, L, embed_dim/num_heads)
query, key, value = torch.chunk(qkv, 3, dim=-1)
query = query.unflatten(-1, [self.num_heads, -1]).transpose(1, 2)
key = key.unflatten(-1, [self.num_heads, -1]).transpose(1, 2)
value = value.unflatten(-1, [self.num_heads, -1]).transpose(1, 2)

into

qkv = self.in_proj(input)

# (N, L, embed_dim*3) → 3 * (N, num_heads, L, embed_dim/num_heads)
qkv = qkv.unflatten(-1, [self.num_heads, -1]).transpose(1, 2)
query, key, value = torch.chunk(qkv, 3, dim=-1)

Versions

PyTorch 2.7.0+cu126 from PyPI.

Full collect_env.py output

Collecting environment information...
PyTorch version: 2.7.0+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.12.0 | packaged by conda-forge | (main, Oct  3 2023, 08:43:22) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-113-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 535.161.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             8
On-line CPU(s) list:                0-7
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) W-2123 CPU @ 3.60GHz
CPU family:                         6
Model:                              85
Thread(s) per core:                 2
Core(s) per socket:                 4
Socket(s):                          1
Stepping:                           4
CPU max MHz:                        3900.0000
CPU min MHz:                        1200.0000
BogoMIPS:                           7200.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single pti intel_ppin ssbd mba ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req md_clear flush_l1d arch_capabilities
Virtualization:                     VT-x
L1d cache:                          128 KiB (4 instances)
L1i cache:                          128 KiB (4 instances)
L2 cache:                           4 MiB (4 instances)
L3 cache:                           8.3 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-7
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX conditional cache flushes, SMT vulnerable
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS; IBPB conditional; STIBP conditional; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; Clear CPU buffers; SMT vulnerable

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect

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

The text was updated successfully, but these errors were encountered:

jbschlosser · 2025-05-12T16:37:25Z

Thanks for the report, @mdeff! Support for a ragged dimension not immediately next to the batch dimension (as in your second example) is newer and buggy wrt input validation. Validation logic here still uses an older definition of "contiguous" that required njt._ragged_idx == 1:

pytorch/torch/nested/_internal/ops.py

Lines 114 to 116 in dc47295

    
           "jt": lambda x: isinstance(x, NestedTensor) 
        
           and x._lengths is None 
        
           and x._ragged_idx == 1,  # ops with "jt" require contiguous JT only

The easiest fix for this particular case is to change the schema string from jt -> jt_all for chunk() here:

pytorch/torch/nested/_internal/ops.py

Line 969 in dc47295

    
           @register_jagged_func(torch.ops.aten.chunk.default, "self: jt, chunks: any, dim: any?")

and error for truly non-contiguous inputs. We'd accept a PR implementing this if you'd like to contribute.

simonzgx · 2025-05-16T06:35:42Z

Hi @jbschlosser , I just submitted a PR for this issue, please have a look.

mdeff · 2025-05-16T13:29:15Z

Thanks for the explanation and guidance, @jbschlosser!
I've been playing with a fix too, at #153706. (I don't think #153689 is the way.)

In general, there seems to be confusion/ambiguity between contiguity (is_contiguous() referring to the map from nD (Nested)Tensor to 1D computer memory) and "hollowness" (_lengths is None referring to a NestedTensor having no holes / unused values). A NestedTensor with holes cannot be contiguous, but it may not have holes while still not being contiguous (e.g., after a transposition). Holes are to be handled at the NestedTensor level of abstraction, while it seems contiguity is mostly a lower-level concern.

In #153237 I identified the contiguity–hollowness missmatch, which was updated in #153529 to mean contiguity. But maybe "hollowness" is more relevant in the __repr__?

Finally, it's assumed throughout the code that _lengths is not None means there are holes but that's too strict. For example:

import torch
values = torch.randn(6, 5)
offsets = torch.tensor([0, 2, 3, 6])
lengths = torch.tensor([2, 1, 3])
x = torch.nested.nested_tensor_from_jagged(values, offsets, lengths)
assert x.lengths() is not None
assert not x.is_contiguous()  # but it has no holes and is contiguous
assert torch.equal(x.lengths(), x.offsets().diff())

I see two fixes:

update the factories to set _lengths = None if torch.equal(inp.lengths(), inp.offsets().diff()) so the rest of the code can continue to assume that an NJT has no hole iff _lengths is None, getting rid of the redundant _lengths tensor,
update all checks for holes (in is_contiguous and other ops) with (inp.lengths() is not None) and not torch.equal(inp.lengths(), inp.offsets().diff()), maybe with a is_hollow()/has_holes() helper.

jbschlosser · 2025-05-16T13:50:58Z

@mdeff Extremely solid analysis!

In general, there seems to be confusion/ambiguity between contiguity (is_contiguous() referring to the map from nD (Nested)Tensor to 1D computer memory) and "hollowness" (_lengths is None referring to a NestedTensor having no holes / unused values). A NestedTensor with holes cannot be contiguous, but it may not have holes while still not being contiguous (e.g., after a transposition). Holes are to be handled at the NestedTensor level of abstraction, while it seems contiguity is mostly a lower-level concern.

100% correct. I'll mention that support for holes was tacked on later, intending to contribute to support for "a ragged view of a dense tensor" usable for kv-cache logic.

Finally, it's assumed throughout the code that _lengths is not None means there are holes but that's too strict.

Also correct. In theory, we -should- allow lengths information that is fully redundant with offsets, but checking lengths is None avoids data-dependent computation and GPU syncs, which helps quite a bit with torch.compile support / performance.

I see two fixes:

update the factories to set _lengths = None if torch.equal(inp.lengths(), inp.offsets().diff()) so the rest of the code can continue to > assume that an NJT has no hole iff _lengths is None, getting rid of the redundant _lengths tensor,

update all checks for holes (in is_contiguous and other ops) with (inp.lengths() is not None) and not torch.equal(inp.lengths(), inp.offsets().diff()), maybe with a is_hollow()/has_holes() helper.

I like 1 better, but even so, it's not ideal to compute the data-dependent, GPU sync incurring logic of torch.equal(inp.lengths(), inp.offsets().diff()) during construction, which is why this hasn't been introduced up until now. So we have taken a "technically too strict" approach wrt lengths to avoid these sorts of checks. I'm not positive what the correct solution is here, but I think introducing the checks will cause problems with torch.compile support / performance (should be easy to see from failing tests in the PR).

But maybe "hollowness" is more relevant in the __repr__?

I think there's a good argument to be made for including this in the __repr__. Or to be very pedantic, the __repr__ should ideally be eval()-able directly to reconstruct the NJT, so contiguity should be left out and other necessary information for reconstruction should be included instead.

mdeff · 2025-05-16T15:08:48Z

@jbschlosser thanks!

I'll mention that support for holes was tacked on later, intending to contribute to support for "a ragged view of a dense tensor" usable for kv-cache logic.

Makes sense! Nice to tackle that with NJTs.

avoids data-dependent computation and GPU syncs, which helps quite a bit with torch.compile support / performance

Wasn't aware of it; thanks for educating me. With that in mind, the current "technically too strict" approach seems to strike a good tradeoff. :)

malfet added the module: nestedtensor NestedTensor tag see issue #25032 label May 9, 2025

jbschlosser added triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module actionable labels May 12, 2025

simonzgx mentioned this issue May 16, 2025

[Torch][NJT] Fix schema check logic for NJT's chunk operation. #153689

Closed

mdeff linked a pull request May 16, 2025 that will close this issue

[Torch][NJT] relax schema checks for ops that support more general NJTs #153706

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[NJT] can only chunk if the 2nd dimension is ragged #153238

[NJT] can only chunk if the 2nd dimension is ragged #153238

[NJT] can only chunk if the 2nd dimension is ragged #153238

[NJT] can only chunk if the 2nd dimension is ragged #153238

Comments

🐛 Describe the bug

Versions