Add alternative distributed segment usm

orb_models/common/models/gns.py

@@ -8,13 +8,19 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.utils.checkpoint import checkpoint
 
 from orb_models.common.atoms.batch.abstract_batch import AbstractAtomBatch
 from orb_models.common.atoms.batch.graph_batch import AtomGraphs
 from orb_models.common.models import base, segment_ops
 from orb_models.common.models.angular import UnitVector
 from orb_models.common.models.embedding import AtomEmbedding, AtomEmbeddingBag
-from orb_models.common.models.nn_util import build_mlp, get_cutoff, mlp_and_layer_norm
+from orb_models.common.models.nn_util import (
+    build_mlp,
+    chunked_apply,
+    get_cutoff,
+    mlp_and_layer_norm,
+)
 
 ConditioningType = Literal["additive", "concatenative", "none"]
 
@@ -76,19 +82,29 @@ def __init__(
         )
 
     def forward(
-        self, node_features: torch.Tensor, edge_features: torch.Tensor
+        self,
+        node_features: torch.Tensor,
+        edge_features: torch.Tensor,
+        chunk_size: int | None = None,
+        checkpoint_edge_block: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Forward pass to encode node and edge features.
 
         Args:
             node_features: Input node features tensor
             edge_features: Input edge features tensor
+            chunk_size: Max edges per chunk in the edge MLP, which is by far the largest
+                activation here. The node MLP is num_nodes-shaped and left whole.
+            checkpoint_edge_block: Recompute the edge MLP in the backward pass instead of
+                storing it. None means "checkpoint iff chunking".
 
         Returns:
             Tuple of (encoded_nodes, encoded_edges)
         """
         encoded_nodes = self._node_fn(node_features)
-        encoded_edges = self._edge_fn(edge_features)
+        encoded_edges = chunked_apply(
+            self._edge_fn, edge_features, chunk_size, checkpoint_chunks=checkpoint_edge_block
+        )
         return encoded_nodes, encoded_edges
 
 
@@ -180,6 +196,36 @@ def conditioning_type(self) -> tuple[ConditioningType, ConditioningType]:
         """The type of conditioning used by the interaction network."""
         return self._node_cond, self._edge_cond
 
+    def _edge_block(
+        self,
+        nodes: torch.Tensor,
+        edges: torch.Tensor,
+        senders: torch.Tensor,
+        receivers: torch.Tensor,
+        receive_attn: torch.Tensor,
+        send_attn: torch.Tensor,
+        segment_sum_impl: Callable,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """The message-passing step for one chunk of edges.
+
+        The gathers and their [num_edges, 3 * latent_dim] concatenation belong in here
+        alongside the MLP: they are its input, so leaving them outside would materialise
+        them at full width and cost more than half of what chunking can save. The
+        attention weights, by contrast, are [num_edges, 1] and stay outside.
+
+        Wrapped in a checkpoint, only `updated_edges` and the two
+        [num_nodes, latent_dim] partial sums survive into the backward pass.
+        """
+        edge_features = torch.cat([edges, nodes[senders], nodes[receivers]], dim=1)
+        updated_edges = self._edge_mlp(edge_features)
+
+        num_segments = nodes.shape[0]
+        return (
+            updated_edges,
+            segment_sum_impl(updated_edges * send_attn, senders, num_segments),
+            segment_sum_impl(updated_edges * receive_attn, receivers, num_segments),
+        )
+
     def forward(
         self,
         nodes: torch.Tensor,
@@ -189,6 +235,10 @@ def forward(
         cutoff: torch.Tensor,
         cond_nodes: torch.Tensor | None = None,
         cond_edges: torch.Tensor | None = None,
+        segment_sum_impl: Callable | None = segment_ops.segment_sum,
+        segment_softmax_impl: Callable | None = segment_ops.segment_softmax,
+        chunk_size: int | None = None,
+        checkpoint_edge_block: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Run interaction network forward pass.
 
@@ -200,7 +250,12 @@ def forward(
             cutoff: Edge cutoff values [num_edges, 1]
             cond_nodes: Optional conditioning for nodes
             cond_edges: Optional conditioning for edges
-
+            segment_sum_impl: Whether to use default or mesh-aware segment sum implementation
+            segment_softmax_impl: Whether to use default or mesh-aware segment softmax implementation
+            chunk_size: Max edges per chunk in the edge block. Bounds the transient that the
+                backward recompute holds; on its own it saves almost nothing.
+            checkpoint_edge_block: Recompute the edge block in the backward pass instead of
+                storing it. None means "checkpoint iff chunking".
         Returns:
             Tuple of (updated_nodes, updated_edges)
         """
@@ -220,13 +275,13 @@ def forward(
 
         if self._attention_gate == "softmax":
             num_segments = nodes.shape[0]
-            receive_attn = segment_ops.segment_softmax(
+            receive_attn = segment_softmax_impl(
                 self._receive_attn(edges),
                 receivers,
                 num_segments,
                 weights=cutoff if self._distance_cutoff else None,
             )
-            send_attn = segment_ops.segment_softmax(
+            send_attn = segment_softmax_impl(
                 self._send_attn(edges),
                 senders,
                 num_segments,
@@ -240,17 +295,42 @@ def forward(
             receive_attn = receive_attn * cutoff
             send_attn = send_attn * cutoff
 
-        sent_attributes = nodes[senders]
-        received_attributes = nodes[receivers]
-        edge_features = torch.cat([edges, sent_attributes, received_attributes], dim=1)
-        updated_edges = self._edge_mlp(edge_features)
+        num_edges = edges.shape[0]
+        chunk_size = min(chunk_size or num_edges, num_edges)
+        if checkpoint_edge_block is None:
+            checkpoint_edge_block = chunk_size < num_edges
+        use_checkpoint = checkpoint_edge_block and torch.is_grad_enabled()
+        # An edgeless graph (e.g. a padding graph) still runs the block once, on empties.
+        starts = range(0, num_edges, chunk_size) if num_edges else [0]
+
+        # The aggregations accumulate as we go, so each chunk's partial sums are freed
+        # once added; the updated edges are collected, as they are all needed at the end.
+        edge_chunks: list[torch.Tensor] = []
+        sent_attributes = edges.new_zeros(nodes.shape[0], self.latent_dim)
+        received_attributes = edges.new_zeros(nodes.shape[0], self.latent_dim)
+        for start in starts:
+            chunk = slice(start, start + chunk_size)
+            block_args = (
+                nodes,
+                edges[chunk],
+                senders[chunk],
+                receivers[chunk],
+                receive_attn[chunk],
+                send_attn[chunk],
+                segment_sum_impl,
+            )
+            if use_checkpoint:
+                chunk_edges, sent, received = checkpoint(
+                    self._edge_block, *block_args, use_reentrant=False
+                )
+            else:
+                chunk_edges, sent, received = self._edge_block(*block_args)
+            edge_chunks.append(chunk_edges)
+            sent_attributes = sent_attributes + sent
+            received_attributes = received_attributes + received
 
-        sent_attributes = segment_ops.segment_sum(
-            updated_edges * send_attn, senders, nodes.shape[0]
-        )
-        received_attributes = segment_ops.segment_sum(
-            updated_edges * receive_attn, receivers, nodes.shape[0]
-        )
+        # torch.cat always copies, so skip it on the single-chunk (default) path.
+        updated_edges = edge_chunks[0] if len(edge_chunks) == 1 else torch.cat(edge_chunks, dim=0)
 
         node_features = torch.cat([nodes, received_attributes, sent_attributes], dim=1)
         updated_nodes = self._node_mlp(node_features)
@@ -473,11 +553,23 @@ def __init__(
             activation=activation,
         )
 
-    def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
+    def forward(
+        self,
+        batch: AtomGraphs,
+        segment_sum_impl: Callable = segment_ops.segment_sum,
+        segment_softmax_impl: Callable = segment_ops.segment_softmax,
+        chunk_size: int | None = None,
+        checkpoint_edge_block: bool | None = None,
+    ) -> dict[str, torch.Tensor]:
         """Encode a graph using molecular GNS.
 
         Args:
             batch: Input molecular graph
+            segment_sum_impl: Whether to use default or mesh-aware segment sum implementation.
+            chunk_size: Max edges per chunk in the encoder's and the interaction networks'
+                edge paths. None disables chunking.
+            checkpoint_edge_block: Recompute those edge paths in the backward pass instead of
+                storing them. None means "checkpoint iff chunking"
 
         Returns:
             Dictionary containing node_features, edge_features, and predictions
@@ -491,7 +583,12 @@ def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
             cond_nodes, cond_edges = None, None
 
         # Encode
-        nodes, edges = self._encoder(node_features, edge_features)
+        nodes, edges = self._encoder(
+            node_features,
+            edge_features,
+            chunk_size=chunk_size,
+            checkpoint_edge_block=checkpoint_edge_block,
+        )
 
         # Process through interaction networks
         cutoff = get_cutoff(batch.edge_features["vectors"].norm(dim=-1))
@@ -504,6 +601,10 @@ def forward(self, batch: AtomGraphs) -> dict[str, torch.Tensor]:
                 cutoff,
                 cond_nodes=cond_nodes,
                 cond_edges=cond_edges,
+                segment_sum_impl=segment_sum_impl,
+                segment_softmax_impl=segment_softmax_impl,
+                chunk_size=chunk_size,
+                checkpoint_edge_block=checkpoint_edge_block,
             )
 
         # Decode

orb_models/common/models/nn_util.py

@@ -7,7 +7,51 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from torch.utils.checkpoint import checkpoint_sequential
+from torch.utils.checkpoint import checkpoint, checkpoint_sequential
+
+
+def chunked_apply(
+    fn: Callable[[torch.Tensor], torch.Tensor],
+    x: torch.Tensor,
+    chunk_size: int | None = None,
+    checkpoint_chunks: bool | None = None,
+) -> torch.Tensor:
+    """Apply a row-wise function to `x`, optionally in chunks and/or recomputed in backward.
+
+    Only valid for `fn` that acts independently on each row of `x` (an MLP, a norm over
+    the feature dim, elementwise ops) — anything that mixes rows, such as attention or a
+    segment reduction, will silently give wrong answers.
+
+    Args:
+        fn: Row-wise callable, e.g. an `nn.Sequential` MLP.
+        x: Input of shape [rows, features].
+        chunk_size: Max rows per chunk. None or >= rows means no chunking.
+        checkpoint_chunks: Recompute each chunk in the backward pass instead of storing it.
+            None (the default) means "checkpoint iff chunking".
+
+    Returns:
+        `fn(x)`, identical up to floating-point non-determinism in the backward reduction.
+    """
+    rows = x.shape[0]
+    # Under no_grad nothing is stored, so neither lever buys anything.
+    if not torch.is_grad_enabled():
+        return fn(x)
+
+    chunked = bool(chunk_size) and chunk_size < rows
+    if checkpoint_chunks is None:
+        checkpoint_chunks = chunked
+
+    if not chunked:
+        return checkpoint(fn, x, use_reentrant=False) if checkpoint_chunks else fn(x)
+
+    assert chunk_size is not None  # implied by `chunked`, but not visible to the checker
+    out = [
+        checkpoint(fn, x[start : start + chunk_size], use_reentrant=False)
+        if checkpoint_chunks
+        else fn(x[start : start + chunk_size])
+        for start in range(0, rows, chunk_size)
+    ]
+    return torch.cat(out, dim=0)
 
 
 class ChargeSpinEmbedding(nn.Module):

orb_models/common/models/segment_ops.py

@@ -1,4 +1,6 @@
 import torch
+from torch.distributed import _functional_collectives as funcol
+from torch.distributed.device_mesh import DeviceMesh
 
 TORCHINT = [torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8]
 
@@ -305,3 +307,94 @@ def split_prediction(pred: torch.Tensor, n_node: torch.Tensor):
         return torch.split(pred, n_node.cpu().tolist(), dim=0)
     else:
         raise ValueError(f"Unexpected length of prediction tensor: {len(pred)}")
+
+
+def distributed_segment_sum(
+    data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int, mesh: DeviceMesh
+) -> torch.Tensor:
+    """Sum local edge shards into replicated node attributes."""
+    local_sum = segment_sum(data, segment_ids, num_segments)
+    return funcol.all_reduce(local_sum, "sum", mesh)
+
+
+def _safe_log(x: torch.Tensor) -> torch.Tensor:
+    positive = x > 0
+    inputs = torch.where(positive, x, torch.ones_like(x))
+    return torch.where(positive, torch.log(inputs), -torch.inf)
+
+
+def _safe_logsumexp(log_terms: torch.Tensor) -> torch.Tensor:
+    maxes = log_terms.amax(dim=0)
+    maxes = torch.where(torch.isfinite(maxes), maxes, torch.zeros_like(maxes))
+    return maxes + _safe_log(torch.sum(torch.exp(log_terms - maxes), dim=0))
+
+
+def segment_softmax_inner(
+    data: torch.Tensor,
+    segment_ids: torch.Tensor,
+    num_segments: int,
+    weights: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Computes log-probabilities and the log denominator for a segment softmax.
+
+    Args:
+        inputs: The input tensor to normalise over segments.
+        segments: The segment indices tensor.
+        num_segments: The number of segments.
+        weights: Optional weights tensor to multiply unnormalised probabilities.
+    """
+    if weights is not None:
+        data = data + _safe_log(weights)
+
+    segment_maxes = segment_max(data, segment_ids, num_segments)
+    # A segment can be all -inf: no rows at all, or every row's weight zero. Its
+    # max is then -inf and `inputs - max` is NaN. Shifting by 0 instead is just as
+    # valid -- there is no magnitude to stabilise -- and stays finite.
+    segment_maxes = torch.where(
+        torch.isfinite(segment_maxes),
+        segment_maxes,
+        torch.zeros_like(segment_maxes),
+    )
+    shifted = data - segment_maxes[segment_ids]
+    log_sum = _safe_log(segment_sum(torch.exp(shifted), segment_ids, num_segments))
+    log_denominator = segment_maxes + log_sum
+    # `log_sum` is -inf exactly for those all -inf segments, and every `shifted` in
+    # one is -inf too. Subtracting 0 leaves the log-probability at -inf, i.e. p = 0,
+    # which is what `safe_division` returns for a zero denominator.
+    log_sum = torch.where(
+        torch.isfinite(log_sum), log_sum, torch.zeros_like(log_sum)
+    )
+    return shifted - log_sum[segment_ids], log_denominator
+
+
+def distributed_segment_softmax(
+    data: torch.Tensor,
+    segment_ids: torch.Tensor,
+    num_segments: int,
+    weights: torch.Tensor | None = None,
+    mesh: DeviceMesh | None = None,
+) -> torch.Tensor:
+    """Segment softmax over rows sharded across mesh, normalised globally.
+
+    Args:
+        inputs: The input tensor to normalise over segments.
+        segments: The segment indices tensor.
+        num_segments: The number of segments.
+        mesh: The device mesh over which rows are sharded.
+        weights: Optional weights tensor to multiply unnormalised probabilities.
+    """
+    log_probs, log_denominator = segment_softmax_inner(
+        data, segment_ids, num_segments, weights
+    )
+
+    gathered = funcol.all_gather_tensor(log_denominator, 0, mesh)
+    total = _safe_logsumexp(
+        gathered.reshape(mesh.size(), *log_denominator.shape)
+    )
+    # A segment with no weight on any rank is -inf on both sides of the
+    # subtraction. Its rows are already at -inf log-probability, so any finite
+    # shift leaves them at p = 0; taking 0 avoids the NaN.
+    rescale = torch.where(
+        torch.isfinite(total), log_denominator - total, torch.zeros_like(total)
+    )
+    return torch.exp(log_probs + rescale[segment_ids])