Design-CP: Context Parallelism for Design of Protein Nanoparticles

Lorenzo Tarricone, Helen Elizabeth Eisenach, Aiko Muraishi, Charlotte Deane

Keywords: Machine Learning, BioML, Protein Design, Parallelism, Protein Nanoparticles

TL;DR: Design-CP enables end-to-end all-atom design of large symmetric protein nanoparticles with RFdiffusion 3 by context-parallel sharding of quadratic activations across GPUs, improving memory feasibility and scaling without changing pretrained weights.

Abstract:

Many all-atom generative protein models can in principle design large multimeric complexes by jointly modelling all chains, but their quadratic token- and atom-pair representations quickly exceed single-GPU memory as the number of chains and residues modelled grows. We introduce Design-CP, two context-parallel (CP) inference strategies for RFdiffusion 3 (1D row-sharding and 2D grid sharding with ring attention) that distribute the quadratic activations across a multi-GPU mesh while preserving pretrained weights. We characterise their scaling when sampling icosahedral assemblies, showing that the maximum feasible asymmetric subunit (ASU) size grows with the expected square-root trend in GPU count and that 2D sharding achieves better wall-clock scaling. Moreover, we show how strong
point-group symmetry constraints make CP usable out of the box for end-to-end, all-atom design of icosahedral nanoparticles, yielding favourable in silico structural and interface metrics. Finally, we demonstrate octahedral nanoparticle design on a small cluster of workstation-grade 16 GB GPUs, illustrating how Design-CP can be a practical path towards democratising large-assembly
protein design.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 68