Align autoencoder splits and add multi-GPU training

train_exposure_autoencoder.py

@@ -50,9 +50,9 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--exposure_cache_dir", required=True)
     parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
+    parser.add_argument("--val_eid_file", default="ukb_val_eid.csv")
     parser.add_argument("--runs_root", default="runs")
     parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--val_fraction", type=float, default=0.05)
     parser.add_argument("--n_embd", type=int, default=120)
     parser.add_argument("--d_model", type=int, default=None)
     parser.add_argument("--n_layers", type=int, default=2)
@@ -72,35 +72,80 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num_workers", type=int, default=4)
     parser.add_argument("--device", default="cuda")
     parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument(
+        "--data_parallel",
+        action="store_true",
+        help="Use torch.nn.DataParallel across multiple CUDA devices.",
+    )
+    parser.add_argument(
+        "--gpu_ids",
+        default=None,
+        help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.",
+    )
     args = parser.parse_args()
-    if not 0.0 < args.val_fraction < 1.0:
-        parser.error("--val_fraction must be between 0 and 1")
     if not 0.0 <= args.mask_ratio < 1.0:
         parser.error("--mask_ratio must be in [0, 1)")
+    if args.gpu_ids:
+        try:
+            args.gpu_ids = [
+                int(part.strip())
+                for part in args.gpu_ids.split(",")
+                if part.strip()
+            ]
+        except ValueError as exc:
+            parser.error("--gpu_ids must be a comma-separated list of integers")
+        if not args.gpu_ids:
+            parser.error("--gpu_ids did not contain any valid CUDA device ids")
+        args.data_parallel = True
     return args
 
 
-def select_rows(
-    cache: ExposureCache, train_eids: set[int], val_fraction: float, seed: int
-) -> tuple[np.ndarray, np.ndarray]:
-    candidate_eids = np.asarray(
-        sorted(set(map(int, cache.eids)) & train_eids), dtype=np.int64
-    )
-    if len(candidate_eids) < 2:
-        raise ValueError("Need at least two training EIDs with cached exposure")
-    rng = np.random.default_rng(seed)
-    rng.shuffle(candidate_eids)
-    n_val = max(1, int(round(len(candidate_eids) * val_fraction)))
-    val_eids = candidate_eids[:n_val]
-    fit_eids = candidate_eids[n_val:]
+def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray:
     valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0
-    fit_event_rows = valid_row & np.isin(cache.eids, fit_eids)
-    val_event_rows = valid_row & np.isin(cache.eids, val_eids)
-    fit_rows = np.unique(np.asarray(cache.row_index[fit_event_rows], dtype=np.int64))
-    val_rows = np.unique(np.asarray(cache.row_index[val_event_rows], dtype=np.int64))
-    if len(fit_rows) == 0 or len(val_rows) == 0:
-        raise ValueError("Training/validation exposure rows are empty after filtering")
-    return fit_rows, val_rows
+    selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64))
+    rows = np.unique(
+        np.asarray(cache.row_index[selected_events], dtype=np.int64)
+    )
+    if len(rows) == 0:
+        raise ValueError(f"{split} exposure rows are empty after EID filtering")
+    return rows
+
+
+def maybe_wrap_data_parallel(
+    model: TimesNetExposureAutoencoder,
+    args: argparse.Namespace,
+    device: torch.device,
+    logger,
+):
+    if not args.data_parallel:
+        return model
+    if device.type != "cuda":
+        raise ValueError("--data_parallel requires --device cuda or cuda:<id>")
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        raise ValueError("--data_parallel requires at least two CUDA devices")
+    primary = (
+        int(device.index)
+        if device.index is not None
+        else int(torch.cuda.current_device())
+    )
+    device_ids = (
+        args.gpu_ids
+        if args.gpu_ids
+        else list(range(torch.cuda.device_count()))
+    )
+    device_ids = [primary, *[idx for idx in device_ids if idx != primary]]
+    if len(device_ids) < 2:
+        raise ValueError("--data_parallel needs at least two device ids")
+    if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids):
+        raise ValueError(f"CUDA device id is out of range: {device_ids}")
+    logger.info(f"Using DataParallel on CUDA devices: {device_ids}")
+    return torch.nn.DataParallel(
+        model, device_ids=device_ids, output_device=primary
+    )
+
+
+def unwrap_model(model) -> TimesNetExposureAutoencoder:
+    return model.module if isinstance(model, torch.nn.DataParallel) else model
 
 
 def channel_stats(
@@ -133,7 +178,7 @@ def masked_mse(
 
 
 def run_epoch(
-    model: TimesNetExposureAutoencoder,
+    model,
     loader: DataLoader,
     device: torch.device,
     stats: tuple[torch.Tensor, ...],
@@ -218,9 +263,12 @@ def main() -> None:
     )
     logger = setup_logging(run_dir)
     cache = ExposureCache(args.exposure_cache_dir)
-    train_rows, val_rows = select_rows(
-        cache, load_eid_file(args.train_eid_file), args.val_fraction, args.seed
-    )
+    train_eids = load_eid_file(args.train_eid_file)
+    val_eids = load_eid_file(args.val_eid_file)
+    if train_eids & val_eids:
+        raise ValueError("train and validation EID files must be disjoint")
+    train_rows = select_rows(cache, train_eids, "Training")
+    val_rows = select_rows(cache, val_eids, "Validation")
     raw_stats = channel_stats(cache, train_rows)
     stats = tuple(
         torch.as_tensor(value, device=device).view(1, 1, -1)
@@ -244,6 +292,7 @@ def main() -> None:
         conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio,
         dropout=args.dropout,
     ).to(device)
+    model = maybe_wrap_data_parallel(model, args, device, logger)
     optimizer = AdamW(
         model.parameters(), lr=args.base_lr,
         weight_decay=args.weight_decay, betas=(0.9, 0.95),
@@ -292,10 +341,11 @@ def main() -> None:
         if val_loss < best_loss:
             best_loss = val_loss
             stale_epochs = 0
+            checkpoint_model = unwrap_model(model)
             torch.save(
                 {
-                    "model_state_dict": model.state_dict(),
-                    "encoder_state_dict": model.encoder.state_dict(),
+                    "model_state_dict": checkpoint_model.state_dict(),
+                    "encoder_state_dict": checkpoint_model.encoder.state_dict(),
                     "model_config": {
                         key: config[key] for key in (
                             "n_embd", "d_model", "n_layers", "top_k",