For the practical, day-to-day reference (which flags to pass, divisibility rules, when to reach for which scheme), see the Hardware & Sharding docs. This post is the "why" and the underlying mental model.

Why shard a video DiT at all​

A text-to-video DiT patchifies a video into a sequence of tokens — one per small spatio-temporal patch — and runs ordinary transformer self-attention over that whole sequence. For a short image, that sequence might be a few hundred tokens. For a few seconds of video at real resolution, it's tens of thousands. Self-attention costs O(sequence_length² × num_heads), so both the weights and the activations of a video DiT can each, independently, be too big for a single TPU chip's HBM (high-bandwidth memory — the on-chip memory an accelerator computes against, analogous to a GPU's VRAM). Two different sharding strategies fix two different halves of that problem.

Tensor parallelism (Megatron-style): shard the weights​

The idea: split each layer's weight matrix across N chips along its output dimension, so each chip only ever holds and computes with 1/N of the parameters. A minimal example, sharding a single dense layer's kernel across a device mesh axis called "tp":

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(jax.devices(), axis_names=("tp",))

# Split the *output* features across chips: each chip holds
# (in_features, out_features / tp_size) of the kernel.
kernel_sharding = NamedSharding(mesh, P(None, "tp"))
kernel = jax.device_put(kernel, kernel_sharding)

def dense(x, kernel):
    return x @ kernel  # each chip computes its own output slice

y = jax.jit(dense)(x, kernel)  # GSPMD inserts any needed communication

This is "column-parallel": each chip computes a different slice of the output, and no communication is needed until something needs the full result back (a "row-parallel" layer right after it, which shards its input dimension instead and sums the partial results with an all-reduce). Chaining column-parallel → row-parallel layers, which is exactly what an attention block's QKV projection → output projection and an MLP's up-projection → down-projection look like, means the all-reduce only has to happen once per block, not once per layer. This is the Megatron-LM scheme, and it's what vidax.core.sharding.build_tpu_mesh/shard_wan_params set up for every DiT and text encoder in vidax.

Tensor parallelism shards weights (and whatever activations are direct projections of them), but every chip still holds the full token sequence. That's fine until the sequence itself — not the weights — is what doesn't fit.

Sequence parallelism (DeepSpeed-Ulysses): shard the tokens​

The idea here is the opposite: keep the full weights on every chip, but give each chip only 1/N of the sequence. The tricky part is self-attention, which needs every token to see every other token — you can't just compute attention locally on each chip's slice. DeepSpeed-Ulysses's answer is to reshuffle data instead of tokens or weights: before attention, an all_to_all redistributes so each chip ends up holding every token but only some of the attention heads; attention runs locally per-head; a second all_to_all reshuffles back to "every chip, its own token slice, every head" for the rest of the block. Sketched out:

# Before: each chip holds (local_seq_chunk, all_heads)
# After all_to_all: each chip holds (all_seq, local_head_chunk)
kv_by_head = jax.lax.all_to_all(kv_by_seq, axis_name="sp",
                                 split_axis=0, concat_axis=1)
out_by_head = local_attention(kv_by_head)  # full sequence, few heads
out_by_seq = jax.lax.all_to_all(out_by_head, axis_name="sp",
                                 split_axis=1, concat_axis=0)

No chip ever recomputes another chip's tokens — it's a pure data reshuffle around one attention op, with ordinary per-chip compute for everything else (FFN, norms, modulation) running on that chip's own sequence slice the whole time. This is what actually shrinks the per-token activation memory that tensor parallelism can't touch.

The two compose: build_tpu_mesh builds a 3-axis (dp, tp, sp) mesh, and vidax's larger models (Wan2.2's A14B) need both at once, since one axis alone doesn't shard enough of either weights or activations to fit.

CogVideoX-1.5 is a second, trickier user of sequence parallelism. Its block has no separate cross-attention: text and video tokens sit in one joint self-attention sequence, so you can't just all_to_all the whole thing — the 226 replicated text tokens would be reshuffled too. The port (sequence_parallel_joint_self_attention) instead sends only the visual q/k/v through the head↔sequence all-to-all and slices the replicated text q/k/v down to each device's local head range before one local attention call. That's what runs CogVideoX-1.5 at its native 1360×768 (~45k visual tokens), whose per-block activations don't fit a v4 chip otherwise — and for CogVideoX the two schemes are mutually exclusive, because the 5B DiT's weights fit replicated per chip and the port never threads weight-sharding through the sequence-parallel path.

JIT compilation: what jax.jit actually does to a loop​

jax.jit doesn't compile your Python function — it traces it once (running it with abstract placeholder values to record what array operations happen) and compiles that trace into a single fused program. This matters a lot for any function containing a Python for loop: the loop itself runs at trace time, in Python, so a 20-iteration loop doesn't produce a 20-iteration compiled program — it produces one compiled program containing 20 iterations' worth of operations, fully unrolled, all of whose intermediate results the compiler may need to keep alive simultaneously.

The VAE decode "hang" that wasn't a hang​

Video VAE decoders in vidax decode a video's latent frames in chunks (a handful of frames at a time), in a loop, because decoding the whole video in one shot doesn't fit in memory. The very first version of this loop looked entirely reasonable:

@jax.jit
def decode_video(vae_params, latent_chunks):
    frames = []
    for chunk in latent_chunks:          # traced, i.e. unrolled!
        frames.append(vae.apply(vae_params, chunk))
    return jnp.concatenate(frames, axis=1)

Run against a real ~20-chunk video, this appeared to hang for 45+ minutes. It wasn't a deadlock — attaching py-spy to the running process showed it stuck deep inside XLA's compiler. @jax.jit on the outer function traces straight through the Python loop, so XLA was being asked to compile one enormous program containing 20 chunks' worth of decoder operations back to back, with no ability to free one chunk's intermediates before the next chunk's began — a compile-time and memory problem disguised as a runtime one.

The fix is to jit only the per-chunk function, and call it from an ordinary (untraced) Python loop:

decode_chunk = jax.jit(vae.apply)          # compiled once, for one shape

frames = []
for chunk in latent_chunks:                # a normal Python loop
    frames.append(decode_chunk(vae_params, chunk))
frames = jnp.concatenate(frames, axis=1)

Every chunk has the same shape, so decode_chunk is compiled exactly once — 20 calls to an already-compiled program, not one enormous from-scratch compilation. Just as importantly, each iteration's intermediate buffers can now be freed once that iteration's call returns, instead of all needing to coexist in one giant program's memory footprint. The rule of thumb this generalizes to, used everywhere in vidax that has a "do the same thing N times" loop (diffusion sampling steps, per-layer weight offloading's chunk loop): jit the repeated unit of work, not the loop around it.

A cross-model theme​

This isn't a Wan-specific fix — it's a pattern, and it keeps showing up wherever a model has some form of chunked, repeated decode step.

LTX-2.5's optional "diffusion" VAE decoder (a transformer that runs local, windowed attention over small 3D neighborhoods, rather than a single deterministic conv decode) hit the exact same shape of problem, one level deeper: a naive version that gathered every attention window at once before running any math didn't just compile slowly, it OOM'd outright — 147GB of intermediate memory for a tiny test shape, since the local window's volume multiplies memory by the product of all three of its axes. Processing one slice of the window at a time in an ordinary Python loop fixed the OOM but reintroduced the other half of this post's lesson: at real resolution, the loop unrolled into a genuine 45+ minute-class compile, once per block, 24 blocks deep. The fix was jax.lax.scan — JAX's built-in "repeat this step N times" primitive, which (unlike a Python for loop inside jax.jit) compiles the per-step body exactly once regardless of how many steps run, closing the compile-time problem the same way jitting only the per-chunk function did above. What remained after that was an ordinary memory-fit problem, closed with the same tensor-parallel sharding already built for every other model in vidax.

HunyuanVideo-1.5's VAE decoder ran into a version of this too — its decode OOM'd on a full-resolution run with peak memory that didn't add up to any single layer's own compute, the same "one fused trace doesn't free each stage's intermediates before the next starts" signature. The fix that landed is the same shape as all the others: on top of the spatial tiling the reference VAE already does (latent-space H/W tiles with a linear cross-fade of the overlaps), each tile is now decoded through a per-decoder-block staged pipeline — decode_stage_level_block / decode_stage_level_upsample, each its own jax.jit call — so one stage's temporaries are freed before the next stage runs, instead of the whole decode being one fused trace. That got the reference's real 121-frame/480p default decoding within budget.

Summary​

See the Hardware & Sharding docs for the practical, flag-by-flag reference, and Weight Offloading for the next technique built on top of this same "jit the unit, loop in Python" idea.

Diffusion Transformers (DiT)​

A video diffusion model's job is to predict, given a noisy video and a timestep, either the noise or the "velocity" needed to denoise it a step further. A DiT does this by patchifying the (spatio-temporal) video into a sequence of tokens and running an ordinary transformer over them, with the timestep injected into every block via adaptive layer norm (AdaLN): the timestep gets embedded and passed through a small MLP that predicts a per-block scale/shift/gate, applied around that block's attention and MLP sublayers. A minimal sketch:

def adaln_block(x, timestep_emb, attn, mlp, norm):
    # One small MLP per block predicts how this block should modulate
    # its own attention and MLP outputs, conditioned on the timestep.
    shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = (
        modulation_mlp(timestep_emb)
    )
    h = norm(x) * (1 + scale_attn) + shift_attn
    x = x + gate_attn * attn(h)

    h = norm(x) * (1 + scale_mlp) + shift_mlp
    x = x + gate_mlp * mlp(h)
    return x

This is the shape every DiT in vidax follows for its diffusion timestep — Wan's, Cosmos-Predict2.5's — though the details (how many separate modulation vectors, whether it's per-sample or per-token, the exact RoPE convention for positions) all differ. One family, Cosmos3, skips AdaLN entirely — more on that below.

Text conditioning: cross-attention​

Text almost always enters through ordinary cross-attention: video tokens as the query, text-encoder tokens as the key/value, inserted as an extra sublayer per block alongside self-attention:

def cross_attend_to_text(video_tokens, text_tokens, cross_attn):
    q = project_q(video_tokens)
    k, v = project_kv(text_tokens)
    return cross_attn(q, k, v)  # video tokens attend over text tokens

What differs across models is which text encoder produces text_tokens — a small T5 variant, or a multi-billion-parameter vision-language model — and, for Cosmos-Predict2.5, how many of that encoder's layers get used (concatenating hidden states from every decoder layer, not just the last).

The alternative: put conditioning tokens in the sequence, not a side branch​

Not every model uses cross-attention for text at all. HunyuanVideo 1.0 and 1.5 instead concatenate their text tokens directly into the same self-attention sequence the video tokens live in, tagged by a learned per-source embedding so the model can tell which tokens came from where. In 1.5 that's two separate towers — a 7B vision-language model and a small glyph/color encoder — plus, for I2V, SigLIP vision tokens; 1.0 is simpler, a single Llama-3-8B tower's hidden states (with a separate pooled CLIP-L vector folded into the AdaLN modulation, not the sequence). Either way there's no separate cross-attention sublayer — conditioning happens because every video token's self-attention already sees the text tokens sitting right next to it in the sequence.

CogVideoX takes the same idea to its minimal form: one joint self-attention per block over [text (226 tokens); video], no cross-attention sublayer and no second tower — just T5-XXL text embeddings prepended to the video sequence, with RoPE applied to the video slice only. Cosmos3's MoT pathway (below) is another variant: text and video tokens share one sequence, just processed by different weight sets depending which pathway they belong to.

Image/video conditioning (I2V)​

Conditioning on a starting image or clip (rather than pure text) shows up in vidax in two genuinely different forms:

Cross-attention to a separate image embedding. Wan2.1's I2V model adds a second cross-attention branch per block, attending to a CLIP embedding of the conditioning image — structurally identical to the text cross-attention above, just with a different key/value source.

Latent substitution. Wan2.2's TI2V-5B and Cosmos-Predict2.5 instead VAE-encode the conditioning frames, then substitute those known latents directly into the noisy input at the corresponding positions, at every sampling step — paired with an extra mask channel marking which frames are "given" versus "to be generated":

def substitute_known_frames(x, known_latents, is_conditioning_frame):
    # is_conditioning_frame: per-frame boolean mask
    return jnp.where(is_conditioning_frame, known_latents, x)

No extra cross-attention branch needed — conditioning is expressed purely through what's already in the input, plus the mask channel telling the model which parts of that input to trust as ground truth.

Mixture-of-Transformers: not sparse MoE​

Cosmos3 is the one model family in vidax that isn't a DiT at all. It's an omnimodal Mixture-of-Transformers (MoT) — a term that's easy to confuse with sparse Mixture-of-Experts (MoE), but describes something different. Sparse MoE routes individual tokens to a subset of expert weights via a learned gate. MoT instead gives each modality its own full weight set, used by every token of that modality, all the time — no routing, no gate:

def mot_layer(tokens, is_text, text_weights, video_weights):
    # every text token always uses text_weights; every video token
    # always uses video_weights -- decided by modality, not learned
    text_out = transformer_block(tokens, text_weights) 
    video_out = transformer_block(tokens, video_weights)
    return jnp.where(is_text, text_out, video_out)

Cosmos3 pairs this with a one-directional attention pattern: a causal text ("understanding") pathway that self-attends normally, and a full-attention "generation" (diffusion) pathway that attends over both its own tokens and the text pathway's — so text conditions video, but video never feeds back into text. And unlike every DiT above, there's no AdaLN anywhere: the timestep is embedded and added directly into the noisy video tokens once, then flows through ordinary pre-norm transformer blocks.

What each model in vidax actually uses​

LTX-2.5 also has a second, genuinely different VAE decoder option: instead of the usual deterministic convolutional decoder, --vae_variant diffusion runs a small transformer that samples its own internal denoising step directly on pixels — a different design from every other model's VAE here.

A few things worth knowing​

A handful of debugging lessons from getting these models running turned out to be useful diagnostic technique, not just "here's a bug we found" — worth keeping in mind for any diffusion model, not just the one that surfaced them.

A low-noise reconstruction probe separates "wrong noise-level handling" from "wrong weights." When Cosmos-Predict2.5's early output came out as a rigid grid of scrambled color blobs, the diagnostic that actually found the cause wasn't staring at attention or RoPE code — it was encoding a real photo with the VAE, adding a small amount of noise, running one forward pass, and decoding. That reconstructed the photo almost perfectly at low noise and degraded into the same grid-like garbage at high noise, which pointed straight at noise-level conditioning (the model was being told the wrong timestep, via a mistakenly-added preconditioning transform) rather than the network's weights or architecture. Any diffusion model producing texture-like or grid-like garbage is worth testing this way before assuming the bug is somewhere more exotic.

Generically-plausible-but-prompt-disconnected output usually means the conditioning signal is weak, not that CFG/precision is wrong. LTX-2.5 once produced video that looked structurally fine but kept drifting toward an unrelated recurring motif regardless of prompt or seed. The cause was a rescale bug in the text-feature extractor that fed the DiT a conditioning signal roughly 7x weaker than intended — not NaN, not wrong-shaped, just quietly under-scaled, so the model's own generative prior dominated. When output is "fine but ignores the prompt," it's worth checking the conditioning path itself before spending time on CFG scale, guidance rescale, or precision — a weak signal produces exactly the kind of plausible-but-generic output that's easy to misread as a model-quality ceiling instead of a bug.

A bf16 text encoder is only safe if the padding is masked. CogVideoX feeds its T5-XXL encoder no attention mask (the reference doesn't either), so the full 226-token padded sequence is attended. T5-XXL's residual stream reaches magnitudes around 1e5 — about two significant digits in bfloat16 — and over 24 layers × 226 unmasked positions the error compounds to 16–37% relative, with JAX-bf16 and torch-bf16 diverging from each other by that much too. Per-block parity against an fp32 input stays at ~1e-7; it's purely the bf16 residual stream that blows up. LTX-Video never hit this because it always passes the mask, so the pad positions get a -inf bias and never enter the sum. The fix isn't bf16 parity — the reference is unstable in bf16 too — it's running that one-time prompt encode in float32 and casting only the output embeddings down.

Prompt formatting requirements matter more for smaller models. Both Cosmos3 checkpoints document that prompts should be "upsampled" into a structured format (subjects, setting, lighting, cinematography, a temporal caption) rather than passed as a short sentence. Nano (16B) tolerates a short prompt reasonably well; Edge (4B) does not — the same short prompt that's fine on Nano produces flat, oversaturated, featureless output on Edge, with nothing in the pipeline (no error, no broken shape) to flag that the prompt itself is the limiting factor. Documented prompt-format requirements are easy to treat as polish; for a smaller model they can be the difference between unusable and excellent output.

See Sharding, parallelism, and JIT on TPUs and Weight offloading for how these architectures actually get run at scale, or the model guides under Model Family Guides for checkpoints and CLI usage.

For the practical reference (which flag, which chunk size, per-model guidance), see the Weight Offloading docs. This post covers the idea, the implementation, and the real cost of using it.

Background: what weight offloading is​

An accelerator's HBM has to hold everything needed for the current computation at once: weights, and everything the forward/backward pass produces along the way (activations — for a diffusion transformer, mostly the per-token modulation and attention tensors; the closest analogy for readers coming from LLM serving is the KV cache, though diffusion models don't have one). When the weights alone are a large fraction of that budget, there's less room left for everything else.

Weight offloading trades HBM for host-to-device bandwidth: keep the full weight tree in ordinary host (CPU) RAM — which is far more abundant and far cheaper than HBM — and copy only the small slice currently needed onto the accelerator, just before it's used, then let that HBM be reclaimed once that slice's computation is done. It isn't a new idea: ZeRO-Offload (and, closer to this exact use case, HuggingFace diffusers' enable_sequential_cpu_offload()) both do a version of the same trade. What makes it worth doing at all is that host-to-device transfer bandwidth is normally much higher than the compute time of the layer you're about to run — if the two overlap well, streaming weights in "for free" while the previous layer computes costs close to nothing.

How it's built in vidax​

Every DiT in vidax repeats one block shape num_layers times — the architectural property (nn.scan-style layer stacking) that makes transformers scale by depth in the first place. That means a single compiled program for "run one block/chunk of blocks" is reusable, verbatim, for every chunk in the model — the compiler never needs to know which chunk of weights it's holding, only that they match the shape/dtype it was compiled for. The loop then looks like a direct application of the "jit the unit, not the loop" rule from the sharding/JIT post:

chunk_apply = jax.jit(apply_one_chunk, donate_argnums=(0,))

for i in range(num_layers // chunk_size):
    chunk_params = jax.device_put(host_params[i], chunk_sharding)  # stream in
    x = chunk_apply(chunk_params, x)                               # compute
    # chunk_params' HBM is reclaimed once the next device_put reuses it
    # (donate_argnums lets the next chunk overwrite the same physical buffer)

chunk_apply is compiled exactly once, since every chunk has an identical signature. donate_argnums tells JAX it's safe to reuse chunk_params's physical HBM for the next chunk's device_put, so the resident footprint stays at one chunk's worth throughout a whole forward pass, not one chunk's worth accumulating on top of the last. --offload_chunk_size controls how many layers get grouped into one such chunk — a knob between "smallest possible HBM footprint" (chunk_size=1) and "fewest, largest transfers" (chunk_size=num_layers, i.e. the whole model).

What it costs​

TPU matmuls accumulate in float32 internally regardless of input dtype, so this technique's cost isn't numerical — it's a wall-clock race between transfer time and compute time. On paper this should be close to free: a 14B-class model's weights, sharded across 4 chips at bf16, are only a few GB per chip — well under a second to transfer at typical TPU host-to-device bandwidth, against tens of seconds of compute per diffusion step.

Measured, it wasn't free. Splitting that one transfer into 40 separate per-layer transfers (one per layer of a 40-layer DiT) means whether they overlap with the previous layer's compute — rather than serializing with it — now matters a great deal. On real hardware, offloading measured roughly 5x slower per step than keeping everything resident, at the chunk size that uses the least memory. Larger chunks (more layers streamed together per transfer, at the cost of a bigger resident buffer) close some but not all of that gap:

The takeaway: treat offloading as a correctness/memory-fit tool for configs that don't fit any other way, not a free option to reach for by default. It's the last lever to pull, after sharding, not the first.

The same idea, five different forms​

Wan2.1: fixing an OOM that wasn't where it looked​

Wan2.1 14B needed offloading at native 720p not because the DiT's own per-step compute didn't fit — a fully-sharded fp32 DiT comfortably fits while the sampling loop runs — but because the DiT's weights were staying HBM-resident for the entire script, competing with an unrelated phase's memory: VAE decode right after sampling ends, VAE encode right before an I2V sampling loop starts. Offloading fixes this by only ever making the weights HBM-resident for the brief window each chunk actually needs them, freeing that memory back up for the next phase.

A14B (Wan2.2): composing with sequence parallelism​

Wan2.2's larger A14B model computes its timestep modulation per token instead of per sample, which makes that one tensor alone multiple GB at high resolution — pure activation memory that offloading, which only targets weight memory, can't shrink no matter the chunk size. Fitting A14B at real resolutions needed offloading and sequence parallelism at once, which meant extending offloading to compose with a sharded token sequence rather than assuming a full replica.

That combination surfaced a genuinely interesting bug, worth sharing as an example of how subtle a sharding bug can look: outputs came out blocky and spatially discontinuous, reproducible even at tiny resolutions. The cause traced back to how a row-parallel layer (one where each chip holds a slice of the input dimension, computes a partial output, and sums partial outputs across chips) handles its bias. Each chip's local matmul call adds the (replicated) bias before the cross-chip sum runs — so summing N chips' outputs summed N copies of the bias instead of one. Invisible with an all-zero bias (every synthetic test used one), but a real per-layer error with trained biases, compounding across 40 layers into a large enough numerical drift to visibly corrupt output. The fix subtracts the bias before the cross-chip sum and adds exactly one copy back after — a pattern worth remembering any time a row-parallel layer has a bias term at all.

Cosmos-Predict2.5: the straightforward case​

Cosmos-Predict2.5's 14B DiT needed the same technique to reach its reference's full 93-frame default (previously reduced to 45 frames without it) — a single DiT with no MoE-style expert switching to compose with, so the implementation is close to the offloading loop shown above, essentially unmodified. With offloading, the full 93 frames run comfortably at 14.7GB peak HBM/chip.

HunyuanVideo 1.0: two chunk pools, one model​

HunyuanVideo 1.0's 13B DiT needs offloading to fit the reference's real 129-frame / 720p default in HBM at all — a plain fit-the-sampling-loop case like Wan2.1 and Cosmos above. The wrinkle is architectural: this DiT is a dual-stream / single-stream MMDiT, 20 "double-stream" blocks followed by 40 "single-stream" blocks, and the two block types have different parameter shapes. One chunk pool can't stream both. So --offload_dit_weights runs two independent chunk loops — --offload_chunk_size_double (default 20, i.e. all double blocks in one chunk) and --offload_chunk_size_single (default 40) — each its own separately-compiled chunk_apply, streamed and freed in sequence. It's the same "jit the unit, stream one chunk at a time" machinery as every other model here, just instantiated once per block family instead of once for the whole stack. The per-step cost is dominated by this offloading tax, not the implementation — HunyuanVideo-1.5, the same block family with zero offloading, is already the slowest non-offloaded row in vidax's benchmark table.

LTX-2.5: offloading for a reason no other model needed it​

LTX-2.5's 22B DiT needed --offload_dit_weights to reach its reference resolution too, but for a genuinely different reason than every case above: its weights were never the bottleneck. Fully resident, they're only ~6.6GB/chip — comfortable headroom. What actually didn't fit was activation memory from tracing the entire 48-block forward pass as one fused jax.jit program, which — confirmed by measurement, not assumed — scaled almost linearly with block count instead of staying roughly constant, meaning per-block intermediates weren't being freed across blocks. Offloading fixes this as a side effect: splitting the block loop into --offload_chunk_size-sized separately-compiled chunks bounds peak activation memory to one chunk's worth, by construction, independent of whether weight streaming itself was ever needed. --offload_chunk_size 8 (the largest divisor of 48 that fits) got the reference's full 1216×704×121 resolution running at both checkpoints.

A useful contrast, though: LTX-2.5's second, optional VAE decoder — a transformer that runs local windowed attention instead of an ordinary conv decode — hit a real memory ceiling of its own that offloading didn't fix. Isolated single-block profiling suggested each block needed only ~12GB, comfortably under budget, but the real multi-device pipeline measured over 32GB for the same block — a genuine, still unexplained gap between isolated and in-context profiling. Chunking further didn't move the number at all. What did was reaching for a tool this repo already had a working pattern for: Megatron-style tensor parallelism across the same mesh the DiT and text encoder already use, which dropped peak memory to a comfortable 14.76GB/chip. The lesson worth generalizing: not every "doesn't fit" problem is a streaming problem. Offloading trades HBM for bandwidth by keeping less resident at once; sharding trades HBM for communication by splitting the work itself across chips. When one doesn't move the number, it's worth checking whether the other, already-built machinery applies instead of continuing to tune the first.

See Sharding, parallelism, and JIT on TPUs for the mesh/JIT background this builds on, and the Weight Offloading docs for the practical, per-model reference.