# train_tokenizer.py """Train the Dreamer 4 video tokenizer: a symmetric encoder/decoder Transformer trained by masked auto-encoding of 224x224 RGB frames. Run from inside ``src/`` (flat imports), e.g. ``torchrun --nproc_per_node=8 train_tokenizer.py``. """ import os import time import random import argparse from pathlib import Path import numpy as np import torch import torch.distributed as dist from torch.amp import autocast from torch.utils.data import DataLoader, DistributedSampler import wandb from task_set import TASK_SET, DOMAINS, UNSEEN_TASK_SET, task_to_domain, compute_task_weights from sharded_frame_dataset import ShardedFrameDataset from train_dynamics import PerDomainAccumulator from model import ( Encoder, Decoder, Tokenizer, temporal_patchify, temporal_unpatchify, recon_loss_from_mae, lpips_on_mae_recon, EmaRms, ) try: import lpips except ImportError: lpips = None torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True def get_dist_info(): rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) return rank, world_size, local_rank def is_rank0() -> bool: return int(os.environ.get("RANK", "0")) == 0 def seed_everything(seed: int): s = int(seed) % (2**32) random.seed(s) np.random.seed(s) torch.manual_seed(s) torch.cuda.manual_seed_all(s) def worker_init_fn(worker_id: int): info = torch.utils.data.get_worker_info() seed_everything(info.seed) def init_distributed() -> tuple[bool, int, int, int]: rank, world_size, local_rank = get_dist_info() ddp = world_size > 1 if ddp: dist.init_process_group(backend="nccl", init_method="env://") torch.cuda.set_device(local_rank) return ddp, rank, world_size, local_rank @torch.no_grad() def log_tokenizer_viz_wandb( *, x_btchw: torch.Tensor, # (B,T,C,H,W) float in [0,1] pred_btnd: torch.Tensor, # (B,T,Np,Dp) float in [0,1] mae_mask_btNp1: torch.Tensor, # (B,T,Np,1) bool True=masked patch: int, step: int, max_items: int = 8, max_T: int = 6, tag: str = "tokenizer/viz", ): B, T, C, H, W = x_btchw.shape Tv = min(T, max_T) Bv = min(B, max_items) # patchify target target_btnd = temporal_patchify(x_btchw[:, :Tv], patch) # (B,Tv,Np,Dp) # panels (patch space) masked_input_btnd = torch.where(mae_mask_btNp1[:, :Tv], torch.zeros_like(target_btnd), target_btnd) recon_masked_btnd = torch.where(mae_mask_btNp1[:, :Tv], pred_btnd[:, :Tv], target_btnd) recon_full_btnd = pred_btnd[:, :Tv] # to image space (B,T,C,H,W) target_img = temporal_unpatchify(target_btnd, H, W, C, patch) masked_img = temporal_unpatchify(masked_input_btnd, H, W, C, patch) rmask_img = temporal_unpatchify(recon_masked_btnd, H, W, C, patch) rfull_img = temporal_unpatchify(recon_full_btnd, H, W, C, patch) def tile_time(x: torch.Tensor) -> torch.Tensor: # (B,T,C,H,W) -> (B,C,H,T*W) x = x[:, :Tv] return x.permute(0, 2, 3, 1, 4).contiguous().view(x.shape[0], C, H, Tv * W) tgt = tile_time(target_img[:Bv]) msk = tile_time(masked_img[:Bv]) rm = tile_time(rmask_img[:Bv]) rf = tile_time(rfull_img[:Bv]) panel = torch.cat([tgt, msk, rm, rf], dim=2) # (Bv,C,4H,Tv*W) big = torch.cat([panel[i] for i in range(Bv)], dim=1) # (C,Bv*4H,Tv*W) big = (big.clamp(0, 1) * 255.0).to(torch.uint8) big_hwc = big.permute(1, 2, 0).cpu().numpy() wandb.log( { tag: wandb.Image( big_hwc, caption="rows=target/masked/recon_masked/recon_full", ), "tokenizer/masked_frac": float(mae_mask_btNp1[:, :Tv].float().mean().item()), }, step=step, ) @torch.no_grad() def log_tokenizer_val_viz_wandb( *, x_btchw: torch.Tensor, # (B,T,C,H,W) float in [0,1] pred_btnd: torch.Tensor, # (B,T,Np,Dp) float in [0,1] patch: int, step: int, max_items: int = 8, max_T: int = 6, tag: str = "val/viz", ): B, T, C, H, W = x_btchw.shape Tv = min(T, max_T) Bv = min(B, max_items) target_btnd = temporal_patchify(x_btchw[:, :Tv], patch) target_img = temporal_unpatchify(target_btnd, H, W, C, patch) recon_img = temporal_unpatchify(pred_btnd[:, :Tv], H, W, C, patch) def tile_time(x: torch.Tensor) -> torch.Tensor: x = x[:, :Tv] return x.permute(0, 2, 3, 1, 4).contiguous().view(x.shape[0], C, H, Tv * W) tgt = tile_time(target_img[:Bv]) rec = tile_time(recon_img[:Bv]) panel = torch.cat([tgt, rec], dim=2) # (Bv,C,2H,Tv*W) big = torch.cat([panel[i] for i in range(Bv)], dim=1) # (C,Bv*2H,Tv*W) big = (big.clamp(0, 1) * 255.0).to(torch.uint8) big_hwc = big.permute(1, 2, 0).cpu().numpy() wandb.log( {tag: wandb.Image(big_hwc, caption="rows=target/recon (full reconstruction)")}, step=step, ) @torch.no_grad() def validate( *, model, val_loader, device, use_amp: bool, lpips_fn, args: argparse.Namespace, max_batches: int, task_idx_to_domain_idx: torch.Tensor = None, local_task_names: list = None, key_prefix: str = "val", ) -> dict: """Full-reconstruction validation (no MAE masking). Returns aggregated metrics. If `local_task_names` is provided (len = number of dataset-local tasks), also computes per-task PSNR. `key_prefix` namespaces the output keys (e.g. "val"). """ was_training = model.training model.eval() _model = model.module if hasattr(model, "module") else model mae = _model.encoder.mae saved_pmin, saved_pmax = mae.p_min, mae.p_max mae.p_min, mae.p_max = 0.0, 0.0 # disable masking for full reconstruction sum_mse = 0.0 sum_lp = 0.0 n_batches = 0 first_x = None first_pred = None # Per-domain MSE sums (rank-local; val loader is not DDP-partitioned). n_domains = len(DOMAINS) dom_mse_sum = torch.zeros(n_domains, device=device, dtype=torch.float64) dom_count = torch.zeros(n_domains, device=device, dtype=torch.float64) # Optional per-task accumulators. Local task indices come straight from the # dataset (0..n_local_tasks-1) so we just bucket per task_idx_batch directly. track_per_task = local_task_names is not None n_local_tasks = len(local_task_names) if track_per_task else 0 task_mse_sum = torch.zeros(n_local_tasks, device=device, dtype=torch.float64) if track_per_task else None task_count = torch.zeros(n_local_tasks, device=device, dtype=torch.float64) if track_per_task else None try: for batch in val_loader: if n_batches >= max_batches: break if isinstance(batch, dict): x = batch["frames"].to(device, non_blocking=True) task_idx_batch = batch["task_idx"].to(device, non_blocking=True).long() else: x = batch.to(device, non_blocking=True) task_idx_batch = None patches = temporal_patchify(x, args.patch) with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp): pred, _, _ = model(patches) if first_x is None: first_x = x first_pred = pred.float() mse = (pred.float() - patches.float()).pow(2).mean() sum_mse += float(mse.item()) if task_idx_batch is not None and task_idx_to_domain_idx is not None: mse_per = (pred.float() - patches.float()).pow(2).mean(dim=(1, 2, 3)) # (B,) domain_ids = task_idx_to_domain_idx[task_idx_batch] dom_mse_sum.index_add_(0, domain_ids, mse_per.to(torch.float64)) dom_count.index_add_(0, domain_ids, torch.ones_like(mse_per, dtype=torch.float64)) if track_per_task: task_mse_sum.index_add_(0, task_idx_batch, mse_per.to(torch.float64)) task_count.index_add_(0, task_idx_batch, torch.ones_like(mse_per, dtype=torch.float64)) if lpips_fn is not None and args.lpips_weight > 0.0: # Full reconstruction LPIPS: pass an all-True mask so `recon_masked == pred`. full_mask = torch.ones_like(patches[..., :1], dtype=torch.bool) lp = lpips_on_mae_recon( lpips_fn, pred, patches, full_mask, H=args.H, W=args.W, C=args.C, patch=args.patch, subsample_frac=args.lpips_frac, ) sum_lp += float(lp.item()) n_batches += 1 finally: mae.p_min, mae.p_max = saved_pmin, saved_pmax if was_training: model.train() if n_batches == 0: return {}, None, None avg_mse = sum_mse / n_batches avg_lp = sum_lp / n_batches psnr = 10.0 * np.log10(1.0 / max(avg_mse, 1e-10)) out = { f"{key_prefix}/mse": avg_mse, f"{key_prefix}/psnr": float(psnr), f"{key_prefix}/n_batches": n_batches, } if lpips_fn is not None and args.lpips_weight > 0.0: out[f"{key_prefix}/lpips"] = avg_lp # Per-domain PSNR on full reconstruction (no MAE masking). for i, name in enumerate(DOMAINS): cnt = float(dom_count[i].item()) if cnt <= 0: continue mse_i = float((dom_mse_sum[i] / dom_count[i]).item()) if not np.isfinite(mse_i): continue out[f"{key_prefix}/domain/{name}/mse"] = mse_i out[f"{key_prefix}/domain/{name}/psnr"] = float(10.0 * np.log10(1.0 / max(mse_i, 1e-10))) # Per-task PSNR (only when local_task_names is provided — typically UNSEEN). if track_per_task: for i, name in enumerate(local_task_names): cnt = float(task_count[i].item()) if cnt <= 0: continue mse_i = float((task_mse_sum[i] / task_count[i]).item()) if not np.isfinite(mse_i): continue out[f"{key_prefix}/per_task/{name}/mse"] = mse_i out[f"{key_prefix}/per_task/{name}/psnr"] = float(10.0 * np.log10(1.0 / max(mse_i, 1e-10))) return out, first_x, first_pred def save_ckpt(path: Path, *, step: int, epoch: int, model, opt, args: argparse.Namespace, rms_state: dict = None): path.parent.mkdir(parents=True, exist_ok=True) # Unwrap DDP and torch.compile so the saved state_dict has plain keys # (no "module." or "_orig_mod." prefixes), making checkpoints portable # across compiled/uncompiled and DDP/non-DDP runs. target = model if hasattr(target, "module"): # DDP wrapper target = target.module if hasattr(target, "_orig_mod"): # torch.compile wrapper target = target._orig_mod obj = { "step": step, "epoch": epoch, "model": target.state_dict(), "opt": opt.state_dict(), "args": vars(args), } if rms_state is not None: obj["rms_state"] = rms_state tmp = path.with_suffix(".tmp") torch.save(obj, tmp) tmp.replace(path) def load_ckpt(path: Path, *, model, opt, rms_objects: dict = None) -> tuple[int, int]: ckpt = torch.load(path, map_location="cpu") state = ckpt["model"] # Unwrap DDP and torch.compile so we load into the raw nn.Module. The # checkpoint stores prefix-free keys (see save_ckpt), so loading into the # OptimizedModule wrapper would fail because it expects "_orig_mod." prefixes. target = model if hasattr(target, "module"): # DDP wrapper target = target.module if hasattr(target, "_orig_mod"): # torch.compile wrapper target = target._orig_mod target.load_state_dict(state, strict=True) opt.load_state_dict(ckpt["opt"]) if rms_objects is not None: for k, v in rms_objects.items(): if k in ckpt.get("rms_state", {}): v.load_state_dict(ckpt["rms_state"][k]) return int(ckpt.get("step", 0)), int(ckpt.get("epoch", 0)) def train(args): assert torch.cuda.is_available() ddp, rank, world_size, local_rank = init_distributed() device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu") seed_everything(args.seed + rank) # ---- data ---- # Also include UNSEEN tasks; ShardedFrameDataset filters to tasks with data # on disk, so runs without UNSEEN data are unaffected. full_task_list = list(TASK_SET) + list(UNSEEN_TASK_SET) tw_list = compute_task_weights( full_task_list, args.task_weighting, targeted_alpha=args.targeted_alpha, ) task_weights = None if args.task_weighting == "valid_starts" else dict(zip(full_task_list, tw_list)) dataset = ShardedFrameDataset( outdirs=args.data_dirs, tasks=full_task_list, seq_len=args.seq_len, iid_sampling=True, cache_size=args.shard_cache_size, samples_per_shard=args.samples_per_shard, ddp_partition=True, task_weights=task_weights, return_task_idx=True, ) # Lookup for per-domain metrics: task_idx (into dataset.tasks) -> domain_idx. task_idx_to_domain_idx = torch.tensor( [DOMAINS.index(task_to_domain(t)) for t in dataset.tasks], dtype=torch.long, device=device, ) sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) if ddp else None loader = DataLoader( dataset, batch_size=args.batch_size, sampler=sampler, shuffle=(sampler is None), num_workers=args.num_workers, pin_memory=True, drop_last=True, persistent_workers=(args.num_workers > 0), prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None, worker_init_fn=worker_init_fn, ) # ---- validation data ---- val_loader = None if args.val_data_dir is not None and args.val_every > 0: val_dataset = ShardedFrameDataset( outdirs=[args.val_data_dir], tasks=TASK_SET, seq_len=args.seq_len, iid_sampling=True, cache_size=args.shard_cache_size, samples_per_shard=args.val_samples_per_shard, # Don't partition val across ranks: val is small, and letting every # rank sample from the full pool maximizes the diversity of # sequences seen (especially the rank-0 first batch used for viz). ddp_partition=False, task_weights=task_weights, return_task_idx=True, ) val_workers = max(1, args.num_workers // 2) val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, # iid_sampling handles randomness num_workers=val_workers, pin_memory=True, drop_last=True, persistent_workers=(val_workers > 0), prefetch_factor=args.prefetch_factor if val_workers > 0 else None, worker_init_fn=worker_init_fn, ) # ---- UNSEEN validation data (optional) ---- val_unseen_loader = None val_unseen_task_idx_to_domain_idx = None val_unseen_local_tasks = None if args.val_unseen_data_dir is not None and args.val_every > 0: val_unseen_dataset = ShardedFrameDataset( outdirs=[args.val_unseen_data_dir], tasks=list(UNSEEN_TASK_SET), seq_len=args.seq_len, iid_sampling=True, cache_size=args.shard_cache_size, samples_per_shard=args.val_samples_per_shard, ddp_partition=False, task_weights=None, # uniform across the (≤10) UNSEEN tasks present on disk return_task_idx=True, ) val_unseen_local_tasks = list(val_unseen_dataset.tasks) val_unseen_task_idx_to_domain_idx = torch.tensor( [DOMAINS.index(task_to_domain(t)) for t in val_unseen_local_tasks], dtype=torch.long, device=device, ) val_unseen_workers = max(1, args.num_workers // 2) val_unseen_loader = DataLoader( val_unseen_dataset, batch_size=args.batch_size, shuffle=False, num_workers=val_unseen_workers, pin_memory=True, drop_last=True, persistent_workers=(val_unseen_workers > 0), prefetch_factor=args.prefetch_factor if val_unseen_workers > 0 else None, worker_init_fn=worker_init_fn, ) if is_rank0(): print(f"[rank0] UNSEEN val loader: {len(val_unseen_local_tasks)} tasks from {args.val_unseen_data_dir}") # ---- model ---- assert args.H % args.patch == 0 and args.W % args.patch == 0 n_patches = (args.H // args.patch) * (args.W // args.patch) d_patch = args.patch * args.patch * args.C assert args.d_model % args.n_heads == 0, "d_model must be divisible by n_heads" enc = Encoder( patch_dim=d_patch, d_model=args.d_model, n_latents=args.n_latents, n_patches=n_patches, n_heads=args.n_heads, depth=args.depth, d_bottleneck=args.d_bottleneck, dropout=args.dropout, mlp_ratio=args.mlp_ratio, time_every=args.time_every, mae_p_min=args.mae_p_min, mae_p_max=args.mae_p_max, ) dec = Decoder( d_bottleneck=args.d_bottleneck, d_model=args.d_model, n_heads=args.n_heads, depth=args.depth, n_latents=args.n_latents, n_patches=n_patches, d_patch=d_patch, dropout=args.dropout, mlp_ratio=args.mlp_ratio, time_every=args.time_every, ) model = Tokenizer(enc, dec).to(device) rms_mse = EmaRms().to(device) rms_lp = EmaRms().to(device) if is_rank0(): print(model) param_count = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Learnable parameters: {param_count:,}") if args.compile: model = torch.compile(model) if ddp: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False ) # ---- optim ---- trainable_params = [p for p in model.parameters() if p.requires_grad] opt = torch.optim.AdamW(trainable_params, lr=args.lr, weight_decay=args.weight_decay) use_amp = torch.cuda.is_available() # ---- lpips ---- if args.lpips_weight > 0.0: assert lpips is not None, "pip install lpips" lpips_fn = lpips.LPIPS(net=args.lpips_net).to(device) lpips_fn.eval() lpips_fn.requires_grad_(False) else: lpips_fn = None # ---- wandb ---- if is_rank0(): _model_for_count = model.module if hasattr(model, "module") else model param_count_total = sum(p.numel() for p in _model_for_count.parameters()) param_count_trainable = sum(p.numel() for p in _model_for_count.parameters() if p.requires_grad) param_count_enc = sum(p.numel() for p in _model_for_count.encoder.parameters() if p.requires_grad) param_count_dec = sum(p.numel() for p in _model_for_count.decoder.parameters() if p.requires_grad) frames_per_opt_step = args.batch_size * world_size * max(1, args.grad_accum) * args.seq_len wandb.init( project=args.wandb_project, name=args.wandb_run_name, entity=args.wandb_entity, mode="online", config={ **vars(args), "run/world_size": world_size, "run/effective_batch_seqs": args.batch_size * world_size * max(1, args.grad_accum), "run/frames_per_opt_step": frames_per_opt_step, "model/param_count_total": param_count_total, "model/param_count_trainable": param_count_trainable, "model/param_count_encoder": param_count_enc, "model/param_count_decoder": param_count_dec, "model/n_patches": n_patches, "model/d_patch": d_patch, }, ) # ---- resume ---- step = 0 start_epoch = 0 ckpt_dir = Path(args.ckpt_dir) if args.resume is not None: step, start_epoch = load_ckpt(Path(args.resume), model=model, opt=opt, rms_objects={"mse": rms_mse, "lp": rms_lp}) if is_rank0(): print(f"[rank0] Resumed from {args.resume} (step={step}, epoch={start_epoch})") # ---- train ---- model.train() t0 = time.time() grad_accum = max(1, int(args.grad_accum)) # Per-domain MSE accumulator (flushed at each log_every boundary). domain_acc = PerDomainAccumulator(n_domains=len(DOMAINS), device=device) while step <= args.max_steps: for epoch in range(start_epoch, 10_000_000): if sampler is not None: sampler.set_epoch(epoch) for batch in loader: if step > args.max_steps: break x = batch["frames"].to(device, non_blocking=True) # (B,T,C,H,W) task_idx_batch = batch["task_idx"].to(device, non_blocking=True).long() patches = temporal_patchify(x, args.patch) if is_rank0() and step % args.log_every == 0: with torch.no_grad(): _model = model.module if hasattr(model, "module") else model z, _ = _model.encoder(patches) zf = z.float() wandb.log({ "debug/z_std": float(zf.std().item()), "debug/z_mean": float(zf.mean().item()), "debug/z_abs_max": float(zf.abs().max().item()), }, step=step) with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp): pred, mae_mask, keep_prob = model(patches) # losses in fp32 (outside autocast). loss_mask = mae_mask mse = recon_loss_from_mae(pred, patches, loss_mask) rms_mse.update(mse) # Per-sample masked MSE for per-domain PSNR tracking. # Mirrors recon_loss_from_mae but sums per-sample instead of globally. with torch.no_grad(): diff_sq = (pred.float() - patches.float()).pow(2) # (B,T,Np,Dp) m = loss_mask.to(dtype=torch.float32) # (B,T,Np,1) sq_per = (diff_sq * m).sum(dim=(1, 2, 3)) # (B,) denom_per = m.sum(dim=(1, 2, 3)) * diff_sq.shape[-1] # (B,) mse_per_sample = sq_per / denom_per.clamp_min(1.0) # (B,) domain_ids = task_idx_to_domain_idx[task_idx_batch] domain_acc.update(mse_per_sample, domain_ids) if lpips_fn is not None and args.lpips_weight > 0.0: lp = lpips_on_mae_recon( lpips_fn, pred, patches, loss_mask, H=args.H, W=args.W, C=args.C, patch=args.patch, subsample_frac=args.lpips_frac ) rms_lp.update(lp) loss = rms_mse.normalize(mse) + args.lpips_weight * rms_lp.normalize(lp) else: lp = torch.zeros((), device=device) loss = rms_mse.normalize(mse) if not torch.isfinite(loss): raise RuntimeError(f"Non-finite loss at step {step}: loss={loss} mse={mse} lp={lp}") loss_to_backprop = loss / grad_accum loss_to_backprop.backward() do_step = ((step + 1) % grad_accum == 0) grad_norm = 0.0 grad_norm_enc = 0.0 grad_norm_dec = 0.0 if do_step: _model = model.module if hasattr(model, "module") else model clip_enc = args.grad_clip_enc if args.grad_clip_enc > 0 else float('inf') clip_dec = args.grad_clip_dec if args.grad_clip_dec > 0 else float('inf') grad_norm_enc = float(torch.nn.utils.clip_grad_norm_( _model.encoder.parameters(), max_norm=clip_enc, ).item()) grad_norm_dec = float(torch.nn.utils.clip_grad_norm_( _model.decoder.parameters(), max_norm=clip_dec, ).item()) grad_norm = (grad_norm_enc ** 2 + grad_norm_dec ** 2) ** 0.5 # warmup if args.warmup_steps > 0 and step < args.warmup_steps: warmup_frac = (step + 1) / args.warmup_steps for pg in opt.param_groups: pg["lr"] = args.lr * warmup_frac opt.step() opt.zero_grad(set_to_none=True) # ---- logging ---- # Flush the per-domain MSE accumulator on every rank at the log # boundary (all_reduce is collective). Only rank 0 logs to wandb. domain_mse_means = None if step % args.log_every == 0: domain_mse_means, domain_mse_counts = domain_acc.flush(ddp=ddp) if is_rank0() and (step % args.log_every == 0): psnr = 10.0 * torch.log10(1.0 / mse.clamp_min(1e-10)) _model = model.module if hasattr(model, "module") else model weight_norm_enc = float(sum(p.float().norm().item() ** 2 for p in _model.encoder.parameters()) ** 0.5) weight_norm_dec = float(sum(p.float().norm().item() ** 2 for p in _model.decoder.parameters()) ** 0.5) loss_mse_normed = float(rms_mse.normalize(mse).item()) loss_lp_normed = float((args.lpips_weight * rms_lp.normalize(lp)).item()) if lpips_fn is not None else 0.0 log_dict = { "loss/total": float(loss.item()), "loss/mse": float(mse.item()), "loss/lpips": float(lp.item()), "loss/mse_normed": loss_mse_normed, "loss/lpips_normed": loss_lp_normed, "stats/psnr": float(psnr.item()), "stats/keep_prob": float(keep_prob.mean().item()), "stats/masked_frac": float(mae_mask.float().mean().item()), "lr": float(opt.param_groups[0]["lr"]), "stats/grad_norm": grad_norm, "stats/grad_norm_enc": grad_norm_enc, "stats/grad_norm_dec": grad_norm_dec, "stats/weight_norm_enc": weight_norm_enc, "stats/weight_norm_dec": weight_norm_dec, "time/hrs": (time.time() - t0) / 3600.0, "rms/mse": rms_mse.rms_val, "rms/lp": rms_lp.rms_val, } if domain_mse_means is not None: for i, name in enumerate(DOMAINS): cnt = float(domain_mse_counts[i].item()) if cnt <= 0: continue mse_i = float(domain_mse_means[i].item()) if not np.isfinite(mse_i): continue psnr_i = 10.0 * np.log10(1.0 / max(mse_i, 1e-10)) log_dict[f"domain/{name}/mse"] = mse_i log_dict[f"domain/{name}/psnr"] = float(psnr_i) wandb.log(log_dict, step=step) if is_rank0() and (step % args.print_every == 0): psnr = 10.0 * torch.log10(1.0 / mse.clamp_min(1e-10)) print( f"step {step:07d} | loss={loss.item():.6f} " f"| mse={mse.item():.6f} | lpips={lp.item():.4f} " f"| psnr={psnr.item():.2f} | keep={keep_prob.mean().item():.3f} | gnorm={grad_norm:.3f}" ) # ---- viz ---- if is_rank0() and args.viz_every > 0 and (step % args.viz_every == 0): log_tokenizer_viz_wandb( x_btchw=x, pred_btnd=pred, mae_mask_btNp1=mae_mask, patch=args.patch, step=step, max_items=args.viz_max_items, max_T=args.viz_max_T, ) # ---- validation ---- if val_loader is not None and args.val_every > 0 and step > 0 and (step % args.val_every == 0): val_metrics, val_x, val_pred = validate( model=model, val_loader=val_loader, device=device, use_amp=use_amp, lpips_fn=lpips_fn, args=args, max_batches=args.val_batches, task_idx_to_domain_idx=task_idx_to_domain_idx, key_prefix="val", ) if is_rank0() and val_metrics: wandb.log(val_metrics, step=step) print( f"step {step:07d} | VAL mse={val_metrics['val/mse']:.6f} " f"| psnr={val_metrics['val/psnr']:.2f}" + (f" | lpips={val_metrics['val/lpips']:.4f}" if 'val/lpips' in val_metrics else "") ) if val_x is not None and args.viz_every > 0: log_tokenizer_val_viz_wandb( x_btchw=val_x, pred_btnd=val_pred, patch=args.patch, step=step, max_items=args.viz_max_items, max_T=args.viz_max_T, tag="val/viz", ) # ---- UNSEEN validation (optional) ---- if val_unseen_loader is not None: val_un_metrics, val_un_x, val_un_pred = validate( model=model, val_loader=val_unseen_loader, device=device, use_amp=use_amp, lpips_fn=lpips_fn, args=args, max_batches=args.val_batches, task_idx_to_domain_idx=val_unseen_task_idx_to_domain_idx, local_task_names=val_unseen_local_tasks, key_prefix="val_unseen", ) if is_rank0() and val_un_metrics: wandb.log(val_un_metrics, step=step) print( f"step {step:07d} | VAL_UNSEEN mse={val_un_metrics['val_unseen/mse']:.6f} " f"| psnr={val_un_metrics['val_unseen/psnr']:.2f}" + (f" | lpips={val_un_metrics['val_unseen/lpips']:.4f}" if 'val_unseen/lpips' in val_un_metrics else "") ) if val_un_x is not None and args.viz_every > 0: log_tokenizer_val_viz_wandb( x_btchw=val_un_x, pred_btnd=val_un_pred, patch=args.patch, step=step, max_items=args.viz_max_items, max_T=args.viz_max_T, tag="val_unseen/viz", ) # ---- ckpt ---- if is_rank0() and args.save_every > 0 and (step % args.save_every == 0) and step > 0 and do_step: rms_state = {"mse": rms_mse.state_dict(), "lp": rms_lp.state_dict()} ckpt_path = ckpt_dir / f"step_{step:07d}.pt" save_ckpt(ckpt_path, step=step, epoch=epoch, model=model, opt=opt, args=args, rms_state=rms_state) # also update a "latest" pointer latest = ckpt_dir / "latest.pt" save_ckpt(latest, step=step, epoch=epoch, model=model, opt=opt, args=args, rms_state=rms_state) # log checkpoint as a wandb artifact (versioned, easy to retrieve from the UI) artifact = wandb.Artifact( name=f"{wandb.run.id}-tokenizer", type="model", metadata={"step": step, "epoch": epoch}, ) artifact.add_file(str(ckpt_path), name=f"step_{step:07d}.pt") wandb.log_artifact(artifact, aliases=["latest", f"step-{step}"]) step += 1 start_epoch = epoch + 1 if step > args.max_steps: break # exit the epoch loop too; otherwise the outer `while` is never re-checked and the dataloader churns idle. if ddp: dist.barrier() dist.destroy_process_group() if __name__ == "__main__": p = argparse.ArgumentParser() # data p.add_argument("--data_dirs", type=str, nargs="+", default=[ "./data/expert-shards", "./data/mixed-large-shards", "./data/mixed-small-shards", "./data/zeros-shards", ]) p.add_argument("--val_data_dir", type=str, default="./data/val-shards", help="optional single directory of preprocessed shards used for validation") p.add_argument("--val_unseen_data_dir", type=str, default=None, help="optional directory of preprocessed shards for UNSEEN-task validation. " "Filtered to UNSEEN_TASK_SET. Logs a parallel `val_unseen/...` namespace " "with per-task PSNR and a side-by-side viz panel under `val_unseen/viz`.") p.add_argument("--val_every", type=int, default=2_000, help="run validation every N steps (0 disables)") p.add_argument("--val_batches", type=int, default=8, help="number of batches per validation pass (per rank)") p.add_argument("--seq_len", type=int, default=24) p.add_argument("--num_workers", type=int, default=6) p.add_argument("--shard_cache_size", type=int, default=12, help="number of shards held in each worker's LRU cache") p.add_argument("--prefetch_factor", type=int, default=4, help="batches prefetched per worker") p.add_argument("--samples_per_shard", type=int, default=16, help="number of sequences drawn from a loaded shard before " "picking a new one (1 = pure iid, larger = much less I/O)") p.add_argument("--val_samples_per_shard", type=int, default=1, help="samples_per_shard for the val loader; defaults to 1 " "for maximum task diversity in val/viz (val is small " "and runs infrequently, so I/O cost is bounded)") p.add_argument("--task_weighting", type=str, default="valid_starts", choices=["valid_starts", "uniform", "targeted"], help="Per-task sampling weighting. 'valid_starts' (default) " "makes P(task) ∝ total valid_starts — skews toward " "long-trajectory domains. 'uniform' " "gives every task equal draw probability. 'targeted' " "splits α uniformly over the targeted-collection tasks " "(SEEN_TASK_SET ∪ UNSEEN_TASK_SET) and (1-α) uniformly " "over the rest.") p.add_argument("--targeted_alpha", type=float, default=0.5, help="(--task_weighting=targeted only) total per-batch " "probability allocated to the targeted-collection tasks.") p.add_argument("--batch_size", type=int, default=12) # image / patching p.add_argument("--H", type=int, default=224) p.add_argument("--W", type=int, default=224) p.add_argument("--C", type=int, default=3) p.add_argument("--patch", type=int, default=14) # model p.add_argument("--d_model", type=int, default=512) p.add_argument("--n_heads", type=int, default=8) p.add_argument("--depth", type=int, default=12) p.add_argument("--n_latents", type=int, default=64) p.add_argument("--d_bottleneck", type=int, default=64) p.add_argument("--dropout", type=float, default=0.05) p.add_argument("--mlp_ratio", type=float, default=4.0) p.add_argument("--time_every", type=int, default=1) p.add_argument("--mae_p_min", type=float, default=0.0) p.add_argument("--mae_p_max", type=float, default=0.9) # optim p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--weight_decay", type=float, default=1e-2) p.add_argument("--max_steps", type=int, default=10_000_000) p.add_argument("--grad_accum", type=int, default=1) p.add_argument("--grad_clip_enc", type=float, default=1.0) p.add_argument("--grad_clip_dec", type=float, default=1.0) p.add_argument("--warmup_steps", type=int, default=0) # lpips p.add_argument("--lpips_weight", type=float, default=0.2) p.add_argument("--lpips_frac", type=float, default=0.5) p.add_argument("--lpips_net", type=str, default="alex", choices=["alex", "vgg", "squeeze"]) # logging / viz p.add_argument("--log_every", type=int, default=200) p.add_argument("--print_every", type=int, default=200) p.add_argument("--viz_every", type=int, default=1_000) p.add_argument("--viz_max_items", type=int, default=4) p.add_argument("--viz_max_T", type=int, default=8) # wandb p.add_argument("--wandb_project", type=str, default="mmbench2-tokenizer") p.add_argument("--wandb_run_name", type=str, default="default") p.add_argument("--wandb_entity", type=str, default=None) # ckpt p.add_argument("--ckpt_dir", type=str, default="./logs/tokenizer_ckpts") p.add_argument("--save_every", type=int, default=10_000) p.add_argument("--resume", type=str, default=None) # misc p.add_argument("--seed", type=int, default=0) p.add_argument("--compile", action="store_true") train(p.parse_args())