908 lines
32 KiB
Python
908 lines
32 KiB
Python
"""
|
|
Train DeepHealth with next-token / next-time-point supervision.
|
|
|
|
The next-step dataset uses observed event histories, including CHECKUP state
|
|
tokens, plus optional gap <NO_EVENT> imputation. UTS training reads out only
|
|
same-time group ends.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterator, List
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
from torch.nn.utils import clip_grad_norm_
|
|
from torch.optim import AdamW
|
|
from torch.utils.data import (
|
|
DataLoader,
|
|
DistributedSampler,
|
|
RandomSampler,
|
|
Sampler,
|
|
)
|
|
from tqdm.auto import tqdm
|
|
|
|
from dataset import HealthDataset, collate_fn
|
|
from losses import build_loss
|
|
from models import DeepHealth
|
|
from readouts import build_readout
|
|
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
|
from train_util import (
|
|
configure_torch_for_training,
|
|
create_unique_run_dir,
|
|
resolve_device,
|
|
save_checkpoint,
|
|
save_config,
|
|
set_optimizer_lr,
|
|
set_seed,
|
|
setup_logging,
|
|
split_dataset,
|
|
split_dataset_by_eid_files,
|
|
)
|
|
|
|
|
|
MODEL_INPUT_KEYS = (
|
|
"event_seq",
|
|
"time_seq",
|
|
"sex",
|
|
"padding_mask",
|
|
)
|
|
|
|
EXPOSURE_INPUT_KEYS = (
|
|
"exposure_embedding",
|
|
)
|
|
|
|
|
|
class ExposureLocalityBatchSampler(Sampler[List[int]]):
|
|
"""Randomized batches with within-buffer sorting by exposure parquet locality."""
|
|
|
|
def __init__(
|
|
self,
|
|
data_source,
|
|
batch_size: int,
|
|
buffer_size: int,
|
|
seed: int,
|
|
drop_last: bool = False,
|
|
) -> None:
|
|
self.data_source = data_source
|
|
self.batch_size = int(batch_size)
|
|
self.buffer_size = max(int(buffer_size), self.batch_size)
|
|
self.seed = int(seed)
|
|
self.drop_last = bool(drop_last)
|
|
self.epoch = 0
|
|
|
|
def __iter__(self) -> Iterator[List[int]]:
|
|
n = len(self.data_source)
|
|
generator = torch.Generator().manual_seed(self.seed + self.epoch)
|
|
self.epoch += 1
|
|
shuffled = torch.randperm(n, generator=generator).tolist()
|
|
for start in range(0, n, self.buffer_size):
|
|
buffer = shuffled[start:start + self.buffer_size]
|
|
buffer.sort(key=self._locality_key)
|
|
for batch_start in range(0, len(buffer), self.batch_size):
|
|
batch = buffer[batch_start:batch_start + self.batch_size]
|
|
if len(batch) < self.batch_size and self.drop_last:
|
|
continue
|
|
yield batch
|
|
|
|
def __len__(self) -> int:
|
|
n = len(self.data_source)
|
|
if self.drop_last:
|
|
return n // self.batch_size
|
|
return (n + self.batch_size - 1) // self.batch_size
|
|
|
|
def _locality_key(self, local_idx: int) -> tuple[int, int, int]:
|
|
dataset = self.data_source
|
|
raw_idx = int(local_idx)
|
|
if hasattr(dataset, "dataset") and hasattr(dataset, "indices"):
|
|
raw_idx = int(dataset.indices[local_idx])
|
|
dataset = dataset.dataset
|
|
sample = getattr(dataset, "samples", [])[raw_idx]
|
|
exposure_index = sample.get("exposure_index")
|
|
exposure_cache = getattr(dataset, "exposure_cache", None)
|
|
if exposure_index is None or exposure_cache is None:
|
|
return (2**31 - 1, 2**31 - 1, raw_idx)
|
|
block_id, block_offset = exposure_cache.locality_key(exposure_index)
|
|
return (block_id, block_offset, raw_idx)
|
|
|
|
|
|
class DistributedExposureLocalityBatchSampler(Sampler[List[int]]):
|
|
"""Shard samples across ranks, then batch each shard by exposure locality."""
|
|
|
|
def __init__(
|
|
self,
|
|
data_source,
|
|
batch_size: int,
|
|
buffer_size: int,
|
|
seed: int,
|
|
rank: int,
|
|
world_size: int,
|
|
) -> None:
|
|
self.data_source = data_source
|
|
self.batch_size = int(batch_size)
|
|
self.buffer_size = max(int(buffer_size), self.batch_size)
|
|
self.distributed_sampler = DistributedSampler(
|
|
data_source,
|
|
num_replicas=world_size,
|
|
rank=rank,
|
|
shuffle=True,
|
|
seed=seed,
|
|
drop_last=False,
|
|
)
|
|
self._key_sampler = ExposureLocalityBatchSampler(
|
|
data_source, batch_size, buffer_size, seed
|
|
)
|
|
|
|
def set_epoch(self, epoch: int) -> None:
|
|
self.distributed_sampler.set_epoch(epoch)
|
|
|
|
def __iter__(self) -> Iterator[List[int]]:
|
|
indices = list(self.distributed_sampler)
|
|
for start in range(0, len(indices), self.buffer_size):
|
|
buffer = indices[start:start + self.buffer_size]
|
|
buffer.sort(key=self._key_sampler._locality_key)
|
|
for batch_start in range(0, len(buffer), self.batch_size):
|
|
yield buffer[batch_start:batch_start + self.batch_size]
|
|
|
|
def __len__(self) -> int:
|
|
return math.ceil(len(self.distributed_sampler) / self.batch_size)
|
|
|
|
|
|
class NextStepTrainingModel(nn.Module):
|
|
"""Keep backbone, readout, and risk head inside the DDP forward boundary."""
|
|
|
|
def __init__(self, model: DeepHealth, readout: nn.Module, readout_name: str):
|
|
super().__init__()
|
|
self.model = model
|
|
self.readout = readout
|
|
self.readout_name = readout_name
|
|
|
|
def forward(
|
|
self,
|
|
event_seq: torch.Tensor,
|
|
time_seq: torch.Tensor,
|
|
sex: torch.Tensor,
|
|
padding_mask: torch.Tensor,
|
|
readout_mask: torch.Tensor | None = None,
|
|
exposure_embedding: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
hidden = self.model(
|
|
event_seq=event_seq,
|
|
time_seq=time_seq,
|
|
sex=sex,
|
|
padding_mask=padding_mask,
|
|
target_mode="next_token",
|
|
exposure_embedding=exposure_embedding,
|
|
)
|
|
if not isinstance(hidden, torch.Tensor):
|
|
raise TypeError("DeepHealth forward must return a hidden-state tensor")
|
|
current_times = time_seq[:, :hidden.size(1)]
|
|
current_padding = padding_mask[:, :hidden.size(1)]
|
|
readout_out = self.readout(
|
|
hidden=hidden,
|
|
time_seq=current_times,
|
|
padding_mask=current_padding,
|
|
readout_mask=readout_mask
|
|
if self.readout_name == "same_time_group_end"
|
|
else None,
|
|
)
|
|
logits = self.model.calc_risk(readout_out.hidden)
|
|
return logits, current_times, readout_out.readout_mask
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Train DeepHealth with next-token/point supervision")
|
|
|
|
parser.add_argument("--data_prefix", type=str, default="ukb")
|
|
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--no_event_interval_years", type=float, default=5.0)
|
|
parser.add_argument("--include_no_event_in_uts_target", action="store_true")
|
|
|
|
parser.add_argument("--train_ratio", type=float, default=0.7)
|
|
parser.add_argument("--val_ratio", type=float, default=0.15)
|
|
parser.add_argument("--test_ratio", type=float, default=0.15)
|
|
parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
|
|
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
|
|
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
|
|
|
|
parser.add_argument("--n_embd", type=int, default=120)
|
|
parser.add_argument("--n_head", type=int, default=10)
|
|
parser.add_argument("--n_hist_layer", type=int, default=12)
|
|
parser.add_argument("--dropout", type=float, default=0.0)
|
|
parser.add_argument("--exposure_cache_dir", type=str, default=None)
|
|
parser.add_argument(
|
|
"--exposure_embeddings_file",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Precomputed exposure embeddings. Defaults to "
|
|
"<exposure_cache_dir>/exposure_embeddings.npy."
|
|
),
|
|
)
|
|
parser.add_argument("--target_mode", type=str, default="uts",
|
|
choices=["delphi2m", "uts"])
|
|
parser.add_argument("--readout_name", type=str, default=None,
|
|
choices=["token", "same_time_group_end", "last_valid"])
|
|
parser.add_argument("--readout_reduce", type=str, default="mean",
|
|
choices=["mean", "sum"])
|
|
parser.add_argument("--t_min", type=float, default=0.0027378507871321013)
|
|
parser.add_argument("--max_exp_input", type=float, default=60.0)
|
|
parser.add_argument("--ce_weight", type=float, default=1.0)
|
|
parser.add_argument("--time_weight", type=float, default=1.0)
|
|
parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true")
|
|
|
|
parser.add_argument("--batch_size", type=int, default=128)
|
|
parser.add_argument("--base_lr", type=float, default=3e-4)
|
|
parser.add_argument("--weight_decay", type=float, default=0.1)
|
|
parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
|
|
parser.add_argument("--grad_clip", type=float, default=1.0)
|
|
parser.add_argument("--max_epochs", type=int, default=200)
|
|
parser.add_argument("--warmup_epochs", type=int, default=10)
|
|
parser.add_argument("--patience", type=int, default=15)
|
|
parser.add_argument("--min_lr_ratio", type=float, default=0.1)
|
|
parser.add_argument("--num_workers", type=int, default=4)
|
|
parser.add_argument(
|
|
"--prefetch_factor",
|
|
type=int,
|
|
default=4,
|
|
help="DataLoader batches prefetched per worker when num_workers > 0.",
|
|
)
|
|
parser.add_argument(
|
|
"--exposure_locality_buffer_size",
|
|
type=int,
|
|
default=4096,
|
|
help=(
|
|
"Training-only shuffle buffer sorted by exposure parquet locality. "
|
|
"Set 0 to use the standard RandomSampler."
|
|
),
|
|
)
|
|
parser.add_argument("--device", type=str, default="cuda")
|
|
parser.add_argument(
|
|
"--amp",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Use CUDA automatic mixed precision.",
|
|
)
|
|
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(
|
|
"--ddp_backend",
|
|
default=None,
|
|
choices=["nccl", "gloo"],
|
|
help="DDP backend. Defaults to nccl for torchrun multi-GPU training.",
|
|
)
|
|
parser.add_argument("--progress_interval", type=int, default=20)
|
|
|
|
args = parser.parse_args()
|
|
use_eid_split = all(
|
|
getattr(args, name)
|
|
for name in ("train_eid_file", "val_eid_file", "test_eid_file")
|
|
)
|
|
if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
|
|
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
|
|
if args.num_workers > 0 and args.prefetch_factor <= 0:
|
|
raise ValueError("prefetch_factor must be positive when num_workers > 0")
|
|
if args.exposure_locality_buffer_size < 0:
|
|
raise ValueError("exposure_locality_buffer_size must be non-negative")
|
|
if args.exposure_embeddings_file and not args.exposure_cache_dir:
|
|
raise ValueError(
|
|
"--exposure_cache_dir is required with --exposure_embeddings_file"
|
|
)
|
|
if args.target_mode == "uts":
|
|
args.readout_name = args.readout_name or "same_time_group_end"
|
|
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
|
|
|
|
|
|
def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
|
|
if epoch < args.warmup_epochs:
|
|
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
|
progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
|
|
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
|
return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
|
|
|
|
|
|
def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
|
|
non_blocking = device.type == "cuda"
|
|
return {
|
|
key: value.to(device, non_blocking=non_blocking)
|
|
if isinstance(value, torch.Tensor)
|
|
else value
|
|
for key, value in batch.items()
|
|
}
|
|
|
|
|
|
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):
|
|
if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)):
|
|
return model.module
|
|
return 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()))
|
|
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 init_distributed(
|
|
args: argparse.Namespace,
|
|
) -> tuple[torch.device, int, int, int]:
|
|
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
|
if world_size == 1:
|
|
return resolve_device(args.device), 0, 0, 1
|
|
if args.data_parallel:
|
|
raise ValueError("--data_parallel cannot be combined with torchrun/DDP")
|
|
if not torch.cuda.is_available():
|
|
raise ValueError("Multi-process next-step training requires CUDA")
|
|
rank = int(os.environ["RANK"])
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
torch.cuda.set_device(local_rank)
|
|
dist.init_process_group(
|
|
backend=args.ddp_backend or "nccl",
|
|
init_method="env://",
|
|
)
|
|
return torch.device("cuda", local_rank), rank, local_rank, world_size
|
|
|
|
|
|
def distributed_run_dir(
|
|
args: argparse.Namespace, rank: int, world_size: int
|
|
) -> tuple[Path, str]:
|
|
payload: list[str | None] = [None, None]
|
|
if rank == 0:
|
|
run_dir, run_name = create_unique_run_dir(
|
|
lambda timestamp: (
|
|
f"absolute_exponential_next_token_{args.target_mode}_"
|
|
f"gap_{args.no_event_interval_years:g}y_"
|
|
f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_"
|
|
f"{timestamp}"
|
|
)
|
|
)
|
|
payload = [str(run_dir), run_name]
|
|
if world_size > 1:
|
|
dist.broadcast_object_list(payload, src=0)
|
|
return Path(str(payload[0])), str(payload[1])
|
|
|
|
|
|
def rank_logger(rank: int, run_dir: Path) -> logging.Logger:
|
|
if rank == 0:
|
|
return setup_logging(run_dir)
|
|
logger = logging.getLogger(f"DeepHealth.rank{rank}")
|
|
logger.handlers.clear()
|
|
logger.addHandler(logging.NullHandler())
|
|
return logger
|
|
|
|
|
|
def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
|
|
return DeepHealth(
|
|
vocab_size=dataset.vocab_size,
|
|
n_embd=args.n_embd,
|
|
n_head=args.n_head,
|
|
n_hist_layer=args.n_hist_layer,
|
|
target_mode="next_token",
|
|
dist_mode="exponential",
|
|
dropout=args.dropout,
|
|
use_exposure_embeddings=args.exposure_embeddings_file is not None,
|
|
)
|
|
|
|
|
|
def build_next_step_readout(args: argparse.Namespace):
|
|
if args.readout_name == "same_time_group_end":
|
|
return build_readout("same_time_group_end", reduce=args.readout_reduce)
|
|
return build_readout(args.readout_name)
|
|
|
|
|
|
def build_next_step_loss(args: argparse.Namespace):
|
|
if args.target_mode == "delphi2m":
|
|
ignored_tokens = {PAD_IDX, CHECKUP_IDX}
|
|
if args.ignore_no_event_in_delphi2m:
|
|
ignored_tokens.add(NO_EVENT_IDX)
|
|
return build_loss(
|
|
"delphi2m",
|
|
ignored_tokens=ignored_tokens,
|
|
t_min=args.t_min,
|
|
max_exp_input=args.max_exp_input,
|
|
ce_weight=args.ce_weight,
|
|
time_weight=args.time_weight,
|
|
)
|
|
return build_loss(
|
|
"uts",
|
|
ignored_idx={PAD_IDX, CHECKUP_IDX},
|
|
t_min=args.t_min,
|
|
max_exp_input=args.max_exp_input,
|
|
)
|
|
|
|
|
|
def compute_next_step_loss(
|
|
args: argparse.Namespace,
|
|
model,
|
|
criterion,
|
|
batch: Dict[str, torch.Tensor],
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]:
|
|
batch_cpu = batch
|
|
input_keys = [*MODEL_INPUT_KEYS, "readout_mask"]
|
|
input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu)
|
|
batch = move_batch_to_device(
|
|
{key: batch_cpu[key] for key in input_keys},
|
|
device,
|
|
)
|
|
model_kwargs = {
|
|
"event_seq": batch["event_seq"],
|
|
"time_seq": batch["time_seq"],
|
|
"sex": batch["sex"],
|
|
"padding_mask": batch["padding_mask"],
|
|
"readout_mask": batch["readout_mask"],
|
|
}
|
|
if "exposure_embedding" in batch:
|
|
model_kwargs["exposure_embedding"] = batch["exposure_embedding"]
|
|
logits, current_times, output_readout_mask = model(**model_kwargs)
|
|
non_blocking = device.type == "cuda"
|
|
targets = {
|
|
"target_event_seq": batch_cpu["target_event_seq"].to(
|
|
device, non_blocking=non_blocking
|
|
),
|
|
"target_time_seq": batch_cpu["target_time_seq"].to(
|
|
device, non_blocking=non_blocking
|
|
),
|
|
}
|
|
if args.target_mode == "uts":
|
|
targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
|
|
device, non_blocking=non_blocking
|
|
)
|
|
targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
|
|
device, non_blocking=non_blocking
|
|
)
|
|
|
|
if args.target_mode == "delphi2m":
|
|
loss, parts = criterion(
|
|
logits=logits,
|
|
target_events=targets["target_event_seq"],
|
|
target_times=targets["target_time_seq"],
|
|
current_times=current_times,
|
|
padding_mask=output_readout_mask,
|
|
return_components=True,
|
|
)
|
|
else:
|
|
loss, parts = criterion(
|
|
logits=logits,
|
|
target_multi_hot=targets["target_multi_hot"],
|
|
target_dt_unique=targets["target_dt_unique"],
|
|
readout_mask=output_readout_mask,
|
|
return_components=True,
|
|
)
|
|
return loss, parts
|
|
|
|
|
|
def run_epoch(
|
|
logger: logging.Logger,
|
|
args: argparse.Namespace,
|
|
model,
|
|
criterion,
|
|
loader: DataLoader,
|
|
optimizer: AdamW | None,
|
|
device: torch.device,
|
|
is_train: bool,
|
|
rank: int,
|
|
scaler: torch.amp.GradScaler,
|
|
amp_enabled: bool,
|
|
) -> float:
|
|
model.train(is_train)
|
|
totals = torch.zeros(3, device=device, dtype=torch.float64)
|
|
parts_sum: Dict[str, torch.Tensor] = {}
|
|
desc = "train" if is_train else "val"
|
|
progress_interval = max(1, int(args.progress_interval))
|
|
|
|
progress = tqdm(
|
|
loader, desc=desc, leave=False, dynamic_ncols=True, disable=rank != 0
|
|
)
|
|
for batch_idx, batch in enumerate(progress):
|
|
with torch.autocast(
|
|
device_type=device.type,
|
|
dtype=torch.float16,
|
|
enabled=amp_enabled,
|
|
):
|
|
loss, parts = compute_next_step_loss(
|
|
args, model, criterion, batch, device
|
|
)
|
|
finite = torch.isfinite(loss).to(dtype=torch.int32)
|
|
if dist.is_initialized():
|
|
dist.all_reduce(finite, op=dist.ReduceOp.MIN)
|
|
if not bool(finite.item()):
|
|
totals[2] += 1
|
|
continue
|
|
if is_train:
|
|
if optimizer is None:
|
|
raise ValueError("optimizer is required for training")
|
|
optimizer.zero_grad(set_to_none=True)
|
|
scaler.scale(loss).backward()
|
|
scaler.unscale_(optimizer)
|
|
if args.grad_clip > 0:
|
|
clip_grad_norm_(model.parameters(), args.grad_clip)
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
|
|
totals[0] += loss.detach().double()
|
|
totals[1] += 1
|
|
for name, value in parts.items():
|
|
parts_sum[name] = (
|
|
parts_sum.get(name, torch.zeros((), device=device))
|
|
+ value.detach()
|
|
)
|
|
if rank == 0 and (batch_idx + 1) % progress_interval == 0:
|
|
progress.set_postfix(
|
|
loss=f"{loss.detach().item():.4f}",
|
|
avg=f"{(totals[0] / totals[1].clamp_min(1)).item():.4f}",
|
|
skipped=int(totals[2].item()),
|
|
)
|
|
|
|
if dist.is_initialized():
|
|
dist.all_reduce(totals, op=dist.ReduceOp.SUM)
|
|
skipped = int(totals[2].item())
|
|
if skipped:
|
|
logger.info(f"Skipped {skipped} rank-batches due to non-finite loss")
|
|
if totals[1].item() == 0:
|
|
return float("inf")
|
|
return float((totals[0] / totals[1]).item())
|
|
|
|
|
|
def build_metadata(
|
|
args: argparse.Namespace,
|
|
dataset: HealthDataset,
|
|
run_name: str,
|
|
train_subset,
|
|
val_subset,
|
|
test_subset,
|
|
) -> Dict[str, Any]:
|
|
return {
|
|
"run_name": run_name,
|
|
"dataset_class": "NextStepHealthDataset",
|
|
"collate_fn": "next_step_collate_fn",
|
|
"model_class": "DeepHealth",
|
|
"model_target_mode": "next_token",
|
|
"target_mode": args.target_mode,
|
|
"dist_mode": "exponential",
|
|
"dataset_metadata": {
|
|
"vocab_size": int(dataset.vocab_size),
|
|
},
|
|
"use_exposure_embeddings": args.exposure_embeddings_file is not None,
|
|
"exposure_cache_dir": args.exposure_cache_dir,
|
|
"exposure_embeddings_file": args.exposure_embeddings_file,
|
|
"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,
|
|
"ddp_world_size": int(os.environ.get("WORLD_SIZE", "1")),
|
|
"split_sizes": {
|
|
"train": int(len(train_subset)),
|
|
"val": int(len(val_subset)),
|
|
"test": int(len(test_subset)),
|
|
},
|
|
"resolved_readout_name": args.readout_name,
|
|
"resolved_loss_name": args.target_mode,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.exposure_cache_dir and args.exposure_embeddings_file is None:
|
|
args.exposure_embeddings_file = str(
|
|
Path(args.exposure_cache_dir) / "exposure_embeddings.npy"
|
|
)
|
|
device, rank, local_rank, world_size = init_distributed(args)
|
|
set_seed(args.seed + rank)
|
|
configure_torch_for_training(device)
|
|
|
|
run_dir, run_name = distributed_run_dir(args, rank, world_size)
|
|
logger = rank_logger(rank, run_dir)
|
|
|
|
logger.info(f"Starting next-step training run: {run_name}")
|
|
logger.info(f"Device: {device}")
|
|
logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
|
|
logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
|
|
logger.info(f"exposure_embeddings_file={args.exposure_embeddings_file}")
|
|
logger.info(
|
|
"DataLoader IO: "
|
|
f"num_workers={args.num_workers}, "
|
|
f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, "
|
|
f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}"
|
|
)
|
|
|
|
dataset = HealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
no_event_interval_years=args.no_event_interval_years,
|
|
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
|
|
exposure_cache_dir=args.exposure_cache_dir,
|
|
exposure_embeddings_file=args.exposure_embeddings_file,
|
|
)
|
|
if dataset.exposure_cache is not None:
|
|
embedding_dim = int(dataset.exposure_cache.embeddings.shape[1])
|
|
if embedding_dim != args.n_embd:
|
|
raise ValueError(
|
|
f"Exposure embedding dim {embedding_dim} must equal "
|
|
f"--n_embd={args.n_embd}"
|
|
)
|
|
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
|
train_subset, val_subset, test_subset = split_dataset_by_eid_files(
|
|
dataset=dataset,
|
|
train_eid_file=args.train_eid_file,
|
|
val_eid_file=args.val_eid_file,
|
|
test_eid_file=args.test_eid_file,
|
|
)
|
|
logger.info(
|
|
"Using eid split files: "
|
|
f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
|
|
)
|
|
else:
|
|
train_subset, val_subset, test_subset = split_dataset(
|
|
dataset=dataset,
|
|
train_ratio=args.train_ratio,
|
|
val_ratio=args.val_ratio,
|
|
test_ratio=args.test_ratio,
|
|
seed=args.seed,
|
|
)
|
|
logger.info(
|
|
f"Using random ratio split: train={args.train_ratio}, "
|
|
f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
|
|
)
|
|
logger.info(
|
|
f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
|
|
)
|
|
|
|
if args.batch_size % world_size != 0:
|
|
raise ValueError(
|
|
f"--batch_size={args.batch_size} must be divisible by "
|
|
f"DDP world size {world_size}"
|
|
)
|
|
local_batch_size = args.batch_size // world_size
|
|
dataloader_kwargs = {
|
|
"collate_fn": collate_fn,
|
|
"num_workers": args.num_workers,
|
|
"pin_memory": device.type == "cuda",
|
|
"persistent_workers": args.num_workers > 0,
|
|
"prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None,
|
|
}
|
|
use_locality_sampler = (
|
|
args.exposure_cache_dir is not None
|
|
and args.exposure_locality_buffer_size > 0
|
|
)
|
|
if use_locality_sampler:
|
|
logger.info("Using exposure-locality batch sampler for training")
|
|
if world_size > 1:
|
|
train_batch_sampler = DistributedExposureLocalityBatchSampler(
|
|
train_subset,
|
|
batch_size=local_batch_size,
|
|
buffer_size=args.exposure_locality_buffer_size,
|
|
seed=args.seed,
|
|
rank=rank,
|
|
world_size=world_size,
|
|
)
|
|
else:
|
|
train_batch_sampler = ExposureLocalityBatchSampler(
|
|
train_subset,
|
|
batch_size=local_batch_size,
|
|
buffer_size=args.exposure_locality_buffer_size,
|
|
seed=args.seed,
|
|
)
|
|
train_loader = DataLoader(
|
|
train_subset,
|
|
batch_sampler=train_batch_sampler,
|
|
**dataloader_kwargs,
|
|
)
|
|
else:
|
|
train_sampler = (
|
|
DistributedSampler(
|
|
train_subset,
|
|
num_replicas=world_size,
|
|
rank=rank,
|
|
shuffle=True,
|
|
seed=args.seed,
|
|
)
|
|
if world_size > 1
|
|
else RandomSampler(
|
|
train_subset,
|
|
generator=torch.Generator().manual_seed(args.seed),
|
|
)
|
|
)
|
|
train_loader = DataLoader(
|
|
train_subset,
|
|
batch_size=local_batch_size,
|
|
sampler=train_sampler,
|
|
**dataloader_kwargs,
|
|
)
|
|
val_sampler = (
|
|
DistributedSampler(
|
|
val_subset, num_replicas=world_size, rank=rank, shuffle=False
|
|
)
|
|
if world_size > 1 else None
|
|
)
|
|
test_sampler = (
|
|
DistributedSampler(
|
|
test_subset, num_replicas=world_size, rank=rank, shuffle=False
|
|
)
|
|
if world_size > 1 else None
|
|
)
|
|
val_loader = DataLoader(
|
|
val_subset,
|
|
batch_size=local_batch_size,
|
|
sampler=val_sampler,
|
|
shuffle=False,
|
|
**dataloader_kwargs,
|
|
)
|
|
test_loader = DataLoader(
|
|
test_subset,
|
|
batch_size=local_batch_size,
|
|
sampler=test_sampler,
|
|
shuffle=False,
|
|
**dataloader_kwargs,
|
|
)
|
|
|
|
backbone = build_model(args, dataset)
|
|
readout = build_next_step_readout(args).to(device)
|
|
model = NextStepTrainingModel(
|
|
backbone, readout, args.readout_name
|
|
).to(device)
|
|
if world_size > 1:
|
|
model = DistributedDataParallel(
|
|
model, device_ids=[local_rank], output_device=local_rank
|
|
)
|
|
logger.info(
|
|
f"Using DDP with {world_size} processes; "
|
|
f"global_batch={args.batch_size}, per_gpu_batch={local_batch_size}"
|
|
)
|
|
else:
|
|
model = maybe_wrap_data_parallel(model, args, device, logger)
|
|
criterion = build_next_step_loss(args)
|
|
optimizer = AdamW(
|
|
model.parameters(),
|
|
lr=args.base_lr,
|
|
betas=tuple(args.betas),
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
amp_enabled = bool(args.amp and device.type == "cuda")
|
|
scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
|
|
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
|
|
|
if rank == 0:
|
|
save_config(
|
|
args,
|
|
run_dir / "train_config.json",
|
|
extra=build_metadata(
|
|
args, dataset, run_name, train_subset, val_subset, test_subset
|
|
),
|
|
)
|
|
|
|
best_val = float("inf")
|
|
patience = 0
|
|
history = []
|
|
best_model_path = run_dir / "best_model.pt"
|
|
start = time.time()
|
|
|
|
for epoch in range(args.max_epochs):
|
|
if hasattr(train_loader.batch_sampler, "set_epoch"):
|
|
train_loader.batch_sampler.set_epoch(epoch)
|
|
elif hasattr(train_loader.sampler, "set_epoch"):
|
|
train_loader.sampler.set_epoch(epoch)
|
|
lr = get_lr(epoch, args, adaptive_lr)
|
|
set_optimizer_lr(optimizer, lr)
|
|
|
|
train_loss = run_epoch(
|
|
logger, args, model, criterion, train_loader, optimizer,
|
|
device, True, rank, scaler, amp_enabled,
|
|
)
|
|
with torch.no_grad():
|
|
val_loss = run_epoch(
|
|
logger, args, model, criterion, val_loader, None,
|
|
device, False, rank, scaler, amp_enabled,
|
|
)
|
|
|
|
is_best = val_loss < best_val
|
|
if is_best:
|
|
best_val = val_loss
|
|
patience = 0
|
|
if rank == 0:
|
|
save_checkpoint(unwrap_model(model).model, best_model_path)
|
|
else:
|
|
patience += 1
|
|
|
|
logger.info(
|
|
f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
|
|
f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
|
|
f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
|
|
f"elapsed={time.time() - start:.1f}s"
|
|
)
|
|
history.append({
|
|
"epoch": epoch + 1,
|
|
"lr": lr,
|
|
"train_loss": train_loss,
|
|
"val_loss": val_loss,
|
|
"best_val_loss": best_val,
|
|
"is_best": int(is_best),
|
|
})
|
|
if patience >= args.patience:
|
|
logger.info(f"Early stopping triggered at epoch {epoch + 1}")
|
|
break
|
|
|
|
if rank == 0:
|
|
with (run_dir / "history.json").open("w", encoding="utf-8") as f:
|
|
json.dump(history, f, indent=2)
|
|
|
|
logger.info("Evaluating best model on next-step test split...")
|
|
if world_size > 1:
|
|
dist.barrier()
|
|
unwrap_model(model).model.load_state_dict(
|
|
torch.load(best_model_path, map_location=device)
|
|
)
|
|
with torch.no_grad():
|
|
test_loss = run_epoch(
|
|
logger, args, model, criterion, test_loader, None,
|
|
device, False, rank, scaler, amp_enabled,
|
|
)
|
|
logger.info(f"Test loss: {test_loss:.6f}")
|
|
logger.info(f"Best checkpoint: {best_model_path}")
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
finally:
|
|
if dist.is_initialized():
|
|
dist.destroy_process_group()
|