fix bug

Dockerfile

@@ -1,5 +1,5 @@
 # Use the BioNeMo Framework image as the base
-FROM nvcr.io/nvidia/clara/bionemo-framework:2.3 AS openfold-bionemo-image
+FROM nvcr.io/nvidia/clara/bionemo-framework:2.7.1 AS openfold-bionemo-image
 
 # Switch to root user to perform system-level operations
 USER root
@@ -27,24 +27,19 @@ WORKDIR /openfold2
 # Install uv
 RUN pip install uv
 
-# Uninstall existing Triton and install Triton 3.3.0
-# Install nvidia-ml-py (replacement for deprecated pynvml) for cuequivariance
-RUN pip uninstall -y triton || true && \
-    pip uninstall -y pynvml || true && \
-    pip install triton==3.3.0 && \
-    pip install nvidia-ml-py
+RUN pip install nvidia-ml-py
 
 # Install cuequivariance and its CUDA operations
-RUN pip install cuequivariance_torch==0.6.1 && \
-    pip install cuequivariance-ops-torch-cu12==0.6.1
+RUN pip install cuequivariance_torch==0.8.0 && \
+    pip install cuequivariance-ops-torch-cu12==0.8.0
 
 RUN uv pip install --upgrade pip && \
     uv pip install --no-cache-dir wheel setuptools && \
     uv pip install --no-cache-dir --no-build-isolation -e . && \
     uv pip install --no-cache-dir \
-        biopython==1.85 \
+        biopython==1.86 \
         mdtraj==1.11.0 \
-        modelcif==1.5 \
+        modelcif==1.6 \
         ml_collections==1.1.0 \
         bionemo-moco==0.0.2.2 \
         "huggingface-hub>=0.24.0,<1.0" \
@@ -78,7 +73,7 @@ ENV PEPTRON_AUTO_DOWNLOAD="true"
 
 EXPOSE 7860
 
-RUN chown 1000 /app
+RUN chown -R 1000:1000 /app
 USER 1000
 
 CMD ["python", "app.py"]

PepTron/esm2/model/attention.py

@@ -54,6 +54,7 @@ class ESM2TEDotProductAttention(TEDotProductAttention):
         k_channels: int | None = None,
         v_channels: int | None = None,
         cp_comm_type: str = "p2p",
+        model_comm_pgs=None,  # Added for BioNeMo 2.7+ / Megatron-Core compatibility
     ):
         """Initialize ESM2TEDotProductAttention."""
         self.config = config
@@ -168,6 +169,7 @@ class ESM2DotProductAttention(DotProductAttention):
         attn_mask_type: AttnMaskType,
         attention_type: str,
         attention_dropout: Optional[float] = None,
+        model_comm_pgs=None,  # Added for BioNeMo 2.7+ / Megatron-Core compatibility
     ) -> None:
         """Initializes the Attention class.
 
@@ -177,6 +179,7 @@ class ESM2DotProductAttention(DotProductAttention):
             attn_mask_type: The type of attention mask to be used.
             attention_type: The type of attention mechanism.
             attention_dropout: The dropout rate for attention weights. Defaults to None.
+            model_comm_pgs: Model communication process groups (for Megatron-Core compatibility).
         """
         super().__init__(
             config=config,

PepTron/esm2/model/model.py

@@ -446,8 +446,17 @@ class ESM2Model(MegatronBioBertModel):
             else:
                 for l_no, layer in enumerate(self.encoder.layers):
                     with self.encoder.offload_context:
-                        layer.use_cudagraph = True
-                        if (len(self.encoder.cuda_graphs) == 0) or (not self.encoder.training):
+                        # Check if use_cudagraph is supported (not available in newer Megatron-Core)
+                        if hasattr(layer, 'use_cudagraph'):
+                            layer.use_cudagraph = True
+                        # Check if cuda_graphs is supported (not available in newer Megatron-Core)
+                        has_cuda_graphs = hasattr(self.encoder, 'cuda_graphs')
+                        use_standard_forward = (
+                            not has_cuda_graphs or 
+                            (has_cuda_graphs and len(self.encoder.cuda_graphs) == 0) or 
+                            not self.encoder.training
+                        )
+                        if use_standard_forward:
                             hidden_states, context = layer(
                                 hidden_states=hidden_states,
                                 attention_mask=attention_mask,

PepTron/peptron/data/data.py

@@ -348,8 +348,44 @@ class OpenFoldDataset(torch.utils.data.Dataset):
 
 class OpenFoldBatchCollator:
     def __call__(self, prots):
-        stack_fn = lambda x: torch.stack(x, dim=0) if isinstance(x[0], torch.Tensor) else x
-        return dict_multimap(stack_fn, prots) 
+        """
+        Collate a list of OpenFold-style feature dicts into a single batch.
+        Many features are padded to a *per-example* fixed size (e.g. predict mode sets
+        crop_size = num_res). When batching multiple proteins, the per-example sizes
+        can differ (L=138 vs L=134), so naive torch.stack fails. Here we pad tensors
+        to the max shape observed in the batch (per key/leaf) and then stack.
+        """
+        def pad_and_stack(xs):
+            if not isinstance(xs[0], torch.Tensor):
+                return xs
+            # Fast path: already uniform shapes.
+            first_shape = tuple(xs[0].shape)
+            if all(tuple(x.shape) == first_shape for x in xs):
+                return torch.stack(xs, dim=0)
+            # Scalars are always stackable.
+            if xs[0].ndim == 0:
+                return torch.stack(xs, dim=0)
+            # Require consistent rank; OpenFold features should meet this.
+            nd = xs[0].ndim
+            if any(x.ndim != nd for x in xs):
+                raise RuntimeError(
+                    f"Cannot collate tensors with different ranks: {[tuple(x.shape) for x in xs]}"
+                )
+            # Pad each dimension to the maximum size in this batch for this leaf.
+            max_shape = [max(int(x.shape[d]) for x in xs) for d in range(nd)]
+            padded = []
+            for x in xs:
+                if list(x.shape) == max_shape:
+                    padded.append(x)
+                    continue
+                # Create an output tensor filled with zeros (safe default for OpenFold features)
+                # and copy the existing values into the top-left slice.
+                out = x.new_zeros(max_shape)
+                slices = tuple(slice(0, int(s)) for s in x.shape)
+                out[slices] = x
+                padded.append(out)
+            return torch.stack(padded, dim=0)
+        return dict_multimap(pad_and_stack, prots)
 
 
 def collate_fn(data_list):