Fixes CP support#4
Open
zyndagj wants to merge 4 commits into
Open
Conversation
…edups
Extends ESMFold2's 2D context-parallel (CP) path beyond the FoldingTrunk to the
MSA encoder, and adds memory/precision optimizations that make the 2402-residue
5xgo complex fold ~2x faster and well under the 80 GB limit on a 2x2 grid.
MSA encoder CP (forward-only inference; the encoder runs the full LxL pair every
recycling loop and previously OOMed on every rank):
- OuterProductMeanDistributed: transpose-based (MSA depth replicated), bit-exact.
- MSAPairWeightedAveragingDistributed: comm="gather" (default, bit-exact) or
"ring" (boltz-style online-softmax).
- MSAEncoderDistributed / MSAEncoderCPWrapper: drop-in for the serial MSAEncoder.
- wrap_model_with_cp(): convenience wrapper for trunk + MSA encoder.
Trunk performance (2x2 grid, 5xgo, vs fp32 baseline 1359 s / 73.8 GB):
- bf16 distributed trunk: cast params + pair to bf16, output back to caller dtype
(boltz-cp runs its CP trunk pair in bf16). Quality-neutral: pLDDT 0.643 vs 0.639.
- ESM-C offload: move the ~12 GB ESM-C 6B LM to CPU after its one-shot
hidden-state computation (restored next call), freeing it for the trunk +
diffusion. This was the dominant cost — resident-but-idle ESM-C kept the
allocator at ~84% occupancy, whose cudaMalloc stalls drove rank desync and ring
spin-wait. Combined with bf16: 636 s / 55.3 GB (2.14x faster, -18.5 GB peak),
pLDDT unchanged (0.640).
- Inference fast path in the replicated Linear/LayerNorm: skip the autograd
Function + ctx bookkeeping and cache replicated weight locals when grad is off.
API: bf16 and ESM-C offload are wrapper arguments, not env vars:
wrap_model_with_cp(model, dm, comm="gather", bf16=True, offload_esmc=True)
wrap_model_with_cp_trunks(model, dm, bf16=True)
offload_esmc sets model._offload_esmc (instance attr, default False on the base
model; opted in by the CP wrapper).
Distributed manager: disable NCCL NVLS by default (NCCL_NVLS_ENABLE=0; opt back in
on NVSwitch + Fabric Manager systems), fall back to lazy init when device_id is
unsupported or eager NCCL connect fails, and normalize get_coordinate() across
PyTorch versions. Fix vendored import paths (projects.huggingface.transformers ->
transformers).
…ence, LM→pair, ESM-C)
Extend 2-D context parallelism from the MSA encoder + trunk to the rest of the
model so per-rank peak memory DROPS as the CP mesh (P×P) grows. Previously only
the Pairformer trunk and MSA encoder were sharded; every other stage ran serial
at full length L on every rank and re-gathered the full L×L pair at its
boundaries, so adding GPUs past 4 did not raise (and could lower) the max input
length. The pair now stays sharded (Shard(0),Shard(1),Shard(2)) end to end.
New distributed components (installed by wrap_model_with_cp via CP-agnostic seams
in modeling_esmfold2.py; serial path unchanged when unwrapped):
- recycle.py (CPRecycleEngine): sharded recycle loop, no per-iteration full-L×L
gather; handles per-loop LM dropout and the parcae tail.
- structure_wrapper.py + diffusion_{transformer,conditioning}.py +
attention_pair_bias.py + atom_to_token.py: distributed diffusion structure head
(conditioned-z + token attention sharded; gather mode, num_diffusion_samples==1).
- single_to_pair.py: sharded LM→pair builder — lm_z built directly as a sharded
DTensor (no full L×L), shard-local per-loop dropout.
- esmc_tp.py: opt-in (tp_esmc=True) TE-native tensor parallelism for the ESM-C
SwiGLU MLP across the CP ranks.
- confidence_wrapper.py + confidence_zbase.py + row_attention_pooling.py:
distributed confidence head (z_base + nested trunk + row-pool + PAE/PDE sharded;
pTM/ipTM reductions reuse a shared serial _finish).
- distogram_wrapper.py: distributed distogram head (symmetrize via transpose comm).
modeling_esmfold2.py gains the CP-tail seams and a behavior-preserving
ConfidenceHead._finish refactor. Validation: spawn-based 2x2 bit-exact unit tests
per component + real-checkpoint fold pLDDT/pTM parity (tests/models/esmfold2/).
Numerically faithful — cross-mesh pLDDT spread <1e-3, ring≡gather.
-SVD fallsback to CPU
The initial pair tensors (z_init, relative-position encoding, token-bond encoding, the initial pair state, and the pair attention mask) were built full [B, L, L, *] on every rank and only sharded at the recycle-loop boundary via distribute_tensor. That made the per-rank init peak scale like a single GPU — several full L×L tensors resident at once, plus rel_pos / token_bonds held full through the whole forward — so CP gave no memory relief for the init phase and OOM'd at short L. Build each of these directly as a 2-D-sharded (Shard0,Shard1,Shard2) DTensor by computing only this rank's (row, col) block, mirroring evolutionaryscale's _sharded_rel_pos_encoding + sliced z_init. Per-token inputs stay replicated; only their O(L²) outer combinations are sharded. - pair_init.py (new): PairInitDistributed builds z_init/rel_pos/token_bonds/ initial-state blocks; build_sharded_pair_mask and build_sharded_distogram_bins helpers (block outer-product / block cdist). - recycle.py: run_loop / parcae_finish accept already-sharded z / z_init / pair_mask DTensors (+ explicit n_orig); no distribute_tensor of a full tensor. - structure_wrapper.py: diffusion conditioning consumes the sharded rel_pos DTensor directly (removes the per-diffusion-step full-rel_pos floor). - msa_wrapper.py: install _cp_pair_init; MSA per-iteration pair mask via the shared build_sharded_pair_mask (was a full [B,L,L] every recycle step). - confidence_wrapper.py / confidence_zbase.py: consume sharded rel_pos / token_bonds / distogram_bins (block cdist) and build the nested-trunk pair mask sharded — removes the end-of-forward gather. - modeling_esmfold2.py: route pair-init + pair mask through the CP seam; drop the transient confidence-boundary gather. Serial path unchanged (isinstance DTensor guards). Every O(L²) tensor on the CP inference path is now sharded at creation. Block builders verified bit-exact vs the serial full construction on a 2×2 grid (pair_init_cp_parity_test.py), including the padded (non-divisible-L) case.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Distributes the entire ESMFold2 forward pass across a context-parallel mesh so that no rank ever materializes a full L×L tensor. This enables inference on sequences that would OOM on a single GPU.
What's distributed:
Supporting fixes:
Validation