Use DataLoader for burden readouts

compute_burden_index_landmarks.py

@@ -10,6 +10,8 @@ from typing import Any, Iterable
 import numpy as np
 import pandas as pd
 import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
 from tqdm.auto import tqdm
 
 from burden_index import (
@@ -294,117 +296,135 @@ def _config_split_indices(
     return make_eval_indices(_Sized(), args, cfg)
 
 
-def _iter_readout_batches(
-    n: int,
-    batch_size: int,
-) -> Iterable[slice]:
-    batch_size = max(1, int(batch_size))
-    for start in range(0, n, batch_size):
-        yield slice(start, min(start + batch_size, n))
+class ReadoutJobIterableDataset(IterableDataset):
+    def __init__(
+        self,
+        *,
+        rows: list[dict[str, Any]],
+        formed_mode: str,
+        horizon: float,
+    ) -> None:
+        super().__init__()
+        self.rows = rows
+        self.formed_mode = str(formed_mode)
+        self.horizon = float(horizon)
+        if self.formed_mode not in {"observed", "model_weighted"}:
+            raise ValueError(f"Unknown formed_mode={self.formed_mode!r}")
+
+    def __iter__(self) -> Iterable[dict[str, torch.Tensor]]:
+        worker = get_worker_info()
+        if worker is None:
+            start, step = 0, 1
+        else:
+            start, step = int(worker.id), int(worker.num_workers)
+
+        for row_idx in range(start, len(self.rows), step):
+            row = self.rows[row_idx]
+            if self.formed_mode == "model_weighted":
+                grid = _build_readout_grid(
+                    event_seq=row["event_seq"],
+                    time_seq=row["time_seq"],
+                    other_type=row["other_type"],
+                    other_time=row["other_time"],
+                    t_query=float(row["t_query"]),
+                )
+                if grid.size > 0:
+                    end_times = np.concatenate(
+                        [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)]
+                    )
+                    row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
+                    valid = row_deltas > 0
+                    for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()):
+                        yield _make_readout_job(row, row_idx, "formed", query_time, delta)
+            if self.horizon > 0:
+                yield _make_readout_job(
+                    row,
+                    row_idx,
+                    "future",
+                    float(row["t_query"]),
+                    self.horizon,
+                )
+
+
+def _make_readout_job(
+    row: dict[str, Any],
+    row_idx: int,
+    kind: str,
+    query_time: float,
+    delta: float,
+) -> dict[str, torch.Tensor]:
+    return {
+        "event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(),
+        "time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(),
+        "sex": torch.tensor(int(row["sex"]), dtype=torch.long),
+        "other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(),
+        "other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(),
+        "other_value_kind": torch.from_numpy(
+            np.asarray(row["other_value_kind"], dtype=np.int64)
+        ).long(),
+        "other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(),
+        "query_time": torch.tensor(float(query_time), dtype=torch.float32),
+        "delta": torch.tensor(float(delta), dtype=torch.float32),
+        "row_idx": torch.tensor(int(row_idx), dtype=torch.long),
+        "kind": torch.tensor(0 if kind == "formed" else 1, dtype=torch.long),
+    }
+
+
+def _collate_readout_jobs(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+    event_seq = pad_sequence(
+        [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    time_seq = pad_sequence(
+        [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_type = pad_sequence(
+        [x["other_type"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_value = pad_sequence(
+        [x["other_value"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_value_kind = pad_sequence(
+        [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_time = pad_sequence(
+        [x["other_time"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    return {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "sex": torch.stack([x["sex"] for x in batch]),
+        "other_type": other_type,
+        "other_value": other_value,
+        "other_value_kind": other_value_kind,
+        "other_time": other_time,
+        "query_time": torch.stack([x["query_time"] for x in batch]),
+        "delta": torch.stack([x["delta"] for x in batch]),
+        "row_idx": torch.stack([x["row_idx"] for x in batch]),
+        "kind": torch.stack([x["kind"] for x in batch]),
+    }
 
 
 @torch.inference_mode()
-def _query_hidden_jobs(
+def _query_hidden_readout_batch(
     *,
     ctx: Any,
-    jobs: list[tuple[dict[str, Any], float]],
+    batch: dict[str, torch.Tensor],
 ) -> torch.Tensor:
-    if not jobs:
-        return torch.empty(0, ctx.model.n_embd, device=ctx.device)
-
-    batch_size = len(jobs)
-    max_event_len = max(int(np.asarray(row["event_seq"]).size) for row, _ in jobs)
-    max_other_len = max(int(np.asarray(row["other_type"]).size) for row, _ in jobs)
-
-    event = np.full((batch_size, max_event_len), PAD_IDX, dtype=np.int64)
-    time = np.zeros((batch_size, max_event_len), dtype=np.float32)
-    other_type = np.zeros((batch_size, max_other_len), dtype=np.int64)
-    other_value = np.zeros((batch_size, max_other_len), dtype=np.float32)
-    other_value_kind = np.zeros((batch_size, max_other_len), dtype=np.int64)
-    other_time = np.zeros((batch_size, max_other_len), dtype=np.float32)
-    sex = np.zeros(batch_size, dtype=np.int64)
-    query_times = np.zeros(batch_size, dtype=np.float32)
-
-    for i, (row, query_time) in enumerate(jobs):
-        event_seq = np.asarray(row["event_seq"], dtype=np.int64)
-        time_seq = np.asarray(row["time_seq"], dtype=np.float32)
-        other_type_seq = np.asarray(row["other_type"], dtype=np.int64)
-        other_value_seq = np.asarray(row["other_value"], dtype=np.float32)
-        other_value_kind_seq = np.asarray(row["other_value_kind"], dtype=np.int64)
-        other_time_seq = np.asarray(row["other_time"], dtype=np.float32)
-
-        event[i, : event_seq.size] = event_seq
-        time[i, : time_seq.size] = time_seq
-        other_type[i, : other_type_seq.size] = other_type_seq
-        other_value[i, : other_value_seq.size] = other_value_seq
-        other_value_kind[i, : other_value_kind_seq.size] = other_value_kind_seq
-        other_time[i, : other_time_seq.size] = other_time_seq
-        sex[i] = int(row["sex"])
-        query_times[i] = np.float32(query_time)
-
-    event_t = torch.from_numpy(event).long().to(ctx.device)
+    event = batch["event_seq"].long().to(ctx.device, non_blocking=True)
     return ctx.model(
-        event_seq=event_t,
-        time_seq=torch.from_numpy(time).float().to(ctx.device),
-        sex=torch.from_numpy(sex).long().to(ctx.device),
-        padding_mask=event_t > PAD_IDX,
-        t_query=torch.from_numpy(query_times).float().to(ctx.device),
-        other_type=torch.from_numpy(other_type).long().to(ctx.device),
-        other_value=torch.from_numpy(other_value).float().to(ctx.device),
-        other_value_kind=torch.from_numpy(other_value_kind).long().to(ctx.device),
-        other_time=torch.from_numpy(other_time).float().to(ctx.device),
+        event_seq=event,
+        time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True),
+        sex=batch["sex"].long().to(ctx.device, non_blocking=True),
+        padding_mask=event > PAD_IDX,
+        t_query=batch["query_time"].float().to(ctx.device, non_blocking=True),
+        other_type=batch["other_type"].long().to(ctx.device, non_blocking=True),
+        other_value=batch["other_value"].float().to(ctx.device, non_blocking=True),
+        other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True),
+        other_time=batch["other_time"].float().to(ctx.device, non_blocking=True),
         target_mode="all_future",
     )
 
-
-def _build_readout_table(
-    *,
-    rows: list[dict[str, Any]],
-    formed_mode: str,
-    horizon: float,
-) -> dict[str, Any]:
-    jobs: list[tuple[dict[str, Any], float]] = []
-    row_indices: list[int] = []
-    kinds: list[str] = []
-    deltas: list[float] = []
-
-    if formed_mode not in {"observed", "model_weighted"}:
-        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
-
-    for row_idx, row in enumerate(rows):
-        if formed_mode == "model_weighted":
-            grid = _build_readout_grid(
-                event_seq=row["event_seq"],
-                time_seq=row["time_seq"],
-                other_type=row["other_type"],
-                other_time=row["other_time"],
-                t_query=float(row["t_query"]),
-            )
-            if grid.size > 0:
-                end_times = np.concatenate(
-                    [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)]
-                )
-                row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
-                valid = row_deltas > 0
-                for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()):
-                    jobs.append((row, float(query_time)))
-                    row_indices.append(row_idx)
-                    kinds.append("formed")
-                    deltas.append(float(delta))
-        if horizon > 0:
-            jobs.append((row, float(row["t_query"])))
-            row_indices.append(row_idx)
-            kinds.append("future")
-            deltas.append(float(horizon))
-
-    return {
-        "jobs": jobs,
-        "row_indices": np.asarray(row_indices, dtype=np.int64),
-        "kinds": np.asarray(kinds, dtype=object),
-        "deltas": np.asarray(deltas, dtype=np.float32),
-    }
-
-
 @torch.inference_mode()
 def _probabilities_from_hidden_torch(
     *,
@@ -442,30 +462,21 @@ def _probabilities_from_hidden_torch(
 
     return -torch.expm1(-rate * exposure)
 
-
 @torch.inference_mode()
-def _readout_probabilities(
+def _readout_probabilities_from_batch(
     *,
     ctx: Any,
-    readout_table: dict[str, Any],
+    batch: dict[str, torch.Tensor],
     union_disease_ids: np.ndarray,
-    readout_batch_size: int,
 ) -> torch.Tensor:
-    jobs = readout_table["jobs"]
-    if not jobs:
-        return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device)
-
-    out = torch.empty((len(jobs), union_disease_ids.size), dtype=torch.float32, device=ctx.device)
-    deltas = np.asarray(readout_table["deltas"], dtype=np.float32)
-    for slc in _iter_readout_batches(len(jobs), readout_batch_size):
-        hidden = _query_hidden_jobs(ctx=ctx, jobs=jobs[slc])
-        out[slc] = _probabilities_from_hidden_torch(
-            ctx=ctx,
-            hidden=hidden,
-            disease_ids=union_disease_ids,
-            deltas=deltas[slc],
-        ).to(dtype=out.dtype)
-    return out
+    hidden = _query_hidden_readout_batch(ctx=ctx, batch=batch)
+    deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False)
+    return _probabilities_from_hidden_torch(
+        ctx=ctx,
+        hidden=hidden,
+        disease_ids=union_disease_ids,
+        deltas=deltas,
+    ).to(dtype=torch.float32)
 
 
 def _observed_formed_for_rows(
@@ -484,94 +495,58 @@ def _observed_formed_for_rows(
     return formed
 
 
-def _reduce_readout_table_to_bi_rows(
+def _apply_readout_batch_to_accumulators(
+    *,
+    batch: dict[str, torch.Tensor],
+    readout_prob: torch.Tensor,
+    survival_by_row: torch.Tensor | None,
+    future_prob_by_row: torch.Tensor,
+    ctx: Any,
+) -> None:
+    if readout_prob.numel() == 0:
+        return
+
+    row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True)
+    kinds = batch["kind"].long().to(ctx.device, non_blocking=True)
+    kind_is_formed = kinds == 0
+    kind_is_future = kinds == 1
+
+    if survival_by_row is not None and bool(kind_is_formed.any().item()):
+        formed_rows = row_indices[kind_is_formed]
+        formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0)
+        if hasattr(survival_by_row, "scatter_reduce_"):
+            survival_by_row.scatter_reduce_(
+                dim=0,
+                index=formed_rows[:, None].expand_as(formed_survival),
+                src=formed_survival,
+                reduce="prod",
+                include_self=True,
+            )
+        else:
+            for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist():
+                survival_by_row[int(row_indices[job_idx].item())] *= (
+                    1.0 - readout_prob[job_idx].clamp(0.0, 1.0)
+                )
+    if bool(kind_is_future.any().item()):
+        future_rows = row_indices[kind_is_future]
+        future_prob_by_row[future_rows] = readout_prob[kind_is_future]
+
+
+def _project_bi_rows(
     *,
     rows: list[dict[str, Any]],
     horizon: float,
     matrices: list[dict[str, Any]],
-    union_disease_ids: np.ndarray,
     formed_mode: str,
-    readout_table: dict[str, Any],
-    readout_prob: torch.Tensor,
+    formed_by_row: torch.Tensor,
+    future_prob_by_row: torch.Tensor,
     ctx: Any,
 ) -> list[dict[str, Any]]:
-    if formed_mode == "observed":
-        formed_by_row = torch.as_tensor(
-            _observed_formed_for_rows(
-                rows=rows,
-                union_disease_ids=union_disease_ids,
-            ),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    elif formed_mode == "model_weighted":
-        formed_by_row = torch.zeros(
-            (len(rows), union_disease_ids.size),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    else:
-        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
-
-    future_prob_by_row = torch.zeros(
-        (len(rows), union_disease_ids.size),
-        dtype=readout_prob.dtype,
-        device=ctx.device,
-    )
-    row_indices = torch.as_tensor(
-        np.asarray(readout_table["row_indices"], dtype=np.int64),
-        dtype=torch.long,
-        device=ctx.device,
-    )
-    kinds = np.asarray(readout_table["kinds"], dtype=object)
-    if formed_mode == "model_weighted":
-        survival_by_row = torch.ones(
-            (len(rows), union_disease_ids.size),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    else:
-        survival_by_row = None
-
-    if readout_prob.numel() > 0:
-        kind_is_formed = torch.as_tensor(
-            np.asarray(kinds == "formed", dtype=np.bool_),
-            dtype=torch.bool,
-            device=ctx.device,
-        )
-        kind_is_future = torch.as_tensor(
-            np.asarray(kinds == "future", dtype=np.bool_),
-            dtype=torch.bool,
-            device=ctx.device,
-        )
-        if survival_by_row is not None and bool(kind_is_formed.any().item()):
-            formed_rows = row_indices[kind_is_formed]
-            formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0)
-            if hasattr(survival_by_row, "scatter_reduce_"):
-                survival_by_row.scatter_reduce_(
-                    dim=0,
-                    index=formed_rows[:, None].expand_as(formed_survival),
-                    src=formed_survival,
-                    reduce="prod",
-                    include_self=True,
-                )
-            else:
-                for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist():
-                    survival_by_row[int(row_indices[job_idx].item())] *= (
-                        1.0 - readout_prob[job_idx].clamp(0.0, 1.0)
-                    )
-        if bool(kind_is_future.any().item()):
-            future_rows = row_indices[kind_is_future]
-            future_prob_by_row[future_rows] = readout_prob[kind_is_future]
-
-    if survival_by_row is not None:
-        formed_by_row = 1.0 - survival_by_row
-
     disease_future_by_row = (1.0 - formed_by_row) * future_prob_by_row
     disease_total_by_row = formed_by_row + disease_future_by_row
     projected: list[dict[str, Any]] = []
     for matrix in matrices:
-        A = torch.as_tensor(matrix["A_union"], dtype=readout_prob.dtype, device=ctx.device)
+        A = torch.as_tensor(matrix["A_union"], dtype=formed_by_row.dtype, device=ctx.device)
         projected.append(
             {
                 "matrix": matrix,
@@ -611,7 +586,7 @@ def _reduce_readout_table_to_bi_rows(
     return out
 
 
-def _compute_bi_from_readout_table(
+def _compute_bi_from_streamed_readouts(
     *,
     rows: list[dict[str, Any]],
     horizon: float,
@@ -619,46 +594,117 @@ def _compute_bi_from_readout_table(
     union_disease_ids: np.ndarray,
     formed_mode: str,
     readout_batch_size: int,
+    num_workers: int,
     ctx: Any,
+    log_prefix: str | None = None,
 ) -> tuple[list[dict[str, Any]], int, dict[str, float]]:
     horizon = float(horizon)
     if horizon < 0:
         raise ValueError(f"horizon must be non-negative, got {horizon}")
-    t0 = time.perf_counter()
-    readout_table = _build_readout_table(
+
+    dtype = torch.float32
+    if formed_mode == "observed":
+        formed_by_row = torch.as_tensor(
+            _observed_formed_for_rows(
+                rows=rows,
+                union_disease_ids=union_disease_ids,
+            ),
+            dtype=dtype,
+            device=ctx.device,
+        )
+        survival_by_row = None
+    elif formed_mode == "model_weighted":
+        survival_by_row = torch.ones(
+            (len(rows), union_disease_ids.size),
+            dtype=dtype,
+            device=ctx.device,
+        )
+        formed_by_row = None
+    else:
+        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
+    future_prob_by_row = torch.zeros(
+        (len(rows), union_disease_ids.size),
+        dtype=dtype,
+        device=ctx.device,
+    )
+
+    readout_jobs = 0
+    n_batches = 0
+    build_readout_sec = 0.0
+    forward_sec = 0.0
+    reduce_sec = 0.0
+    t_loader0 = time.perf_counter()
+    readout_dataset = ReadoutJobIterableDataset(
         rows=rows,
         formed_mode=formed_mode,
         horizon=horizon,
     )
-    t1 = time.perf_counter()
-    readout_prob = _readout_probabilities(
-        ctx=ctx,
-        readout_table=readout_table,
-        union_disease_ids=union_disease_ids,
-        readout_batch_size=readout_batch_size,
+    readout_loader = DataLoader(
+        readout_dataset,
+        batch_size=max(1, int(readout_batch_size)),
+        collate_fn=_collate_readout_jobs,
+        num_workers=max(0, int(num_workers)),
+        pin_memory=ctx.device.type == "cuda",
+        persistent_workers=int(num_workers) > 0,
+        prefetch_factor=2 if int(num_workers) > 0 else None,
     )
-    if ctx.device.type == "cuda":
-        torch.cuda.synchronize(ctx.device)
-    t2 = time.perf_counter()
-    rows_out = _reduce_readout_table_to_bi_rows(
+    build_readout_sec += time.perf_counter() - t_loader0
+
+    for batch in readout_loader:
+        t_forward0 = time.perf_counter()
+        readout_prob = _readout_probabilities_from_batch(
+            ctx=ctx,
+            batch=batch,
+            union_disease_ids=union_disease_ids,
+        )
+        if ctx.device.type == "cuda":
+            torch.cuda.synchronize(ctx.device)
+        forward_sec += time.perf_counter() - t_forward0
+
+        t_reduce0 = time.perf_counter()
+        _apply_readout_batch_to_accumulators(
+            batch=batch,
+            readout_prob=readout_prob,
+            survival_by_row=survival_by_row,
+            future_prob_by_row=future_prob_by_row,
+            ctx=ctx,
+        )
+        if ctx.device.type == "cuda":
+            torch.cuda.synchronize(ctx.device)
+        reduce_sec += time.perf_counter() - t_reduce0
+
+        n_batches += 1
+        readout_jobs += int(batch["row_idx"].numel())
+        if log_prefix and (n_batches == 1 or n_batches % 50 == 0):
+            print(
+                f"{log_prefix} processed {readout_jobs} readout jobs "
+                f"in {n_batches} batches",
+                flush=True,
+            )
+
+    if survival_by_row is not None:
+        formed_by_row = 1.0 - survival_by_row
+    assert formed_by_row is not None
+
+    t_project0 = time.perf_counter()
+    rows_out = _project_bi_rows(
         rows=rows,
         horizon=horizon,
         matrices=matrices,
-        union_disease_ids=union_disease_ids,
         formed_mode=formed_mode,
-        readout_table=readout_table,
-        readout_prob=readout_prob,
+        formed_by_row=formed_by_row,
+        future_prob_by_row=future_prob_by_row,
         ctx=ctx,
     )
     if ctx.device.type == "cuda":
         torch.cuda.synchronize(ctx.device)
-    t3 = time.perf_counter()
+    reduce_sec += time.perf_counter() - t_project0
     timings = {
-        "build_readout_sec": t1 - t0,
-        "forward_sec": t2 - t1,
-        "reduce_sec": t3 - t2,
+        "build_readout_sec": build_readout_sec,
+        "forward_sec": forward_sec,
+        "reduce_sec": reduce_sec,
     }
-    return rows_out, len(readout_table["jobs"]), timings
+    return rows_out, readout_jobs, timings
 
 
 def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
@@ -681,14 +727,16 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
         f"{time.perf_counter() - materialize_start:.2f}s",
         flush=True,
     )
-    out, readout_jobs, timings = _compute_bi_from_readout_table(
+    out, readout_jobs, timings = _compute_bi_from_streamed_readouts(
         rows=rows,
         horizon=payload["horizon"],
         matrices=payload["matrices"],
         union_disease_ids=payload["union_disease_ids"],
         formed_mode=payload["formed_mode"],
         readout_batch_size=int(payload["readout_batch_size"]),
+        num_workers=int(payload["num_workers"]),
         ctx=ctx,
+        log_prefix=f"[BI worker {device}]",
     )
     print(
         f"[BI worker {device}] done: readout_jobs={readout_jobs}, "
@@ -831,6 +879,12 @@ def main() -> None:
             "Increase this to improve GPU utilization if memory allows."
         ),
     )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="DataLoader workers per GPU process for readout job generation.",
+    )
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument(
         "--devices",
@@ -933,14 +987,16 @@ def main() -> None:
             flush=True,
         )
     if len(row_chunks) == 1:
-        all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table(
+        all_rows, total_readout_jobs, timings = _compute_bi_from_streamed_readouts(
             rows=rows,
             horizon=horizon,
             matrices=matrices,
             union_disease_ids=union_disease_ids,
             formed_mode=args.formed_mode,
             readout_batch_size=int(args.readout_batch_size),
+            num_workers=int(args.num_workers),
             ctx=ctx,
+            log_prefix="[BI main]",
         )
         for key, value in timings.items():
             total_timings[key] += float(value)
@@ -957,6 +1013,7 @@ def main() -> None:
                 "matrices": matrices,
                 "union_disease_ids": union_disease_ids,
                 "readout_batch_size": int(args.readout_batch_size),
+                "num_workers": int(args.num_workers),
                 "formed_mode": args.formed_mode,
             }
             for device, chunk_rows in row_chunks
@@ -988,6 +1045,7 @@ def main() -> None:
     print(f"Landmark rows: {len(rows)}")
     print(f"Readout jobs: {total_readout_jobs}")
     print(f"Readout batch size per worker: {int(args.readout_batch_size)}")
+    print(f"DataLoader workers per GPU process: {int(args.num_workers)}")
     print(f"Multiprocessing start method: {args.mp_start_method}")
     print(
         "Timing seconds: "