train.py

import argparse
import os
import multiprocessing as py_mp  # 为 DataLoader 指定上下文
from pathlib import Path
from typing import Optional, List, Tuple, Dict, TypedDict, Set
import math
import pandas as pd  # for noise scaling by YZ vol lookup
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, Subset
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
try:
    from torch.distributed.fsdp import (
        FullyShardedDataParallel as FSDP,
        StateDictType,
        FullStateDictConfig,
    )
except Exception:  # PyTorch 版本或环境不支持 FSDP 时安全降级
    FSDP = None  # type: ignore
    StateDictType = None  # type: ignore
    FullStateDictConfig = None  # type: ignore
import numpy as np
from collections import deque  # unused; kept for compatibility
try:
    from tqdm.auto import tqdm
except Exception:
    tqdm = None
import time
import json

from model import SpacetimeGNNMAM


# 避免 /dev/shm 过小引发的 worker Bus Error：改用基于文件系统的共享策略
try:
    import torch.multiprocessing as mp
    mp.set_sharing_strategy('file_system')
except Exception:
    pass

def is_dist_env() -> bool:
    try:
        return int(os.environ.get("WORLD_SIZE", "1")) > 1
    except Exception:
        return False


def setup_distributed() -> Tuple[bool, int, int, int]:
    """Initialize torch.distributed if WORLD_SIZE>1 via env://.
    Backend can be overridden by env var DIST_BACKEND (nccl/gloo), default nccl.
    Returns: (distributed, world_size, rank, local_rank)
    """
    distributed = is_dist_env()
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    backend = os.environ.get("DIST_BACKEND", "nccl").lower()
    if distributed:
        if not dist.is_initialized():
            dist.init_process_group(backend=backend, init_method="env://")
        rank = dist.get_rank()
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        if torch.cuda.is_available():
            torch.cuda.set_device(local_rank)
    else:
        rank = 0
        local_rank = 0
    return distributed, world_size, rank, local_rank


def is_main_process(rank: int) -> bool:
    return (rank == 0)


def unwrap_model(m: nn.Module) -> nn.Module:
    return m.module if isinstance(m, DDP) else m


def ensure_dir(path: str) -> None:
    Path(path).mkdir(parents=True, exist_ok=True)


def compute_grad_norm_l2(params: List[nn.Parameter]) -> float:
    total = 0.0
    for param in params:
        if param.grad is None:
            continue
        grad = param.grad.detach()
        total += float(grad.float().pow(2).sum().item())
    if total <= 0.0:
        return 0.0
    return math.sqrt(total)


def collect_param_statistics(model: nn.Module) -> List[dict]:
    stats: List[dict] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        data = param.detach()
        if data.numel() == 0:
            continue
        data_f = data.to(dtype=torch.float32)
        stat = {
            "name": name,
            "shape": list(data.shape),
            "numel": int(data.numel()),
            "mean": float(data_f.mean().item()),
            "std": float(data_f.std(unbiased=False).item()),
            "min": float(data_f.min().item()),
            "max": float(data_f.max().item()),
            "l2": float(data_f.norm().item()),
        }
        if param.grad is not None:
            grad = param.grad.detach().to(dtype=torch.float32)
            stat.update({
                "grad_mean": float(grad.mean().item()),
                "grad_std": float(grad.std(unbiased=False).item()),
                "grad_norm": float(grad.norm().item()),
            })
        stats.append(stat)
    return stats


class GradEntry(TypedDict):
    name: str
    grad_norm: float


def collect_layer_grad_snapshot(model: nn.Module) -> Tuple[List[dict], float, GradEntry | None]:
    """
    收集模型当前各层的参数与梯度统计快照。

    Returns:
        records: 每层的统计信息列表
        total_grad_norm: 梯度范数的整体二范数（sqrt(sum(grad_norm^2)))
        max_grad_entry: 梯度最大的层信息 {"name": str, "grad_norm": float} 或 None
    """
    records: List[dict] = []
    total_grad_sq = 0.0
    max_grad_entry: GradEntry | None = None

    with torch.no_grad():
        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue

            param_norm = float(param.norm().item())
            param_mean = float(param.mean().item())
            param_std = float(param.std().item()) if param.numel() > 1 else 0.0
            param_max = float(param.max().item())
            param_min = float(param.min().item())

            if param.grad is not None:
                grad_tensor = param.grad
                grad_norm = float(grad_tensor.norm().item())
                grad_mean = float(grad_tensor.mean().item())
                grad_std = float(grad_tensor.std().item()) if grad_tensor.numel() > 1 else 0.0
                grad_max = float(grad_tensor.max().item())
                grad_min = float(grad_tensor.min().item())
                update_ratio = float(grad_norm / (param_norm + 1e-8))
            else:
                grad_norm = grad_mean = grad_std = grad_max = grad_min = update_ratio = 0.0

            record = {
                "name": name,
                "shape": list(param.shape),
                "param_norm": param_norm,
                "param_mean": param_mean,
                "param_std": param_std,
                "param_range": [param_min, param_max],
                "grad_norm": grad_norm,
                "grad_mean": grad_mean,
                "grad_std": grad_std,
                "grad_range": [grad_min, grad_max],
                "update_ratio": update_ratio,
            }
            records.append(record)

            if grad_norm > 0.0:
                total_grad_sq += grad_norm ** 2
                if max_grad_entry is None or grad_norm > max_grad_entry["grad_norm"]:
                    max_grad_entry = {"name": name, "grad_norm": grad_norm}

    total_grad_norm = math.sqrt(total_grad_sq) if total_grad_sq > 0.0 else 0.0
    return records, total_grad_norm, max_grad_entry


class AdaptiveSpectralNoiseTracker:
    """
    预测任务专用：通过破坏低频特征迫使模型依赖高频内生特征
    核心逻辑：低频特征（GDP、利率）作为条件输入，训练时需自适应增强噪声降低模型对其依赖
    """
    def __init__(self, seq_len: int, feature_dim: int, freq_div: int = 3):
        """
        freq_div控制低频截止点：cutoff = seq_len // freq_div
        - freq_div=3: cutoff=20 (seq_len=60时，前20个频率为低频，较保守)
        - freq_div=10: cutoff=6 (seq_len=60时，前6个频率为低频，太激进)
        """
        self.seq_len = int(max(1, seq_len))
        self.feature_dim = int(max(1, feature_dim))
        freq_bins = self.seq_len // 2 + 1
        freq_div = max(1, int(freq_div))
        self.freq_cutoff = max(1, min(freq_bins, self.seq_len // freq_div if freq_div > 0 else freq_bins))
        self.device = torch.device("cpu")
        self.sensitivity_sum = torch.zeros(feature_dim, dtype=torch.float64, device=self.device)
        self.sensitivity_count = 0
        self.reset()

    def reset(self) -> None:
        F = self.feature_dim
        self.orig_sum = torch.zeros(F, dtype=torch.float64, device=self.device)
        self.orig_sq_sum = torch.zeros(F, dtype=torch.float64, device=self.device)
        self.orig_count = 0

        self.diff_sum = torch.zeros(F, dtype=torch.float64, device=self.device)
        self.diff_sq_sum = torch.zeros(F, dtype=torch.float64, device=self.device)
        self.diff_count = 0

        self.low_power_sum = torch.zeros(F, dtype=torch.float64, device=self.device)
        self.high_power_sum = torch.zeros(F, dtype=torch.float64, device=self.device)

        self.sigma_sum = 0.0
        self.sigma_sq_sum = 0.0
        self.sigma_count = 0

        self.last_stats: Optional[dict] = None
        self.last_config: Optional[dict] = None
        self.last_alpha: Optional[torch.Tensor] = None

    def _update_statistics(self, x: torch.Tensor) -> None:
        # x: [B, N, T, F]
        x_detached = x.detach().to(dtype=torch.float32)
        B, N, T, F = x_detached.shape
        x_cpu = x_detached.view(-1, F).to(device=self.device, dtype=torch.float32)
        orig_sum_batch = x_cpu.sum(dim=0, dtype=torch.float64)
        orig_sq_sum_batch = (x_cpu * x_cpu).sum(dim=0, dtype=torch.float64)
        self.orig_sum += orig_sum_batch
        self.orig_sq_sum += orig_sq_sum_batch
        self.orig_count += x_cpu.shape[0]

        if T > 1:
            diff = torch.diff(x_detached, dim=2)  # [B,N,T-1,F]
            diff_cpu = diff.reshape(-1, F).to(device=self.device, dtype=torch.float32)
            self.diff_sum += diff_cpu.sum(dim=0, dtype=torch.float64)
            self.diff_sq_sum += (diff_cpu * diff_cpu).sum(dim=0, dtype=torch.float64)
            self.diff_count += diff_cpu.shape[0]

        try:
            fft_vals = torch.fft.rfft(x_detached, dim=2)
            power = (fft_vals.abs() ** 2).to(dtype=torch.float32)
            # power: [B,N,freq_bins,F]
            freq_bins = power.shape[2]
            cutoff = min(self.freq_cutoff, freq_bins)
            low_slice = power[:, :, :cutoff, :].sum(dim=(0, 1, 2), dtype=torch.float64)
            if cutoff < freq_bins:
                high_slice = power[:, :, cutoff:, :].sum(dim=(0, 1, 2), dtype=torch.float64)
            else:
                high_slice = torch.zeros_like(low_slice, dtype=torch.float64)
            self.low_power_sum += low_slice.to(self.device)
            self.high_power_sum += high_slice.to(self.device)
        except Exception:
            # FFT 失败时跳过频域累积
            pass

    def _compute_stats(self, robustness_coeff: float = 9.0, epsilon: float = 0.1) -> Optional[dict]:
        """
        预测任务核心：自适应计算每个特征的噪声增强系数alpha
        公式：α_i = 1.0 + C / (k_i + ε)  【完全自适应，无硬截断】
        
        - C: 全局鲁棒性系数（robustness_coeff），默认9.0
        - k_i: 频域能量比（高频/低频），k越小表示越低频
        - ε: 防止除零的小常数，默认0.1
        
        核心逻辑：
        - k_i → 0（极低频，如GDP）：α_i → 1.0 + C/ε ≈ 91x 破坏
        - k_i = 0.5：α_i ≈ 1.0 + 9.0/0.6 ≈ 16x
        - k_i = 5.0：α_i ≈ 1.0 + 9.0/5.1 ≈ 2.8x
        - k_i = 50：α_i ≈ 1.0 + 9.0/50.1 ≈ 1.2x
        - k_i → ∞（极高频）：α_i → 1.0（接近原噪声）
        """
        if self.orig_count == 0:
            return None
        eps = 1e-8
        orig_mean = self.orig_sum / max(1, self.orig_count)
        orig_var = (self.orig_sq_sum / max(1, self.orig_count)) - orig_mean * orig_mean
        orig_var = torch.clamp(orig_var, min=eps)

        if self.diff_count > 0:
            diff_mean = self.diff_sum / max(1, self.diff_count)
            diff_var = (self.diff_sq_sum / max(1, self.diff_count)) - diff_mean * diff_mean
            diff_var = torch.clamp(diff_var, min=eps)
        else:
            diff_var = torch.zeros_like(orig_var)

        diff_ratio = diff_var / torch.clamp(orig_var, min=eps)

        # 计算频域能量比 k = V_high / V_low
        denom = torch.clamp(self.low_power_sum, min=eps)
        k_per_feature = self.high_power_sum / denom
        k_per_feature = torch.clamp(k_per_feature, min=eps)

        # 全局自适应公式：α_i = 1.0 + C / (k_i + ε)
        # 所有特征都参与，无硬截断，完全由k值决定破坏强度
        C = float(max(0.0, robustness_coeff))
        epsilon_safe = float(max(0.01, epsilon))
        
        alpha = 1.0 + C / (k_per_feature + epsilon_safe)

        signal_std = torch.sqrt(orig_var)

        stats = {
            "orig_var": orig_var,
            "diff_var": diff_var,
            "diff_ratio": diff_ratio,
            "k_per_feature": k_per_feature,
            "signal_std": signal_std,
            "alpha": alpha,
            "robustness_coeff": C,
            "epsilon": epsilon_safe,
        }
        return stats

    def apply(self, x: torch.Tensor, sigma_values: torch.Tensor) -> torch.Tensor:
        """
        核心加噪逻辑：xₜ = x₀ + σ * (noise * alpha)
        - alpha[i] 由 _compute_stats 自适应计算
        - 低频特征（k_i小）：alpha大（强破坏）
        - 高频特征（k_i大）：alpha=1.0（保持原样）
        """
        stats = self.last_stats
        if stats is None:
            # 首次调用或未预计算，直接返回原始噪声
            noise = torch.randn_like(x)
            sigma_tensor = sigma_values.detach()
            if sigma_tensor.dim() == 0:
                sigma_tensor = sigma_tensor.view(1)
            sigma_tensor = sigma_tensor.to(device=x.device, dtype=x.dtype)
            sigma_expand = sigma_tensor.view(-1, 1, 1, 1)
            return x + sigma_expand * noise

        # 记录sigma统计
        sigma_tensor = sigma_values.detach()
        if sigma_tensor.dim() == 0:
            sigma_tensor = sigma_tensor.view(1)
        sigma_tensor = sigma_tensor.to(device=x.device, dtype=x.dtype)
        self.sigma_sum += float(sigma_tensor.sum().item())
        self.sigma_sq_sum += float((sigma_tensor * sigma_tensor).sum().item())
        self.sigma_count += int(sigma_tensor.numel())

        # 获取alpha向量 [F]
        alpha = stats["alpha"].to(device=x.device, dtype=x.dtype)  # [F]
        self.last_alpha = alpha.cpu()

        # 生成噪声并应用特征级别的alpha缩放
        noise = torch.randn_like(x)  # [B, N, T, F]
        # alpha形状为[F]，广播到[1,1,1,F]后自动匹配[B,N,T,F]
        noise_scaled = noise * alpha.view(1, 1, 1, -1)
        
        # 最终加噪：xₜ = x₀ + σ * (noise * alpha)
        sigma_expand = sigma_tensor.view(-1, 1, 1, 1)  # [B,1,1,1]
        return x + sigma_expand * noise_scaled
    
    def update_sensitivity(self, x_input: torch.Tensor, model_output: torch.Tensor) -> None:
        """
        计算模型输出对输入特征的敏感度（梯度范数）
        - x_input: [B, N, T, F] 输入特征（需要requires_grad=True）
        - model_output: [B, N, out_dim] 模型输出
        理想情况：低频特征的敏感度应低于0.01（模型不依赖它们）
        """
        if not x_input.requires_grad or model_output is None:
            return
        try:
            # 计算输出对输入的梯度：∂output/∂input
            # 使用output的L2范数作为标量loss进行反向传播
            output_norm = torch.sum(model_output ** 2)
            grads = torch.autograd.grad(
                outputs=output_norm,
                inputs=x_input,
                create_graph=False,
                retain_graph=True,
                only_inputs=True,
            )
            if grads and len(grads) > 0:
                grad_x = grads[0]  # [B, N, T, F]
                # 计算每个特征的梯度范数（跨batch, assets, time求平均）
                grad_norm_per_feature = torch.sqrt(torch.mean(grad_x ** 2, dim=(0, 1, 2)))  # [F]
                self.sensitivity_sum += grad_norm_per_feature.detach().to(self.device, dtype=torch.float64)
                self.sensitivity_count += 1
        except Exception:
            # 梯度计算失败时跳过
            pass
    
    def get_sensitivity_stats(self) -> Optional[dict]:
        """
        返回各特征的敏感度统计（全局自适应，无硬截断）
        """
        if self.sensitivity_count == 0 or self.last_stats is None:
            return None
        avg_sensitivity = self.sensitivity_sum / max(1, self.sensitivity_count)
        alpha_vec = self.last_stats["alpha"].cpu()
        k_feat = self.last_stats["k_per_feature"].cpu()
        
        sens_cpu = avg_sensitivity.cpu().float()
        
        stats: Dict[str, object] = {
            "sensitivity_mean": float(sens_cpu.mean().item()),
            "sensitivity_median": float(torch.median(sens_cpu).item()),
            "sensitivity_max": float(sens_cpu.max().item()),
            "sensitivity_min": float(sens_cpu.min().item()),
        }
        
        # 展示敏感度最高的特征（可能需要增大robustness_coeff）
        top_k = min(5, sens_cpu.numel())
        top_sens_idx = torch.topk(sens_cpu, k=top_k).indices
        top_sensitive_features = []
        for idx in top_sens_idx:
            feat_idx = int(idx.item())
            top_sensitive_features.append({
                "feature": feat_idx,
                "sensitivity": float(sens_cpu[feat_idx].item()),
                "k": float(k_feat[feat_idx].item()),
                "alpha": float(alpha_vec[feat_idx].item()),
            })
        stats["top_sensitive_features"] = top_sensitive_features
        
        return stats

    def get_summary(self, example_count: int = 5, full_dump: bool = False) -> Optional[dict]:
        if self.last_stats is None:
            return None
        stats = self.last_stats
        diff_ratio = stats["diff_ratio"].cpu()
        k_feat = stats["k_per_feature"].cpu()
        alpha_vec = stats["alpha"].cpu()
        signal_std = stats["signal_std"].cpu()

        def summarize_tensor(t: torch.Tensor) -> dict:
            q25 = float(torch.quantile(t, 0.25).item()) if t.numel() > 0 else 0.0
            q50 = float(torch.quantile(t, 0.50).item()) if t.numel() > 0 else 0.0
            q75 = float(torch.quantile(t, 0.75).item()) if t.numel() > 0 else 0.0
            return {
                "mean": float(t.mean().item()) if t.numel() > 0 else 0.0,
                "min": float(t.min().item()) if t.numel() > 0 else 0.0,
                "max": float(t.max().item()) if t.numel() > 0 else 0.0,
                "p25": q25,
                "median": q50,
                "p75": q75,
            }

        summary = {
            "diff_ratio": summarize_tensor(diff_ratio),
            "k_per_feature": summarize_tensor(k_feat),
            "alpha": summarize_tensor(alpha_vec),
            "signal_std": summarize_tensor(signal_std),
            "sigma_mean": float(self.sigma_sum / max(1, self.sigma_count)),
            "sigma_std": float(math.sqrt(max(0.0, self.sigma_sq_sum / max(1, self.sigma_count) - (self.sigma_sum / max(1, self.sigma_count)) ** 2))),
        }

        example_count = max(1, int(example_count))
        # 展示alpha最大和最小的特征
        top_alpha_idx = torch.topk(alpha_vec, k=min(example_count, alpha_vec.numel())).indices
        bottom_alpha_idx = torch.topk(alpha_vec, k=min(example_count, alpha_vec.numel()), largest=False).indices

        def build_example_list(indices: torch.Tensor) -> List[dict]:
            result = []
            for idx in indices:
                feat_idx = int(idx.item())
                result.append({
                    "feature": feat_idx,
                    "k": float(k_feat[feat_idx].item()),
                    "alpha": float(alpha_vec[feat_idx].item()),
                    "signal_std": float(signal_std[feat_idx].item()),
                })
            return result

        summary["top_destroyed_features"] = build_example_list(top_alpha_idx)
        summary["least_destroyed_features"] = build_example_list(bottom_alpha_idx)

        if full_dump:
            summary["diff_ratio_vector"] = [round(float(v), 6) for v in diff_ratio.tolist()]
            summary["k_vector"] = [round(float(v), 6) for v in k_feat.tolist()]
            summary["alpha_vector"] = [round(float(v), 6) for v in alpha_vec.tolist()]
        return summary

class PanelDataset(Dataset):
    def __init__(self, tensor_path: str):
        pack = torch.load(tensor_path, weights_only=False)
        self.X = pack["X"]  # [S, N, T, F]
        self.Y = pack["Y"]  # [S, N, 1]
        self.seq_len = pack["seq_len"]
        self.input_dim = pack["input_dim"]
        self.stocks = pack["stocks"]

    def __len__(self) -> int:
        return self.X.shape[0]

    def __getitem__(self, idx: int):
        # 返回时间位置 idx，便于频域损失基于绝对时间索引累计
        return self.X[idx], self.Y[idx], idx


class StreamingPanelDataset(Dataset):
    def __init__(self, panels_path: str, split: str, exclude_stocks: Optional[List[str]] = None, trim_tail_days: int = 0):
        pack = torch.load(panels_path,weights_only=True)
        # 保持为 Torch Tensor，避免反复 numpy<->torch 转换带来的 CPU 压力
        self.feature_panel = pack["feature_panel"].float().contiguous()  # [D,N,F]
        self.mask_panel = pack.get("mask_panel", torch.ones_like(pack["feature_panel"])).float().contiguous()  # [D,N,F] [mask机制]
        self.target_panel = pack["target_panel"].float().contiguous()    # [D,N]
        self.seq_len = int(pack["seq_len"])                          # T
        self.input_dim = int(pack["input_dim"])                      # F
        self.stocks = pack["stocks"]
        # 可选：按股票排除（用于训练期OOS实体留出）
        if exclude_stocks is not None and len(exclude_stocks) > 0:
            mask = [s not in set(exclude_stocks) for s in self.stocks]
            if any(mask):
                mask_arr = np.array(mask, dtype=bool)
                self.feature_panel = self.feature_panel[:, mask_arr, :]
                self.mask_panel = self.mask_panel[:, mask_arr, :]  # [mask机制] 同步过滤
                self.target_panel = self.target_panel[:, mask_arr]
                self.stocks = [s for s, keep in zip(self.stocks, mask) if keep]
        self.dates = pack["dates"]
        self.train_end = np.datetime64(pack["train_end"]) if "train_end" in pack else None
        self.val_end = np.datetime64(pack["val_end"]) if "val_end" in pack else None
        self.split = split
        self.trim_tail_days = int(max(0, trim_tail_days))
        # 预计算可用时间索引
        self.time_indices = self._build_time_indices()

    def _build_time_indices(self) -> list:
        D = self.feature_panel.shape[0]
        t0 = self.seq_len - 1
        # map dates to np.datetime64
        date_arr = np.array(self.dates, dtype='datetime64[D]')
        if self.train_end is None:
            # fallback: all train
            mask = np.ones(D, dtype=bool)
        else:
            if self.split == "train":
                mask = date_arr <= self.train_end
            elif self.split == "val":
                mask = (date_arr > self.train_end) & (date_arr <= self.val_end)
            else:
                mask = date_arr > self.val_end
        idx = np.where(mask)[0]
        idx = idx[idx >= t0]
        # 裁掉末尾填充日（仅作用于验证/测试 split）
        try:
            if self.trim_tail_days > 0 and self.split in ("val", "test"):
                max_keep = int(max(-1, D - 1 - self.trim_tail_days))
                idx = idx[idx <= max_keep]
        except Exception:
            pass
        # 过滤 target_panel 在该时间步含 NaN 的样本，并打印一次诊断日志
        try:
            # self.target_panel: [D, N]
            finite_t = torch.isfinite(self.target_panel).all(dim=1).cpu().numpy()  # [D]
            keep = finite_t[idx]
            if not np.all(keep):
                # 训练/验证已稳定，无需持续输出
                pass
            idx = idx[keep]
        except Exception:
            pass
        return idx.tolist()

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

    def __getitem__(self, i: int):
        t = self.time_indices[i]
        # window [T,N,F] → [N,T,F]（保持 Torch Tensor）
        win = self.feature_panel[t - self.seq_len + 1:t + 1]  # [T,N,F]
        X = win.permute(1, 0, 2).contiguous()  # [N,T,F]
        # [mask机制] 同步提取mask窗口
        mask_win = self.mask_panel[t - self.seq_len + 1:t + 1]  # [T,N,F]
        mask = mask_win.permute(1, 0, 2).contiguous()  # [N,T,F]
        y = self.target_panel[t]  # [N]
        Y = y.view(-1, 1).contiguous()  # [N,1]
        # 返回绝对时间索引 t，便于频域损失构造 DFT 累计
        return X, Y, int(t), mask  # [mask机制] 返回mask


def _normalize_date(value) -> str:
    ts = pd.to_datetime(value)
    return ts.strftime('%Y-%m-%d')


def _get_dataset_dates(dataset, context: str) -> List[str]:
    if not hasattr(dataset, 'dates'):
        raise AttributeError(f"[{context}] dataset 缺少 dates 属性，无法定位 yz_vol")
    dates = getattr(dataset, 'dates')
    if dates is None or len(dates) == 0:
        raise ValueError(f"[{context}] dataset.dates 为空，无法定位 yz_vol")
    return [ _normalize_date(d) for d in dates ]


def _get_dataset_stocks(dataset, context: str) -> List[str]:
    if not hasattr(dataset, 'stocks'):
        raise AttributeError(f"[{context}] dataset 缺少 stocks 属性，无法对齐 yz_vol")
    stocks = list(getattr(dataset, 'stocks'))
    if len(stocks) == 0:
        raise ValueError(f"[{context}] dataset.stocks 为空，无法对齐 yz_vol")
    return [str(s) for s in stocks]


def _load_yz_table(args: argparse.Namespace) -> pd.DataFrame:
    global YZ_TABLE
    if 'YZ_TABLE' in globals() and isinstance(globals()['YZ_TABLE'], pd.DataFrame):
        return globals()['YZ_TABLE']
    csv_path = str(getattr(args, 'noise_scale_csv', 'data/processed/yz_vol.csv'))
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"未找到 yz_vol CSV: {csv_path}")
    df = pd.read_csv(csv_path)
    if 'date' not in df.columns:
        raise KeyError(f"{csv_path} 中缺少 date 列")
    default_col = 'yz_vol_scale'
    col_name = str(getattr(args, 'noise_scale_col', default_col))
    if col_name not in df.columns:
        raise KeyError(f"{csv_path} 中缺少 {col_name} 列")
    df['date'] = pd.to_datetime(df['date']).dt.strftime('%Y-%m-%d')
    pivot = df.pivot(index='date', columns='order_book_id', values=col_name)
    globals()['YZ_TABLE'] = pivot
    return pivot


def _build_vol_tensor(
    table: pd.DataFrame,
    dates: List[str],
    stocks: List[str],
    indices: torch.Tensor,
    clip: float,
    device: torch.device,
    dtype: torch.dtype,
    context: str,
) -> torch.Tensor:
    idx_cpu = indices.detach().cpu().tolist()
    vol_np = np.ones((len(idx_cpu), len(stocks)), dtype=np.float32)
    for row_idx, idx in enumerate(idx_cpu):
        if idx < 0 or idx >= len(dates):
            raise IndexError(f"[{context}] 索引 {idx} 超出日期范围 0~{len(dates)-1}")
        date_key = dates[idx]
        if date_key not in table.index:
            raise KeyError(f"[{context}] yz_vol 表中缺少日期 {date_key}")
        row = table.loc[date_key]
        values = []
        for stock in stocks:
            if stock not in row.index:
                raise KeyError(f"[{context}] yz_vol 表中日期 {date_key} 缺少股票 {stock}")
            values.append(row[stock])
        arr = np.asarray(values, dtype=np.float64)
        if not np.all(np.isfinite(arr)):
            raise ValueError(f"[{context}] 日期 {date_key} 存在非有限 yz_vol 值: {arr}")
        arr = np.clip(arr, 1e-6, clip)
        vol_np[row_idx, :] = arr.astype(np.float32)
    tensor = torch.from_numpy(vol_np).to(device=device)
    return tensor.to(dtype=dtype)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="训练 GNN-MAM 模型（支持 DDP / ReduceLROnPlateau / 时间滚动交叉验证）")
    parser.add_argument("--config", type=str, default="", help="可选：配置JSON路径；若其中 train.use_tuned_best=false，则按配置覆盖超参")
    parser.add_argument("--data_dir", type=str, default="data/processed", help="预处理数据目录")
    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="模型保存目录")
    parser.add_argument("--resume", type=int, default=0, help="1=从最近的检查点继续训练；0=从头训练")
    parser.add_argument("--resume_path", type=str, default="checkpoints/last_checkpoint.pth", help="resume 时加载的检查点路径（默认 last_checkpoint.pth）")
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--weight_decay", type=float, default=1e-4)
    parser.add_argument("--d_model", type=int, default=128)
    parser.add_argument("--num_heads", type=int, default=8)
    parser.add_argument("--num_layers", type=int, default=4)
    parser.add_argument("--d_ff", type=int, default=512)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--temp", type=float, default=1.0)
    # 移除 top-k/退火相关参数
    # 邻接激活
    parser.add_argument("--adj_activation", type=str, default="entmax", choices=["softmax", "entmax"], help="图邻接归一化：softmax 或 entmax")
    parser.add_argument("--objective", type=str, default="edm", choices=["edm"], help="主训练目标：edm")
    parser.add_argument("--entmax_alpha", type=float, default=1.5, help="entmax 稀疏度超参 α，典型 1.2~1.8，默认1.5")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
    # 分布式参数（遵循 torchrun env://，无需手动传值）
    parser.add_argument("--dist", type=int, default=1, help="若 WORLD_SIZE>1 且 dist=1，则启用分布式DDP")
    parser.add_argument("--dist_mode", type=str, default="ddp", choices=["ddp", "fsdp", "none"], help="分布式模式：ddp（数据并行）/ fsdp（分片数据并行）/ none")
    # Streaming / performance
    parser.add_argument("--use_panels", type=int, default=1, help="使用 panels.pt 流式训练 (1) 或使用预制序列 train.pt (0)")
    parser.add_argument("--grad_accum_steps", type=int, default=8, help="梯度累计步数，用于大等效批量")
    parser.add_argument("--use_amp", type=int, default=0, help="启用自动混合精度训练 (1/0)")
    parser.add_argument("--num_workers", type=int, default=0)
    parser.add_argument("--pin_memory", type=int, default=1)
    parser.add_argument("--prefetch_factor", type=int, default=2)
    parser.add_argument("--progress", type=int, default=0, help="训练时显示进度条 (1/0)")
    # Scheduler
    parser.add_argument("--scheduler_type", type=str, default="plateau", choices=["sgdr", "plateau"], help="学习率调度器类型：sgdr 或 plateau")
    parser.add_argument("--scheduler_patience", type=int, default=3)
    parser.add_argument("--scheduler_factor", type=float, default=0.5)
    parser.add_argument("--sgdr_T_0", type=int, default=50, help="SGDR 初始周期（单位：epoch）")
    parser.add_argument("--sgdr_T_mult", type=int, default=2, help="SGDR 周期放大倍数")
    parser.add_argument("--sgdr_eta_min", type=float, default=1e-6, help="SGDR 最小学习率")
    # Warmup 参数
    parser.add_argument("--warmup_epochs", type=int, default=3, help="Warmup 线性升温的 epoch 数")
    parser.add_argument("--warmup_init_factor", type=float, default=0.1, help="Warmup 初始学习率系数，相对 base lr（0~1）")
    # Time CV
    parser.add_argument("--time_cv_folds", type=int, default=1, help="滚动时间交叉验证的折数（>1 启用）")
    # Logging
    parser.add_argument("--print_shapes_once", type=int, default=1, help="是否在第一批打印一次 X/Y/Pred 形状")
    parser.add_argument("--log_png", type=int, default=1, help="是否保存训练曲线PNG (1/0)")
    parser.add_argument("--plot_dir", type=str, default="plots/train", help="PNG 输出目录")
    # 额外记录：按波动率分桶的 score 范数/损失、样本 realized-vol 分布
    parser.add_argument("--log_vol_buckets", type=int, default=8, help="按 yz_vol_scale 分桶数量（0=禁用）")
    parser.add_argument("--log_vol_bucket_mode", type=str, default="logspace", choices=["logspace", "quantile"], help="vol分桶边界模式：logspace 或 quantile（基于训练分布分位数）")
    parser.add_argument("--log_vol_json", type=str, default="plots/train/vol_stats.jsonl", help="每epoch追加一行JSON：各桶的score范数/损失/样本数等")
    parser.add_argument("--log_vol_png", type=str, default="plots/train/vol_buckets.png", help="每epoch覆盖保存：按桶的曲线/柱状图（若 log_vol_plot_each_epoch=1）")
    parser.add_argument("--log_vol_plot_each_epoch", type=int, default=0, help="是否每个epoch绘制vol分桶PNG(1)；0=仅最终汇总绘制一次")
    parser.add_argument("--log_vol_png_final", type=str, default="plots/train/vol_buckets.png", help="最终汇总PNG路径（训练结束时绘制一次）")
    parser.add_argument("--log_vol_summary_json", type=str, default="plots/train/vol_stats_summary.json", help="最终汇总JSON路径（训练结束时覆盖写出）")
    # 偏差/尺度诊断日志（一次训练尽量定位系统性偏大的原因）
    parser.add_argument("--log_bias_json", type=str, default="plots/train/bias_stats.jsonl", help="每epoch追加：σ/ s_yz/ 组合尺度 与 score/target/y0重建偏差 统计")
    # Early stopping
    parser.add_argument("--early_stop_patience", type=int, default=5, help="早停耐心次数（val不提升时允许的连续epoch数）")
    parser.add_argument("--early_stop_min_delta", type=float, default=0.0, help="判定提升所需的最小改善阈值 (val_loss需较best减少超出该值)")
    # 时间扭曲（VE范式下的 σ(t) 条件化）：t_eff = clamp(t * s_bar^gamma)
    parser.add_argument("--time_warp_gamma", type=float, default=0.5, help="时间嵌入扭曲指数 γ，按 s_yz 的资产均值做 t→t*s_bar^γ (0 表示禁用扭曲)")
    # Save policy
    parser.add_argument("--skip_first_best_saves", type=int, default=0, help="跳过前N次best提升的保存以减少I/O；超过N次后才开始保存")
    # 验证集尾部裁剪，排除预处理保留的 tail_fill 天数对验证的影响
    parser.add_argument("--trim_tail_days_val", type=int, default=0, help="验证集裁掉末尾 k 天（避免 tail_fill 进入验证集）")
    # Saving controls (tuning可禁用保存，仅输出指标/配置)
    parser.add_argument("--save_model", type=int, default=1, help="是否保存最佳权重 (1/0)")
    parser.add_argument("--save_best_metrics_path", type=str, default="", help="将最佳指标写入到该JSON路径（可选）")
    parser.add_argument("--save_best_config_path", type=str, default="", help="将本次使用的模型配置写入到该JSON路径（可选）")
    # EMA 与 σ 策略/分桶日志
    parser.add_argument("--ema", type=int, default=1, help="启用EMA权重 (1/0)")
    parser.add_argument("--ema_decay", type=float, default=0.999, help="EMA衰减系数")
    parser.add_argument("--sigma_schedule", type=str, default="log_uniform", help="噪声采样日程: log_uniform|karras")
    parser.add_argument("--sigma_min", type=float, default=None, help="覆盖 score_sigma_min（可选统一入口）")
    parser.add_argument("--sigma_max", type=float, default=None, help="覆盖 score_sigma_max（可选统一入口）")
    parser.add_argument("--sigma_data", type=float, default=None, help="EDM 预条件化中的 sigma_data（None=使用默认0.025）")
    parser.add_argument("--log_sigma_buckets", type=int, default=8, help="按σ或t分桶记录diffmse")
    parser.add_argument("--log_sigma_json", type=str, default="", help="可选：将每epoch的σ分桶统计写入JSONL路径")
    parser.add_argument("--log_t_loss_png", type=str, default="plots/train/t_loss.png", help="按时间步绘制loss分布PNG路径")
    # 训练过程参数和梯度监控
    parser.add_argument("--log_layer_stats", type=int, default=1, help="是否记录各层参数和梯度统计(1/0)")
    parser.add_argument("--log_layer_stats_path", type=str, default="plots/train/layer_stats.jsonl", help="层级参数梯度记录路径（每epoch一行）")
    parser.add_argument("--log_layer_stats_interval", type=int, default=1, help="记录间隔（每N个epoch记录一次）")
    parser.add_argument("--log_t_loss_json", type=str, default="", help="按时间步loss统计JSONL路径（可选）")
    parser.add_argument("--log_t_grad_json", type=str, default="", help="记录按时间步梯度范数JSONL路径（可选）")
    parser.add_argument("--log_param_stats_json", type=str, default="", help="可选：每隔若干epoch记录参数统计(JSONL)")
    parser.add_argument("--log_param_stats_interval", type=int, default=1, help="参数统计写入的epoch间隔")
    parser.add_argument("--log_spectral_stats_json", type=str, default="", help="可选：记录自适应谱噪声统计和敏感度(JSONL)")
    parser.add_argument("--log_spectral_stats_interval", type=int, default=10, help="谱噪声统计写入的epoch间隔")
    parser.add_argument("--spectral_sensitivity_interval", type=int, default=100, help="敏感度采样间隔（每N个iteration计算一次梯度）")
    parser.add_argument("--diffusion_head_hidden", type=int, default=None, help="Diffusion head隐藏维度（None=使用d_model）")
    parser.add_argument("--diffusion_head_dropout", type=float, default=None, help="Diffusion head dropout（None=使用dropout）")
    # Multi-CLS 聚合（固定 attn 模式，仅暴露查询数与温度）
    parser.add_argument("--num_pool_queries", type=int, default=6)
    parser.add_argument("--pooling_tau", type=float, default=0.7)
    # 越界惩罚：约束预测 log-return 不超过理论涨跌幅
    parser.add_argument("--reg_bound_w", type=float, default=0.0, help="越界惩罚损失权重（超出理论涨跌幅范围时施加二次惩罚）")
    parser.add_argument("--bound_return_upper", type=float, default=0.3444, help="t+2 对 t+1 相对涨幅上限（比例）")
    parser.add_argument("--bound_return_lower", type=float, default=-0.2636, help="t+2 对 t+1 相对跌幅下限（比例，负数）")
    parser.add_argument("--graph_l2_reg", type=float, default=0.0, help="图分支 Frobenius 正则权重，鼓励稀疏高信噪比的 A_learned")
    # （DDPM 已移除）
    parser.add_argument("--pred_consist_w", type=float, default=0.05, help="预测头与去噪重建 y0 的一致性损失权重（保持推理稳定）")
    # Score-Based 训练开关（替代 ε-parameterization）：预测 score 而非噪声
    parser.add_argument("--use_score_model", type=int, default=1, help="启用 SBDM 训练 (1/0)：主损失为 score 匹配")
    parser.add_argument("--score_sigma_min", type=float, default=0.01, help="最小噪声强度，用于 score 目标计算")
    parser.add_argument("--score_sigma_max", type=float, default=1.0, help="最大噪声强度，用于 score 目标计算")
    # 噪声缩放（波动率调制）
    parser.add_argument("--noise_scale_mode", type=str, default="constant", choices=["constant", "yz_vol"], help="训练加噪幅度模式：constant=单高斯；yz_vol=按 Yang-Zhang 波动率缩放")
    parser.add_argument("--noise_scale_csv", type=str, default="data/processed/yz_vol.csv", help="当 noise_scale_mode=yz_vol 时，提供按 (date,stock) 的 yz_vol CSV 路径")
    parser.add_argument("--noise_scale_col", type=str, default="yz_vol_scale", help="noise_scale_csv 中作为缩放的列名，默认使用预处理生成的 yz_vol_scale")
    parser.add_argument("--noise_scale_clip", type=float, default=10.0, help="对缩放系数进行上限截断，避免异常放大")
    parser.add_argument("--adaptive_spectral_noise", type=int, default=0, help="启用特征谱自适应噪声增强 (1/0) [已弃用,请使用spectral_noise_enabled]")
    parser.add_argument("--spectral_noise_enabled", type=int, default=0, help="启用自适应谱噪声 (1/0)")
    # 新的自适应alpha参数
    parser.add_argument("--spectral_robustness_coeff", type=float, default=9.0, help="自适应alpha公式的鲁棒性系数 C (alpha = 1 + C/(k+epsilon))")
    parser.add_argument("--spectral_alpha_epsilon", type=float, default=0.1, help="自适应alpha公式的防除零常数 epsilon")
    parser.add_argument("--spectral_k_threshold", type=float, default=1.0, help="高频特征判定阈值（k>=threshold则alpha=1）")
    # 旧的谱噪声参数（保留向后兼容）
    parser.add_argument("--spectral_noise_sigma_min", type=float, default=0.002, help="谱自适应噪声的 σ 下界 [已弃用]")
    parser.add_argument("--spectral_noise_sigma_max", type=float, default=80.0, help="谱自适应噪声的 σ 上界 [已弃用]")
    parser.add_argument("--spectral_noise_target_snr", type=float, default=3.0, help="谱自适应噪声目标信噪比 S_target [已弃用]")
    parser.add_argument("--spectral_noise_threshold", type=float, default=0.3, help="判定低频特征的 diff_ratio 阈值 [已弃用]")
    parser.add_argument("--spectral_noise_alpha_floor", type=float, default=0.05, help="低频特征噪声因子的最小缩放 [已弃用]")
    parser.add_argument("--spectral_noise_alpha_cap", type=float, default=0.9, help="低频特征噪声因子的最大缩放 [已弃用]")
    parser.add_argument("--spectral_noise_freq_div", type=int, default=3, help="谱能量低频截断的分母（seq_len//freq_div，默认3→cutoff=20）")
    parser.add_argument("--spectral_noise_log_json", type=str, default="", help="自适应谱噪声统计输出JSONL路径")
    parser.add_argument("--spectral_noise_log_interval", type=int, default=1, help="谱噪声统计写入的epoch间隔")
    parser.add_argument("--spectral_noise_log_examples", type=int, default=5, help="谱噪声统计中每类示例特征数")
    parser.add_argument("--spectral_noise_full_dump", type=int, default=0, help="是否输出完整特征向量(0=只输出摘要, 1=完整)")
    # OOS (by stocks) for training holdout
    parser.add_argument("--oos_stock_ratio", type=float, default=0.0, help="训练期按比例留出股票作OOS（不参与训练）")
    parser.add_argument("--oos_stock_list", type=str, default="", help="训练期使用的OOS股票列表文件路径（每行一个股票ID），优先于比例")
    parser.add_argument("--oos_seed", type=int, default=42, help="OOS股票采样的随机种子")
    parser.add_argument("--log_coverage_path", type=str, default="", help="覆盖率分析日志路径")
    parser.add_argument("--sample_coverage_threshold", type=float, default=0.5, help="样本覆盖率阈值")
    # EDM 相关参数
    parser.add_argument("--edm_lambda", type=int, default=0, help="启用 EDM λ(σ) 加权 (1/0)")
    parser.add_argument("--edm_lambda_norm", type=str, default="batch_mean", help="EDM λ 归一化模式")
    parser.add_argument("--edm_lambda_pow", type=float, default=1.0, help="EDM λ 的幂次")
    # yzvol 条件开关
    parser.add_argument("--vol_cond_as_input", type=int, default=1, help="yzvol 作为条件输入网络")
    parser.add_argument("--vol_cond_scale_output", type=int, default=0, help="yzvol 直接乘到输出上")
    # Sensitivity checks (replace SNR monitoring)
    parser.add_argument("--check_sensitivity", type=int, default=1, help="启用特征敏感度阈值检查(1/0)")
    parser.add_argument("--check_sensitivity_features", type=str, default="", help="逗号分隔的特征索引，例如: 0,1,2")
    parser.add_argument("--check_sensitivity_threshold", type=float, default=0.01, help="敏感度阈值，默认0.01")
    parser.add_argument("--check_sensitivity_interval", type=int, default=64, help="每隔多少步进行一次敏感度检查")
    parser.add_argument("--log_grad_clip_json", type=str, default="", help="可选：记录梯度裁剪前范数的JSONL路径")
    parser.add_argument("--log_grad_clip_interval", type=int, default=50, help="梯度裁剪日志写入间隔（每多少次更新记录一次）")
    args = parser.parse_args()
    return args


def evaluate(model: SpacetimeGNNMAM, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    loss_fn = nn.MSELoss()
    total_loss = 0.0
    total_count = 0
    with torch.no_grad():
        for batch in loader:
            if isinstance(batch, (list, tuple)) and len(batch) == 4:
                X, Y, _, mask_b = batch  # [mask机制]
            elif isinstance(batch, (list, tuple)) and len(batch) == 3:
                X, Y, _ = batch
                mask_b = None
            else:
                X, Y = batch
                mask_b = None
            X = X.to(device)
            Y = Y.to(device)
            if mask_b is not None:
                mask_tensor = mask_b.to(device)
            else:
                mask_tensor = None
            feat = model.encode(X, mask=mask_tensor)  # [mask机制]  # pyright: ignore[reportGeneralTypeIssues]
            pred = torch.mean(feat, dim=-1, keepdim=True)
            loss = loss_fn(pred, Y)
            total_loss += loss.item() * X.size(0)
            total_count += X.size(0)
    return total_loss / max(1, total_count)


def build_time_cv_folds(num_samples: int, k: int) -> List[Tuple[List[int], List[int]]]:
    """滚动时间CV：将训练样本按时间等分为 k+1 段。
    第 i 折：train = [0 : cut_i]，val = (cut_i : cut_{i+1}]。
    """
    cuts = [int(num_samples * i / (k + 1)) for i in range(k + 1)] + [num_samples]
    folds = []
    for i in range(1, k + 1):
        train_end = cuts[i]
        val_end = cuts[i + 1]
        if val_end <= train_end:
            continue
        train_idx = list(range(0, train_end))
        val_idx = list(range(train_end, val_end))
        folds.append((train_idx, val_idx))
    return folds


def try_plot_train_curves(plot_dir: str, train_losses: list, val_losses: list, lrs: list) -> None:
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import warnings
        # 设置中文字体（使用系统实际可用的字体）
        plt.rcParams['font.sans-serif'] = ['Noto Sans CJK JP', 'Droid Sans Fallback', 'SimHei', 'Microsoft YaHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib.font_manager')
    except Exception as e:
        print(f"[plot] matplotlib not available, skip plotting: {e}")
        return
    ensure_dir(plot_dir)
    # Loss curves
    try:
        plt.figure(figsize=(6, 4))
        plt.plot(range(1, len(train_losses) + 1), train_losses, label="train_loss")
        plt.plot(range(1, len(val_losses) + 1), val_losses, label="val_loss")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.title("Training and Validation Loss")
        plt.legend()
        # 强制使用整数刻度作为 Epoch 轴，避免出现小数刻度
        try:
            from matplotlib.ticker import MaxNLocator
            ax = plt.gca()
            ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        except Exception:
            pass
        plt.tight_layout()
        plt.savefig(os.path.join(plot_dir, "train_val_loss.png"), dpi=150)
        plt.close()
    except Exception as e:
        print(f"[plot] failed to save loss curve: {e}")
    # LR curve
    try:
        plt.figure(figsize=(6, 4))
        plt.plot(range(1, len(lrs) + 1), lrs, label="lr")
        plt.xlabel("Epoch")
        plt.ylabel("Learning Rate")
        plt.title("Learning Rate by Epoch")
        plt.legend()
        # 强制使用整数刻度作为 Epoch 轴，避免出现小数刻度
        try:
            from matplotlib.ticker import MaxNLocator
            ax = plt.gca()
            ax.xaxis.set_major_locator(MaxNLocator(integer=True))
        except Exception:
            pass
        plt.tight_layout()
        plt.savefig(os.path.join(plot_dir, "lr.png"), dpi=150)
        plt.close()
    except Exception as e:
        print(f"[plot] failed to save lr curve: {e}")

def main():  # pyright: ignore[reportGeneralTypeIssues]
    args = parse_args()
    # 读取配置文件，失败应立即报错
    cfg_path = None
    if isinstance(getattr(args, "config", ""), str) and len(args.config) > 0 and os.path.exists(args.config):
        cfg_path = args.config
    elif os.path.exists("config.json"):
        cfg_path = "config.json"
    if cfg_path is not None:
        with open(cfg_path, "r", encoding="utf-8") as f:
            cfg = json.load(f)
        use_tuned_best = bool(cfg.get("train", {}).get("use_tuned_best", False))
        if use_tuned_best is False:
            tcfg = cfg.get("train", {})
            # 仅当键存在时才覆盖，避免把 None 写进来
            def _ovr(name, cast_fn=lambda x: x):
                val = tcfg.get(name, None)
                if val is not None:
                    setattr(args, name, cast_fn(val))
            _ovr("batch_size", int)
            _ovr("epochs", int)
            _ovr("lr", float)
            _ovr("weight_decay", float)
            _ovr("d_model", int)
            _ovr("num_heads", int)
            _ovr("num_layers", int)
            _ovr("d_ff", int)
            _ovr("dropout", float)
            _ovr("temp", float)
            _ovr("adj_activation", str)
            _ovr("entmax_alpha", float)
            _ovr("objective", str)
            _ovr("grad_accum_steps", int)
            _ovr("use_amp", int)
            _ovr("num_workers", int)
            _ovr("pin_memory", int)
            _ovr("prefetch_factor", int)
            _ovr("scheduler_patience", int)
            _ovr("scheduler_factor", float)
            _ovr("scheduler_type", str)
            _ovr("sgdr_T_0", int)
            _ovr("sgdr_T_mult", int)
            _ovr("sgdr_eta_min", float)
            _ovr("warmup_epochs", int)
            _ovr("warmup_init_factor", float)
            # （DDPM 参数已移除）
            _ovr("pred_consist_w", float)
                # 波动率噪声缩放相关覆盖
            _ovr("noise_scale_mode", str)
            _ovr("noise_scale_csv", str)
            _ovr("noise_scale_col", str)
            _ovr("noise_scale_clip", float)
            _ovr("noise_scale_log1p", int)
            _ovr("noise_scale_row_zscore", int)
            _ovr("noise_scale_row_gamma", float)
            _ovr("noise_scale_z_eps", float)
            _ovr("use_score_model", int)
            _ovr("score_sigma_min", float)
            _ovr("score_sigma_max", float)
                # 新增：扩散噪声与EMA/分桶日志参数透传
            _ovr("sigma_schedule", str)
            _ovr("sigma_min", float)
            _ovr("sigma_max", float)
            _ovr("sigma_data", float)
            _ovr("log_sigma_buckets", int)
            _ovr("log_sigma_json", str)
            _ovr("log_t_loss_json", str)
                # Vol-bucket 可视化/统计 透传
            _ovr("log_vol_buckets", int)
            _ovr("log_vol_bucket_mode", str)
            _ovr("log_vol_json", str)
            _ovr("log_vol_png", str)
            _ovr("ema", int)
            _ovr("ema_decay", float)
            _ovr("time_cv_folds", int)
            _ovr("print_shapes_once", int)
            _ovr("early_stop_patience", int)
            _ovr("early_stop_min_delta", float)
            _ovr("skip_first_best_saves", int)
            _ovr("trim_tail_days_val", int)
            _ovr("save_model", int)
            _ovr("diffusion_head_hidden", int)
            _ovr("diffusion_head_dropout", float)
            _ovr("num_pool_queries", int)
            _ovr("pooling_tau", float)
                # 噪声缩放配置（按波动率）
            _ovr("noise_scale_mode", str)
            _ovr("noise_scale_csv", str)
            _ovr("noise_scale_col", str)
            _ovr("noise_scale_clip", float)
            _ovr("reg_bound_w", float)
            _ovr("bound_return_upper", float)
            _ovr("bound_return_lower", float)
            _ovr("dist_mode", str)
            _ovr("edm_lambda", int)
            _ovr("edm_lambda_norm", str)
            _ovr("edm_lambda_pow", float)
            _ovr("vol_cond_as_input", int)
            _ovr("vol_cond_scale_output", int)
                # Sensitivity checks (replace SNR monitoring)
            _ovr("check_sensitivity", int)
            _ovr("check_sensitivity_features", str)
            _ovr("check_sensitivity_threshold", float)
            _ovr("check_sensitivity_interval", int)
            _ovr("log_grad_clip_json", str)
            _ovr("log_grad_clip_interval", int)
            print(f"[train] 超参已按 {cfg_path} 的 train 段覆盖（use_tuned_best=false）", flush=True)
    if getattr(args, "sigma_data", None) is None:
        args.sigma_data = 0.025
    bound_upper = float(getattr(args, 'bound_return_upper', 0.3444))
    bound_lower = float(getattr(args, 'bound_return_lower', -0.2636))
    if bound_upper <= 0.0:
        raise ValueError(f"bound_return_upper 应为正数，当前={bound_upper}")
    if bound_lower >= 0.0:
        raise ValueError(f"bound_return_lower 应为负数，当前={bound_lower}")
    args.bound_log_upper = math.log(1.0 + bound_upper)
    args.bound_log_lower = math.log(1.0 + bound_lower)
    distributed, world_size, rank, local_rank = setup_distributed()
    ensure_dir(args.checkpoint_dir)
    # 首次训练前清空 σ / t-loss 统计日志，避免续写旧 run
    try:
        if is_main_process(rank):
            sigma_json_path = getattr(args, 'log_sigma_json', '')
            if isinstance(sigma_json_path, str) and len(sigma_json_path) > 0:
                if os.path.exists(sigma_json_path):
                    os.remove(sigma_json_path)
            tloss_json_path = getattr(args, 'log_t_loss_json', '')
            if isinstance(tloss_json_path, str) and len(tloss_json_path) > 0:
                if os.path.exists(tloss_json_path):
                    os.remove(tloss_json_path)
    except Exception as _e:
        if is_main_process(rank):
            print(f"[warn] reset log_sigma_json failed: {_e}", flush=True)
    # 高性能默认
    try:
        torch.set_float32_matmul_precision('high')
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    except Exception:
        pass

    use_panels = args.use_panels == 1
    if use_panels:
        panels_path = os.path.join(args.data_dir, "panels.pt")
        # 构建训练期 OOS 股票留出列表
        # 读取OOS股票列表，失败应立即报错
        exclude_list: Optional[List[str]] = None
        if isinstance(args.oos_stock_list, str) and len(args.oos_stock_list) > 0 and os.path.exists(args.oos_stock_list):
            with open(args.oos_stock_list, "r", encoding="utf-8") as f:
                exclude_list = [line.strip() for line in f if line.strip()]
        elif args.oos_stock_ratio and float(args.oos_stock_ratio) > 0:
            tmp_ds = StreamingPanelDataset(panels_path, split="train")
            stocks_all = tmp_ds.stocks
            rng = np.random.default_rng(int(args.oos_seed))
            k = max(1, int(len(stocks_all) * float(args.oos_stock_ratio)))
            idx = rng.choice(len(stocks_all), size=k, replace=False)
            exclude_list = [stocks_all[i] for i in idx]
            # 持久化到 data_dir 方便回测共用
            oos_path = os.path.join(args.data_dir, "oos_stocks.txt")
            with open(oos_path, "w", encoding="utf-8") as f:
                for s in exclude_list:
                    f.write(f"{s}\n")
            print(f"[train] 写出 OOS 股票列表: {oos_path} (count={len(exclude_list)})", flush=True)

        train_set = StreamingPanelDataset(panels_path, split="train", exclude_stocks=exclude_list, trim_tail_days=0)
        # 严格：验证集同样排除 OOS，并裁掉末尾填充日
        val_trim_k = int(getattr(args, "trim_tail_days_val", 0))
        if val_trim_k > 0:
            print(f"[train] trim_tail_days_val={val_trim_k}: 将从验证集末尾裁掉 {val_trim_k} 天", flush=True)
        val_set = StreamingPanelDataset(panels_path, split="val", exclude_stocks=exclude_list, trim_tail_days=val_trim_k)
        seq_len = train_set.seq_len
        input_dim = train_set.input_dim
        num_stocks = train_set.feature_panel.shape[1]
    else:
        train_set = PanelDataset(os.path.join(args.data_dir, "train.pt"))
        val_set = PanelDataset(os.path.join(args.data_dir, "val.pt"))
        seq_len = train_set.seq_len
        input_dim = train_set.input_dim
        num_stocks = train_set.X.shape[1]

    # 读取特征名（用于回测校验与对齐）
    # 读取feature_names，失败应立即报错而不是静默
    meta_path = os.path.join(args.data_dir, "meta.json")
    with open(meta_path, "r", encoding="utf-8") as f:
        _meta = json.load(f)
    feature_names = _meta.get("features", {}).get("final_features", None)

    # 构造并缓存用于回测一致性校验的缩放参数（严格来源于 test_env.pt，且按训练股票顺序重排）
    art_mean_full: list | None = None
    art_std_full: list | None = None
    art_shapes: dict | None = None
    art_mean_head3: list | None = None
    art_std_head3: list | None = None
    art_scale_sha256: str | None = None
    # 移除try-except，读取失败应立即报错
    env_path = os.path.join(args.data_dir, "test_env.pt")
    if os.path.exists(env_path):
        env_pack = torch.load(env_path, weights_only=False)
        env_mean = env_pack["price_scale_mean"].detach().cpu().view(-1)  # [N_env]
        env_std = env_pack["price_scale_std"].detach().cpu().view(-1)
        env_stocks = env_pack.get("stocks", None)
        if not isinstance(env_stocks, list) or len(env_stocks) == 0:
            raise KeyError("test_env.pt 中缺少 stocks 列表")
        pos_in_env = {str(s): i for i, s in enumerate(env_stocks)}
        train_stocks_here = train_set.stocks if hasattr(train_set, 'stocks') else []
        order_train_to_env = [int(pos_in_env[str(s)]) for s in train_stocks_here]
        mean_reordered = env_mean[order_train_to_env].float().numpy()
        std_reordered = env_std[order_train_to_env].float().numpy()
        art_mean_full = [float(x) for x in mean_reordered.tolist()]
        art_std_full = [float(x) for x in std_reordered.tolist()]
        art_shapes = {"mean_vec": [1, len(art_mean_full), 1], "std_vec": [1, len(art_std_full), 1]}
        art_mean_head3 = art_mean_full[:3]
        art_std_head3 = art_std_full[:3]
        # 计算校验摘要（sha256）
        import hashlib
        h = hashlib.sha256()
        h.update(mean_reordered.tobytes())
        h.update(std_reordered.tobytes())
        art_scale_sha256 = h.hexdigest()
    else:
        if is_main_process(rank):
            print(f"[warn] test_env.pt not found: {env_path}", flush=True)

    # 设备选择：DDP 下强制使用 local_rank 绑定的 GPU
    if distributed and (args.device.startswith("cuda") and torch.cuda.is_available()):
        device = torch.device(f"cuda:{local_rank}")
    else:
        device = torch.device(args.device)

    # 数据集摘要（仅主进程打印）
    if is_main_process(rank):
        try:
            print(
                f"dataset loaded: use_panels={use_panels} seq_len={seq_len} input_dim={input_dim} num_stocks={num_stocks} "
                f"| train_len={len(train_set)} val_len={len(val_set)}",
                flush=True,
            )
        except Exception:
            pass

    # 可选：时间滚动交叉验证（在训练集内部；仅在单机主进程执行，避免分布式重复）
    if (not distributed) and args.time_cv_folds and args.time_cv_folds > 1:
        folds = build_time_cv_folds(len(train_set), args.time_cv_folds)
        print(f"启用时间CV：folds={len(folds)}", flush=True)
        cv_losses = []
        for fid, (tr_idx, va_idx) in enumerate(folds, 1):
            tr_loader = DataLoader(Subset(train_set, tr_idx), batch_size=args.batch_size, shuffle=False)
            va_loader = DataLoader(Subset(train_set, va_idx), batch_size=args.batch_size, shuffle=False)
            cv_epochs = max(3, min(args.epochs, 8))
            model = SpacetimeGNNMAM(
                input_dim=input_dim,
                num_stocks=num_stocks,
                seq_len=seq_len,
                d_model=args.d_model,
                num_heads=args.num_heads,
                num_layers=args.num_layers,
                d_ff=args.d_ff,
                num_classes=1,
                dropout=args.dropout,
                temp=args.temp,
                adj_activation=str(args.adj_activation),
                entmax_alpha=float(args.entmax_alpha),
                num_pool_queries=int(args.num_pool_queries),
                pooling_tau=float(args.pooling_tau),
                diffusion_head_hidden_dim=args.diffusion_head_hidden,
                diffusion_head_dropout=args.diffusion_head_dropout,
            ).to(device)
            optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
            loss_fn = nn.MSELoss()
            # 交叉验证只跑少量epoch以快速估计
            printed = False
            for epoch in range(1, cv_epochs + 1):
                model.train()
                for batch in tr_loader:
                    if isinstance(batch, (list, tuple)) and len(batch) == 4:
                        X, Y, _, mask_cv = batch  # [mask机制]
                    elif isinstance(batch, (list, tuple)) and len(batch) == 3:
                        X, Y, _ = batch
                        mask_cv = None
                    else:
                        X, Y = batch
                        mask_cv = None
                    X = X.to(device)
                    Y = Y.to(device)
                    if mask_cv is not None:
                        mask_cv = mask_cv.to(device)
                    optimizer.zero_grad(set_to_none=True)
                    feat = model.encode(X, mask=mask_cv)  # [mask机制]
                    pred = feat.mean(dim=-1, keepdim=True)
                    if args.print_shapes_once and (not printed):
                        print(f"[CV{fid}] X={tuple(X.shape)} Y={tuple(Y.shape)} pred={tuple(pred.shape)}", flush=True)
                        printed = True
                    loss = loss_fn(pred, Y)
                    loss.backward()
                    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                    optimizer.step()
            val_loss = evaluate(model, va_loader, device)
            print(f"[CV{fid}] val_loss={val_loss:.6f}", flush=True)
            cv_losses.append(val_loss)
        print(f"时间CV平均 val_loss={float(sum(cv_losses)/len(cv_losses)):.6f}", flush=True)

    # 正式训练（用原 train/val 切分），保存 best.pt
    pin = args.pin_memory == 1 and device.type == 'cuda'
    # 分布式采样器
    train_sampler = None
    val_sampler = None
    if distributed and args.dist == 1:
        train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False)
        # 验证集同样使用分布式采样器（确保不同进程读取不同子集）
        val_sampler = DistributedSampler(val_set, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)

    # DataLoader 进程上下文：fork 在部分系统上更易触发共享内存问题，尝试 forkserver（若可用）
    mp_ctx = None
    try:
        if args.num_workers and args.num_workers > 0:
            mp_ctx = py_mp.get_context('forkserver')
    except Exception:
        mp_ctx = None

    train_loader = DataLoader(
        train_set,
        batch_size=args.batch_size,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        drop_last=False,
        pin_memory=pin,
        num_workers=args.num_workers,
        prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
        persistent_workers=True if args.num_workers and args.num_workers > 0 else False,
        multiprocessing_context=mp_ctx if mp_ctx is not None else None,
    )
    # val 使用单进程 worker，禁止传 multiprocessing_context
    val_loader = DataLoader(
        val_set,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=val_sampler,
        drop_last=False,
        pin_memory=pin,
        num_workers=0,
        # 不传 multiprocessing_context，避免 num_workers=0 触发校验错误
    )

    if is_main_process(rank):
        try:
            print(
                f"dataloaders: train_steps={len(train_loader)} val_steps={len(val_loader)} batch_size={args.batch_size}",
                flush=True,
            )
        except Exception:
            pass

    model = SpacetimeGNNMAM(
        input_dim=input_dim,
        num_stocks=num_stocks,
        seq_len=seq_len,
        d_model=args.d_model,
        num_heads=args.num_heads,
        num_layers=args.num_layers,
        d_ff=args.d_ff,
        num_classes=1,
        dropout=args.dropout,
        temp=args.temp,
        adj_activation=str(args.adj_activation),
        entmax_alpha=float(args.entmax_alpha),
        num_pool_queries=int(args.num_pool_queries),
        pooling_tau=float(args.pooling_tau),
        diffusion_head_hidden_dim=args.diffusion_head_hidden,
        diffusion_head_dropout=args.diffusion_head_dropout,
    ).to(device)
    # 通过 config 透传 score head 容量（隐藏维与dropout）；不改变结构，仅改容量
    # DDP/FSDP 包装
    if distributed and args.dist == 1:
        mode = str(getattr(args, "dist_mode", "ddp")).lower()
        if mode == "fsdp":
            if FSDP is None:
                raise RuntimeError("请求使用 FSDP，但当前环境不支持 torch.distributed.fsdp。请升级到 PyTorch>=1.12 且使用 GPU/CUDA 环境。")
            model = FSDP(model, device_id=(local_rank if (device.type == 'cuda') else None))  # type: ignore[arg-type]
        elif mode == "ddp":
            find_unused = False
            model = DDP(
                model,
                device_ids=[local_rank] if device.type == 'cuda' else None,
                output_device=local_rank if device.type == 'cuda' else None,
                broadcast_buffers=True,
                find_unused_parameters=find_unused,
            )
        else:
            # none：不做任何分布式包装（但仍运行在单进程或 torchrun rank=0）
            pass

    # ========== 自适应谱噪声初始化 ==========
    spectral_noise_tracker: Optional[AdaptiveSpectralNoiseTracker] = None
    spectral_precomputed_data: Optional[dict] = None
    
    # 支持两个参数名：spectral_noise_enabled（新）或 adaptive_spectral_noise（旧）
    spectral_enabled = int(getattr(args, 'spectral_noise_enabled', getattr(args, 'adaptive_spectral_noise', 0)))
    if spectral_enabled == 1:
        if is_main_process(rank):
            print("[spectral] 启用自适应谱噪声", flush=True)
        
        # 尝试加载预计算的频域特性
        spectral_precomputed_path = os.path.join(str(getattr(args, 'data_dir', 'data/processed')), "spectral_features.pt")
        if os.path.exists(spectral_precomputed_path):
            if is_main_process(rank):
                print(f"[spectral] 加载预计算频域特性: {spectral_precomputed_path}", flush=True)
            try:
                spectral_precomputed_data = torch.load(spectral_precomputed_path, map_location='cpu', weights_only=False)
                
                # 加载预计算数据（基于频谱质心）
                spectral_centroid = spectral_precomputed_data["spectral_centroid_per_feature"]  # [F]
                k_normalized = spectral_precomputed_data["k_normalized"]  # [F]
                alpha_per_feature = spectral_precomputed_data["alpha_per_feature"]  # [F]
                
                if k_normalized.shape[0] != input_dim:
                    raise ValueError(f"特征维度不匹配：预计算={k_normalized.shape[0]} vs 当前={input_dim}")
                
                if is_main_process(rank):
                    centroid_mean = float(spectral_centroid.mean().item())
                    centroid_median = float(torch.median(spectral_centroid).item())
                    k_mean = float(k_normalized.mean().item())
                    k_median = float(torch.median(k_normalized).item())
                    alpha_mean = float(alpha_per_feature.mean().item())
                    alpha_median = float(torch.median(alpha_per_feature).item())
                    alpha_max = float(alpha_per_feature.max().item())
                    alpha_min = float(alpha_per_feature.min().item())
                    
                    print(f"[spectral] 预计算数据加载成功！特征数={input_dim}", flush=True)
                    print(f"[spectral] 频谱质心分布: mean={centroid_mean:.4f}, median={centroid_median:.4f}", flush=True)
                    print(f"[spectral] k_normalized分布: mean={k_mean:.4f}, median={k_median:.4f}", flush=True)
                    print(f"[spectral] Alpha分布: mean={alpha_mean:.2f}x, median={alpha_median:.2f}x, range=[{alpha_min:.2f}x, {alpha_max:.2f}x]", flush=True)
                    
                    # 展示破坏最强的5个特征
                    top_alpha_idx = torch.topk(alpha_per_feature, k=min(5, input_dim)).indices
                    print(f"[spectral] 破坏最强的5个特征（alpha最大，低频主导）:")
                    for idx in top_alpha_idx:
                        feat_idx = int(idx.item())
                        centroid_val = float(spectral_centroid[feat_idx].item())
                        k_norm = float(k_normalized[feat_idx].item())
                        alpha_val = float(alpha_per_feature[feat_idx].item())
                        print(f"  Feature[{feat_idx}]: centroid={centroid_val:.4f}, k_norm={k_norm:.4f}, alpha={alpha_val:.1f}x")
                    
                    # 展示破坏最弱的5个特征
                    bottom_alpha_idx = torch.topk(alpha_per_feature, k=min(5, input_dim), largest=False).indices
                    print(f"[spectral] 破坏最弱的5个特征（alpha最小，高频主导）:")
                    for idx in bottom_alpha_idx:
                        feat_idx = int(idx.item())
                        centroid_val = float(spectral_centroid[feat_idx].item())
                        k_norm = float(k_normalized[feat_idx].item())
                        alpha_val = float(alpha_per_feature[feat_idx].item())
                        print(f"  Feature[{feat_idx}]: centroid={centroid_val:.4f}, k_norm={k_norm:.2f}, alpha={alpha_val:.2f}x")
            
            except Exception as e_load:
                if is_main_process(rank):
                    print(f"[spectral] WARNING: 加载预计算数据失败: {e_load}", flush=True)
                    print(f"[spectral] 将使用在线计算模式", flush=True)
                spectral_precomputed_data = None
        else:
            if is_main_process(rank):
                print(f"[spectral] 未找到预计算文件: {spectral_precomputed_path}", flush=True)
                print(f"[spectral] 将使用在线计算模式（较慢）", flush=True)
        
        # 如果没有预计算数据，使用在线计算
        if spectral_precomputed_data is None:
            freq_div = int(getattr(args, 'spectral_noise_freq_div', 3))
            spectral_noise_tracker = AdaptiveSpectralNoiseTracker(seq_len, input_dim, freq_div=freq_div)
            if is_main_process(rank):
                freq_bins = seq_len // 2 + 1
                cutoff = seq_len // freq_div
                print(f"[spectral] 在线计算模式 - FFT配置: seq_len={seq_len}, freq_bins={freq_bins}, freq_div={freq_div}, cutoff={cutoff}", flush=True)
        
        # 在线计算模式：遍历训练集预计算特征统计
        if spectral_precomputed_data is None:
            try:
                model.eval()
                precomp_count = 0
                precomp_limit = len(train_loader) if hasattr(train_loader, '__len__') else 999999
                if is_main_process(rank):
                    print(f"[spectral] 预计算批次上限: {precomp_limit}", flush=True)
                with torch.no_grad():
                    for it_pre, batch_pre in enumerate(train_loader, 1):
                        if it_pre > precomp_limit:
                            break
                        if isinstance(batch_pre, (list, tuple)) and len(batch_pre) >= 2:
                            X_pre = batch_pre[0]
                        else:
                            X_pre = batch_pre
                        X_pre = X_pre.to(device, non_blocking=pin)
                        # X_pre: [B, N, T, F]
                        spectral_noise_tracker._update_statistics(X_pre)
                        precomp_count += int(X_pre.size(0))
                        if is_main_process(rank) and it_pre % 50 == 0:
                            print(f"[spectral] 预计算进度: {it_pre}/{precomp_limit} batches, {precomp_count} samples", flush=True)
                model.train()
                if is_main_process(rank):
                    print(f"[spectral] 预计算完成，累计样本数: {precomp_count}", flush=True)
                    # 输出初步统计摘要（使用全局自适应alpha公式）
                    robustness_coeff = float(getattr(args, 'spectral_robustness_coeff', 9.0))
                    alpha_epsilon = float(getattr(args, 'spectral_alpha_epsilon', 0.1))
                    stats_init = spectral_noise_tracker._compute_stats(
                        robustness_coeff=robustness_coeff, 
                        epsilon=alpha_epsilon
                    )
                    if stats_init is not None:
                        spectral_noise_tracker.last_stats = stats_init  # 保存到tracker以供训练使用
                        alpha_vec = stats_init["alpha"].cpu()
                        k_feat = stats_init["k_per_feature"].cpu()
                        
                        # 统计alpha分布
                        alpha_mean = float(alpha_vec.mean().item())
                        alpha_median = float(torch.median(alpha_vec).item())
                        alpha_max = float(alpha_vec.max().item())
                        alpha_min_val = float(alpha_vec.min().item())
                        
                        print(f"[spectral] 全局自适应（无硬截断）：alpha_i = 1.0 + {robustness_coeff:.1f} / (k_i + {alpha_epsilon:.2f})", flush=True)
                        print(f"[spectral] Alpha分布: mean={alpha_mean:.2f}x, median={alpha_median:.2f}x, range=[{alpha_min_val:.2f}x, {alpha_max:.2f}x]", flush=True)
                        
                        # 展示破坏最强的特征（alpha最大，k最小）
                        example_count = min(5, input_dim)
                        top_alpha_idx = torch.topk(alpha_vec, k=example_count).indices
                        print(f"[spectral] 破坏最强的{example_count}个特征（alpha最大）:")
                        for i in top_alpha_idx:
                            feat_idx = int(i.item())
                            print(f"  Feature[{feat_idx}]: k={float(k_feat[feat_idx].item()):.4f}, alpha={float(alpha_vec[feat_idx].item()):.1f}x")
                        
                        # 展示破坏最弱的特征（alpha最小，k最大）
                        bottom_alpha_idx = torch.topk(alpha_vec, k=example_count, largest=False).indices
                        print(f"[spectral] 破坏最弱的{example_count}个特征（alpha最小）:")
                        for i in bottom_alpha_idx:
                            feat_idx = int(i.item())
                            print(f"  Feature[{feat_idx}]: k={float(k_feat[feat_idx].item()):.2f}, alpha={float(alpha_vec[feat_idx].item()):.2f}x")
            except Exception as e_spec:
                if is_main_process(rank):
                    print(f"[spectral] 预计算失败: {e_spec}", flush=True)
                    print(f"[spectral] ERROR: 自适应谱噪声初始化失败，请检查数据和参数！", flush=True)
                raise RuntimeError(f"自适应谱噪声预计算失败: {e_spec}") from e_spec

    # 优化器（移除 τ）
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.lr,
        weight_decay=args.weight_decay
    )
    scheduler_sgdr = None
    scheduler_plateau = None
    if str(getattr(args, "scheduler_type", "plateau")).lower() == "sgdr":
        scheduler_sgdr = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer,
            T_0=int(args.sgdr_T_0),
            T_mult=int(args.sgdr_T_mult),
            eta_min=float(args.sgdr_eta_min),
        )
    else:
        scheduler_plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=args.scheduler_factor,
        patience=args.scheduler_patience
    )
    loss_fn = nn.MSELoss()
    use_amp = (args.use_amp == 1 and device.type == 'cuda')
    # 训练启动时一次性打印AMP/TF32/设备信息
    if is_main_process(rank):
        try:
            tf32_matmul = getattr(torch, 'get_float32_matmul_precision', lambda: 'unknown')()
        except Exception:
            tf32_matmul = 'unknown'
        try:
            tf32_cuda = bool(getattr(torch.backends.cuda.matmul, 'allow_tf32', False))
            tf32_cudnn = bool(getattr(torch.backends.cudnn, 'allow_tf32', False))
        except Exception:
            tf32_cuda, tf32_cudnn = False, False
        print(f"[amp] use_amp={use_amp} device={device} tf32_matmul={tf32_matmul} allow_tf32(cublas/cudnn)={tf32_cuda}/{tf32_cudnn}", flush=True)
    # EMA 状态
    use_ema = (int(getattr(args, 'ema', 1)) == 1)
    ema_decay = float(getattr(args, 'ema_decay', 0.999))
    ema_shadow = {k: v.detach().clone() for k, v in model.state_dict().items()} if use_ema else None

    # =============== Resume: 恢复训练 ===============
    start_epoch = 1
    try:
        if int(getattr(args, 'resume', 0)) == 1:
            ck_path = str(getattr(args, 'resume_path', os.path.join(args.checkpoint_dir, 'last_checkpoint.pth')))
            if os.path.exists(ck_path):
                ck = torch.load(ck_path, map_location='cpu')
                unwrap_model(model).load_state_dict(ck.get('model_state', ck))
                if 'optimizer' in ck:
                    optimizer.load_state_dict(ck['optimizer'])
                if 'scaler' in ck and scaler is not None:
                    try:
                        scaler.load_state_dict(ck['scaler'])
                    except Exception:
                        pass
                if use_ema and 'ema_shadow' in ck and ck['ema_shadow'] is not None:
                    try:
                        ema_shadow = {k: v.to(device) for k, v in ck['ema_shadow'].items()}
                    except Exception:
                        pass
                if 'epoch' in ck:
                    start_epoch = int(max(1, int(ck['epoch']) + 1))
                if is_main_process(rank):
                    print(f"[resume] loaded from {ck_path} (next epoch={start_epoch})", flush=True)
            else:
                if is_main_process(rank):
                    print(f"[resume] path not found: {ck_path}, start from scratch", flush=True)
    except Exception as e:
        if is_main_process(rank):
            print(f"[resume] failed: {e}", flush=True)
    
    # 兼容不同版本的PyTorch AMP API
    scaler = None
    autocast_ctx = None
    
    # 首先检查torch.amp模块是否存在
    has_torch_amp = hasattr(torch, 'amp')
    
    try:
        # 新版本 PyTorch (1.10+): torch.amp
        if has_torch_amp:
            amp_module = getattr(torch, 'amp')
            if hasattr(amp_module, 'autocast') and hasattr(amp_module, 'GradScaler'):
                GradScaler = getattr(amp_module, 'GradScaler')
                autocast = getattr(amp_module, 'autocast')
                
                scaler = GradScaler('cuda', enabled=use_amp)
                # 尝试使用 bf16（更稳定），若不支持则回退到 fp16
                if use_amp and hasattr(torch.cuda, 'is_bf16_supported') and torch.cuda.is_bf16_supported():
                    autocast_ctx = lambda: autocast(device_type='cuda', dtype=torch.bfloat16, enabled=use_amp)
                    if is_main_process(rank):
                        print("[amp] backend=torch.amp dtype=bf16", flush=True)
                else:
                    autocast_ctx = lambda: autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp)
                    if is_main_process(rank):
                        print("[amp] backend=torch.amp dtype=fp16", flush=True)
            else:
                raise AttributeError("torch.amp autocast or GradScaler not available")
        else:
            raise AttributeError("torch.amp module not available")
    except (AttributeError, Exception):
        # 旧版本 PyTorch: torch.cuda.amp
        try:
            if hasattr(torch.cuda, 'amp'):
                cuda_amp = getattr(torch.cuda, 'amp')
                if hasattr(cuda_amp, 'GradScaler') and hasattr(cuda_amp, 'autocast'):
                    GradScaler = getattr(cuda_amp, 'GradScaler')
                    autocast = getattr(cuda_amp, 'autocast')
                    
                    scaler = GradScaler(enabled=use_amp)
                    autocast_ctx = lambda: autocast(enabled=use_amp)
                    if is_main_process(rank):
                        print("[amp] backend=torch.cuda.amp dtype=fp16", flush=True)
                else:
                    raise AttributeError("torch.cuda.amp components not available")
            else:
                raise AttributeError("torch.cuda.amp module not available")
        except Exception as e:
            print(f"[warn] AMP not available, disabling: {e}")
            use_amp = False
            # 创建一个假的scaler和context用于兼容性
            scaler = type('DummyScaler', (), {
                'scale': lambda self, loss: loss,
                'step': lambda self, optimizer: optimizer.step(),
                'update': lambda self: None,
                'unscale_': lambda self, optimizer: None
            })()
            autocast_ctx = lambda: torch.nullcontext()
            if is_main_process(rank):
                print("[amp] disabled -> dtype=fp32", flush=True)
    
    # 最后的安全检查
    if autocast_ctx is None:
        autocast_ctx = lambda: torch.nullcontext()
    if scaler is None:
        # 创建一个假的scaler用于兼容性
        scaler = type('DummyScaler', (), {
            'scale': lambda self, loss: loss,
            'step': lambda self, optimizer: optimizer.step(),
            'update': lambda self: None,
            'unscale_': lambda self, optimizer: None
        })()
    accum = max(1, args.grad_accum_steps)

    best_val = float("inf")
    best_path = os.path.join(args.checkpoint_dir, "best_model.pth")
    patience = int(args.early_stop_patience)
    min_delta = float(args.early_stop_min_delta)
    bad = 0
    improve_count = 0
    saved_any = False

    printed_shapes = False
    global_step = 0
    train_losses = []
    val_losses = []
    lrs = []
    # Warmup 配置（以迭代步为单位实现线性升温）
    base_lr_init = float(args.lr)
    try:
        steps_per_epoch = max(1, len(train_loader))
    except Exception:
        steps_per_epoch = 1
    total_warmup_steps = int(max(0, int(getattr(args, 'warmup_epochs', 0))) * steps_per_epoch)

    # bias 日志现在改为 epoch 级别汇总，不再需要预先清空

    # 梯度裁剪日志初始化
    update_count = 0
    grad_clip_log_path = str(getattr(args, 'log_grad_clip_json', ''))
    try:
        grad_clip_log_interval = int(getattr(args, 'log_grad_clip_interval', 50))
    except Exception:
        grad_clip_log_interval = 0
    grad_clip_logging_enabled = (is_main_process(rank) and isinstance(grad_clip_log_path, str)
                                 and len(grad_clip_log_path) > 0 and grad_clip_log_interval >= 0)
    grad_clip_failure_reported = False
    if grad_clip_logging_enabled:
        ensure_dir(os.path.dirname(grad_clip_log_path))
        try:
            with open(grad_clip_log_path, "w", encoding='utf-8') as f_gc_init:
                f_gc_init.write("")
        except Exception as e_gc_init:
            grad_clip_logging_enabled = False
            if is_main_process(rank):
                print(f"[grad-clip] ERROR: failed to initialize log file '{grad_clip_log_path}': {e_gc_init}", flush=True)
        else:
            if is_main_process(rank):
                if grad_clip_log_interval <= 1:
                    interval_msg = "every update"
                else:
                    interval_msg = f"every {grad_clip_log_interval} updates"
                print(f"[grad-clip] logging enabled -> {grad_clip_log_path} ({interval_msg})", flush=True)

    saved_layer_stats_epochs: Set[int] = set()
    for epoch in range(start_epoch, args.epochs + 1):
        if is_main_process(rank):
            print(f"\n===== Epoch {epoch}/{args.epochs} =====", flush=True)
        # 每个 epoch 重置 σ 分桶统计
        try:
            Kb_epoch = int(getattr(args, 'log_sigma_buckets', 8))
        except Exception:
            Kb_epoch = 0
        if Kb_epoch > 0:
            bucket_sum = [0.0] * Kb_epoch
            bucket_cnt = [0] * Kb_epoch
            bucket_sqsum = [0.0] * Kb_epoch
        # 每个 epoch 重置按"vol桶"的统计
        try:
            Vb = int(getattr(args, 'log_vol_buckets', 0))
        except Exception:
            Vb = 0
        vol_bucket_sum = [0.0] * Vb if Vb > 0 else []         # 加权loss和（与 train_loss 同口径）
        vol_bucket_cnt = [0] * Vb if Vb > 0 else []            # 样本计数（batch聚合平均）
        vol_bucket_snorm = [0.0] * Vb if Vb > 0 else []        # score范数累和（平均每样本的L2）
        # 额外记录本epoch内 s_yz 统计（用于诊断 vol_buckets.png 为0的根因）
        vol_s_min_epoch: float | None = None
        vol_s_max_epoch: float | None = None
        vol_s_mean_sum_epoch: float = 0.0
        vol_s_sample_count_epoch: int = 0
        # bias 统计汇总（epoch级别）
        bias_sigma_sum = 0.0
        bias_syz_sum = 0.0
        bias_y0_sum = 0.0
        bias_step_count = 0
        # 细致日志：记录偏差统计异常与步骤
        bias_err_count: int = 0
        bias_first_err: str | None = None
        # 本epoch使用的桶边界：支持 logspace 与 quantile 两种口径（quantile 使用训练集 yz_vol 分布的分位数，稳定且信息量更高）
        vol_bucket_edges_epoch = None
        vol_bucket_mode_used = 'logspace'
        if Vb > 0:
            mode = str(getattr(args, 'log_vol_bucket_mode', 'logspace')).lower()
            vol_bucket_mode_used = mode
            clip_u_e = float(getattr(args, 'noise_scale_clip', 10.0))
            if mode == 'quantile':
                yz_table = _load_yz_table(args)
                date_list = _get_dataset_dates(train_set, 'train-vol-bucket')
                stocks_here = _get_dataset_stocks(train_set, 'train-vol-bucket')
                sbar_values: List[float] = []
                for dt in date_list:
                    row = yz_table.loc[dt]
                    vals = np.asarray([row[str(s)] for s in stocks_here], dtype=np.float64)
                    if not np.all(np.isfinite(vals)):
                        raise ValueError(f"[train-vol-bucket] 日期 {dt} 存在非有限 yz_vol: {vals}")
                    vals = np.clip(vals, 1e-6, clip_u_e)
                    sbar_values.append(float(vals.mean()))
                if len(sbar_values) == 0:
                    raise ValueError("[train-vol-bucket] 无法从 yz_vol 表中计算任何分位数")
                edges_q = np.quantile(np.asarray(sbar_values, dtype=np.float64), np.linspace(0.0, 1.0, Vb + 1))
                for i in range(1, edges_q.size):
                    if edges_q[i] <= edges_q[i - 1]:
                        edges_q[i] = edges_q[i - 1] + 1e-6
                vol_bucket_edges_epoch = edges_q
            else:
                vol_bucket_edges_epoch = np.exp(np.linspace(np.log(1e-6), np.log(clip_u_e), Vb + 1))
        # 分桶调试计数与首个异常信息（定位 cnt_total=0 的原因）
        vol_bucket_try_batches: int = 0
        vol_bucket_succ_batches: int = 0
        vol_bucket_first_err: str | None = None
        # 每个 epoch 重置按 t 的梯度范数统计（JSONL 输出）
        enable_tgrad = False
        try:
            tgrad_path = str(getattr(args, 'log_t_grad_json', ''))
            enable_tgrad = isinstance(tgrad_path, str) and len(tgrad_path) > 0
        except Exception:
            enable_tgrad = False
        if enable_tgrad:
            Tdiff_cfg_epoch = int(getattr(args, 'diff_steps', 1000)) if hasattr(args, 'diff_steps') else 1000
            t_grad_sum_ep = torch.zeros(Tdiff_cfg_epoch, dtype=torch.float64)
            t_grad_cnt_ep = torch.zeros(Tdiff_cfg_epoch, dtype=torch.int64)
        model.train()
        epoch_t0 = time.time()
        running_loss = 0.0
        running_count = 0
        bound_sum_epoch = 0.0
        bound_count_epoch = 0
        graph_reg_sum_epoch = 0.0
        optimizer.zero_grad(set_to_none=True)
        pbar = None
        if is_main_process(rank) and args.progress == 1 and tqdm is not None:
            pbar = tqdm(total=len(train_loader), desc=f"epoch {epoch}/{args.epochs}", leave=False)
        # DDP 采样器设定 epoch，保证各进程一致洗牌
        if isinstance(train_loader.sampler, DistributedSampler):
            train_loader.sampler.set_epoch(epoch)
        # 累计整轮平均 DiffMSE（ε̂ vs ε），用于日志
        diffmse_sum: float = 0.0
        diffmse_count: int = 0
        steps_in_epoch: int = 0
        model_main = unwrap_model(model)
        graph_params_all: List[nn.Parameter] = []
        other_params_all: List[nn.Parameter] = []
        if hasattr(model_main, "graph_learner"):
            graph_params_all = [p for p in model_main.graph_learner.parameters() if p.requires_grad]
            graph_param_ids = {id(p) for p in graph_params_all}
            other_params_all = [
                p for p in model_main.parameters()
                if p.requires_grad and id(p) not in graph_param_ids
            ]
        else:
            other_params_all = [p for p in model_main.parameters() if p.requires_grad]

        layer_stats_capture = False
        layer_stats_path_epoch: str = ""
        layer_stats_snapshot: Optional[dict] = None
        layer_stats_interval_epoch = 1
        if is_main_process(rank):
            try:
                layer_stats_enabled_flag = int(getattr(args, 'log_layer_stats', 1))
                layer_stats_interval_epoch = max(1, int(getattr(args, 'log_layer_stats_interval', 1)))
                layer_stats_path_candidate = str(getattr(args, 'log_layer_stats_path', 'plots/train/layer_stats.jsonl'))
                if layer_stats_enabled_flag > 0 and isinstance(layer_stats_path_candidate, str) and len(layer_stats_path_candidate) > 0:
                    if epoch % layer_stats_interval_epoch == 0:
                        ensure_dir(os.path.dirname(layer_stats_path_candidate))
                        layer_stats_capture = True
                        layer_stats_path_epoch = layer_stats_path_candidate
            except Exception as e_layer_init:
                if is_main_process(rank):
                    print(f"[layer_stats] WARNING: init failed at epoch {epoch}: {e_layer_init}", flush=True)
                layer_stats_capture = False
        for it, batch in enumerate(train_loader, 1):
            steps_in_epoch += 1
            # 线性 Warmup：将 lr 从 base_lr*warmup_init_factor 线性提升到 base_lr
            if str(getattr(args, 'scheduler_type', 'plateau')).lower() == 'sgdr' and total_warmup_steps > 0 and global_step < total_warmup_steps:
                try:
                    p0 = float(max(0.0, min(1.0, float(getattr(args, 'warmup_init_factor', 0.1)))))
                except Exception:
                    p0 = 0.1
                prog = float(global_step + 1) / float(total_warmup_steps)
                factor = p0 + (1.0 - p0) * prog
                for g in optimizer.param_groups:
                    g['lr'] = base_lr_init * float(factor)
            if isinstance(batch, (list, tuple)) and len(batch) == 4:
                X, Y, TIDX, mask = batch  # [mask机制] 解包mask
            elif isinstance(batch, (list, tuple)) and len(batch) == 3:
                X, Y, TIDX = batch
                mask = None  # 兼容旧数据
            else:
                X, Y = batch  # 兼容旧数据集
                # 若未提供绝对时间索引，采用当前 enumerate 的步号近似（不影响随机迭代的可加性）
                TIDX = torch.arange((it - 1) * X.size(0), it * X.size(0), dtype=torch.long)
                mask = None
            X = X.to(device, non_blocking=pin)
            Y = Y.to(device, non_blocking=pin)
            t_abs = torch.as_tensor(TIDX, device=device, dtype=torch.long)
            if mask is not None:
                mask = mask.to(device, non_blocking=pin)  # [mask机制] 转移到设备
            
            # ========== 自适应谱噪声：破坏低频条件输入 ==========
            sensitivity_sampling = False
            # 方式1：使用预计算的alpha（优先）
            if spectral_precomputed_data is not None:
                try:
                    alpha_per_feature = spectral_precomputed_data["alpha_per_feature"].to(device)  # [F]
                    input_noise_sigma = float(getattr(args, 'sigma_max', 1.0)) * 0.1
                    
                    # 直接应用噪声：X_noisy = X + sigma * (noise * alpha)
                    noise = torch.randn_like(X)  # [B, N, T, F]
                    noise_scaled = noise * alpha_per_feature.view(1, 1, 1, -1)  # 广播alpha到[B,N,T,F]
                    X = X + input_noise_sigma * noise_scaled
                    
                    # 敏感度采样
                    sensitivity_interval = int(getattr(args, 'check_sensitivity_interval', getattr(args, 'spectral_sensitivity_interval', 100)))
                    if sensitivity_interval > 0 and it % sensitivity_interval == 0:
                        sensitivity_sampling = True
                        X.requires_grad_(True)
                except Exception as e_spec_apply:
                    if is_main_process(rank):
                        print(f"[spectral] ERROR: 应用预计算噪声失败 (epoch={epoch}, it={it}): {e_spec_apply}", flush=True)
                    raise RuntimeError(f"自适应谱噪声应用失败: {e_spec_apply}") from e_spec_apply
            
            # 方式2：使用在线计算的tracker（回退）
            elif spectral_noise_tracker is not None and spectral_noise_tracker.last_stats is not None:
                try:
                    input_noise_sigma = float(getattr(args, 'sigma_max', 1.0)) * 0.1
                    X = spectral_noise_tracker.apply(X, torch.tensor(input_noise_sigma, device=X.device))
                    
                    sensitivity_interval = int(getattr(args, 'check_sensitivity_interval', getattr(args, 'spectral_sensitivity_interval', 100)))
                    if sensitivity_interval > 0 and it % sensitivity_interval == 0:
                        sensitivity_sampling = True
                        X.requires_grad_(True)
                except Exception as e_spec_apply:
                    if is_main_process(rank):
                        print(f"[spectral] ERROR: 应用在线计算噪声失败 (epoch={epoch}, it={it}): {e_spec_apply}", flush=True)
                    raise RuntimeError(f"自适应谱噪声应用失败: {e_spec_apply}") from e_spec_apply
            
            with autocast_ctx():
                features = model.encode(X, mask=mask)  # [mask机制] 传递mask
                
                # 计算敏感度（如果本次采样）
                if sensitivity_sampling:
                    try:
                        # 先更新全量统计（不影响断言）
                        if spectral_noise_tracker is not None:
                            spectral_noise_tracker.update_sensitivity(X, features)
                        # 特征级断言检查
                        if int(getattr(args, 'check_sensitivity', 1)) == 1:
                            idx_str = str(getattr(args, 'check_sensitivity_features', '')).strip()
                            if len(idx_str) > 0:
                                # 解析索引
                                parts = [p for p in idx_str.replace(' ', '').split(',') if len(p) > 0]
                                feat_indices = [int(p) for p in parts]
                                F_here = int(X.shape[-1])
                                # 过滤非法索引
                                valid_indices = [i for i in feat_indices if 0 <= i < F_here]
                                if len(valid_indices) != len(feat_indices):
                                    raise ValueError(f"check_sensitivity_features 含非法索引(总特征维={F_here}): {feat_indices}")
                                # 计算梯度范数（对输出范数求导）
                                output_norm = torch.sum(features ** 2)
                                grads = torch.autograd.grad(outputs=output_norm, inputs=X, create_graph=False, retain_graph=True, only_inputs=True)
                                grad_x = grads[0]  # [B,N,T,F]
                                thr = float(getattr(args, 'check_sensitivity_threshold', 0.01))
                                for i_feat in valid_indices:
                                    g = grad_x[..., i_feat]
                                    sens = float(torch.sqrt(torch.mean(g * g)).item())
                                    if not (sens < thr):
                                        raise AssertionError(f"特征{i_feat}敏感度未降至安全阈值 (sens={sens:.6g} >= {thr})")
                    except Exception as _se:
                        # 直接抛出以终止训练，符合断言语义
                        raise
                B = X.size(0)
                batch_size = int(B)
                pred_y_for_bound: Optional[torch.Tensor] = None
                if is_main_process(rank) and args.print_shapes_once and (not printed_shapes):
                    print(f"train X={tuple(X.shape)} Y={tuple(Y.shape)} feat={tuple(features.shape)}", flush=True)
                    printed_shapes = True
                # 目标模式：eps（推荐）或 score（保留兼容）
                objective = "edm"
                
                if objective == 'edm':
                    # ========== EDM 预条件化目标（VE 家族最佳实践）==========
                    # 直接预测去噪后的 y0，避免均值塌缩，支持跨 σ 层级的稳定训练
                    B, N = Y.shape[:2]
                    
                    # σ_data：训练集 log-return 的经验标准差（从 config 读取或自动估计）
                    sigma_data = float(getattr(args, 'sigma_data', 0.025))
                    
                    # 采样 σ：log-uniform（与 eps 分支一致）
                    smin = float(getattr(args, 'sigma_min', 0.01))
                    smax = float(getattr(args, 'sigma_max', 1.0))
                    log_sigma = torch.rand(B, device=X.device) * (math.log(smax) - math.log(smin)) + math.log(smin)
                    sigma = torch.exp(log_sigma).view(B, 1, 1)  # [B,1,1]
                    
                    # 采样噪声 ε~N(0,1)
                    eps = torch.randn_like(Y)  # [B,N,1]
                    
                    # 构造 x_t
                    y0 = Y
                    x_t = y0 + sigma * eps
                    
                    # 时间索引映射（用于模型的时间嵌入）
                    Tdiff = int(getattr(args, 'diff_steps', 1000))
                    t_idx = torch.clamp(((log_sigma - math.log(smin)) / max(1e-9, (math.log(smax) - math.log(smin))) * Tdiff).long() + 1, 1, Tdiff)
                    
                    # EDM 预条件化系数
                    sigma_data_sq = sigma_data ** 2
                    sigma_sq = sigma ** 2
                    c_skip = sigma_data_sq / (sigma_sq + sigma_data_sq)  # [B,1,1]
                    c_out = sigma * sigma_data / torch.sqrt(sigma_sq + sigma_data_sq)  # [B,1,1]
                    c_in = 1.0 / torch.sqrt(sigma_sq + sigma_data_sq)  # [B,1,1]
                    
                    # 模型输入：预条件化的 x_t
                    x_scaled = c_in * x_t  # [B,N,1]
                    
                    # 显式构造波动率条件（若启用 yz_vol）
                    scale_mode_edm = str(getattr(args, 'noise_scale_mode', 'constant')).lower()
                    if scale_mode_edm == 'yz_vol':
                        yz_table = _load_yz_table(args)
                        date_list = _get_dataset_dates(train_set, 'train-edm-vol-cond')
                        stocks_here = _get_dataset_stocks(train_set, 'train-edm-vol-cond')
                        vol_tensor = _build_vol_tensor(
                            yz_table,
                            date_list,
                            stocks_here,
                            t_abs,
                            float(getattr(args, 'noise_scale_clip', 10.0)),
                            X.device,
                            X.dtype,
                            'train-edm-vol-cond',
                        )
                        vol_cond = vol_tensor.unsqueeze(-1)
                    else:
                        vol_cond = None
                    
                    # 获取 yzvol 开关
                    vol_as_input = bool(getattr(args, 'vol_cond_as_input', 1))
                    vol_scale_output = bool(getattr(args, 'vol_cond_scale_output', 0))
                    
                    # 通过 score 头获取网络输出 fθ
                    f_theta = model.forward_score(
                        X, x_scaled, t_idx, 
                        vol_cond=vol_cond,
                        vol_cond_as_input=vol_as_input,
                        vol_cond_scale_output=vol_scale_output
                    )  # [B,N,1]
                    
                    # 诊断: 第一个 epoch 第一个 batch 打印 vol_cond 统计（验证加载是否正确）
                    if is_main_process(rank) and epoch == 1 and it == 1 and vol_cond is not None:
                        vol_np = vol_cond.detach().float().cpu().numpy().squeeze(-1)
                        mode_str = f"as_input={vol_as_input} scale_output={vol_scale_output}"
                        print(f"[yzvol-train] {mode_str} | vol_cond: mean={vol_np.mean():.3f} std={vol_np.std():.3f} range=[{vol_np.min():.3f}, {vol_np.max():.3f}] (should be mean≈1.0)", flush=True)
                    
                    # 组装 y0_hat（EDM 预条件化输出，移除 τ）
                    y0_hat = c_skip * x_t + c_out * f_theta  # [B,N,1]
                    
                    # 损失：MSE(y0_hat, y0)，保留 mask 覆盖权重
                    per_elem = (y0_hat - y0) ** 2  # [B,N,1]
                    # 可选：EDM λ(σ) 加权（提高小σ权重），支持幂次与批内归一化
                    try:
                        if int(getattr(args, 'edm_lambda', 0)) == 1:
                            lam = (sigma_sq + sigma_data_sq) / torch.clamp((sigma * sigma_data), min=1e-12) ** 2  # [B,1,1]
                            pow_k = float(getattr(args, 'edm_lambda_pow', 1.0))
                            if abs(pow_k - 1.0) > 1e-8:
                                lam = torch.pow(lam, pow_k)
                            norm_mode = str(getattr(args, 'edm_lambda_norm', 'batch_mean')).lower()
                            if norm_mode == 'batch_mean':
                                lam = lam / lam.mean()
                            per_elem = per_elem * lam
                    except Exception:
                        pass
                    
                    # [mask机制] 样本权重（覆盖率<阈值的样本权重为0）
                    sample_weight = None
                    if mask is not None:
                        coverage = mask.float().mean(dim=(2,3))  # [B,N]
                        threshold = float(getattr(args, 'sample_coverage_threshold', 0.5))
                        sample_weight = (coverage >= threshold).float().unsqueeze(-1)  # [B,N,1]
                        
                        # [覆盖率分析] 记录覆盖率值
                        log_coverage_path = getattr(args, 'log_coverage_path', '')
                        if is_main_process(rank) and log_coverage_path:
                            with open(log_coverage_path, 'a') as f:
                                np.savetxt(f, coverage.flatten().cpu().numpy(), fmt='%.6f')
                        
                        # 加权损失
                        weighted = per_elem * sample_weight
                        valid_count = sample_weight.sum().clamp_min(1.0)
                        loss = weighted.sum() / valid_count
                    else:
                        loss = per_elem.mean()
                    
                    # 可选：EDM 推荐的权重 λ(σ) = (σ^2+σ_data^2)/(σ·σ_data)^2
                    # 暂不启用，先验证基础版本效果
                    
                    # 按 σ 分桶统计（与 score 分支保持对齐），用于写入 sigma_stats.jsonl 的训练期 buckets
                    try:
                        Kb = int(getattr(args, 'log_sigma_buckets', 8))
                        if Kb > 0:
                            lb_global = math.log(smin)
                            ub_global = math.log(smax)
                            width_global = max(1e-9, (ub_global - lb_global) / max(1, Kb))
                            # 使用 batch 聚合的损失标量近似样本损失（与 score 分支口径一致）
                            diff_mse_here = loss
                            if not isinstance(diff_mse_here, torch.Tensor):
                                diff_mse_here = torch.tensor(float(diff_mse_here))
                            log_s = log_sigma.detach()
                            for i in range(B):
                                bidx = int(min(Kb - 1, max(0, math.floor((float(log_s[i].item()) - lb_global) / width_global))))
                                if 'bucket_sum' not in locals():
                                    bucket_sum = [0.0] * Kb
                                    bucket_cnt = [0] * Kb
                                    bucket_sqsum = [0.0] * Kb
                                v = float(diff_mse_here.detach().float().cpu().item())
                                bucket_sum[bidx] += v
                                bucket_cnt[bidx] += 1
                                bucket_sqsum[bidx] += v * v
                    except Exception:
                        pass

                    # 按 t 统计（与 score 分支保持对齐，用于诊断不同时间步的损失）
                    with torch.no_grad():
                        if 't_loss_sum' not in locals():
                            Tdiff_cfg = int(getattr(args, 'diff_steps', 1000)) if hasattr(args, 'diff_steps') else 1000
                            t_loss_sum = torch.zeros(Tdiff_cfg, dtype=torch.float64)
                            t_loss_cnt = torch.zeros(Tdiff_cfg, dtype=torch.int64)
                        Tdiff_cfg = t_loss_sum.numel()
                        # 统一使用 CPU 做累计，避免设备不一致
                        t_idx_vec = t_idx.view(-1).detach().cpu()
                        per_sample = per_elem.view(B, -1).mean(dim=1).detach().double().cpu()
                        for i in range(B):
                            ti = int(max(1, min(Tdiff_cfg, int(t_idx_vec[i].item())))) - 1
                            t_loss_sum[ti] += per_sample[i]
                            t_loss_cnt[ti] += 1
                    
                    # EDM模式：定义必要变量以兼容后续损失组装
                    base_loss = loss
                    diffmse_sum += float(loss.detach().float().cpu().item()) * X.size(0)
                    diffmse_count += int(X.size(0))
                    pred_y_for_bound = y0_hat
                    
                    # 梯度累积：loss 需除以累积步数
                    loss = loss / float(accum)
                
                # EDM-only: score/eps 分支已彻底移除
                # 越界惩罚损失（仅保留主损失 + 越界约束）
                bound_w = float(getattr(args, 'reg_bound_w', 0.0))
                if bound_w > 0.0 and pred_y_for_bound is not None:
                    upper = torch.tensor(args.bound_log_upper, device=pred_y_for_bound.device, dtype=pred_y_for_bound.dtype)
                    lower = torch.tensor(args.bound_log_lower, device=pred_y_for_bound.device, dtype=pred_y_for_bound.dtype)
                    over = torch.clamp(pred_y_for_bound - upper, min=0.0)
                    under = torch.clamp(lower - pred_y_for_bound, min=0.0)
                    bound_loss = (over * over + under * under).mean()
                    bound_loss_value = float(bound_loss.detach().float().cpu().item())
                else:
                    bound_loss = torch.tensor(0.0, device=X.device)
                    bound_loss_value = 0.0
                bound_sum_epoch += bound_loss_value * batch_size
                bound_count_epoch += batch_size
                graph_reg_w = float(getattr(args, 'graph_l2_reg', 0.0))
                if graph_reg_w > 0.0:
                    if not hasattr(model_main, "graph_learner"):
                        raise RuntimeError("graph_l2_reg>0 但模型缺少 graph_learner")
                    raw_mean_tensor = getattr(model_main.graph_learner, "last_adj_raw_abs_mean", None)
                    if raw_mean_tensor is None:
                        raise RuntimeError("graph_l2_reg>0 但图分支未提供 last_adj_raw_abs_mean")
                    graph_reg_tensor = torch.clamp(raw_mean_tensor - 0.5, min=0.0)
                    graph_reg_value = float(graph_reg_tensor.detach().float().cpu().item())
                else:
                    graph_reg_tensor = base_loss.new_tensor(0.0)
                    graph_reg_value = 0.0
                graph_reg_sum_epoch += graph_reg_value * batch_size

                loss = (base_loss + bound_w * bound_loss + graph_reg_w * graph_reg_tensor) / accum
            # 非有限检查（数值护栏）
            if not torch.isfinite(loss):
                try:
                    b = float(base_loss.detach().float().cpu()) if torch.isfinite(base_loss) else float('nan')
                    bd = float(bound_loss.detach().float().cpu()) if torch.isfinite(bound_loss) else float('nan')
                    if is_main_process(rank):
                        print(f"[warn] non-finite loss detected: base={b:.4g} bound={bd:.4g} | skip step", flush=True)
                except Exception:
                    if is_main_process(rank):
                        print("[warn] non-finite loss detected: skip step", flush=True)
                optimizer.zero_grad(set_to_none=True)
                continue
            scaler.scale(loss).backward()
            if it % accum == 0:
                # 先反缩放再裁剪，避免 AMP 缩放导致的数值问题
                # 不使用try-except，unscale失败应立即报错（否则梯度裁剪无效）
                scaler.unscale_(optimizer)
                graph_grad_norm = compute_grad_norm_l2(graph_params_all) if graph_params_all else 0.0
                other_grad_norm = compute_grad_norm_l2(other_params_all) if other_params_all else 0.0
                if graph_params_all and graph_grad_norm < 1e-7 and other_grad_norm > 1e-5:
                    for param in graph_params_all:
                        if param.grad is not None:
                            param.grad.mul_(2.0)
                total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                if layer_stats_capture:
                    records, snapshot_total_norm, max_grad_entry = collect_layer_grad_snapshot(unwrap_model(model))
                    layer_stats_snapshot = {
                        "records": records,
                        "total_grad_norm": snapshot_total_norm,
                        "max_grad": float(max_grad_entry["grad_norm"]) if max_grad_entry is not None else 0.0,
                        "max_grad_layer": str(max_grad_entry["name"]) if max_grad_entry is not None else "",
                    }
                    if epoch <= 5 and epoch not in saved_layer_stats_epochs:
                        try:
                            snapshot_dir = os.path.join('plots', 'train', 'epoch_snapshots')
                            ensure_dir(snapshot_dir)
                            torch.save(layer_stats_snapshot, os.path.join(snapshot_dir, f"layer_stats_epoch{epoch:02d}.pt"))
                            saved_layer_stats_epochs.add(epoch)
                        except Exception as e_snapshot:
                            if is_main_process(rank):
                                print(f"[layer_stats] WARNING: 保存 epoch {epoch} 快照失败: {e_snapshot}", flush=True)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)
                update_count += 1
                if grad_clip_logging_enabled and (update_count == 1 or grad_clip_log_interval == 0 or update_count % max(1, grad_clip_log_interval) == 0):
                    payload = {
                        "epoch": int(epoch),
                        "update": int(update_count),
                        "batch": int(it),
                        "global_step": int(global_step),
                        "total_norm": float(total_norm),
                        "max_norm": 1.0,
                        "tail_flush": False,
                    }
                    try:
                        with open(grad_clip_log_path, "a", encoding="utf-8") as f_gc:
                            f_gc.write(json.dumps(payload, ensure_ascii=False) + "\n")
                    except Exception as e_gc_append:
                        if not grad_clip_failure_reported and is_main_process(rank):
                            print(f"[grad-clip] ERROR: failed to append log (epoch={epoch}, update={update_count}): {e_gc_append}", flush=True)
                        grad_clip_failure_reported = True
                        grad_clip_logging_enabled = False
                # SGDR：Warmup结束后再推进余弦重启（按迭代步）
                if scheduler_sgdr is not None:
                    frac = (it / max(1, len(train_loader)))
                    pos = (epoch - 1) + frac - float(getattr(args, 'warmup_epochs', 0))
                    if pos >= 0.0 and not (total_warmup_steps > 0 and global_step <= total_warmup_steps):
                        scheduler_sgdr.step(pos)  # type: ignore[arg-type]
            global_step += 1
            running_loss += (loss.item() * accum) * X.size(0)
            running_count += X.size(0)
            if pbar is not None:
                pbar.update(1)
                pbar.set_postfix({"loss": f"{(loss.item()*accum):.4f}"})
            # (旧累加器逻辑已移除)

        # 汇总所有进程的训练损失分子/分母
        train_epoch_loss = running_loss / max(1, running_count)
        # 处理累积步尾部残留：确保所有 batch 都执行 clip+step
        if steps_in_epoch % accum != 0:
            # 不使用try-except，unscale失败应立即报错（否则梯度裁剪无效）
            scaler.unscale_(optimizer)
            graph_grad_norm_tail = compute_grad_norm_l2(graph_params_all) if graph_params_all else 0.0
            other_grad_norm_tail = compute_grad_norm_l2(other_params_all) if other_params_all else 0.0
            if graph_params_all and graph_grad_norm_tail < 1e-7 and other_grad_norm_tail > 1e-5:
                for param in graph_params_all:
                    if param.grad is not None:
                        param.grad.mul_(2.0)
            tail_total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            if layer_stats_capture:
                records, snapshot_total_norm, max_grad_entry = collect_layer_grad_snapshot(unwrap_model(model))
                layer_stats_snapshot = {
                    "records": records,
                    "total_grad_norm": snapshot_total_norm,
                    "max_grad": float(max_grad_entry["grad_norm"]) if max_grad_entry is not None else 0.0,
                    "max_grad_layer": str(max_grad_entry["name"]) if max_grad_entry is not None else "",
                }
                if epoch <= 5 and epoch not in saved_layer_stats_epochs:
                    try:
                        snapshot_dir = os.path.join('plots', 'train', 'epoch_snapshots')
                        ensure_dir(snapshot_dir)
                        torch.save(layer_stats_snapshot, os.path.join(snapshot_dir, f"layer_stats_epoch{epoch:02d}.pt"))
                        saved_layer_stats_epochs.add(epoch)
                    except Exception as e_snapshot:
                        if is_main_process(rank):
                            print(f"[layer_stats] WARNING: 保存 epoch {epoch} 快照失败: {e_snapshot}", flush=True)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)
            update_count += 1
            if grad_clip_logging_enabled and (update_count == 1 or grad_clip_log_interval == 0 or update_count % max(1, grad_clip_log_interval) == 0):
                payload_tail = {
                    "epoch": int(epoch),
                    "update": int(update_count),
                    "batch": int(steps_in_epoch),
                    "global_step": int(global_step),
                    "total_norm": float(tail_total_norm),
                    "max_norm": 1.0,
                    "tail_flush": True,
                }
                try:
                    with open(grad_clip_log_path, "a", encoding="utf-8") as f_gc:
                        f_gc.write(json.dumps(payload_tail, ensure_ascii=False) + "\n")
                except Exception as e_gc_tail:
                    if not grad_clip_failure_reported and is_main_process(rank):
                        print(f"[grad-clip] ERROR: failed to append tail log (epoch={epoch}): {e_gc_tail}", flush=True)
                    grad_clip_failure_reported = True
                    grad_clip_logging_enabled = False
        # σ 分桶日志（控制台与可选JSONL），输出样本占比与桶内 loss 方差
        try:
            Kb = int(getattr(args, 'log_sigma_buckets', 8))
            out_json = str(getattr(args, 'log_sigma_json', ''))
            if 'bucket_sum' in locals() and Kb > 0:
                # 估计每桶的方差（简单近似：var ≈ E[x^2] - (E[x])^2，使用batch级别的 diff_mse_here 近似样本）
                if 'bucket_sqsum' in locals():
                    stats = [
                        (
                            i,
                            (bucket_sum[i] / max(1, bucket_cnt[i])),
                            max(0.0, (bucket_sqsum[i] / max(1, bucket_cnt[i])) - (bucket_sum[i] / max(1, bucket_cnt[i]))**2),
                            bucket_cnt[i]
                        ) for i in range(Kb)
                    ]
                else:
                    stats = [ (i, (bucket_sum[i] / max(1, bucket_cnt[i])), 0.0, bucket_cnt[i]) for i in range(Kb) ]
                if is_main_process(rank):
                    # print(f"[sigma-buckets] K={Kb} (bucket, avg, var, count)={stats}", flush=True)
                    if isinstance(out_json, str) and len(out_json) > 0:
                        try:
                            ensure_dir(os.path.dirname(out_json))
                            payload = {
                                "epoch": int(epoch),
                                "buckets": []
                            }
                            for i in range(Kb):
                                avg = float(bucket_sum[i] / max(1, bucket_cnt[i]))
                                var = float(max(0.0, (bucket_sqsum[i] / max(1, bucket_cnt[i])) - avg*avg)) if 'bucket_sqsum' in locals() else 0.0
                                cnt = int(bucket_cnt[i])
                                frac = float(cnt / max(1, sum(bucket_cnt)))
                                payload["buckets"].append({"bucket": i, "avg": avg, "var": var, "count": cnt, "frac": frac})
                            with open(out_json, "a", encoding="utf-8") as f:
                                f.write(json.dumps(payload, ensure_ascii=False) + "\n")
                        except Exception as _e:
                            print(f"[sigma-buckets] write json failed: {_e}", flush=True)
        except Exception:
            pass

        # 绘制按 t 的 loss 曲线（每epoch覆盖保存）
        try:
            tloss_png = str(getattr(args, 'log_t_loss_png', ''))
            tloss_json = str(getattr(args, 'log_t_loss_json', ''))
            if isinstance(tloss_png, str) and len(tloss_png) > 0:
                import matplotlib
                matplotlib.use("Agg")
                import matplotlib.pyplot as plt
                ensure_dir(os.path.dirname(tloss_png))
                Tdiff_cfg = int(getattr(args, 'diff_steps', 1000)) if hasattr(args, 'diff_steps') else 1000
                t_axis = np.arange(1, Tdiff_cfg + 1)
                tl_avg = np.zeros(Tdiff_cfg, dtype=np.float64)
                if 't_loss_cnt' in locals() and 't_loss_sum' in locals():
                    cnt_np = t_loss_cnt.cpu().numpy()
                    sum_np = t_loss_sum.cpu().numpy()
                    mask = cnt_np > 0
                    tl_avg[mask] = sum_np[mask] / cnt_np[mask]
                    nonzero_bins = int(mask.sum())
                    total_samples = int(cnt_np[mask].sum()) if nonzero_bins > 0 else 0
                    if nonzero_bins > 0:
                        nonzero_idx = np.where(mask)[0]
                        avg_min = float(tl_avg[mask].min())
                        avg_max = float(tl_avg[mask].max())
                        avg_mean = float(tl_avg[mask].mean())
                        first_idxs = (nonzero_idx[:10] + 1).tolist()
                    else:
                        avg_min = avg_max = avg_mean = float('nan')
                        first_idxs = []
                    if isinstance(tloss_json, str) and len(tloss_json) > 0:
                        try:
                            dir_path = os.path.dirname(tloss_json)
                            if len(dir_path) > 0:
                                ensure_dir(dir_path)
                            payload_tloss = {
                                "epoch": int(epoch),
                                "nonzero_bins": nonzero_bins,
                                "total_samples": total_samples,
                                "t_idx_min": int(nonzero_idx[0] + 1) if nonzero_bins > 0 else None,
                                "t_idx_max": int(nonzero_idx[-1] + 1) if nonzero_bins > 0 else None,
                                "avg_min": avg_min,
                                "avg_max": avg_max,
                                "avg_mean": avg_mean,
                                "first_t_indices": first_idxs,
                            }
                            with open(tloss_json, "a", encoding="utf-8") as f_json:
                                f_json.write(json.dumps(payload_tloss, ensure_ascii=False) + "\n")
                        except Exception as _e_tjson:
                            print(f"[t-loss] write json failed: {_e_tjson}", flush=True)
                plt.figure(figsize=(7,4))
                plt.plot(t_axis, tl_avg, label=f"epoch {epoch}")
                plt.xlabel("t (1..Tdiff)")
                plt.ylabel("L_t (avg per-sample)")
                plt.title("Per-t loss curve")
                plt.grid(True, alpha=0.3)
                plt.tight_layout()
                plt.savefig(tloss_png, dpi=150)
                plt.close()
        except Exception:
            pass
        # 写出按 t 的梯度范数 JSONL（每个 epoch 追加一行），并附加 vol 桶统计
        try:
            if enable_tgrad and is_main_process(rank):
                ensure_dir(os.path.dirname(tgrad_path))
                Tdiff_cfg = t_grad_sum_ep.numel()
                avg_list = []
                for i in range(Tdiff_cfg):
                    cnt_i = int(t_grad_cnt_ep[i].item())
                    if cnt_i > 0:
                        avg_i = float(t_grad_sum_ep[i].item() / cnt_i)
                    else:
                        avg_i = 0.0
                    avg_list.append({"t": i + 1, "avg": avg_i, "cnt": cnt_i})
                payload = {"epoch": int(epoch), "grads": avg_list}
                # 追加 vol 桶统计（均为当前epoch聚合平均）
                if Vb > 0 and is_main_process(rank):
                    vol_loss_avg = [ (vol_bucket_sum[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb) ]
                    vol_snorm_avg = [ (vol_bucket_snorm[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb) ]
                    payload["vol_buckets_cnt"] = vol_bucket_cnt
                    payload["vol_buckets_loss_avg"] = vol_loss_avg
                    payload["vol_buckets_snorm_avg"] = vol_snorm_avg
                    try:
                        edges_for_json = vol_bucket_edges_epoch if (isinstance(vol_bucket_edges_epoch, np.ndarray) and vol_bucket_edges_epoch.size == (Vb + 1)) else np.exp(np.linspace(np.log(1e-6), np.log(float(getattr(args,'noise_scale_clip',10.0))), Vb + 1))
                        payload["vol_edges"] = edges_for_json.tolist()
                    except Exception:
                        payload["vol_edges"] = []
                with open(tgrad_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(payload, ensure_ascii=False) + "\n")
        except Exception:
            pass
        # 覆盖保存 "按 vol 桶"的 PNG
        try:
            if args.log_png and is_main_process(rank) and int(getattr(args, 'log_vol_plot_each_epoch', 0)) == 1:
                vpng = str(getattr(args, 'log_vol_png', ''))
                if isinstance(vpng, str) and len(vpng) > 0 and Vb > 0:
                    import matplotlib
                    matplotlib.use("Agg")
                    import matplotlib.pyplot as plt
                    ensure_dir(os.path.dirname(vpng))
                    vol_loss_avg = [ (vol_bucket_sum[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb) ]
                    vol_snorm_avg = [ (vol_bucket_snorm[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb) ]
                    edges = vol_bucket_edges_epoch if (isinstance(vol_bucket_edges_epoch, np.ndarray) and vol_bucket_edges_epoch.size == (Vb + 1)) else np.exp(np.linspace(np.log(1e-6), np.log(float(getattr(args,'noise_scale_clip',10.0))), Vb + 1))
                    centers = 0.5 * (edges[:-1] + edges[1:])
                    plt.figure(figsize=(8,4))
                    ax1 = plt.gca()
                    ax1.plot(centers, vol_loss_avg, label='loss_avg', color='tab:blue')
                    ax1.set_xscale('log')
                    ax1.set_xlabel('yz_vol_scale (bucket center, log)')
                    ax1.set_ylabel('avg weighted loss')
                    ax2 = ax1.twinx()
                    ax2.plot(centers, vol_snorm_avg, label='score_norm_avg', color='tab:orange')
                    ax2.set_xscale('log')
                    ax2.set_ylabel('avg ||score||_2')
                    lines, labels = ax1.get_legend_handles_labels()
                    lines2, labels2 = ax2.get_legend_handles_labels()
                    ax1.legend(lines + lines2, labels + labels2, loc='best')
                    plt.title(f'Vol-bucket stats (epoch {epoch})')
                    plt.tight_layout()
                    plt.savefig(vpng, dpi=150)
                    plt.close()
                    # 追加调试输出，帮助定位"为0"的原因
                    try:
                        cnt_total = int(sum(vol_bucket_cnt))
                        nonzero = int(sum(1 for c in vol_bucket_cnt if c > 0))
                        s_min_ep = (float(vol_s_min_epoch) if vol_s_min_epoch is not None else float('nan'))
                        s_max_ep = (float(vol_s_max_epoch) if vol_s_max_epoch is not None else float('nan'))
                        s_mean_ep = (float(vol_s_mean_sum_epoch / max(1, vol_s_sample_count_epoch)) if vol_s_sample_count_epoch > 0 else float('nan'))
                        dbg_extra = ""
                        try:
                            dbg_extra = f" tries={vol_bucket_try_batches} succ={vol_bucket_succ_batches}"
                            if isinstance(vol_bucket_first_err, str) and len(vol_bucket_first_err) > 0:
                                dbg_extra += f" first_err={vol_bucket_first_err}"
                        except Exception:
                            pass
                        try:
                            e0 = float(edges[0]) if hasattr(edges, '__len__') and len(edges) > 0 else float('nan')
                            eL = float(edges[-1]) if hasattr(edges, '__len__') and len(edges) > 0 else float('nan')
                        except Exception:
                            e0 = float('nan'); eL = float('nan')
                        print(f"[vol-bucket][epoch {epoch}] cnt_total={cnt_total} nonzero_buckets={nonzero}/{Vb} s_mean≈{s_mean_ep:.4g} s_min≈{s_min_ep:.4g} s_max≈{s_max_ep:.4g} mode={vol_bucket_mode_used} edges≈[{e0:.4g}..{eL:.4g}]{dbg_extra} -> {vpng}", flush=True)
                    except Exception:
                        pass
        except Exception:
            pass
        # 写出 vol 桶 JSON（独立于 tgrad 开关），便于离线排查
        try:
            if is_main_process(rank) and Vb > 0:
                vjson = str(getattr(args, 'log_vol_json', ''))
                if isinstance(vjson, str) and len(vjson) > 0:
                    ensure_dir(os.path.dirname(vjson))
                    payload = {
                        "epoch": int(epoch),
                        "cnt": [int(c) for c in vol_bucket_cnt],
                        "loss_avg": [float(vol_bucket_sum[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb)],
                        "snorm_avg": [float(vol_bucket_snorm[i] / max(1, vol_bucket_cnt[i])) for i in range(Vb)],
                        "edges": (vol_bucket_edges_epoch.tolist() if (isinstance(vol_bucket_edges_epoch, np.ndarray) and vol_bucket_edges_epoch.size == (Vb + 1)) else np.exp(np.linspace(np.log(1e-6), np.log(float(getattr(args,'noise_scale_clip',10.0))), Vb + 1)).tolist()),
                        "s_min": float(vol_s_min_epoch) if vol_s_min_epoch is not None else None,
                        "s_max": float(vol_s_max_epoch) if vol_s_max_epoch is not None else None,
                        "s_mean": float(vol_s_mean_sum_epoch / max(1, vol_s_sample_count_epoch)) if vol_s_sample_count_epoch > 0 else None,
                    }
                    mode_json = ('w' if int(epoch) == 1 else 'a')
                    with open(vjson, mode_json, encoding='utf-8') as f:
                        f.write(json.dumps(payload, ensure_ascii=False) + "\n")
        except Exception:
            pass
        if distributed and args.dist == 1:
            t_loss_tensor = torch.tensor([running_loss, running_count], dtype=torch.float64, device=device)
            dist.all_reduce(t_loss_tensor, op=dist.ReduceOp.SUM)
            total_loss_sum = float(t_loss_tensor[0].item())
            total_count_sum = float(t_loss_tensor[1].item())
            train_epoch_loss = total_loss_sum / max(1.0, total_count_sum)
        # 验证：整轮汇总 base loss 与正则项（方差匹配 + 频域匹配）
        model.eval()
        val_total_loss = 0.0
        val_total_count = 0
        # 验证：按 σ 桶累计（用于诊断大σ拖累）
        Kb_val = int(getattr(args, 'log_sigma_buckets', 8))
        val_bucket_sum = None
        val_bucket_cnt = None
        val_bucket_sqsum = None
        # 验证：固定 σ 视图（避免跨σ混合掩盖趋势）
        val_fixed_sum = 0.0
        val_fixed_count = 0
        val_fixed_enabled = False
        # 按股票跨时间的预测/真实累计统计（S、S2、C）
        val_S_pred = None
        val_S2_pred = None
        val_S_true = None
        val_S2_true = None
        val_C = 0
        # 频域：按股票跨时间的 DFT 累计（K,N）
        K_val = int(getattr(args, "freq_bins", 0))
        val_Sr_pred = None
        val_Si_pred = None
        val_Sr_true = None
        val_Si_true = None
        # 相关性评估收集
        preds_all = []
        trues_all = []
        # 验证集 y0_bias 统计（epoch汇总）
        val_bias_sum_epoch = 0.0
        val_bias_cnt_epoch = 0
        # yzvol 有效性诊断：在第一个 epoch 收集 vol_cond 和 Y（分标的）
        if epoch == 1 and is_main_process(rank):
            vol_cond_collect = []  # 收集 vol_cond [B, N, 1]
            y_collect = []         # 收集 Y [B, N, 1]
        with torch.no_grad():
            # 预计算固定 Sigma 网格（5个分位数，对数空间等间距）
            smin = float(getattr(args, 'sigma_min', getattr(args, 'score_sigma_min', 0.01)))
            smax = float(getattr(args, 'sigma_max', getattr(args, 'score_sigma_max', 1.0)))
            val_grid_steps = 5
            val_grid_indices = torch.linspace(0, 1, val_grid_steps, device=device)
            val_log_sigma_grid = math.log(smin) + val_grid_indices * (math.log(smax) - math.log(smin))
            
            for batch_v in val_loader:
                if isinstance(batch_v, (list, tuple)) and len(batch_v) == 4:
                    Xv, Yv, TIDXv, mask_v = batch_v  # [mask机制]
                    mask_v = mask_v.to(device)
                elif isinstance(batch_v, (list, tuple)) and len(batch_v) == 3:
                    Xv, Yv, TIDXv = batch_v
                    mask_v = None  # 兼容旧数据
                else:
                    Xv, Yv = batch_v
                    TIDXv = None
                    mask_v = None
                Xv = Xv.to(device)
                Yv = Yv.to(device)
                Bv = Xv.size(0)
                
                # 对当前 batch，遍历 5 个固定 sigma 计算损失并取平均
                batch_loss_accum = 0.0
                for i_grid in range(val_grid_steps):
                    log_sigma_val = val_log_sigma_grid[i_grid]
                    # 构造全 batch 统一的 sigma
                    log_sigma_v = torch.full((Bv,), log_sigma_val, device=device)
                    sigma_v = torch.exp(log_sigma_v).view(Bv, 1, 1)
                    
                    # 为每个网格点生成确定性噪声（使用固定生成器或简单重置种子太慢，这里既然是对固定网格平均，
                    # 为了覆盖更多噪声模式，我们在网格循环内保持随机 eps，但网格本身固定保证了强度覆盖的完备性）
                    # 或者：为了绝对确定性，应该固定 eps。但在 EDM 评估中，通常是对随机 eps 在不同 sigma 下的期望。
                    # 这里我们保留 eps 的随机性，但固定 sigma 分布，消除了最大的方差源（sigma 的随机采样）。
                    eps_v = torch.randn_like(Yv)
                    
                    # 验证阶段若启用 yz_vol，则与训练一致对 eps_v 做 per-asset 缩放
                    scale_mode_v = str(getattr(args, 'noise_scale_mode', 'constant')).lower()
                    s_sq_v = torch.ones_like(Yv)
                    if scale_mode_v == 'yz_vol':
                        yz_table_val = _load_yz_table(args)
                        date_list_v = _get_dataset_dates(val_set, 'val-noise-scale')
                        stocks_v = _get_dataset_stocks(val_set, 'val-noise-scale')
                        if TIDXv is None:
                            raise ValueError("[val-noise-scale] TIDXv 缺失，无法对齐 yz_vol")
                        if torch.is_tensor(TIDXv):
                            idx_tensor_v = TIDXv.to(device=device, dtype=torch.long)
                        else:
                            idx_tensor_v = torch.as_tensor(TIDXv, device=device, dtype=torch.long)
                        clip_v = float(getattr(args, 'noise_scale_clip', 10.0))
                        scales_v = _build_vol_tensor(
                            yz_table_val,
                            date_list_v,
                            stocks_v,
                            idx_tensor_v,
                            clip_v,
                            Yv.device,
                            Yv.dtype,
                            'val-noise-scale',
                        )
                        eps_v = eps_v * scales_v.unsqueeze(-1)
                        s_sq_v = (scales_v.unsqueeze(-1) * scales_v.unsqueeze(-1)).clamp_min(1e-12)
                    else:
                        scales_v = torch.ones(Xv.size(0), Xv.size(1), device=Yv.device, dtype=Yv.dtype)
                    
                    y0_v = Yv
                    x_t_v = y0_v + sigma_v * eps_v
                    
                    # t 索引用于时间嵌入映射
                    Tdiff_v = int(getattr(args, 'diff_steps', 1000))
                    t_idx_v = torch.clamp(((log_sigma_v - math.log(smin)) / max(1e-9, (math.log(smax) - math.log(smin))) * Tdiff_v).long() + 1, 1, Tdiff_v)
                    
                    # 与训练一致的时间扭曲
                    try:
                        warp_gamma_v = float(getattr(args, 'time_warp_gamma', 0.0))
                    except Exception:
                        warp_gamma_v = 0.0
                    if warp_gamma_v > 0.0 and scale_mode_v == 'yz_vol':
                        with torch.no_grad():
                            t_cont_v = torch.clamp(((log_sigma_v - math.log(smin)) / max(1e-9, (math.log(smax) - math.log(smin)))), 0.0, 1.0).float()
                            s_bar_v = torch.clamp(scales_v.view(Bv, -1).mean(dim=1), min=1e-6)
                            factor_v = torch.pow(s_bar_v, warp_gamma_v)
                            t_eff_v = torch.clamp(t_cont_v * factor_v, 0.0, 1.0)
                            t_idx_v = (t_eff_v * float(Tdiff_v)).long().clamp(1, Tdiff_v)
                    
                    # 显式构造验证期的波动率条件（若启用 yz_vol），形状 [B,N,1]
                    if scale_mode_v == 'yz_vol':
                        vol_cond_v = scales_v.unsqueeze(-1)
                    else:
                        vol_cond_v = None
                    
                    # yzvol 诊断：收集当前 batch 的 vol_cond 和 Y (仅收集第一次 grid pass)
                    if epoch == 1 and is_main_process(rank) and i_grid == 0:
                        if vol_cond_v is not None:
                            vol_cond_collect.append(vol_cond_v.detach().cpu())
                        y_collect.append(Yv.detach().cpu())
                    
                    # EDM 验证路径：直接预测 y0
                    sigma_data_v = float(getattr(args, 'sigma_data', 0.025))
                    sigma_data_sq_v = sigma_data_v ** 2
                    sigma_sq_v = sigma_v ** 2
                    c_skip_v = sigma_data_sq_v / (sigma_sq_v + sigma_data_sq_v)
                    c_out_v = sigma_v * sigma_data_v / torch.sqrt(sigma_sq_v + sigma_data_sq_v)
                    c_in_v = 1.0 / torch.sqrt(sigma_sq_v + sigma_data_sq_v)
                    
                    # 预条件化输入
                    x_scaled_v = c_in_v * x_t_v
                    
                    # 获取 yzvol 开关 (需与训练一致)
                    vol_as_input = bool(getattr(args, 'vol_cond_as_input', 1))
                    vol_scale_output = bool(getattr(args, 'vol_cond_scale_output', 0))

                    f_theta_v = model.forward_score(
                        Xv, x_scaled_v, t_idx_v, 
                        vol_cond=vol_cond_v,
                        vol_cond_as_input=vol_as_input,
                        vol_cond_scale_output=vol_scale_output
                    )
                    # 组装 y0_hat（验证期与训练期一致，移除 τ）
                    y0_hat_v = c_skip_v * x_t_v + c_out_v * f_theta_v
                    
                    # 损失：MSE(y0_hat, y0)
                    per_elem_v = (y0_hat_v - y0_v) ** 2
                    # 可选：验证期也应用 EDM λ(σ) 加权
                    try:
                        if int(getattr(args, 'edm_lambda', 0)) == 1:
                            lam_v = (sigma_sq_v + sigma_data_sq_v) / torch.clamp((sigma_v * sigma_data_v), min=1e-12) ** 2
                            pow_k_v = float(getattr(args, 'edm_lambda_pow', 1.0))
                            if abs(pow_k_v - 1.0) > 1e-8:
                                lam_v = torch.pow(lam_v, pow_k_v)
                            norm_mode_v = str(getattr(args, 'edm_lambda_norm', 'batch_mean')).lower()
                            if norm_mode_v == 'batch_mean':
                                lam_v = lam_v / lam_v.mean()
                            per_elem_v = per_elem_v * lam_v
                    except Exception:
                        pass
                    
                    # 样本覆盖权重（与训练一致）
                    if mask_v is not None:
                        coverage_v = mask_v.float().mean(dim=(2,3))
                        th = float(getattr(args, 'sample_coverage_threshold', 0.5))
                        sample_w_v = (coverage_v >= th).float().unsqueeze(-1)
                        weighted_v = per_elem_v * sample_w_v
                        valid_v = sample_w_v.sum().clamp_min(1.0)
                        bl_grid = weighted_v.sum() / valid_v
                    else:
                        bl_grid = per_elem_v.mean()
                    
                    batch_loss_accum += float(bl_grid.detach().cpu().item())
                
                # 该 batch 的最终损失是 5 个网格点的平均
                bl = torch.tensor(batch_loss_accum / float(val_grid_steps), device=device)
                # (以下逻辑只需执行一次，取最后一次 grid pass 的结果作为代表性输出用于统计)

                # ========== 验证：按 σ 桶累计 ==========
                # 注意：由于现在对每个 batch 遍历了整个 sigma 谱，桶统计会变得更加均匀
                # 我们仅使用最后一次 pass (sigma largest) 或 loop 内全部累计？
                # 为了简单，桶统计仅使用 log_sigma_v (loop 变量) 对应的最后一次值可能偏差较大
                # 因此我们略过桶统计的精细化修改，仅保留最后一次迭代的值用于兼容现有逻辑
                # 或者：如果用户非常看重 sigma 分桶，应该在 loop 内做累计。
                # 鉴于用户主要关注 val_loss 的稳定性，此处简化处理，不改动后续统计代码的缩进，
                # 直接让它们使用最后一次 loop 的变量状态 (sigma max)。
                # 这是一个折中，但对于 val_total_loss 的计算使用了平均值 bl。
                
                # 注意：后续代码如 val_total_loss += bl_value * Bv 使用的是 bl
                # 我们需要确保 bl 是 Tensor 或 float。上面已定义 bl 为 average scalar tensor。
                # ========== 验证：按 σ 桶累计 ==========
                try:
                    if Kb_val > 0:
                        loss_per_sample_v = per_elem_v.view(Bv, -1).mean(dim=1).detach().float().cpu().numpy()  # [B]
                        log_s = log_sigma_v.detach().float().cpu().numpy()  # [B]
                        lb_global = math.log(smin)
                        ub_global = math.log(smax)
                        width_global = max(1e-9, (ub_global - lb_global) / max(1, Kb_val))
                        if val_bucket_sum is None:
                            val_bucket_sum = [0.0] * Kb_val
                            val_bucket_cnt = [0] * Kb_val
                            val_bucket_sqsum = [0.0] * Kb_val
                        for i_b in range(Bv):
                            bidx = int(min(Kb_val - 1, max(0, math.floor((float(log_s[i_b]) - lb_global) / width_global))))
                            vhere = float(loss_per_sample_v[i_b])
                            val_bucket_sum[bidx] += vhere
                            val_bucket_cnt[bidx] += 1
                            val_bucket_sqsum[bidx] += vhere * vhere
                except Exception:
                    pass
                # ========== 验证：固定 σ 视图（EDM） ==========
                try:
                    if objective_v == 'edm':
                        sigma_fix = float(getattr(args, 'sigma_data', 0.025))
                        sigma_fix_t = torch.full_like(sigma_v, sigma_fix)
                        x_t_fix = y0_v + sigma_fix_t * eps_v
                        # t 索引（固定 σ 对应的时间嵌入）
                        log_sigma_fix = torch.full((Bv,), math.log(max(sigma_fix, 1e-12)), device=device, dtype=torch.float32)
                        Tdiff_v = int(getattr(args, 'diff_steps', 1000))
                        t_idx_fix = torch.clamp(((log_sigma_fix - math.log(smin)) / max(1e-9, (math.log(smax) - math.log(smin))) * Tdiff_v).long() + 1, 1, Tdiff_v)
                        # EDM 组装
                        sigma_fix_sq = sigma_fix_t * sigma_fix_t
                        sigma_data_v = float(getattr(args, 'sigma_data', 0.025))
                        sigma_data_sq_v = sigma_data_v * sigma_data_v
                        c_skip_f = sigma_data_sq_v / (sigma_fix_sq + sigma_data_sq_v)
                        c_out_f = sigma_fix_t * sigma_data_v / torch.sqrt(sigma_fix_sq + sigma_data_sq_v)
                        c_in_f = 1.0 / torch.sqrt(sigma_fix_sq + sigma_data_sq_v)
                        x_scaled_f = c_in_f * x_t_fix
                        f_theta_f = model.forward_score(Xv, x_scaled_f, t_idx_fix, vol_cond=vol_cond_v)
                        y0_hat_f = c_skip_f * x_t_fix + c_out_f * f_theta_f
                        per_elem_f = (y0_hat_f - y0_v) * (y0_hat_f - y0_v)
                        if mask_v is not None:
                            coverage_v = mask_v.float().mean(dim=(2,3))
                            th = float(getattr(args, 'sample_coverage_threshold', 0.5))
                            sample_w_v = (coverage_v >= th).float().unsqueeze(-1)
                            weighted_f = per_elem_f * sample_w_v
                            valid_f = sample_w_v.sum().clamp_min(1.0)
                            bl_f = weighted_f.sum() / valid_f
                        else:
                            bl_f = per_elem_f.mean()
                        val_fixed_sum += float(bl_f.detach().cpu().item()) * Bv
                        val_fixed_count += int(Bv)
                        val_fixed_enabled = True
                    else:
                        # 若非 EDM 分支，固定σ视图未执行，打印一次提示（仅首个 batch 以免刷屏）
                        try:
                            if is_main_process(rank) and (epoch == 1 and it == 1):
                                print(f"[val-fixed-sigma] skip: objective_v={objective_v} (only runs under 'edm')", flush=True)
                        except Exception:
                            pass
                except Exception as _e:
                    try:
                        # 输出异常详情，便于排查（不要静默）
                        shape_x = tuple(Xv.shape) if hasattr(Xv, 'shape') else 'NA'
                        shape_y = tuple(Yv.shape) if hasattr(Yv, 'shape') else 'NA'
                        dev = str(device)
                        print(f"[val-fixed-sigma] failed: {_e} | objective_v={objective_v} Bv={Bv} Xv={shape_x} Yv={shape_y} device={dev}", flush=True)
                    except Exception:
                        pass
                # 诊断：打印第一个batch的详细数值
                # 精简日志：不打印验证集 batch 级诊断
                bl_value = float(bl.detach().cpu().item())
                val_total_loss += bl_value * Bv
                val_total_count += int(Bv)
                # 精简日志：不打印每个验证batch的bl值
                # 统计验证集单步重建偏差（与训练口径一致）
                try:
                    y0_bias_v = float((y0_hat_v - y0_v).detach().abs().mean().cpu())
                    val_bias_sum_epoch += y0_bias_v
                    val_bias_cnt_epoch += 1
                except Exception:
                    pass
                p2d = y0_hat_v.detach().float().cpu().squeeze(-1)  # [B,N]
                y2d = Yv.detach().float().cpu().squeeze(-1)  # [B,N]
                if p2d.dim() == 1:  # 退化情况
                    p2d = p2d.view(1, -1)
                    y2d = y2d.view(1, -1)
                bs = p2d.size(0)
                psum = p2d.sum(dim=0)
                psq = (p2d * p2d).sum(dim=0)
                ysum = y2d.sum(dim=0)
                ysq = (y2d * y2d).sum(dim=0)
                if val_S_pred is None:
                    val_S_pred = psum.clone()
                    val_S2_pred = psq.clone()
                    val_S_true = ysum.clone()
                    val_S2_true = ysq.clone()
                else:
                    val_S_pred += psum
                    val_S2_pred += psq
                    val_S_true += ysum
                    val_S2_true += ysq
                val_C += int(bs)
                # 频域 DFT 累计（基于绝对时间索引）
                if float(getattr(args, "reg_freq_w", 0.0)) > 0.0 and K_val > 0:
                    # 绝对时间索引
                    if TIDXv is None:
                        # 若未提供则使用累进索引
                        start = (val_C - bs)
                        t_abs_v = torch.arange(start, start + bs, dtype=torch.long)
                    else:
                        t_abs_v = torch.as_tensor(TIDXv, dtype=torch.long)
                    # 在 CPU 上先计算，后续统一搬到 device 做 all_reduce
                    ks = torch.arange(1, K_val + 1, dtype=torch.float64)
                    C_tot_v = float(len(val_set))
                    theta_v = (2.0 * math.pi) * (t_abs_v.float().unsqueeze(1) * ks.unsqueeze(0)) / C_tot_v  # [B,K]
                    cos_v = torch.cos(theta_v)
                    sin_v = torch.sin(theta_v)
                    p_bn = p2d.double().unsqueeze(1)  # [B,1,N]
                    y_bn = y2d.double().unsqueeze(1)
                    contrib_r_p = (p_bn * cos_v.unsqueeze(-1)).sum(dim=0)  # [K,N]
                    contrib_i_p = (p_bn * sin_v.unsqueeze(-1)).sum(dim=0)
                    contrib_r_y = (y_bn * cos_v.unsqueeze(-1)).sum(dim=0)
                    contrib_i_y = (y_bn * sin_v.unsqueeze(-1)).sum(dim=0)
                    if val_Sr_pred is None:
                        Nn = p2d.size(1)
                        val_Sr_pred = contrib_r_p.clone()
                        val_Si_pred = contrib_i_p.clone()
                        val_Sr_true = contrib_r_y.clone()
                        val_Si_true = contrib_i_y.clone()
                    else:
                        val_Sr_pred += contrib_r_p
                        val_Si_pred += contrib_i_p
                        val_Sr_true += contrib_r_y
                        val_Si_true += contrib_i_y
                # 相关性
                preds_all.append(p2d.reshape(-1))
                trues_all.append(y2d.reshape(-1))
                # 按日收集用于日度横截面 rank-IC（需要 TIDXv 和 val_set.dates）
                try:
                    if TIDXv is not None and hasattr(val_set, 'dates'):
                        if 'daily_pred_dict' not in locals():
                            daily_pred_dict = {}
                            daily_true_dict = {}
                        for i_b in range(p2d.size(0)):
                            try:
                                ti = int(TIDXv[i_b].item())
                                if 0 <= ti < len(val_set.dates):
                                    key = str(pd.to_datetime(val_set.dates[ti]).strftime('%Y-%m-%d'))
                                    # 追加该日的 [N] 向量
                                    pv = p2d[i_b].detach().cpu().numpy()
                                    yv = y2d[i_b].detach().cpu().numpy()
                                    daily_pred_dict.setdefault(key, []).append(pv)
                                    daily_true_dict.setdefault(key, []).append(yv)
                            except Exception:
                                pass
                except Exception:
                    pass
        # 汇总 base loss（跨进程聚合）
        if distributed and args.dist == 1:
            v_loss_tensor = torch.tensor([val_total_loss, val_total_count], dtype=torch.float64, device=device)
            dist.all_reduce(v_loss_tensor, op=dist.ReduceOp.SUM)
            total_v_loss_sum = float(v_loss_tensor[0].item())
            total_v_count_sum = float(v_loss_tensor[1].item())
            val_base_loss = total_v_loss_sum / max(1.0, total_v_count_sum)
        else:
            total_v_loss_sum = val_total_loss
            total_v_count_sum = val_total_count
            val_base_loss = total_v_loss_sum / max(1.0, total_v_count_sum)
        # 写入 bias_epoch.jsonl：将训练与验证指标合并到同一行，并追加诊断字段
        try:
            if is_main_process(rank):
                bj_path = str(getattr(args, 'log_bias_json', ''))
                if isinstance(bj_path, str) and len(bj_path) > 0:
                    bj_epoch_path = bj_path.replace('bias_stats.jsonl', 'bias_epoch.jsonl')
                    ensure_dir(os.path.dirname(bj_epoch_path))
                    train_payload = {
                        "epoch": int(epoch),
                        "sigma_mean": float(bias_sigma_sum / max(1, bias_step_count)) if bias_step_count > 0 else None,
                        "s_yz_mean": float(bias_syz_sum / max(1, bias_step_count)) if bias_step_count > 0 else None,
                        "y0_bias_abs_mean": float(bias_y0_sum / max(1, bias_step_count)) if bias_step_count > 0 else None,
                        "val_y0_bias_abs_mean": float(val_bias_sum_epoch / max(1, val_bias_cnt_epoch)) if val_bias_cnt_epoch > 0 else None,
                        "steps_in_epoch": int(steps_in_epoch),
                        "bias_step_count": int(bias_step_count),
                        "bias_err_count": int(bias_err_count),
                        "bias_first_err": (str(bias_first_err) if (bias_first_err is not None) else None),
                        "used_score_branch": bool(int(getattr(args, 'use_score_model', 0)) == 1)
                    }
                    mode = 'w' if epoch == 1 else 'a'
                    with open(bj_epoch_path, mode, encoding='utf-8') as f:
                        f.write(json.dumps(train_payload, ensure_ascii=False) + "\n")
        except Exception:
            pass
        # 计算"是否在学"的关键指标：相关性、rank-IC、命中率、分层长短差
        r = float('nan'); hit = float('nan')
        long_short = float('nan'); top_mean = float('nan'); bot_mean = float('nan'); ic_mean = float('nan')
        ic_list = []
        try:
            if len(preds_all) > 0 and len(trues_all) > 0:
                # 拼接所有 batch: preds_all 是 List[Tensor[B*N]], 需要重塑为 [总样本数, N]
                # 假设每个 tensor 已经是 [B*N] 形状,需要知道 N
                # 从第一个 batch 推断 N (假设每个 batch 的样本数可能不同,但 N 固定)
                # 实际上 preds_all 存的是 p2d.reshape(-1),即 [B*N] 的一维数组
                # 我们需要重构成 [时间维度, 股票维度]
                
                # 方案: 从 val_set 或外部变量获取股票数量
                try:
                    # 尝试从 val_set 推断
                    if hasattr(val_set, 'feature_panel'):
                        N_stocks = int(val_set.feature_panel.shape[1])
                    elif hasattr(val_set, 'X'):
                        N_stocks = int(val_set.X.shape[1])
                    else:
                        # 从第一个 batch 推断 (假设所有 batch 的 N 相同)
                        if len(preds_all) > 0:
                            first_pred = preds_all[0].detach().cpu().numpy().ravel()
                            # 尝试从模型输出形状推断 (需要知道 batch_size)
                            # 这里作为最后的后备方案,直接使用全局变量
                            import inspect
                            frame = inspect.currentframe()
                            if frame is not None and frame.f_back is not None:
                                outer_locals = frame.f_back.f_locals
                                N_stocks = int(outer_locals.get('num_stocks', 0))
                            else:
                                N_stocks = 0
                        else:
                            N_stocks = 0
                    
                    if N_stocks <= 0:
                        raise ValueError(f"无法推断股票数量: N_stocks={N_stocks}")
                    
                    # 重塑所有预测和真值
                    pred_list = []
                    true_list = []
                    for p_t, t_t in zip(preds_all, trues_all):
                        p_np = p_t.detach().cpu().numpy().ravel().astype(np.float64)
                        t_np = t_t.detach().cpu().numpy().ravel().astype(np.float64)
                        # 确保能整除 N
                        if p_np.size % N_stocks == 0:
                            B_here = p_np.size // N_stocks
                            pred_list.append(p_np.reshape(B_here, N_stocks))
                            true_list.append(t_np.reshape(B_here, N_stocks))
                    
                    if len(pred_list) > 0:
                        pred_2d = np.concatenate(pred_list, axis=0)  # [总时间步, N]
                        true_2d = np.concatenate(true_list, axis=0)  # [总时间步, N]
                        
                        # 1. 全局指标 (保留用于相关性、命中率等)
                        pa = pred_2d.ravel()
                        ta = true_2d.ravel()
                        if (np.isfinite(pa).all() and np.isfinite(ta).all() and (pa.std() > 0) and (ta.std() > 0)):
                            r = float(np.corrcoef(pa, ta)[0, 1])
                        hit = float((np.sign(pa) == np.sign(ta)).mean()) if pa.size == ta.size and pa.size > 0 else float('nan')
                        
                        
                    else:
                        raise ValueError("无法重塑预测数据")
                
                # 分层统计 (top/bot quantile)
                # 确保 pa, ta 已定义
                if 'pa' in locals() and 'ta' in locals():
                    try:
                        q_hi = float(np.quantile(pa, 0.9))
                        q_lo = float(np.quantile(pa, 0.1))
                        top_mask = pa >= q_hi
                        bot_mask = pa <= q_lo
                        top_mean = float(ta[top_mask].mean()) if top_mask.any() else float('nan')
                        bot_mean = float(ta[bot_mask].mean()) if bot_mask.any() else float('nan')
                        if np.isfinite(top_mean) and np.isfinite(bot_mean):
                            long_short = float(top_mean - bot_mean)
                    except Exception:
                        pass
            if 'daily_pred_dict' in locals() and isinstance(daily_pred_dict, dict):
                for k, pv_list in daily_pred_dict.items():
                    try:
                        yv_list = daily_true_dict.get(k, None)
                        if yv_list is None:
                            continue
                        P = np.vstack(pv_list)
                        Tt = np.vstack(yv_list)
                        p_cs = P.mean(axis=0)
                        t_cs = Tt.mean(axis=0)
                        if p_cs.size >= 5 and np.isfinite(p_cs).any() and np.isfinite(t_cs).any():
                            pr = np.argsort(np.argsort(p_cs))
                            tr = np.argsort(np.argsort(t_cs))
                            if pr.std() > 0 and tr.std() > 0:
                                ic = float(np.corrcoef(pr, tr)[0, 1])
                                if np.isfinite(ic):
                                    ic_list.append(ic)
                    except Exception:
                        pass
            if len(ic_list) > 0:
                ic_mean = float(np.mean(ic_list))
        except Exception as e_metrics_prepare:
            print(f"[ERROR] 计算验证诊断指标失败 (epoch {epoch}): {e_metrics_prepare}", flush=True)
            import traceback
            traceback.print_exc()
        # 诊断：验证样本与时间累计统计
        # 移除验证样本计数的诊断打印
        # 诊断：验证样本与时间累计统计（方差匹配统计已移除）
        val_loss = val_base_loss
        # (旧频域匹配逻辑已移除)
        try:
            val_rmse_scalar = float(math.sqrt(max(0.0, float(val_base_loss))))
        except Exception:
            val_rmse_scalar = float('nan')
        epoch_sec = time.time() - epoch_t0
        # 精简日志：已移除 val-breakdown 详细拆解
        val_loss_fixed_sigma_val = None
        val_rmse_fixed_sigma_val = None
        if val_fixed_enabled and val_fixed_count > 0:
            v_fixed = float(val_fixed_sum / max(1, val_fixed_count))
            val_loss_fixed_sigma_val = v_fixed
            val_rmse_fixed_sigma_val = float(math.sqrt(max(0.0, v_fixed)))
        if running_count > 0:
            train_graph_reg_avg = graph_reg_sum_epoch / running_count
        else:
            train_graph_reg_avg = 0.0
        if is_main_process(rank):
            try:
                metrics_path = os.path.join('plots', 'train', 'metrics_epoch.jsonl')
                ensure_dir(os.path.dirname(metrics_path))
                payload = {
                    'epoch': int(epoch),
                    'val_base_loss': float(val_base_loss),
                    'val_loss': float(val_loss),
                    'val_rmse': val_rmse_scalar,
                    'val_loss_fixed_sigma': val_loss_fixed_sigma_val,
                    'val_rmse_fixed_sigma': val_rmse_fixed_sigma_val,
                    'pearson_r': r,
                    'rank_ic_mean': ic_mean,
                    'hitrate': hit,
                    'top_true_mean': top_mean,
                    'bot_true_mean': bot_mean,
                    'long_short_spread': long_short,
                    'train_graph_reg': float(train_graph_reg_avg),
                    'ic_days': int(len(ic_list))
                }
                with open(metrics_path, 'a', encoding='utf-8') as f:
                    f.write(json.dumps(payload, ensure_ascii=False) + "\n")
            except Exception as e_metrics:
                print(f"[ERROR] 写入 metrics_epoch.jsonl 失败 (epoch {epoch}): {e_metrics}", flush=True)
                import traceback
                traceback.print_exc()
            try:
                metrics_msg = (
                    f"[metrics][epoch {epoch}] r={r:.3f} ic={ic_mean:.3f} hit={hit:.3f} "
                    f"spread={long_short:.4f} graph_reg={train_graph_reg_avg:.4g}"
                )
                print(metrics_msg, flush=True)
            except Exception:
                pass
        curr_lr = optimizer.param_groups[0]["lr"]
        if bound_count_epoch > 0:
            train_bound_avg = bound_sum_epoch / bound_count_epoch
        else:
            train_bound_avg = 0.0
        # 计算验证集统计诊断量（std_pred/std_true 变为"按股票跨时间方差"的平均；同时输出方差总和）
        try:
            # 相关性指标：仍按全样本扁平计算
            if preds_all and trues_all:
                p_raw = torch.cat(preds_all).numpy()
                y_raw = torch.cat(trues_all).numpy()
                # 去除 NaN 诊断打印（稳定后不再输出）
                valid_mask = np.isfinite(p_raw) & np.isfinite(y_raw)
                p = p_raw[valid_mask]
                y = y_raw[valid_mask]
                n = p.shape[0]
                if n > 0:
                    if (np.std(p) > 1e-12) and (np.std(y) > 1e-12):
                        corr = float(np.corrcoef(p, y)[0, 1])
                    else:
                        corr = 0.0
                    # rank_ic（Spearman 秩相关）：优先用 scipy；无 scipy 时用 NumPy 回退（含并列平均秩）
                    rank_ic = 0.0
                    try:
                        from scipy.stats import spearmanr  # type: ignore
                        res = spearmanr(p, y)
                        # 兼容对象/元组两种返回类型，避免类型检查告警
                        if isinstance(res, tuple):
                            r0 = res[0]
                            rank_ic = float(r0) if isinstance(r0, (int, float)) else 0.0
                        else:
                            corr_attr = getattr(res, "correlation", None)
                            rank_ic = float(corr_attr) if isinstance(corr_attr, (int, float)) else 0.0
                        if not np.isfinite(rank_ic):
                            rank_ic = 0.0
                    except Exception:
                        # 使用全局 numpy 引用，避免局部作用域问题
                        import numpy
                        def _rank_avg_ties(a):
                            idx = numpy.argsort(a, kind="mergesort")
                            ranks = numpy.empty_like(a, dtype=numpy.float64)
                            n2 = a.shape[0]
                            i2 = 0
                            while i2 < n2:
                                j2 = i2
                                val = a[idx[i2]]
                                while (j2 + 1) < n2 and a[idx[j2 + 1]] == val:
                                    j2 += 1
                                # 平均秩（1-based）
                                avg = 0.5 * (i2 + j2) + 1.0
                                for k2 in range(i2, j2 + 1):
                                    ranks[idx[k2]] = avg
                                i2 = j2 + 1
                            return ranks
                        rp = _rank_avg_ties(p)
                        ry = _rank_avg_ties(y)
                        sp = numpy.std(rp)
                        sy = numpy.std(ry)
                        if sp > 1e-12 and sy > 1e-12:
                            rank_ic = float(numpy.corrcoef(rp, ry)[0, 1])
                        else:
                            rank_ic = 0.0
                    # 验证 RMSE（基于 y0 MSE）
                    val_stats_str = (
                        f"val_stats: corr={corr:.6g} rank_ic={rank_ic:.6g} hit={hit:.6g} val_rmse={val_rmse_scalar:.6g}"
                    )
                else:
                    val_stats_str = "val_stats: corr=nan rank_ic=nan hit=nan val_rmse=nan"
            else:
                val_stats_str = "val_stats: corr=nan rank_ic=nan hit=nan val_rmse=nan"
        except Exception as _e:
            val_stats_str = f"val_stats: unavailable ({_e})"
        # 计算整轮平均 DiffMSE（若启用扩散）
        try:
            epoch_diffmse = float(diffmse_sum / max(1, diffmse_count))
        except Exception:
            epoch_diffmse = 0.0
        # 验证集 diffmse 直接使用主验证循环的 val_base_loss（避免重复计算和随机种子不一致）
        epoch_val_diffmse = val_base_loss
        if is_main_process(rank):
            log_parts = [
                f"epoch {epoch}: time={epoch_sec:.2f}s",
                f"train_loss={train_epoch_loss:.6f}",
                f"val_loss={val_loss:.6f}",
                f"val_rmse={val_rmse_scalar:.6f}",
                f"lr={curr_lr:.6g}",
                f"edm_mse={epoch_diffmse:.4e}",
                f"val_edm_mse={epoch_val_diffmse:.4e}"
            ]
            if float(getattr(args, 'reg_bound_w', 0.0)) > 0.0:
                log_parts.append(f"bound_train={train_bound_avg:.4e}")
            print(' '.join(log_parts))
        print(val_stats_str, flush=True)
        # 附加：写出验证期按σ的桶统计（append到 sigma_stats.jsonl）
        try:
            if is_main_process(rank) and Kb_val > 0:
                out_json = str(getattr(args, 'log_sigma_json', ''))
                if isinstance(out_json, str) and len(out_json) > 0 and (val_bucket_sum is not None):
                    ensure_dir(os.path.dirname(out_json))
                    payload_val = {"epoch": int(epoch), "val_buckets": [], "type": "val"}
                    # 使用与训练相同的全局边界
                    try:
                        lb_global = math.log(smin)
                        ub_global = math.log(smax)
                        width_global = max(1e-9, (ub_global - lb_global) / max(1, Kb_val))
                        edges_list = [float(lb_global + i * width_global) for i in range(Kb_val + 1)]
                        payload_val["log_sigma_edges"] = edges_list
                    except Exception:
                        pass
                    for i in range(Kb_val):
                        cnt = int(val_bucket_cnt[i])
                        avg = float(val_bucket_sum[i] / max(1, cnt))
                        var = float(max(0.0, (val_bucket_sqsum[i] / max(1, cnt)) - avg * avg))
                        payload_val["val_buckets"].append({"bucket": i, "avg": avg, "var": var, "count": cnt})
                    with open(out_json, "a", encoding="utf-8") as f:
                        f.write(json.dumps(payload_val, ensure_ascii=False) + "\n")
        except Exception:
            pass
        # 调度器步进
        if scheduler_plateau is not None:
            scheduler_plateau.step(val_loss)
        lrs.append(optimizer.param_groups[0]["lr"])
        train_losses.append(train_epoch_loss)
        val_losses.append(val_loss)

        # 训练/验证写入统一在验证结束处合并输出（此处不再单独写训练行）

        # 参数统计日志（JSONL，每隔若干epoch记录一次，便于跟踪梯度/权重）
        try:
            param_stats_path = str(getattr(args, 'log_param_stats_json', ''))
            param_stats_interval = int(getattr(args, 'log_param_stats_interval', 1))
            if is_main_process(rank) and isinstance(param_stats_path, str) and len(param_stats_path) > 0 and epoch % param_stats_interval == 0:
                ensure_dir(os.path.dirname(param_stats_path))
                param_stats = collect_param_statistics(unwrap_model(model))
                param_payload = {
                    "epoch": int(epoch),
                    "params": param_stats,
                }
                with open(param_stats_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(param_payload, ensure_ascii=False) + "\n")
        except Exception:
            pass
        
        # 层级参数和梯度统计（JSONL，每隔若干epoch记录一次，监控训练健康度）
        if layer_stats_capture and is_main_process(rank):
            try:
                if layer_stats_snapshot is None:
                    print(f"[layer_stats] WARNING: Epoch {epoch} 无可用梯度快照（可能未触发优化步骤）", flush=True)
                else:
                    layer_payload = {
                        "epoch": int(epoch),
                        "train_loss": float(train_epoch_loss),
                        "val_loss": float(val_loss),
                        "lr": float(curr_lr),
                        "total_grad_norm": float(layer_stats_snapshot.get("total_grad_norm", 0.0)),
                        "max_grad": float(layer_stats_snapshot.get("max_grad", 0.0)),
                        "max_grad_layer": layer_stats_snapshot.get("max_grad_layer", ""),
                        "layers": layer_stats_snapshot.get("records", []),
                    }
                    mode = "w" if epoch == start_epoch else "a"
                    with open(layer_stats_path_epoch, mode, encoding="utf-8") as f:
                        f.write(json.dumps(layer_payload, ensure_ascii=False) + "\n")
                    print(
                        f"[layer_stats] Epoch {epoch}: total_grad_norm={layer_payload['total_grad_norm']:.4e}, "
                        f"max_grad={layer_payload['max_grad']:.4e} ({layer_payload['max_grad_layer']})",
                        flush=True,
                    )
            except Exception as e_layer:
                print(f"[layer_stats] WARNING: Failed to log layer stats: {e_layer}", flush=True)
        
        # 自适应谱噪声统计日志（JSONL，每隔若干epoch记录一次）
        try:
            spectral_stats_path = str(getattr(args, 'log_spectral_stats_json', ''))
            spectral_stats_interval = int(getattr(args, 'log_spectral_stats_interval', 10))
            if (is_main_process(rank) and spectral_noise_tracker is not None and 
                isinstance(spectral_stats_path, str) and len(spectral_stats_path) > 0 and 
                epoch % spectral_stats_interval == 0):
                ensure_dir(os.path.dirname(spectral_stats_path))
                
                # 获取谱噪声摘要
                summary = spectral_noise_tracker.get_summary(example_count=5, full_dump=False)
                # 获取敏感度统计
                sensitivity_stats = spectral_noise_tracker.get_sensitivity_stats()
                
                spectral_payload = {
                    "epoch": int(epoch),
                    "spectral_summary": summary,
                    "sensitivity": sensitivity_stats,
                }
                with open(spectral_stats_path, "a", encoding="utf-8") as f:
                    f.write(json.dumps(spectral_payload, ensure_ascii=False) + "\n")
                
                # 打印敏感度摘要
                if sensitivity_stats is not None:
                    sens_mean = sensitivity_stats["sensitivity_mean"]
                    sens_median = sensitivity_stats["sensitivity_median"]
                    sens_max = sensitivity_stats["sensitivity_max"]
                    print(f"[spectral] Epoch {epoch} - Sensitivity: mean={sens_mean:.4f}, median={sens_median:.4f}, max={sens_max:.4f} (target<0.01)", flush=True)
                    if sens_mean > 0.01:
                        print(f"[spectral] WARNING: 整体敏感度过高 (mean={sens_mean:.4f}>0.01)，建议增大robustness_coeff", flush=True)
        except Exception:
            pass

        # 仅主进程进行最佳模型选择与保存
        improved = bool(val_loss < (best_val - min_delta))
        if improved:
            best_val = val_loss
            bad = 0
            improve_count += 1
            if is_main_process(rank):
                if improve_count <= int(args.skip_first_best_saves):
                    print(f"[save] skip best #{improve_count} due to skip_first_best_saves={args.skip_first_best_saves}", flush=True)
                else:
                    if args.save_model == 1:
                        # 兼容 FSDP：在 rank0 汇聚完整 state_dict 再保存
                        if FSDP is not None and isinstance(model, FSDP):  # type: ignore[arg-type]
                            assert StateDictType is not None and FullStateDictConfig is not None
                            cfg_full = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
                            with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg_full):  # type: ignore[attr-defined]
                                full_state = model.state_dict()
                            # 记录训练期 artifacts 用于回测时一致性断言
                            artifacts = {}
                            try:
                                artifacts["train_stocks"] = (train_set.stocks if hasattr(train_set, 'stocks') else [])
                                # 严格：使用 test_env.pt 重排后的缩放参数
                                artifacts["mean_std_shapes"] = art_shapes
                                artifacts["mean_head3"] = art_mean_head3
                                artifacts["std_head3"] = art_std_head3
                                artifacts["scale_sha256"] = art_scale_sha256
                            except Exception:
                                pass
                            payload = {
                                "model_state": (full_state if int(getattr(args, 'ema', 1)) == 0 else full_state),
                                "config": {
                                    "input_dim": input_dim,
                                    "num_stocks": num_stocks,
                                    "seq_len": seq_len,
                                    "lr": args.lr,
                                    "d_model": args.d_model,
                                    "num_heads": args.num_heads,
                                    "num_layers": args.num_layers,
                                    "d_ff": args.d_ff,
                                    "dropout": args.dropout,
                                    "temp": args.temp,
                                    "feature_names": feature_names,
                                    "train": {
                                        "objective": "edm",
                                        "sigma_min": getattr(args, 'sigma_min', 0.01),
                                        "sigma_max": getattr(args, 'sigma_max', 1.0),
                                        "sigma_data": getattr(args, 'sigma_data', None),
                                        "edm_lambda": getattr(args, 'edm_lambda', 1),
                                        "edm_lambda_norm": getattr(args, 'edm_lambda_norm', 'batch_mean'),
                                        "edm_lambda_pow": getattr(args, 'edm_lambda_pow', 1.0),
                                        "vol_cond_as_input": getattr(args, 'vol_cond_as_input', 1),
                                        "vol_cond_scale_output": getattr(args, 'vol_cond_scale_output', 0),
                                    },
                                },
                                "artifacts": artifacts,
                                "best_val_loss": float(best_val),
                            }
                            torch.save(payload, best_path)
                        else:
                            state_to_save = unwrap_model(model).state_dict()
                            if int(getattr(args, 'ema', 1)) == 1 and 'ema_shadow' in locals() and (ema_shadow is not None):
                                state_to_save = ema_shadow
                            # 记录训练期 artifacts 用于回测时一致性断言
                            artifacts = {}
                            try:
                                artifacts["train_stocks"] = (train_set.stocks if hasattr(train_set, 'stocks') else [])
                                artifacts["mean_std_shapes"] = art_shapes
                                artifacts["mean_head3"] = art_mean_head3
                                artifacts["std_head3"] = art_std_head3
                                artifacts["scale_sha256"] = art_scale_sha256
                            except Exception:
                                pass
                            torch.save({
                                "model_state": state_to_save,
                                "config": {
                                    "input_dim": input_dim,
                                    "num_stocks": num_stocks,
                                    "seq_len": seq_len,
                                    "lr": args.lr,
                                    "d_model": args.d_model,
                                    "num_heads": args.num_heads,
                                    "num_layers": args.num_layers,
                                    "d_ff": args.d_ff,
                                    "dropout": args.dropout,
                                    "temp": args.temp,
                                    "feature_names": feature_names,
                                    "train": {
                                        "objective": "edm",
                                        "sigma_min": getattr(args, 'sigma_min', 0.01),
                                        "sigma_max": getattr(args, 'sigma_max', 1.0),
                                        "sigma_data": getattr(args, 'sigma_data', None),
                                        "edm_lambda": getattr(args, 'edm_lambda', 1),
                                        "edm_lambda_pow": getattr(args, 'edm_lambda_pow', 1.0),
                                        "vol_cond_as_input": getattr(args, 'vol_cond_as_input', 1),
                                        "vol_cond_scale_output": getattr(args, 'vol_cond_scale_output', 0),
                                    },
                                },
                                "artifacts": artifacts,
                                "best_val_loss": float(best_val),
                            }, best_path)
                        saved_any = True
                        # 保存最近一次检查点（支持 resume）
                        # 不使用try-except，checkpoint保存失败应立即报错
                        last_path = os.path.join(args.checkpoint_dir, 'last_checkpoint.pth')
                        torch.save({
                            'model_state': unwrap_model(model).state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'scaler': (scaler.state_dict() if scaler is not None and hasattr(scaler, 'state_dict') else None),
                            'ema_shadow': (ema_shadow if use_ema else None),
                            'epoch': int(epoch),
                        }, last_path)
                        print(f"saved best model -> {best_path}", flush=True)
                    else:
                        print("[save] best improved but save_model=0, skip saving weights", flush=True)
        else:
            bad += 1
            if bad >= patience:
                if is_main_process(rank):
                    print("Early stopping", flush=True)
                break

    # 训练结束：汇总 vol 分桶（跨全训练期间的 JSONL 明细，生成一次最终 PNG 与汇总 JSON）
    try:
        if is_main_process(rank) and int(getattr(args, 'log_vol_buckets', 0)) > 0:
            # 读取 vol_stats.jsonl 全部行并做均值/加权汇总
            vjson = str(getattr(args, 'log_vol_json', ''))
            out_png = str(getattr(args, 'log_vol_png_final', getattr(args, 'log_vol_png', 'plots/train/vol_buckets.png')))
            out_sum = str(getattr(args, 'log_vol_summary_json', 'plots/train/vol_stats_summary.json'))
            if isinstance(vjson, str) and len(vjson) > 0 and os.path.exists(vjson):
                cnt_list = []
                loss_list = []
                snorm_list = []
                edges_ref = None
                with open(vjson, 'r', encoding='utf-8') as f:
                    for ln in f:
                        try:
                            obj = json.loads(ln)
                            cnt_list.append(obj.get('cnt', []))
                            loss_list.append(obj.get('loss_avg', []))
                            snorm_list.append(obj.get('snorm_avg', []))
                            if edges_ref is None and 'edges' in obj:
                                edges_ref = obj['edges']
                        except Exception:
                            pass
                if cnt_list:
                    import numpy as _np
                    C = _np.array(cnt_list, dtype=_np.float64)
                    L = _np.array(loss_list, dtype=_np.float64)
                    S = _np.array(snorm_list, dtype=_np.float64)
                    # 简单平均（也可做样本数加权平均）
                    C_avg = _np.nanmean(C, axis=0).tolist()
                    L_avg = _np.nanmean(L, axis=0).tolist()
                    S_avg = _np.nanmean(S, axis=0).tolist()
                    try:
                        ensure_dir(os.path.dirname(out_sum))
                    except Exception:
                        pass
                    with open(out_sum, 'w', encoding='utf-8') as f:
                        json.dump({"cnt_avg": C_avg, "loss_avg": L_avg, "snorm_avg": S_avg, "edges": edges_ref or []}, f, ensure_ascii=False, indent=2)
                    # 画最终一次 PNG
                    try:
                        import matplotlib
                        matplotlib.use("Agg")
                        import matplotlib.pyplot as plt
                        ensure_dir(os.path.dirname(out_png))
                        edges = _np.array(edges_ref) if edges_ref is not None else None
                        if edges is None:
                            # 回退边界
                            Vb = int(getattr(args, 'log_vol_buckets', 8))
                            clip_u = float(getattr(args, 'noise_scale_clip', 10.0))
                            edges = _np.exp(_np.linspace(_np.log(1e-6), _np.log(clip_u), Vb + 1))
                        centers = 0.5 * (edges[:-1] + edges[1:])
                        plt.figure(figsize=(8,4))
                        ax1 = plt.gca()
                        ax1.plot(centers, L_avg, label='loss_avg', color='tab:blue')
                        ax1.set_xscale('log')
                        ax1.set_xlabel('yz_vol_scale (bucket center, log)')
                        ax1.set_ylabel('avg weighted loss (epoch-avg)')
                        ax2 = ax1.twinx()
                        ax2.plot(centers, S_avg, label='score_norm_avg', color='tab:orange')
                        ax2.set_xscale('log')
                        ax2.set_ylabel('avg ||score||_2 (epoch-avg)')
                        lines, labels = ax1.get_legend_handles_labels()
                        lines2, labels2 = ax2.get_legend_handles_labels()
                        ax1.legend(lines + lines2, labels + labels2, loc='best')
                        plt.title('Vol-bucket stats (final summary)')
                        plt.tight_layout()
                        plt.savefig(out_png, dpi=150)
                        plt.close()
                        print(f"[vol-bucket] wrote final summary: {out_png} & {out_sum}", flush=True)
                    except Exception as _e:
                        print(f"[vol-bucket] final plotting failed: {_e}")
    except Exception:
        pass

    if is_main_process(rank):
        print(f"best val loss: {best_val:.6f}", flush=True)
    if args.save_model == 1 and not saved_any:
        print("[warn] no model saved: skip_first_best_saves may be too large or no improvement over baseline", flush=True)

    # 可选：输出最佳指标与配置 JSON（用于调优阶段不保存权重的评分与注入）
    if is_main_process(rank) and isinstance(args.save_best_metrics_path, str) and len(args.save_best_metrics_path) > 0:
        try:
            with open(args.save_best_metrics_path, "w", encoding="utf-8") as f:
                json.dump({"best_val_loss": float(best_val)}, f, ensure_ascii=False, indent=2)
            print(f"[write] best metrics -> {args.save_best_metrics_path}", flush=True)
        except Exception as e:
            print(f"[write] best metrics failed: {e}", flush=True)

    if is_main_process(rank) and isinstance(args.save_best_config_path, str) and len(args.save_best_config_path) > 0:
        try:
            cfg_payload = {
                "input_dim": input_dim,
                "num_stocks": num_stocks,
                "seq_len": seq_len,
                "lr": args.lr,
                "d_model": args.d_model,
                "num_heads": args.num_heads,
                "num_layers": args.num_layers,
                "d_ff": args.d_ff,
                "dropout": args.dropout,
                "temp": args.temp,
                "best_val_loss": float(best_val),
                "feature_names": feature_names,
            }
            with open(args.save_best_config_path, "w", encoding="utf-8") as f:
                json.dump(cfg_payload, f, ensure_ascii=False, indent=2)
            print(f"[write] best config -> {args.save_best_config_path}", flush=True)
        except Exception as e:
            print(f"[write] best config failed: {e}", flush=True)

    if is_main_process(rank) and args.log_png:
        try_plot_train_curves(args.plot_dir, train_losses, val_losses, lrs)

    # 优雅退出分布式
    if distributed and dist.is_initialized():
        try:
            dist.barrier()
            dist.destroy_process_group()
        except Exception:
            pass


if __name__ == "__main__":
    main()