Wang et al. (2022), Baidu — ARXIV-2211.00235
Key takeaways
- The problem is a batch size that cannot grow. AlphaFold2 trains at a total batch of 128 with one protein per device. Data parallelism therefore stops at 128 devices, and the paper keeps those hyperparameters fixed rather than change them. The Evoformer accounts for 62% of step time in initial training and 78% in fine-tuning (HelixFold, 8 A100s).
- Parallel Evoformer is a one-line architectural change. In AlphaFold2 the outer product mean, which carries information from the MSA track to the pair track, sits mid-block. The pair stack therefore waits for the MSA stack. AlphaFold-Multimer moved it to the start of the block. This paper moves it to the end, so both stacks evolve independently within a block and exchange information only at the boundary. Trained from scratch on HelixFold with 8 A100s, 120,000 PDB proteins and 40,000 steps, all three block variants reach “competitive” TM-score and lDDT-Cα on CASP14 (87 domains) and CAMEO (371 targets). The comparison is a figure, with no numbers in the text. Step time is unchanged: 5.00 s for the original block and for the parallel one.
- Branch parallelism puts one stack on each GPU. One device computes the MSA stack and the outer product mean; the other computes the pair stack. The results are exchanged by broadcast, and gradients by broadcast and all-reduce. Activations are not split, so each device’s kernels keep the computational intensity they had. The method scales only to as many devices as there are branches, two here, and only pays when the branches cost about the same.
- Measured gains, BF16, total batch 128. Throughput improves by 37.7% (initial training) and 40.4% (fine-tuning) on UniFold, and by 38.6% and 33.8% on HelixFold. FP32 initial training improves by 30.1% and 42.4%. End to end, UniFold falls from 5.80 to 4.18 days (+38.67%) and HelixFold from 6.69 to 4.88 days (+36.93%), but on 256 A100s against 128. On this record’s arithmetic, doubling the GPUs for about 1.38× throughput is roughly 69% parallel efficiency. The gain is faster wall-clock time, not cheaper training.
- Against dynamic axial parallelism (FastFold’s DAP), it depends on shape. At initial-training sizes, splitting activations across 2 GPUs is slower than not splitting: -3.9% per layer in FastFold, -2.2% in the authors’ Paddle re-implementation, and -4.7% end to end. Branch parallelism on 2 GPUs gives +67.5% per layer and +38.5% end to end. At fine-tuning sizes the order flips: DAP=2 gives +48.1% and BP=2 gives +33.8%. On the same 8 GPUs, DAP=4 with BP=2 gives +188.7%, against +125.0% for DAP=8. Tensor and pipeline parallelism are not compared; the authors cite FastFold’s comparison and the implementation cost.
Standing in the anthology
It extends LIT-583, AlphaFold 2. It takes AlphaFold2’s Evoformer block and its training recipe unchanged and changes one thing: where the outer product mean sits, which removes the dependency between the two tracks. It is compared against LIT-583 in the sense that matters for that change: the original Evoformer block was trained from scratch as one arm of the accuracy comparison and timed as one arm of the speed table, and it matched the parallel block on both. For end-to-end training time, though, AlphaFold2’s own 128-TPUv3 figure is quoted rather than rerun. The record holds AlphaFold 2 for what it predicts and why (SOTA-353, SOTA-354, THEORY-084); this is the first paper in it on what AlphaFold costs to train.
The general lesson is the one in LIT-088, reached from another
direction. ViT-22B, like PaLM, runs attention and MLP in parallel within a
block so that their matrix multiplies fuse. This paper makes two branches
parallel within a block so that they can run on different devices. In both
cases a sequential dependency inside the block turns out to cost little or
no accuracy, and removing it buys speed. The record has practices for the
usual answers when one device is not enough, micro-batched pipelines
(SOTA-017) and sequence parallelism beside tensor parallelism (SOTA-058).
Those answers assume a batch or a layer large enough to split. Here the
batch is one protein and the kernels are small, and the paper’s DAP
comparison shows what splitting activations costs in that regime. As
ADR-059 directs for a method found in AlphaFold, it is filed by its kind
(distributed training) first and under biomolecular-modeling as well.
Unread — no NOTE.