2 posts tagged with "Infrastructure"

Video DiTs are big in two different ways at once: the weights can be too large for one chip, and — because self-attention runs over every patch in every frame — a single forward pass's activations can be too large for one chip even when the weights fit fine. This post is background on the two sharding strategies vidax uses to deal with that (tensor parallelism, sequence parallelism), and on a JAX-specific gotcha around jax.jit and Python loops that's caused more real bugs, in more of vidax's models, than anything else in this codebase.

Some of vidax's models don't fit fully on-chip at the resolutions their reference implementations were trained at, even after sharding weights and activations across every available chip (Sharding, parallelism, and JIT on TPUs). Weight offloading is the technique that closes that last gap: instead of sharding, don't keep the weights resident on the accelerator at all — leave them in host RAM and stream them in just in time.