Pin BatchData so DataLoader pin_memory takes effect

docs/train_folds.md

@@ -359,3 +359,4 @@ hound_train_folds --backend gcp --gcp_data_dir gs://<bucket>/<prefix> \
 
 - Cross-fold training requires ~N times more computation (N = number of folds)
 - This document uses `--backend local` because it will work for everyone, but jobs can also go to GCP Batch with `--backend gcp` (see [gcprunner](gcprunner.md)), or to Slurm with `--backend slurm`, which requires the non-public `slurmrunner` package and is the default when it is installed.
+- Set `train.num_workers` high enough to keep the GPU fed. Each worker reads and decompresses one example at a time, so long sequences need many: 16 for 786 kb on an H100 at batch 4.

src/baskerville/trainer.py

@@ -931,14 +931,14 @@ def _mlm_species_id(self, label, context):
 
     def _mlm_unpack(self, batch):
         """Pull sequence / species / masks off a BatchData and move to device."""
-        x = batch.sequence.to(self.device)
+        x = batch.sequence.to(self.device, non_blocking=True)
         label = batch.species_label
         exon_mask = batch.exon_mask
         repeat_mask = batch.repeat_mask
         if exon_mask is not None:
-            exon_mask = exon_mask.to(self.device)
+            exon_mask = exon_mask.to(self.device, non_blocking=True)
         if repeat_mask is not None:
-            repeat_mask = repeat_mask.to(self.device)
+            repeat_mask = repeat_mask.to(self.device, non_blocking=True)
         return x, label, exon_mask, repeat_mask
 
     def _optimizer_step(self, loss):
@@ -1239,17 +1239,17 @@ def _unpack_batch(self, batch_data) -> tuple:
             y = y.to(self.device).float()
             return x, y, None, None, None
 
-        # Handle BatchData format
-        x = batch_data.sequence.to(self.device)
+        # Handle BatchData format (pinned by the DataLoader, so copies are async)
+        x = batch_data.sequence.to(self.device, non_blocking=True)
         y = batch_data.coverage_targets
         if y is not None:
-            y = y.to(self.device).float()
+            y = y.to(self.device, non_blocking=True).float()
 
         yg = gene_presence = gene_out_mask = None
         if batch_data.has_gene:
-            yg = batch_data.gene_targets.to(self.device).float()
-            gene_presence = batch_data.gene_presence.to(self.device)
-            gene_out_mask = batch_data.gene_out_mask.to(self.device)
+            yg = batch_data.gene_targets.to(self.device, non_blocking=True).float()
+            gene_presence = batch_data.gene_presence.to(self.device, non_blocking=True)
+            gene_out_mask = batch_data.gene_out_mask.to(self.device, non_blocking=True)
 
         return x, y, yg, gene_presence, gene_out_mask
 

src/baskerville/types.py

@@ -1,6 +1,6 @@
 """Structured data types for coverage and gene prediction."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, fields, replace
 from typing import Optional
 
 import torch
@@ -74,6 +74,17 @@ def has_gene(self) -> bool:
         """Check if gene data is present."""
         return self.gene_targets is not None
 
+    def pin_memory(self) -> "BatchData":
+        """Page-lock every tensor; DataLoader(pin_memory=True) calls this per batch."""
+        return replace(
+            self,
+            **{
+                f.name: getattr(self, f.name).pin_memory()
+                for f in fields(self)
+                if getattr(self, f.name) is not None
+            },
+        )
+
     @staticmethod
     def collate(batch: list["BatchData"]) -> "BatchData":
         """Collate a list of BatchData objects into a batched BatchData.

src/tests/test_dataset.py

@@ -458,3 +458,18 @@ def test_dataloader_with_genes(self, mock_gene_data_dir):
         assert batch.gene_targets.dtype == torch.float16
         assert batch.gene_presence.dtype == torch.bool
         assert batch.gene_out_mask.dtype == torch.bool
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="pinning needs CUDA")
+def test_batchdata_pin_memory():
+    """DataLoader(pin_memory=True) pins every BatchData tensor, leaving Nones."""
+    items = [
+        BatchData(sequence=torch.zeros(4, 8), coverage_targets=torch.ones(2, 4))
+        for _ in range(3)
+    ]
+    loader = DataLoader(
+        items, batch_size=3, collate_fn=BatchData.collate, pin_memory=True
+    )
+    batch = next(iter(loader))
+    assert batch.sequence.is_pinned() and batch.coverage_targets.is_pinned()
+    assert batch.gene_targets is None