HeyGen Speeds Up Avatar IV Video Generation on Google Cloud TPUs by 1.86×

HeyGen, an AI video generation platform, partnered with Google Cloud to dramatically accelerate Avatar IV—its neural model for creating talking-head video content. By porting the system to an eight-chip Trillium (v6e) TPU cluster and applying targeted hardware optimizations, the team achieved a 1.86× speedup while maintaining identical output quality and reducing costs by as much as 25% per minute of video generated compared to their existing GPU infrastructure.

The Challenge: Streaming Under Deadline#

Avatar IV turns a single photograph and audio track into a photorealistic talking person. The system operates chunk by chunk, rendering video segments in real time: if a chunk arrives late, the playback stalls. This hard deadline shaped every optimization. The full pipeline comprises three components working in sequence on each chunk—two diffusion transformers (totaling over 36 GB of weights) and a VAE decoder that converts learned representations into pixels, ultimately producing 720p or 1080p video at 25 frames per second.

Porting the GPU-native model to TPUs involved torchax, a PyTorch frontend running atop JAX. The production code ran unmodified, dispatched to JAX arrays and compiled by XLA, allowing both GPU and TPU backends to use identical model code.

Parallelism and Memory: FSDP and SparseCore#

The two transformers exceed the 32 GB of memory available on a single Trillium chip, forcing fully sharded data parallelism (FSDP) across all eight chips. Trillium's SparseCore co-processor handles weight gathering asynchronously—a critical advantage that hides memory movement behind compute on the critical path. Ulysses sequence parallelism split the video sequence itself across the same mesh, combining both sharding strategies seamlessly.

Three Optimization Frontiers#

The team's engineering work progressed through six major milestones, each reducing time-per-chunk. After establishing custom attention kernels and tuning XLA compiler flags in the opening phase, three hardware-specific bottlenecks emerged:

Pipelining Collective Communications#

Ulysses parallelism requires two all-to-all communication passes per self-attention layer—one to exchange sequence shards for head shards, and one to swap back. These collectives initially blocked the entire pipeline, consuming 85–90% of the mesh's bandwidth but sitting synchronously on the critical path. The fix: split attention heads into independent groups, each with its own all-to-all pattern. This allowed later groups' transfers to hide behind earlier groups' compute, moving communication off the critical path and without changing wire throughput. The challenge lay in preserving mask structure in the super-resolution stage's sparse attention, where head grouping had to respect exact token ordering.