Batch StreamingStepsDataset internally instead of per-row collate

The dataset yielded one row at a time, forcing DataLoader's default
collate to Python-loop over every row to assemble each batch. That
loop scales with batch size and was pinning a CPU core at 100% while
the GPU sat idle. Now the dataset yields whole batches via vectorized
numpy slicing, used with DataLoader(batch_size=None).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

giant/cli.py

@@ -175,22 +175,26 @@ def train(
         files=files, split_events=train_events,
         pdg_map=pdg_map, mat_map=mat_map,
         cond_normalizer=cond_norm, target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle_buffer=shuffle_buffer, shuffle=True,
     )
     val_ds = StreamingStepsDataset(
         files=files, split_events=val_events,
         pdg_map=pdg_map, mat_map=mat_map,
         cond_normalizer=cond_norm, target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle=False,
     )
 
+    # Dataset yields whole batches already, so batch_size=None tells DataLoader
+    # to pass them through instead of re-collating row-by-row in Python.
     pin = _device.type == "cuda"
     train_loader = DataLoader(
-        train_ds, batch_size=t["batch_size"],
+        train_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )
     val_loader = DataLoader(
-        val_ds, batch_size=t["batch_size"],
+        val_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )
 

giant/data/dataset.py

@@ -71,6 +71,10 @@ class StreamingStepsDataset(IterableDataset):
 
     Never loads more than `shuffle_buffer` rows into RAM simultaneously.
     Files are split evenly across DataLoader workers via worker_info.
+
+    Yields whole batches (use with `DataLoader(..., batch_size=None)`)
+    rather than single rows, so the batch is assembled with vectorized
+    numpy slicing instead of a per-row Python loop in the default collate.
     """
 
     def __init__(
@@ -81,6 +85,7 @@ class StreamingStepsDataset(IterableDataset):
         mat_map: dict[int, int],
         cond_normalizer: Normalizer,
         target_normalizer: Normalizer,
+        batch_size: int,
         shuffle_buffer: int = 65536,
         shuffle: bool = True,
     ) -> None:
@@ -91,7 +96,9 @@ class StreamingStepsDataset(IterableDataset):
         self.mat_map = mat_map
         self.cond_normalizer = cond_normalizer
         self.target_normalizer = target_normalizer
-        self.shuffle_buffer = shuffle_buffer
+        self.batch_size = batch_size
+        # Buffer must hold at least one batch or we could never emit one.
+        self.shuffle_buffer = max(shuffle_buffer, batch_size)
         self.shuffle = shuffle
 
     def __iter__(self):
@@ -126,27 +133,44 @@ class StreamingStepsDataset(IterableDataset):
                 buf_tgt.append(target)
                 buf_n += len(cond_cont)
 
-                if not self.shuffle or buf_n >= self.shuffle_buffer:
-                    yield from self._flush(buf_cont, buf_cat, buf_tgt)
-                    buf_cont, buf_cat, buf_tgt, buf_n = [], [], [], 0
+                if buf_n >= self.shuffle_buffer:
+                    buf_cont, buf_cat, buf_tgt, buf_n = yield from self._flush(
+                        buf_cont, buf_cat, buf_tgt, final=False
+                    )
 
         if buf_n > 0:
-            yield from self._flush(buf_cont, buf_cat, buf_tgt)
+            yield from self._flush(buf_cont, buf_cat, buf_tgt, final=True)
 
     def _flush(
         self,
         buf_cont: list[np.ndarray],
         buf_cat: list[np.ndarray],
         buf_tgt: list[np.ndarray],
+        final: bool,
     ):
+        """Yield full batches of `batch_size`; carry any remainder back to the caller.
+
+        All batching is done via vectorized numpy slicing (no per-row Python loop).
+        """
         cont = np.concatenate(buf_cont)
         cat = np.concatenate(buf_cat)
         tgt = np.concatenate(buf_tgt)
         if self.shuffle:
             idx = np.random.permutation(len(cont))
             cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
-        t_cont = torch.from_numpy(cont).float()
-        t_cat = torch.from_numpy(cat).long()
-        t_tgt = torch.from_numpy(tgt).float()
-        for i in range(len(cont)):
-            yield t_cont[i], t_cat[i], t_tgt[i]
+
+        bs = self.batch_size
+        n = len(cont)
+        n_full = n // bs if not final else (n + bs - 1) // bs
+        for start in range(0, n_full * bs, bs):
+            end = min(start + bs, n)
+            yield (
+                torch.from_numpy(cont[start:end]).float(),
+                torch.from_numpy(cat[start:end]).long(),
+                torch.from_numpy(tgt[start:end]).float(),
+            )
+
+        if final:
+            return [], [], [], 0
+        rem = n_full * bs
+        return [cont[rem:]], [cat[rem:]], [tgt[rem:]], n - rem

scripts/train.py

@@ -154,6 +154,7 @@ def main() -> None:
         mat_map=mat_map,
         cond_normalizer=cond_norm,
         target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle_buffer=args.shuffle_buffer,
         shuffle=True,
     )
@@ -164,16 +165,19 @@ def main() -> None:
         mat_map=mat_map,
         cond_normalizer=cond_norm,
         target_normalizer=tgt_norm,
+        batch_size=t["batch_size"],
         shuffle=False,
     )
 
+    # Dataset yields whole batches already, so batch_size=None tells DataLoader
+    # to pass them through instead of re-collating row-by-row in Python.
     pin = device.type == "cuda"
     train_loader = DataLoader(
-        train_ds, batch_size=t["batch_size"],
+        train_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )
     val_loader = DataLoader(
-        val_ds, batch_size=t["batch_size"],
+        val_ds, batch_size=None,
         num_workers=t["num_workers"], pin_memory=pin,
     )