Draft of dynamic batch size based on user set token limit

src/boltz/data/module/inferencev2.py

@@ -6,7 +6,7 @@
 import pytorch_lightning as pl
 import torch
 from torch import Tensor
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Sampler
 
 from boltz.data import const
 from boltz.data.crop.affinity import AffinityCropper
@@ -24,6 +24,112 @@
 )
 
 
+class SortedBatchSampler(Sampler):
+    """Yield batches of indices sorted by total residue count.
+
+    Sorting inputs by size before batching minimises the wasted compute from
+    padding shorter inputs to the length of the longest one in the batch.
+    """
+
+    def __init__(self, manifest: Manifest, batch_size: int) -> None:
+        self.batch_size = batch_size
+        # Sort indices by ascending total residue count so similar-length
+        # inputs end up in the same batch.
+        self.sorted_indices = sorted(
+            range(len(manifest.records)),
+            key=lambda i: sum(c.num_residues for c in manifest.records[i].chains),
+        )
+        # PyTorch Lightning inspects batch_sampler.sampler to detect shuffling.
+        self.sampler = self.sorted_indices
+
+    def __iter__(self):
+        for start in range(0, len(self.sorted_indices), self.batch_size):
+            yield self.sorted_indices[start : start + self.batch_size]
+
+    def __len__(self) -> int:
+        return (len(self.sorted_indices) + self.batch_size - 1) // self.batch_size
+
+
+class TokenBudgetBatchSampler(Sampler):
+    """Yield variable-size batches bounded by a token-budget.
+
+    Inputs are sorted ascending by residue count.  A greedy scan then packs
+    consecutive inputs into a batch while the constraint
+
+        batch_size * N_pad_bucketed^2 <= token_budget^2
+
+    is satisfied, where N_pad_bucketed is the largest input in the pending
+    batch rounded up to the nearest bucket_size.  Rounding to a fixed grid
+    of sequence lengths avoids triggering a new torch.compile kernel trace
+    for every distinct padding length.
+
+    Parameters
+    ----------
+    manifest : Manifest
+        Input manifest (read-only).
+    token_budget : int
+        Effective sequence-length budget.  A single input of this many
+        residues will always be processed at batch size 1.  Smaller inputs
+        are grouped into larger batches so that
+        batch_size * N_max^2 ~= token_budget^2.
+    bucket_size : int, optional
+        Round padded sequence lengths to this multiple to reduce the number
+        of distinct tensor shapes seen by torch.compile.  Default 64.
+    max_batch_size : int, optional
+        Hard cap on batch size regardless of budget.  Default 256.
+    """
+
+    def __init__(
+        self,
+        manifest: Manifest,
+        token_budget: int,
+        bucket_size: int = 64,
+        max_batch_size: int = 256,
+    ) -> None:
+        self.token_budget = token_budget
+        self.bucket_size = bucket_size
+        self.max_batch_size = max_batch_size
+
+        sizes = [
+            sum(c.num_residues for c in r.chains)
+            for r in manifest.records
+        ]
+        self.sorted_indices = sorted(range(len(sizes)), key=lambda i: sizes[i])
+        self.sizes = sizes
+        # PyTorch Lightning inspects batch_sampler.sampler for shuffling detection.
+        self.sampler = self.sorted_indices
+
+    def _bucket(self, n: int) -> int:
+        """Round n up to the nearest multiple of bucket_size."""
+        return ((n + self.bucket_size - 1) // self.bucket_size) * self.bucket_size
+
+    def __iter__(self):
+        batch: list = []
+        n_max_bucketed: int = 0
+
+        for idx in self.sorted_indices:
+            n = self._bucket(self.sizes[idx])
+            candidate_max = max(n_max_bucketed, n)
+            candidate_bs = len(batch) + 1
+
+            over_budget = candidate_bs * candidate_max ** 2 > self.token_budget ** 2
+            over_cap = candidate_bs > self.max_batch_size
+
+            if batch and (over_budget or over_cap):
+                yield batch
+                batch = [idx]
+                n_max_bucketed = n
+            else:
+                batch.append(idx)
+                n_max_bucketed = candidate_max
+
+        if batch:
+            yield batch
+
+    def __len__(self) -> int:
+        return sum(1 for _ in self.__iter__())
+
+
 def load_input(
     record: Record,
     target_dir: Path,
@@ -325,6 +431,7 @@ def __init__(
         mol_dir: Path,
         num_workers: int,
         batch_size: int = 1,
+        token_budget: Optional[int] = None,
         constraints_dir: Optional[Path] = None,
         template_dir: Optional[Path] = None,
         extra_mols_dir: Optional[Path] = None,
@@ -347,6 +454,13 @@ def __init__(
             The number of workers to use.
         batch_size : int
             The number of inputs to process in parallel per batch. Default is 1.
+            Ignored when token_budget is set.
+        token_budget : Optional[int]
+            If set, use TokenBudgetBatchSampler instead of SortedBatchSampler.
+            Inputs are batched dynamically so that batch_size * N_max^2 <=
+            token_budget^2.  A value equal to the largest expected input will
+            always yield batch_size=1 for that input and larger batches for
+            smaller inputs.  Overrides batch_size.
         constraints_dir : Optional[Path]
             The path to the constraints directory.
         template_dir : Optional[Path]
@@ -360,6 +474,7 @@ def __init__(
         super().__init__()
         self.num_workers = num_workers
         self.batch_size = batch_size
+        self.token_budget = token_budget
         self.manifest = manifest
         self.target_dir = target_dir
         self.msa_dir = msa_dir
@@ -390,14 +505,32 @@ def predict_dataloader(self) -> DataLoader:
             override_method=self.override_method,
             affinity=self.affinity,
         )
-        return DataLoader(
-            dataset,
-            batch_size=1,
-            num_workers=self.num_workers,
-            pin_memory=True,
-            shuffle=False,
-            collate_fn=collate,
-        )
+        if self.token_budget is not None:
+            sampler = TokenBudgetBatchSampler(self.manifest, self.token_budget)
+        elif self.batch_size > 1:
+            # Only sort by size when batching multiple inputs together — sorting
+            # changes the processing order, which shifts global RNG state and
+            # breaks reproducibility vs. the default sequential ordering.
+            sampler = SortedBatchSampler(self.manifest, self.batch_size)
+        else:
+            sampler = None  # let DataLoader use default sequential order
+        if sampler is not None:
+            return DataLoader(
+                dataset,
+                batch_sampler=sampler,
+                num_workers=self.num_workers,
+                pin_memory=True,
+                collate_fn=collate,
+            )
+        else:
+            return DataLoader(
+                dataset,
+                batch_size=1,
+                shuffle=False,
+                num_workers=self.num_workers,
+                pin_memory=True,
+                collate_fn=collate,
+            )
 
     def transfer_batch_to_device(
         self,

src/boltz/main.py

@@ -916,6 +916,17 @@ def cli() -> None:
     help="The number of inputs to process in parallel per batch. Default is 1.",
     default=1,
 )
+@click.option(
+    "--token_budget",
+    type=int,
+    help=(
+        "Dynamic batching budget (residues). When set, batch size is chosen "
+        "automatically per batch so that batch_size * N_max^2 <= token_budget^2. "
+        "Overrides --batch_size for boltz2. A good starting value is the token "
+        "count of your largest input."
+    ),
+    default=None,
+)
 @click.option(
     "--override",
     is_flag=True,
@@ -1065,6 +1076,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
     output_format: Literal["pdb", "mmcif"] = "mmcif",
     num_workers: int = 2,
     batch_size: int = 1,
+    token_budget: Optional[int] = None,
     override: bool = False,
     seed: Optional[int] = None,
     use_msa_server: bool = False,
@@ -1276,19 +1288,29 @@ def predict(  # noqa: C901, PLR0915, PLR0912
 
         # Create data module
         if model == "boltz2":
+            if token_budget is not None:
+                click.echo(
+                    f"Dynamic batching enabled: token_budget={token_budget}. "
+                    "--batch_size is ignored."
+                )
             data_module = Boltz2InferenceDataModule(
                 manifest=processed.manifest,
                 target_dir=processed.targets_dir,
                 msa_dir=processed.msa_dir,
                 mol_dir=mol_dir,
                 num_workers=num_workers,
                 batch_size=batch_size,
+                token_budget=token_budget,
                 constraints_dir=processed.constraints_dir,
                 template_dir=processed.template_dir,
                 extra_mols_dir=processed.extra_mols_dir,
                 override_method=method,
             )
         else:
+            if batch_size > 1:
+                click.echo(
+                    "Warning: --batch_size is ignored for boltz1 (always uses 1)."
+                )
             data_module = BoltzInferenceDataModule(
                 manifest=processed.manifest,
                 target_dir=processed.targets_dir,
@@ -1371,6 +1393,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             mol_dir=mol_dir,
             num_workers=num_workers,
             batch_size=batch_size,
+            token_budget=token_budget,
             constraints_dir=processed.constraints_dir,
             template_dir=processed.template_dir,
             extra_mols_dir=processed.extra_mols_dir,