Support multi-GPU next-step training

train_next_step.py

@@ -174,6 +174,17 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument(
+        "--data_parallel",
+        action="store_true",
+        help="Use torch.nn.DataParallel across multiple CUDA devices.",
+    )
+    parser.add_argument(
+        "--gpu_ids",
+        type=str,
+        default=None,
+        help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.",
+    )
     parser.add_argument("--progress_interval", type=int, default=20)
 
     args = parser.parse_args()
@@ -192,6 +203,14 @@ def parse_args() -> argparse.Namespace:
         args.include_no_event_in_uts_target = True
     else:
         args.readout_name = args.readout_name or "token"
+    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:
+            raise ValueError("--gpu_ids must be a comma-separated list of integers") from exc
+        if not args.gpu_ids:
+            raise ValueError("--gpu_ids did not contain any valid CUDA device ids")
+        args.data_parallel = True
     return args
 
 
@@ -213,6 +232,41 @@ def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -
     }
 
 
+def _cuda_device_index(device: torch.device) -> int:
+    if device.type != "cuda":
+        raise ValueError("CUDA device is required for multi-GPU training")
+    if device.index is not None:
+        return int(device.index)
+    current = torch.cuda.current_device()
+    return int(current)
+
+
+def unwrap_model(model):
+    return model.module if isinstance(model, torch.nn.DataParallel) else model
+
+
+def maybe_wrap_data_parallel(
+    model: DeepHealth,
+    args: argparse.Namespace,
+    device: torch.device,
+    logger: logging.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 = _cuda_device_index(device)
+    device_ids = args.gpu_ids if args.gpu_ids else list(range(torch.cuda.device_count()))
+    if primary not in device_ids:
+        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")
+    logger.info(f"Using DataParallel on CUDA devices: {device_ids}")
+    return torch.nn.DataParallel(model, device_ids=device_ids, output_device=primary)
+
+
 def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
@@ -303,14 +357,19 @@ def compute_next_step_loss(
         "sex": batch["sex"],
         "padding_mask": batch["padding_mask"],
         "target_mode": "next_token",
-        "return_output": True,
     }
     if "exposure_daily" in batch:
         model_kwargs["exposure_daily"] = batch["exposure_daily"]
         model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
-    model_out = model(**model_kwargs)
-    if not isinstance(model_out, DeepHealthOutput):
-        raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
+    hidden = model(**model_kwargs)
+    if not isinstance(hidden, torch.Tensor):
+        raise TypeError("DeepHealth forward must return a hidden-state tensor")
+    model_out = DeepHealthOutput(
+        hidden=hidden,
+        time_seq=batch["time_seq"][:, : hidden.size(1)],
+        padding_mask=batch["padding_mask"][:, : hidden.size(1)],
+        event_len=int(hidden.size(1)),
+    )
     targets = build_augmented_next_step_targets(
         batch_cpu=batch_cpu,
         model_out=model_out,
@@ -324,7 +383,7 @@ def compute_next_step_loss(
         if args.readout_name == "same_time_group_end"
         else None,
     )
-    logits = model.calc_risk(readout_out.hidden)
+    logits = unwrap_model(model).calc_risk(readout_out.hidden)
 
     if args.target_mode == "delphi2m":
         loss, parts = criterion(
@@ -438,6 +497,8 @@ def build_metadata(
         "num_workers": int(args.num_workers),
         "prefetch_factor": int(args.prefetch_factor),
         "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
+        "data_parallel": bool(args.data_parallel),
+        "gpu_ids": args.gpu_ids,
         "split_sizes": {
             "train": int(len(train_subset)),
             "val": int(len(val_subset)),
@@ -557,6 +618,7 @@ def main() -> None:
     )
 
     model = build_model(args, dataset).to(device)
+    model = maybe_wrap_data_parallel(model, args, device, logger)
     readout = build_next_step_readout(args).to(device)
     criterion = build_next_step_loss(args)
     optimizer = AdamW(
@@ -591,7 +653,7 @@ def main() -> None:
         if is_best:
             best_val = val_loss
             patience = 0
-            save_checkpoint(model, best_model_path)
+            save_checkpoint(unwrap_model(model), best_model_path)
         else:
             patience += 1
 
@@ -617,7 +679,7 @@ def main() -> None:
         json.dump(history, f, indent=2)
 
     logger.info("Evaluating best model on next-step test split...")
-    model.load_state_dict(torch.load(best_model_path, map_location=device))
+    unwrap_model(model).load_state_dict(torch.load(best_model_path, map_location=device))
     with torch.no_grad():
         test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False)
     logger.info(f"Test loss: {test_loss:.6f}")