refactor: isolate Delphi2M next-token pipeline

train_util.py

@@ -6,6 +6,7 @@ import sys
 import time
 import csv
 from datetime import datetime
+import math
 from pathlib import Path
 from typing import Any, Dict, Iterable, Tuple
 
@@ -141,6 +142,31 @@ def resolve_device(device_arg: str) -> torch.device:
     raise ValueError(f"Unsupported device: {device_arg}")
 
 
+def get_lr(epoch: int, args: Any, 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 split_dataset(
     dataset: HealthDataset,
     train_ratio: float,
@@ -291,15 +317,6 @@ def split_all_future_datasets_by_eid_files(
     )
 
 
-def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
-    return AdamW(
-        model.parameters(),
-        lr=args.base_lr,
-        betas=tuple(args.betas),
-        weight_decay=args.weight_decay,
-    )
-
-
 def get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]:
     """Return stable total and trainable parameter counts."""
     return {