Memory chunking of the derivative transpose (chunk_size)#

The probe-derivative transpose 𝒥ᵀ (back-projecting residual jets into a variation gradient — tv_probe_derivatives_transpose, T3Tangent.probe_derivatives_transpose, their uniform twins, and the fitting 𝒥ᵀr step) has one stage whose peak memory can dwarf everything else: the variation gradient assembly. The chunk_size parameter bounds that peak. This note explains the memory model, how to set chunk_size, and how it composes with multi-device sharding.

Why the assembly is the bottleneck#

Every stage of the transpose is linear in the sample-stack size |W| (the number of probes), but they fall into two very different per-W-row costs:

  • The edge-variable jets (mu, nu, sigma_tilde, tau_tilde, deta_tilde, …) — resident arrays with per-row cost ~ (order+1)·K·r. This is the necessary floor: it is the data.

  • The assembly intermediate — the only stage whose per-row cost is super-linear in the ranks: it forms frame-pairwise products (mu ⊗ xi, xi ⊗ nu, mu ⊗ nu), each ~ (order+1)²·(rank·rank). That (order+1)²·r² factor is what turns a few-GB pipeline into hundreds of GB.

So the assembly is the only thing that runs out of memory, and chunk_size is the only knob that touches it. (The forward 𝒥 assembles via a plain lift with no r² blow-up, so it needs no chunking — chunk_size is a transpose-only concept.)

What chunk_size does#

The assembly’s peak is exactly linear in |W|, so it is computed in slices of chunk_size probes and the partial results are combined — added when W is summed (sum_over_probes=True, the Gauss–Newton gradient) or concatenated when W is kept (sum_over_probes=False). The peak drops from |W|·(per-row) to chunk_size·(per-row), and the result is bit-for-bit the dense assembly (an exact reorganization, not an approximation).

# same numbers, bounded memory:
JTr_dense   = T3Tangent.probe_derivatives_transpose(r, ww, pp, frame, order, chunk_size=None)
JTr_chunked = T3Tangent.probe_derivatives_transpose(r, ww, pp, frame, order, chunk_size=512)
# JTr_chunked == JTr_dense, but the chunked run peaks at ~512/|W| of the memory

Semantics:

  • chunk_size = None (or any value ≥ |W|) → the dense assembly, no chunking.

  • chunk_size = <int> → chunk into slices of that many probes. The default is a small, safe fixed value.

  • Chunking engages only on the uniform + JAX path (a stacked supercore under jit), because only there does the sequential lax.scan/lax.map actually force the intermediates to be freed one slice at a time. On the ragged path or with NumPy, arrays are freed eagerly and the call falls back to the dense assembly regardless of chunk_size.

A fixed chunk_size bounds the chunk count, not the bytes#

Peak memory is chunk_size · (per-row bytes), and the per-row bytes grow like (order+1)²·r² (TT gradient) or (order+1)·nU·N (Tucker gradient, the term that dominates when the ambient dimension N ≫ n). So a single fixed default is only “safe for moderate problems”: comfortable at r = 128, but heavier at large ranks or large N. To actually bound the memory, choose chunk_size from the problem shapes (next section).

Choosing chunk_size#

The recommended policy is balance: pick chunk_size so the assembly’s peak is comparable to the edge-variable memory that is already necessarily resident. Then the assembly is never the tallest pole — if the rest of the pipeline fits, so does the assembly (total peak ≈ 2× the necessary floor). This is device-agnostic: it needs no knowledge of the device’s memory, only the problem shapes.

The estimator. Rather than guess, call estimate_chunk_size once (eagerly, outside jit) with the same shapes you used to build the problem, and pass the integer it returns as chunk_size. It measures the assembly’s true per-row cost with XLA’s own scratch accounting (memory_analysis) — no ~20×-off analytic formula — and returns the largest chunk whose peak stays comparable to the resident jets:

from t3toolbox.backend.sampling_derivatives import estimate_chunk_size, max_chunk_size_within

cs = estimate_chunk_size(
    mode_shapes=(N_1, ..., N_d),     # ambient dims (as passed to TuckerTensorTrain.randn)
    tucker_ranks=(nU_1, ..., nU_d),
    tt_ranks=(r_0, ..., r_d),        # the d+1 TT bonds
    order=K,
    n_probes=len_W,                  # number of probes
    n_tangent=1,                     # tangent-stack size (a batch of tangents at one frame)
    n_shards=1,                      # W split across this many devices -> sizes the LOCAL shard
    dtype=np.float32,                # float64 under x64
)
JTr = T3Tangent.probe_derivatives_transpose(r, ww, pp, frame, K, sum_over_probes=True, chunk_size=cs)

It sizes the larger of the TT-core (r² legs) and Tucker (nO, N legs) gradients, so it stays safe when N ≫ n; and when the whole assembly already fits the balance it returns chunk_size = |W| (the transpose then runs dense, no chunking overhead). The first call compiles (~2 s); results are cached by shape.

If instead you want to fill a known device — the “use my whole GPU” policy — use max_chunk_size_within(..., target_bytes=...), which returns the largest chunk_size whose assembly peak stays under an absolute byte cap (e.g. a fraction of device memory).

For sharded W, pass n_shards (or read it from the input array’s .sharding eagerly) so the returned chunk sizes each device’s shard; combine with the shard_map recipe below.

In fitting you get this for free. The optimizers (newton_cg, mc_sgd, adam, gradient_descent) and probe_derivatives_kind take a chunk_size argument defaulting to 'auto': for a uniform probe_derivatives fit they call estimate_chunk_size once with the shapes read off x0 (and the minibatch size for the minibatch optimizers), so a large-|W| fit stops OOMing with no action from you. Pass an int or None to override; ragged and non-probe_derivatives fits ignore it (nothing to chunk).

Sharding over W#

For very large |W| the sample stack is often sharded across devices. Chunking and sharding compose — but only through shard_map, and the reason is worth understanding.

chunk_size chunks within a device; sharding splits across devices. The transpose is a map-reduce over W (independent per-probe work, then a sum), which is exactly the shard_map shape:

Only the W-carrying arrays go through shard_map; the replicated frame and the scalar order are closed over (a scalar cannot take a PartitionSpec):

import jax
from functools import partial
from jax.sharding import PartitionSpec as P

@partial(jax.shard_map, mesh=mesh,
         in_specs=(P('w', None),   # ztildes: W sharded
                   P('w', None),   # ww:      W sharded
                   P('w', None)),  # pp:      W sharded
         out_specs=P())            # summed gradient: replicated
def jtr(ztildes, ww, pp):
    dU, dG = tv_probe_derivatives_transpose(          # frame (replicated), order: closed over
        ztildes, ww, pp, frame, order, sum_over_probes=True)   # chunk_size chunks the LOCAL shard
    return jax.tree.map(lambda a: jax.lax.psum(a, 'w'), (dU, dG))   # finish the W-sum across shards

Inside shard_map the body sees each device’s local shard as an ordinary array, so chunk_size chunks the local W with no cross-device traffic, and the add-reducer for sum_over_probes=True produces a per-shard partial that a single cheap psum finishes. (For sum_over_probes=False the output keeps W; use out_specs=P('w', ...) and drop the psum.) You only need to know which inputs carry W — the residual jets, ww, pp; the frame does not.

Pitfall: do not feed a globally-sharded W array to a plain jit#

If instead you shard W with the automatic partitioner and call the transpose under a plain jit (no shard_map), the per-chunk dynamic slice on the sharded axis forces XLA to all-gather the whole W axis onto every device — silently undoing the sharding (per-device memory jumps back to the full |W|) and adding communication. This is because the chunk loop’s slice offset is a runtime value, so the partitioner cannot prove which shard each window hits and conservatively replicates. Use shard_map (above) to shard W; it is the only way to get per-shard chunking without the all-gather.