vjp3 + array refs fails when some of the gradients are masked

[qGentry]

Description

Hi, I keep playing around with the vjp3 + array refs for zero-memory-overhead gradient accumulation. I recently reported the problem about scan + reverse:
#32411 (comment)
I brought it to JAX 0.7.2 via monkey patching and now full transformer training is working properly. Now I moved to more advanced training setups, like training extra heads on top of the transformer while freezing the backbone (via optax). So in forward pass, we're doing full model forward pass, get output embeddings and then feed it to smaller model to predict main model's distribution. But masking some keys in model state's optimizer, we ensure that during the backward pass we only train heads and not the backbone model so jax optimizes away entire main model backward pass.

For the backbone models, we're not using scan to iterate over layers. But for the heads, we're using scan.
We got the following error:

 File "/papyrax/papyrax/training/trainer.py", line 128, in _loop_body
    grad_fn.with_refs(grads_accum)(jnp.ones_like(loss))
ValueError: Invalid shape for `addupdate`. Ref shape: (14336, 4096). Expected shape: (14336, 4096). Value shape: (4096, 14336). Transforms: ().

Error happens here in backward_pass3 function:

            for x, ct in zip(primals, cts_out):
              if isinstance(x, GradAccum):
                x.accum(ct) <- here

primals
[JitTracer<bfloat16[8,8192,4096]>, <jax._src.interpreters.ad.RefAccum object at 0x7fa7e0342bd0>]
cts_out
(None, JitTracer<bfloat16[4096,14336]>)

My guess would be that at some points training state tree is being flattened in a list and then some entires are filtered out while gradient update expect original flat tree.

It is also interesting that it looks like I'm getting this error during tracing, even before before lowering and compiling.

I tried reproducing this error via simple script but no luck yet.

System info (python version, jaxlib version, accelerator, etc.)

root@computeinstance-e00p9411jfafcfq06n:/papyrax# python3 -c 'import jax; jax.print_environment_info()'
jax:    0.7.2
jaxlib: 0.7.2
numpy:  2.1.3
python: 3.11.13 (main, Jun  4 2025, 08:57:30) [GCC 13.3.0]
device info: NVIDIA H100 80GB HBM3-8, 8 local devices"
process_count: 1
platform: uname_result(system='Linux', node='computeinstance-e00p9411jfafcfq06n', release='5.15.0-152-generic', version='#162-Ubuntu SMP Wed Jul 23 09:48:42 UTC 2025', machine='x86_64')
XLA_PYTHON_CLIENT_MEM_FRACTION=.85
JAX_COMPILATION_CACHE_DIR=/mnt/llm/cache/jax_compilation_cache

$ nvidia-smi
Tue Oct  7 15:18:00 2025       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.163.01             Driver Version: 550.163.01     CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 80GB HBM3          On  |   00000000:8D:00.0 Off |                    0 |
| N/A   36C    P0            120W /  700W |     558MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA H100 80GB HBM3          On  |   00000000:91:00.0 Off |                    0 |
| N/A   32C    P0            122W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   2  NVIDIA H100 80GB HBM3          On  |   00000000:95:00.0 Off |                    0 |
| N/A   31C    P0            120W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   3  NVIDIA H100 80GB HBM3          On  |   00000000:99:00.0 Off |                    0 |
| N/A   36C    P0            123W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   4  NVIDIA H100 80GB HBM3          On  |   00000000:AB:00.0 Off |                    0 |
| N/A   37C    P0            123W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   5  NVIDIA H100 80GB HBM3          On  |   00000000:AF:00.0 Off |                    0 |
| N/A   32C    P0            120W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   6  NVIDIA H100 80GB HBM3          On  |   00000000:B3:00.0 Off |                    0 |
| N/A   31C    P0            121W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
|   7  NVIDIA H100 80GB HBM3          On  |   00000000:B7:00.0 Off |                    0 |
| N/A   37C    P0            124W /  700W |     538MiB /  81559MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+

[qGentry]

Another example of values not being aligned. This time we have the backbone models with scan and heads with scan. Sorry for the pour readability.

cts_out
[None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, JitTracer<bfloat16[16,8192,4096]>, JitTracer<bfloat16[3,16,8192,4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096,8,128]>, JitTracer<bfloat16[32,128,4096]>, JitTracer<bfloat16[4096,32,128]>, JitTracer<bfloat16[4096,8,128]>, JitTracer<bfloat16[4096,14336]>, JitTracer<bfloat16[14336,4096]>, JitTracer<bfloat16[4096,14336]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[8192,4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[3,16,8192,4096]>]

primals
[JitTracer<bfloat16[3,16,8192,4096]>, JitTracer<bfloat16[3,16,8192,4096]>, JitTracer<bfloat16[16,8192,4096]>, JitTracer<int32[3,16,8192]>, JitTracer<uint32[16,8192]>, JitTracer<int32[64]>, JitTracer<int32[64]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096,8,128]>, JitTracer<bfloat16[32,128,4096]>, JitTracer<bfloat16[4096,32,128]>, JitTracer<bfloat16[4096,8,128]>, JitTracer<bfloat16[4096,14336]>, JitTracer<bfloat16[14336,4096]>, JitTracer<bfloat16[4096,14336]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[8192,4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096]>, JitTracer<bfloat16[4096]>, JitTracer<int32[]>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881de890>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dfe50>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd950>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dc750>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dfb10>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd850>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd810>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dcb90>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd450>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dc310>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dcad0>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd0d0>, <jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd090>, <jax._src.interpreters.ad.ValAccum object at 0x7fc6881df050>, <jax._src.interpreters.ad.ValAccum object at 0x7fc6881de950>, <jax._src.interpreters.ad.ValAccum object at 0x7fc6881df4d0>]
special variables
function variables
00 =
JitTracer<bfloat16[3,16,8192,4096]>
01 =
JitTracer<bfloat16[3,16,8192,4096]>
02 =
JitTracer<bfloat16[16,8192,4096]>
03 =
JitTracer<int32[3,16,8192]>
04 =
JitTracer<uint32[16,8192]>
05 =
JitTracer<int32[64]>
06 =
JitTracer<int32[64]>
07 =
JitTracer<bfloat16[4096]>
08 =
JitTracer<bfloat16[4096,8,128]>
09 =
JitTracer<bfloat16[32,128,4096]>
10 =
JitTracer<bfloat16[4096,32,128]>
11 =
JitTracer<bfloat16[4096,8,128]>
12 =
JitTracer<bfloat16[4096,14336]>
13 =
JitTracer<bfloat16[14336,4096]>
14 =
JitTracer<bfloat16[4096,14336]>
15 =
JitTracer<bfloat16[4096]>
16 =
JitTracer<bfloat16[8192,4096]>
17 =
JitTracer<bfloat16[4096]>
18 =
JitTracer<bfloat16[4096]>
19 =
JitTracer<bfloat16[4096]>
20 =
JitTracer<int32[]>
21 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881de890>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
ref =
JitTracer<Ref{bfloat16[4096]}>
22 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dfe50>
special variables
function variables
aval =
ShapedArray(bfloat16[4096,8,128])
ref =
JitTracer<Ref{bfloat16[4096,8,128]}>
23 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd950>
special variables
function variables
aval =
ShapedArray(bfloat16[32,128,4096])
ref =
JitTracer<Ref{bfloat16[32,128,4096]}>
24 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dc750>
special variables
function variables
aval =
ShapedArray(bfloat16[4096,32,128])
ref =
JitTracer<Ref{bfloat16[4096,32,128]}>
25 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dfb10>
special variables
function variables
aval =
ShapedArray(bfloat16[4096,8,128])
ref =
JitTracer<Ref{bfloat16[4096,8,128]}>
26 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd850>
special variables
function variables
aval =
ShapedArray(bfloat16[4096,14336])
ref =
JitTracer<Ref{bfloat16[4096,14336]}>
27 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd810>
special variables
function variables
aval =
ShapedArray(bfloat16[14336,4096])
ref =
JitTracer<Ref{bfloat16[14336,4096]}>
28 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dcb90>
special variables
function variables
aval =
ShapedArray(bfloat16[4096,14336])
ref =
JitTracer<Ref{bfloat16[4096,14336]}>
29 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd450>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
ref =
JitTracer<Ref{bfloat16[4096]}>
30 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dc310>
special variables
function variables
aval =
ShapedArray(bfloat16[8192,4096])
ref =
JitTracer<Ref{bfloat16[8192,4096]}>
31 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dcad0>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
ref =
JitTracer<Ref{bfloat16[4096]}>
32 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd0d0>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
ref =
JitTracer<Ref{bfloat16[4096]}>
33 =
<jax._src.interpreters.ad.RefAccum object at 0x7fc6881dd090>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
ref =
JitTracer<Ref{bfloat16[4096]}>
34 =
<jax._src.interpreters.ad.ValAccum object at 0x7fc6881df050>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
val =
JitTracer<bfloat16[4096]>
35 =
<jax._src.interpreters.ad.ValAccum object at 0x7fc6881de950>
special variables
function variables
aval =
ShapedArray(bfloat16[4096])
val =
JitTracer<bfloat16[4096]>
36 =
<jax._src.interpreters.ad.ValAccum object at 0x7fc6881df4d0>
special variables
function variables
aval =
ShapedArray(bfloat16[3,16,8192,4096])
val =
Zero(bfloat16[3,16,8192,4096])

[qGentry]

Also, seems like disabling scan for the heads solves this issue

[qGentry]

@mattjj
I just tried latest nightly JAX with your latest fixes around the vjp3 but it didn't help
#32583

root@computeinstance-e00p9411jfafcfq06n:/papyrax# pip show jax
Name: jax
Version: 0.8.0.dev20251014