From a756c3dec4f23cad3b018e94ba3962637c3c1442 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Fri, 9 Oct 2026 10:04:39 -0700 Subject: [PATCH] Pin BatchData so DataLoader pin_memory takes effect DataLoader(pin_memory=True) pins only objects that expose pin_memory(), so BatchData batches were copied host->device synchronously from pageable memory. Add BatchData.pin_memory() and copy with non_blocking=True. Note in train_folds.md that num_workers must scale with example size. Co-Authored-By: Claude Opus 5.5 --- docs/train_folds.md | 1 + src/baskerville/trainer.py | 18 +++++++++--------- src/baskerville/types.py | 13 ++++++++++++- src/tests/test_dataset.py | 15 +++++++++++++++ 4 files changed, 37 insertions(+), 10 deletions(-) diff --git a/docs/train_folds.md b/docs/train_folds.md index 7060358..50e08c4 100644 --- a/docs/train_folds.md +++ b/docs/train_folds.md @@ -359,3 +359,4 @@ hound_train_folds --backend gcp --gcp_data_dir gs:/// \ - 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. diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py index 594a321..e451425 100644 --- a/src/baskerville/trainer.py +++ b/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 diff --git a/src/baskerville/types.py b/src/baskerville/types.py index 4ed0869..a83e8f3 100644 --- a/src/baskerville/types.py +++ b/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. diff --git a/src/tests/test_dataset.py b/src/tests/test_dataset.py index cdb2817..1d06668 100644 --- a/src/tests/test_dataset.py +++ b/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