3 posts tagged with "Engineering Notes"

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.

vidax now ports nine model families, and no two of them make exactly the same architectural choices. Rather than write a separate deep-dive per model, this is a tour of the handful of ideas that, combined differently, account for most of the differences: the Diffusion Transformer (DiT) itself, how text and image conditioning get into it, and Mixture-of-Transformers — an architecture that isn't a DiT variant at all.

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.