I am working on a legal-tech AI that needs to process 100k+ token sequences. Standard Data Parallelism fails, and even ZeRO-3 struggles with the activation memory at that length. Is DeepSpeed Ulysses the right path? I've read it uses sequence parallelism to split the workload. How does it compare to FlashAttention-2 in terms of memory saving for long-context deep learning?
3 answers
DeepSpeed Ulysses is specifically designed for this scenario. Unlike standard parallelism, Ulysses splits the sequence dimension across GPUs. It uses an "All-to-All" communication pattern to reshuffle data before and after the attention computation. The beauty of it is that it is complementary to FlashAttention-2. You use FlashAttention to save memory on a single GPU, and you use Ulysses to scale that across multiple GPUs. This combination is what allowed teams to break the 100k and even 1M token barrier. It is much more communication-efficient than older sequence parallel methods.
Does Ulysses require any specific interconnect hardware like InfiniBand to be effective, or will standard Ethernet work for a small cluster?
We switched to Ulysses for our document analysis model and the memory scaling was linear. It's much easier to configure than Pipeline Parallelism for long sequences.
Gary is spot on. Pipeline Parallelism introduces "bubbles" and complexity that Ulysses simply avoids by keeping the model layers intact on each device while sharding the data.
Edward, while it works on Ethernet, the All-to-All communication is very heavy on bandwidth. If you're on a 10Gbps Ethernet, you'll likely see the GPUs idling while waiting for data. For long contexts with DeepSpeed Ulysses, you really want NVLink or InfiniBand to ensure that the communication doesn't completely eclipse the actual computation time.