��
b��
E
C/intern9/huhongkai/hs300_factor_lab/src/train_seq_gru_ddp_memmap.pyڇ
և
C/intern9/huhongkai/hs300_factor_lab/src/train_seq_gru_ddp_memmap.py *�     y_parts = []
     w_parts = []
-    seq_offsets_cpu: torch.Tensor | None = None
     for part in batch.parts:
+        span_x = part.span_x.to(device, non_blocking=copy_non_blocking)
+        local_end = part.local_end.to(device, non_blocking=copy_non_blocking)
         yb = part.y.to(device, non_blocking=copy_non_blocking)
         wb = part.w.to(device, non_blocking=copy_non_blocking)
-        span_rows = int(part.span_x.shape[0])
-        batch_rows = int(part.local_end.shape[0])
-        use_cpu_materialize = (
-            device.type == "cuda"
-            and part.span_x.device.type == "cpu"
-            and span_rows > max(1, batch_rows) * 2
-        )
-        if use_cpu_materialize:
-            if seq_offsets_cpu is None:
-                seq_offsets_cpu = seq_offsets.detach().to(device="cpu")
-            seq_idx_cpu = part.local_end[:, None] + seq_offsets_cpu[None, :]
-            xb_cpu = part.span_x[seq_idx_cpu].float()
-            if copy_non_blocking:
-                xb_cpu = xb_cpu.pin_memory()
-            xb = xb_cpu.to(device, non_blocking=copy_non_blocking)
-        else:
-            span_x = part.span_x.to(device, non_blocking=copy_non_blocking)
-            local_end = part.local_end.to(device, non_blocking=copy_non_blocking)
-            seq_idx = local_end[:, None] + seq_offsets[None, :]
-            xb = span_x[seq_idx].float()
+        seq_idx = local_end[:, None] + seq_offsets[None, :]
+        xb = span_x[seq_idx].float()
         xb = (xb - feat_mean) / feat_std
         # Keep GRU inputs in a standard dense layout so cuDNN does not fall back2�import argparse
import csv
import hashlib
import json
import math
import os
import queue
import random
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import nullcontext
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from accelerate import Accelerator

from common import ensure_dir, load_config, load_split_view_meta, resolve_named_day_dirs, resolve_row_root
from splitview_time_utils import load_split_view_source_field_slice


def set_global_seed(seed: int) -> None:
    seed = int(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def build_epoch_lr_schedule(train_cfg: dict, epochs: int) -> List[float]:
    base_lr = float(train_cfg["learning_rate"])
    raw_values = train_cfg.get("lr_epoch_values")
    if raw_values is None:
        return [base_lr for _ in range(int(epochs))]
    values = [float(x) for x in raw_values]
    if len(values) != int(epochs):
        raise ValueError(
            f"train.lr_epoch_values length ({len(values)}) must match epochs ({int(epochs)})."
        )
    return values


def set_optimizer_lr(optimizer, lr: float) -> None:
    lr = float(lr)
    for group in optimizer.param_groups:
        group["lr"] = lr


class AccelerateRuntime:
    def __init__(self, use_amp: bool):
        self.accelerator = Accelerator(mixed_precision="fp16" if bool(use_amp) else "no")
        self.device = self.accelerator.device
        self.process_index = int(self.accelerator.process_index)
        self.num_processes = int(self.accelerator.num_processes)
        self.is_main_process = bool(self.accelerator.is_main_process)

    def autocast(self):
        return self.accelerator.autocast()

    def backward(self, loss: torch.Tensor) -> None:
        self.accelerator.backward(loss)

    def step_optimizer(self, optimizer) -> None:
        optimizer.step()

    def prepare(self, model, optimizer):
        return self.accelerator.prepare(model, optimizer)

    def unwrap_model(self, model):
        return self.accelerator.unwrap_model(model)

    def reduce(self, tensor: torch.Tensor, reduction: str = "sum") -> torch.Tensor:
        return self.accelerator.reduce(tensor, reduction=reduction)

    def gather(self, tensor: torch.Tensor) -> torch.Tensor:
        return self.accelerator.gather(tensor)

    def wait_for_everyone(self) -> None:
        self.accelerator.wait_for_everyone()

    def close(self) -> None:
        return None


class NativeDDPRuntime:
    def __init__(self, use_amp: bool, enable_static_graph: bool = True):
        if torch.cuda.is_available():
            local_rank = int(os.environ.get("LOCAL_RANK", 0))
            torch.cuda.set_device(local_rank)
            self.device = torch.device("cuda", local_rank)
        else:
            self.device = torch.device("cpu")
        requested_world = int(os.environ.get("WORLD_SIZE", "1"))
        self._distributed = requested_world > 1
        if self._distributed and not dist.is_initialized():
            backend = "nccl" if self.device.type == "cuda" else "gloo"
            dist.init_process_group(backend=backend, timeout=timedelta(seconds=7200))
        if dist.is_initialized():
            self.process_index = int(dist.get_rank())
            self.num_processes = int(dist.get_world_size())
        else:
            self.process_index = 0
            self.num_processes = 1
        self.is_main_process = self.process_index == 0
        self.use_amp = bool(use_amp) and self.device.type == "cuda"
        self.enable_static_graph = bool(enable_static_graph)
        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)

    def autocast(self):
        if self.device.type == "cuda":
            return torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp)
        return nullcontext()

    def backward(self, loss: torch.Tensor) -> None:
        if self.use_amp:
            self.scaler.scale(loss).backward()
        else:
            loss.backward()

    def step_optimizer(self, optimizer) -> None:
        if self.use_amp:
            self.scaler.step(optimizer)
            self.scaler.update()
        else:
            optimizer.step()

    def prepare(self, model, optimizer):
        model = model.to(self.device)
        if self.num_processes > 1:
            ddp_kwargs = {
                "broadcast_buffers": False,
                "gradient_as_bucket_view": True,
            }
            if self.device.type == "cuda":
                ddp_kwargs["device_ids"] = [self.device.index]
                ddp_kwargs["output_device"] = self.device.index
            if self.enable_static_graph:
                try:
                    model = DDP(model, static_graph=True, **ddp_kwargs)
                except TypeError:
                    model = DDP(model, **ddp_kwargs)
            else:
                model = DDP(model, **ddp_kwargs)
        return model, optimizer

    def unwrap_model(self, model):
        return model.module if isinstance(model, DDP) else model

    def reduce(self, tensor: torch.Tensor, reduction: str = "sum") -> torch.Tensor:
        if self.num_processes <= 1:
            return tensor
        out = tensor.clone()
        reduction_name = str(reduction).lower()
        if reduction_name == "sum":
            op = dist.ReduceOp.SUM
            dist.all_reduce(out, op=op)
        elif reduction_name == "mean":
            dist.all_reduce(out, op=dist.ReduceOp.SUM)
            out = out / float(self.num_processes)
        elif reduction_name == "max":
            dist.all_reduce(out, op=dist.ReduceOp.MAX)
        else:
            raise ValueError(f"Unsupported reduction: {reduction}")
        return out

    def gather(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.num_processes <= 1:
            return tensor
        gather_list = [torch.empty_like(tensor) for _ in range(self.num_processes)]
        dist.all_gather(gather_list, tensor)
        return torch.cat(gather_list, dim=0)

    def wait_for_everyone(self) -> None:
        if self.num_processes > 1:
            if self.device.type == "cuda":
                dist.barrier(device_ids=[self.device.index])
            else:
                dist.barrier()

    def close(self) -> None:
        if dist.is_initialized():
            dist.destroy_process_group()


def build_runtime(runtime_backend: str, use_amp: bool, enable_static_graph: bool = True):
    backend = str(runtime_backend).strip().lower()
    if backend == "native":
        return NativeDDPRuntime(use_amp=use_amp, enable_static_graph=enable_static_graph)
    if backend == "accelerate":
        return AccelerateRuntime(use_amp=use_amp)
    raise ValueError(f"Unsupported runtime_backend: {runtime_backend}")


@dataclass
class DayStore:
    name: str
    n_rows: int
    n_factors: int
    x: np.memmap
    y: np.memmap
    w: np.memmap
    sym_start: np.memmap
    sym_end: np.memmap
    need_clean: bool


_SPLIT_X_ARRAY_CACHE: Dict[str, np.ndarray] = {}
DEFAULT_MIN_TIMECODE = 93_000_000
_USE_SOURCE_RAW_WEIGHTS = False


def configure_split_view_weight_loading(use_source_raw_weights: bool) -> None:
    global _USE_SOURCE_RAW_WEIGHTS
    _USE_SOURCE_RAW_WEIGHTS = bool(use_source_raw_weights)


def _resolve_source_x_path(day_dir: Path, meta: dict) -> Path | None:
    explicit = str(meta.get("source_x_path") or "").strip()
    if explicit:
        return Path(explicit)
    source_root = str(meta.get("source_root") or "").strip()
    source_split = str(meta.get("source_split") or "").strip()
    if source_root and source_split:
        return Path(source_root) / source_split / "x.npy"
    return None


def _load_cached_x_npy(path: Path) -> np.ndarray:
    key = str(path.resolve())
    arr = _SPLIT_X_ARRAY_CACHE.get(key)
    if arr is None:
        arr = np.load(path, mmap_mode="r", allow_pickle=False)
        _SPLIT_X_ARRAY_CACHE[key] = arr
    return arr


def list_ready_days(row_root: Path) -> List[Path]:
    days = []
    for p in sorted(row_root.iterdir()):
        if p.is_dir() and (p / "_SUCCESS").exists():
            days.append(p)
    return days


def split_days(days: List[Path], ratio: float) -> Tuple[List[Path], List[Path]]:
    n = len(days)
    if n < 2:
        return days, days
    cut = max(1, min(n - 1, int(n * ratio)))
    return days[:cut], days[cut:]


def _day_str_from_dir(day_dir: Path) -> str:
    name = day_dir.name
    if len(name) >= 8:
        day = name[:8]
        if day.isdigit():
            return day
    return ""


def filter_days_by_date(days: List[Path], start_date: str, end_date: str) -> List[Path]:
    s = (start_date or "").strip()
    e = (end_date or "").strip()
    if not s and not e:
        return list(days)
    if s and (len(s) != 8 or not s.isdigit()):
        raise ValueError(f"Invalid start_date: {s}")
    if e and (len(e) != 8 or not e.isdigit()):
        raise ValueError(f"Invalid end_date: {e}")
    lo = s if s else "00000000"
    hi = e if e else "99999999"
    if lo > hi:
        raise ValueError(f"Invalid date window: {lo} > {hi}")
    out: List[Path] = []
    for d in days:
        day = _day_str_from_dir(d)
        if day and lo <= day <= hi:
            out.append(d)
    return out


def limit_days(days: List[Path], day_limit: int) -> List[Path]:
    if day_limit <= 0 or day_limit >= len(days):
        return days
    return days[:day_limit]


def load_day_meta(day_dir: Path) -> dict:
    return json.loads((day_dir / "meta.json").read_text(encoding="utf-8"))


def load_day_store(day_dir: Path, prefer_fp16: bool) -> DayStore:
    meta = load_day_meta(day_dir)
    n_rows = int(meta["n_rows"])
    n_factors = int(meta["n_factors"])
    source_x_path = _resolve_source_x_path(day_dir, meta)
    if source_x_path is not None:
        row_start = int(meta.get("source_row_start", 0))
        row_stop = int(meta.get("source_row_stop", row_start + n_rows))
        x_all = _load_cached_x_npy(source_x_path)
        x = x_all[row_start:row_stop]
        need_clean = True
    else:
        fp16_path = day_dir / "x_top200_f16_filled.memmap"
        if prefer_fp16 and fp16_path.exists():
            x = np.memmap(fp16_path, mode="r", dtype=np.float16, shape=(n_rows, n_factors))
            need_clean = False
        else:
            x = np.memmap(day_dir / "x_top200_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows, n_factors))
            need_clean = True
    y = np.memmap(day_dir / "y_sum_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    if bool(_USE_SOURCE_RAW_WEIGHTS):
        try:
            w = np.asarray(load_split_view_source_field_slice(day_dir, "w"), dtype=np.float32)
        except Exception:
            w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    else:
        w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    sym_n = int(meta["symbol_count_meta"])
    sym_start = np.memmap(day_dir / "symbol_start_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    sym_end = np.memmap(day_dir / "symbol_end_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    return DayStore(day_dir.name, n_rows, n_factors, x, y, w, sym_start, sym_end, need_clean)


def load_day_symbol_bounds(day_dir: Path) -> Tuple[np.memmap, np.memmap]:
    meta = load_day_meta(day_dir)
    sym_n = int(meta["symbol_count_meta"])
    sym_start = np.memmap(day_dir / "symbol_start_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    sym_end = np.memmap(day_dir / "symbol_end_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    return sym_start, sym_end


def filter_valid_end_indices(
    day_dir: Path,
    end_indices: np.ndarray,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> np.ndarray:
    ends = np.asarray(end_indices, dtype=np.int64)
    if ends.size == 0:
        return ends
    mask = np.ones((int(ends.shape[0]),), dtype=bool)
    if bool(require_positive_weight):
        meta = load_day_meta(day_dir)
        n_rows = int(meta["n_rows"])
        cache_w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
        weight_mask = np.asarray(cache_w[ends], dtype=np.float32)
        mask &= np.isfinite(weight_mask) & (weight_mask > 0.0)
    if int(min_timecode) > 0:
        dt_slice = np.asarray(load_split_view_source_field_slice(day_dir, "datetime"), dtype=np.int64)
        dt_end = np.asarray(dt_slice[ends], dtype=np.int64)
        mask &= dt_end >= int(min_timecode)
    return ends[mask]


def build_valid_end_indices_from_bounds(
    sym_start: np.ndarray,
    sym_end: np.ndarray,
    seq_len: int,
    sample_stride: int,
) -> np.ndarray:
    stride = int(sample_stride)
    starts = np.asarray(sym_start, dtype=np.int64) + int(seq_len) - 1
    # symbol_end_idx is exclusive in this memmap layout.
    ends = np.asarray(sym_end, dtype=np.int64) - 1
    valid_mask = starts <= ends
    if not np.any(valid_mask):
        return np.empty((0,), dtype=np.int64)
    starts = starts[valid_mask]
    ends = ends[valid_mask]
    lengths = ((ends - starts) // stride + 1).astype(np.int64, copy=False)
    total = int(lengths.sum())
    if total <= 0:
        return np.empty((0,), dtype=np.int64)
    repeated_starts = np.repeat(starts, lengths)
    group_offsets = np.repeat(np.cumsum(lengths, dtype=np.int64) - lengths, lengths)
    intra_offsets = np.arange(total, dtype=np.int64) - group_offsets
    return repeated_starts + intra_offsets * stride


def build_valid_end_indices(store: DayStore, seq_len: int, sample_stride: int) -> np.ndarray:
    return build_valid_end_indices_from_bounds(
        store.sym_start,
        store.sym_end,
        seq_len=seq_len,
        sample_stride=sample_stride,
    )


def build_end_index_cache_path(
    cache_dir: Path,
    day_dir: Path,
    seq_len: int,
    sample_stride: int,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> Path:
    key_src = "\n".join(
        [
            str(day_dir.parent),
            str(day_dir.name),
            str(int(seq_len)),
            str(int(sample_stride)),
            str(int(min_timecode)),
            str(int(bool(require_positive_weight))),
        ]
    )
    key = hashlib.sha1(key_src.encode("utf-8")).hexdigest()[:16]
    return cache_dir / f"{day_dir.name}_seq{int(seq_len)}_stride{int(sample_stride)}_{key}.npy"


def load_end_index_cache(path: Path) -> np.ndarray:
    arr = np.load(path, allow_pickle=False)
    return np.asarray(arr, dtype=np.int64)


def wait_for_end_index_cache(path: Path, timeout_sec: float = 600.0) -> np.ndarray:
    deadline = time.time() + float(timeout_sec)
    last_err = None
    while time.time() < deadline:
        if path.exists() and path.stat().st_size > 0:
            try:
                return load_end_index_cache(path)
            except Exception as exc:  # pragma: no cover - transient partial-write case
                last_err = exc
        time.sleep(0.2)
    if last_err is not None:
        raise TimeoutError(f"Timed out waiting for end-index cache {path}: {last_err}") from last_err
    raise TimeoutError(f"Timed out waiting for end-index cache {path}")


def load_or_build_valid_end_indices(
    day_dir: Path,
    seq_len: int,
    sample_stride: int,
    cache_dir: Path | None,
    cache_writer: bool,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> np.ndarray:
    if cache_dir is None:
        sym_start, sym_end = load_day_symbol_bounds(day_dir)
        ends = build_valid_end_indices_from_bounds(sym_start, sym_end, seq_len=seq_len, sample_stride=sample_stride)
        return filter_valid_end_indices(
            day_dir,
            ends,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
    cache_path = build_end_index_cache_path(
        cache_dir,
        day_dir,
        seq_len=seq_len,
        sample_stride=sample_stride,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    if cache_path.exists() and cache_path.stat().st_size > 0:
        return load_end_index_cache(cache_path)
    if not cache_writer:
        return wait_for_end_index_cache(cache_path)
    sym_start, sym_end = load_day_symbol_bounds(day_dir)
    ends = build_valid_end_indices_from_bounds(sym_start, sym_end, seq_len=seq_len, sample_stride=sample_stride)
    ends = filter_valid_end_indices(
        day_dir,
        ends,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    tmp_path = cache_path.with_suffix(f".tmp.{int(time.time() * 1000)}.{os.getpid()}.npy")
    np.save(tmp_path, ends)
    tmp_path.replace(cache_path)
    return ends


def truncate_end_indices(end_indices: List[np.ndarray], max_samples: int) -> List[np.ndarray]:
    if max_samples <= 0:
        return end_indices
    remaining = int(max_samples)
    trimmed: List[np.ndarray] = []
    for ends in end_indices:
        if remaining <= 0:
            trimmed.append(np.empty((0,), dtype=np.int64))
            continue
        take = min(int(ends.shape[0]), remaining)
        trimmed.append(ends[:take])
        remaining -= take
    return trimmed


def split_contiguous_even(total: int, rank: int, world_size: int) -> Tuple[int, int]:
    if total <= 0:
        return 0, 0
    start = (total * rank) // max(1, world_size)
    end = (total * (rank + 1)) // max(1, world_size)
    return int(start), int(end)


@dataclass(frozen=True)
class BatchSlice:
    day_i: int
    start: int
    stop: int


@dataclass(frozen=True)
class BatchPlanItem:
    parts: Tuple[BatchSlice, ...]
    pad_size: int = 0


@dataclass(frozen=True)
class PackedSeqBatchPart:
    span_x: torch.Tensor
    local_end: torch.Tensor
    y: torch.Tensor
    w: torch.Tensor


@dataclass(frozen=True)
class PackedSeqBatch:
    parts: Tuple[PackedSeqBatchPart, ...]


@dataclass(frozen=True)
class MaterializedSeqBatch:
    xb: torch.Tensor
    y: torch.Tensor
    w: torch.Tensor


class SeqMemmapBatchLoader:
    def __init__(
        self,
        day_dirs: List[Path],
        seq_len: int,
        sample_stride: int,
        batch_size: int,
        shuffle: bool,
        max_samples: int = -1,
        prefer_fp16: bool = True,
        rank: int = 0,
        world_size: int = 1,
        seed: int = 20260312,
        pin_memory: bool = True,
        loader_threads: int = 1,
        prefetch_batches: int = 1,
        pad_last_batch: bool = False,
        batch_overlap: int = 0,
        index_cache_dir: Path | None = None,
        min_timecode: int = -1,
        require_positive_weight: bool = False,
    ):
        self.seq_len = seq_len
        self.batch_size = int(batch_size)
        self.batch_overlap = max(0, int(batch_overlap))
        if self.batch_overlap >= self.batch_size:
            raise ValueError(
                f"batch_overlap must be smaller than batch_size, got overlap={self.batch_overlap} batch_size={self.batch_size}"
            )
        self.batch_stride = self.batch_size - self.batch_overlap
        self.shuffle = bool(shuffle)
        self.rank = int(rank)
        self.world_size = int(world_size)
        self.seed = int(seed)
        self.day_dirs = list(day_dirs)
        self.prefer_fp16 = bool(prefer_fp16)
        self.pin_memory = bool(pin_memory) and torch.cuda.is_available()
        self.loader_threads = max(1, int(loader_threads))
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.pad_last_batch = bool(pad_last_batch)
        self.epoch = 0
        self.index_cache_dir = index_cache_dir
        self.min_timecode = int(min_timecode)
        self.require_positive_weight = bool(require_positive_weight)
        self.stores: List[DayStore | None] = [None for _ in self.day_dirs]
        end_indices = [
            load_or_build_valid_end_indices(
                d,
                seq_len=seq_len,
                sample_stride=sample_stride,
                cache_dir=self.index_cache_dir,
                cache_writer=(self.rank == 0),
                min_timecode=self.min_timecode,
                require_positive_weight=self.require_positive_weight,
            )
            for d in self.day_dirs
        ]
        self.end_indices: List[np.ndarray] = truncate_end_indices(end_indices, max_samples=max_samples)
        self.total_samples = int(sum(int(ends.shape[0]) for ends in self.end_indices))
        self.seq_offsets = np.arange(-(self.seq_len - 1), 1, dtype=np.int64)
        self.rank_sample_counts = self._compute_rank_sample_counts()
        self.rank_total_samples = [int(sum(day_counts)) for day_counts in self.rank_sample_counts]
        self.rank_unique_batches = [
            self._compute_batch_count(total) for total in self.rank_total_samples
        ]
        self.total_batches = max(self.rank_unique_batches, default=0)

    def __len__(self) -> int:
        return self.total_batches

    def set_epoch(self, epoch: int) -> None:
        self.epoch = int(epoch)

    def _get_store(self, day_i: int) -> DayStore:
        store = self.stores[day_i]
        if store is None:
            store = load_day_store(self.day_dirs[day_i], prefer_fp16=self.prefer_fp16)
            self.stores[day_i] = store
        return store

    def _compute_rank_sample_counts(self) -> List[List[int]]:
        counts: List[List[int]] = [[] for _ in range(self.world_size)]
        for ends in self.end_indices:
            n = int(ends.shape[0])
            for rank in range(self.world_size):
                start, stop = split_contiguous_even(n, rank, self.world_size)
                counts[rank].append(max(0, stop - start))
        return counts

    def _compute_batch_count(self, total: int) -> int:
        total = int(total)
        if total <= 0:
            return 0
        if total <= self.batch_size:
            return 1
        return 1 + int(math.ceil((total - self.batch_size) / self.batch_stride))

    def _tail_parts(self, parts: List[BatchSlice], keep: int) -> List[BatchSlice]:
        keep = max(0, int(keep))
        if keep <= 0:
            return []
        out: List[BatchSlice] = []
        remaining = keep
        for part in reversed(parts):
            part_len = int(part.stop - part.start)
            if part_len <= 0:
                continue
            take = min(remaining, part_len)
            out.append(BatchSlice(day_i=part.day_i, start=part.stop - take, stop=part.stop))
            remaining -= take
            if remaining == 0:
                break
        if remaining != 0:
            raise RuntimeError(f"Failed to preserve batch overlap keep={keep}, remaining={remaining}")
        out.reverse()
        return out

    def _build_rank_plan(self) -> List[BatchPlanItem]:
        # Rotate the contiguous shard every epoch so each rank sees different regions over time.
        shard_rank = (self.rank + self.epoch) % max(1, self.world_size) if self.shuffle else self.rank
        items: List[BatchPlanItem] = []
        current_parts: List[BatchSlice] = []
        filled = 0
        for day_i, ends in enumerate(self.end_indices):
            total = int(ends.shape[0])
            start, stop = split_contiguous_even(total, shard_rank, self.world_size)
            if stop <= start:
                continue
            cursor = int(start)
            stop = int(stop)
            while cursor < stop:
                need = self.batch_size - filled
                take = min(need, stop - cursor)
                current_parts.append(BatchSlice(day_i=day_i, start=cursor, stop=cursor + take))
                cursor += take
                filled += take
                if filled == self.batch_size:
                    emitted_parts = tuple(current_parts)
                    items.append(BatchPlanItem(parts=emitted_parts, pad_size=0))
                    current_parts = self._tail_parts(list(emitted_parts), self.batch_overlap)
                    filled = self.batch_overlap
        if current_parts:
            pad_size = self.batch_size - filled if self.pad_last_batch else 0
            items.append(BatchPlanItem(parts=tuple(current_parts), pad_size=pad_size))
        if not items:
            return []
        if self.shuffle and len(items) > 1:
            rng = np.random.default_rng(self.seed + self.epoch)
            if self.batch_overlap > 0:
                shift = int(rng.integers(len(items)))
                if shift > 0:
                    items = items[shift:] + items[:shift]
            else:
                order = rng.permutation(len(items))
                items = [items[int(i)] for i in order.tolist()]
        if len(items) < self.total_batches:
            base = list(items)
            pad_idx = 0
            while len(items) < self.total_batches:
                items.append(base[pad_idx % len(base)])
                pad_idx += 1
        return items

    def _load_batch_part(self, part: BatchSlice) -> PackedSeqBatchPart:
        day_i, start, stop = part.day_i, part.start, part.stop
        store = self._get_store(day_i)
        batch_end = self.end_indices[day_i][start:stop]
        if batch_end.size == 0:
            raise RuntimeError(f"Empty batch slice for day_i={day_i}, start={start}, stop={stop}")
        span_start = int(batch_end[0]) - self.seq_len + 1
        span_end = int(batch_end[-1])
        span_x_dtype = np.float32 if store.need_clean else store.x.dtype
        span_x = np.array(store.x[span_start : span_end + 1], dtype=span_x_dtype, copy=True)
        if store.need_clean:
            np.nan_to_num(span_x, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        local_end = np.ascontiguousarray(batch_end - span_start, dtype=np.int64)
        y = np.array(store.y[batch_end], dtype=np.float32, copy=True)
        w = np.array(store.w[batch_end], dtype=np.float32, copy=True)
        np.nan_to_num(y, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        np.nan_to_num(w, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        np.maximum(w, 0.0, out=w)
        span_xb = torch.from_numpy(np.ascontiguousarray(span_x))
        local_endb = torch.from_numpy(local_end)
        yb = torch.from_numpy(np.ascontiguousarray(y))
        wb = torch.from_numpy(np.ascontiguousarray(w))
        if self.pin_memory:
            span_xb = span_xb.pin_memory()
            local_endb = local_endb.pin_memory()
            yb = yb.pin_memory()
            wb = wb.pin_memory()
        return PackedSeqBatchPart(span_x=span_xb, local_end=local_endb, y=yb, w=wb)

    def _pad_batch_part(self, part: PackedSeqBatchPart, pad_size: int) -> PackedSeqBatchPart:
        if pad_size <= 0:
            return part
        local_end = torch.cat([part.local_end, part.local_end[-1:].repeat(int(pad_size))], dim=0)
        y = torch.cat([part.y, part.y[-1:].repeat(int(pad_size))], dim=0)
        w = torch.cat([part.w, part.w[-1:].repeat(int(pad_size))], dim=0)
        if self.pin_memory:
            local_end = local_end.pin_memory()
            y = y.pin_memory()
            w = w.pin_memory()
        return PackedSeqBatchPart(span_x=part.span_x, local_end=local_end, y=y, w=w)

    def _load_batch_from_plan_item(self, item: BatchPlanItem) -> PackedSeqBatch:
        parts = [self._load_batch_part(part) for part in item.parts]
        if not parts:
            raise RuntimeError("Empty batch plan item")
        if item.pad_size > 0:
            parts[-1] = self._pad_batch_part(parts[-1], item.pad_size)
        return PackedSeqBatch(parts=tuple(parts))

    def __iter__(self):
        plan = self._build_rank_plan()
        if not plan:
            return
        if self.loader_threads <= 1 and self.prefetch_batches <= 1:
            for item in plan:
                yield self._load_batch_from_plan_item(item)
            return

        submit_ahead = max(self.prefetch_batches, self.loader_threads)
        futures: List[Future] = []
        next_idx = 0
        with ThreadPoolExecutor(max_workers=self.loader_threads) as executor:
            while next_idx < len(plan) and len(futures) < submit_ahead:
                futures.append(executor.submit(self._load_batch_from_plan_item, plan[next_idx]))
                next_idx += 1
            while futures:
                fut = futures.pop(0)
                yield fut.result()
                if next_idx < len(plan):
                    futures.append(executor.submit(self._load_batch_from_plan_item, plan[next_idx]))
                    next_idx += 1


def move_batch_to_device(
    batch: PackedSeqBatch,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> PackedSeqBatch:
    moved_parts = []
    for part in batch.parts:
        moved_parts.append(
            PackedSeqBatchPart(
                span_x=part.span_x.to(device, non_blocking=copy_non_blocking),
                local_end=part.local_end.to(device, non_blocking=copy_non_blocking),
                y=part.y.to(device, non_blocking=copy_non_blocking),
                w=part.w.to(device, non_blocking=copy_non_blocking),
            )
        )
    return PackedSeqBatch(parts=tuple(moved_parts))


class DeviceTransferPrefetchLoader:
    def __init__(
        self,
        loader,
        device: torch.device,
        prefetch_batches: int = 2,
        copy_non_blocking: bool = False,
    ):
        self.loader = loader
        self.device = device
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.copy_non_blocking = bool(copy_non_blocking)

    def __len__(self) -> int:
        return len(self.loader)

    def __getattr__(self, name: str):
        return getattr(self.loader, name)

    def set_epoch(self, epoch: int) -> None:
        if hasattr(self.loader, "set_epoch"):
            self.loader.set_epoch(epoch)

    def __iter__(self):
        if self.device.type != "cuda":
            for batch in self.loader:
                yield batch
            return
        result_q: queue.Queue = queue.Queue(maxsize=self.prefetch_batches)
        sentinel = object()
        device = self.device
        stop_event = threading.Event()

        def put_result(item, event) -> bool:
            while not stop_event.is_set():
                try:
                    result_q.put((item, event), timeout=0.1)
                    return True
                except queue.Full:
                    continue
            return False

        def worker():
            try:
                torch.cuda.set_device(device)
                stream = torch.cuda.Stream(device=device)
                for batch in self.loader:
                    if stop_event.is_set():
                        break
                    with torch.cuda.stream(stream):
                        moved = move_batch_to_device(
                            batch,
                            device=device,
                            copy_non_blocking=self.copy_non_blocking,
                        )
                        event = torch.cuda.Event()
                        event.record(stream)
                    if not put_result(moved, event):
                        return
                put_result(sentinel, None)
            except Exception as exc:  # pragma: no cover - worker thread failure propagation
                put_result(exc, None)

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        current_stream = torch.cuda.current_stream(device)
        try:
            while True:
                item, event = result_q.get()
                if item is sentinel:
                    break
                if isinstance(item, Exception):
                    raise item
                current_stream.wait_event(event)
                for part in item.parts:
                    part.span_x.record_stream(current_stream)
                    part.local_end.record_stream(current_stream)
                    part.y.record_stream(current_stream)
                    part.w.record_stream(current_stream)
                yield item
        finally:
            stop_event.set()
            thread.join()


class MaterializeDevicePrefetchLoader:
    def __init__(
        self,
        loader,
        device: torch.device,
        seq_offsets: torch.Tensor,
        feat_mean: torch.Tensor,
        feat_std: torch.Tensor,
        prefetch_batches: int = 2,
        copy_non_blocking: bool = False,
    ):
        self.loader = loader
        self.device = device
        self.seq_offsets = seq_offsets
        self.feat_mean = feat_mean
        self.feat_std = feat_std
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.copy_non_blocking = bool(copy_non_blocking)

    def __len__(self) -> int:
        return len(self.loader)

    def __getattr__(self, name: str):
        return getattr(self.loader, name)

    def set_epoch(self, epoch: int) -> None:
        if hasattr(self.loader, "set_epoch"):
            self.loader.set_epoch(epoch)

    def __iter__(self):
        if self.device.type != "cuda":
            for batch in self.loader:
                xb, yb, wb = materialize_batch(
                    batch,
                    seq_offsets=self.seq_offsets,
                    feat_mean=self.feat_mean,
                    feat_std=self.feat_std,
                    device=self.device,
                    copy_non_blocking=self.copy_non_blocking,
                )
                yield MaterializedSeqBatch(xb=xb, y=yb, w=wb)
            return
        result_q: queue.Queue = queue.Queue(maxsize=self.prefetch_batches)
        sentinel = object()
        device = self.device
        stop_event = threading.Event()

        def put_result(item, event) -> bool:
            while not stop_event.is_set():
                try:
                    result_q.put((item, event), timeout=0.1)
                    return True
                except queue.Full:
                    continue
            return False

        def worker():
            try:
                torch.cuda.set_device(device)
                stream = torch.cuda.Stream(device=device)
                for batch in self.loader:
                    if stop_event.is_set():
                        break
                    with torch.cuda.stream(stream):
                        xb, yb, wb = materialize_batch(
                            batch,
                            seq_offsets=self.seq_offsets,
                            feat_mean=self.feat_mean,
                            feat_std=self.feat_std,
                            device=device,
                            copy_non_blocking=self.copy_non_blocking,
                        )
                        event = torch.cuda.Event()
                        event.record(stream)
                    prefetched = MaterializedSeqBatch(xb=xb, y=yb, w=wb)
                    if not put_result(prefetched, event):
                        return
                put_result(sentinel, None)
            except Exception as exc:  # pragma: no cover - worker thread failure propagation
                put_result(exc, None)

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        current_stream = torch.cuda.current_stream(device)
        try:
            while True:
                item, event = result_q.get()
                if item is sentinel:
                    break
                if isinstance(item, Exception):
                    raise item
                current_stream.wait_event(event)
                item.xb.record_stream(current_stream)
                item.y.record_stream(current_stream)
                item.w.record_stream(current_stream)
                yield item
        finally:
            stop_event.set()
            thread.join()


class ParallelCNN1D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        branch_channels = out_channels // 3
        self.conv3 = nn.Conv1d(in_channels, branch_channels, kernel_size=3, padding=1)
        self.conv5 = nn.Conv1d(in_channels, branch_channels, kernel_size=5, padding=2)
        self.conv7 = nn.Conv1d(in_channels, out_channels - 2 * branch_channels, kernel_size=7, padding=3)
        self.act = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_t = x.transpose(1, 2)
        c3 = self.conv3(x_t)
        c5 = self.conv5(x_t)
        c7 = self.conv7(x_t)
        out = torch.cat([c3, c5, c7], dim=1)
        return self.act(out).transpose(1, 2)


class FeatureChannelGate(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 64, init_bias: float = 2.0):
        super().__init__()
        hidden_dim = max(1, int(hidden_dim))
        self.norm = nn.LayerNorm(input_dim)
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.act = nn.SiLU()
        self.fc2 = nn.Linear(hidden_dim, input_dim)
        nn.init.zeros_(self.fc2.weight)
        nn.init.constant_(self.fc2.bias, float(init_bias))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pooled = x.mean(dim=1)
        gate = self.fc2(self.act(self.fc1(self.norm(pooled))))
        gate = torch.sigmoid(gate).unsqueeze(1)
        return x * gate


class GRURegressor(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        num_layers: int,
        dropout: float,
        pooling: str = "last",
        bidirectional: bool = False,
        use_cnn1d: bool = False,
        input_gate_hidden_dim: int = 0,
        input_gate_bias: float = 2.0,
    ):
        super().__init__()
        self.pooling = str(pooling).lower()
        if self.pooling not in {"last", "attn"}:
            raise ValueError(f"Unsupported pooling: {pooling}")
        self.bidirectional = bool(bidirectional)
        self.use_cnn1d = bool(use_cnn1d)
        self.input_gate_hidden_dim = max(0, int(input_gate_hidden_dim))
        self.output_dim = int(hidden_dim) * (2 if self.bidirectional else 1)
        if self.input_gate_hidden_dim > 0:
            self.input_gate = FeatureChannelGate(
                input_dim=input_dim,
                hidden_dim=self.input_gate_hidden_dim,
                init_bias=float(input_gate_bias),
            )
        else:
            self.input_gate = nn.Identity()

        if self.use_cnn1d:
            self.cnn = ParallelCNN1D(input_dim, hidden_dim)
            rnn_input_dim = hidden_dim
        else:
            self.cnn = nn.Identity()
            rnn_input_dim = input_dim

        self.rnn = nn.GRU(
            input_size=rnn_input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            dropout=dropout if num_layers > 1 else 0.0,
            batch_first=True,
            bidirectional=self.bidirectional,
        )
        if self.pooling == "attn":
            self.attn_norm = nn.LayerNorm(self.output_dim)
            self.attn_proj = nn.Linear(self.output_dim, 1)
        head_hidden_dim = max(1, self.output_dim // 2)
        self.head = nn.Sequential(
            nn.LayerNorm(self.output_dim),
            nn.Linear(self.output_dim, head_hidden_dim),
            nn.ReLU(),
            nn.Linear(head_hidden_dim, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_gate(x)
        if self.use_cnn1d:
            x = self.cnn(x)
        out, _ = self.rnn(x)
        if self.pooling == "attn":
            score = self.attn_proj(self.attn_norm(out)).squeeze(-1)
            weight = torch.softmax(score, dim=1).unsqueeze(-1)
            pooled = (out * weight).sum(dim=1)
        else:
            pooled = out[:, -1, :]
        return self.head(pooled).squeeze(-1)


def weighted_mse(pred: torch.Tensor, y: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    w = torch.clamp(w, min=0.0)
    w_norm = w / torch.clamp(w.mean(), min=1e-6)
    return (w_norm * (pred - y) ** 2).mean()


def build_val_epoch_checkpoint_path(out_root: Path, epoch: int) -> Path:
    return out_root / f"gru_seq_memmap_val_epoch_{int(epoch):03d}.pt"


def refresh_top_val_checkpoints(
    top_val_checkpoints: List[dict],
    candidate_epoch: int,
    candidate_val_metrics: dict,
    keep_topk: int,
    out_root: Path,
) -> Tuple[List[dict], bool, List[dict]]:
    keep_topk = max(1, int(keep_topk))
    candidate = {
        "epoch": int(candidate_epoch),
        "val_ic": float(candidate_val_metrics["ic"]),
        "val_unweighted_ic": float(candidate_val_metrics.get("unweighted_ic", candidate_val_metrics["ic"])),
        "val_weighted_ic": float(candidate_val_metrics.get("weighted_ic", candidate_val_metrics["ic"])),
        "val_loss": float(candidate_val_metrics["loss"]),
        "path": str(build_val_epoch_checkpoint_path(out_root, candidate_epoch)),
    }
    updated = list(top_val_checkpoints)
    updated.append(candidate)
    updated.sort(key=lambda item: (-float(item.get("val_weighted_ic", item["val_ic"])), int(item["epoch"])))
    kept = [dict(item) for item in updated[:keep_topk]]
    dropped = [dict(item) for item in updated[keep_topk:]]
    kept_epochs = {int(item["epoch"]) for item in kept}
    entered_topk = int(candidate_epoch) in kept_epochs
    for rank_i, item in enumerate(kept, start=1):
        item["rank"] = int(rank_i)
    return kept, entered_topk, dropped


class CorrStats:
    def __init__(self):
        self.buf: torch.Tensor | None = None

    def update(self, pred: torch.Tensor, y: torch.Tensor, w: torch.Tensor | None = None):
        p = pred.detach().float().reshape(-1)
        t = y.detach().float().reshape(-1)
        mask = torch.isfinite(p) & torch.isfinite(t)
        if w is not None:
            ww = w.detach().float().reshape(-1)
            mask = mask & torch.isfinite(ww) & (ww > 0)
        else:
            ww = None
        zero = torch.zeros_like(p)
        p = torch.where(mask, p, zero).to(dtype=torch.float64)
        t = torch.where(mask, t, zero).to(dtype=torch.float64)
        d = p - t
        if ww is None:
            weight = mask.to(dtype=torch.float64)
        else:
            weight = torch.where(mask, ww, zero).to(dtype=torch.float64)
        sums = torch.stack(
            [
                mask.to(dtype=torch.float64).sum(),
                p.sum(),
                t.sum(),
                (p * p).sum(),
                (t * t).sum(),
                (p * t).sum(),
                torch.abs(d).sum(),
                (d * d).sum(),
                weight.sum(),
                (weight * p).sum(),
                (weight * t).sum(),
                (weight * p * p).sum(),
                (weight * t * t).sum(),
                (weight * p * t).sum(),
                (weight * torch.abs(d)).sum(),
                (weight * d * d).sum(),
            ]
        )
        if self.buf is None:
            self.buf = sums
        else:
            self.buf = self.buf + sums

    def to_tensor(self, device: torch.device) -> torch.Tensor:
        if self.buf is None:
            return torch.zeros((16,), dtype=torch.float64, device=device)
        return self.buf.to(device=device, dtype=torch.float64)


def corr_from_sums(sum_x: float, sum_y: float, sum_xx: float, sum_yy: float, sum_xy: float, denom_weight: float) -> float:
    if (not np.isfinite(denom_weight)) or denom_weight <= 0:
        return float("nan")
    mean_x = sum_x / denom_weight
    mean_y = sum_y / denom_weight
    var_x = max(sum_xx / denom_weight - mean_x * mean_x, 1e-12)
    var_y = max(sum_yy / denom_weight - mean_y * mean_y, 1e-12)
    cov_xy = sum_xy / denom_weight - mean_x * mean_y
    return float(cov_xy / math.sqrt(var_x * var_y))


def metrics_from_tensor(t: torch.Tensor) -> dict:
    values = [float(x) for x in t.tolist()]
    if len(values) < 8:
        raise ValueError(f"metrics_from_tensor expects at least 8 values, got {len(values)}")
    n, sum_p, sum_y, sum_pp, sum_yy, sum_py, sum_abs, sum_sq = values[:8]
    if (not np.isfinite(n)) or n <= 0:
        return {
            "mse": float("nan"),
            "rmse": float("nan"),
            "mae": float("nan"),
            "ic": float("nan"),
            "unweighted_mse": float("nan"),
            "unweighted_rmse": float("nan"),
            "unweighted_mae": float("nan"),
            "unweighted_ic": float("nan"),
            "weighted_mse": float("nan"),
            "weighted_rmse": float("nan"),
            "weighted_mae": float("nan"),
            "weighted_ic": float("nan"),
            "weight_sum": 0.0,
            "n": 0,
        }
    mse = sum_sq / n
    mae = sum_abs / n
    ic = corr_from_sums(sum_p, sum_y, sum_pp, sum_yy, sum_py, n)
    if len(values) >= 16:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy, sum_wabs, sum_wsq = values[8:16]
    elif len(values) >= 14:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy = values[8:14]
        sum_wabs, sum_wsq = sum_abs, sum_sq
    else:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy = n, sum_p, sum_y, sum_pp, sum_yy, sum_py
        sum_wabs, sum_wsq = sum_abs, sum_sq
    weighted_ic = corr_from_sums(sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy, weight_sum)
    weighted_mse = (sum_wsq / weight_sum) if weight_sum > 0 else float("nan")
    weighted_mae = (sum_wabs / weight_sum) if weight_sum > 0 else float("nan")
    weighted_rmse = math.sqrt(weighted_mse) if np.isfinite(weighted_mse) and weighted_mse >= 0 else float("nan")
    return {
        "mse": float(weighted_mse),
        "rmse": float(weighted_rmse),
        "mae": float(weighted_mae),
        "ic": float(weighted_ic),
        "unweighted_mse": float(mse),
        "unweighted_rmse": float(math.sqrt(mse)),
        "unweighted_mae": float(mae),
        "unweighted_ic": float(ic),
        "weighted_mse": float(weighted_mse),
        "weighted_rmse": float(weighted_rmse),
        "weighted_mae": float(weighted_mae),
        "weighted_ic": float(weighted_ic),
        "weight_sum": float(weight_sum),
        "n": int(n),
    }


def compute_feature_stats(
    train_days: List[Path],
    prefer_fp16: bool,
    sample_stride: int = 50,
    chunk_sample_rows: int = 250_000,
) -> Tuple[np.ndarray, np.ndarray]:
    s = None
    s2 = None
    n = 0
    for day in train_days:
        store = load_day_store(day, prefer_fp16=prefer_fp16)
        chunk_span = max(int(sample_stride), int(sample_stride) * max(1, int(chunk_sample_rows)))
        for start in range(0, store.n_rows, chunk_span):
            stop = min(store.n_rows, start + chunk_span)
            # Stream smaller memmap slices to reduce pressure and avoid giant one-shot reads.
            arr = np.array(store.x[start:stop:sample_stride], dtype=np.float32, copy=True)
            if arr.size == 0:
                continue
            arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
            if s is None:
                s = arr.sum(axis=0, dtype=np.float64)
                s2 = (arr * arr).sum(axis=0, dtype=np.float64)
            else:
                s += arr.sum(axis=0, dtype=np.float64)
                s2 += (arr * arr).sum(axis=0, dtype=np.float64)
            n += arr.shape[0]
    mean = (s / max(n, 1)).astype(np.float32)
    var = (s2 / max(n, 1) - mean.astype(np.float64) ** 2).astype(np.float32)
    std = np.sqrt(np.clip(var, 1e-8, None)).astype(np.float32)
    return mean, std


def build_feature_stats_cache_path(
    output_root: Path,
    row_root: Path,
    train_days: List[Path],
    prefer_fp16: bool,
    sample_stride: int,
) -> Path:
    cache_dir = ensure_dir(output_root / "training_seq" / "_feature_stats_cache")
    key_src = "\n".join(
        [
            str(row_root),
            str(bool(prefer_fp16)),
            str(int(sample_stride)),
            *[d.name for d in train_days],
        ]
    )
    key = hashlib.sha1(key_src.encode("utf-8")).hexdigest()[:16]
    return cache_dir / f"feature_stats_{key}.npz"


def load_feature_stats_cache(path: Path) -> Tuple[np.ndarray, np.ndarray]:
    with np.load(path) as arr:
        mean = arr["mean"].astype(np.float32, copy=False)
        std = arr["std"].astype(np.float32, copy=False)
    return mean, std


def wait_for_feature_stats_cache(path: Path, timeout_seconds: int = 7200) -> Tuple[np.ndarray, np.ndarray]:
    deadline = time.time() + max(1, int(timeout_seconds))
    last_err: Exception | None = None
    while time.time() < deadline:
        if path.exists() and path.stat().st_size > 0:
            try:
                return load_feature_stats_cache(path)
            except Exception as exc:  # pragma: no cover - transient partial-write case
                last_err = exc
        time.sleep(2.0)
    if last_err is not None:
        raise TimeoutError(f"Timed out waiting for feature stats cache {path}: {last_err}") from last_err
    raise TimeoutError(f"Timed out waiting for feature stats cache {path}")


def evaluate(
    model,
    loader,
    feat_mean,
    feat_std,
    seq_offsets,
    accelerator,
    split_name: str = "eval",
    progress_parts: int = 4,
    copy_non_blocking: bool = False,
    collect_timing: bool = False,
) -> dict:
    model.eval()
    stats = CorrStats()
    timing_labels = [
        "batch_wait_sec",
        "materialize_sec",
        "forward_sec",
        "stats_update_sec",
        "step_total_sec",
    ]
    timing_sums = np.zeros((len(timing_labels),), dtype=np.float64)
    local_batches = 0
    total_batches = 0
    try:
        total_batches = len(loader)
    except TypeError:
        total_batches = 0
    progress_step = max(1, total_batches // max(1, int(progress_parts))) if total_batches > 0 else 0
    with torch.no_grad():
        loader_iter = iter(loader)
        batch_i = 0
        while True:
            t_step0 = time.perf_counter()
            t0 = time.perf_counter()
            try:
                batch = next(loader_iter)
            except StopIteration:
                break
            t1 = time.perf_counter()
            xb, yb, wb = resolve_model_batch(
                batch,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                device=accelerator.device,
                copy_non_blocking=copy_non_blocking,
            )
            if collect_timing:
                sync_if_needed(accelerator.device)
            t2 = time.perf_counter()
            with accelerator.autocast():
                pred = model(xb)
                loss = weighted_mse(pred, yb, wb)
            if collect_timing:
                sync_if_needed(accelerator.device)
            t3 = time.perf_counter()
            stats.update(pred, yb, wb)
            if collect_timing:
                sync_if_needed(accelerator.device)
            t4 = time.perf_counter()
            local_batches += 1
            batch_i += 1
            if collect_timing:
                timing_sums += np.asarray(
                    [
                        t1 - t0,
                        t2 - t1,
                        t3 - t2,
                        t4 - t3,
                        t4 - t_step0,
                    ],
                    dtype=np.float64,
                )
            if (
                accelerator.is_main_process
                and total_batches > 0
                and (batch_i % progress_step == 0 or batch_i == total_batches)
            ):
                print(f"[{split_name}] progress {batch_i}/{total_batches}", flush=True)
    stats_t = accelerator.reduce(stats.to_tensor(accelerator.device), reduction="sum")
    m = metrics_from_tensor(stats_t)
    m["loss"] = float(m.get("weighted_mse", float("nan")))
    if collect_timing:
        m["_timing"] = summarize_timing_across_ranks(
            accelerator=accelerator,
            labels=timing_labels,
            count=local_batches,
            sums=timing_sums,
            count_key="measured_batches",
            global_batch=int(loader.batch_size) * int(accelerator.num_processes),
        )
    return m


def materialize_batch(
    batch: PackedSeqBatch,
    seq_offsets: torch.Tensor,
    feat_mean: torch.Tensor,
    feat_std: torch.Tensor,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    xb_parts = []
    y_parts = []
    w_parts = []
    seq_offsets_cpu: torch.Tensor | None = None
    for part in batch.parts:
        yb = part.y.to(device, non_blocking=copy_non_blocking)
        wb = part.w.to(device, non_blocking=copy_non_blocking)
        span_rows = int(part.span_x.shape[0])
        batch_rows = int(part.local_end.shape[0])
        use_cpu_materialize = (
            device.type == "cuda"
            and part.span_x.device.type == "cpu"
            and span_rows > max(1, batch_rows) * 2
        )
        if use_cpu_materialize:
            if seq_offsets_cpu is None:
                seq_offsets_cpu = seq_offsets.detach().to(device="cpu")
            seq_idx_cpu = part.local_end[:, None] + seq_offsets_cpu[None, :]
            xb_cpu = part.span_x[seq_idx_cpu].float()
            if copy_non_blocking:
                xb_cpu = xb_cpu.pin_memory()
            xb = xb_cpu.to(device, non_blocking=copy_non_blocking)
        else:
            span_x = part.span_x.to(device, non_blocking=copy_non_blocking)
            local_end = part.local_end.to(device, non_blocking=copy_non_blocking)
            seq_idx = local_end[:, None] + seq_offsets[None, :]
            xb = span_x[seq_idx].float()
        xb = (xb - feat_mean) / feat_std
        # Keep GRU inputs in a standard dense layout so cuDNN does not fall back
        # to a slower path for batches assembled from wide memmap spans.
        xb = xb.contiguous()
        xb_parts.append(xb)
        y_parts.append(yb)
        w_parts.append(wb)
    if len(xb_parts) == 1:
        return xb_parts[0], y_parts[0], w_parts[0]
    return torch.cat(xb_parts, dim=0), torch.cat(y_parts, dim=0), torch.cat(w_parts, dim=0)


def resolve_model_batch(
    batch,
    seq_offsets: torch.Tensor,
    feat_mean: torch.Tensor,
    feat_std: torch.Tensor,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if isinstance(batch, MaterializedSeqBatch):
        return batch.xb, batch.y, batch.w
    return materialize_batch(
        batch,
        seq_offsets=seq_offsets,
        feat_mean=feat_mean,
        feat_std=feat_std,
        device=device,
        copy_non_blocking=copy_non_blocking,
    )


def sync_if_needed(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def warmup_collective(accelerator) -> float:
    if int(getattr(accelerator, "num_processes", 1)) <= 1:
        return 0.0
    t0 = time.perf_counter()
    dummy = torch.zeros((1,), dtype=torch.float32, device=accelerator.device)
    _ = accelerator.reduce(dummy, reduction="sum")
    sync_if_needed(accelerator.device)
    return float(time.perf_counter() - t0)


def summarize_timing_across_ranks(
    accelerator: Accelerator,
    labels: List[str],
    count: int,
    sums: np.ndarray,
    count_key: str,
    global_batch: int | None = None,
) -> dict:
    local = torch.tensor([float(count), *sums.tolist()], dtype=torch.float64, device=accelerator.device)[None, :]
    gathered = accelerator.gather(local)
    empty = {
        count_key: int(count),
        "per_rank": [],
        "global_step_sec_estimate": float("nan"),
        "global_samples_per_sec_estimate": float("nan"),
    }
    if not accelerator.is_main_process:
        return empty
    gathered_np = gathered.detach().cpu().numpy()
    per_rank = []
    max_step_sec = 0.0
    step_label = labels[-1] if labels else None
    for rank_i, row in enumerate(gathered_np):
        rank_count = max(1.0, float(row[0]))
        stats = {"rank": int(rank_i), count_key: int(row[0])}
        for label_i, label in enumerate(labels, start=1):
            stats[label] = float(row[label_i] / rank_count)
        if step_label is not None:
            max_step_sec = max(max_step_sec, float(stats[step_label]))
        per_rank.append(stats)
    out = {
        count_key: int(count),
        "per_rank": per_rank,
        "global_step_sec_estimate": float(max_step_sec) if step_label is not None else float("nan"),
        "global_samples_per_sec_estimate": float("nan"),
    }
    if step_label is not None and global_batch is not None:
        out["global_samples_per_sec_estimate"] = float(global_batch / max(max_step_sec, 1e-9))
    return out


def profile_train_steps(
    model,
    optimizer,
    loader,
    accelerator,
    feat_mean,
    feat_std,
    seq_offsets,
    copy_non_blocking: bool,
    warmup_steps: int,
    profile_steps: int,
    grad_accum_steps: int = 1,
) -> dict:
    total_batches = 0
    try:
        total_batches = len(loader)
    except TypeError:
        total_batches = 0
    warmup = max(0, int(warmup_steps))
    active = max(0, int(profile_steps))
    if total_batches <= 0 or active <= 0:
        return {
            "warmup_steps": warmup,
            "profile_steps": active,
            "measured_steps": 0,
            "per_rank": [],
            "global_step_sec_estimate": float("nan"),
            "global_samples_per_sec_estimate": float("nan"),
        }
    run_steps = min(total_batches, warmup + active)
    measured_steps = max(0, run_steps - warmup)
    labels = [
        "batch_wait_sec",
        "materialize_sec",
        "forward_sec",
        "backward_sec",
        "optimizer_step_sec",
        "step_total_sec",
    ]
    sums = np.zeros((len(labels),), dtype=np.float64)
    model.train()
    train_iter = iter(loader)
    warmup_collective(accelerator)
    if accelerator.device.type == "cuda":
        sync_if_needed(accelerator.device)
        torch.cuda.reset_peak_memory_stats(accelerator.device)
    accum_steps = max(1, int(grad_accum_steps))
    accum_micro = 0
    optimizer.zero_grad(set_to_none=True)
    for step_i in range(1, run_steps + 1):
        if accelerator.device.type == "cuda" and step_i == warmup + 1:
            sync_if_needed(accelerator.device)
            torch.cuda.reset_peak_memory_stats(accelerator.device)
        t_step0 = time.perf_counter()
        t0 = time.perf_counter()
        batch = next(train_iter)
        t1 = time.perf_counter()
        accum_micro += 1
        last_batch = step_i == run_steps
        should_step = accum_micro >= accum_steps or last_batch
        group_total = min(accum_steps, accum_micro + max(0, run_steps - step_i))
        xb, yb, wb = resolve_model_batch(
            batch,
            seq_offsets=seq_offsets,
            feat_mean=feat_mean,
            feat_std=feat_std,
            device=accelerator.device,
            copy_non_blocking=copy_non_blocking,
        )
        sync_if_needed(accelerator.device)
        t2 = time.perf_counter()
        sync_ctx = model.no_sync() if hasattr(model, "no_sync") and not should_step else nullcontext()
        with sync_ctx:
            with accelerator.autocast():
                pred = model(xb)
                loss_raw = weighted_mse(pred, yb, wb)
                loss = loss_raw / float(group_total)
            sync_if_needed(accelerator.device)
            t3 = time.perf_counter()
            accelerator.backward(loss)
            sync_if_needed(accelerator.device)
            t4 = time.perf_counter()
        if should_step:
            accelerator.step_optimizer(optimizer)
            optimizer.zero_grad(set_to_none=True)
            accum_micro = 0
        sync_if_needed(accelerator.device)
        t5 = time.perf_counter()
        if step_i > warmup:
            sums += np.asarray(
                [
                    t1 - t0,
                    t2 - t1,
                    t3 - t2,
                    t4 - t3,
                    t5 - t4,
                    t5 - t_step0,
                ],
                dtype=np.float64,
            )
        if accelerator.is_main_process and (step_i == run_steps or step_i % max(1, run_steps // 4) == 0):
            print(f"[profile-train] step {step_i}/{run_steps} loss={loss_raw.detach().float().item():.6f}", flush=True)
    peak_allocated_bytes = 0.0
    peak_reserved_bytes = 0.0
    if accelerator.device.type == "cuda":
        sync_if_needed(accelerator.device)
        peak_allocated_bytes = float(torch.cuda.max_memory_allocated(accelerator.device))
        peak_reserved_bytes = float(torch.cuda.max_memory_reserved(accelerator.device))
    global_batch = int(loader.batch_size) * int(accelerator.num_processes)
    summary = summarize_timing_across_ranks(
        accelerator=accelerator,
        labels=labels,
        count=measured_steps,
        sums=sums,
        count_key="measured_steps",
        global_batch=global_batch,
    )
    summary["warmup_steps"] = warmup
    summary["profile_steps"] = active
    if accelerator.device.type == "cuda":
        local_mem = torch.tensor(
            [peak_allocated_bytes, peak_reserved_bytes],
            dtype=torch.float64,
            device=accelerator.device,
        )[None, :]
        gathered_mem = accelerator.gather(local_mem)
        if accelerator.is_main_process:
            gathered_mem_np = gathered_mem.detach().cpu().numpy()
            peak_allocated_max = 0.0
            peak_reserved_max = 0.0
            for rank_i, row in enumerate(gathered_mem_np):
                allocated_bytes = float(row[0])
                reserved_bytes = float(row[1])
                peak_allocated_max = max(peak_allocated_max, allocated_bytes)
                peak_reserved_max = max(peak_reserved_max, reserved_bytes)
                if rank_i < len(summary.get("per_rank", [])):
                    summary["per_rank"][rank_i]["peak_allocated_bytes"] = allocated_bytes
                    summary["per_rank"][rank_i]["peak_reserved_bytes"] = reserved_bytes
                    summary["per_rank"][rank_i]["peak_allocated_gib"] = allocated_bytes / float(1024**3)
                    summary["per_rank"][rank_i]["peak_reserved_gib"] = reserved_bytes / float(1024**3)
            summary["peak_allocated_bytes_max"] = peak_allocated_max
            summary["peak_reserved_bytes_max"] = peak_reserved_max
            summary["peak_allocated_gib_max"] = peak_allocated_max / float(1024**3)
            summary["peak_reserved_gib_max"] = peak_reserved_max / float(1024**3)
    return summary


def run(args):
    run_t0 = time.perf_counter()
    phase_timings: Dict[str, float] = {}
    phase_details: Dict[str, dict] = {}
    data_prepare_t0 = time.perf_counter()
    set_global_seed(args.seed)
    cfg = load_config(args.config)
    train_cfg = dict(cfg["train"])
    if args.override_epochs > 0:
        train_cfg["epochs"] = args.override_epochs
    if args.override_lr > 0:
        train_cfg["learning_rate"] = args.override_lr
    if args.override_batch_size > 0:
        train_cfg["batch_size"] = args.override_batch_size

    row_root = resolve_row_root(cfg, args.cache_name, row_root_override=args.row_root_override)
    profile_phases = bool(args.profile_phases)
    profile_warmup_steps = max(0, int(args.profile_warmup_steps))
    profile_steps = max(0, int(args.profile_steps))
    stop_after_profile = bool(args.stop_after_profile)
    profile_only = bool(stop_after_profile and profile_steps > 0)
    grad_accum_steps = max(1, int(args.grad_accum_steps))
    split_view_meta = load_split_view_meta(row_root)
    if split_view_meta is not None:
        train_days = resolve_named_day_dirs(row_root, split_view_meta.get("train_days"))
        valid_days = resolve_named_day_dirs(row_root, split_view_meta.get("valid_days"))
        test_days = resolve_named_day_dirs(row_root, split_view_meta.get("test_days"))
        train_days = limit_days(train_days, args.train_day_limit)
        valid_days = limit_days(valid_days, args.valid_day_limit)
        test_days = limit_days(test_days, args.test_day_limit)
        train_pool_days = list(train_days) + list(valid_days)
        days = list(train_days) + list(valid_days) + list(test_days)
        split_ratio = float(args.valid_split_ratio) if args.valid_split_ratio > 0.0 else float(train_cfg["train_split_ratio"])
        train_cfg["train_split_ratio"] = split_ratio
    else:
        days = list_ready_days(row_root)
        use_explicit_split = bool(
            args.train_pool_start_date or args.train_pool_end_date or args.test_start_date or args.test_end_date
        )
        if args.max_days > 0 and not use_explicit_split:
            days = days[: args.max_days]
        if args.test_start_date or args.test_end_date:
            test_days = filter_days_by_date(days, args.test_start_date, args.test_end_date)
        else:
            test_days = []
        test_day_names = {d.name for d in test_days}
        raw_train_pool_days = filter_days_by_date(days, args.train_pool_start_date, args.train_pool_end_date)
        if args.train_pool_start_date or args.train_pool_end_date:
            train_pool_days = [d for d in raw_train_pool_days if d.name not in test_day_names]
        else:
            train_pool_days = [d for d in days if d.name not in test_day_names]
        split_ratio = float(args.valid_split_ratio) if args.valid_split_ratio > 0.0 else float(train_cfg["train_split_ratio"])
        train_cfg["train_split_ratio"] = split_ratio
        train_days, valid_days = split_days(train_pool_days, split_ratio)
        train_days = limit_days(train_days, args.train_day_limit)
        valid_days = limit_days(valid_days, args.valid_day_limit)
        test_days = limit_days(test_days, args.test_day_limit)
    if not train_pool_days:
        raise RuntimeError("train_pool_days is empty after applying date filters.")
    if len(train_days) == 0 or len(valid_days) == 0:
        raise RuntimeError("Need both train and valid day split.")
    print(
        f"[data] days_total={len(days)} train_pool_days={len(train_pool_days)} "
        f"train_days={len(train_days)} valid_days={len(valid_days)} test_days={len(test_days)}",
        flush=True,
    )
    print(f"[row-root] {row_root}", flush=True)
    if args.train_pool_start_date or args.train_pool_end_date:
        print(
            f"[data] train_pool_date_window=[{args.train_pool_start_date or 'min'}, "
            f"{args.train_pool_end_date or 'max'}]",
            flush=True,
        )
    if args.test_start_date or args.test_end_date:
        print(
            f"[data] test_date_window=[{args.test_start_date or 'min'}, {args.test_end_date or 'max'}]",
            flush=True,
        )
    data_prepare_local_sec = time.perf_counter() - data_prepare_t0
    accelerator_init_t0 = time.perf_counter()
    accelerator = build_runtime(
        args.runtime_backend,
        bool(args.use_amp),
        enable_static_graph=(grad_accum_steps == 1),
    )
    phase_timings["data_prepare_sec"] = float(data_prepare_local_sec)
    phase_timings["accelerator_init_sec"] = float(time.perf_counter() - accelerator_init_t0)
    loader_threads = max(1, int(args.num_workers))
    eval_batch_size = int(args.eval_batch_size) if int(args.eval_batch_size) > 0 else int(train_cfg["batch_size"])
    device_transfer_mode = str(args.device_transfer_mode).strip().lower()
    device_transfer_prefetch_batches = max(1, int(args.device_transfer_prefetch_batches))
    configure_split_view_weight_loading(bool(args.use_source_raw_weights))
    min_timecode = int(args.min_timecode)
    require_positive_weight = bool(args.require_positive_weight)
    # Keep host tensors pinned even on the threaded transfer path so H2D copies
    # can overlap with compute on the prefetch stream.
    use_pinned_transfer = bool(train_cfg.get("pin_memory", False))
    if int(args.loader_prefetch_batches) > 0:
        loader_prefetch = max(1, int(args.loader_prefetch_batches))
    else:
        loader_prefetch = max(1, int(train_cfg.get("prefetch_factor", 2)))
    train_batch_overlap = max(0, int(args.train_batch_overlap))
    skip_train_eval = bool(args.skip_train_eval)
    skip_test = bool(args.skip_test)
    print(
        f"[device-transfer] mode={device_transfer_mode} loader_pin_memory={int(use_pinned_transfer)} "
        f"prefetch_batches={device_transfer_prefetch_batches}",
        flush=True,
    )
    print(
        f"[sample-filter] min_timecode={min_timecode} require_positive_weight={int(require_positive_weight)} "
        f"use_source_raw_weights={int(bool(args.use_source_raw_weights))}",
        flush=True,
    )
    if train_batch_overlap > 0:
        print(
            f"[train-loader] batch_overlap={train_batch_overlap} batch_stride={int(train_cfg['batch_size']) - train_batch_overlap}",
            flush=True,
        )
    end_index_cache_dir = ensure_dir(Path(cfg["paths"]["output_root"]) / "training_seq" / "_end_index_cache")
    loader_init_t0 = time.perf_counter()
    train_loader = SeqMemmapBatchLoader(
        train_days,
        seq_len=args.seq_len,
        sample_stride=args.sample_stride,
        batch_size=int(train_cfg["batch_size"]),
        shuffle=True,
        max_samples=args.max_samples,
        prefer_fp16=bool(args.prefer_fp16),
        rank=accelerator.process_index,
        world_size=accelerator.num_processes,
        seed=args.seed,
        pin_memory=use_pinned_transfer,
        loader_threads=loader_threads,
        prefetch_batches=loader_prefetch,
        pad_last_batch=(accelerator.num_processes > 1),
        batch_overlap=train_batch_overlap,
        index_cache_dir=end_index_cache_dir,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    phase_timings["train_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    train_eval_loader = (
        SeqMemmapBatchLoader(
            train_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=args.max_samples,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if not profile_only and not skip_train_eval and train_batch_overlap > 0
        else None
    )
    phase_timings["train_eval_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    valid_loader = (
        SeqMemmapBatchLoader(
            valid_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=args.max_samples,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if not profile_only
        else None
    )
    phase_timings["valid_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    test_loader = (
        SeqMemmapBatchLoader(
            test_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=-1,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if test_days and not skip_test and not profile_only
        else None
    )
    phase_timings["test_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    stats_stride = max(20, args.sample_stride)
    feature_stats_cache = build_feature_stats_cache_path(
        output_root=Path(cfg["paths"]["output_root"]),
        row_root=row_root,
        train_days=train_days,
        prefer_fp16=bool(args.prefer_fp16),
        sample_stride=stats_stride,
    )
    feature_stats_t0 = time.perf_counter()
    if accelerator.is_main_process:
        if feature_stats_cache.exists() and feature_stats_cache.stat().st_size > 0:
            print(f"[feature-stats] reused_from={feature_stats_cache}", flush=True)
            feat_mean_np, feat_std_np = load_feature_stats_cache(feature_stats_cache)
        else:
            print(
                f"[feature-stats] computing cache={feature_stats_cache} sample_stride={stats_stride}",
                flush=True,
            )
            feat_mean_np, feat_std_np = compute_feature_stats(
                train_days, prefer_fp16=bool(args.prefer_fp16), sample_stride=stats_stride
            )
            tmp_path = feature_stats_cache.with_suffix(f".tmp.{int(time.time())}.npz")
            np.savez(tmp_path, mean=feat_mean_np, std=feat_std_np)
            tmp_path.replace(feature_stats_cache)
            print(f"[feature-stats] saved_cache={feature_stats_cache}", flush=True)
    else:
        print(f"[feature-stats] waiting_for_cache={feature_stats_cache}", flush=True)
        feat_mean_np, feat_std_np = wait_for_feature_stats_cache(feature_stats_cache)
        print(f"[feature-stats] loaded_cache={feature_stats_cache}", flush=True)
    accelerator.wait_for_everyone()
    phase_timings["feature_stats_sec"] = float(time.perf_counter() - feature_stats_t0)

    model_init_t0 = time.perf_counter()
    model = GRURegressor(
        input_dim=200,
        hidden_dim=args.hidden_dim,
        num_layers=args.num_layers,
        dropout=args.dropout,
        pooling=args.pooling,
        bidirectional=bool(args.bidirectional),
        use_cnn1d=bool(getattr(args, "use_cnn1d", 0)),
        input_gate_hidden_dim=int(getattr(args, "input_gate_hidden_dim", 0)),
        input_gate_bias=float(getattr(args, "input_gate_bias", 2.0)),
    )
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=float(train_cfg["learning_rate"]), weight_decay=float(args.weight_decay)
    )
    phase_timings["model_optimizer_init_sec"] = float(time.perf_counter() - model_init_t0)

    prepare_t0 = time.perf_counter()
    if test_loader is not None:
        model, optimizer = accelerator.prepare(model, optimizer)
    else:
        model, optimizer = accelerator.prepare(model, optimizer)
    phase_timings["accelerator_prepare_sec"] = float(time.perf_counter() - prepare_t0)
    tensor_init_t0 = time.perf_counter()
    feat_mean = torch.from_numpy(feat_mean_np).to(accelerator.device)[None, None, :]
    feat_std = torch.from_numpy(feat_std_np).to(accelerator.device)[None, None, :]
    seq_offsets = torch.arange(-(args.seq_len - 1), 1, dtype=torch.long, device=accelerator.device)
    phase_timings["device_tensor_init_sec"] = float(time.perf_counter() - tensor_init_t0)
    device_prefetch_t0 = time.perf_counter()
    if device_transfer_mode == "thread_prefetch":
        train_loader = DeviceTransferPrefetchLoader(
            train_loader,
            device=accelerator.device,
            prefetch_batches=device_transfer_prefetch_batches,
            copy_non_blocking=use_pinned_transfer,
        )
        if train_eval_loader is not None:
            train_eval_loader = DeviceTransferPrefetchLoader(
                train_eval_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if valid_loader is not None:
            valid_loader = DeviceTransferPrefetchLoader(
                valid_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if test_loader is not None:
            test_loader = DeviceTransferPrefetchLoader(
                test_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
    elif device_transfer_mode == "materialize_prefetch":
        train_loader = MaterializeDevicePrefetchLoader(
            train_loader,
            device=accelerator.device,
            seq_offsets=seq_offsets,
            feat_mean=feat_mean,
            feat_std=feat_std,
            prefetch_batches=device_transfer_prefetch_batches,
            copy_non_blocking=use_pinned_transfer,
        )
        if train_eval_loader is not None:
            train_eval_loader = MaterializeDevicePrefetchLoader(
                train_eval_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if valid_loader is not None:
            valid_loader = MaterializeDevicePrefetchLoader(
                valid_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if test_loader is not None:
            test_loader = MaterializeDevicePrefetchLoader(
                test_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
    phase_timings["device_prefetch_wrap_sec"] = float(time.perf_counter() - device_prefetch_t0)

    output_init_t0 = time.perf_counter()
    out_root = ensure_dir(Path(cfg["paths"]["output_root"]) / "training_seq" / args.run_name)
    log_path = out_root / "train_log.csv"
    best_model_path = out_root / "gru_seq_memmap_best_val_ic.pt"
    save_topk_val_checkpoints = max(
        1,
        int(getattr(args, "save_topk_val_checkpoints", 1)),
        int(getattr(args, "test_checkpoint_rank", 1)),
    )
    selected_test_checkpoint_rank = max(1, int(getattr(args, "test_checkpoint_rank", 1)))
    if accelerator.is_main_process and not profile_only:
        with log_path.open("w", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow(
                [
                    "epoch",
                    "lr",
                    "train_loss",
                    "train_ic",
                    "train_unweighted_ic",
                    "train_weighted_ic",
                    "train_rmse",
                    "train_unweighted_rmse",
                    "val_loss",
                    "val_ic",
                    "val_unweighted_ic",
                    "val_weighted_ic",
                    "val_rmse",
                    "val_unweighted_rmse",
                    "val_mae",
                    "is_best",
                    "pooling",
                ]
            )
    phase_timings["output_init_sec"] = float(time.perf_counter() - output_init_t0)

    epochs = int(train_cfg["epochs"])
    epoch_lrs = build_epoch_lr_schedule(train_cfg, epochs=epochs)
    total_train_batches = 0
    try:
        total_train_batches = len(train_loader)
    except TypeError:
        total_train_batches = 0
    train_progress_step = (
        max(1, total_train_batches // max(1, int(args.train_progress_splits))) if total_train_batches > 0 else 0
    )
    best_epoch = 0
    best_val_m: Dict[str, float] | None = None
    final_val_m: Dict[str, float] | None = None
    top_val_checkpoints: List[dict] = []
    first_train_step_local: np.ndarray | None = None
    collective_warmup_sec = 0.0
    collective_warmup_done = False
    train_eval_timing: dict | None = None
    valid_eval_timing: dict | None = None
    test_eval_timing: dict | None = None
    if profile_only:
        train_loader.set_epoch(1)
        profile_summary = profile_train_steps(
            model=model,
            optimizer=optimizer,
            loader=train_loader,
            accelerator=accelerator,
            feat_mean=feat_mean,
            feat_std=feat_std,
            seq_offsets=seq_offsets,
            copy_non_blocking=use_pinned_transfer,
            warmup_steps=profile_warmup_steps,
            profile_steps=profile_steps,
            grad_accum_steps=grad_accum_steps,
        )
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            summary_path = out_root / "step_profile_summary.json"
            summary = {
                "run_name": args.run_name,
                "row_root": str(row_root),
                "world_size": int(accelerator.num_processes),
                "batch_size_per_rank": int(train_cfg["batch_size"]),
                "train_batch_overlap": train_batch_overlap,
                "runtime_backend": args.runtime_backend,
                "device_transfer_mode": device_transfer_mode,
                "min_timecode": int(min_timecode),
                "require_positive_weight": bool(require_positive_weight),
                "use_source_raw_weights": bool(args.use_source_raw_weights),
                "seq_len": int(args.seq_len),
                "sample_stride": int(args.sample_stride),
                "profile": profile_summary,
            }
            summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
            print(
                f"[profile-train] saved={summary_path} global_step_sec={profile_summary['global_step_sec_estimate']:.6f} "
                f"global_samples_per_sec={profile_summary['global_samples_per_sec_estimate']:.2f}",
                flush=True,
            )
        accelerator.close()
        return
    train_loop_t0 = time.perf_counter()
    for epoch in range(1, epochs + 1):
        current_lr = float(epoch_lrs[epoch - 1])
        set_optimizer_lr(optimizer, current_lr)
        epoch_train_t0 = time.perf_counter()
        model.train()
        train_loader.set_epoch(epoch)
        if accelerator.is_main_process:
            print(f"[epoch {epoch}] lr={current_lr:.8f}", flush=True)
        train_iter = iter(train_loader)
        if not collective_warmup_done:
            collective_warmup_sec = warmup_collective(accelerator)
            collective_warmup_done = True
            phase_timings["collective_warmup_sec"] = float(collective_warmup_sec)
        batch_i = 0
        accum_micro = 0
        optimizer.zero_grad(set_to_none=True)
        while True:
            fetch_t0 = time.perf_counter()
            try:
                batch = next(train_iter)
            except StopIteration:
                break
            fetch_t1 = time.perf_counter()
            batch_i += 1
            accum_micro += 1
            last_batch = total_train_batches > 0 and batch_i == total_train_batches
            should_step = accum_micro >= grad_accum_steps or last_batch
            group_total = min(grad_accum_steps, accum_micro + max(0, total_train_batches - batch_i))
            measure_first_step = bool(profile_phases and first_train_step_local is None and epoch == 1)
            xb, yb, wb = resolve_model_batch(
                batch,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                device=accelerator.device,
                copy_non_blocking=use_pinned_transfer,
            )
            if measure_first_step:
                sync_if_needed(accelerator.device)
                mat_t = time.perf_counter()
            sync_ctx = model.no_sync() if hasattr(model, "no_sync") and not should_step else nullcontext()
            with sync_ctx:
                with accelerator.autocast():
                    pred = model(xb)
                    loss_raw = weighted_mse(pred, yb, wb)
                    loss = loss_raw / float(group_total)
                if measure_first_step:
                    sync_if_needed(accelerator.device)
                    fwd_t = time.perf_counter()
                accelerator.backward(loss)
                if measure_first_step:
                    sync_if_needed(accelerator.device)
                    bwd_t = time.perf_counter()
            if should_step:
                accelerator.step_optimizer(optimizer)
                optimizer.zero_grad(set_to_none=True)
                accum_micro = 0
            if measure_first_step:
                sync_if_needed(accelerator.device)
                step_t = time.perf_counter()
                first_train_step_local = np.asarray(
                    [
                        fetch_t1 - fetch_t0,
                        mat_t - fetch_t1,
                        fwd_t - mat_t,
                        bwd_t - fwd_t,
                        step_t - bwd_t,
                        step_t - fetch_t0,
                    ],
                    dtype=np.float64,
                )
            if (
                accelerator.is_main_process
                and total_train_batches > 0
                and (batch_i % train_progress_step == 0 or batch_i == total_train_batches)
            ):
                print(
                    f"[epoch {epoch}] train progress {batch_i}/{total_train_batches} "
                    f"loss={loss_raw.detach().float().item():.6f}",
                    flush=True,
                )
        phase_timings[f"epoch_{epoch}_train_sec"] = float(time.perf_counter() - epoch_train_t0)

        if skip_train_eval:
            train_m = {
                "loss": float("nan"),
                "ic": float("nan"),
                "unweighted_ic": float("nan"),
                "weighted_ic": float("nan"),
                "rmse": float("nan"),
                "mae": float("nan"),
                "weight_sum": 0.0,
                "n": 0,
            }
            if accelerator.is_main_process:
                print(f"[epoch {epoch}] train-eval skipped", flush=True)
        else:
            train_eval_t0 = time.perf_counter()
            train_m = evaluate(
                model,
                train_eval_loader if train_eval_loader is not None else train_loader,
                feat_mean,
                feat_std,
                seq_offsets,
                accelerator,
                split_name=f"epoch {epoch} train-eval",
                progress_parts=args.eval_progress_splits,
                copy_non_blocking=use_pinned_transfer,
                collect_timing=profile_phases,
            )
            phase_timings[f"epoch_{epoch}_train_eval_sec"] = float(time.perf_counter() - train_eval_t0)
            train_eval_timing = train_m.pop("_timing", None)
        valid_eval_t0 = time.perf_counter()
        val_m = evaluate(
            model,
            valid_loader,
            feat_mean,
            feat_std,
            seq_offsets,
            accelerator,
            split_name=f"epoch {epoch} valid",
            progress_parts=args.eval_progress_splits,
            copy_non_blocking=use_pinned_transfer,
            collect_timing=profile_phases,
        )
        phase_timings[f"epoch_{epoch}_valid_eval_sec"] = float(time.perf_counter() - valid_eval_t0)
        valid_eval_timing = val_m.pop("_timing", None)
        final_val_m = dict(val_m)
        top_val_checkpoints, entered_topk, dropped_topk = refresh_top_val_checkpoints(
            top_val_checkpoints=top_val_checkpoints,
            candidate_epoch=epoch,
            candidate_val_metrics=val_m,
            keep_topk=save_topk_val_checkpoints,
            out_root=out_root,
        )
        is_best = 0
        current_weighted_ic = float(val_m.get("weighted_ic", val_m["ic"]))
        best_weighted_ic = float(best_val_m.get("weighted_ic", best_val_m["ic"])) if best_val_m is not None else float("-inf")
        if best_val_m is None or current_weighted_ic > best_weighted_ic:
            best_val_m = dict(val_m)
            best_epoch = epoch
            is_best = 1
            if accelerator.is_main_process:
                torch.save(accelerator.unwrap_model(model).state_dict(), best_model_path)

        if accelerator.is_main_process:
            if entered_topk:
                torch.save(
                    accelerator.unwrap_model(model).state_dict(),
                    build_val_epoch_checkpoint_path(out_root, epoch),
                )
            for dropped_item in dropped_topk:
                dropped_path = Path(str(dropped_item["path"]))
                if dropped_path.exists():
                    dropped_path.unlink()
            with log_path.open("a", newline="", encoding="utf-8") as f:
                csv.writer(f).writerow(
                    [
                        epoch,
                        current_lr,
                        train_m["loss"],
                        train_m["ic"],
                        train_m.get("unweighted_ic", float("nan")),
                        train_m.get("weighted_ic", float("nan")),
                        train_m["rmse"],
                        train_m.get("unweighted_rmse", float("nan")),
                        val_m["loss"],
                        val_m["ic"],
                        val_m.get("unweighted_ic", float("nan")),
                        val_m.get("weighted_ic", float("nan")),
                        val_m["rmse"],
                        val_m.get("unweighted_rmse", float("nan")),
                        val_m["mae"],
                        is_best,
                        str(args.pooling),
                    ]
                )
            print(
                f"[epoch {epoch}] "
                f"lr={current_lr:.8f} "
                f"train_loss={train_m['loss']:.6f} train_ic={train_m['ic']:.6f} "
                f"train_uic={train_m.get('unweighted_ic', float('nan')):.6f} train_rmse={train_m['rmse']:.6f} "
                f"val_loss={val_m['loss']:.6f} val_ic={val_m['ic']:.6f} "
                f"val_uic={val_m.get('unweighted_ic', float('nan')):.6f} val_rmse={val_m['rmse']:.6f} "
                f"best_epoch={best_epoch}"
            )
    phase_timings["fit_loop_sec"] = float(time.perf_counter() - train_loop_t0)
    if profile_phases and first_train_step_local is not None:
        phase_details["first_train_step"] = summarize_timing_across_ranks(
            accelerator=accelerator,
            labels=[
                "batch_wait_sec",
                "materialize_sec",
                "forward_sec",
                "backward_sec",
                "optimizer_step_sec",
                "step_total_sec",
            ],
            count=1,
            sums=first_train_step_local,
            count_key="measured_steps",
            global_batch=int(train_cfg["batch_size"]) * int(accelerator.num_processes),
        )
    if collective_warmup_done:
        phase_details["collective_warmup_sec"] = {"local_sec": float(collective_warmup_sec)}
    if train_eval_timing is not None:
        phase_details["train_eval_timing"] = train_eval_timing
    if valid_eval_timing is not None:
        phase_details["valid_eval_timing"] = valid_eval_timing

    def evaluate_saved_checkpoint(checkpoint_entry: dict, split_name: str) -> Tuple[Dict[str, float], float, dict | None]:
        eval_t0 = time.perf_counter()
        state_dict = torch.load(str(checkpoint_entry["path"]), map_location="cpu", weights_only=True)
        accelerator.unwrap_model(model).load_state_dict(state_dict)
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            print(
                f"[{split_name}] evaluating val_rank={checkpoint_entry['rank']} "
                f"epoch={checkpoint_entry['epoch']} val_ic={checkpoint_entry['val_ic']:.6f} "
                f"val_wic={checkpoint_entry.get('val_weighted_ic', checkpoint_entry['val_ic']):.6f}",
                flush=True,
            )
        metrics = evaluate(
            model,
            test_loader,
            feat_mean,
            feat_std,
            seq_offsets,
            accelerator,
            split_name=split_name,
            progress_parts=args.eval_progress_splits,
            copy_non_blocking=use_pinned_transfer,
            collect_timing=profile_phases,
        )
        elapsed = float(time.perf_counter() - eval_t0)
        timing = metrics.pop("_timing", None)
        return metrics, elapsed, timing

    accelerator.wait_for_everyone()
    test_m_best: Dict[str, float] | None = None
    test_m_selected: Dict[str, float] | None = None
    selected_test_checkpoint: dict | None = None
    if test_loader is not None and not skip_test:
        if not top_val_checkpoints:
            raise RuntimeError("No validation checkpoints were recorded for test evaluation.")
        if selected_test_checkpoint_rank > len(top_val_checkpoints):
            raise RuntimeError(
                f"Requested test_checkpoint_rank={selected_test_checkpoint_rank} but only "
                f"{len(top_val_checkpoints)} validation checkpoints are available."
            )
        best_test_checkpoint = top_val_checkpoints[0]
        selected_test_checkpoint = top_val_checkpoints[selected_test_checkpoint_rank - 1]
        test_m_best, test_best_sec, test_eval_timing = evaluate_saved_checkpoint(
            best_test_checkpoint,
            split_name="test best-val",
        )
        phase_timings["test_eval_best_val_sec"] = float(test_best_sec)
        if selected_test_checkpoint_rank == 1:
            test_m_selected = dict(test_m_best)
            phase_timings["test_eval_selected_val_rank_sec"] = float(test_best_sec)
            phase_timings["test_eval_sec"] = float(test_best_sec)
        else:
            test_m_selected, test_selected_sec, test_eval_timing_selected = evaluate_saved_checkpoint(
                selected_test_checkpoint,
                split_name=f"test val-rank-{selected_test_checkpoint_rank}",
            )
            phase_timings["test_eval_selected_val_rank_sec"] = float(test_selected_sec)
            phase_timings["test_eval_sec"] = float(test_best_sec + test_selected_sec)
            if test_eval_timing_selected is not None:
                phase_details["test_eval_timing_selected_val_rank"] = test_eval_timing_selected
        if accelerator.is_main_process:
            print(
                f"[test best-val] rmse={test_m_best['rmse']:.6f} mae={test_m_best['mae']:.6f} "
                f"ic={test_m_best['ic']:.6f} uic={test_m_best.get('unweighted_ic', float('nan')):.6f} "
                f"wic={test_m_best.get('weighted_ic', float('nan')):.6f}",
                flush=True,
            )
            if selected_test_checkpoint_rank != 1 and test_m_selected is not None:
                print(
                    f"[test val-rank-{selected_test_checkpoint_rank}] "
                    f"rmse={test_m_selected['rmse']:.6f} mae={test_m_selected['mae']:.6f} "
                    f"ic={test_m_selected['ic']:.6f} uic={test_m_selected.get('unweighted_ic', float('nan')):.6f} "
                    f"wic={test_m_selected.get('weighted_ic', float('nan')):.6f}",
                    flush=True,
                )
    elif test_loader is not None and accelerator.is_main_process:
        print("[test] skipped", flush=True)
    if test_eval_timing is not None:
        phase_details["test_eval_timing"] = test_eval_timing
    final_save_t0 = time.perf_counter()
    if accelerator.is_main_process:
        torch.save(accelerator.unwrap_model(model).state_dict(), out_root / "gru_seq_memmap_ddp.pt")
    accelerator.wait_for_everyone()
    phase_timings["final_save_sec"] = float(time.perf_counter() - final_save_t0)
    phase_timings["total_run_sec"] = float(time.perf_counter() - run_t0)
    if accelerator.is_main_process:
        summary = {
            "run_name": args.run_name,
            "cache_name": args.cache_name,
            "max_days": args.max_days,
            "train_pool_start_date": args.train_pool_start_date,
            "train_pool_end_date": args.train_pool_end_date,
            "test_start_date": args.test_start_date,
            "test_end_date": args.test_end_date,
            "valid_split_ratio": split_ratio,
            "selection_metric": "val_weighted_ic",
            "train_day_count": len(train_days),
            "valid_day_count": len(valid_days),
            "test_day_count": len(test_days),
            "train_samples": int(train_loader.total_samples),
            "valid_samples": int(valid_loader.total_samples),
            "test_samples": int(test_loader.total_samples) if test_loader is not None else 0,
            "seed": int(args.seed),
            "seq_len": args.seq_len,
            "sample_stride": args.sample_stride,
            "min_timecode": int(min_timecode),
            "require_positive_weight": bool(require_positive_weight),
            "use_source_raw_weights": bool(args.use_source_raw_weights),
            "hidden_dim": args.hidden_dim,
            "num_layers": args.num_layers,
            "bidirectional": bool(args.bidirectional),
            "use_cnn1d": bool(getattr(args, "use_cnn1d", 0)),
            "input_gate_hidden_dim": int(getattr(args, "input_gate_hidden_dim", 0)),
            "input_gate_bias": float(getattr(args, "input_gate_bias", 2.0)),
            "weight_decay": args.weight_decay,
            "prefer_fp16": bool(args.prefer_fp16),
            "use_amp": bool(args.use_amp),
            "pooling": str(args.pooling),
            "loader_backend": "threaded_seq_batch_loader",
            "loader_threads": loader_threads,
            "loader_prefetch_batches": loader_prefetch,
            "train_batch_overlap": train_batch_overlap,
            "train_batch_stride": int(train_cfg["batch_size"]) - train_batch_overlap,
            "eval_batch_size_per_rank": eval_batch_size,
            "grad_accum_steps": grad_accum_steps,
            "effective_global_batch": int(train_cfg["batch_size"]) * int(accelerator.num_processes) * grad_accum_steps,
            "runtime_backend": args.runtime_backend,
            "device_transfer_mode": device_transfer_mode,
            "device_transfer_prefetch_batches": device_transfer_prefetch_batches,
            "pinned_nonblocking_transfer": use_pinned_transfer,
            "skip_train_eval": skip_train_eval,
            "skip_test": skip_test,
            "train_cfg": train_cfg,
            "epoch_learning_rates": [float(x) for x in epoch_lrs],
            "loss": "weighted_mse",
            "metric_main": "val_weighted_ic",
            "test_selection_rule": f"val_weighted_ic_rank_{selected_test_checkpoint_rank}",
            "save_topk_val_checkpoints": int(save_topk_val_checkpoints),
            "best_epoch_by_val_ic": best_epoch,
            "best_epoch_by_val_weighted_ic": best_epoch,
            "best_val_metrics": best_val_m,
            "final_epoch_val_metrics": final_val_m,
            "top_val_checkpoints": top_val_checkpoints,
            "selected_test_checkpoint_rank": int(selected_test_checkpoint_rank),
            "selected_test_checkpoint_epoch": int(selected_test_checkpoint["epoch"]) if selected_test_checkpoint is not None else None,
            "selected_test_model_path": str(selected_test_checkpoint["path"]) if selected_test_checkpoint is not None else None,
            "test_metrics_at_best_val": test_m_best,
            "test_metrics_at_selected_val_rank": test_m_selected,
            "best_model_path": str(best_model_path),
            "phase_timings": phase_timings,
            "phase_details": phase_details,
        }
        (out_root / "training_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
        np.savez(out_root / "feature_stats.npz", mean=feat_mean_np, std=feat_std_np)
        if profile_phases:
            phase_summary = {
                "run_name": args.run_name,
                "row_root": str(row_root),
                "world_size": int(accelerator.num_processes),
                "batch_size_per_rank": int(train_cfg["batch_size"]),
                "train_batch_overlap": train_batch_overlap,
                "runtime_backend": args.runtime_backend,
                "device_transfer_mode": device_transfer_mode,
                "min_timecode": int(min_timecode),
                "require_positive_weight": bool(require_positive_weight),
                "use_source_raw_weights": bool(args.use_source_raw_weights),
                "phase_timings": phase_timings,
                "phase_details": phase_details,
            }
            phase_path = out_root / "phase_profile_summary.json"
            phase_path.write_text(json.dumps(phase_summary, ensure_ascii=False, indent=2), encoding="utf-8")
            print(f"[phase-profile] saved={phase_path}", flush=True)
    accelerator.close()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="/intern9/huhongkai/hs300_factor_lab/configs/experiment_2024_2025_memmap.json")
    p.add_argument("--run-name", default="seq_gru_memmap_run")
    p.add_argument("--cache-name", default="top200_eps20_rows_2024_2025")
    p.add_argument(
        "--row-root-override",
        type=str,
        default="",
        help="Optional row_memmap root or exact cache dir. Useful for local staged cache under /dev/shm.",
    )
    p.add_argument("--max-days", type=int, default=-1)
    p.add_argument("--max-samples", type=int, default=-1)
    p.add_argument("--train-day-limit", type=int, default=-1)
    p.add_argument("--valid-day-limit", type=int, default=-1)
    p.add_argument("--test-day-limit", type=int, default=-1)
    p.add_argument("--train-pool-start-date", type=str, default="")
    p.add_argument("--train-pool-end-date", type=str, default="")
    p.add_argument("--test-start-date", type=str, default="")
    p.add_argument("--test-end-date", type=str, default="")
    p.add_argument("--valid-split-ratio", type=float, default=-1.0)
    p.add_argument("--seed", type=int, default=20260312)
    p.add_argument("--seq-len", type=int, default=60)
    p.add_argument("--sample-stride", type=int, default=10)
    p.add_argument("--train-batch-overlap", type=int, default=0)
    p.add_argument("--eval-batch-size", type=int, default=0)
    p.add_argument("--grad-accum-steps", type=int, default=1)
    p.add_argument("--runtime-backend", type=str, default="native", choices=["accelerate", "native"])
    p.add_argument(
        "--device-transfer-mode",
        type=str,
        default="thread_prefetch",
        choices=["direct", "thread_prefetch", "materialize_prefetch"],
    )
    p.add_argument("--device-transfer-prefetch-batches", type=int, default=2)
    p.add_argument("--loader-prefetch-batches", type=int, default=0)
    p.add_argument("--prefer-fp16", type=int, default=1)
    p.add_argument("--override-epochs", type=int, default=-1)
    p.add_argument("--override-lr", type=float, default=-1.0)
    p.add_argument("--override-batch-size", type=int, default=-1)
    p.add_argument("--hidden-dim", type=int, default=256)
    p.add_argument("--num-layers", type=int, default=2)
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--pooling", type=str, default="last", choices=["last", "attn"])
    p.add_argument("--bidirectional", type=int, default=0, help="1 to enable bidirectional GRU")
    p.add_argument("--use-cnn1d", type=int, default=0, help="1 to use parallel 1D-CNN before GRU")
    p.add_argument("--input-gate-hidden-dim", type=int, default=0, help=">0 to enable feature/channel gate before GRU")
    p.add_argument("--input-gate-bias", type=float, default=2.0, help="Initial bias for feature/channel gate")
    p.add_argument("--weight-decay", type=float, default=1e-5)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--use-amp", type=int, default=1, help="1 to enable fp16 mixed precision")
    p.add_argument("--train-progress-splits", type=int, default=10)
    p.add_argument("--eval-progress-splits", type=int, default=4)
    p.add_argument("--skip-train-eval", type=int, default=1, help="1 to skip full train-set evaluation")
    p.add_argument("--skip-test", type=int, default=0, help="1 to skip final test evaluation")
    p.add_argument(
        "--use-source-raw-weights",
        type=int,
        default=1,
        help="1 to load raw continuous weights from the source split npy instead of cache-side binary masks.",
    )
    p.add_argument(
        "--min-timecode",
        type=int,
        default=DEFAULT_MIN_TIMECODE,
        help="Only keep sequence end points whose source datetime >= this HHMMSSmmm timecode.",
    )
    p.add_argument(
        "--require-positive-weight",
        type=int,
        default=1,
        help="1 to drop sequence end points whose cached/source weight is not positive.",
    )
    p.add_argument("--profile-phases", type=int, default=0)
    p.add_argument("--profile-warmup-steps", type=int, default=0)
    p.add_argument("--profile-steps", type=int, default=0)
    p.add_argument("--stop-after-profile", type=int, default=0)
    p.add_argument("--save-topk-val-checkpoints", type=int, default=1)
    p.add_argument("--test-checkpoint-rank", type=int, default=1)
    args = p.parse_args()
    run(args)


if __name__ == "__main__":
    main()
:�import argparse
import csv
import hashlib
import json
import math
import os
import queue
import random
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import nullcontext
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from accelerate import Accelerator

from common import ensure_dir, load_config, load_split_view_meta, resolve_named_day_dirs, resolve_row_root
from splitview_time_utils import load_split_view_source_field_slice


def set_global_seed(seed: int) -> None:
    seed = int(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def build_epoch_lr_schedule(train_cfg: dict, epochs: int) -> List[float]:
    base_lr = float(train_cfg["learning_rate"])
    raw_values = train_cfg.get("lr_epoch_values")
    if raw_values is None:
        return [base_lr for _ in range(int(epochs))]
    values = [float(x) for x in raw_values]
    if len(values) != int(epochs):
        raise ValueError(
            f"train.lr_epoch_values length ({len(values)}) must match epochs ({int(epochs)})."
        )
    return values


def set_optimizer_lr(optimizer, lr: float) -> None:
    lr = float(lr)
    for group in optimizer.param_groups:
        group["lr"] = lr


class AccelerateRuntime:
    def __init__(self, use_amp: bool):
        self.accelerator = Accelerator(mixed_precision="fp16" if bool(use_amp) else "no")
        self.device = self.accelerator.device
        self.process_index = int(self.accelerator.process_index)
        self.num_processes = int(self.accelerator.num_processes)
        self.is_main_process = bool(self.accelerator.is_main_process)

    def autocast(self):
        return self.accelerator.autocast()

    def backward(self, loss: torch.Tensor) -> None:
        self.accelerator.backward(loss)

    def step_optimizer(self, optimizer) -> None:
        optimizer.step()

    def prepare(self, model, optimizer):
        return self.accelerator.prepare(model, optimizer)

    def unwrap_model(self, model):
        return self.accelerator.unwrap_model(model)

    def reduce(self, tensor: torch.Tensor, reduction: str = "sum") -> torch.Tensor:
        return self.accelerator.reduce(tensor, reduction=reduction)

    def gather(self, tensor: torch.Tensor) -> torch.Tensor:
        return self.accelerator.gather(tensor)

    def wait_for_everyone(self) -> None:
        self.accelerator.wait_for_everyone()

    def close(self) -> None:
        return None


class NativeDDPRuntime:
    def __init__(self, use_amp: bool, enable_static_graph: bool = True):
        if torch.cuda.is_available():
            local_rank = int(os.environ.get("LOCAL_RANK", 0))
            torch.cuda.set_device(local_rank)
            self.device = torch.device("cuda", local_rank)
        else:
            self.device = torch.device("cpu")
        requested_world = int(os.environ.get("WORLD_SIZE", "1"))
        self._distributed = requested_world > 1
        if self._distributed and not dist.is_initialized():
            backend = "nccl" if self.device.type == "cuda" else "gloo"
            dist.init_process_group(backend=backend, timeout=timedelta(seconds=7200))
        if dist.is_initialized():
            self.process_index = int(dist.get_rank())
            self.num_processes = int(dist.get_world_size())
        else:
            self.process_index = 0
            self.num_processes = 1
        self.is_main_process = self.process_index == 0
        self.use_amp = bool(use_amp) and self.device.type == "cuda"
        self.enable_static_graph = bool(enable_static_graph)
        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)

    def autocast(self):
        if self.device.type == "cuda":
            return torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp)
        return nullcontext()

    def backward(self, loss: torch.Tensor) -> None:
        if self.use_amp:
            self.scaler.scale(loss).backward()
        else:
            loss.backward()

    def step_optimizer(self, optimizer) -> None:
        if self.use_amp:
            self.scaler.step(optimizer)
            self.scaler.update()
        else:
            optimizer.step()

    def prepare(self, model, optimizer):
        model = model.to(self.device)
        if self.num_processes > 1:
            ddp_kwargs = {
                "broadcast_buffers": False,
                "gradient_as_bucket_view": True,
            }
            if self.device.type == "cuda":
                ddp_kwargs["device_ids"] = [self.device.index]
                ddp_kwargs["output_device"] = self.device.index
            if self.enable_static_graph:
                try:
                    model = DDP(model, static_graph=True, **ddp_kwargs)
                except TypeError:
                    model = DDP(model, **ddp_kwargs)
            else:
                model = DDP(model, **ddp_kwargs)
        return model, optimizer

    def unwrap_model(self, model):
        return model.module if isinstance(model, DDP) else model

    def reduce(self, tensor: torch.Tensor, reduction: str = "sum") -> torch.Tensor:
        if self.num_processes <= 1:
            return tensor
        out = tensor.clone()
        reduction_name = str(reduction).lower()
        if reduction_name == "sum":
            op = dist.ReduceOp.SUM
            dist.all_reduce(out, op=op)
        elif reduction_name == "mean":
            dist.all_reduce(out, op=dist.ReduceOp.SUM)
            out = out / float(self.num_processes)
        elif reduction_name == "max":
            dist.all_reduce(out, op=dist.ReduceOp.MAX)
        else:
            raise ValueError(f"Unsupported reduction: {reduction}")
        return out

    def gather(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.num_processes <= 1:
            return tensor
        gather_list = [torch.empty_like(tensor) for _ in range(self.num_processes)]
        dist.all_gather(gather_list, tensor)
        return torch.cat(gather_list, dim=0)

    def wait_for_everyone(self) -> None:
        if self.num_processes > 1:
            if self.device.type == "cuda":
                dist.barrier(device_ids=[self.device.index])
            else:
                dist.barrier()

    def close(self) -> None:
        if dist.is_initialized():
            dist.destroy_process_group()


def build_runtime(runtime_backend: str, use_amp: bool, enable_static_graph: bool = True):
    backend = str(runtime_backend).strip().lower()
    if backend == "native":
        return NativeDDPRuntime(use_amp=use_amp, enable_static_graph=enable_static_graph)
    if backend == "accelerate":
        return AccelerateRuntime(use_amp=use_amp)
    raise ValueError(f"Unsupported runtime_backend: {runtime_backend}")


@dataclass
class DayStore:
    name: str
    n_rows: int
    n_factors: int
    x: np.memmap
    y: np.memmap
    w: np.memmap
    sym_start: np.memmap
    sym_end: np.memmap
    need_clean: bool


_SPLIT_X_ARRAY_CACHE: Dict[str, np.ndarray] = {}
DEFAULT_MIN_TIMECODE = 93_000_000
_USE_SOURCE_RAW_WEIGHTS = False


def configure_split_view_weight_loading(use_source_raw_weights: bool) -> None:
    global _USE_SOURCE_RAW_WEIGHTS
    _USE_SOURCE_RAW_WEIGHTS = bool(use_source_raw_weights)


def _resolve_source_x_path(day_dir: Path, meta: dict) -> Path | None:
    explicit = str(meta.get("source_x_path") or "").strip()
    if explicit:
        return Path(explicit)
    source_root = str(meta.get("source_root") or "").strip()
    source_split = str(meta.get("source_split") or "").strip()
    if source_root and source_split:
        return Path(source_root) / source_split / "x.npy"
    return None


def _load_cached_x_npy(path: Path) -> np.ndarray:
    key = str(path.resolve())
    arr = _SPLIT_X_ARRAY_CACHE.get(key)
    if arr is None:
        arr = np.load(path, mmap_mode="r", allow_pickle=False)
        _SPLIT_X_ARRAY_CACHE[key] = arr
    return arr


def list_ready_days(row_root: Path) -> List[Path]:
    days = []
    for p in sorted(row_root.iterdir()):
        if p.is_dir() and (p / "_SUCCESS").exists():
            days.append(p)
    return days


def split_days(days: List[Path], ratio: float) -> Tuple[List[Path], List[Path]]:
    n = len(days)
    if n < 2:
        return days, days
    cut = max(1, min(n - 1, int(n * ratio)))
    return days[:cut], days[cut:]


def _day_str_from_dir(day_dir: Path) -> str:
    name = day_dir.name
    if len(name) >= 8:
        day = name[:8]
        if day.isdigit():
            return day
    return ""


def filter_days_by_date(days: List[Path], start_date: str, end_date: str) -> List[Path]:
    s = (start_date or "").strip()
    e = (end_date or "").strip()
    if not s and not e:
        return list(days)
    if s and (len(s) != 8 or not s.isdigit()):
        raise ValueError(f"Invalid start_date: {s}")
    if e and (len(e) != 8 or not e.isdigit()):
        raise ValueError(f"Invalid end_date: {e}")
    lo = s if s else "00000000"
    hi = e if e else "99999999"
    if lo > hi:
        raise ValueError(f"Invalid date window: {lo} > {hi}")
    out: List[Path] = []
    for d in days:
        day = _day_str_from_dir(d)
        if day and lo <= day <= hi:
            out.append(d)
    return out


def limit_days(days: List[Path], day_limit: int) -> List[Path]:
    if day_limit <= 0 or day_limit >= len(days):
        return days
    return days[:day_limit]


def load_day_meta(day_dir: Path) -> dict:
    return json.loads((day_dir / "meta.json").read_text(encoding="utf-8"))


def load_day_store(day_dir: Path, prefer_fp16: bool) -> DayStore:
    meta = load_day_meta(day_dir)
    n_rows = int(meta["n_rows"])
    n_factors = int(meta["n_factors"])
    source_x_path = _resolve_source_x_path(day_dir, meta)
    if source_x_path is not None:
        row_start = int(meta.get("source_row_start", 0))
        row_stop = int(meta.get("source_row_stop", row_start + n_rows))
        x_all = _load_cached_x_npy(source_x_path)
        x = x_all[row_start:row_stop]
        need_clean = True
    else:
        fp16_path = day_dir / "x_top200_f16_filled.memmap"
        if prefer_fp16 and fp16_path.exists():
            x = np.memmap(fp16_path, mode="r", dtype=np.float16, shape=(n_rows, n_factors))
            need_clean = False
        else:
            x = np.memmap(day_dir / "x_top200_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows, n_factors))
            need_clean = True
    y = np.memmap(day_dir / "y_sum_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    if bool(_USE_SOURCE_RAW_WEIGHTS):
        try:
            w = np.asarray(load_split_view_source_field_slice(day_dir, "w"), dtype=np.float32)
        except Exception:
            w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    else:
        w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
    sym_n = int(meta["symbol_count_meta"])
    sym_start = np.memmap(day_dir / "symbol_start_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    sym_end = np.memmap(day_dir / "symbol_end_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    return DayStore(day_dir.name, n_rows, n_factors, x, y, w, sym_start, sym_end, need_clean)


def load_day_symbol_bounds(day_dir: Path) -> Tuple[np.memmap, np.memmap]:
    meta = load_day_meta(day_dir)
    sym_n = int(meta["symbol_count_meta"])
    sym_start = np.memmap(day_dir / "symbol_start_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    sym_end = np.memmap(day_dir / "symbol_end_idx_int64.memmap", mode="r", dtype=np.int64, shape=(sym_n,))
    return sym_start, sym_end


def filter_valid_end_indices(
    day_dir: Path,
    end_indices: np.ndarray,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> np.ndarray:
    ends = np.asarray(end_indices, dtype=np.int64)
    if ends.size == 0:
        return ends
    mask = np.ones((int(ends.shape[0]),), dtype=bool)
    if bool(require_positive_weight):
        meta = load_day_meta(day_dir)
        n_rows = int(meta["n_rows"])
        cache_w = np.memmap(day_dir / "w_float32.memmap", mode="r", dtype=np.float32, shape=(n_rows,))
        weight_mask = np.asarray(cache_w[ends], dtype=np.float32)
        mask &= np.isfinite(weight_mask) & (weight_mask > 0.0)
    if int(min_timecode) > 0:
        dt_slice = np.asarray(load_split_view_source_field_slice(day_dir, "datetime"), dtype=np.int64)
        dt_end = np.asarray(dt_slice[ends], dtype=np.int64)
        mask &= dt_end >= int(min_timecode)
    return ends[mask]


def build_valid_end_indices_from_bounds(
    sym_start: np.ndarray,
    sym_end: np.ndarray,
    seq_len: int,
    sample_stride: int,
) -> np.ndarray:
    stride = int(sample_stride)
    starts = np.asarray(sym_start, dtype=np.int64) + int(seq_len) - 1
    # symbol_end_idx is exclusive in this memmap layout.
    ends = np.asarray(sym_end, dtype=np.int64) - 1
    valid_mask = starts <= ends
    if not np.any(valid_mask):
        return np.empty((0,), dtype=np.int64)
    starts = starts[valid_mask]
    ends = ends[valid_mask]
    lengths = ((ends - starts) // stride + 1).astype(np.int64, copy=False)
    total = int(lengths.sum())
    if total <= 0:
        return np.empty((0,), dtype=np.int64)
    repeated_starts = np.repeat(starts, lengths)
    group_offsets = np.repeat(np.cumsum(lengths, dtype=np.int64) - lengths, lengths)
    intra_offsets = np.arange(total, dtype=np.int64) - group_offsets
    return repeated_starts + intra_offsets * stride


def build_valid_end_indices(store: DayStore, seq_len: int, sample_stride: int) -> np.ndarray:
    return build_valid_end_indices_from_bounds(
        store.sym_start,
        store.sym_end,
        seq_len=seq_len,
        sample_stride=sample_stride,
    )


def build_end_index_cache_path(
    cache_dir: Path,
    day_dir: Path,
    seq_len: int,
    sample_stride: int,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> Path:
    key_src = "\n".join(
        [
            str(day_dir.parent),
            str(day_dir.name),
            str(int(seq_len)),
            str(int(sample_stride)),
            str(int(min_timecode)),
            str(int(bool(require_positive_weight))),
        ]
    )
    key = hashlib.sha1(key_src.encode("utf-8")).hexdigest()[:16]
    return cache_dir / f"{day_dir.name}_seq{int(seq_len)}_stride{int(sample_stride)}_{key}.npy"


def load_end_index_cache(path: Path) -> np.ndarray:
    arr = np.load(path, allow_pickle=False)
    return np.asarray(arr, dtype=np.int64)


def wait_for_end_index_cache(path: Path, timeout_sec: float = 600.0) -> np.ndarray:
    deadline = time.time() + float(timeout_sec)
    last_err = None
    while time.time() < deadline:
        if path.exists() and path.stat().st_size > 0:
            try:
                return load_end_index_cache(path)
            except Exception as exc:  # pragma: no cover - transient partial-write case
                last_err = exc
        time.sleep(0.2)
    if last_err is not None:
        raise TimeoutError(f"Timed out waiting for end-index cache {path}: {last_err}") from last_err
    raise TimeoutError(f"Timed out waiting for end-index cache {path}")


def load_or_build_valid_end_indices(
    day_dir: Path,
    seq_len: int,
    sample_stride: int,
    cache_dir: Path | None,
    cache_writer: bool,
    min_timecode: int = -1,
    require_positive_weight: bool = False,
) -> np.ndarray:
    if cache_dir is None:
        sym_start, sym_end = load_day_symbol_bounds(day_dir)
        ends = build_valid_end_indices_from_bounds(sym_start, sym_end, seq_len=seq_len, sample_stride=sample_stride)
        return filter_valid_end_indices(
            day_dir,
            ends,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
    cache_path = build_end_index_cache_path(
        cache_dir,
        day_dir,
        seq_len=seq_len,
        sample_stride=sample_stride,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    if cache_path.exists() and cache_path.stat().st_size > 0:
        return load_end_index_cache(cache_path)
    if not cache_writer:
        return wait_for_end_index_cache(cache_path)
    sym_start, sym_end = load_day_symbol_bounds(day_dir)
    ends = build_valid_end_indices_from_bounds(sym_start, sym_end, seq_len=seq_len, sample_stride=sample_stride)
    ends = filter_valid_end_indices(
        day_dir,
        ends,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    tmp_path = cache_path.with_suffix(f".tmp.{int(time.time() * 1000)}.{os.getpid()}.npy")
    np.save(tmp_path, ends)
    tmp_path.replace(cache_path)
    return ends


def truncate_end_indices(end_indices: List[np.ndarray], max_samples: int) -> List[np.ndarray]:
    if max_samples <= 0:
        return end_indices
    remaining = int(max_samples)
    trimmed: List[np.ndarray] = []
    for ends in end_indices:
        if remaining <= 0:
            trimmed.append(np.empty((0,), dtype=np.int64))
            continue
        take = min(int(ends.shape[0]), remaining)
        trimmed.append(ends[:take])
        remaining -= take
    return trimmed


def split_contiguous_even(total: int, rank: int, world_size: int) -> Tuple[int, int]:
    if total <= 0:
        return 0, 0
    start = (total * rank) // max(1, world_size)
    end = (total * (rank + 1)) // max(1, world_size)
    return int(start), int(end)


@dataclass(frozen=True)
class BatchSlice:
    day_i: int
    start: int
    stop: int


@dataclass(frozen=True)
class BatchPlanItem:
    parts: Tuple[BatchSlice, ...]
    pad_size: int = 0


@dataclass(frozen=True)
class PackedSeqBatchPart:
    span_x: torch.Tensor
    local_end: torch.Tensor
    y: torch.Tensor
    w: torch.Tensor


@dataclass(frozen=True)
class PackedSeqBatch:
    parts: Tuple[PackedSeqBatchPart, ...]


@dataclass(frozen=True)
class MaterializedSeqBatch:
    xb: torch.Tensor
    y: torch.Tensor
    w: torch.Tensor


class SeqMemmapBatchLoader:
    def __init__(
        self,
        day_dirs: List[Path],
        seq_len: int,
        sample_stride: int,
        batch_size: int,
        shuffle: bool,
        max_samples: int = -1,
        prefer_fp16: bool = True,
        rank: int = 0,
        world_size: int = 1,
        seed: int = 20260312,
        pin_memory: bool = True,
        loader_threads: int = 1,
        prefetch_batches: int = 1,
        pad_last_batch: bool = False,
        batch_overlap: int = 0,
        index_cache_dir: Path | None = None,
        min_timecode: int = -1,
        require_positive_weight: bool = False,
    ):
        self.seq_len = seq_len
        self.batch_size = int(batch_size)
        self.batch_overlap = max(0, int(batch_overlap))
        if self.batch_overlap >= self.batch_size:
            raise ValueError(
                f"batch_overlap must be smaller than batch_size, got overlap={self.batch_overlap} batch_size={self.batch_size}"
            )
        self.batch_stride = self.batch_size - self.batch_overlap
        self.shuffle = bool(shuffle)
        self.rank = int(rank)
        self.world_size = int(world_size)
        self.seed = int(seed)
        self.day_dirs = list(day_dirs)
        self.prefer_fp16 = bool(prefer_fp16)
        self.pin_memory = bool(pin_memory) and torch.cuda.is_available()
        self.loader_threads = max(1, int(loader_threads))
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.pad_last_batch = bool(pad_last_batch)
        self.epoch = 0
        self.index_cache_dir = index_cache_dir
        self.min_timecode = int(min_timecode)
        self.require_positive_weight = bool(require_positive_weight)
        self.stores: List[DayStore | None] = [None for _ in self.day_dirs]
        end_indices = [
            load_or_build_valid_end_indices(
                d,
                seq_len=seq_len,
                sample_stride=sample_stride,
                cache_dir=self.index_cache_dir,
                cache_writer=(self.rank == 0),
                min_timecode=self.min_timecode,
                require_positive_weight=self.require_positive_weight,
            )
            for d in self.day_dirs
        ]
        self.end_indices: List[np.ndarray] = truncate_end_indices(end_indices, max_samples=max_samples)
        self.total_samples = int(sum(int(ends.shape[0]) for ends in self.end_indices))
        self.seq_offsets = np.arange(-(self.seq_len - 1), 1, dtype=np.int64)
        self.rank_sample_counts = self._compute_rank_sample_counts()
        self.rank_total_samples = [int(sum(day_counts)) for day_counts in self.rank_sample_counts]
        self.rank_unique_batches = [
            self._compute_batch_count(total) for total in self.rank_total_samples
        ]
        self.total_batches = max(self.rank_unique_batches, default=0)

    def __len__(self) -> int:
        return self.total_batches

    def set_epoch(self, epoch: int) -> None:
        self.epoch = int(epoch)

    def _get_store(self, day_i: int) -> DayStore:
        store = self.stores[day_i]
        if store is None:
            store = load_day_store(self.day_dirs[day_i], prefer_fp16=self.prefer_fp16)
            self.stores[day_i] = store
        return store

    def _compute_rank_sample_counts(self) -> List[List[int]]:
        counts: List[List[int]] = [[] for _ in range(self.world_size)]
        for ends in self.end_indices:
            n = int(ends.shape[0])
            for rank in range(self.world_size):
                start, stop = split_contiguous_even(n, rank, self.world_size)
                counts[rank].append(max(0, stop - start))
        return counts

    def _compute_batch_count(self, total: int) -> int:
        total = int(total)
        if total <= 0:
            return 0
        if total <= self.batch_size:
            return 1
        return 1 + int(math.ceil((total - self.batch_size) / self.batch_stride))

    def _tail_parts(self, parts: List[BatchSlice], keep: int) -> List[BatchSlice]:
        keep = max(0, int(keep))
        if keep <= 0:
            return []
        out: List[BatchSlice] = []
        remaining = keep
        for part in reversed(parts):
            part_len = int(part.stop - part.start)
            if part_len <= 0:
                continue
            take = min(remaining, part_len)
            out.append(BatchSlice(day_i=part.day_i, start=part.stop - take, stop=part.stop))
            remaining -= take
            if remaining == 0:
                break
        if remaining != 0:
            raise RuntimeError(f"Failed to preserve batch overlap keep={keep}, remaining={remaining}")
        out.reverse()
        return out

    def _build_rank_plan(self) -> List[BatchPlanItem]:
        # Rotate the contiguous shard every epoch so each rank sees different regions over time.
        shard_rank = (self.rank + self.epoch) % max(1, self.world_size) if self.shuffle else self.rank
        items: List[BatchPlanItem] = []
        current_parts: List[BatchSlice] = []
        filled = 0
        for day_i, ends in enumerate(self.end_indices):
            total = int(ends.shape[0])
            start, stop = split_contiguous_even(total, shard_rank, self.world_size)
            if stop <= start:
                continue
            cursor = int(start)
            stop = int(stop)
            while cursor < stop:
                need = self.batch_size - filled
                take = min(need, stop - cursor)
                current_parts.append(BatchSlice(day_i=day_i, start=cursor, stop=cursor + take))
                cursor += take
                filled += take
                if filled == self.batch_size:
                    emitted_parts = tuple(current_parts)
                    items.append(BatchPlanItem(parts=emitted_parts, pad_size=0))
                    current_parts = self._tail_parts(list(emitted_parts), self.batch_overlap)
                    filled = self.batch_overlap
        if current_parts:
            pad_size = self.batch_size - filled if self.pad_last_batch else 0
            items.append(BatchPlanItem(parts=tuple(current_parts), pad_size=pad_size))
        if not items:
            return []
        if self.shuffle and len(items) > 1:
            rng = np.random.default_rng(self.seed + self.epoch)
            if self.batch_overlap > 0:
                shift = int(rng.integers(len(items)))
                if shift > 0:
                    items = items[shift:] + items[:shift]
            else:
                order = rng.permutation(len(items))
                items = [items[int(i)] for i in order.tolist()]
        if len(items) < self.total_batches:
            base = list(items)
            pad_idx = 0
            while len(items) < self.total_batches:
                items.append(base[pad_idx % len(base)])
                pad_idx += 1
        return items

    def _load_batch_part(self, part: BatchSlice) -> PackedSeqBatchPart:
        day_i, start, stop = part.day_i, part.start, part.stop
        store = self._get_store(day_i)
        batch_end = self.end_indices[day_i][start:stop]
        if batch_end.size == 0:
            raise RuntimeError(f"Empty batch slice for day_i={day_i}, start={start}, stop={stop}")
        span_start = int(batch_end[0]) - self.seq_len + 1
        span_end = int(batch_end[-1])
        span_x_dtype = np.float32 if store.need_clean else store.x.dtype
        span_x = np.array(store.x[span_start : span_end + 1], dtype=span_x_dtype, copy=True)
        if store.need_clean:
            np.nan_to_num(span_x, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        local_end = np.ascontiguousarray(batch_end - span_start, dtype=np.int64)
        y = np.array(store.y[batch_end], dtype=np.float32, copy=True)
        w = np.array(store.w[batch_end], dtype=np.float32, copy=True)
        np.nan_to_num(y, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        np.nan_to_num(w, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
        np.maximum(w, 0.0, out=w)
        span_xb = torch.from_numpy(np.ascontiguousarray(span_x))
        local_endb = torch.from_numpy(local_end)
        yb = torch.from_numpy(np.ascontiguousarray(y))
        wb = torch.from_numpy(np.ascontiguousarray(w))
        if self.pin_memory:
            span_xb = span_xb.pin_memory()
            local_endb = local_endb.pin_memory()
            yb = yb.pin_memory()
            wb = wb.pin_memory()
        return PackedSeqBatchPart(span_x=span_xb, local_end=local_endb, y=yb, w=wb)

    def _pad_batch_part(self, part: PackedSeqBatchPart, pad_size: int) -> PackedSeqBatchPart:
        if pad_size <= 0:
            return part
        local_end = torch.cat([part.local_end, part.local_end[-1:].repeat(int(pad_size))], dim=0)
        y = torch.cat([part.y, part.y[-1:].repeat(int(pad_size))], dim=0)
        w = torch.cat([part.w, part.w[-1:].repeat(int(pad_size))], dim=0)
        if self.pin_memory:
            local_end = local_end.pin_memory()
            y = y.pin_memory()
            w = w.pin_memory()
        return PackedSeqBatchPart(span_x=part.span_x, local_end=local_end, y=y, w=w)

    def _load_batch_from_plan_item(self, item: BatchPlanItem) -> PackedSeqBatch:
        parts = [self._load_batch_part(part) for part in item.parts]
        if not parts:
            raise RuntimeError("Empty batch plan item")
        if item.pad_size > 0:
            parts[-1] = self._pad_batch_part(parts[-1], item.pad_size)
        return PackedSeqBatch(parts=tuple(parts))

    def __iter__(self):
        plan = self._build_rank_plan()
        if not plan:
            return
        if self.loader_threads <= 1 and self.prefetch_batches <= 1:
            for item in plan:
                yield self._load_batch_from_plan_item(item)
            return

        submit_ahead = max(self.prefetch_batches, self.loader_threads)
        futures: List[Future] = []
        next_idx = 0
        with ThreadPoolExecutor(max_workers=self.loader_threads) as executor:
            while next_idx < len(plan) and len(futures) < submit_ahead:
                futures.append(executor.submit(self._load_batch_from_plan_item, plan[next_idx]))
                next_idx += 1
            while futures:
                fut = futures.pop(0)
                yield fut.result()
                if next_idx < len(plan):
                    futures.append(executor.submit(self._load_batch_from_plan_item, plan[next_idx]))
                    next_idx += 1


def move_batch_to_device(
    batch: PackedSeqBatch,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> PackedSeqBatch:
    moved_parts = []
    for part in batch.parts:
        moved_parts.append(
            PackedSeqBatchPart(
                span_x=part.span_x.to(device, non_blocking=copy_non_blocking),
                local_end=part.local_end.to(device, non_blocking=copy_non_blocking),
                y=part.y.to(device, non_blocking=copy_non_blocking),
                w=part.w.to(device, non_blocking=copy_non_blocking),
            )
        )
    return PackedSeqBatch(parts=tuple(moved_parts))


class DeviceTransferPrefetchLoader:
    def __init__(
        self,
        loader,
        device: torch.device,
        prefetch_batches: int = 2,
        copy_non_blocking: bool = False,
    ):
        self.loader = loader
        self.device = device
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.copy_non_blocking = bool(copy_non_blocking)

    def __len__(self) -> int:
        return len(self.loader)

    def __getattr__(self, name: str):
        return getattr(self.loader, name)

    def set_epoch(self, epoch: int) -> None:
        if hasattr(self.loader, "set_epoch"):
            self.loader.set_epoch(epoch)

    def __iter__(self):
        if self.device.type != "cuda":
            for batch in self.loader:
                yield batch
            return
        result_q: queue.Queue = queue.Queue(maxsize=self.prefetch_batches)
        sentinel = object()
        device = self.device
        stop_event = threading.Event()

        def put_result(item, event) -> bool:
            while not stop_event.is_set():
                try:
                    result_q.put((item, event), timeout=0.1)
                    return True
                except queue.Full:
                    continue
            return False

        def worker():
            try:
                torch.cuda.set_device(device)
                stream = torch.cuda.Stream(device=device)
                for batch in self.loader:
                    if stop_event.is_set():
                        break
                    with torch.cuda.stream(stream):
                        moved = move_batch_to_device(
                            batch,
                            device=device,
                            copy_non_blocking=self.copy_non_blocking,
                        )
                        event = torch.cuda.Event()
                        event.record(stream)
                    if not put_result(moved, event):
                        return
                put_result(sentinel, None)
            except Exception as exc:  # pragma: no cover - worker thread failure propagation
                put_result(exc, None)

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        current_stream = torch.cuda.current_stream(device)
        try:
            while True:
                item, event = result_q.get()
                if item is sentinel:
                    break
                if isinstance(item, Exception):
                    raise item
                current_stream.wait_event(event)
                for part in item.parts:
                    part.span_x.record_stream(current_stream)
                    part.local_end.record_stream(current_stream)
                    part.y.record_stream(current_stream)
                    part.w.record_stream(current_stream)
                yield item
        finally:
            stop_event.set()
            thread.join()


class MaterializeDevicePrefetchLoader:
    def __init__(
        self,
        loader,
        device: torch.device,
        seq_offsets: torch.Tensor,
        feat_mean: torch.Tensor,
        feat_std: torch.Tensor,
        prefetch_batches: int = 2,
        copy_non_blocking: bool = False,
    ):
        self.loader = loader
        self.device = device
        self.seq_offsets = seq_offsets
        self.feat_mean = feat_mean
        self.feat_std = feat_std
        self.prefetch_batches = max(1, int(prefetch_batches))
        self.copy_non_blocking = bool(copy_non_blocking)

    def __len__(self) -> int:
        return len(self.loader)

    def __getattr__(self, name: str):
        return getattr(self.loader, name)

    def set_epoch(self, epoch: int) -> None:
        if hasattr(self.loader, "set_epoch"):
            self.loader.set_epoch(epoch)

    def __iter__(self):
        if self.device.type != "cuda":
            for batch in self.loader:
                xb, yb, wb = materialize_batch(
                    batch,
                    seq_offsets=self.seq_offsets,
                    feat_mean=self.feat_mean,
                    feat_std=self.feat_std,
                    device=self.device,
                    copy_non_blocking=self.copy_non_blocking,
                )
                yield MaterializedSeqBatch(xb=xb, y=yb, w=wb)
            return
        result_q: queue.Queue = queue.Queue(maxsize=self.prefetch_batches)
        sentinel = object()
        device = self.device
        stop_event = threading.Event()

        def put_result(item, event) -> bool:
            while not stop_event.is_set():
                try:
                    result_q.put((item, event), timeout=0.1)
                    return True
                except queue.Full:
                    continue
            return False

        def worker():
            try:
                torch.cuda.set_device(device)
                stream = torch.cuda.Stream(device=device)
                for batch in self.loader:
                    if stop_event.is_set():
                        break
                    with torch.cuda.stream(stream):
                        xb, yb, wb = materialize_batch(
                            batch,
                            seq_offsets=self.seq_offsets,
                            feat_mean=self.feat_mean,
                            feat_std=self.feat_std,
                            device=device,
                            copy_non_blocking=self.copy_non_blocking,
                        )
                        event = torch.cuda.Event()
                        event.record(stream)
                    prefetched = MaterializedSeqBatch(xb=xb, y=yb, w=wb)
                    if not put_result(prefetched, event):
                        return
                put_result(sentinel, None)
            except Exception as exc:  # pragma: no cover - worker thread failure propagation
                put_result(exc, None)

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        current_stream = torch.cuda.current_stream(device)
        try:
            while True:
                item, event = result_q.get()
                if item is sentinel:
                    break
                if isinstance(item, Exception):
                    raise item
                current_stream.wait_event(event)
                item.xb.record_stream(current_stream)
                item.y.record_stream(current_stream)
                item.w.record_stream(current_stream)
                yield item
        finally:
            stop_event.set()
            thread.join()


class ParallelCNN1D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        branch_channels = out_channels // 3
        self.conv3 = nn.Conv1d(in_channels, branch_channels, kernel_size=3, padding=1)
        self.conv5 = nn.Conv1d(in_channels, branch_channels, kernel_size=5, padding=2)
        self.conv7 = nn.Conv1d(in_channels, out_channels - 2 * branch_channels, kernel_size=7, padding=3)
        self.act = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_t = x.transpose(1, 2)
        c3 = self.conv3(x_t)
        c5 = self.conv5(x_t)
        c7 = self.conv7(x_t)
        out = torch.cat([c3, c5, c7], dim=1)
        return self.act(out).transpose(1, 2)


class FeatureChannelGate(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 64, init_bias: float = 2.0):
        super().__init__()
        hidden_dim = max(1, int(hidden_dim))
        self.norm = nn.LayerNorm(input_dim)
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.act = nn.SiLU()
        self.fc2 = nn.Linear(hidden_dim, input_dim)
        nn.init.zeros_(self.fc2.weight)
        nn.init.constant_(self.fc2.bias, float(init_bias))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pooled = x.mean(dim=1)
        gate = self.fc2(self.act(self.fc1(self.norm(pooled))))
        gate = torch.sigmoid(gate).unsqueeze(1)
        return x * gate


class GRURegressor(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        num_layers: int,
        dropout: float,
        pooling: str = "last",
        bidirectional: bool = False,
        use_cnn1d: bool = False,
        input_gate_hidden_dim: int = 0,
        input_gate_bias: float = 2.0,
    ):
        super().__init__()
        self.pooling = str(pooling).lower()
        if self.pooling not in {"last", "attn"}:
            raise ValueError(f"Unsupported pooling: {pooling}")
        self.bidirectional = bool(bidirectional)
        self.use_cnn1d = bool(use_cnn1d)
        self.input_gate_hidden_dim = max(0, int(input_gate_hidden_dim))
        self.output_dim = int(hidden_dim) * (2 if self.bidirectional else 1)
        if self.input_gate_hidden_dim > 0:
            self.input_gate = FeatureChannelGate(
                input_dim=input_dim,
                hidden_dim=self.input_gate_hidden_dim,
                init_bias=float(input_gate_bias),
            )
        else:
            self.input_gate = nn.Identity()

        if self.use_cnn1d:
            self.cnn = ParallelCNN1D(input_dim, hidden_dim)
            rnn_input_dim = hidden_dim
        else:
            self.cnn = nn.Identity()
            rnn_input_dim = input_dim

        self.rnn = nn.GRU(
            input_size=rnn_input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            dropout=dropout if num_layers > 1 else 0.0,
            batch_first=True,
            bidirectional=self.bidirectional,
        )
        if self.pooling == "attn":
            self.attn_norm = nn.LayerNorm(self.output_dim)
            self.attn_proj = nn.Linear(self.output_dim, 1)
        head_hidden_dim = max(1, self.output_dim // 2)
        self.head = nn.Sequential(
            nn.LayerNorm(self.output_dim),
            nn.Linear(self.output_dim, head_hidden_dim),
            nn.ReLU(),
            nn.Linear(head_hidden_dim, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_gate(x)
        if self.use_cnn1d:
            x = self.cnn(x)
        out, _ = self.rnn(x)
        if self.pooling == "attn":
            score = self.attn_proj(self.attn_norm(out)).squeeze(-1)
            weight = torch.softmax(score, dim=1).unsqueeze(-1)
            pooled = (out * weight).sum(dim=1)
        else:
            pooled = out[:, -1, :]
        return self.head(pooled).squeeze(-1)


def weighted_mse(pred: torch.Tensor, y: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    w = torch.clamp(w, min=0.0)
    w_norm = w / torch.clamp(w.mean(), min=1e-6)
    return (w_norm * (pred - y) ** 2).mean()


def build_val_epoch_checkpoint_path(out_root: Path, epoch: int) -> Path:
    return out_root / f"gru_seq_memmap_val_epoch_{int(epoch):03d}.pt"


def refresh_top_val_checkpoints(
    top_val_checkpoints: List[dict],
    candidate_epoch: int,
    candidate_val_metrics: dict,
    keep_topk: int,
    out_root: Path,
) -> Tuple[List[dict], bool, List[dict]]:
    keep_topk = max(1, int(keep_topk))
    candidate = {
        "epoch": int(candidate_epoch),
        "val_ic": float(candidate_val_metrics["ic"]),
        "val_unweighted_ic": float(candidate_val_metrics.get("unweighted_ic", candidate_val_metrics["ic"])),
        "val_weighted_ic": float(candidate_val_metrics.get("weighted_ic", candidate_val_metrics["ic"])),
        "val_loss": float(candidate_val_metrics["loss"]),
        "path": str(build_val_epoch_checkpoint_path(out_root, candidate_epoch)),
    }
    updated = list(top_val_checkpoints)
    updated.append(candidate)
    updated.sort(key=lambda item: (-float(item.get("val_weighted_ic", item["val_ic"])), int(item["epoch"])))
    kept = [dict(item) for item in updated[:keep_topk]]
    dropped = [dict(item) for item in updated[keep_topk:]]
    kept_epochs = {int(item["epoch"]) for item in kept}
    entered_topk = int(candidate_epoch) in kept_epochs
    for rank_i, item in enumerate(kept, start=1):
        item["rank"] = int(rank_i)
    return kept, entered_topk, dropped


class CorrStats:
    def __init__(self):
        self.buf: torch.Tensor | None = None

    def update(self, pred: torch.Tensor, y: torch.Tensor, w: torch.Tensor | None = None):
        p = pred.detach().float().reshape(-1)
        t = y.detach().float().reshape(-1)
        mask = torch.isfinite(p) & torch.isfinite(t)
        if w is not None:
            ww = w.detach().float().reshape(-1)
            mask = mask & torch.isfinite(ww) & (ww > 0)
        else:
            ww = None
        zero = torch.zeros_like(p)
        p = torch.where(mask, p, zero).to(dtype=torch.float64)
        t = torch.where(mask, t, zero).to(dtype=torch.float64)
        d = p - t
        if ww is None:
            weight = mask.to(dtype=torch.float64)
        else:
            weight = torch.where(mask, ww, zero).to(dtype=torch.float64)
        sums = torch.stack(
            [
                mask.to(dtype=torch.float64).sum(),
                p.sum(),
                t.sum(),
                (p * p).sum(),
                (t * t).sum(),
                (p * t).sum(),
                torch.abs(d).sum(),
                (d * d).sum(),
                weight.sum(),
                (weight * p).sum(),
                (weight * t).sum(),
                (weight * p * p).sum(),
                (weight * t * t).sum(),
                (weight * p * t).sum(),
                (weight * torch.abs(d)).sum(),
                (weight * d * d).sum(),
            ]
        )
        if self.buf is None:
            self.buf = sums
        else:
            self.buf = self.buf + sums

    def to_tensor(self, device: torch.device) -> torch.Tensor:
        if self.buf is None:
            return torch.zeros((16,), dtype=torch.float64, device=device)
        return self.buf.to(device=device, dtype=torch.float64)


def corr_from_sums(sum_x: float, sum_y: float, sum_xx: float, sum_yy: float, sum_xy: float, denom_weight: float) -> float:
    if (not np.isfinite(denom_weight)) or denom_weight <= 0:
        return float("nan")
    mean_x = sum_x / denom_weight
    mean_y = sum_y / denom_weight
    var_x = max(sum_xx / denom_weight - mean_x * mean_x, 1e-12)
    var_y = max(sum_yy / denom_weight - mean_y * mean_y, 1e-12)
    cov_xy = sum_xy / denom_weight - mean_x * mean_y
    return float(cov_xy / math.sqrt(var_x * var_y))


def metrics_from_tensor(t: torch.Tensor) -> dict:
    values = [float(x) for x in t.tolist()]
    if len(values) < 8:
        raise ValueError(f"metrics_from_tensor expects at least 8 values, got {len(values)}")
    n, sum_p, sum_y, sum_pp, sum_yy, sum_py, sum_abs, sum_sq = values[:8]
    if (not np.isfinite(n)) or n <= 0:
        return {
            "mse": float("nan"),
            "rmse": float("nan"),
            "mae": float("nan"),
            "ic": float("nan"),
            "unweighted_mse": float("nan"),
            "unweighted_rmse": float("nan"),
            "unweighted_mae": float("nan"),
            "unweighted_ic": float("nan"),
            "weighted_mse": float("nan"),
            "weighted_rmse": float("nan"),
            "weighted_mae": float("nan"),
            "weighted_ic": float("nan"),
            "weight_sum": 0.0,
            "n": 0,
        }
    mse = sum_sq / n
    mae = sum_abs / n
    ic = corr_from_sums(sum_p, sum_y, sum_pp, sum_yy, sum_py, n)
    if len(values) >= 16:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy, sum_wabs, sum_wsq = values[8:16]
    elif len(values) >= 14:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy = values[8:14]
        sum_wabs, sum_wsq = sum_abs, sum_sq
    else:
        weight_sum, sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy = n, sum_p, sum_y, sum_pp, sum_yy, sum_py
        sum_wabs, sum_wsq = sum_abs, sum_sq
    weighted_ic = corr_from_sums(sum_wp, sum_wy, sum_wpp, sum_wyy, sum_wpy, weight_sum)
    weighted_mse = (sum_wsq / weight_sum) if weight_sum > 0 else float("nan")
    weighted_mae = (sum_wabs / weight_sum) if weight_sum > 0 else float("nan")
    weighted_rmse = math.sqrt(weighted_mse) if np.isfinite(weighted_mse) and weighted_mse >= 0 else float("nan")
    return {
        "mse": float(weighted_mse),
        "rmse": float(weighted_rmse),
        "mae": float(weighted_mae),
        "ic": float(weighted_ic),
        "unweighted_mse": float(mse),
        "unweighted_rmse": float(math.sqrt(mse)),
        "unweighted_mae": float(mae),
        "unweighted_ic": float(ic),
        "weighted_mse": float(weighted_mse),
        "weighted_rmse": float(weighted_rmse),
        "weighted_mae": float(weighted_mae),
        "weighted_ic": float(weighted_ic),
        "weight_sum": float(weight_sum),
        "n": int(n),
    }


def compute_feature_stats(
    train_days: List[Path],
    prefer_fp16: bool,
    sample_stride: int = 50,
    chunk_sample_rows: int = 250_000,
) -> Tuple[np.ndarray, np.ndarray]:
    s = None
    s2 = None
    n = 0
    for day in train_days:
        store = load_day_store(day, prefer_fp16=prefer_fp16)
        chunk_span = max(int(sample_stride), int(sample_stride) * max(1, int(chunk_sample_rows)))
        for start in range(0, store.n_rows, chunk_span):
            stop = min(store.n_rows, start + chunk_span)
            # Stream smaller memmap slices to reduce pressure and avoid giant one-shot reads.
            arr = np.array(store.x[start:stop:sample_stride], dtype=np.float32, copy=True)
            if arr.size == 0:
                continue
            arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
            if s is None:
                s = arr.sum(axis=0, dtype=np.float64)
                s2 = (arr * arr).sum(axis=0, dtype=np.float64)
            else:
                s += arr.sum(axis=0, dtype=np.float64)
                s2 += (arr * arr).sum(axis=0, dtype=np.float64)
            n += arr.shape[0]
    mean = (s / max(n, 1)).astype(np.float32)
    var = (s2 / max(n, 1) - mean.astype(np.float64) ** 2).astype(np.float32)
    std = np.sqrt(np.clip(var, 1e-8, None)).astype(np.float32)
    return mean, std


def build_feature_stats_cache_path(
    output_root: Path,
    row_root: Path,
    train_days: List[Path],
    prefer_fp16: bool,
    sample_stride: int,
) -> Path:
    cache_dir = ensure_dir(output_root / "training_seq" / "_feature_stats_cache")
    key_src = "\n".join(
        [
            str(row_root),
            str(bool(prefer_fp16)),
            str(int(sample_stride)),
            *[d.name for d in train_days],
        ]
    )
    key = hashlib.sha1(key_src.encode("utf-8")).hexdigest()[:16]
    return cache_dir / f"feature_stats_{key}.npz"


def load_feature_stats_cache(path: Path) -> Tuple[np.ndarray, np.ndarray]:
    with np.load(path) as arr:
        mean = arr["mean"].astype(np.float32, copy=False)
        std = arr["std"].astype(np.float32, copy=False)
    return mean, std


def wait_for_feature_stats_cache(path: Path, timeout_seconds: int = 7200) -> Tuple[np.ndarray, np.ndarray]:
    deadline = time.time() + max(1, int(timeout_seconds))
    last_err: Exception | None = None
    while time.time() < deadline:
        if path.exists() and path.stat().st_size > 0:
            try:
                return load_feature_stats_cache(path)
            except Exception as exc:  # pragma: no cover - transient partial-write case
                last_err = exc
        time.sleep(2.0)
    if last_err is not None:
        raise TimeoutError(f"Timed out waiting for feature stats cache {path}: {last_err}") from last_err
    raise TimeoutError(f"Timed out waiting for feature stats cache {path}")


def evaluate(
    model,
    loader,
    feat_mean,
    feat_std,
    seq_offsets,
    accelerator,
    split_name: str = "eval",
    progress_parts: int = 4,
    copy_non_blocking: bool = False,
    collect_timing: bool = False,
) -> dict:
    model.eval()
    stats = CorrStats()
    timing_labels = [
        "batch_wait_sec",
        "materialize_sec",
        "forward_sec",
        "stats_update_sec",
        "step_total_sec",
    ]
    timing_sums = np.zeros((len(timing_labels),), dtype=np.float64)
    local_batches = 0
    total_batches = 0
    try:
        total_batches = len(loader)
    except TypeError:
        total_batches = 0
    progress_step = max(1, total_batches // max(1, int(progress_parts))) if total_batches > 0 else 0
    with torch.no_grad():
        loader_iter = iter(loader)
        batch_i = 0
        while True:
            t_step0 = time.perf_counter()
            t0 = time.perf_counter()
            try:
                batch = next(loader_iter)
            except StopIteration:
                break
            t1 = time.perf_counter()
            xb, yb, wb = resolve_model_batch(
                batch,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                device=accelerator.device,
                copy_non_blocking=copy_non_blocking,
            )
            if collect_timing:
                sync_if_needed(accelerator.device)
            t2 = time.perf_counter()
            with accelerator.autocast():
                pred = model(xb)
                loss = weighted_mse(pred, yb, wb)
            if collect_timing:
                sync_if_needed(accelerator.device)
            t3 = time.perf_counter()
            stats.update(pred, yb, wb)
            if collect_timing:
                sync_if_needed(accelerator.device)
            t4 = time.perf_counter()
            local_batches += 1
            batch_i += 1
            if collect_timing:
                timing_sums += np.asarray(
                    [
                        t1 - t0,
                        t2 - t1,
                        t3 - t2,
                        t4 - t3,
                        t4 - t_step0,
                    ],
                    dtype=np.float64,
                )
            if (
                accelerator.is_main_process
                and total_batches > 0
                and (batch_i % progress_step == 0 or batch_i == total_batches)
            ):
                print(f"[{split_name}] progress {batch_i}/{total_batches}", flush=True)
    stats_t = accelerator.reduce(stats.to_tensor(accelerator.device), reduction="sum")
    m = metrics_from_tensor(stats_t)
    m["loss"] = float(m.get("weighted_mse", float("nan")))
    if collect_timing:
        m["_timing"] = summarize_timing_across_ranks(
            accelerator=accelerator,
            labels=timing_labels,
            count=local_batches,
            sums=timing_sums,
            count_key="measured_batches",
            global_batch=int(loader.batch_size) * int(accelerator.num_processes),
        )
    return m


def materialize_batch(
    batch: PackedSeqBatch,
    seq_offsets: torch.Tensor,
    feat_mean: torch.Tensor,
    feat_std: torch.Tensor,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    xb_parts = []
    y_parts = []
    w_parts = []
    for part in batch.parts:
        span_x = part.span_x.to(device, non_blocking=copy_non_blocking)
        local_end = part.local_end.to(device, non_blocking=copy_non_blocking)
        yb = part.y.to(device, non_blocking=copy_non_blocking)
        wb = part.w.to(device, non_blocking=copy_non_blocking)
        seq_idx = local_end[:, None] + seq_offsets[None, :]
        xb = span_x[seq_idx].float()
        xb = (xb - feat_mean) / feat_std
        # Keep GRU inputs in a standard dense layout so cuDNN does not fall back
        # to a slower path for batches assembled from wide memmap spans.
        xb = xb.contiguous()
        xb_parts.append(xb)
        y_parts.append(yb)
        w_parts.append(wb)
    if len(xb_parts) == 1:
        return xb_parts[0], y_parts[0], w_parts[0]
    return torch.cat(xb_parts, dim=0), torch.cat(y_parts, dim=0), torch.cat(w_parts, dim=0)


def resolve_model_batch(
    batch,
    seq_offsets: torch.Tensor,
    feat_mean: torch.Tensor,
    feat_std: torch.Tensor,
    device: torch.device,
    copy_non_blocking: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if isinstance(batch, MaterializedSeqBatch):
        return batch.xb, batch.y, batch.w
    return materialize_batch(
        batch,
        seq_offsets=seq_offsets,
        feat_mean=feat_mean,
        feat_std=feat_std,
        device=device,
        copy_non_blocking=copy_non_blocking,
    )


def sync_if_needed(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def warmup_collective(accelerator) -> float:
    if int(getattr(accelerator, "num_processes", 1)) <= 1:
        return 0.0
    t0 = time.perf_counter()
    dummy = torch.zeros((1,), dtype=torch.float32, device=accelerator.device)
    _ = accelerator.reduce(dummy, reduction="sum")
    sync_if_needed(accelerator.device)
    return float(time.perf_counter() - t0)


def summarize_timing_across_ranks(
    accelerator: Accelerator,
    labels: List[str],
    count: int,
    sums: np.ndarray,
    count_key: str,
    global_batch: int | None = None,
) -> dict:
    local = torch.tensor([float(count), *sums.tolist()], dtype=torch.float64, device=accelerator.device)[None, :]
    gathered = accelerator.gather(local)
    empty = {
        count_key: int(count),
        "per_rank": [],
        "global_step_sec_estimate": float("nan"),
        "global_samples_per_sec_estimate": float("nan"),
    }
    if not accelerator.is_main_process:
        return empty
    gathered_np = gathered.detach().cpu().numpy()
    per_rank = []
    max_step_sec = 0.0
    step_label = labels[-1] if labels else None
    for rank_i, row in enumerate(gathered_np):
        rank_count = max(1.0, float(row[0]))
        stats = {"rank": int(rank_i), count_key: int(row[0])}
        for label_i, label in enumerate(labels, start=1):
            stats[label] = float(row[label_i] / rank_count)
        if step_label is not None:
            max_step_sec = max(max_step_sec, float(stats[step_label]))
        per_rank.append(stats)
    out = {
        count_key: int(count),
        "per_rank": per_rank,
        "global_step_sec_estimate": float(max_step_sec) if step_label is not None else float("nan"),
        "global_samples_per_sec_estimate": float("nan"),
    }
    if step_label is not None and global_batch is not None:
        out["global_samples_per_sec_estimate"] = float(global_batch / max(max_step_sec, 1e-9))
    return out


def profile_train_steps(
    model,
    optimizer,
    loader,
    accelerator,
    feat_mean,
    feat_std,
    seq_offsets,
    copy_non_blocking: bool,
    warmup_steps: int,
    profile_steps: int,
    grad_accum_steps: int = 1,
) -> dict:
    total_batches = 0
    try:
        total_batches = len(loader)
    except TypeError:
        total_batches = 0
    warmup = max(0, int(warmup_steps))
    active = max(0, int(profile_steps))
    if total_batches <= 0 or active <= 0:
        return {
            "warmup_steps": warmup,
            "profile_steps": active,
            "measured_steps": 0,
            "per_rank": [],
            "global_step_sec_estimate": float("nan"),
            "global_samples_per_sec_estimate": float("nan"),
        }
    run_steps = min(total_batches, warmup + active)
    measured_steps = max(0, run_steps - warmup)
    labels = [
        "batch_wait_sec",
        "materialize_sec",
        "forward_sec",
        "backward_sec",
        "optimizer_step_sec",
        "step_total_sec",
    ]
    sums = np.zeros((len(labels),), dtype=np.float64)
    model.train()
    train_iter = iter(loader)
    warmup_collective(accelerator)
    if accelerator.device.type == "cuda":
        sync_if_needed(accelerator.device)
        torch.cuda.reset_peak_memory_stats(accelerator.device)
    accum_steps = max(1, int(grad_accum_steps))
    accum_micro = 0
    optimizer.zero_grad(set_to_none=True)
    for step_i in range(1, run_steps + 1):
        if accelerator.device.type == "cuda" and step_i == warmup + 1:
            sync_if_needed(accelerator.device)
            torch.cuda.reset_peak_memory_stats(accelerator.device)
        t_step0 = time.perf_counter()
        t0 = time.perf_counter()
        batch = next(train_iter)
        t1 = time.perf_counter()
        accum_micro += 1
        last_batch = step_i == run_steps
        should_step = accum_micro >= accum_steps or last_batch
        group_total = min(accum_steps, accum_micro + max(0, run_steps - step_i))
        xb, yb, wb = resolve_model_batch(
            batch,
            seq_offsets=seq_offsets,
            feat_mean=feat_mean,
            feat_std=feat_std,
            device=accelerator.device,
            copy_non_blocking=copy_non_blocking,
        )
        sync_if_needed(accelerator.device)
        t2 = time.perf_counter()
        sync_ctx = model.no_sync() if hasattr(model, "no_sync") and not should_step else nullcontext()
        with sync_ctx:
            with accelerator.autocast():
                pred = model(xb)
                loss_raw = weighted_mse(pred, yb, wb)
                loss = loss_raw / float(group_total)
            sync_if_needed(accelerator.device)
            t3 = time.perf_counter()
            accelerator.backward(loss)
            sync_if_needed(accelerator.device)
            t4 = time.perf_counter()
        if should_step:
            accelerator.step_optimizer(optimizer)
            optimizer.zero_grad(set_to_none=True)
            accum_micro = 0
        sync_if_needed(accelerator.device)
        t5 = time.perf_counter()
        if step_i > warmup:
            sums += np.asarray(
                [
                    t1 - t0,
                    t2 - t1,
                    t3 - t2,
                    t4 - t3,
                    t5 - t4,
                    t5 - t_step0,
                ],
                dtype=np.float64,
            )
        if accelerator.is_main_process and (step_i == run_steps or step_i % max(1, run_steps // 4) == 0):
            print(f"[profile-train] step {step_i}/{run_steps} loss={loss_raw.detach().float().item():.6f}", flush=True)
    peak_allocated_bytes = 0.0
    peak_reserved_bytes = 0.0
    if accelerator.device.type == "cuda":
        sync_if_needed(accelerator.device)
        peak_allocated_bytes = float(torch.cuda.max_memory_allocated(accelerator.device))
        peak_reserved_bytes = float(torch.cuda.max_memory_reserved(accelerator.device))
    global_batch = int(loader.batch_size) * int(accelerator.num_processes)
    summary = summarize_timing_across_ranks(
        accelerator=accelerator,
        labels=labels,
        count=measured_steps,
        sums=sums,
        count_key="measured_steps",
        global_batch=global_batch,
    )
    summary["warmup_steps"] = warmup
    summary["profile_steps"] = active
    if accelerator.device.type == "cuda":
        local_mem = torch.tensor(
            [peak_allocated_bytes, peak_reserved_bytes],
            dtype=torch.float64,
            device=accelerator.device,
        )[None, :]
        gathered_mem = accelerator.gather(local_mem)
        if accelerator.is_main_process:
            gathered_mem_np = gathered_mem.detach().cpu().numpy()
            peak_allocated_max = 0.0
            peak_reserved_max = 0.0
            for rank_i, row in enumerate(gathered_mem_np):
                allocated_bytes = float(row[0])
                reserved_bytes = float(row[1])
                peak_allocated_max = max(peak_allocated_max, allocated_bytes)
                peak_reserved_max = max(peak_reserved_max, reserved_bytes)
                if rank_i < len(summary.get("per_rank", [])):
                    summary["per_rank"][rank_i]["peak_allocated_bytes"] = allocated_bytes
                    summary["per_rank"][rank_i]["peak_reserved_bytes"] = reserved_bytes
                    summary["per_rank"][rank_i]["peak_allocated_gib"] = allocated_bytes / float(1024**3)
                    summary["per_rank"][rank_i]["peak_reserved_gib"] = reserved_bytes / float(1024**3)
            summary["peak_allocated_bytes_max"] = peak_allocated_max
            summary["peak_reserved_bytes_max"] = peak_reserved_max
            summary["peak_allocated_gib_max"] = peak_allocated_max / float(1024**3)
            summary["peak_reserved_gib_max"] = peak_reserved_max / float(1024**3)
    return summary


def run(args):
    run_t0 = time.perf_counter()
    phase_timings: Dict[str, float] = {}
    phase_details: Dict[str, dict] = {}
    data_prepare_t0 = time.perf_counter()
    set_global_seed(args.seed)
    cfg = load_config(args.config)
    train_cfg = dict(cfg["train"])
    if args.override_epochs > 0:
        train_cfg["epochs"] = args.override_epochs
    if args.override_lr > 0:
        train_cfg["learning_rate"] = args.override_lr
    if args.override_batch_size > 0:
        train_cfg["batch_size"] = args.override_batch_size

    row_root = resolve_row_root(cfg, args.cache_name, row_root_override=args.row_root_override)
    profile_phases = bool(args.profile_phases)
    profile_warmup_steps = max(0, int(args.profile_warmup_steps))
    profile_steps = max(0, int(args.profile_steps))
    stop_after_profile = bool(args.stop_after_profile)
    profile_only = bool(stop_after_profile and profile_steps > 0)
    grad_accum_steps = max(1, int(args.grad_accum_steps))
    split_view_meta = load_split_view_meta(row_root)
    if split_view_meta is not None:
        train_days = resolve_named_day_dirs(row_root, split_view_meta.get("train_days"))
        valid_days = resolve_named_day_dirs(row_root, split_view_meta.get("valid_days"))
        test_days = resolve_named_day_dirs(row_root, split_view_meta.get("test_days"))
        train_days = limit_days(train_days, args.train_day_limit)
        valid_days = limit_days(valid_days, args.valid_day_limit)
        test_days = limit_days(test_days, args.test_day_limit)
        train_pool_days = list(train_days) + list(valid_days)
        days = list(train_days) + list(valid_days) + list(test_days)
        split_ratio = float(args.valid_split_ratio) if args.valid_split_ratio > 0.0 else float(train_cfg["train_split_ratio"])
        train_cfg["train_split_ratio"] = split_ratio
    else:
        days = list_ready_days(row_root)
        use_explicit_split = bool(
            args.train_pool_start_date or args.train_pool_end_date or args.test_start_date or args.test_end_date
        )
        if args.max_days > 0 and not use_explicit_split:
            days = days[: args.max_days]
        if args.test_start_date or args.test_end_date:
            test_days = filter_days_by_date(days, args.test_start_date, args.test_end_date)
        else:
            test_days = []
        test_day_names = {d.name for d in test_days}
        raw_train_pool_days = filter_days_by_date(days, args.train_pool_start_date, args.train_pool_end_date)
        if args.train_pool_start_date or args.train_pool_end_date:
            train_pool_days = [d for d in raw_train_pool_days if d.name not in test_day_names]
        else:
            train_pool_days = [d for d in days if d.name not in test_day_names]
        split_ratio = float(args.valid_split_ratio) if args.valid_split_ratio > 0.0 else float(train_cfg["train_split_ratio"])
        train_cfg["train_split_ratio"] = split_ratio
        train_days, valid_days = split_days(train_pool_days, split_ratio)
        train_days = limit_days(train_days, args.train_day_limit)
        valid_days = limit_days(valid_days, args.valid_day_limit)
        test_days = limit_days(test_days, args.test_day_limit)
    if not train_pool_days:
        raise RuntimeError("train_pool_days is empty after applying date filters.")
    if len(train_days) == 0 or len(valid_days) == 0:
        raise RuntimeError("Need both train and valid day split.")
    print(
        f"[data] days_total={len(days)} train_pool_days={len(train_pool_days)} "
        f"train_days={len(train_days)} valid_days={len(valid_days)} test_days={len(test_days)}",
        flush=True,
    )
    print(f"[row-root] {row_root}", flush=True)
    if args.train_pool_start_date or args.train_pool_end_date:
        print(
            f"[data] train_pool_date_window=[{args.train_pool_start_date or 'min'}, "
            f"{args.train_pool_end_date or 'max'}]",
            flush=True,
        )
    if args.test_start_date or args.test_end_date:
        print(
            f"[data] test_date_window=[{args.test_start_date or 'min'}, {args.test_end_date or 'max'}]",
            flush=True,
        )
    data_prepare_local_sec = time.perf_counter() - data_prepare_t0
    accelerator_init_t0 = time.perf_counter()
    accelerator = build_runtime(
        args.runtime_backend,
        bool(args.use_amp),
        enable_static_graph=(grad_accum_steps == 1),
    )
    phase_timings["data_prepare_sec"] = float(data_prepare_local_sec)
    phase_timings["accelerator_init_sec"] = float(time.perf_counter() - accelerator_init_t0)
    loader_threads = max(1, int(args.num_workers))
    eval_batch_size = int(args.eval_batch_size) if int(args.eval_batch_size) > 0 else int(train_cfg["batch_size"])
    device_transfer_mode = str(args.device_transfer_mode).strip().lower()
    device_transfer_prefetch_batches = max(1, int(args.device_transfer_prefetch_batches))
    configure_split_view_weight_loading(bool(args.use_source_raw_weights))
    min_timecode = int(args.min_timecode)
    require_positive_weight = bool(args.require_positive_weight)
    # Keep host tensors pinned even on the threaded transfer path so H2D copies
    # can overlap with compute on the prefetch stream.
    use_pinned_transfer = bool(train_cfg.get("pin_memory", False))
    if int(args.loader_prefetch_batches) > 0:
        loader_prefetch = max(1, int(args.loader_prefetch_batches))
    else:
        loader_prefetch = max(1, int(train_cfg.get("prefetch_factor", 2)))
    train_batch_overlap = max(0, int(args.train_batch_overlap))
    skip_train_eval = bool(args.skip_train_eval)
    skip_test = bool(args.skip_test)
    print(
        f"[device-transfer] mode={device_transfer_mode} loader_pin_memory={int(use_pinned_transfer)} "
        f"prefetch_batches={device_transfer_prefetch_batches}",
        flush=True,
    )
    print(
        f"[sample-filter] min_timecode={min_timecode} require_positive_weight={int(require_positive_weight)} "
        f"use_source_raw_weights={int(bool(args.use_source_raw_weights))}",
        flush=True,
    )
    if train_batch_overlap > 0:
        print(
            f"[train-loader] batch_overlap={train_batch_overlap} batch_stride={int(train_cfg['batch_size']) - train_batch_overlap}",
            flush=True,
        )
    end_index_cache_dir = ensure_dir(Path(cfg["paths"]["output_root"]) / "training_seq" / "_end_index_cache")
    loader_init_t0 = time.perf_counter()
    train_loader = SeqMemmapBatchLoader(
        train_days,
        seq_len=args.seq_len,
        sample_stride=args.sample_stride,
        batch_size=int(train_cfg["batch_size"]),
        shuffle=True,
        max_samples=args.max_samples,
        prefer_fp16=bool(args.prefer_fp16),
        rank=accelerator.process_index,
        world_size=accelerator.num_processes,
        seed=args.seed,
        pin_memory=use_pinned_transfer,
        loader_threads=loader_threads,
        prefetch_batches=loader_prefetch,
        pad_last_batch=(accelerator.num_processes > 1),
        batch_overlap=train_batch_overlap,
        index_cache_dir=end_index_cache_dir,
        min_timecode=min_timecode,
        require_positive_weight=require_positive_weight,
    )
    phase_timings["train_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    train_eval_loader = (
        SeqMemmapBatchLoader(
            train_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=args.max_samples,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if not profile_only and not skip_train_eval and train_batch_overlap > 0
        else None
    )
    phase_timings["train_eval_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    valid_loader = (
        SeqMemmapBatchLoader(
            valid_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=args.max_samples,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if not profile_only
        else None
    )
    phase_timings["valid_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    loader_init_t0 = time.perf_counter()
    test_loader = (
        SeqMemmapBatchLoader(
            test_days,
            seq_len=args.seq_len,
            sample_stride=args.sample_stride,
            batch_size=eval_batch_size,
            shuffle=False,
            max_samples=-1,
            prefer_fp16=bool(args.prefer_fp16),
            rank=accelerator.process_index,
            world_size=accelerator.num_processes,
            seed=args.seed,
            pin_memory=use_pinned_transfer,
            loader_threads=loader_threads,
            prefetch_batches=loader_prefetch,
            pad_last_batch=False,
            batch_overlap=0,
            index_cache_dir=end_index_cache_dir,
            min_timecode=min_timecode,
            require_positive_weight=require_positive_weight,
        )
        if test_days and not skip_test and not profile_only
        else None
    )
    phase_timings["test_loader_init_sec"] = float(time.perf_counter() - loader_init_t0)
    stats_stride = max(20, args.sample_stride)
    feature_stats_cache = build_feature_stats_cache_path(
        output_root=Path(cfg["paths"]["output_root"]),
        row_root=row_root,
        train_days=train_days,
        prefer_fp16=bool(args.prefer_fp16),
        sample_stride=stats_stride,
    )
    feature_stats_t0 = time.perf_counter()
    if accelerator.is_main_process:
        if feature_stats_cache.exists() and feature_stats_cache.stat().st_size > 0:
            print(f"[feature-stats] reused_from={feature_stats_cache}", flush=True)
            feat_mean_np, feat_std_np = load_feature_stats_cache(feature_stats_cache)
        else:
            print(
                f"[feature-stats] computing cache={feature_stats_cache} sample_stride={stats_stride}",
                flush=True,
            )
            feat_mean_np, feat_std_np = compute_feature_stats(
                train_days, prefer_fp16=bool(args.prefer_fp16), sample_stride=stats_stride
            )
            tmp_path = feature_stats_cache.with_suffix(f".tmp.{int(time.time())}.npz")
            np.savez(tmp_path, mean=feat_mean_np, std=feat_std_np)
            tmp_path.replace(feature_stats_cache)
            print(f"[feature-stats] saved_cache={feature_stats_cache}", flush=True)
    else:
        print(f"[feature-stats] waiting_for_cache={feature_stats_cache}", flush=True)
        feat_mean_np, feat_std_np = wait_for_feature_stats_cache(feature_stats_cache)
        print(f"[feature-stats] loaded_cache={feature_stats_cache}", flush=True)
    accelerator.wait_for_everyone()
    phase_timings["feature_stats_sec"] = float(time.perf_counter() - feature_stats_t0)

    model_init_t0 = time.perf_counter()
    model = GRURegressor(
        input_dim=200,
        hidden_dim=args.hidden_dim,
        num_layers=args.num_layers,
        dropout=args.dropout,
        pooling=args.pooling,
        bidirectional=bool(args.bidirectional),
        use_cnn1d=bool(getattr(args, "use_cnn1d", 0)),
        input_gate_hidden_dim=int(getattr(args, "input_gate_hidden_dim", 0)),
        input_gate_bias=float(getattr(args, "input_gate_bias", 2.0)),
    )
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=float(train_cfg["learning_rate"]), weight_decay=float(args.weight_decay)
    )
    phase_timings["model_optimizer_init_sec"] = float(time.perf_counter() - model_init_t0)

    prepare_t0 = time.perf_counter()
    if test_loader is not None:
        model, optimizer = accelerator.prepare(model, optimizer)
    else:
        model, optimizer = accelerator.prepare(model, optimizer)
    phase_timings["accelerator_prepare_sec"] = float(time.perf_counter() - prepare_t0)
    tensor_init_t0 = time.perf_counter()
    feat_mean = torch.from_numpy(feat_mean_np).to(accelerator.device)[None, None, :]
    feat_std = torch.from_numpy(feat_std_np).to(accelerator.device)[None, None, :]
    seq_offsets = torch.arange(-(args.seq_len - 1), 1, dtype=torch.long, device=accelerator.device)
    phase_timings["device_tensor_init_sec"] = float(time.perf_counter() - tensor_init_t0)
    device_prefetch_t0 = time.perf_counter()
    if device_transfer_mode == "thread_prefetch":
        train_loader = DeviceTransferPrefetchLoader(
            train_loader,
            device=accelerator.device,
            prefetch_batches=device_transfer_prefetch_batches,
            copy_non_blocking=use_pinned_transfer,
        )
        if train_eval_loader is not None:
            train_eval_loader = DeviceTransferPrefetchLoader(
                train_eval_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if valid_loader is not None:
            valid_loader = DeviceTransferPrefetchLoader(
                valid_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if test_loader is not None:
            test_loader = DeviceTransferPrefetchLoader(
                test_loader,
                device=accelerator.device,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
    elif device_transfer_mode == "materialize_prefetch":
        train_loader = MaterializeDevicePrefetchLoader(
            train_loader,
            device=accelerator.device,
            seq_offsets=seq_offsets,
            feat_mean=feat_mean,
            feat_std=feat_std,
            prefetch_batches=device_transfer_prefetch_batches,
            copy_non_blocking=use_pinned_transfer,
        )
        if train_eval_loader is not None:
            train_eval_loader = MaterializeDevicePrefetchLoader(
                train_eval_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if valid_loader is not None:
            valid_loader = MaterializeDevicePrefetchLoader(
                valid_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
        if test_loader is not None:
            test_loader = MaterializeDevicePrefetchLoader(
                test_loader,
                device=accelerator.device,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                prefetch_batches=device_transfer_prefetch_batches,
                copy_non_blocking=use_pinned_transfer,
            )
    phase_timings["device_prefetch_wrap_sec"] = float(time.perf_counter() - device_prefetch_t0)

    output_init_t0 = time.perf_counter()
    out_root = ensure_dir(Path(cfg["paths"]["output_root"]) / "training_seq" / args.run_name)
    log_path = out_root / "train_log.csv"
    best_model_path = out_root / "gru_seq_memmap_best_val_ic.pt"
    save_topk_val_checkpoints = max(
        1,
        int(getattr(args, "save_topk_val_checkpoints", 1)),
        int(getattr(args, "test_checkpoint_rank", 1)),
    )
    selected_test_checkpoint_rank = max(1, int(getattr(args, "test_checkpoint_rank", 1)))
    if accelerator.is_main_process and not profile_only:
        with log_path.open("w", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow(
                [
                    "epoch",
                    "lr",
                    "train_loss",
                    "train_ic",
                    "train_unweighted_ic",
                    "train_weighted_ic",
                    "train_rmse",
                    "train_unweighted_rmse",
                    "val_loss",
                    "val_ic",
                    "val_unweighted_ic",
                    "val_weighted_ic",
                    "val_rmse",
                    "val_unweighted_rmse",
                    "val_mae",
                    "is_best",
                    "pooling",
                ]
            )
    phase_timings["output_init_sec"] = float(time.perf_counter() - output_init_t0)

    epochs = int(train_cfg["epochs"])
    epoch_lrs = build_epoch_lr_schedule(train_cfg, epochs=epochs)
    total_train_batches = 0
    try:
        total_train_batches = len(train_loader)
    except TypeError:
        total_train_batches = 0
    train_progress_step = (
        max(1, total_train_batches // max(1, int(args.train_progress_splits))) if total_train_batches > 0 else 0
    )
    best_epoch = 0
    best_val_m: Dict[str, float] | None = None
    final_val_m: Dict[str, float] | None = None
    top_val_checkpoints: List[dict] = []
    first_train_step_local: np.ndarray | None = None
    collective_warmup_sec = 0.0
    collective_warmup_done = False
    train_eval_timing: dict | None = None
    valid_eval_timing: dict | None = None
    test_eval_timing: dict | None = None
    if profile_only:
        train_loader.set_epoch(1)
        profile_summary = profile_train_steps(
            model=model,
            optimizer=optimizer,
            loader=train_loader,
            accelerator=accelerator,
            feat_mean=feat_mean,
            feat_std=feat_std,
            seq_offsets=seq_offsets,
            copy_non_blocking=use_pinned_transfer,
            warmup_steps=profile_warmup_steps,
            profile_steps=profile_steps,
            grad_accum_steps=grad_accum_steps,
        )
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            summary_path = out_root / "step_profile_summary.json"
            summary = {
                "run_name": args.run_name,
                "row_root": str(row_root),
                "world_size": int(accelerator.num_processes),
                "batch_size_per_rank": int(train_cfg["batch_size"]),
                "train_batch_overlap": train_batch_overlap,
                "runtime_backend": args.runtime_backend,
                "device_transfer_mode": device_transfer_mode,
                "min_timecode": int(min_timecode),
                "require_positive_weight": bool(require_positive_weight),
                "use_source_raw_weights": bool(args.use_source_raw_weights),
                "seq_len": int(args.seq_len),
                "sample_stride": int(args.sample_stride),
                "profile": profile_summary,
            }
            summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
            print(
                f"[profile-train] saved={summary_path} global_step_sec={profile_summary['global_step_sec_estimate']:.6f} "
                f"global_samples_per_sec={profile_summary['global_samples_per_sec_estimate']:.2f}",
                flush=True,
            )
        accelerator.close()
        return
    train_loop_t0 = time.perf_counter()
    for epoch in range(1, epochs + 1):
        current_lr = float(epoch_lrs[epoch - 1])
        set_optimizer_lr(optimizer, current_lr)
        epoch_train_t0 = time.perf_counter()
        model.train()
        train_loader.set_epoch(epoch)
        if accelerator.is_main_process:
            print(f"[epoch {epoch}] lr={current_lr:.8f}", flush=True)
        train_iter = iter(train_loader)
        if not collective_warmup_done:
            collective_warmup_sec = warmup_collective(accelerator)
            collective_warmup_done = True
            phase_timings["collective_warmup_sec"] = float(collective_warmup_sec)
        batch_i = 0
        accum_micro = 0
        optimizer.zero_grad(set_to_none=True)
        while True:
            fetch_t0 = time.perf_counter()
            try:
                batch = next(train_iter)
            except StopIteration:
                break
            fetch_t1 = time.perf_counter()
            batch_i += 1
            accum_micro += 1
            last_batch = total_train_batches > 0 and batch_i == total_train_batches
            should_step = accum_micro >= grad_accum_steps or last_batch
            group_total = min(grad_accum_steps, accum_micro + max(0, total_train_batches - batch_i))
            measure_first_step = bool(profile_phases and first_train_step_local is None and epoch == 1)
            xb, yb, wb = resolve_model_batch(
                batch,
                seq_offsets=seq_offsets,
                feat_mean=feat_mean,
                feat_std=feat_std,
                device=accelerator.device,
                copy_non_blocking=use_pinned_transfer,
            )
            if measure_first_step:
                sync_if_needed(accelerator.device)
                mat_t = time.perf_counter()
            sync_ctx = model.no_sync() if hasattr(model, "no_sync") and not should_step else nullcontext()
            with sync_ctx:
                with accelerator.autocast():
                    pred = model(xb)
                    loss_raw = weighted_mse(pred, yb, wb)
                    loss = loss_raw / float(group_total)
                if measure_first_step:
                    sync_if_needed(accelerator.device)
                    fwd_t = time.perf_counter()
                accelerator.backward(loss)
                if measure_first_step:
                    sync_if_needed(accelerator.device)
                    bwd_t = time.perf_counter()
            if should_step:
                accelerator.step_optimizer(optimizer)
                optimizer.zero_grad(set_to_none=True)
                accum_micro = 0
            if measure_first_step:
                sync_if_needed(accelerator.device)
                step_t = time.perf_counter()
                first_train_step_local = np.asarray(
                    [
                        fetch_t1 - fetch_t0,
                        mat_t - fetch_t1,
                        fwd_t - mat_t,
                        bwd_t - fwd_t,
                        step_t - bwd_t,
                        step_t - fetch_t0,
                    ],
                    dtype=np.float64,
                )
            if (
                accelerator.is_main_process
                and total_train_batches > 0
                and (batch_i % train_progress_step == 0 or batch_i == total_train_batches)
            ):
                print(
                    f"[epoch {epoch}] train progress {batch_i}/{total_train_batches} "
                    f"loss={loss_raw.detach().float().item():.6f}",
                    flush=True,
                )
        phase_timings[f"epoch_{epoch}_train_sec"] = float(time.perf_counter() - epoch_train_t0)

        if skip_train_eval:
            train_m = {
                "loss": float("nan"),
                "ic": float("nan"),
                "unweighted_ic": float("nan"),
                "weighted_ic": float("nan"),
                "rmse": float("nan"),
                "mae": float("nan"),
                "weight_sum": 0.0,
                "n": 0,
            }
            if accelerator.is_main_process:
                print(f"[epoch {epoch}] train-eval skipped", flush=True)
        else:
            train_eval_t0 = time.perf_counter()
            train_m = evaluate(
                model,
                train_eval_loader if train_eval_loader is not None else train_loader,
                feat_mean,
                feat_std,
                seq_offsets,
                accelerator,
                split_name=f"epoch {epoch} train-eval",
                progress_parts=args.eval_progress_splits,
                copy_non_blocking=use_pinned_transfer,
                collect_timing=profile_phases,
            )
            phase_timings[f"epoch_{epoch}_train_eval_sec"] = float(time.perf_counter() - train_eval_t0)
            train_eval_timing = train_m.pop("_timing", None)
        valid_eval_t0 = time.perf_counter()
        val_m = evaluate(
            model,
            valid_loader,
            feat_mean,
            feat_std,
            seq_offsets,
            accelerator,
            split_name=f"epoch {epoch} valid",
            progress_parts=args.eval_progress_splits,
            copy_non_blocking=use_pinned_transfer,
            collect_timing=profile_phases,
        )
        phase_timings[f"epoch_{epoch}_valid_eval_sec"] = float(time.perf_counter() - valid_eval_t0)
        valid_eval_timing = val_m.pop("_timing", None)
        final_val_m = dict(val_m)
        top_val_checkpoints, entered_topk, dropped_topk = refresh_top_val_checkpoints(
            top_val_checkpoints=top_val_checkpoints,
            candidate_epoch=epoch,
            candidate_val_metrics=val_m,
            keep_topk=save_topk_val_checkpoints,
            out_root=out_root,
        )
        is_best = 0
        current_weighted_ic = float(val_m.get("weighted_ic", val_m["ic"]))
        best_weighted_ic = float(best_val_m.get("weighted_ic", best_val_m["ic"])) if best_val_m is not None else float("-inf")
        if best_val_m is None or current_weighted_ic > best_weighted_ic:
            best_val_m = dict(val_m)
            best_epoch = epoch
            is_best = 1
            if accelerator.is_main_process:
                torch.save(accelerator.unwrap_model(model).state_dict(), best_model_path)

        if accelerator.is_main_process:
            if entered_topk:
                torch.save(
                    accelerator.unwrap_model(model).state_dict(),
                    build_val_epoch_checkpoint_path(out_root, epoch),
                )
            for dropped_item in dropped_topk:
                dropped_path = Path(str(dropped_item["path"]))
                if dropped_path.exists():
                    dropped_path.unlink()
            with log_path.open("a", newline="", encoding="utf-8") as f:
                csv.writer(f).writerow(
                    [
                        epoch,
                        current_lr,
                        train_m["loss"],
                        train_m["ic"],
                        train_m.get("unweighted_ic", float("nan")),
                        train_m.get("weighted_ic", float("nan")),
                        train_m["rmse"],
                        train_m.get("unweighted_rmse", float("nan")),
                        val_m["loss"],
                        val_m["ic"],
                        val_m.get("unweighted_ic", float("nan")),
                        val_m.get("weighted_ic", float("nan")),
                        val_m["rmse"],
                        val_m.get("unweighted_rmse", float("nan")),
                        val_m["mae"],
                        is_best,
                        str(args.pooling),
                    ]
                )
            print(
                f"[epoch {epoch}] "
                f"lr={current_lr:.8f} "
                f"train_loss={train_m['loss']:.6f} train_ic={train_m['ic']:.6f} "
                f"train_uic={train_m.get('unweighted_ic', float('nan')):.6f} train_rmse={train_m['rmse']:.6f} "
                f"val_loss={val_m['loss']:.6f} val_ic={val_m['ic']:.6f} "
                f"val_uic={val_m.get('unweighted_ic', float('nan')):.6f} val_rmse={val_m['rmse']:.6f} "
                f"best_epoch={best_epoch}"
            )
    phase_timings["fit_loop_sec"] = float(time.perf_counter() - train_loop_t0)
    if profile_phases and first_train_step_local is not None:
        phase_details["first_train_step"] = summarize_timing_across_ranks(
            accelerator=accelerator,
            labels=[
                "batch_wait_sec",
                "materialize_sec",
                "forward_sec",
                "backward_sec",
                "optimizer_step_sec",
                "step_total_sec",
            ],
            count=1,
            sums=first_train_step_local,
            count_key="measured_steps",
            global_batch=int(train_cfg["batch_size"]) * int(accelerator.num_processes),
        )
    if collective_warmup_done:
        phase_details["collective_warmup_sec"] = {"local_sec": float(collective_warmup_sec)}
    if train_eval_timing is not None:
        phase_details["train_eval_timing"] = train_eval_timing
    if valid_eval_timing is not None:
        phase_details["valid_eval_timing"] = valid_eval_timing

    def evaluate_saved_checkpoint(checkpoint_entry: dict, split_name: str) -> Tuple[Dict[str, float], float, dict | None]:
        eval_t0 = time.perf_counter()
        state_dict = torch.load(str(checkpoint_entry["path"]), map_location="cpu", weights_only=True)
        accelerator.unwrap_model(model).load_state_dict(state_dict)
        accelerator.wait_for_everyone()
        if accelerator.is_main_process:
            print(
                f"[{split_name}] evaluating val_rank={checkpoint_entry['rank']} "
                f"epoch={checkpoint_entry['epoch']} val_ic={checkpoint_entry['val_ic']:.6f} "
                f"val_wic={checkpoint_entry.get('val_weighted_ic', checkpoint_entry['val_ic']):.6f}",
                flush=True,
            )
        metrics = evaluate(
            model,
            test_loader,
            feat_mean,
            feat_std,
            seq_offsets,
            accelerator,
            split_name=split_name,
            progress_parts=args.eval_progress_splits,
            copy_non_blocking=use_pinned_transfer,
            collect_timing=profile_phases,
        )
        elapsed = float(time.perf_counter() - eval_t0)
        timing = metrics.pop("_timing", None)
        return metrics, elapsed, timing

    accelerator.wait_for_everyone()
    test_m_best: Dict[str, float] | None = None
    test_m_selected: Dict[str, float] | None = None
    selected_test_checkpoint: dict | None = None
    if test_loader is not None and not skip_test:
        if not top_val_checkpoints:
            raise RuntimeError("No validation checkpoints were recorded for test evaluation.")
        if selected_test_checkpoint_rank > len(top_val_checkpoints):
            raise RuntimeError(
                f"Requested test_checkpoint_rank={selected_test_checkpoint_rank} but only "
                f"{len(top_val_checkpoints)} validation checkpoints are available."
            )
        best_test_checkpoint = top_val_checkpoints[0]
        selected_test_checkpoint = top_val_checkpoints[selected_test_checkpoint_rank - 1]
        test_m_best, test_best_sec, test_eval_timing = evaluate_saved_checkpoint(
            best_test_checkpoint,
            split_name="test best-val",
        )
        phase_timings["test_eval_best_val_sec"] = float(test_best_sec)
        if selected_test_checkpoint_rank == 1:
            test_m_selected = dict(test_m_best)
            phase_timings["test_eval_selected_val_rank_sec"] = float(test_best_sec)
            phase_timings["test_eval_sec"] = float(test_best_sec)
        else:
            test_m_selected, test_selected_sec, test_eval_timing_selected = evaluate_saved_checkpoint(
                selected_test_checkpoint,
                split_name=f"test val-rank-{selected_test_checkpoint_rank}",
            )
            phase_timings["test_eval_selected_val_rank_sec"] = float(test_selected_sec)
            phase_timings["test_eval_sec"] = float(test_best_sec + test_selected_sec)
            if test_eval_timing_selected is not None:
                phase_details["test_eval_timing_selected_val_rank"] = test_eval_timing_selected
        if accelerator.is_main_process:
            print(
                f"[test best-val] rmse={test_m_best['rmse']:.6f} mae={test_m_best['mae']:.6f} "
                f"ic={test_m_best['ic']:.6f} uic={test_m_best.get('unweighted_ic', float('nan')):.6f} "
                f"wic={test_m_best.get('weighted_ic', float('nan')):.6f}",
                flush=True,
            )
            if selected_test_checkpoint_rank != 1 and test_m_selected is not None:
                print(
                    f"[test val-rank-{selected_test_checkpoint_rank}] "
                    f"rmse={test_m_selected['rmse']:.6f} mae={test_m_selected['mae']:.6f} "
                    f"ic={test_m_selected['ic']:.6f} uic={test_m_selected.get('unweighted_ic', float('nan')):.6f} "
                    f"wic={test_m_selected.get('weighted_ic', float('nan')):.6f}",
                    flush=True,
                )
    elif test_loader is not None and accelerator.is_main_process:
        print("[test] skipped", flush=True)
    if test_eval_timing is not None:
        phase_details["test_eval_timing"] = test_eval_timing
    final_save_t0 = time.perf_counter()
    if accelerator.is_main_process:
        torch.save(accelerator.unwrap_model(model).state_dict(), out_root / "gru_seq_memmap_ddp.pt")
    accelerator.wait_for_everyone()
    phase_timings["final_save_sec"] = float(time.perf_counter() - final_save_t0)
    phase_timings["total_run_sec"] = float(time.perf_counter() - run_t0)
    if accelerator.is_main_process:
        summary = {
            "run_name": args.run_name,
            "cache_name": args.cache_name,
            "max_days": args.max_days,
            "train_pool_start_date": args.train_pool_start_date,
            "train_pool_end_date": args.train_pool_end_date,
            "test_start_date": args.test_start_date,
            "test_end_date": args.test_end_date,
            "valid_split_ratio": split_ratio,
            "selection_metric": "val_weighted_ic",
            "train_day_count": len(train_days),
            "valid_day_count": len(valid_days),
            "test_day_count": len(test_days),
            "train_samples": int(train_loader.total_samples),
            "valid_samples": int(valid_loader.total_samples),
            "test_samples": int(test_loader.total_samples) if test_loader is not None else 0,
            "seed": int(args.seed),
            "seq_len": args.seq_len,
            "sample_stride": args.sample_stride,
            "min_timecode": int(min_timecode),
            "require_positive_weight": bool(require_positive_weight),
            "use_source_raw_weights": bool(args.use_source_raw_weights),
            "hidden_dim": args.hidden_dim,
            "num_layers": args.num_layers,
            "bidirectional": bool(args.bidirectional),
            "use_cnn1d": bool(getattr(args, "use_cnn1d", 0)),
            "input_gate_hidden_dim": int(getattr(args, "input_gate_hidden_dim", 0)),
            "input_gate_bias": float(getattr(args, "input_gate_bias", 2.0)),
            "weight_decay": args.weight_decay,
            "prefer_fp16": bool(args.prefer_fp16),
            "use_amp": bool(args.use_amp),
            "pooling": str(args.pooling),
            "loader_backend": "threaded_seq_batch_loader",
            "loader_threads": loader_threads,
            "loader_prefetch_batches": loader_prefetch,
            "train_batch_overlap": train_batch_overlap,
            "train_batch_stride": int(train_cfg["batch_size"]) - train_batch_overlap,
            "eval_batch_size_per_rank": eval_batch_size,
            "grad_accum_steps": grad_accum_steps,
            "effective_global_batch": int(train_cfg["batch_size"]) * int(accelerator.num_processes) * grad_accum_steps,
            "runtime_backend": args.runtime_backend,
            "device_transfer_mode": device_transfer_mode,
            "device_transfer_prefetch_batches": device_transfer_prefetch_batches,
            "pinned_nonblocking_transfer": use_pinned_transfer,
            "skip_train_eval": skip_train_eval,
            "skip_test": skip_test,
            "train_cfg": train_cfg,
            "epoch_learning_rates": [float(x) for x in epoch_lrs],
            "loss": "weighted_mse",
            "metric_main": "val_weighted_ic",
            "test_selection_rule": f"val_weighted_ic_rank_{selected_test_checkpoint_rank}",
            "save_topk_val_checkpoints": int(save_topk_val_checkpoints),
            "best_epoch_by_val_ic": best_epoch,
            "best_epoch_by_val_weighted_ic": best_epoch,
            "best_val_metrics": best_val_m,
            "final_epoch_val_metrics": final_val_m,
            "top_val_checkpoints": top_val_checkpoints,
            "selected_test_checkpoint_rank": int(selected_test_checkpoint_rank),
            "selected_test_checkpoint_epoch": int(selected_test_checkpoint["epoch"]) if selected_test_checkpoint is not None else None,
            "selected_test_model_path": str(selected_test_checkpoint["path"]) if selected_test_checkpoint is not None else None,
            "test_metrics_at_best_val": test_m_best,
            "test_metrics_at_selected_val_rank": test_m_selected,
            "best_model_path": str(best_model_path),
            "phase_timings": phase_timings,
            "phase_details": phase_details,
        }
        (out_root / "training_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
        np.savez(out_root / "feature_stats.npz", mean=feat_mean_np, std=feat_std_np)
        if profile_phases:
            phase_summary = {
                "run_name": args.run_name,
                "row_root": str(row_root),
                "world_size": int(accelerator.num_processes),
                "batch_size_per_rank": int(train_cfg["batch_size"]),
                "train_batch_overlap": train_batch_overlap,
                "runtime_backend": args.runtime_backend,
                "device_transfer_mode": device_transfer_mode,
                "min_timecode": int(min_timecode),
                "require_positive_weight": bool(require_positive_weight),
                "use_source_raw_weights": bool(args.use_source_raw_weights),
                "phase_timings": phase_timings,
                "phase_details": phase_details,
            }
            phase_path = out_root / "phase_profile_summary.json"
            phase_path.write_text(json.dumps(phase_summary, ensure_ascii=False, indent=2), encoding="utf-8")
            print(f"[phase-profile] saved={phase_path}", flush=True)
    accelerator.close()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="/intern9/huhongkai/hs300_factor_lab/configs/experiment_2024_2025_memmap.json")
    p.add_argument("--run-name", default="seq_gru_memmap_run")
    p.add_argument("--cache-name", default="top200_eps20_rows_2024_2025")
    p.add_argument(
        "--row-root-override",
        type=str,
        default="",
        help="Optional row_memmap root or exact cache dir. Useful for local staged cache under /dev/shm.",
    )
    p.add_argument("--max-days", type=int, default=-1)
    p.add_argument("--max-samples", type=int, default=-1)
    p.add_argument("--train-day-limit", type=int, default=-1)
    p.add_argument("--valid-day-limit", type=int, default=-1)
    p.add_argument("--test-day-limit", type=int, default=-1)
    p.add_argument("--train-pool-start-date", type=str, default="")
    p.add_argument("--train-pool-end-date", type=str, default="")
    p.add_argument("--test-start-date", type=str, default="")
    p.add_argument("--test-end-date", type=str, default="")
    p.add_argument("--valid-split-ratio", type=float, default=-1.0)
    p.add_argument("--seed", type=int, default=20260312)
    p.add_argument("--seq-len", type=int, default=60)
    p.add_argument("--sample-stride", type=int, default=10)
    p.add_argument("--train-batch-overlap", type=int, default=0)
    p.add_argument("--eval-batch-size", type=int, default=0)
    p.add_argument("--grad-accum-steps", type=int, default=1)
    p.add_argument("--runtime-backend", type=str, default="native", choices=["accelerate", "native"])
    p.add_argument(
        "--device-transfer-mode",
        type=str,
        default="thread_prefetch",
        choices=["direct", "thread_prefetch", "materialize_prefetch"],
    )
    p.add_argument("--device-transfer-prefetch-batches", type=int, default=2)
    p.add_argument("--loader-prefetch-batches", type=int, default=0)
    p.add_argument("--prefer-fp16", type=int, default=1)
    p.add_argument("--override-epochs", type=int, default=-1)
    p.add_argument("--override-lr", type=float, default=-1.0)
    p.add_argument("--override-batch-size", type=int, default=-1)
    p.add_argument("--hidden-dim", type=int, default=256)
    p.add_argument("--num-layers", type=int, default=2)
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--pooling", type=str, default="last", choices=["last", "attn"])
    p.add_argument("--bidirectional", type=int, default=0, help="1 to enable bidirectional GRU")
    p.add_argument("--use-cnn1d", type=int, default=0, help="1 to use parallel 1D-CNN before GRU")
    p.add_argument("--input-gate-hidden-dim", type=int, default=0, help=">0 to enable feature/channel gate before GRU")
    p.add_argument("--input-gate-bias", type=float, default=2.0, help="Initial bias for feature/channel gate")
    p.add_argument("--weight-decay", type=float, default=1e-5)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--use-amp", type=int, default=1, help="1 to enable fp16 mixed precision")
    p.add_argument("--train-progress-splits", type=int, default=10)
    p.add_argument("--eval-progress-splits", type=int, default=4)
    p.add_argument("--skip-train-eval", type=int, default=1, help="1 to skip full train-set evaluation")
    p.add_argument("--skip-test", type=int, default=0, help="1 to skip final test evaluation")
    p.add_argument(
        "--use-source-raw-weights",
        type=int,
        default=1,
        help="1 to load raw continuous weights from the source split npy instead of cache-side binary masks.",
    )
    p.add_argument(
        "--min-timecode",
        type=int,
        default=DEFAULT_MIN_TIMECODE,
        help="Only keep sequence end points whose source datetime >= this HHMMSSmmm timecode.",
    )
    p.add_argument(
        "--require-positive-weight",
        type=int,
        default=1,
        help="1 to drop sequence end points whose cached/source weight is not positive.",
    )
    p.add_argument("--profile-phases", type=int, default=0)
    p.add_argument("--profile-warmup-steps", type=int, default=0)
    p.add_argument("--profile-steps", type=int, default=0)
    p.add_argument("--stop-after-profile", type=int, default=0)
    p.add_argument("--save-topk-val-checkpoints", type=int, default=1)
    p.add_argument("--test-checkpoint-rank", type=int, default=1)
    args = p.parse_args()
    run(args)


if __name__ == "__main__":
    main()
B^The file /intern9/huhongkai/hs300_factor_lab/src/train_seq_gru_ddp_memmap.py has been updated.