dump_ee_phaseA_features_split.py

import os
import json
import math
import time
import random
import datetime
import numpy as np
import torch
import torch.distributed as dist
from tqdm import tqdm
from torch.utils.data import DataLoader
from transformers import HfArgumentParser, AutoConfig
from sklearn.model_selection import train_test_split
import yaml
from datasets import concatenate_datasets

from src.arguments import ModelArguments, DataArguments, TrainingArguments
from src.data.collator.eval_collator import MultimodalEvalDataCollator
from src.data.eval_dataset.base_eval_dataset import AutoEvalPairDataset, generate_cand_dataset
from src.model.model_cut_layer_AOP_add_text_cut import MMEBModel
from src.model.processor import get_backbone_name, load_processor
from src.utils import batch_to_device, print_master


# ---------------- Utils ----------------
def _parse_bool(v: str, default=False):
    if v is None:
        return default
    v = v.strip().lower()
    return v in {"1", "true", "yes", "y", "t", "on"}


def _parse_int(v: str, default=None):
    try:
        return int(v) if v is not None else default
    except Exception:
        return default


def _parse_float(v: str, default=None):
    try:
        return float(v) if v is not None else default
    except Exception:
        return default


def get_env_aop_config():
    enabled = _parse_bool(os.environ.get("AOP_ENABLED"), False)
    apply_to = os.environ.get("AOP_APPLY", "qry").strip().lower()
    layer_idx = _parse_int(os.environ.get("AOP_LAYER"), None)
    mode = os.environ.get("AOP_MODE", "ratio").strip().lower()
    prune_vision = _parse_bool(os.environ.get("AOP_PRUNE_VISION"), True)
    prune_text = _parse_bool(os.environ.get("AOP_PRUNE_TEXT"), False)
    keep_ratio_v = _parse_float(os.environ.get("AOP_KEEP_RATIO_VISION"), None)
    keep_ratio_t = _parse_float(os.environ.get("AOP_KEEP_RATIO_TEXT"), None)
    attn_agg = os.environ.get("AOP_ATTENTION_AGG", "mean").strip().lower()
    ee_layer = _parse_int(os.environ.get("EE_LAYER"), None)

    return {
        "enabled": enabled,
        "apply_to": apply_to,
        "layer_idx": layer_idx,
        "mode": mode,
        "prune_vision": prune_vision,
        "prune_text": prune_text,
        "keep_ratio_vision": keep_ratio_v,
        "keep_ratio_text": keep_ratio_t,
        "attn_agg": attn_agg,
        "ee_layer": ee_layer,
    }


def pad_dataset_to_divisible(dataset, world_size):
    num_samples = len(dataset)
    if num_samples % world_size == 0:
        return dataset, num_samples
    num_to_add = world_size - (num_samples % world_size)
    padding_data = dataset.select([i % len(dataset) for i in range(num_to_add)])
    padded_dataset = concatenate_datasets([dataset, padding_data])
    return padded_dataset, num_samples + num_to_add


# ---------- Candidate encode in one pass (mid+last) ----------
@torch.no_grad()
def encode_candidates_both_layers(model: MMEBModel, loader: DataLoader, training_args: TrainingArguments, mid_layer: int):
    model.eval()
    all_mid, all_last, all_ids = [], [], []
    for inputs, infos in tqdm(loader, desc="[DUMP] Cands[BOTH]", disable=False):
        inputs = batch_to_device(inputs, training_args.device)
        # cand 侧不启用 AOP
        aop_cfg = getattr(model.encoder, "aop_prune_config", None)
        if isinstance(aop_cfg, dict) and aop_cfg:
            aop_off = dict(aop_cfg)
            aop_off["enabled"] = False
            setattr(model.encoder, "aop_prune_config", aop_off)

        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
            out = model.encoder(
                **inputs,
                return_dict=True,
                output_hidden_states=True,
                stop_at_layer=None,
                compute_lm_head=False,
            )
        hs_list = out.hidden_states
        assert hs_list is not None and len(hs_list) > mid_layer, "hidden_states too short for mid_layer"
        mid_hs, last_hs = hs_list[mid_layer], hs_list[-1]
        am = inputs.get("attention_mask", None)
        if am is not None and hasattr(am, "device") and am.device != mid_hs.device:
            am = am.to(mid_hs.device)
        reps_mid = model._pooling(mid_hs, am).detach().float().cpu()
        reps_last = model._pooling(last_hs, am).detach().float().cpu()
        all_mid.append(reps_mid)
        all_last.append(reps_last)
        all_ids.extend([info["cand_name"] for info in infos])
        
        # 恢复 AOP
        if isinstance(aop_cfg, dict) and aop_cfg:
            setattr(model.encoder, "aop_prune_config", aop_cfg)
            
    cand_mid = torch.cat(all_mid, dim=0).numpy()
    cand_last = torch.cat(all_last, dim=0).numpy()
    return cand_mid, cand_last, all_ids


# ---------- Build Phase A features ----------
@torch.no_grad()
def build_phaseA_features_global(
    reps_mid_t: torch.Tensor,  # [B,D] GPU
    cand_mid_t: torch.Tensor,  # [Nc,D] GPU
    am_mid: torch.Tensor,  # [B,L]
    input_ids: torch.Tensor,  # [B,L]
    cfg,  # model.encoder.config
    topk: int = 200,
    temp: float = 0.05,
):
    device = reps_mid_t.device
    B = reps_mid_t.size(0)
    # 相似度类
    scores_t = reps_mid_t @ cand_mid_t.T
    k = min(topk, scores_t.size(1))
    vals_t, _ = torch.topk(scores_t, k=k, dim=1)
    s1 = vals_t[:, 0]
    s2 = vals_t[:, 1] if k >= 2 else torch.zeros_like(s1)
    margin = s1 - s2
    p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=1)
    H = -(p_t * (torch.log(p_t + 1e-12))).sum(dim=1) / math.log(max(k, 1))
    sum_p2 = (p_t**2).sum(dim=1)

    # 长度比例
    am = am_mid.to(torch.bool)
    iid = input_ids
    image_token_id = getattr(cfg, "image_token_id", None)
    video_token_id = getattr(cfg, "video_token_id", None)
    bos_id = getattr(cfg, "bos_token_id", None)
    eos_id = getattr(cfg, "eos_token_id", None)
    pad_id = getattr(cfg, "pad_token_id", None)
    is_image = (iid == image_token_id) if (image_token_id is not None and image_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
    is_video = (iid == video_token_id) if (video_token_id is not None and video_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
    is_vision = (is_image | is_video) & am

    is_special = torch.zeros_like(iid, dtype=torch.bool)
    for tid in [bos_id, eos_id, pad_id]:
        if tid is not None and tid >= 0:
            is_special |= (iid == tid)
    is_text = am & (~is_vision) & (~is_special)

    L_vis = is_vision.sum(dim=1).float()
    L_txt = is_text.sum(dim=1).float()
    L_tot = am.sum(dim=1).float().clamp(min=1.0)
    r_vis = L_vis / L_tot
    r_txt = L_txt / L_tot

    # 类型 one-hot
    is_I = ((L_vis > 0) & (L_txt == 0)).float()
    is_T = ((L_txt > 0) & (L_vis == 0)).float()
    is_IT = ((L_txt > 0) & (L_vis > 0)).float()

    feats = torch.stack([s1, s2, margin, H, sum_p2, L_txt, L_vis, L_tot, r_txt, r_vis, is_I, is_T, is_IT], dim=1)
    return feats  # [B,13]


@torch.no_grad()
def build_phaseA_features_local(
    reps_mid_t: torch.Tensor,  # [B,D]
    cand_mid_t: torch.Tensor,  # [Nc,D]
    am_mid: torch.Tensor,  # [B,L]
    input_ids: torch.Tensor,  # [B,L]
    cfg,
    per_sample_rows: list,  # list[list[int]]
    topk: int = 200,
    temp: float = 0.05,
):
    device = reps_mid_t.device
    B = reps_mid_t.size(0)
    s1_list, s2_list, H_list, sum_p2_list = [], [], [], []
    for b in range(B):
        rows = per_sample_rows[b]
        if len(rows) == 0:
            s1_list.append(torch.tensor(0.0, device=device))
            s2_list.append(torch.tensor(0.0, device=device))
            H_list.append(torch.tensor(1.0, device=device))
            sum_p2_list.append(torch.tensor(0.0, device=device))
            continue
        cmat = cand_mid_t[rows]
        sv = (reps_mid_t[b:b+1] @ cmat.T)[0]
        k = min(topk, sv.size(0))
        vals, _ = torch.topk(sv, k=k, dim=0)
        s1_list.append(vals[0])
        s2_list.append(vals[1] if k >= 2 else torch.tensor(0.0, device=device, dtype=vals.dtype))
        p = torch.softmax(vals / max(temp, 1e-6), dim=0)
        H_list.append((-(p * (torch.log(p + 1e-12))).sum() / math.log(max(k, 1))))
        sum_p2_list.append((p**2).sum())
    s1 = torch.stack(s1_list)
    s2 = torch.stack(s2_list)
    H = torch.stack(H_list)
    sum_p2 = torch.stack(sum_p2_list)
    margin = s1 - s2

    am = am_mid.to(torch.bool)
    iid = input_ids
    image_token_id = getattr(cfg, "image_token_id", None)
    video_token_id = getattr(cfg, "video_token_id", None)
    bos_id = getattr(cfg, "bos_token_id", None)
    eos_id = getattr(cfg, "eos_token_id", None)
    pad_id = getattr(cfg, "pad_token_id", None)
    is_image = (iid == image_token_id) if (image_token_id is not None and image_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
    is_video = (iid == video_token_id) if (video_token_id is not None and video_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
    is_vision = (is_image | is_video) & am

    is_special = torch.zeros_like(iid, dtype=torch.bool)
    for tid in [bos_id, eos_id, pad_id]:
        if tid is not None and tid >= 0:
            is_special |= (iid == tid)
    is_text = am & (~is_vision) & (~is_special)

    L_vis = is_vision.sum(dim=1).float()
    L_txt = is_text.sum(dim=1).float()
    L_tot = am.sum(dim=1).float().clamp(min=1.0)
    r_vis = L_vis / L_tot
    r_txt = L_txt / L_tot

    is_I = ((L_vis > 0) & (L_txt == 0)).float()
    is_T = ((L_txt > 0) & (L_vis == 0)).float()
    is_IT = ((L_txt > 0) & (L_vis > 0)).float()

    feats = torch.stack([s1, s2, margin, H, sum_p2, L_txt, L_vis, L_tot, r_txt, r_vis, is_I, is_T, is_IT], dim=1)
    return feats  # [B,13]


# ---------- Label(y_exit) ----------
def compute_label_top1_equal_global(scores_mid: np.ndarray, scores_last: np.ndarray) -> np.ndarray:
    top1_mid = scores_mid.argmax(axis=1)
    top1_last = scores_last.argmax(axis=1)
    return (top1_mid == top1_last).astype(np.int32)


def compute_label_top1_equal_local(scores_mid_list, scores_last_list):
    y = []
    for sv_mid, sv_last in zip(scores_mid_list, scores_last_list):
        if sv_mid.size == 0 or sv_last.size == 0:
            y.append(0)
        else:
            y.append(int(int(sv_mid.argmax()) == int(sv_last.argmax())))
    return np.array(y, dtype=np.int32)


# ---------------- Main dump ----------------
def main():
    # DDP init
    if "RANK" in os.environ and dist.is_available() and not dist.is_initialized():
        dist.init_process_group(backend="nccl", timeout=datetime.timedelta(minutes=60))
    local_rank = dist.get_rank() if dist.is_initialized() else 0

    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    os.makedirs(data_args.encode_output_path, exist_ok=True)

    # Load model
    hf_config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
    if not getattr(model_args, "model_backbone", None):
        model_backbone = get_backbone_name(hf_config=hf_config, model_type=model_args.model_type)
        setattr(model_args, 'model_backbone', model_backbone)
        setattr(training_args, 'model_backbone', model_backbone)

    if local_rank == 0:
        processor = load_processor(model_args, data_args)
        model = MMEBModel.load(model_args, is_trainable=False, processor=processor)
    if dist.is_initialized():
        dist.barrier()
    if local_rank != 0:
        processor = load_processor(model_args, data_args)
        time.sleep(random.randint(2 * local_rank, 3 * local_rank))
        model = MMEBModel.load(model_args, is_trainable=False, processor=processor)

    model.eval()
    model = model.to(training_args.device, dtype=torch.bfloat16)

    # Configs
    ee_layer = int(os.environ.get("EE_LAYER", os.environ.get("AOP_LAYER", "12")))
    feat_topk = int(os.environ.get("EE_FEAT_TOPK", "200"))
    force_no_aop = os.environ.get("DUMP_EXIT_NO_AOP", "1").strip().lower() in {"1", "true", "yes", "on"}
    
    # 切分比例
    TRAIN_RATIO = 0.1
    VAL_RATIO = 0.1
    # Test ratio = 1 - 0.7 - 0.15 = 0.15

    with open(data_args.dataset_config, 'r', encoding='utf-8') as yf:
        dataset_configs = yaml.safe_load(yf)

    for dataset_name, task_cfg in dataset_configs.items():
        if dist.is_initialized(): dist.barrier()
        print_master(f"\n[DUMP] Processing {dataset_name} ...")

        if data_args.data_basedir:
            for key in ["image_root", "video_root", "frame_root", "clip_root", "data_path"]:
                if task_cfg.get(key):
                    task_cfg[key] = os.path.join(data_args.data_basedir, task_cfg[key])

        # 1. 加载全量数据
        full_qry, corpus = AutoEvalPairDataset.instantiate(model_args=model_args, data_args=data_args, **task_cfg)
        full_cand = generate_cand_dataset(full_qry, corpus)

        # 2. [CRITICAL] 生成切分索引 & 全局索引映射
        # 这一步确保我们在 Dump 时知道每个样本在全量数据集中的绝对位置 (Global Index)
        total_len = len(full_qry)
        all_indices = np.arange(total_len)
        
        # 切分逻辑
        train_idxs, temp_idxs = train_test_split(
            all_indices, train_size=TRAIN_RATIO, random_state=42, shuffle=True
        )
        val_relative_ratio = VAL_RATIO / (1.0 - TRAIN_RATIO)
        val_idxs, test_idxs = train_test_split(
            temp_idxs, train_size=val_relative_ratio, random_state=42, shuffle=True
        )
        
        print_master(f"[DUMP] Split sizes -> Train: {len(train_idxs)}, Val: {len(val_idxs)}, Test: {len(test_idxs)}")

        # 构造 Splits 字典，包含 Subset 和对应的 Global Indices
        splits = {
            "train": {"ds": full_qry.select(train_idxs), "indices": train_idxs},
            "val":   {"ds": full_qry.select(val_idxs),   "indices": val_idxs},
            "test":  {"ds": full_qry.select(test_idxs),  "indices": test_idxs}
        }

        # 3. 准备 Candidates
        cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand")
        cand_loader = DataLoader(
            full_cand, batch_size=training_args.per_device_eval_batch_size, 
            collate_fn=cand_collator, num_workers=training_args.dataloader_num_workers
        )
        cand_mid_np, cand_last_np, cand_ids = encode_candidates_both_layers(model, cand_loader, training_args, mid_layer=ee_layer)
        cand_id2row = {str(cid): i for i, cid in enumerate(cand_ids)}
        device = training_args.device
        cand_mid_t = torch.from_numpy(cand_mid_np).to(device=device, dtype=torch.bfloat16)
        cand_last_t = None

        # 4. 遍历 Split Dump
        sum_feat, sum2_feat, n_feat = None, None, 0
        scaler_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_phaseA_scaler.json")

        for split_name, split_info in splits.items():
            qry_dataset = split_info["ds"]
            global_indices = split_info["indices"]
            
            if len(qry_dataset) == 0: continue
            
            # [DDP Sharding] 同时切分 Dataset 和 Indices
            if dist.is_initialized():
                world_size = dist.get_world_size()
                per_rank = len(qry_dataset) // world_size
                start_idx = local_rank * per_rank
                end_idx = start_idx + per_rank
                
                if start_idx >= len(qry_dataset):
                    local_dataset = qry_dataset.select([])
                    local_indices = []
                else:
                    local_dataset = qry_dataset.select(range(start_idx, end_idx))
                    local_indices = global_indices[start_idx : end_idx]
            else:
                local_dataset = qry_dataset
                local_indices = global_indices

            qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry")
            qry_loader = DataLoader(
                local_dataset,
                batch_size=training_args.per_device_eval_batch_size,
                collate_fn=qry_collator,
                num_workers=training_args.dataloader_num_workers,
                shuffle=False # 必须 False 才能对齐 local_indices
            )

            feat_out_path_rank = os.path.join(data_args.encode_output_path, f"{dataset_name}_{split_name}_features.jsonl.rank{local_rank}")
            print_master(f" -> Dump {split_name} features to {feat_out_path_rank} ...")

            # 用于追踪当前 Batch 在 local_indices 中的游标
            cursor = 0

            with open(feat_out_path_rank, "w", encoding="utf-8") as fout:
                for inputs, infos in tqdm(qry_loader, desc=f"[{split_name.upper()}]", disable=(local_rank!=0)):
                    inputs = batch_to_device(inputs, device)
                    B = inputs["input_ids"].size(0)

                    # [CRITICAL] 获取当前 Batch 对应的 Global ID
                    batch_global_ids = local_indices[cursor : cursor + B]
                    cursor += B

                    # --- A. Forward Mid ---
                    aop_cfg_cur = getattr(model.encoder, "aop_prune_config", None)
                    orig_aop = None
                    if force_no_aop and isinstance(aop_cfg_cur, dict):
                        orig_aop = dict(aop_cfg_cur)
                        aop_off = dict(aop_cfg_cur)
                        aop_off["enabled"] = False
                        setattr(model.encoder, "aop_prune_config", aop_off)

                    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                        out_mid = model.encoder(
                            **inputs, return_dict=True, output_hidden_states=False,
                            stop_at_layer=int(ee_layer), compute_lm_head=False,
                            return_intermediate_state=True
                        )
                    if orig_aop is not None: setattr(model.encoder, "aop_prune_config", orig_aop)

                    # [FIX] Boolean tensor fix
                    hs_mid = getattr(out_mid, "last_hidden_state", None)
                    if hs_mid is None: hs_mid = out_mid.hidden_states[-1]
                    am_mid = getattr(out_mid, "attention_mask", None)
                    if am_mid is None: am_mid = inputs.get("attention_mask")
                    if hasattr(am_mid, "device") and am_mid.device != hs_mid.device: am_mid = am_mid.to(hs_mid.device)
                    reps_mid_t = model._pooling(hs_mid, am_mid).detach().to(device=device, dtype=torch.bfloat16)

                    # --- B. Build Feats ---
                    rank_global = task_cfg.get("eval_type", "global") == "global"
                    if rank_global:
                        feats_t = build_phaseA_features_global(reps_mid_t, cand_mid_t, am_mid, inputs["input_ids"], model.encoder.config, topk=feat_topk)
                    else:
                        rows_list = []
                        for b_idx in range(B):
                            cand_local = infos[b_idx]["cand_names"]
                            rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                            rows = [r for r in rows if r >= 0]
                            rows_list.append(rows)
                        feats_t = build_phaseA_features_local(reps_mid_t, cand_mid_t, am_mid, inputs["input_ids"], model.encoder.config, rows_list, topk=feat_topk)
                    feats_np = feats_t.detach().float().cpu().numpy()

                    # --- C. Forward Last ---
                    interm = getattr(out_mid, "intermediate_state", None)
                    resume_state = {
                        "hidden_states": interm["hidden_states"].detach(),
                        "attention_mask": interm["attention_mask"].detach(),
                        "position_ids": interm["position_ids"].detach(),
                        "vision_mask": interm.get("vision_mask"),
                        "text_mask": interm.get("text_mask"),
                        "next_layer_idx": int(interm["next_layer_idx"])
                    }
                    aop_cfg_cur = getattr(model.encoder, "aop_prune_config", None)
                    orig_aop2 = None
                    if force_no_aop and isinstance(aop_cfg_cur, dict):
                        orig_aop2 = dict(aop_cfg_cur)
                        aop_off2 = dict(aop_cfg_cur)
                        aop_off2["enabled"] = False
                        setattr(model.encoder, "aop_prune_config", aop_off2)

                    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                        out_last = model.encoder(
                            return_dict=True, output_hidden_states=False, stop_at_layer=None,
                            resume_state=resume_state, compute_lm_head=False
                        )
                    if orig_aop2 is not None: setattr(model.encoder, "aop_prune_config", orig_aop2)

                    # [FIX] Boolean tensor fix
                    hs_last = getattr(out_last, "last_hidden_state", None)
                    if hs_last is None: hs_last = out_last.hidden_states[-1]
                    am_last = getattr(out_last, "attention_mask", None)
                    if am_last is None: am_last = resume_state["attention_mask"]
                    if hasattr(am_last, "device") and am_last.device != hs_last.device: am_last = am_last.to(hs_last.device)
                    reps_last_t = model._pooling(hs_last, am_last).detach().to(device=device, dtype=torch.bfloat16)

                    # --- Label Logic ---
                    if rank_global:
                        if cand_last_t is None: cand_last_t = torch.from_numpy(cand_last_np).to(device=device, dtype=torch.bfloat16)
                        sim_mid = (reps_mid_t @ cand_mid_t.T).detach().float().cpu().numpy()
                        sim_last = (reps_last_t @ cand_last_t.T).detach().float().cpu().numpy()
                        y = compute_label_top1_equal_global(sim_mid, sim_last)
                    else:
                        y_list = []
                        for b_idx in range(B):
                            cand_local = infos[b_idx]["cand_names"]
                            rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                            rows = [r for r in rows if r >= 0]
                            if not rows: 
                                y_list.append(0)
                                continue
                            c_mid = cand_mid_t[rows]
                            if cand_last_t is None: cand_last_t = torch.from_numpy(cand_last_np).to(device=device, dtype=torch.bfloat16)
                            c_last = cand_last_t[rows]
                            sv_mid = (reps_mid_t[b_idx:b_idx+1] @ c_mid.T)[0].detach().float().cpu().numpy()
                            sv_last = (reps_last_t[b_idx:b_idx+1] @ c_last.T)[0].detach().float().cpu().numpy()
                            y_list.append(int(int(sv_mid.argmax()) == int(sv_last.argmax())))
                        y = np.array(y_list, dtype=np.int32)
                    
                    # --- D. Write ---
                    # Update Scaler Stats (Train Only)
                    if split_name == "train":
                        if sum_feat is None:
                            sum_feat = feats_np.sum(axis=0)
                            sum2_feat = (feats_np**2).sum(axis=0)
                        else:
                            sum_feat += feats_np.sum(axis=0)
                            sum2_feat += (feats_np**2).sum(axis=0)
                        n_feat += feats_np.shape[0]
                    
                    L_txt = feats_np[:, 5]
                    L_vis = feats_np[:, 6]
                    types = np.where((L_vis > 0) & (L_txt == 0), "I", np.where((L_txt > 0) & (L_vis == 0), "T", "IT"))

                    for b_idx in range(B):
                        row = {
                            "dataset": dataset_name,
                            "split": split_name,
                            "qid": int(batch_global_ids[b_idx]), # Global Index as QID
                            "type": str(types[b_idx]),
                            "feats": feats_np[b_idx].tolist(),
                            "y_exit": int(y[b_idx]),
                        }
                        fout.write(json.dumps(row, ensure_ascii=False) + "\n")

        # 5. Save Scaler
        if dist.is_initialized():
            dist.barrier()
            stats_vec = torch.tensor(
                np.concatenate([sum_feat, sum2_feat, [n_feat]]) if n_feat > 0 else np.zeros(13*2+1), 
                device=device, dtype=torch.float64
            )
            dist.all_reduce(stats_vec, op=dist.ReduceOp.SUM)
            sum_feat_all = stats_vec[:13].cpu().numpy()
            sum2_feat_all = stats_vec[13:26].cpu().numpy()
            n_feat_all = stats_vec[26].item()
        else:
            sum_feat_all = sum_feat
            sum2_feat_all = sum2_feat
            n_feat_all = n_feat

        if local_rank == 0 and n_feat_all > 0:
            mean = (sum_feat_all / n_feat_all).tolist()
            var = (sum2_feat_all / n_feat_all - (sum_feat_all / n_feat_all) ** 2)
            std = [float(max(1e-6, math.sqrt(max(0.0, v)))) for v in var.tolist()]
            with open(scaler_path, "w", encoding="utf-8") as f:
                json.dump({"mean": mean, "std": std, "in_dim": len(mean), "n_samples": n_feat_all, "dataset": dataset_name}, f, indent=2)
            print_master(f"[DUMP] {dataset_name} Scaler saved -> {scaler_path}")

    if dist.is_initialized(): dist.barrier()

if __name__ == "__main__":
    main()