# HeyGen and Google Cloud: Bringing Avatar IV to TPUs Revolution in AI video generation: 1.86x performance boost achieved - **Source:** [Google Developers Blog](https://brifing.si/source/1e914d4a-9164-4a1a-82cb-3d169fbad877) - **Topic:** [Data & AI](https://brifing.si/c/data-ai) - **Original article:** - **On Brifing:** https://brifing.si/s/4c36e7ba-7d0a-4458-9a63-1d208e1b9865 ## HeyGen'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. ## Google 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. ## Technical 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. ## Model 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. ## Parallelism 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. ## Ulysses 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. ## Sparse 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. ## Performance 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. ## Cost 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. ## Quality 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. ## Industry 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. ## Sources covering this story - https://developers.googleblog.com/heygen-x-google-cloud-bringing-avatar-iv-to-tpus/ — --- _Summary written by Brifing from the original reporting. Copyright in the original article belongs to its publisher; read the full piece at the link above._