AR / train_v8.py

Upload folder using huggingface_hub

394918c verified 4 months ago

18.6 kB

import os
import argparse
import logging
from datetime import datetime

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm
import torch.nn.functional as F

from models.auto import StackedTrajectorySpatialTextRefiner
from models.vae import TripoSGVaeWrapper
from models.clip import FrozenCLIPTextEncoder
from utils.autodataloader import get_mesh_dataloader
from utils.loss import TokenL2Loss


class TextProjectionLayer(nn.Module):
    def __init__(self, in_dim=512, out_dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, in_dim),
            nn.ReLU(),
            nn.Linear(in_dim, out_dim),
        )
        self.residual = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity()

    def forward(self, x):
        return self.net(x) + self.residual(x)


def setup_logger(run_dir, rank):
    logger = logging.getLogger("TrainLogger")
    if logger.hasHandlers():
        logger.handlers.clear()

    if rank == 0:
        logger.setLevel(logging.INFO)
        formatter = logging.Formatter("%(asctime)s - %(message)s")

        log_file = os.path.join(run_dir, "train.log")
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
    else:
        logger.setLevel(logging.WARNING)

    return logger


def count_parameters(model, name, logger):
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"[{name}] trainable parameters: {trainable:,}")
    return trainable

def log_grad_stats(model, name, logger, topk=None):
    rows = []

    for param_name, param in model.named_parameters():
        if not param.requires_grad:
            continue

        if param.grad is None:
            rows.append((param_name, "grad=None"))
            continue

        grad = param.grad.detach()
        grad_norm = grad.norm(2).item()
        grad_mean_abs = grad.abs().mean().item()
        grad_max_abs = grad.abs().max().item()
        grad_min_abs = grad.abs().min().item()
        zero_ratio = (grad == 0).float().mean().item()

        rows.append(
            (
                param_name,
                grad_norm,
                grad_mean_abs,
                grad_max_abs,
                grad_min_abs,
                zero_ratio,
            )
        )

    logger.info(f"--- gradient statistics: {name} ---")

    for row in rows[:topk] if topk is not None else rows:
        if row[1] == "grad=None":
            logger.info(f"{name}.{row[0]} | grad=None")
        else:
            param_name, grad_norm, grad_mean_abs, grad_max_abs, grad_min_abs, zero_ratio = row
            logger.info(
                f"{name}.{param_name} | "
                f"norm={grad_norm:.3e} | "
                f"mean_abs={grad_mean_abs:.3e} | "
                f"max_abs={grad_max_abs:.3e} | "
                f"min_abs={grad_min_abs:.3e} | "
                f"zero_ratio={zero_ratio:.3f}"
            )


def train(args):
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl", device_id=device)
    rank = dist.get_rank()

    run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = os.path.join(args.save_dir, run_name)
    if rank == 0:
        os.makedirs(run_dir, exist_ok=True)
    dist.barrier()

    logger = setup_logger(run_dir, rank)

    vae = TripoSGVaeWrapper(device=device)
    vae.eval()
    for param in vae.parameters():
        param.requires_grad = False

    text_encoder = FrozenCLIPTextEncoder(args.clip_model_path, device=device)
    text_proj = TextProjectionLayer(in_dim=512, out_dim=args.token_dim).to(device)

    auto_model = StackedTrajectorySpatialTextRefiner(
        num_groups=args.num_refine_groups,
        temporal_kwargs={
            "num_tokens": args.num_tokens,
            "token_dim": args.token_dim,
            "max_seq_len": 16,
            "num_blocks": 6,
            "compress_layers": 3,
        },
        spatial_kwargs={
            "num_tokens": args.num_tokens,
            "token_dim": args.token_dim,
            "num_blocks": 3,
        },
        text_kwargs={
            "token_dim": args.token_dim,
            "text_dim": args.token_dim,
        },
    ).to(device)

    text_proj = DDP(text_proj, device_ids=[local_rank])
    auto_model = DDP(auto_model, device_ids=[local_rank])

    optimizer = optim.Adam(
        list(auto_model.parameters()) + list(text_proj.parameters()),
        lr=args.lr,
    )

    logger.info("--- parameter statistics ---")
    count_parameters(auto_model, "AutoRefiner", logger)
    count_parameters(text_proj, "TextProjection", logger)

    start_epoch = 0
    global_step = 0

    if args.auto_resume:
        if args.last_ckpt and os.path.exists(args.last_ckpt):
            logger.info(f"Resuming training from checkpoint: {args.last_ckpt}")
            checkpoint = torch.load(args.last_ckpt, map_location="cpu")

            auto_model.module.load_state_dict(checkpoint["auto"])
            text_proj.module.load_state_dict(checkpoint["proj"])

            if "optimizer" in checkpoint:
                optimizer.load_state_dict(checkpoint["optimizer"])
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if isinstance(v, torch.Tensor):
                            state[k] = v.to(device)

            if "epoch" in checkpoint:
                start_epoch = checkpoint["epoch"] + 1
            if "step" in checkpoint:
                global_step = checkpoint["step"]

            logger.info(f"Resume finished. Continue from Epoch {start_epoch + 1}, Global Step {global_step}.")
        else:
            logger.warning(
                f"auto_resume=True, but no valid checkpoint was found at: {args.last_ckpt}. "
                "Training will start from scratch."
            )

    criterion = TokenL2Loss()

    data_dir = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/shards/"
    dataloader = get_mesh_dataloader(
        data_dir=data_dir,
        batch_size=args.batch_size,
        num_workers=12,
    )

    bar_format = "{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]"

    for epoch in range(start_epoch, args.epochs):
        auto_model.train()
        text_proj.train()

        pbar = tqdm(
            dataloader,
            desc=f"Epoch {epoch + 1:03d}",
            disable=(rank != 0),
            ncols=1200,
            bar_format=bar_format,
        )

        for batch_idx, batch in enumerate(pbar):
            optimizer.zero_grad()
            vertices, faces = batch["vertices"], batch["faces"]

            with torch.no_grad():
                latents_batch = []
                for i in range(len(vertices)):
                    sample_latents = []
                    for t in range(vertices[i].shape[0]):
                        v_t = vertices[i][t].detach().clone().to(device=device, dtype=torch.float32)
                        f_t = faces[i].detach().clone().to(device=device, dtype=torch.float32)
                        latent = vae.encode_mesh(v_t, f_t)
                        if isinstance(latent, np.ndarray):
                            latent = torch.from_numpy(latent).to(device)
                        sample_latents.append(latent.squeeze(0))
                    latents_batch.append(torch.stack(sample_latents, dim=0))
                latents = torch.stack(latents_batch, dim=0)

            text_tokens = text_encoder(list(batch["caption"]))      # [B, L, 512]
            text_embed = text_proj(text_tokens) 
            # print(f"caption: {batch['caption']}")    
            # print(f"text_tokens.shape:{text_tokens.shape}")                    # [B, L, token_dim]
            # print(f"text_embed.shape:{text_embed.shape}")

            history = latents[:, :-1, :, :]
            target = latents[:, 1:, :, :]
            

            shortest_path_matrix = torch.zeros(
                (len(batch["vertices"]), args.num_tokens, args.num_tokens),
                device=device,
                dtype=torch.long,
            )

            refined_pred = auto_model(
                history,
                text_embed=text_embed,
                shortest_path_matrix=shortest_path_matrix,
            )

            loss = criterion(refined_pred, target)
            total_loss = loss["total"]

            if torch.isnan(total_loss):
                print(f"\n[!] NaN detected at Epoch {epoch + 1}, Batch {batch_idx + 1}")
                raise ValueError("Training stopped: NaN loss detected.")

            total_loss.backward()

            with torch.no_grad():
                debug_seed = 42

                # ------------------------------------------------------------
                # 准备同一个 surface_tensor
                # 这样可以排除 preprocess_mesh / trimesh.sample_surface 的随机性。
                # ------------------------------------------------------------
                surface_tensor = vae.preprocess_mesh(
                    v_t,
                    f_t,
                    num_samples=51200,
                    seed=debug_seed,
                )

                if surface_tensor is None:
                    print("[Debug] surface_tensor is None, skip VAE debug.")
                else:
                    # ============================================================
                    # 实验 1：同一个 surface_tensor，固定 VAE seed encode 两次
                    # ============================================================
                    enc1_fixed = vae.vae.encode(surface_tensor, seed=debug_seed)
                    enc2_fixed = vae.vae.encode(surface_tensor, seed=debug_seed)

                    z1_fixed = enc1_fixed.latent_dist.mode().to(vae.dtype).squeeze(0)
                    z2_fixed = enc2_fixed.latent_dist.mode().to(vae.dtype).squeeze(0)

                    fixed_mse = F.mse_loss(z1_fixed, z2_fixed).item()
                    fixed_cos = F.cosine_similarity(
                        z1_fixed.reshape(-1, z1_fixed.shape[-1]),
                        z2_fixed.reshape(-1, z2_fixed.shape[-1]),
                        dim=-1,
                    ).mean().item()

                    print("[Exp1] same surface_tensor + fixed VAE seed")
                    print("  mse:", fixed_mse)
                    print("  same-index cosine:", fixed_cos)

                    # ============================================================
                    # 实验 2：同一个 surface_tensor，不固定 VAE seed encode 两次
                    # ============================================================
                    enc1_random = vae.vae.encode(surface_tensor)
                    enc2_random = vae.vae.encode(surface_tensor)

                    z1_random = enc1_random.latent_dist.mode().to(vae.dtype).squeeze(0)
                    z2_random = enc2_random.latent_dist.mode().to(vae.dtype).squeeze(0)

                    random_mse = F.mse_loss(z1_random, z2_random).item()
                    random_diag_cos = F.cosine_similarity(
                        z1_random.reshape(-1, z1_random.shape[-1]),
                        z2_random.reshape(-1, z2_random.shape[-1]),
                        dim=-1,
                    ).mean().item()

                    sim_random = torch.matmul(
                        F.normalize(z1_random, dim=-1),
                        F.normalize(z2_random, dim=-1).T,
                    )  # [512, 512]

                    random_best_values, random_best_indices = sim_random.max(dim=-1)

                    print("[Exp2] same surface_tensor + random VAE seed")
                    print("  mse:", random_mse)
                    print("  same-index cosine:", random_diag_cos)
                    print("  best-match cosine mean:", random_best_values.mean().item())
                    print("  best-match cosine max:", random_best_values.max().item())
                    print("  best-match cosine min:", random_best_values.min().item())
                    print("  best-match cosine std:", random_best_values.std().item())

                    print("  first 20 token best matches: token_i -> token_j, cosine")
                    for i in range(min(20, z1_random.shape[0])):
                        print(
                            f"    {i:03d} -> {random_best_indices[i].item():03d}, "
                            f"{random_best_values[i].item():.4f}"
                        )

                    # ============================================================
                    # 实验 3：相邻帧 history[0,0] 和 target[0,0] 的 token set 匹配
                    # 每个 history token 都去 target 的 512 个 token 中找最高相似度。
                    # ============================================================
                    h = history[0, 0]  # [512, D]
                    y = target[0, 0]   # [512, D]

                    sim_adjacent = torch.matmul(
                        F.normalize(h, dim=-1),
                        F.normalize(y, dim=-1).T,
                    )  # [512, 512]

                    diag_values = sim_adjacent.diag()
                    best_values, best_indices = sim_adjacent.max(dim=-1)

                    print("[Exp3] adjacent frames token matching: history[0,0] -> target[0,0]")
                    print("  diag cosine mean:", diag_values.mean().item())
                    print("  diag cosine max:", diag_values.max().item())
                    print("  diag cosine min:", diag_values.min().item())
                    print("  diag cosine std:", diag_values.std().item())

                    print("  best-match cosine mean:", best_values.mean().item())
                    print("  best-match cosine max:", best_values.max().item())
                    print("  best-match cosine min:", best_values.min().item())
                    print("  best-match cosine std:", best_values.std().item())

                    same_index_ratio = (best_indices == torch.arange(best_indices.numel(), device=best_indices.device)).float().mean()
                    unique_match_ratio = best_indices.unique().numel() / best_indices.numel()

                    print("  best match is same index ratio:", same_index_ratio.item())
                    print("  unique matched target token ratio:", unique_match_ratio)

                    print("  first 20 adjacent-frame best matches: history_token_i -> target_token_j, cosine")
                    for i in range(min(20, h.shape[0])):
                        print(
                            f"    {i:03d} -> {best_indices[i].item():03d}, "
                            f"{best_values[i].item():.4f}, "
                            f"diag={diag_values[i].item():.4f}"
                        )

            if rank == 0 and global_step % args.grad_log_interval == 0:
                log_grad_stats(auto_model.module, "AutoRefiner", logger)
                log_grad_stats(text_proj.module, "TextProjection", logger)
                for handler in logger.handlers:
                    handler.flush()

            torch.nn.utils.clip_grad_norm_(
                list(auto_model.parameters()) + list(text_proj.parameters()),
                max_norm=1.0,
            )

            optimizer.step()

            global_step += 1

            if rank == 0:
                if global_step % 500 == 0:
                    ckpt_path = os.path.join(run_dir, f"model_step_{global_step}.pth")
                    torch.save(
                        {
                            "auto": auto_model.module.state_dict(),
                            "proj": text_proj.module.state_dict(),
                            "optimizer": optimizer.state_dict(),
                            "epoch": epoch,
                            "step": global_step,
                        },
                        ckpt_path,
                    )
                    logger.info(f"Checkpoint saved at Step {global_step}: {ckpt_path}")

                if (batch_idx + 1) % 10 == 0:
                    log_msg = (
                        f"Epoch {epoch + 1:03d} | Batch {batch_idx + 1:04d} | "
                        f"Total: {total_loss.item():.4f} | "
                        f"L1: {loss['l1'].item():.4f} | L2: {loss['l2'].item():.4f}"
                    )
                    logger.info(log_msg)
                    for handler in logger.handlers:
                        handler.flush()

                pbar.set_postfix({"Loss": f"{total_loss.item():.3f}"})

        if rank == 0 and (epoch + 1) % args.save_every_k_epochs == 0:
            ckpt_path = os.path.join(run_dir, f"model_epoch_{epoch + 1}.pth")
            torch.save(
                {
                    "auto": auto_model.module.state_dict(),
                    "proj": text_proj.module.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "epoch": epoch,
                    "step": global_step,
                },
                ckpt_path,
            )
            logger.info(f"Model saved: {ckpt_path}")

    dist.destroy_process_group()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_tokens", type=int, default=512)
    parser.add_argument("--token_dim", type=int, default=64)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--save_every_k_epochs", type=int, default=1)
    parser.add_argument("--save_dir", type=str, default="./checkpoints")
    parser.add_argument(
        "--clip_model_path",
        type=str,
        default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32",
    )
    parser.add_argument("--num_refine_groups", type=int, default=3)

    parser.add_argument("--auto_resume", action="store_true", help="whether to resume training from a checkpoint")
    parser.add_argument("--last_ckpt", type=str, default="", help="checkpoint path to resume from (.pth)")
    parser.add_argument("--grad_log_interval", type=int, default=100000)

    args = parser.parse_args()
    train(args)

	import os
	import argparse
	import logging
	from datetime import datetime

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.optim as optim
	import torch.distributed as dist
	from torch.nn.parallel import DistributedDataParallel as DDP
	from tqdm import tqdm
	import torch.nn.functional as F

	from models.auto import StackedTrajectorySpatialTextRefiner
	from models.vae import TripoSGVaeWrapper
	from models.clip import FrozenCLIPTextEncoder
	from utils.autodataloader import get_mesh_dataloader
	from utils.loss import TokenL2Loss


	class TextProjectionLayer(nn.Module):
	def __init__(self, in_dim=512, out_dim=64):
	super().__init__()
	self.net = nn.Sequential(
	nn.Linear(in_dim, in_dim),
	nn.ReLU(),
	nn.Linear(in_dim, out_dim),
	)
	self.residual = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity()

	def forward(self, x):
	return self.net(x) + self.residual(x)


	def setup_logger(run_dir, rank):
	logger = logging.getLogger("TrainLogger")
	if logger.hasHandlers():
	logger.handlers.clear()

	if rank == 0:
	logger.setLevel(logging.INFO)
	formatter = logging.Formatter("%(asctime)s - %(message)s")

	log_file = os.path.join(run_dir, "train.log")
	file_handler = logging.FileHandler(log_file)
	file_handler.setFormatter(formatter)
	logger.addHandler(file_handler)

	console_handler = logging.StreamHandler()
	console_handler.setFormatter(formatter)
	logger.addHandler(console_handler)
	else:
	logger.setLevel(logging.WARNING)

	return logger


	def count_parameters(model, name, logger):
	trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
	logger.info(f"[{name}] trainable parameters: {trainable:,}")
	return trainable

	def log_grad_stats(model, name, logger, topk=None):
	rows = []

	for param_name, param in model.named_parameters():
	if not param.requires_grad:
	continue

	if param.grad is None:
	rows.append((param_name, "grad=None"))
	continue

	grad = param.grad.detach()
	grad_norm = grad.norm(2).item()
	grad_mean_abs = grad.abs().mean().item()
	grad_max_abs = grad.abs().max().item()
	grad_min_abs = grad.abs().min().item()
	zero_ratio = (grad == 0).float().mean().item()

	rows.append(
	(
	param_name,
	grad_norm,
	grad_mean_abs,
	grad_max_abs,
	grad_min_abs,
	zero_ratio,
	)
	)

	logger.info(f"--- gradient statistics: {name} ---")

	for row in rows[:topk] if topk is not None else rows:
	if row[1] == "grad=None":
	logger.info(f"{name}.{row[0]} \| grad=None")
	else:
	param_name, grad_norm, grad_mean_abs, grad_max_abs, grad_min_abs, zero_ratio = row
	logger.info(
	f"{name}.{param_name} \| "
	f"norm={grad_norm:.3e} \| "
	f"mean_abs={grad_mean_abs:.3e} \| "
	f"max_abs={grad_max_abs:.3e} \| "
	f"min_abs={grad_min_abs:.3e} \| "
	f"zero_ratio={zero_ratio:.3f}"
	)


	def train(args):
	local_rank = int(os.environ["LOCAL_RANK"])
	torch.cuda.set_device(local_rank)
	device = torch.device("cuda", local_rank)
	dist.init_process_group(backend="nccl", device_id=device)
	rank = dist.get_rank()

	run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
	run_dir = os.path.join(args.save_dir, run_name)
	if rank == 0:
	os.makedirs(run_dir, exist_ok=True)
	dist.barrier()

	logger = setup_logger(run_dir, rank)

	vae = TripoSGVaeWrapper(device=device)
	vae.eval()
	for param in vae.parameters():
	param.requires_grad = False

	text_encoder = FrozenCLIPTextEncoder(args.clip_model_path, device=device)
	text_proj = TextProjectionLayer(in_dim=512, out_dim=args.token_dim).to(device)

	auto_model = StackedTrajectorySpatialTextRefiner(
	num_groups=args.num_refine_groups,
	temporal_kwargs={
	"num_tokens": args.num_tokens,
	"token_dim": args.token_dim,
	"max_seq_len": 16,
	"num_blocks": 6,
	"compress_layers": 3,
	},
	spatial_kwargs={
	"num_tokens": args.num_tokens,
	"token_dim": args.token_dim,
	"num_blocks": 3,
	},
	text_kwargs={
	"token_dim": args.token_dim,
	"text_dim": args.token_dim,
	},
	).to(device)

	text_proj = DDP(text_proj, device_ids=[local_rank])
	auto_model = DDP(auto_model, device_ids=[local_rank])

	optimizer = optim.Adam(
	list(auto_model.parameters()) + list(text_proj.parameters()),
	lr=args.lr,
	)

	logger.info("--- parameter statistics ---")
	count_parameters(auto_model, "AutoRefiner", logger)
	count_parameters(text_proj, "TextProjection", logger)

	start_epoch = 0
	global_step = 0

	if args.auto_resume:
	if args.last_ckpt and os.path.exists(args.last_ckpt):
	logger.info(f"Resuming training from checkpoint: {args.last_ckpt}")
	checkpoint = torch.load(args.last_ckpt, map_location="cpu")

	auto_model.module.load_state_dict(checkpoint["auto"])
	text_proj.module.load_state_dict(checkpoint["proj"])

	if "optimizer" in checkpoint:
	optimizer.load_state_dict(checkpoint["optimizer"])
	for state in optimizer.state.values():
	for k, v in state.items():
	if isinstance(v, torch.Tensor):
	state[k] = v.to(device)

	if "epoch" in checkpoint:
	start_epoch = checkpoint["epoch"] + 1
	if "step" in checkpoint:
	global_step = checkpoint["step"]

	logger.info(f"Resume finished. Continue from Epoch {start_epoch + 1}, Global Step {global_step}.")
	else:
	logger.warning(
	f"auto_resume=True, but no valid checkpoint was found at: {args.last_ckpt}. "
	"Training will start from scratch."
	)

	criterion = TokenL2Loss()

	data_dir = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/shards/"
	dataloader = get_mesh_dataloader(
	data_dir=data_dir,
	batch_size=args.batch_size,
	num_workers=12,
	)

	bar_format = "{l_bar}{bar:30}\| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]"

	for epoch in range(start_epoch, args.epochs):
	auto_model.train()
	text_proj.train()

	pbar = tqdm(
	dataloader,
	desc=f"Epoch {epoch + 1:03d}",
	disable=(rank != 0),
	ncols=1200,
	bar_format=bar_format,
	)

	for batch_idx, batch in enumerate(pbar):
	optimizer.zero_grad()
	vertices, faces = batch["vertices"], batch["faces"]

	with torch.no_grad():
	latents_batch = []
	for i in range(len(vertices)):
	sample_latents = []
	for t in range(vertices[i].shape[0]):
	v_t = vertices[i][t].detach().clone().to(device=device, dtype=torch.float32)
	f_t = faces[i].detach().clone().to(device=device, dtype=torch.float32)
	latent = vae.encode_mesh(v_t, f_t)
	if isinstance(latent, np.ndarray):
	latent = torch.from_numpy(latent).to(device)
	sample_latents.append(latent.squeeze(0))
	latents_batch.append(torch.stack(sample_latents, dim=0))
	latents = torch.stack(latents_batch, dim=0)

	text_tokens = text_encoder(list(batch["caption"])) # [B, L, 512]
	text_embed = text_proj(text_tokens)
	# print(f"caption: {batch['caption']}")
	# print(f"text_tokens.shape:{text_tokens.shape}") # [B, L, token_dim]
	# print(f"text_embed.shape:{text_embed.shape}")

	history = latents[:, :-1, :, :]
	target = latents[:, 1:, :, :]


	shortest_path_matrix = torch.zeros(
	(len(batch["vertices"]), args.num_tokens, args.num_tokens),
	device=device,
	dtype=torch.long,
	)

	refined_pred = auto_model(
	history,
	text_embed=text_embed,
	shortest_path_matrix=shortest_path_matrix,
	)

	loss = criterion(refined_pred, target)
	total_loss = loss["total"]

	if torch.isnan(total_loss):
	print(f"\n[!] NaN detected at Epoch {epoch + 1}, Batch {batch_idx + 1}")
	raise ValueError("Training stopped: NaN loss detected.")

	total_loss.backward()

	with torch.no_grad():
	debug_seed = 42

	# ------------------------------------------------------------
	# 准备同一个 surface_tensor
	# 这样可以排除 preprocess_mesh / trimesh.sample_surface 的随机性。
	# ------------------------------------------------------------
	surface_tensor = vae.preprocess_mesh(
	v_t,
	f_t,
	num_samples=51200,
	seed=debug_seed,
	)

	if surface_tensor is None:
	print("[Debug] surface_tensor is None, skip VAE debug.")
	else:
	# ============================================================
	# 实验 1：同一个 surface_tensor，固定 VAE seed encode 两次
	# ============================================================
	enc1_fixed = vae.vae.encode(surface_tensor, seed=debug_seed)
	enc2_fixed = vae.vae.encode(surface_tensor, seed=debug_seed)

	z1_fixed = enc1_fixed.latent_dist.mode().to(vae.dtype).squeeze(0)
	z2_fixed = enc2_fixed.latent_dist.mode().to(vae.dtype).squeeze(0)

	fixed_mse = F.mse_loss(z1_fixed, z2_fixed).item()
	fixed_cos = F.cosine_similarity(
	z1_fixed.reshape(-1, z1_fixed.shape[-1]),
	z2_fixed.reshape(-1, z2_fixed.shape[-1]),
	dim=-1,
	).mean().item()

	print("[Exp1] same surface_tensor + fixed VAE seed")
	print(" mse:", fixed_mse)
	print(" same-index cosine:", fixed_cos)

	# ============================================================
	# 实验 2：同一个 surface_tensor，不固定 VAE seed encode 两次
	# ============================================================
	enc1_random = vae.vae.encode(surface_tensor)
	enc2_random = vae.vae.encode(surface_tensor)

	z1_random = enc1_random.latent_dist.mode().to(vae.dtype).squeeze(0)
	z2_random = enc2_random.latent_dist.mode().to(vae.dtype).squeeze(0)

	random_mse = F.mse_loss(z1_random, z2_random).item()
	random_diag_cos = F.cosine_similarity(
	z1_random.reshape(-1, z1_random.shape[-1]),
	z2_random.reshape(-1, z2_random.shape[-1]),
	dim=-1,
	).mean().item()

	sim_random = torch.matmul(
	F.normalize(z1_random, dim=-1),
	F.normalize(z2_random, dim=-1).T,
	) # [512, 512]

	random_best_values, random_best_indices = sim_random.max(dim=-1)

	print("[Exp2] same surface_tensor + random VAE seed")
	print(" mse:", random_mse)
	print(" same-index cosine:", random_diag_cos)
	print(" best-match cosine mean:", random_best_values.mean().item())
	print(" best-match cosine max:", random_best_values.max().item())
	print(" best-match cosine min:", random_best_values.min().item())
	print(" best-match cosine std:", random_best_values.std().item())

	print(" first 20 token best matches: token_i -> token_j, cosine")
	for i in range(min(20, z1_random.shape[0])):
	print(
	f" {i:03d} -> {random_best_indices[i].item():03d}, "
	f"{random_best_values[i].item():.4f}"
	)

	# ============================================================
	# 实验 3：相邻帧 history[0,0] 和 target[0,0] 的 token set 匹配
	# 每个 history token 都去 target 的 512 个 token 中找最高相似度。
	# ============================================================
	h = history[0, 0] # [512, D]
	y = target[0, 0] # [512, D]

	sim_adjacent = torch.matmul(
	F.normalize(h, dim=-1),
	F.normalize(y, dim=-1).T,
	) # [512, 512]

	diag_values = sim_adjacent.diag()
	best_values, best_indices = sim_adjacent.max(dim=-1)

	print("[Exp3] adjacent frames token matching: history[0,0] -> target[0,0]")
	print(" diag cosine mean:", diag_values.mean().item())
	print(" diag cosine max:", diag_values.max().item())
	print(" diag cosine min:", diag_values.min().item())
	print(" diag cosine std:", diag_values.std().item())

	print(" best-match cosine mean:", best_values.mean().item())
	print(" best-match cosine max:", best_values.max().item())
	print(" best-match cosine min:", best_values.min().item())
	print(" best-match cosine std:", best_values.std().item())

	same_index_ratio = (best_indices == torch.arange(best_indices.numel(), device=best_indices.device)).float().mean()
	unique_match_ratio = best_indices.unique().numel() / best_indices.numel()

	print(" best match is same index ratio:", same_index_ratio.item())
	print(" unique matched target token ratio:", unique_match_ratio)

	print(" first 20 adjacent-frame best matches: history_token_i -> target_token_j, cosine")
	for i in range(min(20, h.shape[0])):
	print(
	f" {i:03d} -> {best_indices[i].item():03d}, "
	f"{best_values[i].item():.4f}, "
	f"diag={diag_values[i].item():.4f}"
	)

	if rank == 0 and global_step % args.grad_log_interval == 0:
	log_grad_stats(auto_model.module, "AutoRefiner", logger)
	log_grad_stats(text_proj.module, "TextProjection", logger)
	for handler in logger.handlers:
	handler.flush()

	torch.nn.utils.clip_grad_norm_(
	list(auto_model.parameters()) + list(text_proj.parameters()),
	max_norm=1.0,
	)

	optimizer.step()

	global_step += 1

	if rank == 0:
	if global_step % 500 == 0:
	ckpt_path = os.path.join(run_dir, f"model_step_{global_step}.pth")
	torch.save(
	{
	"auto": auto_model.module.state_dict(),
	"proj": text_proj.module.state_dict(),
	"optimizer": optimizer.state_dict(),
	"epoch": epoch,
	"step": global_step,
	},
	ckpt_path,
	)
	logger.info(f"Checkpoint saved at Step {global_step}: {ckpt_path}")

	if (batch_idx + 1) % 10 == 0:
	log_msg = (
	f"Epoch {epoch + 1:03d} \| Batch {batch_idx + 1:04d} \| "
	f"Total: {total_loss.item():.4f} \| "
	f"L1: {loss['l1'].item():.4f} \| L2: {loss['l2'].item():.4f}"
	)
	logger.info(log_msg)
	for handler in logger.handlers:
	handler.flush()

	pbar.set_postfix({"Loss": f"{total_loss.item():.3f}"})

	if rank == 0 and (epoch + 1) % args.save_every_k_epochs == 0:
	ckpt_path = os.path.join(run_dir, f"model_epoch_{epoch + 1}.pth")
	torch.save(
	{
	"auto": auto_model.module.state_dict(),
	"proj": text_proj.module.state_dict(),
	"optimizer": optimizer.state_dict(),
	"epoch": epoch,
	"step": global_step,
	},
	ckpt_path,
	)
	logger.info(f"Model saved: {ckpt_path}")

	dist.destroy_process_group()


	if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("--num_tokens", type=int, default=512)
	parser.add_argument("--token_dim", type=int, default=64)
	parser.add_argument("--batch_size", type=int, default=4)
	parser.add_argument("--lr", type=float, default=5e-4)
	parser.add_argument("--epochs", type=int, default=100)
	parser.add_argument("--save_every_k_epochs", type=int, default=1)
	parser.add_argument("--save_dir", type=str, default="./checkpoints")
	parser.add_argument(
	"--clip_model_path",
	type=str,
	default="/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/pretrain/clip-vit-base-patch32",
	)
	parser.add_argument("--num_refine_groups", type=int, default=3)

	parser.add_argument("--auto_resume", action="store_true", help="whether to resume training from a checkpoint")
	parser.add_argument("--last_ckpt", type=str, default="", help="checkpoint path to resume from (.pth)")
	parser.add_argument("--grad_log_interval", type=int, default=100000)

	args = parser.parse_args()
	train(args)