refactor: isolate Delphi2M next-token pipeline

train_all_future.py

@@ -37,8 +37,10 @@ from train_util import (
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    get_lr,
     get_model_parameter_counts,
     load_extra_info_types_file,
+    move_batch_to_device,
     resolve_device,
     save_checkpoint,
     save_config,
@@ -135,24 +137,6 @@ def parse_args() -> argparse.Namespace:
     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 build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
@@ -214,7 +198,6 @@ def compute_all_future_loss(
         other_value=batch["other_value"],
         other_value_kind=batch["other_value_kind"],
         other_time=batch["other_time"],
-        target_mode="all_future",
     )
     logits = model.calc_risk(hidden)