From 60eba48440914c7437cffa72c990c3b572598d27 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 15 Jun 2026 08:28:09 -0700 Subject: [PATCH 1/4] ESMFold2: context-parallel MSA encoder + bf16/ESM-C-offload trunk speedups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../models/esmfold2/distributed/__init__.py | 18 +- .../models/esmfold2/distributed/comm.py | 2 +- .../models/esmfold2/distributed/manager.py | 43 ++- .../distributed/model/layers/layernorm.py | 43 ++- .../distributed/model/layers/linear.py | 35 ++- .../distributed/model/layers/msa_encoder.py | 210 ++++++++++++++ .../model/layers/outer_product_mean.py | 146 ++++++++++ .../model/layers/pair_averaging.py | 270 ++++++++++++++++++ .../distributed/model/layers/pairformer.py | 16 +- .../model/layers/triangular_mult.py | 40 ++- .../esmfold2/distributed/msa_wrapper.py | 220 ++++++++++++++ .../models/esmfold2/distributed/utils.py | 58 +++- .../models/esmfold2/modeling_esmfold2.py | 16 ++ 13 files changed, 1066 insertions(+), 51 deletions(-) create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/msa_encoder.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/outer_product_mean.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/pair_averaging.py create mode 100644 src/transformers/models/esmfold2/distributed/msa_wrapper.py diff --git a/src/transformers/models/esmfold2/distributed/__init__.py b/src/transformers/models/esmfold2/distributed/__init__.py index 1854a7aa0a..5e8f86711d 100644 --- a/src/transformers/models/esmfold2/distributed/__init__.py +++ b/src/transformers/models/esmfold2/distributed/__init__.py @@ -21,13 +21,21 @@ """2D context-parallel distributed extensions for ESMFold2.""" -from projects.huggingface.transformers.models.esmfold2.distributed.manager import ( +from transformers.models.esmfold2.distributed.manager import ( DistributedManager, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.pairformer import ( +from transformers.models.esmfold2.distributed.model.layers.msa_encoder import ( + MSAEncoderDistributed, +) +from transformers.models.esmfold2.distributed.model.layers.pairformer import ( FoldingTrunkDistributed, ) -from projects.huggingface.transformers.models.esmfold2.distributed.utils import ( +from transformers.models.esmfold2.distributed.msa_wrapper import ( + MSAEncoderCPWrapper, + wrap_model_with_cp, + wrap_model_with_cp_msa_encoder, +) +from transformers.models.esmfold2.distributed.utils import ( TrunkCPWrapper, wrap_model_with_cp_trunks, ) @@ -35,6 +43,10 @@ __all__ = [ "DistributedManager", "FoldingTrunkDistributed", + "MSAEncoderCPWrapper", + "MSAEncoderDistributed", "TrunkCPWrapper", + "wrap_model_with_cp", + "wrap_model_with_cp_msa_encoder", "wrap_model_with_cp_trunks", ] diff --git a/src/transformers/models/esmfold2/distributed/comm.py b/src/transformers/models/esmfold2/distributed/comm.py index 6a8f803ebb..643cb06ec2 100644 --- a/src/transformers/models/esmfold2/distributed/comm.py +++ b/src/transformers/models/esmfold2/distributed/comm.py @@ -26,7 +26,7 @@ import torch import torch.distributed as dist -from projects.huggingface.transformers.models.esmfold2.distributed.utils import ( +from transformers.models.esmfold2.distributed.utils import ( LayoutMap, get_group_rank_from_axial_shift, ) diff --git a/src/transformers/models/esmfold2/distributed/manager.py b/src/transformers/models/esmfold2/distributed/manager.py index f2b2c4b7e1..37d125d9bf 100644 --- a/src/transformers/models/esmfold2/distributed/manager.py +++ b/src/transformers/models/esmfold2/distributed/manager.py @@ -21,14 +21,13 @@ import os -from copy import deepcopy from math import prod from typing import Any, Dict, Optional, OrderedDict, Union from warnings import warn import torch -from projects.huggingface.transformers.models.esmfold2.distributed.utils import ( +from transformers.models.esmfold2.distributed.utils import ( LayoutMap, LayoutRightMap, ) @@ -220,6 +219,11 @@ def _setup( manager._backend = backend # type: ignore[assignment] if manager.device.type == "cuda" and backend == "nccl": + # Disable NVLS (NVLink SHARP multicast) by default — it requires + # NVSwitch hardware + Fabric Manager, and probing for it via + # cuMulticastCreate raises a hard CUDA error (401) when unavailable. + # Users on NVSwitch systems can opt in by setting NCCL_NVLS_ENABLE=1. + os.environ.setdefault("NCCL_NVLS_ENABLE", "0") try: torch.distributed.init_process_group( manager.backend, @@ -228,7 +232,11 @@ def _setup( device_id=manager.device, **kwargs_init_pg, ) - except TypeError: + except (TypeError, RuntimeError): + # TypeError: older PyTorch doesn't accept device_id. + # RuntimeError: device_id triggers eager NCCL connect which + # can fail (e.g. NVLS/NVSwitch not available, CUDA error 401). + # Fall back to lazy (non-eager) init in both cases. torch.distributed.init_process_group( manager.backend, rank=manager.rank, @@ -383,9 +391,11 @@ def create_grid_group(grid_group_sizes: _GridGroupSizesType) -> None: name_subgroups, shape_subgroups, suffix_mesh="subgroups" ) layout = DistributedManager._state["_layout_device_mesh_subgroups"] - coords = DistributedManager._state[ + # get_coordinate() returned a list in older PyTorch and a tuple in + # newer versions; normalise to list so axis assignment works. + coords = list(DistributedManager._state[ "_device_mesh_subgroups" - ].get_coordinate() + ].get_coordinate()) for name_group, subgroup_names in group2subgroup.items(): DistributedManager._state["_subgroups"][name_group] = [ DistributedManager._state["_group"][n] for n in subgroup_names @@ -397,7 +407,7 @@ def create_grid_group(grid_group_sizes: _GridGroupSizesType) -> None: DistributedManager._state["_group_rank"][n] for n in subgroup_names ] axes_subgroup = group2subgroup_axes[name_group] - slices = deepcopy(coords) + slices = coords.copy() for axis in axes_subgroup: slices[axis] = slice(None) layout_subgroup = layout[tuple(slices)] @@ -562,15 +572,18 @@ def initialize( def cleanup(): if DistributedManager._state.get("_group", {}) != {}: if torch.distributed.is_initialized(): - if ( - DistributedManager._state["_device"].type == "cuda" - and torch.cuda.is_available() - ): - torch.distributed.barrier( - device_ids=[DistributedManager._state["_local_rank"]] - ) - else: - torch.distributed.barrier() + try: + if ( + DistributedManager._state["_device"].type == "cuda" + and torch.cuda.is_available() + ): + torch.distributed.barrier( + device_ids=[DistributedManager._state["_local_rank"]] + ) + else: + torch.distributed.barrier() + except Exception: + pass torch.distributed.destroy_process_group() else: DistributedManager._state = {} diff --git a/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py b/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py index 88fec07fa0..83849dbf81 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/layernorm.py @@ -194,12 +194,41 @@ def __init__(self, layer_local: nn.LayerNorm, device_mesh: DeviceMesh) -> None: else: self._reduce_group = dist.group.WORLD + # Local-tensor cache for the inference fast path (see LinearParamsReplicated). + self._w_local: Optional[torch.Tensor] = None + self._b_local: Optional[torch.Tensor] = None + + def _inference_locals(self): + w = self.weight + cur_dtype = None if w is None else w.dtype + if (w is not None and self._w_local is None) or ( + self._w_local is not None and self._w_local.dtype != cur_dtype + ): + self._w_local = w.to_local() + self._b_local = self.bias.to_local() if self.bias is not None else None + return self._w_local, self._b_local + def forward(self, x: DTensor) -> DTensor: - return _LayerNormParamsReplicatedImpl.apply( # type: ignore[return-value] - x, - self.normalized_shape, - self.weight, - self.bias, - self.eps, - self._reduce_group, + if torch.is_grad_enabled(): + return _LayerNormParamsReplicatedImpl.apply( # type: ignore[return-value] + x, + self.normalized_shape, + self.weight, + self.bias, + self.eps, + self._reduce_group, + ) + # Inference fast path: skip the autograd.Function machinery + ctx stores, + # reuse cached replicated param locals. Bit-identical to the Function's + # forward (same to_local → F.layer_norm → from_local). + w_local, b_local = self._inference_locals() + out_local = F.layer_norm( + x.to_local(), self.normalized_shape, w_local, b_local, self.eps + ) + return DTensor.from_local( # type: ignore[return-value] + out_local, + device_mesh=x.device_mesh, + placements=x.placements, + shape=x.shape, + stride=x.stride(), ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/linear.py b/src/transformers/models/esmfold2/distributed/model/layers/linear.py index ba195f0509..575a364799 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/linear.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/linear.py @@ -31,7 +31,7 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Replicate, distribute_tensor -from projects.huggingface.transformers.models.esmfold2.distributed.utils import ( +from transformers.models.esmfold2.distributed.utils import ( update_exhaustive_strides, ) @@ -193,7 +193,36 @@ def __init__( else: self._reduce_group = dist.group.WORLD + # Cache of the replicated params' local tensors for the inference fast + # path (populated lazily on first inference forward, after any post- + # construction dtype cast such as TrunkCPWrapper's .to(bf16)). + self._w_local: Optional[Tensor] = None + self._b_local: Optional[Tensor] = None + + def _inference_locals(self) -> tuple[Tensor, Optional[Tensor]]: + # Refresh if dtype changes (covers a later .to(bf16)); weights are static + # within an inference run so this caches after the first call. + if self._w_local is None or self._w_local.dtype != self.weight.dtype: + self._w_local = self.weight.to_local() + self._b_local = self.bias.to_local() if self.bias is not None else None + return self._w_local, self._b_local + def forward(self, x: DTensor) -> DTensor: - return _LinearParamsReplicatedImpl.apply( # type: ignore[return-value] - x, self.weight, self.bias, self._reduce_group, self.avg_reduce + if torch.is_grad_enabled(): + return _LinearParamsReplicatedImpl.apply( # type: ignore[return-value] + x, self.weight, self.bias, self._reduce_group, self.avg_reduce + ) + # Inference fast path: skip the autograd.Function.apply machinery + ctx + # bookkeeping and reuse cached replicated weight locals. Bit-identical to + # the Function's forward (same to_local → F.linear → from_local). + w_local, b_local = self._inference_locals() + out_local = F.linear(x.to_local(), w_local, b_local) + shape_output = x.shape[:-1] + (self.weight.shape[0],) + stride_output = update_exhaustive_strides(x.shape, x.stride(), shape_output) + return DTensor.from_local( # type: ignore[return-value] + out_local, + device_mesh=x.device_mesh, + placements=x.placements, + shape=shape_output, + stride=stride_output, ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/msa_encoder.py b/src/transformers/models/esmfold2/distributed/model/layers/msa_encoder.py new file mode 100644 index 0000000000..0cbd87cfe1 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/msa_encoder.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed MSAEncoder block / stack for ESMFold2 (inference-only). + +Mirrors the serial ``MSAEncoderBlock`` (modeling_esmfold2.py) exactly: + + pair = pair + outer_product_mean(m, msa_mask) + if not final: + m = m + msa_pair_weighted_averaging(m, pair, pair_mask) + m = m + msa_transition(m) + pair = pair + tri_mul_out(pair, mask=pair_mask) + pair = pair + tri_mul_in(pair, mask=pair_mask) + pair = pair + pair_transition(pair) + +All residuals are explicit (added here). The serial block uses +``modeling_esmfold2.PairTransition`` for both ``msa_transition`` and +``pair_transition`` — that class returns ``ffn(norm(x))`` WITHOUT a residual, +so we use the bare ``MSATransitionDistributed`` (not ``TransitionDistributed``, +which folds the residual in) and add the residual in the block. +""" + +from typing import Optional + +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor + +from transformers.models.esmfold2.distributed.comm import Ring2DComm +from transformers.models.esmfold2.distributed.manager import DistributedManager +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.outer_product_mean import ( + OuterProductMeanDistributed, +) +from transformers.models.esmfold2.distributed.model.layers.pair_averaging import ( + MSAPairWeightedAveragingDistributed, +) +from transformers.models.esmfold2.distributed.model.layers.triangular_mult import ( + TriangleMultiplicativeBlockDistributed, +) +from transformers.models.esmfold2.modeling_esmfold2 import ( + MSAEncoder as SerialMSAEncoder, +) +from transformers.models.esmfold2.modeling_esmfold2 import ( + MSAEncoderBlock as SerialMSAEncoderBlock, +) +from transformers.models.esmfold2.modeling_esmfold2 import ( + PairTransition as SerialPairTransition, +) + + +class MSATransitionDistributed(nn.Module): + """Bare LayerNorm + SwiGLU FFN (no residual) on a sharded representation. + + Matches ``modeling_esmfold2.PairTransition.forward`` which returns + ``ffn(norm(x))``; the residual is added by the calling block. + """ + + def __init__(self, layer: SerialPairTransition, device_mesh: DeviceMesh) -> None: + super().__init__() + if not isinstance(layer, SerialPairTransition): + raise TypeError( + f"layer must be PairTransition, got {type(layer).__name__}" + ) + self.norm = LayerNormParamsReplicated(layer.norm, device_mesh) + self.w12 = LinearParamsReplicated(layer.ffn.w12, device_mesh) + self.w3 = LinearParamsReplicated(layer.ffn.w3, device_mesh) + self.hidden_features = layer.ffn.hidden_features + + def forward(self, x: DTensor) -> DTensor: + normed = self.norm(x) + x12 = self.w12(normed) + x1, x2 = x12.split(self.hidden_features, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +class MSAEncoderBlockDistributed(nn.Module): + """Distributed MSAEncoderBlock. + + Parameters + ---------- + layer: + Serial MSAEncoderBlock to distribute. + dist_manager: + DistributedManager with the CP group and subgroups set up. + """ + + def __init__( + self, + layer: SerialMSAEncoderBlock, + dist_manager: DistributedManager, + comm: str = "gather", + ) -> None: + super().__init__() + if not isinstance(layer, SerialMSAEncoderBlock): + raise TypeError( + f"layer must be MSAEncoderBlock, got {type(layer).__name__}" + ) + mesh = dist_manager.device_mesh_subgroups + self.is_final_block = layer.is_final_block + + self.outer_product_mean = OuterProductMeanDistributed( + layer.outer_product_mean, dist_manager + ) + if not self.is_final_block: + self.msa_pair_weighted_averaging = MSAPairWeightedAveragingDistributed( + layer.msa_pair_weighted_averaging, dist_manager, comm=comm + ) + self.msa_transition = MSATransitionDistributed(layer.msa_transition, mesh) + + ring_comm_out = Ring2DComm( + dist_manager.group["cp"], + dist_manager.subgroups["cp"][0], + dist_manager.layout_subgroups["cp"], + ) + ring_comm_in = Ring2DComm( + dist_manager.group["cp"], + dist_manager.subgroups["cp"][0], + dist_manager.layout_subgroups["cp"], + ) + self.tri_mul_out = TriangleMultiplicativeBlockDistributed( + layer.tri_mul_out._engine, mesh, ring_comm_out + ) + self.tri_mul_in = TriangleMultiplicativeBlockDistributed( + layer.tri_mul_in._engine, mesh, ring_comm_in + ) + # Serial pair_transition is a (residual-free) PairTransition. + self.pair_transition = MSATransitionDistributed(layer.pair_transition, mesh) + + def forward( + self, + m: DTensor, + pair: DTensor, + msa_attention_mask: DTensor, + pair_attention_mask: Optional[DTensor] = None, + ) -> tuple[DTensor, DTensor]: + pair = pair + self.outer_product_mean(m, msa_attention_mask) + if not self.is_final_block: + m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask) + m = m + self.msa_transition(m) + pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask) + pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask) + pair = pair + self.pair_transition(pair) + return m, pair + + +class MSAEncoderDistributed(nn.Module): + """Distributed MSAEncoder: ModuleList of MSAEncoderBlockDistributed. + + Parameters + ---------- + encoder: + Serial MSAEncoder module. + dist_manager: + DistributedManager with the CP group and subgroups set up. + """ + + def __init__( + self, + encoder: SerialMSAEncoder, + dist_manager: DistributedManager, + comm: str = "gather", + ) -> None: + super().__init__() + if not isinstance(encoder, SerialMSAEncoder): + raise TypeError( + f"encoder must be MSAEncoder, got {type(encoder).__name__}" + ) + self.blocks = nn.ModuleList( + [ + MSAEncoderBlockDistributed(block, dist_manager, comm=comm) + for block in encoder.blocks # type: ignore[arg-type] + ] + ) + + def forward( + self, + m: DTensor, + pair: DTensor, + msa_attention_mask: DTensor, + pair_attention_mask: DTensor, + ) -> DTensor: + for block in self.blocks: + m, pair = block(m, pair, msa_attention_mask, pair_attention_mask) + return pair diff --git a/src/transformers/models/esmfold2/distributed/model/layers/outer_product_mean.py b/src/transformers/models/esmfold2/distributed/model/layers/outer_product_mean.py new file mode 100644 index 0000000000..0736f79974 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/outer_product_mean.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed OuterProductMean for ESMFold2's MSA encoder (inference-only). + +The serial op maps an MSA representation m (B, L, M, d_msa) into a pair update +(B, L, L, d_pair): + + outer[b,i,j] = sum_m a[b,i,m] (x) b[b,j,m] # einsum "bimc,bjmd->bijcd" + z = Wout(outer) / n_valid # divide_outer_before_proj=False + +Under 2D context parallelism the pair z is sharded ``(Shard(0), Shard(1), +Shard(2))`` — row token i on cp_axis_0, col token j on cp_axis_1. The MSA m is +sharded ``(Shard(0), Shard(1), Replicate())`` — token L on cp_axis_0, MSA depth +M replicated on cp_axis_1. + +Because M is replicated, the contraction over m is fully local. The tile owned +by rank (p, q) needs row block p (held locally) and col block q. Col block q is +exactly the row block held by the transpose peer (q, p), so a single transpose +of the b operand (and the mask) suffices — no ring rotation is required (unlike +boltz-cp, which shards the contracted dimension and therefore must ring-reduce). +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor, Shard + +from transformers.models.esmfold2.distributed.comm import TransposeComm +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + OuterProductMean as SerialOuterProductMean, +) + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class OuterProductMeanDistributed(nn.Module): + """Distributed (transpose-based) OuterProductMean. + + Parameters + ---------- + layer: + The serial OuterProductMean to distribute. + dist_manager: + DistributedManager with the CP group / subgroups set up. + """ + + def __init__(self, layer: SerialOuterProductMean, dist_manager) -> None: + super().__init__() + if not isinstance(layer, SerialOuterProductMean): + raise TypeError( + f"layer must be OuterProductMean, got {type(layer).__name__}" + ) + self.device_mesh = dist_manager.device_mesh_subgroups + self.d_hidden = layer.d_hidden + self.divide_outer_before_proj = layer.divide_outer_before_proj + + self.norm = LayerNormParamsReplicated(layer.norm, self.device_mesh) + self.W = LinearParamsReplicated(layer.W, self.device_mesh) + self.Wout = LinearParamsReplicated(layer.Wout, self.device_mesh) + + # Transpose (i,j) <-> (j,i) to fetch the column-token block. + self.transpose = TransposeComm( + dist_manager.group["cp"], dist_manager.layout_subgroups["cp"] + ) + + def forward(self, m: DTensor, msa_attention_mask: DTensor) -> DTensor: + mesh = self.device_mesh + dh = self.d_hidden + L = m.shape[1] + + m_norm = self.norm(m) + x = self.W(m_norm) # (B, L, M, 2*d_hidden), placements (S0, S1, R) + x = x * msa_attention_mask.unsqueeze(-1).to(x.dtype) + + x_local = x.to_local() + a_local, b_local = torch.chunk(x_local, 2, dim=-1) + a_local = a_local.contiguous() + b_local = b_local.contiguous() + mask_local = msa_attention_mask.to_local().to(a_local.dtype) # (B, sL, M) + + # Fetch the column-token block of b and the mask via a single transpose. + packed = torch.cat([b_local, mask_local.unsqueeze(-1)], dim=-1).contiguous() + recv = self.transpose.enqueue_to_dispatch(packed) + self.transpose.wait_until_finished() + b_q = recv[..., :dh].contiguous() + mask_q = recv[..., dh].contiguous() # (B, sL, M) + + outer = torch.einsum("bimc,bjmd->bijcd", a_local, b_q).flatten(-2).contiguous() + n_valid_local = ( + torch.einsum("bim,bjm->bij", mask_local, mask_q) + .unsqueeze(-1) + .clamp(min=1.0) + .contiguous() + ) + + b_size = outer.shape[0] + c_out = outer.shape[-1] + z_shape = torch.Size((b_size, L, L, c_out)) + z_dt = DTensor.from_local( + outer, + device_mesh=mesh, + placements=[Shard(0), Shard(1), Shard(2)], + shape=z_shape, + stride=_contiguous_strides(tuple(z_shape)), + ) + nv_shape = torch.Size((b_size, L, L, 1)) + n_valid_dt = DTensor.from_local( + n_valid_local, + device_mesh=mesh, + placements=[Shard(0), Shard(1), Shard(2)], + shape=nv_shape, + stride=_contiguous_strides(tuple(nv_shape)), + ) + + if self.divide_outer_before_proj: + return self.Wout(z_dt / n_valid_dt) + return self.Wout(z_dt) / n_valid_dt diff --git a/src/transformers/models/esmfold2/distributed/model/layers/pair_averaging.py b/src/transformers/models/esmfold2/distributed/model/layers/pair_averaging.py new file mode 100644 index 0000000000..9151695d68 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/pair_averaging.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed MSAPairWeightedAveraging for ESMFold2's MSA encoder (inference). + +The serial op (AF3 Algorithm 10) updates the MSA representation m using the +pair representation as attention bias: + + bias = compute_bias(pair) # (B, L, L, h) + attn = softmax_j(masked(bias)) # over col token j + out[b,i,m,h,d] = (sum_j attn[b,i,j,h] v[b,j,m,h,d]) * gate[b,i,m,h,d] + return Wout(out) + +Under 2D CP the pair (and bias) is sharded ``(Shard(0), Shard(1), Shard(2))`` +and m / v / gate are sharded ``(Shard(0), Shard(1), Replicate())`` (token L on +cp_axis_0, MSA depth M replicated). The softmax is over j = cp_axis_1, which is +sharded, and the value index is the (replicated) MSA depth with token on +cp_axis_0. + +Two communication strategies (selectable via ``comm``): + +* ``"gather"`` (default): ``DTensor.redistribute`` gathers the full j of the + bias/mask and the full token axis of v, then the attention is computed + locally over the full j in natural order — bit-exact with the serial op. The + gathered buffers are small (``L²·h`` and ``L·M·c``, far below the sharded + pair), so this is the better choice for typical grids / MSA depths. + +* ``"ring"`` (boltz-style): never materialises the full j. ``v`` is transposed + so block ``q`` aligns with ``bias[block p, block q]``, then both ring along + the column axis (``comm_row``) accumulating with an online softmax + (``tiled_softmax_attention_update``). Pays off only at larger grids (n >= 3) + with deep MSAs, where the gathered buffers would be large. The online-softmax + reassociation makes it close-but-not-bit-exact vs. the serial op. +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate, Shard + +from transformers.models.esmfold2.distributed.comm import Ring2DComm +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.distributed.utils import ( + tiled_softmax_attention_update, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + MSAPairWeightedAveraging as SerialMSAPairWeightedAveraging, +) + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class MSAPairWeightedAveragingDistributed(nn.Module): + """Distributed MSAPairWeightedAveraging (inference-only). + + Parameters + ---------- + layer: + The serial MSAPairWeightedAveraging to distribute. + dist_manager: + DistributedManager with the CP group / subgroups set up. + comm: + ``"gather"`` (default, bit-exact all-gather) or ``"ring"`` (boltz-style + online-softmax ring). + """ + + def __init__( + self, + layer: SerialMSAPairWeightedAveraging, + dist_manager, + comm: str = "gather", + ) -> None: + super().__init__() + if not isinstance(layer, SerialMSAPairWeightedAveraging): + raise TypeError( + "layer must be MSAPairWeightedAveraging, got " + f"{type(layer).__name__}" + ) + if comm not in ("gather", "ring"): + raise ValueError(f"comm must be 'gather' or 'ring', got {comm!r}") + self.device_mesh = dist_manager.device_mesh_subgroups + self.comm_mode = comm + self.n_heads = layer.n_heads + self.head_width = layer.head_width + + self.norm_single = LayerNormParamsReplicated(layer.norm_single, self.device_mesh) + # compute_bias is nn.Sequential(LayerNorm(d_pair), Linear(d_pair, n_heads)) + self.bias_norm = LayerNormParamsReplicated( + layer.compute_bias[0], self.device_mesh + ) + self.bias_lin = LinearParamsReplicated(layer.compute_bias[1], self.device_mesh) + self.Wv = LinearParamsReplicated(layer.Wv, self.device_mesh) + self.Wgate = LinearParamsReplicated(layer.Wgate, self.device_mesh) + self.Wout = LinearParamsReplicated(layer.Wout, self.device_mesh) + + if comm == "ring": + # ring_v owns the value transpose + value column-ring; ring_bias + # owns an independent bias column-ring (same schedule, separate + # comm handles so both can be in flight together). + self.ring_v = Ring2DComm( + dist_manager.group["cp"], + dist_manager.subgroups["cp"][0], + dist_manager.layout_subgroups["cp"], + ) + self.ring_bias = Ring2DComm( + dist_manager.group["cp"], + dist_manager.subgroups["cp"][0], + dist_manager.layout_subgroups["cp"], + ) + + def forward( + self, m: DTensor, pair: DTensor, pair_attention_mask: DTensor + ) -> DTensor: + if self.comm_mode == "ring": + return self._forward_ring(m, pair, pair_attention_mask) + return self._forward_gather(m, pair, pair_attention_mask) + + # -- shared projections ------------------------------------------------ + def _project(self, m: DTensor, pair: DTensor): + msa_normed = self.norm_single(m) + v_dt = self.Wv(msa_normed) # (B, L, M, h*dh), (S0, S1, R) + gate_dt = torch.sigmoid(self.Wgate(msa_normed)) # (B, L, M, h*dh) + bias_dt = self.bias_lin(self.bias_norm(pair)) # (B, L, L, h), (S0, S1, S2) + return v_dt, gate_dt, bias_dt + + def _finish(self, m: DTensor, o_local: torch.Tensor, gate_local: torch.Tensor): + h, dh = self.n_heads, self.head_width + b_size, s_l, m_depth = o_local.shape[0], o_local.shape[1], o_local.shape[2] + o_local = (o_local * gate_local).reshape(b_size, s_l, m_depth, h * dh).contiguous() + out_shape = torch.Size((b_size, m.shape[1], m_depth, h * dh)) + out_dt = DTensor.from_local( + o_local, + device_mesh=self.device_mesh, + placements=[Shard(0), Shard(1), Replicate()], + shape=out_shape, + stride=_contiguous_strides(tuple(out_shape)), + ) + return self.Wout(out_dt) + + # -- gather strategy (default, bit-exact) ------------------------------ + def _forward_gather( + self, m: DTensor, pair: DTensor, pair_attention_mask: DTensor + ) -> DTensor: + mesh = self.device_mesh + h, dh = self.n_heads, self.head_width + v_dt, gate_dt, bias_dt = self._project(m, pair) + + bias_full = ( + bias_dt.redistribute(mesh, [Shard(0), Shard(1), Replicate()]) + .to_local() + .contiguous() + ) # (B, sL, L, h) + mask_full = ( + pair_attention_mask.redistribute(mesh, [Shard(0), Shard(1), Replicate()]) + .to_local() + .contiguous() + ) # (B, sL, L) + bias_full = bias_full.masked_fill(~mask_full.unsqueeze(-1).bool(), -1e5) + attn = torch.softmax(bias_full, dim=-2) # softmax over j + + v_full = ( + v_dt.redistribute(mesh, [Shard(0), Replicate(), Replicate()]) + .to_local() + .contiguous() + ) # (B, L, M, h*dh) + + b_size, s_l = attn.shape[0], attn.shape[1] + l_full, m_depth = v_full.shape[1], v_full.shape[2] + v_full = v_full.reshape(b_size, l_full, m_depth, h, dh) + gate_local = gate_dt.to_local().reshape(b_size, s_l, m_depth, h, dh) + + o_local = torch.einsum("bijh,bjmhd->bimhd", attn, v_full) + return self._finish(m, o_local, gate_local) + + # -- ring strategy (boltz-style online softmax) ------------------------ + def _forward_ring( + self, m: DTensor, pair: DTensor, pair_attention_mask: DTensor + ) -> DTensor: + h, dh = self.n_heads, self.head_width + n = self.device_mesh.size(1) # CP axis size + v_dt, gate_dt, bias_dt = self._project(m, pair) + + # Local blocks: bias[i in p, j in q], v[block p], mask[i in p, j in q]. + bias_local = bias_dt.to_local() # (B, sLi, sLj, h) + mask_local = pair_attention_mask.to_local() # (B, sLi, sLj) + # Pre-mask: each block's mask is co-located with its bias and rings + # along with it, so masking once before the ring is correct. + bias_local = bias_local.masked_fill( + ~mask_local.unsqueeze(-1).bool(), -1e5 + ).contiguous() + v_local = v_dt.to_local().contiguous() # (B, sLp, M, h*dh) + + # Transpose v: block p -> block q (aligns v with the local bias block). + v_q = self.ring_v.comm_2d_trans.enqueue_to_dispatch(v_local) + self.ring_v.comm_2d_trans.wait_until_finished() + + b_size, s_li = bias_local.shape[0], bias_local.shape[1] + m_depth = v_q.shape[2] + bias_buf = [bias_local, torch.empty_like(bias_local)] + v_buf = [v_q.contiguous(), torch.empty_like(v_q)] + i_ready, i_recv = 0, 1 + o = lse = amax = None + + for k in range(n): + b_blk = bias_buf[i_ready] + v_blk = v_buf[i_ready] + if k < n - 1: + bias_buf[i_recv] = self.ring_bias.comm_row.enqueue_to_dispatch( + b_blk, bias_buf[i_recv] + ) + v_buf[i_recv] = self.ring_v.comm_row.enqueue_to_dispatch( + v_blk, v_buf[i_recv] + ) + + amax_blk = b_blk.amax(dim=2, keepdim=True) # (B, sLi, 1, h) + lse_blk = torch.logsumexp(b_blk - amax_blk, dim=2, keepdim=True) + p = torch.softmax(b_blk, dim=2) # (B, sLi, sLj, h) + v_blk_r = v_blk.reshape(b_size, v_blk.shape[1], m_depth, h, dh) + o_blk = torch.einsum("bijh,bjmhd->bimhd", p, v_blk_r) # (B, sLi, M, h, dh) + + # Arrange so the softmax-reduced axes (M, dh) are the trailing + # feature and the per-(i, h) lse/amax broadcast over them. + o_blk2 = o_blk.permute(0, 1, 3, 2, 4).reshape(b_size, s_li, h, m_depth * dh) + lse_blk2 = lse_blk.permute(0, 1, 3, 2).reshape(b_size, s_li, h, 1) + amax_blk2 = amax_blk.permute(0, 1, 3, 2).reshape(b_size, s_li, h, 1) + o, lse, amax = tiled_softmax_attention_update( + o_blk2, lse_blk2, amax_blk2, o, lse, amax + ) + + if k < n - 1: + self.ring_bias.comm_row.wait_until_finished() + self.ring_v.comm_row.wait_until_finished() + i_ready ^= 1 + i_recv ^= 1 + + # (B, sLi, h, M*dh) -> (B, sLi, M, h, dh) + o_local = ( + o.reshape(b_size, s_li, h, m_depth, dh) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + gate_local = gate_dt.to_local().reshape(b_size, s_li, m_depth, h, dh) + return self._finish(m, o_local, gate_local) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py b/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py index 038f1471b1..ddf5e5c746 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/pairformer.py @@ -37,28 +37,28 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor -from projects.huggingface.transformers.models.esmfold2.distributed.comm import ( +from transformers.models.esmfold2.distributed.comm import ( Ring2DComm, ) -from projects.huggingface.transformers.models.esmfold2.distributed.manager import ( +from transformers.models.esmfold2.distributed.manager import ( DistributedManager, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.layernorm import ( +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( LayerNormParamsReplicated, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.linear import ( +from transformers.models.esmfold2.distributed.model.layers.linear import ( LinearParamsReplicated, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.triangular_mult import ( +from transformers.models.esmfold2.distributed.model.layers.triangular_mult import ( TriangleMultiplicativeBlockDistributed, ) -from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( +from transformers.models.esmfold2.modeling_esmfold2_common import ( FoldingTrunk as SerialFoldingTrunk, ) -from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( +from transformers.models.esmfold2.modeling_esmfold2_common import ( PairUpdateBlock as SerialPairUpdateBlock, ) -from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( +from transformers.models.esmfold2.modeling_esmfold2_common import ( Transition as SerialTransition, ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/triangular_mult.py b/src/transformers/models/esmfold2/distributed/model/layers/triangular_mult.py index da22c40b27..4ffc7bf90e 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/triangular_mult.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/triangular_mult.py @@ -35,31 +35,43 @@ - Each step computes a local matmul; results are accumulated """ +import os from enum import Enum, auto from typing import Tuple import torch +import torch.distributed as dist from torch import nn from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Shard -from projects.huggingface.transformers.models.esmfold2.distributed.comm import ( +from transformers.models.esmfold2.distributed.comm import ( Ring2DComm, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.layernorm import ( +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( LayerNormParamsReplicated, ) -from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.linear import ( +from transformers.models.esmfold2.distributed.model.layers.linear import ( LinearParamsReplicated, ) -from projects.huggingface.transformers.models.esmfold2.distributed.utils import ( +from transformers.models.esmfold2.distributed.utils import ( update_exhaustive_strides, ) -from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( +from transformers.models.esmfold2.modeling_esmfold2_common import ( TriangleMultiplicativeBlock as SerialTriangleMultiplicativeBlock, ) +# Optional cross-rank re-alignment before the ring comm (cuda.synchronize + +# barrier on the CP group at each tri-mul). Motivated by the hypothesis that the +# CP trunk's time is NCCL P2P spin-wait from rank desync. A controlled A/B +# (CP_RANK_SYNC 0 vs 1, 5 loops) showed only a ~9% mean change, WITHIN the +# run-to-run noise (trunk ~45–66 s/call) — i.e. desync is NOT the dominant cost. +# Kept as an opt-in experiment knob, DEFAULT OFF. Correctness-neutral (touches no +# tensors). Enable with CP_RANK_SYNC=1. +_CP_RANK_SYNC = os.environ.get("CP_RANK_SYNC", "0") == "1" + + class _Direction(Enum): Outgoing = auto() Incoming = auto() @@ -142,7 +154,11 @@ def _distributed_bmm( i_ready ^= 1 i_recv ^= 1 - out = torch.zeros_like(lhs_buffer[i_ready]) + # Accumulator initialised from the first matmul rather than a pre-allocated + # zeros buffer: avoids one large per-call allocation (~370 MB bf16) and is + # more correct — the matmul output is (B,D,n,m), whereas zeros_like(lhs) is + # (B,D,n,k) and only matched when the local shard was square (n==m==k). + out: torch.Tensor | None = None comm.comm_row_init.wait_until_finished() comm.comm_col_init.wait_until_finished() @@ -157,7 +173,8 @@ def _distributed_bmm( rhs_buffer[i_recv] = comm.comm_col.enqueue_to_dispatch( rhs_ready, rhs_buffer[i_recv] ) - out = out + torch.matmul(lhs_ready, rhs_ready) + prod = torch.matmul(lhs_ready, rhs_ready) + out = prod if out is None else out + prod if k_step < comm.group_layout.shape[1] - 1: comm.comm_row.wait_until_finished() comm.comm_col.wait_until_finished() @@ -399,6 +416,15 @@ def forward(self, pair: DTensor, mask: DTensor | None = None) -> DTensor: ------- DTensor of same shape and placements as pair. """ + # Re-align ranks before the ring comm to prevent NCCL P2P spin-wait + # (see _CP_RANK_SYNC above). The cuda.synchronize() throttles the CPU so + # it can't race ahead enqueuing unbounded async P2P + large allocations + # (the dominant cause of the pile-up); the barrier then aligns ranks + # cross-process. Correctness-neutral (no tensors touched). + if _CP_RANK_SYNC and dist.is_initialized(): + torch.cuda.synchronize() + dist.barrier(group=self.ring_comm.group_2d) + # 1. Layer-normalise input normalized = self.norm_start(pair) diff --git a/src/transformers/models/esmfold2/distributed/msa_wrapper.py b/src/transformers/models/esmfold2/distributed/msa_wrapper.py new file mode 100644 index 0000000000..a9fdcba721 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/msa_wrapper.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""End-to-end CP runtime for ESMFold2's MSAEncoder. + +``MSAEncoderCPWrapper`` is a drop-in replacement for the serial ``MSAEncoder`` +(same plain-tensor in/out signature) that shards the MSA encoder's pair-space +work across the 2D CP grid, mirroring ``TrunkCPWrapper`` for the folding trunk. +Without it, the serial MSA encoder materialises the full L×L pair on every rank +(the dominant cost for large complexes) before the trunk even runs. +""" + +from math import lcm + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import Replicate, Shard, distribute_tensor + + +class MSAEncoderCPWrapper(nn.Module): + """Drop-in replacement for ``MSAEncoder`` that runs distributed. + + The small per-residue embedding (``embed`` + ``project_inputs``) runs serial + per rank (cheap), then m / pair / masks are zero-padded to a multiple of the + CP shard factor, distributed as DTensors, processed by + ``MSAEncoderDistributed``, and the gathered pair is sliced back. + """ + + def __init__( + self, serial_encoder: nn.Module, dist_manager, comm: str = "gather" + ) -> None: + super().__init__() + from transformers.models.esmfold2.distributed.manager import ( + DistributedManager, + ) + from transformers.models.esmfold2.distributed.model.layers.msa_encoder import ( + MSAEncoderDistributed, + ) + from transformers.models.esmfold2.modeling_esmfold2 import ( + MSAEncoder as SerialMSAEncoder, + ) + + if not isinstance(serial_encoder, SerialMSAEncoder): + raise TypeError( + f"expected MSAEncoder, got {type(serial_encoder).__name__}" + ) + if not isinstance(dist_manager, DistributedManager): + raise TypeError( + f"expected DistributedManager, got {type(dist_manager).__name__}" + ) + + # The distributed path ignores the serial chunk-size knob. + serial_encoder.set_chunk_size(None) + + # The embedding projections stay serial (tiny); reference them directly. + self.embed = serial_encoder.embed + self.project_inputs = serial_encoder.project_inputs + self.dist_encoder = MSAEncoderDistributed(serial_encoder, dist_manager, comm=comm) + + self.dist_manager = dist_manager + self.device_mesh = dist_manager.device_mesh_subgroups + # device_mesh is (dp, cp_axis_0, cp_axis_1) + self.cp_axis_0 = self.device_mesh.size(1) + self.cp_axis_1 = self.device_mesh.size(2) + self.shard_factor = lcm(self.cp_axis_0, self.cp_axis_1) + + def set_chunk_size(self, _chunk_size: int | None) -> None: + return + + def set_kernel_backend(self, _backend: str | None) -> None: + return + + def forward( + self, + x_pair: torch.Tensor, + x_inputs: torch.Tensor, + msa_oh: torch.Tensor, + has_deletion: torch.Tensor, + deletion_value: torch.Tensor, + msa_attention_mask: torch.Tensor, + ) -> torch.Tensor: + # Serial embedding (matches MSAEncoder.forward), per rank on full tensors. + m_feat = torch.cat( + [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 + ) + m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) + + n = x_pair.shape[1] + pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor + if pad: + m = F.pad(m, (0, 0, 0, 0, 0, pad)) # pad L (dim 1) + x_pair = F.pad(x_pair, (0, 0, 0, pad, 0, pad)) # pad both L dims + msa_attention_mask = F.pad(msa_attention_mask, (0, 0, 0, pad)) # pad L + + # Build the pair attention mask from the (padded) token mask; padded + # rows/cols are zero so they contribute nothing. + tok_mask = msa_attention_mask[:, :, 0].bool() + pair_mask = (tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)).to(m.dtype) + + mesh = self.device_mesh + m_dt = distribute_tensor( + m.contiguous(), mesh, [Shard(0), Shard(1), Replicate()] + ) + pair_dt = distribute_tensor( + x_pair.contiguous(), mesh, [Shard(0), Shard(1), Shard(2)] + ) + msa_mask_dt = distribute_tensor( + msa_attention_mask.to(m.dtype).contiguous(), + mesh, + [Shard(0), Shard(1), Replicate()], + ) + pair_mask_dt = distribute_tensor( + pair_mask.contiguous(), mesh, [Shard(0), Shard(1), Shard(2)] + ) + + out_dt = self.dist_encoder(m_dt, pair_dt, msa_mask_dt, pair_mask_dt) + out = out_dt.full_tensor() + if pad: + out = out[:, :n, :n, :] + return out + + +def wrap_model_with_cp_msa_encoder( + model: nn.Module, dist_manager, comm: str = "gather" +) -> list[str]: + """Replace every ``MSAEncoder`` submodule with ``MSAEncoderCPWrapper``. + + Walks ``model.named_modules()`` and rebinds each attribute that points at a + serial ``MSAEncoder``. Returns the list of replaced submodule paths. Models + without an MSA encoder (``model.msa_encoder is None``) are left untouched. + + ``comm`` selects the MSA pair-weighted-averaging communication strategy: + ``"gather"`` (default, bit-exact all-gather) or ``"ring"`` (boltz-style + online-softmax ring, for large grids / deep MSAs). + """ + mesh = dist_manager.device_mesh_subgroups + cp0, cp1 = mesh.size(1), mesh.size(2) + if cp0 != cp1: + raise ValueError( + f"CP grid must be square (cp_axis_0 == cp_axis_1), got {cp0}×{cp1}" + ) + + from transformers.models.esmfold2.modeling_esmfold2 import ( + MSAEncoder as SerialMSAEncoder, + ) + + targets: list[tuple[str, nn.Module, str, nn.Module]] = [] + for parent_name, parent in model.named_modules(): + for child_name, child in parent.named_children(): + if child is None: + continue + if isinstance(child, SerialMSAEncoder): + full = f"{parent_name}.{child_name}" if parent_name else child_name + targets.append((full, parent, child_name, child)) + + replaced: list[str] = [] + for full, parent, child_name, child in targets: + wrapped = MSAEncoderCPWrapper(child, dist_manager, comm=comm).to( + device=dist_manager.device, dtype=next(child.parameters()).dtype + ) + setattr(parent, child_name, wrapped) + replaced.append(full) + return replaced + + +def wrap_model_with_cp( + model: nn.Module, + dist_manager, + comm: str = "gather", + bf16: bool = True, + offload_esmc: bool = True, +) -> list[str]: + """Wrap both the folding trunk(s) and the MSA encoder for 2D CP. + + Convenience over calling ``wrap_model_with_cp_trunks`` and + ``wrap_model_with_cp_msa_encoder`` separately. Returns the combined list of + replaced submodule paths. + + ``comm`` selects the MSA pair-weighted-averaging communication strategy: + ``"gather"`` (default, bit-exact all-gather) or ``"ring"`` (boltz-style + online-softmax ring). It affects only the MSA encoder; the folding trunk + and the OuterProductMean (whose contracted MSA-depth axis is replicated) + are unaffected. + + ``bf16`` (default True): run the distributed trunk in bf16 — quality-neutral, + ~1.3-1.6x faster, lower peak VRAM (see ``wrap_model_with_cp_trunks``). + + ``offload_esmc`` (default True): offload the ESM-C LM (~12 GB) to CPU after + its one-shot use, freeing it for the trunk/diffusion. Sets + ``model._offload_esmc``. Validated with bf16: 2.14x end-to-end + 18.5 GB lower + peak vs fp32, pLDDT unchanged. Set False for high-throughput tiny-fold or + concurrent-fold use (the per-fold ~12 GB CPU<->GPU transfer isn't worth it). + """ + from transformers.models.esmfold2.distributed.utils import ( + wrap_model_with_cp_trunks, + ) + + model._offload_esmc = offload_esmc + replaced = wrap_model_with_cp_trunks(model, dist_manager, bf16=bf16) + replaced += wrap_model_with_cp_msa_encoder(model, dist_manager, comm=comm) + return replaced diff --git a/src/transformers/models/esmfold2/distributed/utils.py b/src/transformers/models/esmfold2/distributed/utils.py index a4ba6e7e60..8632b9c0d4 100644 --- a/src/transformers/models/esmfold2/distributed/utils.py +++ b/src/transformers/models/esmfold2/distributed/utils.py @@ -377,18 +377,20 @@ class TrunkCPWrapper(nn.Module): # across the CP grid; everything else stays serial per rank. """ - def __init__(self, serial_trunk: nn.Module, dist_manager) -> None: + def __init__( + self, serial_trunk: nn.Module, dist_manager, bf16: bool = True + ) -> None: super().__init__() # Lazy imports: this module is imported by manager.py and the # distributed layers, so importing pairformer / model_common at # module level would create a cycle. - from projects.huggingface.transformers.models.esmfold2.distributed.manager import ( + from transformers.models.esmfold2.distributed.manager import ( DistributedManager, ) - from projects.huggingface.transformers.models.esmfold2.distributed.model.layers.pairformer import ( + from transformers.models.esmfold2.distributed.model.layers.pairformer import ( FoldingTrunkDistributed, ) - from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( + from transformers.models.esmfold2.modeling_esmfold2_common import ( FoldingTrunk as SerialFoldingTrunk, ) @@ -409,6 +411,23 @@ def __init__(self, serial_trunk: nn.Module, dist_manager) -> None: serial_trunk.set_chunk_size(None) # type: ignore[operator] self.dist_trunk = FoldingTrunkDistributed(serial_trunk, dist_manager) + + # bf16 trunk: the serial recycle path runs the trunk with no autocast on + # an fp32 pair, and the distributed ring does plain fp32 torch.matmul + # (TF32 off) → fp32 tensor-core-less matmul + huge fp32 intermediates at + # ~92% memory occupancy dominate wall-clock. boltz-cp runs its CP trunk + # entirely in bf16 (bf16-mixed; tri-mul contraction accumulated in bf16), + # so bf16 here has direct production precedent. We cast params to bf16 + # ("bf16-true") rather than autocast, because the distributed tri-mul's + # custom_fwd disables autocast. The pair is cast to bf16 at the boundary + # in forward() and the output cast back to the caller's dtype. + # NOTE: ESMFold2's *serial* reference upcasts the tri-mul contraction to + # fp32; validated quality-neutral (pLDDT 0.643 vs 0.639). Controlled by + # the ``bf16`` arg (threaded from wrap_model_with_cp[_trunks]). + self._trunk_bf16 = bf16 + if self._trunk_bf16: + self.dist_trunk = self.dist_trunk.to(torch.bfloat16) + self.dist_manager = dist_manager self.device_mesh = dist_manager.device_mesh_subgroups # device_mesh is (dp, cp_axis_0, cp_axis_1) @@ -430,6 +449,15 @@ def forward( self, pair: torch.Tensor, pair_attention_mask: torch.Tensor | None = None ) -> torch.Tensor: N = pair.shape[1] + orig_dtype = pair.dtype + # Run the distributed trunk in bf16 (params cast in __init__). Cast the + # pair and visibility mask to bf16 so the tri-mul ops stay bf16×bf16 + # (no fp32 promotion); the gathered output is cast back to orig_dtype. + if self._trunk_bf16: + pair = pair.to(torch.bfloat16) + if pair_attention_mask is not None: + pair_attention_mask = pair_attention_mask.to(torch.bfloat16) + pad = (self.shard_factor - N % self.shard_factor) % self.shard_factor if pad: # F.pad pads from the last dim backward; pair is (B, N, N, d_pair). @@ -450,19 +478,35 @@ def forward( out_dt = self.dist_trunk(pair_dt, pair_attention_mask=mask_dt) out = out_dt.full_tensor() + if self._trunk_bf16: + out = out.to(orig_dtype) if pad: out = out[:, :N, :N, :] return out -def wrap_model_with_cp_trunks(model: nn.Module, dist_manager) -> list[str]: +def wrap_model_with_cp_trunks( + model: nn.Module, dist_manager, bf16: bool = True +) -> list[str]: """Replace every ``FoldingTrunk`` submodule with ``TrunkCPWrapper``. Walks ``model.named_modules()`` and rebinds each attribute that points at a serial ``FoldingTrunk``. Returns the list of replaced submodule paths so callers (e.g. spawned workers) can log what got wrapped. + + ``bf16`` (default True): run the distributed trunk in bf16 (params + pair cast + to bf16, output back to the caller's dtype) — ~1.3-1.6x faster and lower peak + VRAM, quality-neutral (pLDDT 0.643 vs 0.639). Pass False for an fp32 trunk + (e.g. tight bit-exact parity checks). """ - from projects.huggingface.transformers.models.esmfold2.modeling_esmfold2_common import ( + mesh = dist_manager.device_mesh_subgroups + cp0, cp1 = mesh.size(1), mesh.size(2) + if cp0 != cp1: + raise ValueError( + f"CP grid must be square (cp_axis_0 == cp_axis_1), got {cp0}×{cp1}" + ) + + from transformers.models.esmfold2.modeling_esmfold2_common import ( FoldingTrunk as SerialFoldingTrunk, ) @@ -475,7 +519,7 @@ def wrap_model_with_cp_trunks(model: nn.Module, dist_manager) -> list[str]: replaced: list[str] = [] for full, parent, child_name, child in targets: - wrapped = TrunkCPWrapper(child, dist_manager).to( + wrapped = TrunkCPWrapper(child, dist_manager, bf16=bf16).to( device=dist_manager.device, dtype=next(child.parameters()).dtype ) setattr(parent, child_name, wrapped) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 06b590e770..4dd33609a7 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -536,6 +536,14 @@ def __init__(self, config: ESMFold2Config) -> None: ) self._esmc: nn.Module | None = None self._esmc_fp8: bool = False # set by load_esmc(fp8=True) + # When True, ESM-C is offloaded to CPU after its one-shot hidden-state + # computation (restored on the next call), freeing its ~12 GB for the + # recycling trunk + diffusion — the freed blocks stay in the caching pool + # so the trunk reuses them without cudaMalloc, relieving the allocator + # pressure that drove rank desync / spin-wait. Opt-in (the CP entry point + # ``wrap_model_with_cp(offload_esmc=...)`` sets it). Validated with CP: + # 2.14x end-to-end + 18.5 GB lower peak vs fp32, pLDDT unchanged. + self._offload_esmc: bool = False pf = config.folding_trunk self.folding_trunk = FoldingTrunk( @@ -725,6 +733,10 @@ def _compute_lm_hidden_states( lm_mask_pct: float = 0.0, ) -> Tensor: assert self._esmc is not None + # Restore ESM-C to the compute device if it was offloaded to CPU after a + # previous fold (see self._offload_esmc). No-op when already resident. + if self._offload_esmc: + self._esmc.to(self.device) # fp8 TE kernels require prod(shape[:-1]) % 8 == 0. pad_to = 8 if self._esmc_fp8 else None with _lm_precision_context(self._esmc_fp8): @@ -982,6 +994,10 @@ def forward( if lm_hidden_states is not None: lm_z = self.language_model(lm_hidden_states.detach()) del lm_hidden_states + # ESM-C is done for this forward — offload to free its ~12 GB for the + # trunk/diffusion (freed blocks stay in the caching pool for reuse). + if self._offload_esmc and self._esmc is not None: + self._esmc.to("cpu") pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() From 56f5ac808ba114ee343b44b64acc9df3c5950d2d Mon Sep 17 00:00:00 2001 From: Greg Zynda Date: Mon, 29 Jun 2026 22:49:52 -0700 Subject: [PATCH 2/4] =?UTF-8?q?ESMFold2:=20distribute=20the=20full=20CP=20?= =?UTF-8?q?pipeline=20(recycle,=20diffusion,=20confidence,=20LM=E2=86=92pa?= =?UTF-8?q?ir,=20ESM-C)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../distributed/confidence_wrapper.py | 190 ++++++++++ .../esmfold2/distributed/distogram_wrapper.py | 93 +++++ .../models/esmfold2/distributed/esmc_tp.py | 129 +++++++ .../distributed/model/layers/atom_to_token.py | 192 ++++++++++ .../model/layers/attention_pair_bias.py | 114 ++++++ .../model/layers/confidence_zbase.py | 137 ++++++++ .../model/layers/diffusion_conditioning.py | 127 +++++++ .../model/layers/diffusion_transformer.py | 186 ++++++++++ .../model/layers/row_attention_pooling.py | 120 +++++++ .../model/layers/single_to_pair.py | 200 +++++++++++ .../esmfold2/distributed/msa_wrapper.py | 172 ++++++++- .../models/esmfold2/distributed/recycle.py | 327 ++++++++++++++++++ .../esmfold2/distributed/structure_wrapper.py | 254 ++++++++++++++ .../models/esmfold2/distributed/utils.py | 22 +- .../models/esmfold2/modeling_esmfold2.py | 213 ++++++++++-- 15 files changed, 2421 insertions(+), 55 deletions(-) create mode 100644 src/transformers/models/esmfold2/distributed/confidence_wrapper.py create mode 100644 src/transformers/models/esmfold2/distributed/distogram_wrapper.py create mode 100644 src/transformers/models/esmfold2/distributed/esmc_tp.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/atom_to_token.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/attention_pair_bias.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/diffusion_conditioning.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/diffusion_transformer.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/row_attention_pooling.py create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/single_to_pair.py create mode 100644 src/transformers/models/esmfold2/distributed/recycle.py create mode 100644 src/transformers/models/esmfold2/distributed/structure_wrapper.py diff --git a/src/transformers/models/esmfold2/distributed/confidence_wrapper.py b/src/transformers/models/esmfold2/distributed/confidence_wrapper.py new file mode 100644 index 0000000000..d8337a33c7 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/confidence_wrapper.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed confidence head for ESMFold2 (#6, inference-only). + +Keeps the pair sharded ``(Shard(0), Shard(1), Shard(2))`` through the expensive +front half — z_base build, the nested (already CP-wrapped) FoldingTrunk, row +pooling, and the PAE/PDE heads — so the full ``L×L×d_pair`` pair (and the trunk +activations) is never resident on any rank. Only small tensors are gathered: + + * ``single`` (B, L, d_single) after row pooling, and + * ``pae_logits`` / ``pde_logits`` (B, L, L, bins) for the output contract, + +after which the **serial** ``ConfidenceHead._finish`` runs the atom-space +pLDDT/resolved heads and the pTM/ipTM/pair_chains reductions on the gathered +tensors — one source of truth for that fiddly code, replicated and cheap (it is +per-token / single-channel-L², not the d_pair pair). + +The wrapper consumes the sharded pair DTensor directly from the recycle engine's +CP tail (no ``full_tensor()`` re-gather of ``z``), which is what removes the +confidence-phase peak. +""" + +from math import lcm + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from transformers.models.esmfold2.distributed.model.layers.confidence_zbase import ( + ConfidenceZBaseDistributed, +) +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.row_attention_pooling import ( + RowAttentionPoolingDistributed, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + gather_rep_atom_coords, +) + +_PAIR = [Shard(0), Shard(1), Shard(2)] + + +class ConfidenceHeadCPWrapper(nn.Module): + """Distributed wrapper around the serial ``ConfidenceHead``. + + The nested ``folding_trunk`` is already replaced by a ``TrunkCPWrapper`` via + ``wrap_model_with_cp_trunks`` (so ``forward_sharded`` is available). This + wrapper holds the distributed z_base / row-pool / pae-pde-head pieces and + delegates the back-half to ``head._finish``. + """ + + def __init__(self, head, dist_manager) -> None: + super().__init__() + self.head = head # serial ConfidenceHead (nested trunk already CP-wrapped) + self.device_mesh = dist_manager.device_mesh_subgroups + self.shard_factor = lcm(self.device_mesh.size(1), self.device_mesh.size(2)) + + self.zbase = ConfidenceZBaseDistributed(head, dist_manager) + self.row_pool = RowAttentionPoolingDistributed( + head.row_attention_pooling, dist_manager + ) + self.pae_ln = LayerNormParamsReplicated(head.pae_ln, self.device_mesh) + self.pae_head = LinearParamsReplicated(head.pae_head, self.device_mesh) + self.pde_ln = LayerNormParamsReplicated(head.pde_ln, self.device_mesh) + self.pde_head = LinearParamsReplicated(head.pde_head, self.device_mesh) + + def forward_sharded( + self, + z_dt: DTensor, + n_orig: int, + *, + s_inputs: torch.Tensor, + x_pred: torch.Tensor, + distogram_atom_idx: torch.Tensor, + token_attention_mask: torch.Tensor, + atom_to_token: torch.Tensor, + atom_attention_mask: torch.Tensor, + asym_id: torch.Tensor, + mol_type: torch.Tensor, + num_diffusion_samples: int = 1, + relative_position_encoding: torch.Tensor | None = None, + token_bonds_encoding: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + if num_diffusion_samples != 1: + raise NotImplementedError( + "distributed confidence head supports num_diffusion_samples==1" + ) + mesh = self.device_mesh + Lpad = z_dt.shape[1] + pad = Lpad - n_orig + pair_dtype = torch.float32 + + # --- distogram bins (full, replicated; small single-channel L²) --------- + rep_idx = distogram_atom_idx.long() + rep_coords = gather_rep_atom_coords(x_pred, rep_idx) # (B, n, 3) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + distogram_bins = ( + (rep_distances.unsqueeze(-1) > self.head.boundaries).sum(dim=-1).long() + ) + + # --- pad aux tensors to the shard factor (z_dt is already padded) ------- + def _pad_pair(t): + return F.pad(t, (0, 0, 0, pad, 0, pad)) if pad else t + + s_in_p = F.pad(s_inputs, (0, 0, 0, pad)) if pad else s_inputs + bins_p = F.pad(distogram_bins, (0, pad, 0, pad)) if pad else distogram_bins + relpos_p = ( + _pad_pair(relative_position_encoding) + if relative_position_encoding is not None + else None + ) + tokb_p = ( + _pad_pair(token_bonds_encoding) + if token_bonds_encoding is not None + else None + ) + mask_p = F.pad(token_attention_mask.float(), (0, pad)) if pad else ( + token_attention_mask.float() + ) + + # --- z_base (sharded) --------------------------------------------------- + pair_dt = self.zbase( + z_dt.to(pair_dtype), s_in_p, bins_p, relpos_p, tokb_p + ) # (S0, S1, S2) fp32, [B, Lpad, Lpad, d_pair] + + # --- nested FoldingTrunk (add-back, like serial pair.add_(trunk(pair))) - + pair_mask = (mask_p[:, :, None] * mask_p[:, None, :]).contiguous() # (B,Lpad,Lpad) + with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): + pair_mask_dt = distribute_tensor( + pair_mask.to(torch.bfloat16), mesh, _PAIR + ) + delta_dt = self.head.folding_trunk.forward_sharded( + pair_dt.to(torch.bfloat16), pair_attention_mask=pair_mask_dt + ) + pair_dt = pair_dt + delta_dt.to(pair_dtype) + + # --- row pooling -> single (gather small) ------------------------------- + single_dt = self.row_pool(pair_dt, mask_p) # (S0,S1,R) [B,Lpad,d_single] + single = single_dt.full_tensor()[:, :n_orig, :].contiguous() + + # --- PAE / PDE heads on the sharded pair; gather (sliced) logits -------- + pae_logits_dt = self.pae_head(self.pae_ln(pair_dt)) + pde_logits_dt = self.pde_head(self.pde_ln(pair_dt)) + del pair_dt + pae_logits = pae_logits_dt.full_tensor()[:, :n_orig, :n_orig, :].contiguous() + del pae_logits_dt + pde_logits = pde_logits_dt.full_tensor()[:, :n_orig, :n_orig, :].contiguous() + del pde_logits_dt + + # --- serial back-half (atom-space + pTM/ipTM), shared with forward ------ + return self.head._finish( + single=single, + pae_logits=pae_logits, + pde_logits=pde_logits, + x_pred=x_pred, + distogram_atom_idx=distogram_atom_idx, + token_attention_mask=token_attention_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atom_attention_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=num_diffusion_samples, + ) diff --git a/src/transformers/models/esmfold2/distributed/distogram_wrapper.py b/src/transformers/models/esmfold2/distributed/distogram_wrapper.py new file mode 100644 index 0000000000..5440e3eef4 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/distogram_wrapper.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed distogram head for ESMFold2 (inference-only). + +The serial op is a single channel-wise Linear over the symmetrized pair: + + distogram_logits = distogram_head(z + z.transpose(-2, -3)) # (B, L, L, bins) + +``z.transpose(-2, -3)`` swaps the two token axes (rows i <-> cols j). Under 2D CP +the pair is sharded ``(Shard(0), Shard(1), Shard(2))``; the (p, q) rank's +transposed tile ``zᵀ[i∈p, j∈q] = z[j∈q, i∈p]`` lives on the transpose peer +``(q, p)`` — one ``TransposeComm`` fetches that tile, and a local axis-swap +arranges it. The Linear is per-pair channel-wise → local. Only the small +``bins``-channel logits are gathered (vs. the full ``d_pair`` pair), so the full +``z`` (256 ch) is never re-gathered for the distogram. +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor, Shard + +from transformers.models.esmfold2.distributed.comm import TransposeComm +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) + +_PAIR = [Shard(0), Shard(1), Shard(2)] + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class DistogramHeadCPWrapper(nn.Module): + """Distributed distogram head. ``forward_sharded`` takes the sharded pair + DTensor (padded) + the original length and returns the full (sliced) + ``distogram_logits`` — gathering only the ``bins``-channel output.""" + + def __init__(self, distogram_head: nn.Linear, dist_manager) -> None: + super().__init__() + if not isinstance(distogram_head, nn.Linear): + raise TypeError(f"distogram_head must be nn.Linear, got {type(distogram_head).__name__}") + self.device_mesh = dist_manager.device_mesh_subgroups + self.head = LinearParamsReplicated(distogram_head, self.device_mesh) + self.transpose = TransposeComm( + dist_manager.group["cp"], dist_manager.layout_subgroups["cp"] + ) + + def forward_sharded(self, z_dt: DTensor, n_orig: int) -> torch.Tensor: + mesh = self.device_mesh + # Match serial: distogram runs on z.float(). + z_local = z_dt.to_local().float().contiguous() # (B, sLi, sLj, c) = z[i,j] + + recv = self.transpose.enqueue_to_dispatch(z_local) + self.transpose.wait_until_finished() + # recv = z's (q,p) tile = z[a∈q, b∈p]; zᵀ[i∈p, j∈q] = z[j,i] = recv[j, i]. + zT_local = recv.transpose(1, 2).contiguous() # (B, sLi, sLj, c) + + sym_local = z_local + zT_local + L = z_dt.shape[1] + b_size, _, _, c = sym_local.shape + full_shape = torch.Size((b_size, L, L, c)) + sym_dt = DTensor.from_local( + sym_local, + device_mesh=mesh, + placements=_PAIR, + shape=full_shape, + stride=_contiguous_strides(tuple(full_shape)), + ) + logits_dt = self.head(sym_dt) # (S0, S1, S2) [B, L, L, bins] + return logits_dt.full_tensor()[:, :n_orig, :n_orig, :].contiguous() diff --git a/src/transformers/models/esmfold2/distributed/esmc_tp.py b/src/transformers/models/esmfold2/distributed/esmc_tp.py new file mode 100644 index 0000000000..55286d98b7 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/esmc_tp.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Tensor-parallel ESM-C over the CP ranks (inference-only). + +The replicated ESM-C 6B forward is the dominant per-rank floor (T0 re-hook: +~19.8 GB live, identical on every rank). Tensor parallelism shards its weights +AND intra-block activations across the CP ranks. + +**Phase 1 (this module): the SwiGLU MLP** — ~67% of ESM-C's params. Each +``UnifiedTransformerBlock.ffn`` is a Transformer Engine ``LayerNormMLP``; TE has +native TP (``set_parallel_mode=True`` → column-parallel fc1 with the SwiGLU +gate/up split handled internally + row-parallel fc2 + an internal all-reduce), so +the wrap is: construct a TP ``LayerNormMLP`` over the CP process group and copy the +full weights in sharded. The MLP's ``ffn_hidden`` (6912) divides 4/9/16, so this +works at every CP grid (unlike head-parallel attention, blocked by n_heads=40). + +Validated bit-exact (modulo bf16 rounding) under ``torch.inference_mode`` — the +column/row shard + internal all-reduce reproduce the serial output; the ESM-C +forward already runs under the model's ``@torch.inference_mode`` so TE-TP (plain +sharded matmuls + NCCL, no autograd hooks) is inference-safe. + +Attention TP (qkv column + out_proj row + the full-width QK-LayerNorm all-reduce, +constrained to TP=4 by the 40 heads) is a later phase — see ``fix_9x_ESMC_TP.md``. +""" + +import torch +import torch.nn as nn + + +def _tp_layernorm_mlp(old, tp_group, tp_size: int, rank: int): + """Build a TE LayerNormMLP TP-shard of ``old`` (a serial te.LayerNormMLP), + holding only this rank's slice. Bit-exact (bf16) vs the full module.""" + import transformer_engine.pytorch as te + + hidden = old.layer_norm_weight.shape[0] + ffn = old.fc2_weight.shape[1] # fc2 = [hidden, ffn] + if ffn % tp_size: + raise ValueError( + f"ESM-C ffn_hidden={ffn} not divisible by tp_size={tp_size}" + ) + dtype = old.fc1_weight.dtype + device = old.fc1_weight.device + + new = te.LayerNormMLP( + hidden, + ffn, + eps=getattr(old, "eps", 1e-5), + normalization=getattr(old, "normalization", "LayerNorm"), + activation=getattr(old, "activation", "swiglu"), + bias=False, # ESM-C trained with bias=False + zero_centered_gamma=getattr(old, "zero_centered_gamma", False), + set_parallel_mode=True, + tp_group=tp_group, + tp_size=tp_size, + sequence_parallel=False, + params_dtype=dtype, + ).to(device=device, dtype=dtype) + new.eval() + + fl = ffn // tp_size + with torch.no_grad(): + new.layer_norm_weight.copy_(old.layer_norm_weight) + new.layer_norm_bias.copy_(old.layer_norm_bias) + # fc1 full = [2*ffn, hidden] laid out [gate(ffn); up(ffn)]; column-parallel + # SwiGLU shard = cat([gate[r], up[r]]) (confirmed bit-exact by the spike). + gate, up = old.fc1_weight[:ffn], old.fc1_weight[ffn:] + new.fc1_weight.copy_( + torch.cat([gate[rank * fl:(rank + 1) * fl], + up[rank * fl:(rank + 1) * fl]], dim=0) + ) + # fc2 full = [hidden, ffn]; row-parallel shards the input (ffn) dim. + new.fc2_weight.copy_(old.fc2_weight[:, rank * fl:(rank + 1) * fl]) + return new + + +def tp_shard_esmc_mlp(model: nn.Module, dist_manager, tp_group=None) -> int: + """Replace every ESM-C block's SwiGLU MLP with a TE tensor-parallel shard + over ``tp_group`` (default: the full CP process group). Returns the number of + blocks sharded. No-op (returns 0) if ESM-C isn't loaded. + + Each rank ends up holding ``1/tp_size`` of every MLP weight; the per-block + forward all-reduces internally so the block output stays replicated — the + ESM-C hidden states reach all CP ranks unchanged for the #14 lm_z builder. + """ + esmc = getattr(model, "_esmc", None) + if esmc is None: + return 0 + blocks = esmc.transformer.blocks + + if tp_group is None: + tp_group = dist_manager.group["cp"] + tp_size = torch.distributed.get_world_size(tp_group) + rank = torch.distributed.get_rank(tp_group) + if tp_size == 1: + return 0 + + n = 0 + for block in blocks: + ffn = getattr(block, "ffn", None) + # Only the TE LayerNormMLP path is supported (the accelerated build); + # the pure-PyTorch fallback would need a DTensor TP and isn't on this path. + if ffn is None or type(ffn).__name__ != "LayerNormMLP": + raise TypeError( + f"expected block.ffn to be a TE LayerNormMLP, got " + f"{type(ffn).__name__ if ffn is not None else None}" + ) + block.ffn = _tp_layernorm_mlp(ffn, tp_group, tp_size, rank) + del ffn + n += 1 + return n diff --git a/src/transformers/models/esmfold2/distributed/model/layers/atom_to_token.py b/src/transformers/models/esmfold2/distributed/model/layers/atom_to_token.py new file mode 100644 index 0000000000..3e003c75f2 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/atom_to_token.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed atom<->token gather/scatter for ESMFold2's diffusion module. + +ESMFold2 maps atoms to tokens with an **index** (``atom_to_token: [B, A]`` int64), +not boltz-cp's one-hot ``[B, A, n_tokens]`` matrix, so these are index-based +gather/scatter (``torch.gather`` / ``scatter_reduce``) rather than a one-hot +matmul — strictly less memory (no ``A×L`` materialisation). + +Sharding convention (matches boltz-cp's 1D sequence reprs): a sequence tensor +(atoms or tokens) on the ``(dp, cp_axis_0, cp_axis_1)`` mesh has placements +``(Shard(0), Shard(1), Replicate())`` — the sequence axis (dim 1) is split across +``cp_axis_0`` and replicated across ``cp_axis_1``. (The pair stays 2-D sharded; +the diffusion's heavy L×L work is handled separately by the ring attention.) + +Inference-only (no autograd) — ESMFold2's CP path runs under +``@torch.inference_mode``. + +These are the atom<->token data-movement primitives. The atom encoder/decoder's +sliding-window attention (``swa_window_size``) imposes a separate constraint when +sharding the atom axis — atom shards must align to window boundaries (shard at a +multiple of the window) or exchange halos — handled where the atom transformer is +distributed, not here. +""" + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor, Replicate, Shard + +from transformers.models.esmfold2.modeling_esmfold2_common import ( + gather_token_to_atom, +) + +_SEQ_PL = [Shard(0), Shard(1), Replicate()] # (dp, cp_axis_0=seq, cp_axis_1) +_SEQ_GATHERED_PL = [Shard(0), Replicate(), Replicate()] # seq gathered on cp_axis_0 + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +def dist_gather_token_to_atom( + token_dt: DTensor, atom_to_token_idx_dt: DTensor +) -> DTensor: + """Broadcast per-token features to per-atom features (sharded). + + Parameters + ---------- + token_dt: + Token features ``(B, L, d)`` with placements ``(Shard(0), Shard(1), + Replicate())`` (token axis split on ``cp_axis_0``). + atom_to_token_idx_dt: + Per-atom global token index ``(B, A)`` int64, same placements (atom axis + split on ``cp_axis_0``). + + Returns + ------- + Atom features ``(B, A, d)`` with placements ``(Shard(0), Shard(1), + Replicate())`` (atom axis split on ``cp_axis_0``). + """ + mesh = token_dt.device_mesh + # Gather the full token axis onto every cp_axis_0 rank (token repr is the + # cheap 1-D L×d, not L×L), so each rank can index its local atoms. + token_full = token_dt.redistribute(mesh, _SEQ_GATHERED_PL).to_local().contiguous() + idx_local = atom_to_token_idx_dt.to_local() + atom_local = gather_token_to_atom(token_full, idx_local) + + b = atom_local.shape[0] + a_global = atom_to_token_idx_dt.shape[1] + d = atom_local.shape[-1] + shape = torch.Size((b, a_global, d)) + return DTensor.from_local( + atom_local.contiguous(), + device_mesh=mesh, + placements=_SEQ_PL, + shape=shape, + stride=_contiguous_strides(tuple(shape)), + ) + + +def _local_scatter_sum_count( + atom_local: torch.Tensor, + idx_local: torch.Tensor, + n_tokens: int, + mask_local: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Local (per-rank) scatter-add of atom features into token bins, plus a + per-token contributing-atom count. Masked atoms are routed to a throwaway + bin ``n_tokens`` and dropped.""" + b, _, d = atom_local.shape + idx_use = idx_local + n_out = n_tokens + if mask_local is not None: + idx_use = torch.where(mask_local.bool(), idx_local, n_tokens) + n_out = n_tokens + 1 + + idx_e = idx_use.unsqueeze(-1).expand(b, idx_use.shape[1], d) + s = torch.zeros(b, n_out, d, device=atom_local.device, dtype=atom_local.dtype) + s.scatter_add_(1, idx_e, atom_local) + c = torch.zeros(b, n_out, 1, device=atom_local.device, dtype=atom_local.dtype) + c.scatter_add_( + 1, + idx_use.unsqueeze(-1), + torch.ones(b, idx_use.shape[1], 1, device=atom_local.device, dtype=atom_local.dtype), + ) + return s[:, :n_tokens], c[:, :n_tokens] + + +def dist_scatter_atom_to_token( + atom_dt: DTensor, + atom_to_token_idx_dt: DTensor, + n_tokens: int, + atom_mask_dt: DTensor | None = None, +) -> DTensor: + """Aggregate per-atom features to per-token features (mean), sharded. + + Bit-equivalent to the serial ``scatter_atom_to_token`` (mean over the atoms + of each token, empty tokens → 0): each rank scatter-adds its local atom shard + into a full-L sum + count, the two are all-reduced across ``cp_axis_0``, then + divided to a mean and re-sharded onto the token axis. + + Parameters + ---------- + atom_dt: + Atom features ``(B, A, d)`` with placements ``(Shard(0), Shard(1), + Replicate())`` (atom axis split on ``cp_axis_0``). + atom_to_token_idx_dt: + Per-atom global token index ``(B, A)`` int64, same placements. + n_tokens: + Global token count ``L``. + atom_mask_dt: + Optional per-atom bool mask ``(B, A)``, same placements. + + Returns + ------- + Token features ``(B, L, d)`` with placements ``(Shard(0), Shard(1), + Replicate())`` (token axis split on ``cp_axis_0``). + """ + mesh = atom_dt.device_mesh + atom_local = atom_dt.to_local() + idx_local = atom_to_token_idx_dt.to_local() + mask_local = atom_mask_dt.to_local() if atom_mask_dt is not None else None + + s_local, c_local = _local_scatter_sum_count( + atom_local, idx_local, n_tokens, mask_local + ) + + # Sum the per-rank partials across the cp_axis_0 sequence-shard group. The + # mesh-dim-1 process group is exactly that axis; cp_axis_1 holds replicas, so + # each cp_axis_1 column independently reconstructs the same full token tensor. + cp0_group = mesh.get_group(1) + if dist.is_initialized() and dist.get_world_size(cp0_group) > 1: + dist.all_reduce(s_local, op=dist.ReduceOp.SUM, group=cp0_group) + dist.all_reduce(c_local, op=dist.ReduceOp.SUM, group=cp0_group) + + mean_full = s_local / c_local.clamp(min=1.0) # empty tokens -> 0 + + # mean_full is the full (B, L, d) tensor, identical on every cp_axis_0 rank in + # a column. Wrap as gathered-on-cp_axis_0, then reshard to the token axis + # (Replicate -> Shard is a local chunk, no further communication). + b, d = mean_full.shape[0], mean_full.shape[-1] + shape = torch.Size((b, n_tokens, d)) + full_dt = DTensor.from_local( + mean_full.contiguous(), + device_mesh=mesh, + placements=_SEQ_GATHERED_PL, + shape=shape, + stride=_contiguous_strides(tuple(shape)), + ) + return full_dt.redistribute(mesh, _SEQ_PL) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/attention_pair_bias.py b/src/transformers/models/esmfold2/distributed/model/layers/attention_pair_bias.py new file mode 100644 index 0000000000..f64d1c0cd8 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/attention_pair_bias.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed token self-attention with a 2-D-sharded pair bias (diffusion). + +This is the *core attention mechanism* for distributing ESMFold2's diffusion +token transformer (``AttentionPairBias``). It shards the expensive ``L×L`` +attention over the query-row (``cp_axis_0``) axis, which is the diffusion's +memory ceiling; the cheap 1-D token reprs and the atom encoder/decoder stay +replicated (see ``fix_9x_#2.md``). + +Two strategies, mirroring the MSA pair-averaging design: + +* ``"gather"`` (implemented here): query rows are sharded on ``cp_axis_0``; + ``k``/``v`` and the bias *row* are gathered along the key/column axis, then a + single local softmax attention runs over the full key axis — **bit-exact** + with the serial op. Per-rank attention/bias is ``L²·H / P`` (sharded on + ``cp_axis_0``), down from the full ``L²``. The gathered ``k``/``v`` are the + cheap 1-D ``L·H·D``. +* ``"ring"`` (future): also shard the key axis on ``cp_axis_1`` and ring k/v/z + with online softmax (``AttentionPairBiasComm`` + ``tiled_softmax_attention_update``), + removing the full-key-axis gather. A memory optimisation over ``"gather"``. + +Inference-only. The ``q``/``k``/``v`` token reprs use placements +``(Shard(0), Shard(1), Replicate())`` (rows on ``cp_axis_0``); the pair bias uses +``(Shard(0), Shard(1), Shard(2))`` (the trunk's 2-D pair sharding). +""" + +import torch +from torch.distributed.tensor import DTensor, Replicate, Shard + +_ROW_PL = [Shard(0), Shard(1), Replicate()] # token rows on cp_axis_0 +_ROW_GATHERED_PL = [Shard(0), Replicate(), Replicate()] # key axis gathered +_PAIR_PL = [Shard(0), Shard(1), Shard(2)] # 2-D pair sharding +_PAIR_ROWS_PL = [Shard(0), Shard(1), Replicate()] # bias rows on cp0, cols gathered + + +def attention_pair_bias_gather( + q_dt: DTensor, + k_dt: DTensor, + v_dt: DTensor, + bias_dt: DTensor, + scale: float, + key_mask_dt: DTensor | None = None, +) -> DTensor: + """Self-attention with an additive per-head pair bias, "gather" strategy. + + Bit-exact with the serial reference:: + + logits[b,i,j,h] = scale * (q[b,i,h,:] . k[b,j,h,:]) + bias[b,i,j,h] + attn = softmax_j(logits) # over keys j + o[b,i,h,:] = sum_j attn[b,i,j,h] v[b,j,h,:] + + Parameters + ---------- + q_dt, k_dt, v_dt: + Token reprs ``(B, L, H, D)`` with placements ``(Shard(0), Shard(1), + Replicate())`` — query/key rows sharded on ``cp_axis_0``. + bias_dt: + Per-head pair bias ``(B, L, L, H)`` with placements ``(Shard(0), + Shard(1), Shard(2))`` — the trunk's 2-D pair sharding. + scale: + Attention logit scale (``head_dim**-0.5``). + key_mask_dt: + Optional key mask ``(B, L)`` (``True`` = keep), placements ``(Shard(0), + Shard(1), Replicate())``. Masked keys get ``-inf`` logit. + + Returns + ------- + Output ``(B, L, H, D)`` with placements ``(Shard(0), Shard(1), Replicate())`` + (query rows on ``cp_axis_0``). + """ + mesh = q_dt.device_mesh + + # Query rows stay sharded on cp_axis_0; gather the key axis (cheap 1-D k/v) + # and the bias *row* (rows block stays on cp_axis_0, columns gathered). + q_local = q_dt.to_local() # (B, Lq_local, H, D) + k_full = k_dt.redistribute(mesh, _ROW_GATHERED_PL).to_local() # (B, L, H, D) + v_full = v_dt.redistribute(mesh, _ROW_GATHERED_PL).to_local() # (B, L, H, D) + bias_row = bias_dt.redistribute(mesh, _PAIR_ROWS_PL).to_local() # (B, Lq_local, L, H) + + # logits (B, Lq_local, L, H) + logits = torch.einsum("bihd,bjhd->bijh", q_local, k_full) * scale + logits = logits + bias_row.to(logits.dtype) + + if key_mask_dt is not None: + key_mask = key_mask_dt.redistribute(mesh, _ROW_GATHERED_PL).to_local() # (B, L) + neg = torch.finfo(logits.dtype).min + logits = logits + torch.where( + key_mask.bool()[:, None, :, None], 0.0, neg + ).to(logits.dtype) + + attn = torch.softmax(logits, dim=2).to(v_full.dtype) # over keys j + o_local = torch.einsum("bijh,bjhd->bihd", attn, v_full) # (B, Lq_local, H, D) + + return DTensor.from_local(o_local.contiguous(), device_mesh=mesh, placements=_ROW_PL) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py b/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py new file mode 100644 index 0000000000..443767b6ec --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed z_base / s->z builder for the ESMFold2 confidence head (#6). + +Reproduces the serial pre-trunk pair construction +(``modeling_esmfold2.py:ConfidenceHead.forward`` lines 187-216) keeping the pair a +sharded ``(Shard(0), Shard(1), Shard(2))`` DTensor: + + z_base = z_norm(z) [+ rel_pos] [+ token_bonds] + + s_to_z(s)[:, :, None] # row i, broadcast over j + + s_to_z_transpose(s)[:, None, :] # col j, broadcast over i + + s_to_z_prod_out(prod_in1(s)[:, :, None] * prod_in2(s)[:, None, :]) + pair = z_base + dist_bin_pairwise_embed(distogram_bins) + +Because ``s_inputs`` (hence ``s = s_inputs_norm(s_inputs)`` and all the per-token +projections) is replicated, the row term needs only the local cp_axis_0 rows and +the column term only the local cp_axis_1 cols — both obtained by slicing the cheap +full ``[B, L, d_pair]`` projection two ways (``distribute_tensor`` of a replicated +tensor = a local slice, no communication). So the whole build is **local and +bit-exact** (no transpose, no reduction): every term lands on the same +``(sLi, sLj)`` tile the pair owns. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) + +_PAIR = [Shard(0), Shard(1), Shard(2)] +_ROW = [Shard(0), Shard(1), Replicate()] # token L sharded on cp_axis_0 (rows i) +_COL = [Shard(0), Replicate(), Shard(1)] # token L sharded on cp_axis_1 (cols j) + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class ConfidenceZBaseDistributed(nn.Module): + """Builds the confidence head's pre-trunk pair as a sharded DTensor. + + Reads the projection submodules off the serial ``ConfidenceHead``; only the + pair-channel ``z_norm`` is wrapped (it runs on the sharded ``z``). The + per-token projections run on the replicated ``s_inputs`` exactly as serial. + """ + + def __init__(self, layer, dist_manager) -> None: + super().__init__() + self.device_mesh = dist_manager.device_mesh_subgroups + self.z_norm = LayerNormParamsReplicated(layer.z_norm, self.device_mesh) + # Per-token / pair-channel ops: replicated params, run as-is. + self.s_inputs_norm = layer.s_inputs_norm + self.s_to_z = layer.s_to_z + self.s_to_z_transpose = layer.s_to_z_transpose + self.s_to_z_prod_in1 = layer.s_to_z_prod_in1 + self.s_to_z_prod_in2 = layer.s_to_z_prod_in2 + self.s_to_z_prod_out = layer.s_to_z_prod_out + self.dist_bin_pairwise_embed = layer.dist_bin_pairwise_embed + + def forward( + self, + z: DTensor, + s_inputs: torch.Tensor, + distogram_bins: torch.Tensor, + relative_position_encoding: torch.Tensor | None = None, + token_bonds_encoding: torch.Tensor | None = None, + ) -> DTensor: + mesh = self.device_mesh + L = z.shape[1] + + s = self.s_inputs_norm(s_inputs) # full [B, L, d_inputs], replicated + + def row(t): # local cp_axis_0 rows + return distribute_tensor(t.contiguous(), mesh, _ROW).to_local() + + def col(t): # local cp_axis_1 cols + return distribute_tensor(t.contiguous(), mesh, _COL).to_local() + + a_row = row(self.s_to_z(s)) # (B, sLi, d_pair) + b_col = col(self.s_to_z_transpose(s)) # (B, sLj, d_pair) + p1_row = row(self.s_to_z_prod_in1(s)) # (B, sLi, d_pair) + p2_col = col(self.s_to_z_prod_in2(s)) # (B, sLj, d_pair) + + zb = self.z_norm(z).to_local() # (B, sLi, sLj, d_pair) + if relative_position_encoding is not None: + zb = zb + distribute_tensor( + relative_position_encoding.contiguous(), mesh, _PAIR + ).to_local() + if token_bonds_encoding is not None: + zb = zb + distribute_tensor( + token_bonds_encoding.contiguous(), mesh, _PAIR + ).to_local() + + zb = zb + a_row[:, :, None, :] + b_col[:, None, :, :] + prod = p1_row[:, :, None, :] * p2_col[:, None, :, :] # (B, sLi, sLj, d_pair) + zb = zb + F.linear(prod, self.s_to_z_prod_out.weight) + + bins_local = distribute_tensor( + distogram_bins.contiguous(), mesh, _PAIR + ).to_local() # (B, sLi, sLj) + zb = zb + self.dist_bin_pairwise_embed(bins_local) + + b_size = zb.shape[0] + d_pair = zb.shape[-1] + full_shape = torch.Size((b_size, L, L, d_pair)) + return DTensor.from_local( + zb.contiguous(), + device_mesh=mesh, + placements=_PAIR, + shape=full_shape, + stride=_contiguous_strides(tuple(full_shape)), + ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/diffusion_conditioning.py b/src/transformers/models/esmfold2/distributed/model/layers/diffusion_conditioning.py new file mode 100644 index 0000000000..9aca3c730d --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/diffusion_conditioning.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed DiffusionConditioning for ESMFold2's structure head. + +The conditioning produces, from the trunk pair ``z_trunk``, the per-head bias +source ``z`` (``B, L, L, c_z``) that biases the diffusion token attention — the +**dominant** diffusion tensor (``L²·c_z`` ≈ 4.6 GB at L=3000, far above the +per-layer attention scores). To get any diffusion memory win this ``z`` must be +sharded, so this module produces it as a 2-D-sharded DTensor +``(Shard(0), Shard(1), Shard(2))`` rather than the full tensor. + +The ``z`` path (``z_input_norm`` → ``z_proj`` → ``z_transitions``) is entirely +pair-channel-wise (``TransitionLayer`` is LayerNorm + SwiGLU over the last dim), +so it runs **as-is on the local pair shard** using the serial submodule's +identically-replicated params. The ``s`` path (token single repr, cheap 1-D) runs +**replicated/full** exactly as serial. + +Caller owns padding: ``z_trunk`` / ``rel_pos`` must already be padded so the token +axis is a multiple of the CP shard factor. Inference-only. +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor, Shard + +from transformers.models.esmfold2.modeling_esmfold2_common import ( + DiffusionConditioning as SerialDiffusionConditioning, +) + +_PAIR_PL = [Shard(0), Shard(1), Shard(2)] + + +class DiffusionConditioningDistributed(nn.Module): + """Distributed ``DiffusionConditioning``: full ``s``, 2-D-sharded ``z``. + + Thin shell over the serial module — reuses its submodules, running the ``z`` + path on the local pair shard and the ``s`` path on the full tensor. Bit-exact + with serial (the ``z`` path is pointwise in the pair indices). + """ + + def __init__(self, layer: SerialDiffusionConditioning, device_mesh) -> None: + super().__init__() + if not isinstance(layer, SerialDiffusionConditioning): + raise TypeError( + f"layer must be DiffusionConditioning, got {type(layer).__name__}" + ) + self.layer = layer + self.device_mesh = device_mesh + + def forward( + self, + t_hat: torch.Tensor, + s_inputs: torch.Tensor, + z_trunk_dt: DTensor, + rel_pos_dt: DTensor, + sigma_data: float | None = None, + num_diffusion_samples: int = 1, + inference_cache: dict | None = None, + ) -> tuple[torch.Tensor, DTensor]: + """Mirrors ``DiffusionConditioning.forward`` (z cached across rollout). + + ``z_trunk_dt`` / ``rel_pos_dt`` are 2-D-sharded DTensors (already padded); + returns ``(s_full, z_sharded_dt)``. + """ + layer = self.layer + mesh = self.device_mesh + sigma = layer.sigma_data if sigma_data is None else float(sigma_data) + # base (unexpanded) batch from s_inputs — z_trunk_dt is None on the cached + # path, and s_inputs always carries the base batch. + base_batch = s_inputs.shape[0] + target_batch = base_batch * num_diffusion_samples + + # --- z path (cached), on the local pair shard --- + if inference_cache is not None and "z_cp" in inference_cache: + z_dt = inference_cache["z_cp"] + else: + zt_local = z_trunk_dt.to_local().to(torch.float32) + zr_local = rel_pos_dt.to_local().to(torch.float32) + z_local = torch.cat([zt_local, zr_local], dim=-1) + z_local = layer.z_proj(layer.z_input_norm(z_local)) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for block in layer.z_transitions: + z_local = z_local + block(z_local) + z_dt = DTensor.from_local( + z_local.contiguous(), device_mesh=mesh, placements=_PAIR_PL + ) + if inference_cache is not None: + inference_cache["z_cp"] = z_dt + + # --- s path (full / replicated), identical to serial --- + s_inputs_eff = s_inputs + if s_inputs_eff.shape[0] != target_batch: + s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0) + s = layer.s_proj(layer.s_input_norm(s_inputs_eff.to(torch.float32))) + + t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1) + if t.numel() == 1: + t = t.expand(target_batch) + elif t.shape[0] != target_batch: + t = t.repeat_interleave(num_diffusion_samples, 0) + t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20)) + n = layer.fourier(t_noise) + n = layer.noise_proj(layer.noise_norm(n)) + s = s + n.unsqueeze(1) + for block in layer.s_transitions: + s = s + block(s) + + return s, z_dt diff --git a/src/transformers/models/esmfold2/distributed/model/layers/diffusion_transformer.py b/src/transformers/models/esmfold2/distributed/model/layers/diffusion_transformer.py new file mode 100644 index 0000000000..a2bcf37d04 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/diffusion_transformer.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed diffusion token transformer for ESMFold2. + +Distributes ``DiffusionTransformer`` (``AttentionPairBias`` + +``ConditionedTransitionBlock`` per block) by sharding the token row axis on +``cp_axis_0`` (the diffusion memory ceiling is the L×L token attention). + +Design: every per-token operation — AdaLN, the q/k/v/gate/out projections, the +output gate, and the whole transition block — is *channel-wise* on the +row-sharded token repr, so it runs **as-is on the local shard** using the serial +submodule's (identically replicated) parameters. Only the attention score/context +(``q·k`` over the full key axis + the 2-D-sharded pair bias) needs cross-rank +communication, handled by :func:`attention_pair_bias_gather`. This keeps the code +a thin shell over the serial layer and bit-exact with it (modulo the gather, +which is lossless). + +Token reprs (``a``, ``s``) are DTensors ``(B, L, d)`` with placements +``(Shard(0), Shard(1), Replicate())`` (rows on ``cp_axis_0``); the pair ``z`` is +``(B, L, L, d_pair)`` with ``(Shard(0), Shard(1), Shard(2))``. Inference-only, +single diffusion sample (``num_diffusion_samples`` handled by the caller/sampler). +""" + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate, Shard + +from transformers.models.esmfold2.distributed.model.layers.attention_pair_bias import ( + attention_pair_bias_gather, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + AttentionPairBias as SerialAttentionPairBias, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + ConditionedTransitionBlock as SerialConditionedTransitionBlock, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + DiffusionTransformer as SerialDiffusionTransformer, +) + +_ROW_PL = [Shard(0), Shard(1), Replicate()] +_PAIR_PL = [Shard(0), Shard(1), Shard(2)] + + +def _row_dtensor(local: torch.Tensor, mesh) -> DTensor: + return DTensor.from_local(local.contiguous(), device_mesh=mesh, placements=_ROW_PL) + + +class AttentionPairBiasDistributed(nn.Module): + """Distributed ``AttentionPairBias`` (standard, non-fused path). + + Holds the serial layer; runs its per-token ops on local shards and routes the + attention through :func:`attention_pair_bias_gather`. + """ + + def __init__(self, layer: SerialAttentionPairBias) -> None: + super().__init__() + if not isinstance(layer, SerialAttentionPairBias): + raise TypeError( + f"layer must be AttentionPairBias, got {type(layer).__name__}" + ) + self.layer = layer + self.num_heads = layer.num_heads + self.head_dim = layer.head_dim + self.scale = layer.scale + self.use_conditioning = hasattr(layer, "adaln") + self.has_pair = hasattr(layer, "pair_bias_proj") + + def forward( + self, + a_dt: DTensor, + s_dt: DTensor | None, + z_dt: DTensor, + key_mask_dt: DTensor | None = None, + ) -> DTensor: + mesh = a_dt.device_mesh + layer = self.layer + h, hd = self.num_heads, self.head_dim + + a_local = a_dt.to_local() + b, l_loc = a_local.shape[0], a_local.shape[1] + + # --- per-token ops on the local shard (channel-wise) --- + if s_dt is not None and self.use_conditioning: + s_local = s_dt.to_local() + x_local = layer.adaln(a_local, s_local) + else: + s_local = None + x_local = layer.pre_norm(a_local) + + q_local = layer.q_proj(x_local).view(b, l_loc, h, hd) + k_local, v_local = layer.kv_proj(x_local).chunk(2, dim=-1) + k_local = k_local.reshape(b, l_loc, h, hd) + v_local = v_local.reshape(b, l_loc, h, hd) + g_local = torch.sigmoid(layer.g_proj(x_local)).view(b, l_loc, h, hd) + + # --- pair bias on the local (2-D-sharded) pair shard --- + z_local = z_dt.to_local() # (B, rows_loc, cols_loc, d_pair) + bias_local = layer.pair_bias_proj(layer.pair_norm(z_local)) # (.., H) + bias_dt = DTensor.from_local( + bias_local.contiguous(), device_mesh=mesh, placements=_PAIR_PL + ) + + # --- distributed attention (gather strategy) --- + q_dt = _row_dtensor(q_local, mesh) + k_dt = _row_dtensor(k_local, mesh) + v_dt = _row_dtensor(v_local, mesh) + o_dt = attention_pair_bias_gather( + q_dt, k_dt, v_dt, bias_dt, self.scale, key_mask_dt=key_mask_dt + ) + + # --- gate + output projection on the local shard --- + ctx_local = (g_local * o_dt.to_local()).reshape(b, l_loc, h * hd) + out_local = layer.out_proj(ctx_local) + if s_local is not None and self.use_conditioning: + out_local = torch.sigmoid(layer.out_gate(s_local)) * out_local + return _row_dtensor(out_local, mesh) + + +class ConditionedTransitionDistributed(nn.Module): + """Distributed ``ConditionedTransitionBlock`` — purely per-token, so just the + serial block run on the local shard.""" + + def __init__(self, layer: SerialConditionedTransitionBlock) -> None: + super().__init__() + if not isinstance(layer, SerialConditionedTransitionBlock): + raise TypeError( + f"layer must be ConditionedTransitionBlock, got {type(layer).__name__}" + ) + self.layer = layer + + def forward(self, a_dt: DTensor, s_dt: DTensor | None) -> DTensor: + mesh = a_dt.device_mesh + a_local = a_dt.to_local() + s_local = s_dt.to_local() if s_dt is not None else None + return _row_dtensor(self.layer(a_local, s_local), mesh) + + +class DiffusionTransformerDistributed(nn.Module): + """Distributed ``DiffusionTransformer``: per block ``x = x + attn(x,s,z); x = + x + transition(x,s)``, matching the serial loop.""" + + def __init__(self, transformer: SerialDiffusionTransformer) -> None: + super().__init__() + if not isinstance(transformer, SerialDiffusionTransformer): + raise TypeError( + f"transformer must be DiffusionTransformer, got {type(transformer).__name__}" + ) + self.attn_blocks = nn.ModuleList( + [AttentionPairBiasDistributed(b) for b in transformer.attn_blocks] + ) + self.transition_blocks = nn.ModuleList( + [ConditionedTransitionDistributed(b) for b in transformer.transition_blocks] + ) + + def forward( + self, + a_dt: DTensor, + s_dt: DTensor | None, + z_dt: DTensor, + key_mask_dt: DTensor | None = None, + ) -> DTensor: + x = a_dt + for attn, transition in zip(self.attn_blocks, self.transition_blocks): + x = x + attn(x, s_dt, z_dt, key_mask_dt=key_mask_dt) + x = x + transition(x, s_dt) + return x diff --git a/src/transformers/models/esmfold2/distributed/model/layers/row_attention_pooling.py b/src/transformers/models/esmfold2/distributed/model/layers/row_attention_pooling.py new file mode 100644 index 0000000000..1b7e3e5934 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/row_attention_pooling.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed RowAttentionPooling for the ESMFold2 confidence head (#6, inference). + +The serial op (``modeling_esmfold2_common.py:RowAttentionPooling``) pools the pair +``z`` (B, L, L, d_pair) over the column axis j into a single repr (B, L, d_single): + + scores = attn_proj(z).squeeze(-1) # (B, L, L) + scores = scores + mask_bias_over_j # -1e9 where col j is padding + weights = softmax(scores, dim=-1) # over j + pooled = einsum("bnm,bnmd->bnd", weights, z) # sum over j + return out_proj(pooled) # (B, L, d_single) + +Under 2D CP the pair is sharded ``(Shard(0), Shard(1), Shard(2))`` (row i on +cp_axis_0, col j on cp_axis_1). The softmax and the weighted sum are both over +j = cp_axis_1, so we do a **distributed softmax + weighted-sum reduction along the +column group** — never gathering ``z``'s columns: + + * attn_proj is channel-wise → local scores (B, sLi, sLj); + * softmax over the full j: all-reduce MAX then SUM over cp_axis_1 (online stats); + * pooled: local ``einsum`` over the local j, then all-reduce SUM over cp_axis_1. + +The result is row-sharded on cp_axis_0, replicated on cp_axis_1 — returned as a +``(Shard(0), Shard(1), Replicate())`` single-repr DTensor. +""" + +import torch +import torch.nn as nn +import torch.distributed as dist +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + RowAttentionPooling as SerialRowAttentionPooling, +) + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class RowAttentionPoolingDistributed(nn.Module): + """Distributed RowAttentionPooling (inference-only). + + ``forward`` takes the pair DTensor ``z`` ``(Shard(0), Shard(1), Shard(2))`` and + the full token mask ``(B, L)``; returns the single-repr DTensor + ``(Shard(0), Shard(1), Replicate())`` of shape ``(B, L, d_single)``. + """ + + def __init__(self, layer: SerialRowAttentionPooling, dist_manager) -> None: + super().__init__() + if not isinstance(layer, SerialRowAttentionPooling): + raise TypeError( + f"layer must be RowAttentionPooling, got {type(layer).__name__}" + ) + self.device_mesh = dist_manager.device_mesh_subgroups + self.attn_proj = LinearParamsReplicated(layer.attn_proj, self.device_mesh) + self.out_proj = LinearParamsReplicated(layer.out_proj, self.device_mesh) + # device_mesh dims: (dp, cp_axis_0, cp_axis_1); j is cp_axis_1 (col group). + self.col_group = self.device_mesh.get_group(2) + + def forward(self, z: DTensor, mask: torch.Tensor) -> DTensor: + mesh = self.device_mesh + # Column-shard the token mask to match z's local columns j (cp_axis_1). + mask_cols = distribute_tensor( + mask.contiguous(), mesh, [Shard(0), Replicate(), Shard(1)] + ).to_local() # (B, sLj) + + scores = self.attn_proj(z).to_local().squeeze(-1) # (B, sLi, sLj) + neg = torch.full_like(scores, -1e9) + scores = torch.where(mask_cols[:, None, :].bool(), scores, neg) + + # Distributed softmax over the full j (cp_axis_1). + local_max = scores.amax(dim=-1, keepdim=True) # (B, sLi, 1) + dist.all_reduce(local_max, op=dist.ReduceOp.MAX, group=self.col_group) + e = torch.exp(scores - local_max) + local_sum = e.sum(dim=-1, keepdim=True) # (B, sLi, 1) + dist.all_reduce(local_sum, op=dist.ReduceOp.SUM, group=self.col_group) + weights = e / local_sum # (B, sLi, sLj) — globally normalized over j + + # Weighted sum over the local j, then reduce across the column group. + z_local = z.to_local() # (B, sLi, sLj, d_pair) + pooled = torch.einsum("bij,bijd->bid", weights, z_local).contiguous() + dist.all_reduce(pooled, op=dist.ReduceOp.SUM, group=self.col_group) + # pooled is now the full-j pool: row-sharded on cp_axis_0, identical across + # cp_axis_1 -> wrap as (Shard(0), Shard(1), Replicate()). + b_size, s_li, d_pair = pooled.shape + full_shape = torch.Size((b_size, z.shape[1], d_pair)) + pooled_dt = DTensor.from_local( + pooled, + device_mesh=mesh, + placements=[Shard(0), Shard(1), Replicate()], + shape=full_shape, + stride=_contiguous_strides(tuple(full_shape)), + ) + return self.out_proj(pooled_dt) # (B, L, d_single), (S0, S1, R) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/single_to_pair.py b/src/transformers/models/esmfold2/distributed/model/layers/single_to_pair.py new file mode 100644 index 0000000000..992d19a449 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/single_to_pair.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed LM->pair builder for ESMFold2 (Fix #14, inference-only). + +The serial ``LanguageModelShim`` turns the ESM-C hidden states ``[B, L, 81, +d_model]`` into the pair representation ``lm_z`` ``[B, L, L, d_z]``. The L×L +blow-up happens entirely inside ``SingleToPair``: + + x = downproject(x) # [B, L, dp] + feat = cat([x_i * x_j, x_i - x_j], dim=-1) # [B, L, L, 2*dp] <-- the floor + lm_z = output_mlp(feat) # [B, L, L, d_z] + +Materialising that full ``[B, L, L, *]`` on every rank is the binding per-rank +peak measured by the T0 phase sweep (the ``language_model`` phase, identical at +2×2 and 4×4 -> unsharded). This module builds ``lm_z`` directly as a 2-D-sharded +``(Shard(0), Shard(1), Shard(2))`` DTensor (row token i on cp_axis_0, col token j +on cp_axis_1), so the L×L tensor is never resident full on any rank. + +The outer op ``f(x_i, x_j)`` mirrors ``OuterProductMeanDistributed``: the row +block ``x_i`` is held locally; the column block ``x_j`` is fetched with a single +``TransposeComm`` (the (q,p) peer's row block is exactly rank (p,q)'s column +block on a square grid). All channel-wise ops (downproject, output_mlp, the +trailing LayerNorm) run on the local shard via replicated params — only the +outer op needs communication. +""" + +from math import lcm + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +from transformers.models.esmfold2.distributed.comm import TransposeComm +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.distributed.model.layers.linear import ( + LinearParamsReplicated, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + LanguageModelShim as SerialLanguageModelShim, + SingleToPair as SerialSingleToPair, +) + +_PAIR_PLACEMENTS = [Shard(0), Shard(1), Shard(2)] +_TOKEN_PLACEMENTS = [Shard(0), Shard(1), Replicate()] + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +class SingleToPairDistributed(nn.Module): + """Distributed (transpose-based) ``SingleToPair``. + + ``forward`` takes a per-token DTensor ``x`` with placements + ``(Shard(0), Shard(1), Replicate())`` (token L sharded on cp_axis_0, + replicated on cp_axis_1) and returns the pair DTensor with placements + ``(Shard(0), Shard(1), Shard(2))``. + """ + + def __init__(self, layer: SerialSingleToPair, dist_manager) -> None: + super().__init__() + if not isinstance(layer, SerialSingleToPair): + raise TypeError( + f"layer must be SingleToPair, got {type(layer).__name__}" + ) + self.device_mesh = dist_manager.device_mesh_subgroups + + self.downproject = LinearParamsReplicated(layer.downproject, self.device_mesh) + # output_mlp = Sequential(Linear, GELU, Linear) + self.out_in = LinearParamsReplicated(layer.output_mlp[0], self.device_mesh) + self.gelu = layer.output_mlp[1] + self.out_proj = LinearParamsReplicated(layer.output_mlp[2], self.device_mesh) + + # (i,j) <-> (j,i) transpose to fetch the column-token block. + self.transpose = TransposeComm( + dist_manager.group["cp"], dist_manager.layout_subgroups["cp"] + ) + + def forward(self, x: DTensor) -> DTensor: + mesh = self.device_mesh + L = x.shape[1] + + x = self.downproject(x) # (B, L, dp), placements (S0, S1, R) + + # Row block i is local; fetch column block j from the transpose peer. + a_local = x.to_local().contiguous() # (B, sL_i, dp) + b_q = self.transpose.enqueue_to_dispatch(a_local) + self.transpose.wait_until_finished() + b_q = b_q.contiguous() # (B, sL_j, dp) + + # cat([x_i * x_j, x_i - x_j]) on the local (sL_i, sL_j) tile — mirrors the + # serial cat([x.unsqueeze(2) * x.unsqueeze(1), x.unsqueeze(2) - x.unsqueeze(1)]). + ai = a_local.unsqueeze(2) # (B, sL_i, 1, dp) + bj = b_q.unsqueeze(1) # (B, 1, sL_j, dp) + feat_local = torch.cat([ai * bj, ai - bj], dim=-1).contiguous() + + b_size = feat_local.shape[0] + c_feat = feat_local.shape[-1] + feat_shape = torch.Size((b_size, L, L, c_feat)) + feat_dt = DTensor.from_local( + feat_local, + device_mesh=mesh, + placements=_PAIR_PLACEMENTS, + shape=feat_shape, + stride=_contiguous_strides(tuple(feat_shape)), + ) + + out = self.out_in(feat_dt) # (B, L, L, out_dim), (S0, S1, S2) + # GELU is pointwise → run on the local shard, preserve the DTensor metadata. + gelu_local = self.gelu(out.to_local()) + out = DTensor.from_local( + gelu_local.contiguous(), + device_mesh=mesh, + placements=_PAIR_PLACEMENTS, + shape=out.shape, + stride=out.stride(), + ) + return self.out_proj(out) + + +class LanguageModelShimDistributed(nn.Module): + """Distributed ``LanguageModelShim`` producing a sharded ``lm_z`` DTensor. + + The per-token front end (``base_z_linear`` + the ``base_z_combine`` softmax + mix) is cheap (per-token, ~MB) and runs replicated on each rank exactly as in + the serial shim. Only the L×L ``SingleToPair`` blow-up and the trailing + LayerNorm are distributed. + + ``forward`` returns ``(lm_z_dt, n_orig)``: a padded, 2-D-sharded pair DTensor + and the original (pre-pad) token count. The token axis is padded to + ``shard_factor = lcm(cp_axis_0, cp_axis_1)`` so the pair shards evenly — the + same padding contract the recycle engine uses for ``z`` / ``z_init``. Padded + rows/cols carry no signal (masked downstream by the padded pair mask; sliced + to ``n_orig`` after any gather), so they need not be explicitly zeroed. + """ + + def __init__(self, layer: SerialLanguageModelShim, dist_manager) -> None: + super().__init__() + if not isinstance(layer, SerialLanguageModelShim): + raise TypeError( + f"layer must be LanguageModelShim, got {type(layer).__name__}" + ) + self.device_mesh = dist_manager.device_mesh_subgroups + # device_mesh is (dp, cp_axis_0, cp_axis_1); pad to lcm so the pair shards evenly. + self.shard_factor = lcm( + self.device_mesh.size(1), self.device_mesh.size(2) + ) + + # Per-token front end runs replicated (kept as the serial submodules). + self.base_z_linear = layer.base_z_linear + self.base_z_combine = layer.base_z_combine + + # base_z_mlp = Sequential(SingleToPair, LayerNorm(d_z)) + self.single_to_pair = SingleToPairDistributed( + layer.base_z_mlp[0], dist_manager + ) + self.norm = LayerNormParamsReplicated(layer.base_z_mlp[1], self.device_mesh) + + def forward(self, hidden_states: torch.Tensor) -> tuple[DTensor, int]: + # Per-token front end (replicated, full L — cheap, no L×L here). + lm_single = self.base_z_linear(hidden_states) # [B, L, 81, d_z] + weights = self.base_z_combine.softmax(0) # [81] + lm_single = (weights @ lm_single).squeeze(-2) # [B, L, d_z] + + n_orig = lm_single.shape[1] + pad = (self.shard_factor - n_orig % self.shard_factor) % self.shard_factor + if pad: + lm_single = F.pad(lm_single, (0, 0, 0, pad)) + + x_dt = distribute_tensor( + lm_single.contiguous(), self.device_mesh, _TOKEN_PLACEMENTS + ) + lm_z = self.single_to_pair(x_dt) # (S0, S1, S2), [B, Lpad, Lpad, d_z] + lm_z = self.norm(lm_z) + return lm_z, n_orig diff --git a/src/transformers/models/esmfold2/distributed/msa_wrapper.py b/src/transformers/models/esmfold2/distributed/msa_wrapper.py index a9fdcba721..c686332d5b 100644 --- a/src/transformers/models/esmfold2/distributed/msa_wrapper.py +++ b/src/transformers/models/esmfold2/distributed/msa_wrapper.py @@ -33,7 +33,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import Replicate, Shard, distribute_tensor +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor class MSAEncoderCPWrapper(nn.Module): @@ -46,7 +46,11 @@ class MSAEncoderCPWrapper(nn.Module): """ def __init__( - self, serial_encoder: nn.Module, dist_manager, comm: str = "gather" + self, + serial_encoder: nn.Module, + dist_manager, + comm: str = "gather", + bf16: bool = True, ) -> None: super().__init__() from transformers.models.esmfold2.distributed.manager import ( @@ -76,6 +80,18 @@ def __init__( self.project_inputs = serial_encoder.project_inputs self.dist_encoder = MSAEncoderDistributed(serial_encoder, dist_manager, comm=comm) + # bf16 MSA encoder: the recycle pair is fp32 and the MSA encoder otherwise + # runs fp32 (its distributed tri-mul's custom_fwd disables autocast), which + # made it the largest recycle-loop peak contributor (~16 GB @ L1422). Cast + # the heavy dist_encoder to bf16 (mirrors TrunkCPWrapper's bf16=True, which + # was validated quality-neutral); the pair is cast to bf16 at the boundary + # in forward[_sharded] and the output cast back to the caller's dtype. The + # tiny embed/project_inputs are left fp32 (shared with the serial model; + # m is cast to bf16 after embedding). + self._bf16 = bf16 + if self._bf16: + self.dist_encoder = self.dist_encoder.to(torch.bfloat16) + self.dist_manager = dist_manager self.device_mesh = dist_manager.device_mesh_subgroups # device_mesh is (dp, cp_axis_0, cp_axis_1) @@ -89,26 +105,45 @@ def set_chunk_size(self, _chunk_size: int | None) -> None: def set_kernel_backend(self, _backend: str | None) -> None: return - def forward( + def forward_sharded( self, - x_pair: torch.Tensor, + x_pair_dt: DTensor, x_inputs: torch.Tensor, msa_oh: torch.Tensor, has_deletion: torch.Tensor, deletion_value: torch.Tensor, msa_attention_mask: torch.Tensor, - ) -> torch.Tensor: + ) -> DTensor: + """Run the distributed MSA encoder on an already-sharded pair DTensor. + + ``x_pair_dt`` must already be padded to a multiple of + ``self.shard_factor`` and distributed ``[Shard(0), Shard(1), Shard(2)]`` + on ``self.device_mesh``. The MSA inputs (``m`` and the masks) are still + embedded serially per rank (cheap), padded to match the pair's padded + length, and distributed here. Returns the pair-space output as a DTensor + — **no** ``full_tensor()`` gather. + + This is the entry point a CP orchestrator uses to keep the pair sharded + across the MSA→trunk boundary; the plain-tensor :meth:`forward` is a thin + adapter around it. + """ # Serial embedding (matches MSAEncoder.forward), per rank on full tensors. m_feat = torch.cat( [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 ) m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) - n = x_pair.shape[1] - pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor + # The incoming pair is already padded; pad m / masks to match it so the + # token (L) axis lines up shard-for-shard. + n = m.shape[1] + n_padded = x_pair_dt.shape[1] + pad = n_padded - n + if pad < 0: + raise ValueError( + f"sharded pair length {n_padded} < MSA token length {n}" + ) if pad: m = F.pad(m, (0, 0, 0, 0, 0, pad)) # pad L (dim 1) - x_pair = F.pad(x_pair, (0, 0, 0, pad, 0, pad)) # pad both L dims msa_attention_mask = F.pad(msa_attention_mask, (0, 0, 0, pad)) # pad L # Build the pair attention mask from the (padded) token mask; padded @@ -116,13 +151,19 @@ def forward( tok_mask = msa_attention_mask[:, :, 0].bool() pair_mask = (tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)).to(m.dtype) + # Cast the heavy inputs to bf16 so the (bf16) dist_encoder runs bf16 end to + # end; restore the caller's dtype on the output. + orig_dtype = x_pair_dt.dtype + if self._bf16: + m = m.to(torch.bfloat16) + x_pair_dt = x_pair_dt.to(torch.bfloat16) + msa_attention_mask = msa_attention_mask.to(torch.bfloat16) + pair_mask = pair_mask.to(torch.bfloat16) + mesh = self.device_mesh m_dt = distribute_tensor( m.contiguous(), mesh, [Shard(0), Shard(1), Replicate()] ) - pair_dt = distribute_tensor( - x_pair.contiguous(), mesh, [Shard(0), Shard(1), Shard(2)] - ) msa_mask_dt = distribute_tensor( msa_attention_mask.to(m.dtype).contiguous(), mesh, @@ -132,7 +173,37 @@ def forward( pair_mask.contiguous(), mesh, [Shard(0), Shard(1), Shard(2)] ) - out_dt = self.dist_encoder(m_dt, pair_dt, msa_mask_dt, pair_mask_dt) + out_dt = self.dist_encoder(m_dt, x_pair_dt, msa_mask_dt, pair_mask_dt) + if self._bf16: + out_dt = out_dt.to(orig_dtype) + return out_dt + + def forward( + self, + x_pair: torch.Tensor, + x_inputs: torch.Tensor, + msa_oh: torch.Tensor, + has_deletion: torch.Tensor, + deletion_value: torch.Tensor, + msa_attention_mask: torch.Tensor, + ) -> torch.Tensor: + n = x_pair.shape[1] + pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor + if pad: + x_pair = F.pad(x_pair, (0, 0, 0, pad, 0, pad)) # pad both L dims + + pair_dt = distribute_tensor( + x_pair.contiguous(), self.device_mesh, [Shard(0), Shard(1), Shard(2)] + ) + + out_dt = self.forward_sharded( + pair_dt, + x_inputs, + msa_oh, + has_deletion, + deletion_value, + msa_attention_mask, + ) out = out_dt.full_tensor() if pad: out = out[:, :n, :n, :] @@ -140,7 +211,7 @@ def forward( def wrap_model_with_cp_msa_encoder( - model: nn.Module, dist_manager, comm: str = "gather" + model: nn.Module, dist_manager, comm: str = "gather", bf16: bool = True ) -> list[str]: """Replace every ``MSAEncoder`` submodule with ``MSAEncoderCPWrapper``. @@ -174,7 +245,7 @@ def wrap_model_with_cp_msa_encoder( replaced: list[str] = [] for full, parent, child_name, child in targets: - wrapped = MSAEncoderCPWrapper(child, dist_manager, comm=comm).to( + wrapped = MSAEncoderCPWrapper(child, dist_manager, comm=comm, bf16=bf16).to( device=dist_manager.device, dtype=next(child.parameters()).dtype ) setattr(parent, child_name, wrapped) @@ -188,6 +259,8 @@ def wrap_model_with_cp( comm: str = "gather", bf16: bool = True, offload_esmc: bool = True, + wrap_structure: bool = True, + tp_esmc: bool = False, ) -> list[str]: """Wrap both the folding trunk(s) and the MSA encoder for 2D CP. @@ -216,5 +289,74 @@ def wrap_model_with_cp( model._offload_esmc = offload_esmc replaced = wrap_model_with_cp_trunks(model, dist_manager, bf16=bf16) - replaced += wrap_model_with_cp_msa_encoder(model, dist_manager, comm=comm) + replaced += wrap_model_with_cp_msa_encoder( + model, dist_manager, comm=comm, bf16=bf16 + ) + + # Install the CP recycle orchestrator so the pair stays sharded across the + # recycle loop (no per-iteration full_tensor() round-trips). The model's + # _run_one_loop delegates to this when present (CP-agnostic seam). Only the + # main model owns the loop (parcae_input_norm + _run_one_loop); guard so + # wrapping a bare submodule is a no-op. + if hasattr(model, "_run_one_loop") and hasattr(model, "parcae_input_norm"): + from transformers.models.esmfold2.distributed.recycle import ( + CPRecycleEngine, + ) + + model._cp_recycle_engine = CPRecycleEngine(model, dist_manager) + + # Install the distributed LM->pair builder (#14) so lm_z is produced as a + # sharded DTensor instead of a full L×L tensor on every rank (the binding + # per-rank peak per the T0 phase sweep). The model's forward delegates via + # the _cp_language_model seam when present. + lm = getattr(model, "language_model", None) + if lm is not None: + from transformers.models.esmfold2.distributed.model.layers.single_to_pair import ( + LanguageModelShimDistributed, + ) + + model._cp_language_model = LanguageModelShimDistributed(lm, dist_manager) + + # Distributed confidence head (#6): consumes the sharded pair from the CP + # tail (no full-L×L z re-gather). Its nested FoldingTrunk is already + # CP-wrapped by wrap_model_with_cp_trunks above. + conf = getattr(model, "confidence_head", None) + if conf is not None: + from transformers.models.esmfold2.distributed.confidence_wrapper import ( + ConfidenceHeadCPWrapper, + ) + + model._cp_confidence_head = ConfidenceHeadCPWrapper(conf, dist_manager) + + # Distributed distogram head: distogram_head(z + zᵀ) off the sharded pair + # (one transpose + local Linear), gathering only the bins-channel logits. + disto = getattr(model, "distogram_head", None) + if disto is not None: + from transformers.models.esmfold2.distributed.distogram_wrapper import ( + DistogramHeadCPWrapper, + ) + + model._cp_distogram_head = DistogramHeadCPWrapper(disto, dist_manager) + + # Distribute the diffusion structure head (conditioned-z + token attention + # sharded; atom encoder/decoder replicated). The serial sample() loop keeps + # driving the (now wrapped) diffusion_module. + if wrap_structure: + from transformers.models.esmfold2.distributed.structure_wrapper import ( + wrap_model_with_cp_structure_head, + ) + + replaced += wrap_model_with_cp_structure_head(model, dist_manager) + + # Tensor-parallel ESM-C (the dominant replicated floor). Phase 1: shard the + # SwiGLU MLP (~67% of ESM-C weights) over the CP ranks via TE-native TP. + if tp_esmc: + from transformers.models.esmfold2.distributed.esmc_tp import ( + tp_shard_esmc_mlp, + ) + + n_tp = tp_shard_esmc_mlp(model, dist_manager) + if n_tp: + replaced.append(f"_esmc.transformer.blocks[*].ffn (TP x{n_tp})") + return replaced diff --git a/src/transformers/models/esmfold2/distributed/recycle.py b/src/transformers/models/esmfold2/distributed/recycle.py new file mode 100644 index 0000000000..54f82ff4a8 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/recycle.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Context-parallel orchestrator for ESMFold2's recycle loop. + +``CPRecycleEngine`` owns the sharded version of ``ESMFold2Model._run_one_loop``. +The model file stays CP-agnostic: ``_run_one_loop`` simply delegates to an +installed engine (see the ``_cp_recycle_engine`` seam there), which is set by +``wrap_model_with_cp``. + +Why an orchestrator: the serial loop materialises the full ``L×L`` pair on every +rank at *each* module boundary (MSA encoder, LM encoder, trunk) via the wrappers' +``forward()`` (``distribute_tensor`` → … → ``full_tensor``). Keeping the running +pair sharded across the whole loop — calling each wrapper's ``forward_sharded`` +and running the parcae combine on local shards — removes those per-iteration +gathers. The pair is gathered exactly once, when the loop returns (until the +diffusion / confidence heads are also distributed, at which point even that +final gather can go). + +All per-iteration tensors stay DTensors with placements +``(Shard(0), Shard(1), Shard(2))``; the channel-dim parcae combine is computed on +``to_local()`` shards (channel is replicated, so it needs no communication) and +re-wrapped — mirroring ``LinearParamsReplicated.forward``. Validated bit-exact +vs. the serial path on a 2×2 grid. +""" + +from math import lcm + +import torch +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from transformers.models.esmfold2.distributed.model.layers.layernorm import ( + LayerNormParamsReplicated, +) +from transformers.models.esmfold2.modeling_esmfold2_common import ( + NUM_RES_TYPES, + maybe_subsample_msa, +) + +_PLACEMENTS = [Shard(0), Shard(1), Shard(2)] + + +class CPRecycleEngine: + """Sharded drop-in for ``ESMFold2Model._run_one_loop``. + + Parameters + ---------- + model: + The ESMFold2 model being wrapped. Used to read ``parcae_input_norm`` (to + build a replicated-param copy) at construction; the (already CP-wrapped) + ``folding_trunk`` / ``msa_encoder`` / ``lm_encoder`` and config are read + per call so the engine always sees the current submodules. + dist_manager: + The DistributedManager providing the CP device mesh. + """ + + def __init__(self, model, dist_manager) -> None: + self.dist_manager = dist_manager + self.device_mesh = dist_manager.device_mesh_subgroups + # device_mesh is (dp, cp_axis_0, cp_axis_1) + self.cp_axis_0 = self.device_mesh.size(1) + self.cp_axis_1 = self.device_mesh.size(2) + self.shard_factor = lcm(self.cp_axis_0, self.cp_axis_1) + + # Replicated-param copy of the channel-dim LayerNorm: a plain nn.LayerNorm + # cannot run on a DTensor (mixed plain-tensor/DTensor params raise), so the + # parcae input norm is wrapped to hold replicated DTensor params. + self.parcae_input_norm = LayerNormParamsReplicated( + model.parcae_input_norm, self.device_mesh + ) + + def parcae_finish(self, model, z_dt: DTensor, pair_mask: torch.Tensor) -> DTensor: + """Sharded post-recycle parcae tail, keeping the pair sharded. + + Mirrors the serial ``z = parcae_readout(z); z = parcae_coda(z, mask)``: + ``parcae_readout`` is a channel-wise Linear run on the local shard + (weight replicated, no token mixing); ``parcae_coda`` is a CP-wrapped + ``FoldingTrunk`` driven via ``forward_sharded`` (no gather/re-distribute). + ``z_dt`` is the 2-D-sharded (padded) pair from ``run_loop(gather=False)``; + ``pair_mask`` is the full (unpadded) token pair mask. Returns sharded z.""" + mesh = self.device_mesh + n = pair_mask.shape[-1] + pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor + + # parcae_readout: channel Linear on the local shard. + z_local = z_dt.to_local() + z_local = F.linear(z_local, model.parcae_readout.weight) + z_dt = DTensor.from_local(z_local.contiguous(), mesh, _PLACEMENTS) + + # parcae_coda: CP-wrapped trunk → forward_sharded (needs a padded, sharded mask). + pm = F.pad(pair_mask, (0, pad, 0, pad)) if pad else pair_mask + pair_mask_dt = distribute_tensor( + pm.to(z_dt.dtype).contiguous(), mesh, _PLACEMENTS + ) + return model.parcae_coda.forward_sharded( + z_dt, pair_attention_mask=pair_mask_dt + ) + + def _prepare_msa_iter(self, model, msa_inputs, tok_mask): + """Per-iteration MSA feature prep — a faithful copy of the serial block + in ``ESMFold2Model._run_one_loop`` (kept here so the engine doesn't + mutate the serial inference path). Returns plain tensors; sharding + happens inside ``MSAEncoderCPWrapper.forward_sharded``.""" + msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa( + msa_inputs["msa"], + msa_inputs["msa_attention_mask"], + msa_inputs["has_deletion"], + msa_inputs["deletion_value"], + max_depth=msa_inputs["max_depth"], + enabled=msa_inputs["subsample_enabled"], + ) + B_msa, M, L_msa = msa_i.shape + msa_oh = F.one_hot( + msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES + ).float() + msa_attn = ( + mask_i.permute(0, 2, 1).float() + if mask_i is not None + else tok_mask[:, :, None].expand(-1, -1, M).float() + ) + # Bias-free MSAEncoder.embed requires zeroed padding. + msa_oh = msa_oh * msa_attn.unsqueeze(-1) + hd = ( + hd_i.permute(0, 2, 1).float() + if hd_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + dv = ( + dv_i.permute(0, 2, 1).float() + if dv_i is not None + else torch.zeros(B_msa, L_msa, M, device=msa_i.device) + ) + return msa_oh, msa_attn, hd, dv, msa_inputs["x_inputs"] + + def run_loop( + self, + model, + *, + z: torch.Tensor, + z_init: torch.Tensor, + lm_z: torch.Tensor | None, + msa_inputs: dict | None, + pair_mask: torch.Tensor, + a: torch.Tensor, + b_mat: torch.Tensor, + tok_mask: torch.Tensor, + total_steps: int, + gather: bool = True, + ): + """Sharded equivalent of ``_run_one_loop``. + + With ``gather=True`` (default) returns the gathered full tensor (drop-in + for the serial loop). With ``gather=False`` returns + ``(z_dt, n_orig)`` — the **sharded** ``(Shard0, Shard1, Shard2)`` pair + DTensor (still padded to the shard factor) and the original token length + ``n_orig`` — so a CP tail can keep the pair sharded through + parcae / distogram / structure / confidence instead of re-gathering. The + padded rows/cols carry no signal (masked); the caller slices to + ``n_orig`` after any later gather.""" + lm_cfg = model.config.lm_encoder + lm_dropout_p = getattr(lm_cfg, "lm_dropout", 0.0) + # Per-loop LM dropout: the serial path applies a fresh F.dropout over the + # *full* lm_z each iteration (training=True even under eval). We reproduce + # it the same way — dropout on the full (unpadded) lm_z, then pad + + # distribute — so with the run's synchronized RNG (all ranks share the + # seed and draw in the same order) every rank gets the identical mask, and + # it matches the serial / gather-fallback path. Drawing the mask on a + # per-rank shard instead would desync the ranks. + per_loop_lm_dropout = ( + lm_z is not None + and getattr(lm_cfg, "per_loop_lm_dropout", False) + and lm_dropout_p > 0.0 + ) + + mesh = self.device_mesh + loop_dtype = z.dtype + N = z.shape[1] + pad = (self.shard_factor - N % self.shard_factor) % self.shard_factor + + def _pad_pair(t: torch.Tensor) -> torch.Tensor: # (B, L, L, C) + return F.pad(t, (0, 0, 0, pad, 0, pad)) if pad else t + + def _pad_mask(t: torch.Tensor) -> torch.Tensor: # (B, L, L) + return F.pad(t, (0, pad, 0, pad)) if pad else t + + def _distribute_lm(lm_full: torch.Tensor) -> DTensor: + return distribute_tensor( + _pad_pair(lm_full.to(z_init.dtype)).contiguous(), mesh, _PLACEMENTS + ) + + # Distribute the persistent loop tensors ONCE (vs. per-boundary in serial). + z_dt = distribute_tensor(_pad_pair(z).contiguous(), mesh, _PLACEMENTS) + z_init_dt = distribute_tensor(_pad_pair(z_init).contiguous(), mesh, _PLACEMENTS) + pair_mask_dt = distribute_tensor( + _pad_mask(pair_mask).to(loop_dtype).contiguous(), mesh, _PLACEMENTS + ) + + # lm_z may arrive as a sharded (Shard0,Shard1,Shard2) DTensor (#14 — + # produced full-L×L-free by LanguageModelShimDistributed) or as a full + # tensor (serial-produced, e.g. the unit tests). Either way the loop needs a + # per-iteration sharded DTensor. + lm_is_dt = isinstance(lm_z, DTensor) + # Per-rank generator for SHARD-LOCAL dropout (sharded lm_z only): each rank + # drops its own tile with a distinct, fresh-per-loop mask. This is NOT + # bit-exact with the serial full-L×L dropout mask (torch's sequential philox + # can't be cheaply tiled) — a valid-but-different stochastic sample, like the + # ring MSA path; validated by end-to-end pLDDT/pTM parity. Crucially it uses + # a *separate* generator, leaving the global RNG untouched so the MSA + # subsample below stays synchronized ACROSS ranks (all ranks pick the same + # subset — required for the distributed MSA encoder to be consistent). + lm_local_static: torch.Tensor | None = None + lm_shape = lm_stride = None + lm_drop_gen: torch.Generator | None = None + lm_z_dt_static: DTensor | None = None + if lm_is_dt: + lm_cast = lm_z.to(z_init.dtype) + if per_loop_lm_dropout: + lm_local_static = lm_cast.to_local() + lm_shape, lm_stride = lm_cast.shape, lm_cast.stride() + lm_drop_gen = torch.Generator(device=lm_local_static.device) + lm_drop_gen.manual_seed(0x5F14 + torch.distributed.get_rank()) + else: + lm_z_dt_static = lm_cast + elif lm_z is not None and not per_loop_lm_dropout: + lm_z_dt_static = _distribute_lm(lm_z) + + overwrite = model.config.msa_encoder_overwrite + + for _ in range(total_steps): + # (1) LM dropout FIRST, matching the serial loop's RNG-draw order + # (dropout before MSA subsample) so the synchronized RNG stays in + # lockstep with the serial / gather-fallback path. + if lm_z is None: + lm_z_i = None + elif not per_loop_lm_dropout: + lm_z_i = lm_z_dt_static + elif lm_is_dt: + # Shard-local dropout on the local tile (see generator note above): + # keep with prob (1-p), scale by 1/(1-p) — matches F.dropout. + keep = 1.0 - lm_dropout_p + mask = ( + torch.rand( + lm_local_static.shape, + generator=lm_drop_gen, + device=lm_local_static.device, + dtype=lm_local_static.dtype, + ) + >= lm_dropout_p + ) + drop_local = lm_local_static * mask / keep + lm_z_i = DTensor.from_local( + drop_local.contiguous(), + device_mesh=mesh, + placements=_PLACEMENTS, + shape=lm_shape, + stride=lm_stride, + ) + else: + lm_i_full = F.dropout(lm_z, p=lm_dropout_p, training=True) + lm_z_i = _distribute_lm(lm_i_full) + + refined_lm_z: DTensor | None = None + if lm_z_i is not None and model.lm_encoder is not None: + refined_lm_z = model.lm_encoder.forward_sharded( + lm_z_i, pair_attention_mask=pair_mask_dt + ) + + z_inject = z_init_dt + if lm_z_i is not None and model.lm_encoder is None: + z_inject = z_inject + lm_z_i.to(z_inject.dtype) + + # (2) MSA prep (RNG draw) AFTER LM dropout, matching serial order. + if model.msa_encoder is not None and msa_inputs is not None: + msa_oh, msa_attn, hd, dv, x_inputs = self._prepare_msa_iter( + model, msa_inputs, tok_mask + ) + msa_pair = model.msa_encoder.forward_sharded( + z_inject, x_inputs, msa_oh, hd, dv, msa_attn + ).to(z_inject.dtype) + z_inject = msa_pair if overwrite else (z_inject + msa_pair) + + if refined_lm_z is not None: + z_inject = z_inject + refined_lm_z.to(z_inject.dtype) + + # parcae combine (channel-dim): norm via replicated params, then the + # affine recurrence on local shards (channel is replicated, so a*z and + # F.linear are purely local — raw F.linear on a DTensor would try to + # flatten the sharded token axis and fail). + inj_dt = self.parcae_input_norm(z_inject) + z_local = z_dt.to_local() + out_local = a * z_local + F.linear( + inj_dt.to_local().to(z_local.dtype), b_mat + ) + z_dt = DTensor.from_local(out_local, mesh, _PLACEMENTS) + + z_dt = model.folding_trunk.forward_sharded( + z_dt, pair_attention_mask=pair_mask_dt + ) + + if not gather: + # Keep the pair sharded for an end-to-end sharded tail; padded to the + # shard factor. Caller slices to N after any later gather. + return z_dt, N + + z_full = z_dt.full_tensor() + if pad: + z_full = z_full[:, :N, :N, :] + return z_full diff --git a/src/transformers/models/esmfold2/distributed/structure_wrapper.py b/src/transformers/models/esmfold2/distributed/structure_wrapper.py new file mode 100644 index 0000000000..e217ef67e8 --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/structure_wrapper.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed diffusion structure head for ESMFold2. + +``DiffusionModuleCPWrapper`` is a drop-in for ``DiffusionModule`` that shards the +per-step diffusion working set: the conditioned pair ``z`` (via +``DiffusionConditioningDistributed``) and the token transformer's L×L attention +(via ``DiffusionTransformerDistributed``). The atom encoder/decoder and the +single-repr steps run **replicated** on full tensors (cheap 1-D / windowed work). +``wrap_model_with_cp_structure_head`` installs it by swapping +``structure_head.diffusion_module`` — the serial ``sample()`` loop (noise +schedule, augmentation) is untouched and keeps driving it. + +Scope: ``num_diffusion_samples == 1`` (multi-sample batch expansion of the pair +bias is a TODO). The trunk pair ``z_trunk`` still arrives **full** from the +recycle gather, so it remains a per-step floor; the conditioned ``z`` and the +attention are what get sharded here. Removing the ``z_trunk`` floor needs the +pair kept sharded across the recycle→parcae→structure-head boundary (follow-on). +Inference-only. +""" + +from math import lcm + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +_PAIR_PL = [Shard(0), Shard(1), Shard(2)] +_ROW_PL = [Shard(0), Shard(1), Replicate()] + + +class DiffusionModuleCPWrapper(nn.Module): + """Drop-in for ``DiffusionModule`` that distributes conditioning + token attn.""" + + def __init__(self, serial_module: nn.Module, dist_manager) -> None: + super().__init__() + from transformers.models.esmfold2.distributed.model.layers.diffusion_conditioning import ( + DiffusionConditioningDistributed, + ) + from transformers.models.esmfold2.distributed.model.layers.diffusion_transformer import ( + DiffusionTransformerDistributed, + ) + from transformers.models.esmfold2.modeling_esmfold2_common import ( + DiffusionModule as SerialDiffusionModule, + ) + + if not isinstance(serial_module, SerialDiffusionModule): + raise TypeError( + f"expected DiffusionModule, got {type(serial_module).__name__}" + ) + self.serial = serial_module + self.device_mesh = dist_manager.device_mesh_subgroups + self.shard_factor = lcm(self.device_mesh.size(1), self.device_mesh.size(2)) + self.cond = DiffusionConditioningDistributed( + serial_module.conditioning, self.device_mesh + ) + self.token_transformer = DiffusionTransformerDistributed( + serial_module.token_transformer + ) + + def forward( + self, + *, + x_noisy, + t_hat, + ref_pos, + ref_charge, + ref_mask, + ref_element, + ref_atom_name_chars, + ref_space_uid, + tok_idx, + s_inputs, + s_trunk, + z_trunk, + relative_position_encoding, + asym_id, + residue_index, + entity_id, + token_index, + sym_id, + sigma_data=None, + token_attention_mask=None, + num_diffusion_samples: int = 1, + return_token_repr: bool = False, + return_atom_repr: bool = False, + inference_cache=None, + ): + if num_diffusion_samples != 1: + raise NotImplementedError( + "DiffusionModuleCPWrapper supports num_diffusion_samples==1 " + "(pair-bias batch expansion not yet sharded)." + ) + m = self.serial + mesh = self.device_mesh + bsz = x_noisy.shape[0] + sigma = m.sigma_data if sigma_data is None else float(sigma_data) + t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1) + if t.numel() == 1: + t = t.expand(bsz) + + # z_trunk may arrive as a full tensor (pad here) or as an already-padded, + # already-sharded DTensor from the CP tail (#7) — in which case the token + # axis is its padded length and the other (token-space) tensors are padded + # to match. Padded token rows carry no atoms (atom_to_token < L), so they + # don't affect the atom-space coordinate output. + zt_is_dtensor = isinstance(z_trunk, DTensor) + if zt_is_dtensor: + L = s_inputs.shape[1] # original token length + pad = z_trunk.shape[1] - L # to the pre-padded sharded length + else: + L = z_trunk.shape[1] + pad = (self.shard_factor - L % self.shard_factor) % self.shard_factor + + def _pad_pair(x): # (B, L, L, C) + return F.pad(x, (0, 0, 0, pad, 0, pad)) if pad else x + + def _pad_tok(x): # (B, L, C) + return F.pad(x, (0, 0, 0, pad)) if pad else x + + def _pad_mask(x): # (B, L) + return F.pad(x, (0, pad)) if pad else x + + # Step 1: conditioning -> full s, sharded (padded) z. Only materialise + + # distribute the full z_trunk on the cache miss (first step); afterwards + # the sharded z is reused from the cache. + z_cached = inference_cache is not None and "z_cp" in inference_cache + if z_cached: + zt_dt = rp_dt = None + elif zt_is_dtensor: + zt_dt = z_trunk # already padded + sharded; do not re-distribute + rp_dt = distribute_tensor( + _pad_pair(relative_position_encoding).contiguous(), mesh, _PAIR_PL + ) + else: + zt_dt = distribute_tensor(_pad_pair(z_trunk).contiguous(), mesh, _PAIR_PL) + rp_dt = distribute_tensor( + _pad_pair(relative_position_encoding).contiguous(), mesh, _PAIR_PL + ) + s, z_dt = self.cond( + t_hat=t, + s_inputs=s_inputs, + z_trunk_dt=zt_dt, + rel_pos_dt=rp_dt, + sigma_data=sigma, + num_diffusion_samples=num_diffusion_samples, + inference_cache=inference_cache, + ) + + # Step 2: normalise noisy coords (replicated) + denom = torch.sqrt(t * t + sigma * sigma) + r_noisy = x_noisy / denom[:, None, None] + + # Step 3: atom encoder (replicated, full) + a, q_skip, c_skip, p_skip, enc_int = m.atom_encoder( + ref_pos=ref_pos, + atom_attention_mask=ref_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=tok_idx, + r_l=r_noisy, + s_i=s_trunk, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + inference_cache=inference_cache, + ) + + # Step 4: add conditioned s (replicated) + a = a + m.s_to_token(m.s_step_norm(s)) + + # Step 5: token transformer (distributed). Pad token axis; row-shard a/s/ + # mask; z_dt already 2-D sharded at padded length; gather + slice a back. + a_dt = distribute_tensor(_pad_tok(a).contiguous(), mesh, _ROW_PL) + s_dt = distribute_tensor(_pad_tok(s).contiguous(), mesh, _ROW_PL) + mask_dt = None + if token_attention_mask is not None: + mask_dt = distribute_tensor( + _pad_mask(token_attention_mask.to(a.dtype)).contiguous(), mesh, _ROW_PL + ) + a_dt = self.token_transformer(a_dt, s_dt, z_dt, key_mask_dt=mask_dt) + a = a_dt.full_tensor() + if pad: + a = a[:, :L, :] + + # Step 6: token norm (replicated) + a = m.token_norm(a) + + # Step 7: atom decoder (replicated, full) + r_update, dec_int = m.atom_decoder( + a_i=a, + q_l=q_skip, + c_l=c_skip, + p_lm=p_skip, + atom_to_token=tok_idx, + atom_attention_mask=ref_mask, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + ) + + # Step 8: denoised output (replicated) + sigma2, t2 = sigma * sigma, t * t + out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy + out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update + + atom_intermediates = None + if return_atom_repr: + all_ints = enc_int + dec_int + if all_ints: + atom_intermediates = torch.stack(all_ints, dim=2) + return { + "x_denoised": out, + "token_repr": a if return_token_repr else None, + "atom_intermediates": atom_intermediates, + } + + +def wrap_model_with_cp_structure_head(model: nn.Module, dist_manager) -> list[str]: + """Replace each ``DiffusionStructureHead``'s ``diffusion_module`` with the CP + wrapper. The serial ``sample()`` loop keeps driving it. Returns replaced paths.""" + from transformers.models.esmfold2.modeling_esmfold2_common import ( + DiffusionModule as SerialDiffusionModule, + ) + + replaced: list[str] = [] + for name, module in model.named_modules(): + dm_child = getattr(module, "diffusion_module", None) + if isinstance(dm_child, SerialDiffusionModule): + module.diffusion_module = DiffusionModuleCPWrapper(dm_child, dist_manager).to( + device=dist_manager.device, dtype=next(dm_child.parameters()).dtype + ) + replaced.append(f"{name}.diffusion_module" if name else "diffusion_module") + return replaced diff --git a/src/transformers/models/esmfold2/distributed/utils.py b/src/transformers/models/esmfold2/distributed/utils.py index 8632b9c0d4..147d056e74 100644 --- a/src/transformers/models/esmfold2/distributed/utils.py +++ b/src/transformers/models/esmfold2/distributed/utils.py @@ -29,7 +29,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import Shard, distribute_tensor +from torch.distributed.tensor import DTensor, Shard, distribute_tensor class LayoutMap: @@ -445,6 +445,24 @@ def set_kernel_backend(self, _backend: str | None) -> None: def set_chunk_size(self, _chunk_size: int | None) -> None: return + def forward_sharded( + self, pair_dt: DTensor, pair_attention_mask: DTensor | None = None + ) -> DTensor: + """Run the distributed trunk on an already-sharded pair DTensor. + + ``pair_dt`` (and the optional mask) must already be padded to a multiple + of ``self.shard_factor``, cast to the trunk dtype, and distributed with + placements ``[Shard(0), Shard(1), Shard(2)]`` on ``self.device_mesh``. + Returns a DTensor with the same placements — **no** ``full_tensor()`` + gather. + + This is the entry point a CP orchestrator uses to keep the pair sharded + across module boundaries (e.g. the recycle loop), instead of paying a + ``full_tensor()`` + ``distribute_tensor`` round-trip on every call. The + plain-tensor :meth:`forward` is a thin adapter around it. + """ + return self.dist_trunk(pair_dt, pair_attention_mask=pair_attention_mask) + def forward( self, pair: torch.Tensor, pair_attention_mask: torch.Tensor | None = None ) -> torch.Tensor: @@ -476,7 +494,7 @@ def forward( [Shard(0), Shard(1), Shard(2)], ) - out_dt = self.dist_trunk(pair_dt, pair_attention_mask=mask_dt) + out_dt = self.forward_sharded(pair_dt, pair_attention_mask=mask_dt) out = out_dt.full_tensor() if self._trunk_bf16: out = out.to(orig_dtype) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 4dd33609a7..31214ebaf8 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -226,6 +226,55 @@ def forward( del pair_delta single = self.row_attention_pooling(pair, mask) + pae_logits = self.pae_head(self.pae_ln(pair)) + pde_logits = self.pde_head(self.pde_ln(pair)) + return self._finish( + single=single, + pae_logits=pae_logits, + pde_logits=pde_logits, + x_pred=x_pred, + distogram_atom_idx=distogram_atom_idx, + token_attention_mask=token_attention_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atom_attention_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=num_diffusion_samples, + ) + + def _finish( + self, + *, + single: Tensor, + pae_logits: Tensor, + pde_logits: Tensor, + x_pred: Tensor, + distogram_atom_idx: Tensor, + token_attention_mask: Tensor, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + asym_id: Tensor, + mol_type: Tensor, + num_diffusion_samples: int, + ) -> dict[str, Tensor]: + """Atom-space (pLDDT/resolved) + PAE/PDE + pTM/ipTM back-half. + + Split out of ``forward`` so the CP wrapper (#6) can run the distributed + pair front (z_base → trunk → row-pool → pae/pde heads), gather the small + ``single`` / ``pae_logits`` / ``pde_logits``, and reuse this exact + (serial, replicated) reduction code — one source of truth for the fiddly + pTM/ipTM/pair_chains logic.""" + x_pred_flat = self._flatten_sample_axis(x_pred) + atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples) + atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) + rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() + mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) + Bm = single.shape[0] + rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + atom_mask_f = atom_mask_m.float() s_at_atoms = gather_token_to_atom(single, atom_to_token_m) s_at_atoms_ln = self.plddt_ln(s_at_atoms) @@ -275,12 +324,8 @@ def forward( plddt_ca = plddt_per_atom.gather(1, rep_idx_m) - # PAE - pae_logits = self.pae_head(self.pae_ln(pair)) + # PAE / PDE (logits computed by the caller — serial front or CP wrapper). pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach() - - # PDE - pde_logits = self.pde_head(self.pde_ln(pair)) pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach() # Resolved (per-atom binary). @@ -775,6 +820,27 @@ def _run_one_loop( tok_mask: Tensor, total_steps: int, ) -> Tensor: + # CP orchestrator seam: when a context-parallel recycle engine has been + # installed (by ``wrap_model_with_cp``), delegate the whole loop to it so + # the pair stays sharded across iterations instead of round-tripping + # through ``full_tensor()`` at every module boundary. Plain ``getattr`` + # default keeps this file CP-agnostic (no distributed import). The engine + # returns a gathered full tensor, so the caller is unchanged. + _cp_engine = getattr(self, "_cp_recycle_engine", None) + if _cp_engine is not None: + return _cp_engine.run_loop( + self, + z=z, + z_init=z_init, + lm_z=lm_z, + msa_inputs=_msa_inputs, + pair_mask=pair_mask, + a=a, + b_mat=b_mat, + tok_mask=tok_mask, + total_steps=total_steps, + ) + # Helper method (not inline) so per-iter locals free on return — # otherwise leaks ~2 GB L²×c_z into distogram/sample scope. # training=True forces dropout under eval(), matching the per-loop @@ -992,7 +1058,19 @@ def forward( ) lm_z: Tensor | None = None if lm_hidden_states is not None: - lm_z = self.language_model(lm_hidden_states.detach()) + # CP (#14): when a distributed LM->pair builder is installed, emit + # lm_z as a sharded (Shard0,Shard1,Shard2) DTensor so the full L×L + # pair is never resident on any rank (the binding peak per the T0 + # phase sweep). Consumed by the recycle engine, installed together + # with it. Plain getattr keeps this file CP-agnostic. + # Only emit a sharded lm_z when the recycle engine (its sole + # consumer) is active; the serial _run_one_loop expects a full + # tensor. Keeps the gather-per-iteration fallback path valid. + _cp_lm = getattr(self, "_cp_language_model", None) + if _cp_lm is not None and getattr(self, "_cp_recycle_engine", None) is not None: + lm_z, _ = _cp_lm(lm_hidden_states.detach()) + else: + lm_z = self.language_model(lm_hidden_states.detach()) del lm_hidden_states # ESM-C is done for this forward — offload to free its ~12 GB for the # trunk/diffusion (freed blocks stay in the caching pool for reuse). @@ -1025,28 +1103,64 @@ def forward( subsample_enabled=msa_subsample_at_inference, ) - # Method call (not inline loop) frees per-iter L²×c_z locals. - z = self._run_one_loop( - z=z, - z_init=z_init, - lm_z=lm_z, - _msa_inputs=_msa_inputs, - pair_mask=pair_mask, - a=a, - b_mat=b_mat, - tok_mask=tok_mask, - total_steps=total_steps, - ) - del z_init, lm_z, _msa_inputs, a, b_mat + # CP tail (#7): when a recycle engine is installed, keep the pair + # SHARDED through the recycle loop + parcae so the full L×L pair is + # not resident during the expensive structure phase. The pair lives as + # ``_z_dt`` (DTensor); it is gathered only transiently for the heads + # that don't yet consume sharded input (distogram, confidence). Plain + # ``getattr`` keeps this file CP-agnostic. + _cp_engine = getattr(self, "_cp_recycle_engine", None) + _z_dt = None + _n_orig = z.shape[1] + if _cp_engine is not None: + _z_dt, _n_orig = _cp_engine.run_loop( + self, + z=z, + z_init=z_init, + lm_z=lm_z, + msa_inputs=_msa_inputs, + pair_mask=pair_mask, + a=a, + b_mat=b_mat, + tok_mask=tok_mask, + total_steps=total_steps, + gather=False, + ) + del z_init, lm_z, _msa_inputs, a, b_mat + _z_dt = _cp_engine.parcae_finish(self, _z_dt, pair_mask) + z = None # pair lives sharded in _z_dt; gathered transiently below + else: + # Method call (not inline loop) frees per-iter L²×c_z locals. + z = self._run_one_loop( + z=z, + z_init=z_init, + lm_z=lm_z, + _msa_inputs=_msa_inputs, + pair_mask=pair_mask, + a=a, + b_mat=b_mat, + tok_mask=tok_mask, + total_steps=total_steps, + ) + del z_init, lm_z, _msa_inputs, a, b_mat - z = self.parcae_readout(z) - z = self.parcae_coda(z, pair_attention_mask=pair_mask) + z = self.parcae_readout(z) + z = self.parcae_coda(z, pair_attention_mask=pair_mask) - z = z.float() - distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + z = z.float() + _cp_disto = getattr(self, "_cp_distogram_head", None) + if _z_dt is not None and _cp_disto is not None: + # Distributed: symmetrize + head off the sharded pair, no full-z gather. + distogram_logits = _cp_disto.forward_sharded(_z_dt, _n_orig) + else: + if _z_dt is not None: + z = _z_dt.full_tensor()[:, :_n_orig, :_n_orig, :].float() + distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + if _z_dt is not None: + del z # free the full pair before the (sharded) structure phase structure_output = self.structure_head.sample( - z_trunk=z, + z_trunk=(_z_dt if _z_dt is not None else z), s_inputs=x_inputs, s_trunk=None, relative_position_encoding=relative_position_encoding, @@ -1074,20 +1188,43 @@ def forward( output: dict[str, Tensor] = {"distogram_logits": distogram_logits} output["sample_atom_coords"] = sample_coords - confidence_output = self.confidence_head( - s_inputs=x_inputs.detach(), - z=z.detach().float(), - x_pred=sample_coords.detach(), - distogram_atom_idx=disto_idx, - token_attention_mask=tok_mask, - atom_to_token=atom_to_token, - atom_attention_mask=atm_mask, - asym_id=asym_id, - mol_type=mol_type, - num_diffusion_samples=n_samples, - relative_position_encoding=relative_position_encoding.detach(), - token_bonds_encoding=token_bonds_encoding.detach(), - ) + # Confidence head (#6): on the CP tail, run it distributed straight off the + # sharded pair (no full-L×L re-gather of z) when the wrapper is installed; + # otherwise re-gather transiently and run serially. CP-agnostic getattr. + _cp_conf = getattr(self, "_cp_confidence_head", None) + if _z_dt is not None and _cp_conf is not None: + confidence_output = _cp_conf.forward_sharded( + _z_dt, + _n_orig, + s_inputs=x_inputs.detach(), + x_pred=sample_coords.detach(), + distogram_atom_idx=disto_idx, + token_attention_mask=tok_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atm_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=n_samples, + relative_position_encoding=relative_position_encoding.detach(), + token_bonds_encoding=token_bonds_encoding.detach(), + ) + else: + if _z_dt is not None: + z = _z_dt.full_tensor()[:, :_n_orig, :_n_orig, :].float() + confidence_output = self.confidence_head( + s_inputs=x_inputs.detach(), + z=z.detach().float(), + x_pred=sample_coords.detach(), + distogram_atom_idx=disto_idx, + token_attention_mask=tok_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atm_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=n_samples, + relative_position_encoding=relative_position_encoding.detach(), + token_bonds_encoding=token_bonds_encoding.detach(), + ) output.update(confidence_output) output["atom_pad_mask"] = ( atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask From 67a26d2d2f68c6fd81ea5fe363eaa122b1bbe17f Mon Sep 17 00:00:00 2001 From: Greg Zynda Date: Wed, 1 Jul 2026 08:16:36 -0700 Subject: [PATCH 3/4] -Empty cache after ESMC offload -SVD fallsback to CPU --- .../models/esmfold2/modeling_esmfold2.py | 5 ++++- .../models/esmfold2/modeling_esmfold2_common.py | 12 +++++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 31214ebaf8..912b678090 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -1073,9 +1073,12 @@ def forward( lm_z = self.language_model(lm_hidden_states.detach()) del lm_hidden_states # ESM-C is done for this forward — offload to free its ~12 GB for the - # trunk/diffusion (freed blocks stay in the caching pool for reuse). + # trunk/diffusion. Return the freed blocks to the driver (not just the + # caching pool) so cuSOLVER's out-of-pool cudaMalloc — for its handle in + # the diffusion rigid-align SVD — can succeed near the OOM boundary. if self._offload_esmc and self._esmc is not None: self._esmc.to("cpu") + torch.cuda.empty_cache() pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() diff --git a/src/transformers/models/esmfold2/modeling_esmfold2_common.py b/src/transformers/models/esmfold2/modeling_esmfold2_common.py index 4c27da1823..a746aa4c29 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2_common.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2_common.py @@ -1751,7 +1751,17 @@ def _weighted_rigid_align( xgt_c = x_gt - mu_gt H = torch.einsum("bni,bnj->bij", w * xgt_c, x_c) H32 = H.float() - U, _, Vh = torch.linalg.svd(H32, driver="gesvd" if H32.is_cuda else None) + try: + U, _, Vh = torch.linalg.svd(H32, driver="gesvd" if H32.is_cuda else None) + except RuntimeError: + # Near the OOM boundary cuSOLVER can fail to allocate its handle + # (cusolverDnCreate) even though this 3x3 SVD is trivial. Fall back to + # a CPU SVD of the tiny matrix — instant, and only reached when the GPU + # path would otherwise crash. driver=None (gesvd is cuSOLVER-only). + if not H32.is_cuda: + raise + Ucpu, _, Vhcpu = torch.linalg.svd(H32.cpu(), driver=None) + U, Vh = Ucpu.to(H32.device), Vhcpu.to(H32.device) det = torch.linalg.det(U @ Vh) ones = torch.ones_like(det) R = (U @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ Vh).to( From 942df6894333134d06c597105b57945173bb108b Mon Sep 17 00:00:00 2001 From: Greg Zynda Date: Wed, 8 Jul 2026 18:32:52 -0700 Subject: [PATCH 4/4] =?UTF-8?q?ESMFold2:=20shard=20the=20CP=20pair-init=20?= =?UTF-8?q?so=20no=20full=20L=C3=97L=20is=20resident=20per=20rank?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../distributed/confidence_wrapper.py | 33 +- .../model/layers/confidence_zbase.py | 26 +- .../distributed/model/layers/pair_init.py | 374 ++++++++++++++++++ .../esmfold2/distributed/msa_wrapper.py | 27 +- .../models/esmfold2/distributed/recycle.py | 56 ++- .../esmfold2/distributed/structure_wrapper.py | 18 +- .../models/esmfold2/modeling_esmfold2.py | 68 +++- 7 files changed, 535 insertions(+), 67 deletions(-) create mode 100644 src/transformers/models/esmfold2/distributed/model/layers/pair_init.py diff --git a/src/transformers/models/esmfold2/distributed/confidence_wrapper.py b/src/transformers/models/esmfold2/distributed/confidence_wrapper.py index d8337a33c7..668cf25aa7 100644 --- a/src/transformers/models/esmfold2/distributed/confidence_wrapper.py +++ b/src/transformers/models/esmfold2/distributed/confidence_wrapper.py @@ -44,11 +44,15 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.tensor import DTensor, Shard, distribute_tensor +from torch.distributed.tensor import DTensor from transformers.models.esmfold2.distributed.model.layers.confidence_zbase import ( ConfidenceZBaseDistributed, ) +from transformers.models.esmfold2.distributed.model.layers.pair_init import ( + build_sharded_distogram_bins, + build_sharded_pair_mask, +) from transformers.models.esmfold2.distributed.model.layers.layernorm import ( LayerNormParamsReplicated, ) @@ -62,8 +66,6 @@ gather_rep_atom_coords, ) -_PAIR = [Shard(0), Shard(1), Shard(2)] - class ConfidenceHeadCPWrapper(nn.Module): """Distributed wrapper around the serial ``ConfidenceHead``. @@ -115,22 +117,22 @@ def forward_sharded( pad = Lpad - n_orig pair_dtype = torch.float32 - # --- distogram bins (full, replicated; small single-channel L²) --------- + # --- distogram bins: sharded block cdist (never the full [B, N, N]) ----- rep_idx = distogram_atom_idx.long() - rep_coords = gather_rep_atom_coords(x_pred, rep_idx) # (B, n, 3) - rep_distances = torch.cdist( - rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" - ) - distogram_bins = ( - (rep_distances.unsqueeze(-1) > self.head.boundaries).sum(dim=-1).long() - ) + rep_coords = gather_rep_atom_coords(x_pred, rep_idx) # (B, n, 3) — replicated, cheap + bins_p = build_sharded_distogram_bins( + rep_coords, self.head.boundaries, mesh + ) # (S0, S1, S2) [B, Lpad, Lpad] # --- pad aux tensors to the shard factor (z_dt is already padded) ------- def _pad_pair(t): + # An already-sharded DTensor is padded (built at the shard factor) — + # pass it through; only a full tensor needs padding here. + if isinstance(t, DTensor): + return t return F.pad(t, (0, 0, 0, pad, 0, pad)) if pad else t s_in_p = F.pad(s_inputs, (0, 0, 0, pad)) if pad else s_inputs - bins_p = F.pad(distogram_bins, (0, pad, 0, pad)) if pad else distogram_bins relpos_p = ( _pad_pair(relative_position_encoding) if relative_position_encoding is not None @@ -151,11 +153,10 @@ def _pad_pair(t): ) # (S0, S1, S2) fp32, [B, Lpad, Lpad, d_pair] # --- nested FoldingTrunk (add-back, like serial pair.add_(trunk(pair))) - - pair_mask = (mask_p[:, :, None] * mask_p[:, None, :]).contiguous() # (B,Lpad,Lpad) + # Pair mask built SHARDED from the (padded) token mask — the per-block outer + # product, never a full [B, Lpad, Lpad] (matches the recycle / MSA masks). with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): - pair_mask_dt = distribute_tensor( - pair_mask.to(torch.bfloat16), mesh, _PAIR - ) + pair_mask_dt = build_sharded_pair_mask(mask_p, mesh).to(torch.bfloat16) delta_dt = self.head.folding_trunk.forward_sharded( pair_dt.to(torch.bfloat16), pair_attention_mask=pair_mask_dt ) diff --git a/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py b/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py index 443767b6ec..9148f5f7ab 100644 --- a/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py +++ b/src/transformers/models/esmfold2/distributed/model/layers/confidence_zbase.py @@ -106,23 +106,31 @@ def col(t): # local cp_axis_1 cols p1_row = row(self.s_to_z_prod_in1(s)) # (B, sLi, d_pair) p2_col = col(self.s_to_z_prod_in2(s)) # (B, sLj, d_pair) + # rel_pos / token_bonds may arrive already sharded (built block-local by + # PairInitDistributed) — use the local shard directly; only a full tensor + # needs the distribute-then-take-local round trip. + def _pair_local(t): + if isinstance(t, DTensor): + return t.to_local() + return distribute_tensor(t.contiguous(), mesh, _PAIR).to_local() + zb = self.z_norm(z).to_local() # (B, sLi, sLj, d_pair) if relative_position_encoding is not None: - zb = zb + distribute_tensor( - relative_position_encoding.contiguous(), mesh, _PAIR - ).to_local() + zb = zb + _pair_local(relative_position_encoding) if token_bonds_encoding is not None: - zb = zb + distribute_tensor( - token_bonds_encoding.contiguous(), mesh, _PAIR - ).to_local() + zb = zb + _pair_local(token_bonds_encoding) zb = zb + a_row[:, :, None, :] + b_col[:, None, :, :] prod = p1_row[:, :, None, :] * p2_col[:, None, :, :] # (B, sLi, sLj, d_pair) zb = zb + F.linear(prod, self.s_to_z_prod_out.weight) - bins_local = distribute_tensor( - distogram_bins.contiguous(), mesh, _PAIR - ).to_local() # (B, sLi, sLj) + # distogram_bins may arrive already sharded (block cdist by + # build_sharded_distogram_bins) — take the local shard directly. + bins_local = ( + distogram_bins.to_local() + if isinstance(distogram_bins, DTensor) + else distribute_tensor(distogram_bins.contiguous(), mesh, _PAIR).to_local() + ) # (B, sLi, sLj) zb = zb + self.dist_bin_pairwise_embed(bins_local) b_size = zb.shape[0] diff --git a/src/transformers/models/esmfold2/distributed/model/layers/pair_init.py b/src/transformers/models/esmfold2/distributed/model/layers/pair_init.py new file mode 100644 index 0000000000..a90cb12a3c --- /dev/null +++ b/src/transformers/models/esmfold2/distributed/model/layers/pair_init.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +"""Distributed sharded pair-initialisation for ESMFold2 (inference-only). + +The serial forward builds the initial pair tensors — ``z_init`` (an outer sum of +two per-token projections), the relative-position encoding, and the token-bond +encoding — as full ``[B, L, L, d_pair]`` tensors on *every* rank, then only +shards them at the recycle-loop boundary (``distribute_tensor`` inside +``CPRecycleEngine.run_loop``). That makes 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 for the whole forward), so CP buys no memory relief for +the init phase and OOMs at short L. + +``PairInitDistributed`` builds each of these directly as a 2-D-sharded +``(Shard(0), Shard(1), Shard(2))`` DTensor by computing only this rank's +``(row, col)`` block — the full L×L tensor is never resident. Per-token inputs +(``x_inputs`` and the index tensors) are cheap ``[B, L, *]`` and stay replicated; +only their O(L²) outer combinations are sharded. This mirrors evolutionaryscale's +``_sharded_rel_pos_encoding`` + sliced ``z_init`` (``ESMCFoldModelCP``), wrapped +as DTensors so the recycle engine / structure-head wrapper consume them with no +gather. + +The channel-wise submodules (``z_init_1`` / ``z_init_2``, the ``token_bonds`` +Linear, and ``rel_pos.embed``) are the model's own layers, run on the local +block — their weights are identical on every rank, so the block result equals the +corresponding slice of the serial full tensor. +""" + +from math import lcm, sqrt + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.tensor import DTensor, Shard + +from transformers.models.esmfold2.modeling_esmfold2_common import ( + ResIdxAsymIdSymIdEntityIdEncoding, +) + +_PAIR_PL = [Shard(0), Shard(1), Shard(2)] + +# Per-rank RNG seed base for the sharded initial pair-state draw. Kept off the +# global generator so the recycle loop's MSA-subsample RNG stays synchronized +# across ranks (see ``init_pair_state``). +_Z_INIT_SEED = 0x2717 + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return tuple(strides) + + +def _block_geometry(n_orig: int, device_mesh): + """This rank's pair-block geometry: ``(lpad, s0, s1, row_slice, col_slice)``. + + The token axis is padded to ``shard_factor = lcm(cp_axis_0, cp_axis_1)``; the + rank holds rows ``row_slice`` (its cp_axis_0 coordinate) × cols ``col_slice`` + (its cp_axis_1 coordinate) of the ``[lpad, lpad]`` pair. + """ + cp0 = device_mesh.size(1) + cp1 = device_mesh.size(2) + shard_factor = lcm(cp0, cp1) + coord = device_mesh.get_coordinate() + assert coord is not None, "device mesh has no coordinate for this rank" + row_rank, col_rank = int(coord[1]), int(coord[2]) + pad = (shard_factor - n_orig % shard_factor) % shard_factor + lpad = n_orig + pad + s0 = lpad // cp0 + s1 = lpad // cp1 + row = slice(row_rank * s0, (row_rank + 1) * s0) + col = slice(col_rank * s1, (col_rank + 1) * s1) + return lpad, s0, s1, row, col + + +def build_sharded_pair_mask(tok_mask: torch.Tensor, device_mesh) -> DTensor: + """This rank's ``[B, s0, s1]`` block of the outer-product pair mask as a + ``(Shard0, Shard1, Shard2)`` DTensor — never the full ``[B, L, L]``. + + Mirrors evolutionaryscale's sharded ``_cp_pair_mask``: outer product of the + row / col slices of the per-token mask (``tok_mask[:, row]`` ⊗ + ``tok_mask[:, col]``), value-identical to + ``distribute_tensor(tok_mask[:,:,None] * tok_mask[:,None,:])`` but with no full + L×L intermediate. The token axis is padded to + ``shard_factor = lcm(cp_axis_0, cp_axis_1)`` (padded positions → 0, matching + the serial zero-pad). Shared by ``PairInitDistributed.pair_mask`` (recycle + pair mask) and the MSA encoder's ``forward_sharded`` (per-iteration mask). + + ``tok_mask`` is ``[B, L]`` (any numeric / bool dtype); the returned DTensor is + float32 — cast at the call site to the consumer dtype. + """ + b = tok_mask.shape[0] + lpad, _, _, row, col = _block_geometry(tok_mask.shape[1], device_mesh) + pad = lpad - tok_mask.shape[1] + + tm = tok_mask.float() + if pad: + tm = F.pad(tm, (0, pad)) + block = tm[:, row].unsqueeze(2) * tm[:, col].unsqueeze(1) # [B, s0, s1] + + shape = torch.Size((b, lpad, lpad)) + return DTensor.from_local( + block.contiguous(), + device_mesh=device_mesh, + placements=_PAIR_PL, + stride=_contiguous_strides(tuple(shape)), + shape=shape, + ) + + +def build_sharded_distogram_bins( + rep_coords: torch.Tensor, boundaries: torch.Tensor, device_mesh +) -> DTensor: + """This rank's ``[B, s0, s1]`` block of the distance-distogram bins as a + ``(Shard0, Shard1, Shard2)`` DTensor — ``cdist`` on the row / col coordinate + blocks only, never the full ``[B, N, N]`` distance matrix. + + ``rep_coords`` is ``[B, N, 3]`` per-token representative-atom coordinates + (replicated, cheap). ``boundaries`` is the distance-bin edge buffer. Padded + rows/cols → bin 0, matching the serial ``F.pad(distogram_bins)`` zero-fill. + There is no evolutionaryscale analogue — its confidence head computes the full + ``[B, N, N]`` cdist (the head is not CP-distributed). + """ + b = rep_coords.shape[0] + n = rep_coords.shape[1] + lpad, _, _, row, col = _block_geometry(n, device_mesh) + pad = lpad - n + + coords = F.pad(rep_coords, (0, 0, 0, pad)) if pad else rep_coords + row_c = coords[:, row] # [B, s0, 3] + col_c = coords[:, col] # [B, s1, 3] + dist = torch.cdist(row_c, col_c, compute_mode="donot_use_mm_for_euclid_dist") + bins = (dist.unsqueeze(-1) > boundaries).sum(dim=-1).long() # [B, s0, s1] + if pad: + rr = torch.arange(row.start, row.stop, device=rep_coords.device) < n + cc = torch.arange(col.start, col.stop, device=rep_coords.device) < n + bins = bins * (rr[:, None] & cc[None, :]) # padded positions → bin 0 + + shape = torch.Size((b, lpad, lpad)) + return DTensor.from_local( + bins.contiguous(), + device_mesh=device_mesh, + placements=_PAIR_PL, + stride=_contiguous_strides(tuple(shape)), + shape=shape, + ) + + +class PairInitDistributed(nn.Module): + """Builds ``z_init`` / ``rel_pos`` / ``token_bonds`` as 2-D-sharded DTensors. + + ``forward`` returns ``(z_init_dt, rel_pos_dt, token_bonds_dt, n_orig)``: the + three pair tensors as padded ``(Shard0, Shard1, Shard2)`` DTensors (token axis + padded to ``shard_factor = lcm(cp_axis_0, cp_axis_1)`` — the same padding + contract the recycle engine and LM->pair builder use) plus the original + (pre-pad) token count. ``z_init_dt`` already folds in ``rel_pos`` + + ``token_bonds`` (matching the serial ``z_init`` sum); ``rel_pos_dt`` / + ``token_bonds_dt`` are returned separately for the structure / confidence + heads. + """ + + def __init__(self, model, dist_manager) -> None: + super().__init__() + self.device_mesh = dist_manager.device_mesh_subgroups + # device_mesh is (dp, cp_axis_0, cp_axis_1); rows shard on cp_axis_0, cols on cp_axis_1. + self.cp_axis_0 = self.device_mesh.size(1) + self.cp_axis_1 = self.device_mesh.size(2) + self.shard_factor = lcm(self.cp_axis_0, self.cp_axis_1) + coord = self.device_mesh.get_coordinate() + assert coord is not None, "device mesh has no coordinate for this rank" + self.row_rank = int(coord[1]) + self.col_rank = int(coord[2]) + + # Model's own channel-wise layers (run on the local block; weights replicated). + self.z_init_1 = model.z_init_1 + self.z_init_2 = model.z_init_2 + self.token_bonds = model.token_bonds + self.rel_pos: ResIdxAsymIdSymIdEntityIdEncoding = model.rel_pos + + def _block_valid_mask( + self, n_orig: int, lpad: int, device: torch.device + ) -> torch.Tensor | None: + """This rank's ``[s0, s1]`` boolean mask marking (row, col) positions that + are *real* tokens (global index < ``n_orig``), or ``None`` if there is no + CP padding. Used to zero the padded rows/cols so the sharded pair matches + the serial ``_pad_pair`` (zero-fill) contract — padded positions must not + leak nonzero values into the masked trunk. Note this is a *position* mask + (CP padding only), NOT the token attention mask: in-range but attention- + masked tokens stay nonzero here, exactly as in the serial full tensor.""" + if lpad - n_orig <= 0: + return None + s0 = lpad // self.cp_axis_0 + s1 = lpad // self.cp_axis_1 + r0 = self.row_rank * s0 + c0 = self.col_rank * s1 + row_valid = torch.arange(r0, r0 + s0, device=device) < n_orig + col_valid = torch.arange(c0, c0 + s1, device=device) < n_orig + return row_valid[:, None] & col_valid[None, :] # [s0, s1] bool + + def _wrap(self, local: torch.Tensor, b: int, lpad: int, d: int) -> DTensor: + shape = torch.Size((b, lpad, lpad, d)) + return DTensor.from_local( + local.contiguous(), + device_mesh=self.device_mesh, + placements=_PAIR_PL, + shape=shape, + stride=_contiguous_strides(tuple(shape)), + ) + + def _sharded_rel_pos( + self, + residue_index: torch.Tensor, + asym_id: torch.Tensor, + sym_id: torch.Tensor, + entity_id: torch.Tensor, + token_index: torch.Tensor, + row: slice, + col: slice, + ) -> torch.Tensor: + """This rank's ``[B, s0, s1, d_pair]`` block of the relative-position + encoding. Byte-for-byte the arithmetic of + ``ResIdxAsymIdSymIdEntityIdEncoding.forward`` but with the index tensors + sliced to the row / col block before the outer differences (so no full + ``[B, L, L, *]`` one-hot intermediate).""" + rp = self.rel_pos + bins_r = rp.n_relative_residx_bins + bins_c = rp.n_relative_chain_bins + + ri_r, ri_c = residue_index[:, row], residue_index[:, col] + ai_r, ai_c = asym_id[:, row], asym_id[:, col] + si_r, si_c = sym_id[:, row], sym_id[:, col] + ei_r, ei_c = entity_id[:, row], entity_id[:, col] + ti_r, ti_c = token_index[:, row], token_index[:, col] + + same_chain = ai_r.unsqueeze(2) == ai_c.unsqueeze(1) + same_res = ri_r.unsqueeze(2) == ri_c.unsqueeze(1) + same_ent = ei_r.unsqueeze(2) == ei_c.unsqueeze(1) + + d_res = torch.clip(ri_r.unsqueeze(2) - ri_c.unsqueeze(1) + bins_r, 0, 2 * bins_r) + d_res = torch.where(same_chain, d_res, 2 * bins_r + 1) + a_res = F.one_hot(d_res, 2 * bins_r + 2) + + d_tok = torch.clip(ti_r.unsqueeze(2) - ti_c.unsqueeze(1) + bins_r, 0, 2 * bins_r) + d_tok = torch.where(same_chain & same_res, d_tok, 2 * bins_r + 1) + a_tok = F.one_hot(d_tok, 2 * bins_r + 2) + + d_ch = torch.clip(si_r.unsqueeze(2) - si_c.unsqueeze(1) + bins_c, 0, 2 * bins_c) + d_ch = torch.where(same_chain, 2 * bins_c + 1, d_ch) + a_ch = F.one_hot(d_ch, 2 * bins_c + 2) + + feats = torch.cat( + [a_res.float(), a_tok.float(), same_ent.float().unsqueeze(-1), a_ch.float()], + dim=-1, + ) + return rp.embed(feats) # [B, s0, s1, d_pair] + + def forward( + self, + x_inputs: torch.Tensor, + residue_index: torch.Tensor, + asym_id: torch.Tensor, + sym_id: torch.Tensor, + entity_id: torch.Tensor, + token_index: torch.Tensor, + token_bonds: torch.Tensor, + ) -> tuple[DTensor, DTensor, DTensor, int]: + b = x_inputs.shape[0] + n_orig = x_inputs.shape[1] + pad = (self.shard_factor - n_orig % self.shard_factor) % self.shard_factor + lpad = n_orig + pad + s0 = lpad // self.cp_axis_0 + s1 = lpad // self.cp_axis_1 + row = slice(self.row_rank * s0, (self.row_rank + 1) * s0) + col = slice(self.col_rank * s1, (self.col_rank + 1) * s1) + + # --- z_init: outer SUM of two per-token projections, sliced to the block. + # z1/z2 are [B, L, d] (per-token, cheap — not L×L), so building them full and + # slicing costs O(L*d), not O(L²). + z1 = self.z_init_1(x_inputs) + z2 = self.z_init_2(x_inputs) + d = z1.shape[-1] + if pad: + z1 = F.pad(z1, (0, 0, 0, pad)) + z2 = F.pad(z2, (0, 0, 0, pad)) + z_init_local = z1[:, row].unsqueeze(2) + z2[:, col].unsqueeze(1) # [B, s0, s1, d] + + # --- relative-position encoding (index tensors sliced to row/col block) --- + def _pad_idx(idx: torch.Tensor) -> torch.Tensor: + return F.pad(idx, (0, pad)) if pad else idx + + rel_local = self._sharded_rel_pos( + _pad_idx(residue_index), + _pad_idx(asym_id), + _pad_idx(sym_id), + _pad_idx(entity_id), + _pad_idx(token_index), + row, + col, + ) # [B, s0, s1, d] + + # --- token-bond encoding: slice the (valid) input block, then pad the small + # block to [s0, s1] — never materialises a full padded L×L×1 tensor. + r0 = self.row_rank * s0 + c0 = self.col_rank * s1 + tb_v = token_bonds[:, r0 : min(r0 + s0, n_orig), c0 : min(c0 + s1, n_orig), :].float() + pr = s0 - tb_v.shape[1] + pc = s1 - tb_v.shape[2] + if pr or pc: + tb_v = F.pad(tb_v, (0, 0, 0, pc, 0, pr)) + tb_local = self.token_bonds(tb_v) # [B, s0, s1, d] + + z_comb = z_init_local + rel_local + tb_local + # Zero the CP-padded rows/cols so the sharded pair matches the serial + # zero-padding (else e.g. z_init_local[pad_row, j] = z2[j] would leak). + vmask = self._block_valid_mask(n_orig, lpad, x_inputs.device) + if vmask is not None: + z_comb = z_comb * vmask[None, :, :, None].to(z_comb.dtype) + rel_local = rel_local * vmask[None, :, :, None].to(rel_local.dtype) + tb_local = tb_local * vmask[None, :, :, None].to(tb_local.dtype) + + z_init_dt = self._wrap(z_comb, b, lpad, d) + rel_pos_dt = self._wrap(rel_local, b, lpad, d) + token_bonds_dt = self._wrap(tb_local, b, lpad, d) + return z_init_dt, rel_pos_dt, token_bonds_dt, n_orig + + def pair_mask(self, tok_mask: torch.Tensor) -> DTensor: + """This rank's sharded ``[B, s0, s1]`` block of the pair attention mask + (see :func:`build_sharded_pair_mask`).""" + return build_sharded_pair_mask(tok_mask, self.device_mesh) + + def init_pair_state(self, z_init_dt: DTensor, n_orig: int) -> DTensor: + """Sharded random initial pair state (the serial ``_init_pair_state``). + + Draws this rank's local block from a *per-rank* generator (distinct seed) + so (a) the full L×L state is never resident and (b) the global RNG — which + the recycle loop's MSA subsample relies on to stay in lockstep across + ranks — is untouched. Not bit-exact with the serial full-tensor + ``trunc_normal_`` draw (a valid-but-different iid sample, and ``clamp`` vs + resample at the ±3σ bound), matching the shard-local-dropout contract + already used in the recycle engine; validate by end-to-end pLDDT/pTM + parity. CP-padded rows/cols are zeroed to match the serial ``_pad_pair``. + """ + ref = z_init_dt.to_local() + b, lpad, _, d = z_init_dt.shape + std = sqrt(2.0 / (5.0 * ref.shape[-1])) + gen = torch.Generator(device=ref.device) + gen.manual_seed(_Z_INIT_SEED + torch.distributed.get_rank()) + state = torch.empty(ref.shape, dtype=torch.float32, device=ref.device) + state.normal_(0.0, std, generator=gen).clamp_(-3 * std, 3 * std) + state = state.to(dtype=ref.dtype) + vmask = self._block_valid_mask(int(n_orig), int(lpad), ref.device) + if vmask is not None: + state = state * vmask[None, :, :, None].to(state.dtype) + return self._wrap(state, int(b), int(lpad), int(d)) diff --git a/src/transformers/models/esmfold2/distributed/msa_wrapper.py b/src/transformers/models/esmfold2/distributed/msa_wrapper.py index c686332d5b..96ea4fccf1 100644 --- a/src/transformers/models/esmfold2/distributed/msa_wrapper.py +++ b/src/transformers/models/esmfold2/distributed/msa_wrapper.py @@ -146,11 +146,6 @@ def forward_sharded( m = F.pad(m, (0, 0, 0, 0, 0, pad)) # pad L (dim 1) msa_attention_mask = F.pad(msa_attention_mask, (0, 0, 0, pad)) # pad L - # Build the pair attention mask from the (padded) token mask; padded - # rows/cols are zero so they contribute nothing. - tok_mask = msa_attention_mask[:, :, 0].bool() - pair_mask = (tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)).to(m.dtype) - # Cast the heavy inputs to bf16 so the (bf16) dist_encoder runs bf16 end to # end; restore the caller's dtype on the output. orig_dtype = x_pair_dt.dtype @@ -158,7 +153,6 @@ def forward_sharded( m = m.to(torch.bfloat16) x_pair_dt = x_pair_dt.to(torch.bfloat16) msa_attention_mask = msa_attention_mask.to(torch.bfloat16) - pair_mask = pair_mask.to(torch.bfloat16) mesh = self.device_mesh m_dt = distribute_tensor( @@ -169,10 +163,17 @@ def forward_sharded( mesh, [Shard(0), Shard(1), Replicate()], ) - pair_mask_dt = distribute_tensor( - pair_mask.contiguous(), mesh, [Shard(0), Shard(1), Shard(2)] + # Pair attention mask built SHARDED from the (padded) token mask — the + # per-block outer product, never a full [B, L, L] (matches evolutionaryscale + # _cp_pair_mask). Padded rows/cols are zero, so they contribute nothing. + from transformers.models.esmfold2.distributed.model.layers.pair_init import ( + build_sharded_pair_mask, ) + pair_mask_dt = build_sharded_pair_mask( + msa_attention_mask[:, :, 0], mesh + ).to(m.dtype) + out_dt = self.dist_encoder(m_dt, x_pair_dt, msa_mask_dt, pair_mask_dt) if self._bf16: out_dt = out_dt.to(orig_dtype) @@ -305,6 +306,16 @@ def wrap_model_with_cp( model._cp_recycle_engine = CPRecycleEngine(model, dist_manager) + # Sharded pair-init builder: constructs z_init / rel_pos / token_bonds as + # 2-D-sharded DTensors so the full L×L pair is never resident on a rank + # during the init phase (the dominant per-rank peak that made CP OOM at + # short L). The model's forward delegates via the _cp_pair_init seam. + from transformers.models.esmfold2.distributed.model.layers.pair_init import ( + PairInitDistributed, + ) + + model._cp_pair_init = PairInitDistributed(model, dist_manager) + # Install the distributed LM->pair builder (#14) so lm_z is produced as a # sharded DTensor instead of a full L×L tensor on every rank (the binding # per-rank peak per the T0 phase sweep). The model's forward delegates via diff --git a/src/transformers/models/esmfold2/distributed/recycle.py b/src/transformers/models/esmfold2/distributed/recycle.py index 54f82ff4a8..c8842a01cf 100644 --- a/src/transformers/models/esmfold2/distributed/recycle.py +++ b/src/transformers/models/esmfold2/distributed/recycle.py @@ -96,10 +96,9 @@ def parcae_finish(self, model, z_dt: DTensor, pair_mask: torch.Tensor) -> DTenso (weight replicated, no token mixing); ``parcae_coda`` is a CP-wrapped ``FoldingTrunk`` driven via ``forward_sharded`` (no gather/re-distribute). ``z_dt`` is the 2-D-sharded (padded) pair from ``run_loop(gather=False)``; - ``pair_mask`` is the full (unpadded) token pair mask. Returns sharded z.""" + ``pair_mask`` is either the full (unpadded) token pair mask or an already + padded + sharded DTensor (from ``PairInitDistributed``). Returns sharded z.""" mesh = self.device_mesh - n = pair_mask.shape[-1] - pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor # parcae_readout: channel Linear on the local shard. z_local = z_dt.to_local() @@ -107,10 +106,15 @@ def parcae_finish(self, model, z_dt: DTensor, pair_mask: torch.Tensor) -> DTenso z_dt = DTensor.from_local(z_local.contiguous(), mesh, _PLACEMENTS) # parcae_coda: CP-wrapped trunk → forward_sharded (needs a padded, sharded mask). - pm = F.pad(pair_mask, (0, pad, 0, pad)) if pad else pair_mask - pair_mask_dt = distribute_tensor( - pm.to(z_dt.dtype).contiguous(), mesh, _PLACEMENTS - ) + if isinstance(pair_mask, DTensor): + pair_mask_dt = pair_mask.to(z_dt.dtype) + else: + n = pair_mask.shape[-1] + pad = (self.shard_factor - n % self.shard_factor) % self.shard_factor + pm = F.pad(pair_mask, (0, pad, 0, pad)) if pad else pair_mask + pair_mask_dt = distribute_tensor( + pm.to(z_dt.dtype).contiguous(), mesh, _PLACEMENTS + ) return model.parcae_coda.forward_sharded( z_dt, pair_attention_mask=pair_mask_dt ) @@ -165,9 +169,17 @@ def run_loop( tok_mask: torch.Tensor, total_steps: int, gather: bool = True, + n_orig: int | None = None, ): """Sharded equivalent of ``_run_one_loop``. + ``z`` / ``z_init`` may arrive either as full plain tensors (distributed + here, the serial-fallback contract) or as already-padded, already-sharded + ``(Shard0, Shard1, Shard2)`` DTensors built by ``PairInitDistributed`` — in + which case the full L×L pair is never resident on any rank and ``n_orig`` + (the pre-pad token count) must be supplied so the returned length is the + original, not the padded, count. + With ``gather=True`` (default) returns the gathered full tensor (drop-in for the serial loop). With ``gather=False`` returns ``(z_dt, n_orig)`` — the **sharded** ``(Shard0, Shard1, Shard2)`` pair @@ -193,7 +205,16 @@ def run_loop( mesh = self.device_mesh loop_dtype = z.dtype - N = z.shape[1] + # z / z_init may be pre-sharded DTensors (PairInitDistributed) or full + # tensors (serial fallback). For the DTensor case N is the ORIGINAL token + # count (passed as n_orig); z.shape[1] would be the padded length. + z_is_dt = isinstance(z, DTensor) + z_init_is_dt = isinstance(z_init, DTensor) + if z_is_dt: + assert n_orig is not None, "n_orig required when z is a pre-sharded DTensor" + N = n_orig + else: + N = z.shape[1] pad = (self.shard_factor - N % self.shard_factor) % self.shard_factor def _pad_pair(t: torch.Tensor) -> torch.Tensor: # (B, L, L, C) @@ -208,11 +229,22 @@ def _distribute_lm(lm_full: torch.Tensor) -> DTensor: ) # Distribute the persistent loop tensors ONCE (vs. per-boundary in serial). - z_dt = distribute_tensor(_pad_pair(z).contiguous(), mesh, _PLACEMENTS) - z_init_dt = distribute_tensor(_pad_pair(z_init).contiguous(), mesh, _PLACEMENTS) - pair_mask_dt = distribute_tensor( - _pad_mask(pair_mask).to(loop_dtype).contiguous(), mesh, _PLACEMENTS + # Pre-sharded DTensors (already padded to shard_factor) are used as-is — + # no full L×L is ever materialised on a rank. + z_dt = z if z_is_dt else distribute_tensor(_pad_pair(z).contiguous(), mesh, _PLACEMENTS) + z_init_dt = ( + z_init + if z_init_is_dt + else distribute_tensor(_pad_pair(z_init).contiguous(), mesh, _PLACEMENTS) ) + # pair_mask may arrive as a pre-sharded (padded) DTensor (PairInitDistributed) + # or a full tensor (serial fallback). + if isinstance(pair_mask, DTensor): + pair_mask_dt = pair_mask.to(loop_dtype) + else: + pair_mask_dt = distribute_tensor( + _pad_mask(pair_mask).to(loop_dtype).contiguous(), mesh, _PLACEMENTS + ) # lm_z may arrive as a sharded (Shard0,Shard1,Shard2) DTensor (#14 — # produced full-L×L-free by LanguageModelShimDistributed) or as a full diff --git a/src/transformers/models/esmfold2/distributed/structure_wrapper.py b/src/transformers/models/esmfold2/distributed/structure_wrapper.py index e217ef67e8..3324265379 100644 --- a/src/transformers/models/esmfold2/distributed/structure_wrapper.py +++ b/src/transformers/models/esmfold2/distributed/structure_wrapper.py @@ -144,19 +144,25 @@ def _pad_mask(x): # (B, L) # Step 1: conditioning -> full s, sharded (padded) z. Only materialise + # distribute the full z_trunk on the cache miss (first step); afterwards # the sharded z is reused from the cache. + # relative_position_encoding may arrive already padded + sharded (from + # PairInitDistributed) — use it directly instead of re-materialising the + # full L×L tensor to distribute it. + def _rp_to_dt(): + if isinstance(relative_position_encoding, DTensor): + return relative_position_encoding + return distribute_tensor( + _pad_pair(relative_position_encoding).contiguous(), mesh, _PAIR_PL + ) + z_cached = inference_cache is not None and "z_cp" in inference_cache if z_cached: zt_dt = rp_dt = None elif zt_is_dtensor: zt_dt = z_trunk # already padded + sharded; do not re-distribute - rp_dt = distribute_tensor( - _pad_pair(relative_position_encoding).contiguous(), mesh, _PAIR_PL - ) + rp_dt = _rp_to_dt() else: zt_dt = distribute_tensor(_pad_pair(z_trunk).contiguous(), mesh, _PAIR_PL) - rp_dt = distribute_tensor( - _pad_pair(relative_position_encoding).contiguous(), mesh, _PAIR_PL - ) + rp_dt = _rp_to_dt() s, z_dt = self.cond( t_hat=t, s_inputs=s_inputs, diff --git a/src/transformers/models/esmfold2/modeling_esmfold2.py b/src/transformers/models/esmfold2/modeling_esmfold2.py index 912b678090..b29f6fb48f 100644 --- a/src/transformers/models/esmfold2/modeling_esmfold2.py +++ b/src/transformers/models/esmfold2/modeling_esmfold2.py @@ -1025,19 +1025,41 @@ def forward( atom_to_token=atom_to_token, ) - z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( - x_inputs - ).unsqueeze(1) - - relative_position_encoding = self.rel_pos( - residue_index=residue_index, - asym_id=asym_id, - sym_id=sym_id, - entity_id=entity_id, - token_index=token_index, - ) - token_bonds_encoding = self.token_bonds(token_bonds.float()) - z_init = z_init + relative_position_encoding + token_bonds_encoding + # CP (sharded pair-init): build z_init / rel_pos / token_bonds directly + # as 2-D-sharded DTensors so the full L×L pair is never resident on any + # rank (the init-phase peak that made CP OOM at short L). Installed with + # the recycle engine; plain getattr keeps this file CP-agnostic. + _cp_pair_init = getattr(self, "_cp_pair_init", None) + if _cp_pair_init is not None: + ( + z_init, + relative_position_encoding, + token_bonds_encoding, + _pi_n_orig, + ) = _cp_pair_init( + x_inputs, + residue_index, + asym_id, + sym_id, + entity_id, + token_index, + token_bonds, + ) + else: + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2( + x_inputs + ).unsqueeze(1) + + relative_position_encoding = self.rel_pos( + residue_index=residue_index, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + token_index=token_index, + ) + token_bonds_encoding = self.token_bonds(token_bonds.float()) + z_init = z_init + relative_position_encoding + token_bonds_encoding + _pi_n_orig = z_init.shape[1] if ( lm_hidden_states is None @@ -1080,9 +1102,20 @@ def forward( self._esmc.to("cpu") torch.cuda.empty_cache() - pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() + # Pair attention mask: sharded block under CP (outer product of the + # sliced row/col token-mask, matching evolutionaryscale's _cp_pair_mask + # — never the full L×L), else the serial full outer product. + if _cp_pair_init is not None: + pair_mask = _cp_pair_init.pair_mask(tok_mask) + else: + pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() - z = self._init_pair_state(z_init) + # Initial pair state: sharded random block under CP (never full L×L), + # else the serial full draw. + if _cp_pair_init is not None: + z = _cp_pair_init.init_pair_state(z_init, _pi_n_orig) + else: + z = self._init_pair_state(z_init) a, b = self._discretized_dynamics() a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype) @@ -1114,7 +1147,9 @@ def forward( # ``getattr`` keeps this file CP-agnostic. _cp_engine = getattr(self, "_cp_recycle_engine", None) _z_dt = None - _n_orig = z.shape[1] + # z may be a sharded DTensor (padded) under CP, so use the pre-pad + # token count tracked at pair-init rather than z.shape[1]. + _n_orig = _pi_n_orig if _cp_engine is not None: _z_dt, _n_orig = _cp_engine.run_loop( self, @@ -1128,6 +1163,7 @@ def forward( tok_mask=tok_mask, total_steps=total_steps, gather=False, + n_orig=_pi_n_orig, ) del z_init, lm_z, _msa_inputs, a, b_mat _z_dt = _cp_engine.parcae_finish(self, _z_dt, pair_mask)