feat: GPU float16 monkey-patches to fix TITAN OOM on large IMPACT slides

README-commands.md

@@ -247,6 +247,19 @@ The following models are currently supported,
 | MADELEINE      | MADELEINE_SLIDE | CONCH1_5 | 🔒 gated |
 | CHIEF          | CHIEF_SLIDE     | CTRANSPATH | local ckpt |
 
+`model_kwargs={...}` forwards extra constructor arguments to patch encoders, and
+`slide_model_kwargs={...}` forwards them to slide encoders. `TITAN_SLIDE` applies
+its GPU OOM patch by default (`patch_oom=true`), which also pins the validated
+TITAN revision. Disable it only when testing upstream behavior:
+
+```bash
+aggregate_slide_features \
+    patch_features_h5_path=patch_features.h5 \
+    output_h5_path=slide_features.h5 \
+    slide_model_type=TITAN_SLIDE \
+    slide_model_kwargs={patch_oom:false}
+```
+
 OpenCLIP is used by default, with the default model being [QuiltNet-B-16-PMB](https://huggingface.co/wisdomik/QuiltNet-B-16-PMB).  Use the `model_type` parameter to specify a different model.
 To use H-Optimus-0, for example,
 
@@ -536,4 +549,3 @@ Each input file `<stem>.<ext>` produces `output_dir/<stem>.tiff`. Pass
 | `downscale_by` | `1` | Integer downsample factor (e.g. `2` converts a 40× slide to 20×). |
 | `num_workers` | `1` | Parallel workers for batch mode (`0` = all CPUs). |
 | `bigtiff` | `false` | Write BigTIFF format (required for files > ~4 GB). |
-

README.md

@@ -280,6 +280,47 @@ For running Mussel at scale on a compute cluster, see
 pipeline that wraps the Mussel CLI tools and handles job scheduling, parallelism,
 and output management across large slide cohorts.
 
+## Development
+
+### Docker/Apptainer Containers with Flash Attention
+
+Mussel supports building Docker containers with **flash-attn 2.0** for accelerated attention in the CONCH1.5 patch encoder and the Prov-GigaPath slide encoder. Flash attention provides ~30-50% speedup on patch encoding (~20% overall TITAN pipeline improvement).
+
+**Building the flash-attn container:**
+
+```bash
+# Build Docker image with flash-attn backend
+make docker-build BACKEND=fastattn
+# Or manually:
+docker build --build-arg BACKEND=fastattn -t mussel:fastattn .
+
+# Convert to Apptainer SIF for HPC deployment
+make sif
+# Or manually:
+apptainer build --force mussel-fastattn.sif docker-daemon://mussel:fastattn
+```
+
+**Key details:**
+- The `[fastattn]` extra in `pyproject.toml` installs:
+  - torch 2.11.0+cu126 (CUDA 12.6)
+  - flash-attn 2.6.3 (custom manylinux_2_28 wheels for Rocky 8 compatibility)
+  - xformers 0.0.35
+- Flash attention is required for **CONCH1.5** (patch encoding) and **Prov-GigaPath** (slide encoding); both require CUDA compute capability ≥ 8.0 (A100, H100, etc.)
+- For V100 (compute 7.0) or CPU, the code automatically falls back to PyTorch SDPA
+- TITAN slide encoder uses `SDPBackend.EFFICIENT_ATTENTION` (Phase 1 optimization); flash-attn accelerates CONCH1.5 patch encoding and GigaPath slide encoding
+
+**Using with mussel-nf:**
+
+Copy the SIF to your mussel-nf repo and use the `apptainer_fastattn` profile:
+
+```bash
+cp mussel-fastattn.sif /path/to/mussel-nf/
+cd /path/to/mussel-nf
+nextflow run main.nf -profile cluster,slurm,apptainer_fastattn ...
+```
+
+The profile automatically uses the flash-attn container for `FEATURIZE` tasks.
+
 ## License
 This code is made available under the GPLv3 License and is available for non-commercial academic purposes.
 Forked from CLAM, © [Mahmood Lab](http://www.mahmoodlab.org).

mussel/cli/aggregate_slide_features.py

@@ -1,11 +1,11 @@
 import logging
-from dataclasses import dataclass
-from typing import List, Optional
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
 
 import hydra
 from hydra.conf import HelpConf, HydraConf
 from hydra.core.config_store import ConfigStore
-from omegaconf import MISSING
+from omegaconf import MISSING, DictConfig
 
 from mussel.models import ModelType
 from mussel.utils import aggregate_slide_features, resolve_aggregation_method
@@ -24,6 +24,9 @@ class AggregateSlideFeaturesConfig:
         aggregation_method (str): Method for aggregation: 'identity' (no aggregation), 'mean', 'max', or 'model'.
         slide_model_type (Optional[ModelType]): Type of slide encoder model (when aggregation_method="model").
         slide_model_path (Optional[str]): Path to slide encoder model weights.
+        slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the slide model constructor.
+            TITAN_SLIDE applies its OOM patch by default via patch_oom=True; pass
+            slide_model_kwargs={patch_oom:false} to disable that patch and revision pin.
         use_gpu (bool): Whether to use GPU for computation.
         gpu_device_id (Optional[int]): Specific GPU device ID to use.
         gpu_device_ids (Optional[List[int]]): List of GPU device IDs for multi-GPU inference.
@@ -37,12 +40,17 @@ class AggregateSlideFeaturesConfig:
     aggregation_method: str = "identity"
     slide_model_type: Optional[ModelType] = None
     slide_model_path: Optional[str] = None
+    slide_model_kwargs: Dict[str, Any] = field(default_factory=dict)
     use_gpu: bool = True
     gpu_device_id: Optional[int] = None
     gpu_device_ids: Optional[List[int]] = None
     ssl_verify: bool = True  # Whether to verify SSL certificates for remote operations
     embedding_precision: str = "float32"
 
+    def __post_init__(self):
+        if isinstance(self.slide_model_kwargs, DictConfig):
+            self.slide_model_kwargs = dict(self.slide_model_kwargs)
+
 
 desc_doc = """== ${hydra.help.app_name} ==
 
@@ -104,6 +112,13 @@ def main(cfg: AggregateSlideFeaturesConfig):
             patch_features_h5_path=patch_features.h5 \\
             output_h5_path=slide_features.h5 \\
             slide_model_type=TITAN_SLIDE
+
+        # Disable TITAN's default OOM monkey-patch/revision pin, e.g. to test upstream
+        aggregate_slide_features \\
+            patch_features_h5_path=patch_features.h5 \\
+            output_h5_path=slide_features.h5 \\
+            slide_model_type=TITAN_SLIDE \\
+            slide_model_kwargs={patch_oom:false}
     """
     logger.info("Starting slide feature aggregation")
     logger.info(f"Input: {cfg.patch_features_h5_path}")
@@ -131,6 +146,7 @@ def main(cfg: AggregateSlideFeaturesConfig):
         gpu_device_id=cfg.gpu_device_id,
         gpu_device_ids=cfg.gpu_device_ids,
         embedding_precision=cfg.embedding_precision,
+        slide_model_kwargs=cfg.slide_model_kwargs,
     )
 
     logger.info(f"Slide features saved to: {cfg.output_h5_path}")

mussel/cli/extract_features.py

@@ -1,27 +1,30 @@
 import logging
 import os
 import ssl
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 import h5py
 import hydra
 import numpy as np
 import torch
 from hydra.conf import HelpConf, HydraConf
 from hydra.core.config_store import ConfigStore
-from omegaconf import MISSING
+from omegaconf import MISSING, DictConfig
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from mussel.datasets import WholeSlideImageTileCoordDataset
 from mussel.models import ModelType, get_model_factory
-from mussel.utils import (aggregate_slide_features_batch,
-                          ensure_directory_exists,
-                          extract_patch_features_batch,
-                          get_slide_ids_from_paths, resolve_remote_paths,
-                          save_features)
+from mussel.utils import (
+    aggregate_slide_features_batch,
+    ensure_directory_exists,
+    extract_patch_features_batch,
+    get_slide_ids_from_paths,
+    resolve_remote_paths,
+    save_features,
+)
 from mussel.utils.feature_extract import _numpy_to_torch
 from mussel.utils.file import save_hdf5, save_torch_tensor
 from mussel.utils.ml import collate_features
@@ -76,6 +79,10 @@ class ExtractFeaturesConfig:
         aggregation_method (str): Aggregation method: identity (single-step), mean/max/model (two-step).
         slide_model_type (Optional[ModelType]): Type of slide encoder model (when aggregation_method="model").
         slide_model_path (Optional[str]): Path to slide encoder model weights.
+        model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the patch model constructor.
+        slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to the slide model constructor.
+            TITAN_SLIDE applies its OOM patch by default via patch_oom=True; pass
+            slide_model_kwargs={patch_oom:false} to disable that patch and revision pin.
         ssl_verify (bool): Whether to verify SSL certificates when downloading models or accessing remote resources (default: True).
 
     Processing Parameters:
@@ -115,6 +122,8 @@ class ExtractFeaturesConfig:
     aggregation_method: str = "identity"
     slide_model_type: Optional[ModelType] = None
     slide_model_path: Optional[str] = None
+    model_kwargs: Dict[str, Any] = field(default_factory=dict)
+    slide_model_kwargs: Dict[str, Any] = field(default_factory=dict)
     ssl_verify: bool = True  # Whether to verify SSL certificates for remote operations
     # Processing parameters
     batch_size: int = 64
@@ -125,6 +134,12 @@ class ExtractFeaturesConfig:
     is_test_run: bool = False
     embedding_precision: str = "float32"
 
+    def __post_init__(self):
+        if isinstance(self.model_kwargs, DictConfig):
+            self.model_kwargs = dict(self.model_kwargs)
+        if isinstance(self.slide_model_kwargs, DictConfig):
+            self.slide_model_kwargs = dict(self.slide_model_kwargs)
+
 
 desc_doc = """== ${hydra.help.app_name} ==
 
@@ -223,6 +238,8 @@ def _main_single(cfg: ExtractFeaturesConfig):
         aggregation_method=cfg.aggregation_method,
         slide_model_type=cfg.slide_model_type,
         slide_model_path=cfg.slide_model_path,
+        model_kwargs=cfg.model_kwargs,
+        slide_model_kwargs=cfg.slide_model_kwargs,
         embedding_precision=cfg.embedding_precision,
     )
 
@@ -292,6 +309,7 @@ def _main_batch(cfg: ExtractFeaturesConfig):
             pin_memory=True,
             is_test_run=cfg.is_test_run,
             embedding_precision=cfg.embedding_precision,
+            model_kwargs=cfg.model_kwargs,
         )
 
         # Aggregate to slide level using batch processing
@@ -310,6 +328,7 @@ def _main_batch(cfg: ExtractFeaturesConfig):
             gpu_device_id=cfg.gpu_device_id,
             gpu_device_ids=cfg.gpu_device_ids,
             slide_batch_size=cfg.slide_batch_size,
+            slide_model_kwargs=cfg.slide_model_kwargs,
         )
 
         # Defense-in-depth: verify at least one output file was created.
@@ -340,6 +359,7 @@ def _main_batch(cfg: ExtractFeaturesConfig):
             pin_memory=True,
             is_test_run=cfg.is_test_run,
             embedding_precision=cfg.embedding_precision,
+            model_kwargs=cfg.model_kwargs,
         )
 
         # Save as PT format for consistency

mussel/cli/tessellate_extract_features.py

@@ -153,6 +153,12 @@ class TessellateExtractFeaturesConfig:
                 * GIGAPATH_SLIDE requires GIGAPATH patch encoder
                 * TITAN_SLIDE requires CONCH1_5 patch encoder
             - The required patch encoder is automatically paired and run as needed
+        model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to patch model constructors.
+        slide_model_kwargs (Dict[str, Any]): Extra keyword arguments forwarded to slide model constructors.
+            TITAN_SLIDE applies GPU float16/expand-based OOM monkey-patches and a
+            validated revision pin by default (patch_oom=True). Pass
+            slide_model_kwargs={patch_oom:false} to disable that patch when testing
+            upstream TITAN behavior.
         model_dir (Optional[str]): Directory containing pre-downloaded models.
             - When specified, the system looks for model subdirectories named after each model type
             - For example: /mnt/batch_models/GIGAPATH_SLIDE, /mnt/batch_models/CONCH1_5
@@ -216,6 +222,8 @@ class TessellateExtractFeaturesConfig:
     prefilter_model_type: Optional[ModelType] = None  # No prefilter by default
     prefilter_model_path: Optional[str] = None  # Path to prefilter model file
     model_type: Any = None  # Can be ModelType or List[ModelType]
+    model_kwargs: Dict[str, Any] = field(default_factory=dict)
+    slide_model_kwargs: Dict[str, Any] = field(default_factory=dict)
     model_dir: Optional[str] = None  # Directory containing pre-downloaded models
     pre_download_models: bool = False  # Whether to pre-download models to model_dir
     # Single mode visualization
@@ -272,6 +280,10 @@ def __post_init__(self):
         # Convert DictConfig to regular dict for model_batch_sizes
         if isinstance(self.model_batch_sizes, DictConfig):
             self.model_batch_sizes = dict(self.model_batch_sizes)
+        if isinstance(self.model_kwargs, DictConfig):
+            self.model_kwargs = dict(self.model_kwargs)
+        if isinstance(self.slide_model_kwargs, DictConfig):
+            self.slide_model_kwargs = dict(self.slide_model_kwargs)
 
         # Convert DictConfig to regular dict for model_embedding_precision
         if isinstance(self.model_embedding_precision, DictConfig):
@@ -910,6 +922,7 @@ def _main_batch(
                 num_workers=cfg.num_workers,
                 pin_memory=True,
                 is_test_run=False,
+                model_kwargs=cfg.model_kwargs,
             )
         except Exception as e:
             logger.error(f"Error during batch feature extraction: {e}")
@@ -1079,6 +1092,7 @@ def _main_batch(
                     slide_batch_size=cfg.slide_batch_size,
                     max_slide_patches=cfg.max_slide_patches,
                     embedding_precision=_resolve_precision(cfg, cfg.slide_model_type),
+                    slide_model_kwargs=cfg.slide_model_kwargs,
                 )
             except Exception as e:
                 logger.error(f"Error during batch aggregation: {e}")
@@ -1163,6 +1177,7 @@ def _main_batch(
             num_workers=cfg.num_workers,
             pin_memory=True,
             is_test_run=False,
+            model_kwargs=cfg.model_kwargs,
         )
 
         # Save PT files and coords-only H5

mussel/cli/tessellate_extract_features_common.py

@@ -155,6 +155,7 @@ def _tessellate_and_filter(
             batch_size=cfg.batch_size,
             num_workers=cfg.num_workers,
             gpu_device_ids=cfg.gpu_device_ids,
+            model_kwargs=getattr(cfg, "model_kwargs", {}),
         )
 
         logger.info(f"Filtering features: {slide_path}")
@@ -317,6 +318,7 @@ def process_slide_tessellation_and_filtering(
             aggregation_method="identity",
             slide_model_type=None,
             slide_model_path=None,
+            model_kwargs=getattr(cfg, "model_kwargs", {}),
         )
         return {
             "intermediate_h5_path": intermediate_h5_path,
@@ -344,6 +346,8 @@ def process_slide_tessellation_and_filtering(
             aggregation_method=cfg.aggregation_method,
             slide_model_type=getattr(cfg, "slide_model_type", None),
             slide_model_path=slide_model_path,
+            model_kwargs=getattr(cfg, "model_kwargs", {}),
+            slide_model_kwargs=getattr(cfg, "slide_model_kwargs", {}),
         )
         return {
             "final_coords": final_coords,