【DEBUG 日志 | 分布式评估 AUC 乱飞】

import torch, torch.distributed as dist

def gather_step_wrong(pred, label):
    # pred: [B, 1], label: [B]
    ws = dist.get_world_size()
    pred_list  = [torch.zeros_like(pred)  for _ in range(ws)]
    label_list = [torch.zeros_like(label) for _ in range(ws)]

    # ❌ 直接 all_gather：要求每个 rank 张量同形状，否则会截断/复用缓存
    dist.all_gather(pred_list,  pred)     # <-- 尾批 B 不同 => 错位
    dist.all_gather(label_list, label)

    pred_all  = torch.cat(pred_list,  dim=0)
    label_all = torch.cat(label_list, dim=0)
    return pred_all, label_all

print(rank, pred.shape[0], label.shape[0])

# ddp_metric_gather.py —— 可直接复用
import torch, torch.distributed as dist

def gather_varlen_tensor(x: torch.Tensor, dim=0):
    """
    变长安全 all_gather：返回 rank0 上拼接后的张量；其他 rank 返回 None
    """
    assert x.is_cuda, "put tensors on CUDA for NCCL"
    world = dist.get_world_size()
    rank  = dist.get_rank()

    # 1) 同步各自真实长度
    len_local = torch.tensor([x.size(dim)], device=x.device, dtype=torch.int64)
    lens = [torch.zeros_like(len_local) for _ in range(world)]
    dist.all_gather(lens, len_local)
    lens = torch.stack(lens).squeeze(-1)  # [world]
    max_len = int(lens.max().item())

    # 2) 按 max_len padding 到同形状
    pad_shape = list(x.shape)
    pad_shape[dim] = max_len - x.size(dim)
    pad = torch.zeros(pad_shape, device=x.device, dtype=x.dtype)
    x_pad = torch.cat([x, pad], dim=dim)

    # 3) all_gather 到各自的缓冲区
    gather_list = [torch.zeros_like(x_pad) for _ in range(world)]
    dist.all_gather(gather_list, x_pad)

    # 4) 仅在 rank0 回切并拼接（按 lens 截断）
    if rank == 0:
        parts = []
        for r in range(world):
            end = int(lens[r].item())
            slc = [slice(None)] * x.dim()
            slc[dim] = slice(0, end)
            parts.append(gather_list[r][tuple(slc)])
        return torch.cat(parts, dim=dim)
    else:
        return None

@torch.no_grad()
def gather_preds_labels(pred: torch.Tensor, label: torch.Tensor):
    # pred [B,1] / [B,C]；label [B] / [B,C]
    pred_all  = gather_varlen_tensor(pred,  dim=0)
    label_all = gather_varlen_tensor(label, dim=0)
    # 仅 rank0 计算指标，其他 rank 返回 None
    if dist.get_rank() == 0:
        return pred_all.detach().cpu(), label_all.detach().cpu()
    return None, None

model.eval()
with torch.inference_mode():
    preds_local, labels_local = [], []
    for batch in val_loader:
        x, y = batch["img"].cuda(non_blocking=True), batch["label"].cuda(non_blocking=True)
        logits = model(x)
        prob   = torch.sigmoid(logits).squeeze(-1)  # [B]
        preds_local.append(prob)
        labels_local.append(y.float())

    pred = torch.cat(preds_local, dim=0)
    lab  = torch.cat(labels_local, dim=0)

    # 变长安全汇总
    pred_all, lab_all = gather_preds_labels(pred, lab)

    if dist.get_rank() == 0:
        from sklearn.metrics import roc_auc_score, average_precision_score
        auc = roc_auc_score(lab_all.numpy(), pred_all.numpy())
        ap  = average_precision_score(lab_all.numpy(), pred_all.numpy())
        print(f"[Global] AUC={auc:.4f} AP={ap:.4f}")