HeyGen and Google Cloud: Bringing Avatar IV to TPUs
Revolution in AI video generation: 1.86x performance boost achieved
1HeyGen's AI Video Platform
HeyGen offers an avatar-based AI video generation platform. Its Avatar IV technology utilizes a diffusion model with over 18 billion parameters. The system can produce talking-head videos from just a single photo and audio recording. Available via web and API, the platform delivers 720p or 1080p video streams at 25 fps.
2Google Cloud Collaboration
HeyGen partnered with Google Cloud's AI infrastructure specialists to optimize Avatar IV. Migrated to an eight-chip Trillium (v6e) host system, the model became 1.86x faster than its initial version. All optimizations were made against strict latency deadlines for seamless video streaming, while maintaining output quality standards throughout the process.
3Technical Challenges
Three major technical barriers were identified: exposed all-to-all collectives in the mesh, partial blocks in the sparse attention grid, and a serial dependency in the softmax inner loop. These posed significant challenges in porting the model to TPUs. Through collaboration with Google Cloud, specialized solutions and compiler contracts were developed to overcome each obstacle.
4Model Architecture
Avatar IV comprises three core models: a diffusion transformer for audio-conditioned motion generation, a second transformer for super-resolution, and a VAE decoder converting latents to pixels. Originally written for GPUs, the model was ported to JAX using torchax (a PyTorch frontend), enabling the same code to run on both GPU and TPU stacks with hardware-specific optimizations.
5Parallelism and Sharding
The transformers' combined 36GB+ weights were distributed across eight Trillium chips (each with 32GB HBM) using FSDP-sharding. Trillium's SparseCore co-processor handles weight gathers asynchronously. Production traces show these gathers hidden behind computation, effectively removing weight movement from the critical path without performance penalties.
6Ulysses Sequence Parallelism
Ulysses sequence parallelism requires bidirectional all-to-all communication in each self-attention operation. Initially these collectives ran synchronously, consuming 85-90% of mesh bandwidth. The Google Cloud team implemented a pipelined solution where attention heads split into independent groups, enabling asynchronous transfers that were removed from the critical path.
7Sparse Attention Optimization
The super-resolution stage's sparse attention operates across tens of thousands of tokens with a windowed sparsity pattern. Its masks are naturally frame-aligned, meaning active regions precisely match frame boundaries. Post-optimization, this stage's computational load saw significant reduction while maintaining output quality.
8Performance Metrics
Six key milestones marked the optimization journey. Final results showed a 1.86x speedup, halving per-chunk processing time. The most substantial improvement came initially through custom attention kernels, sequence-parallel layout, and XLA flag tuning. Subsequent refinements addressed more nuanced bottlenecks in the pipeline.
9Cost Efficiency
The optimized TPU pipeline delivers performance comparable to an 8×H100 production setup while achieving up to 25% better cost efficiency per minute of generated video. This advantage proves particularly valuable for businesses requiring high-volume video generation at scale.
10Quality Assurance
All optimizations preserved the original model's output quality standards. Each change passed identical quality gates before approval. Interestingly, some optimizations that showed promise in isolated tests failed to maintain performance at full pipeline depth and were consequently discarded.
11Industry Impact
The HeyGen-Google Cloud collaboration sets a new benchmark for AI video generation. TPU optimizations make high-quality, real-time avatar videos more accessible than ever. This breakthrough holds revolutionary potential across digital content creation, e-learning, virtual assistants, and beyond, reshaping expectations for synthetic media.