"""Standalone eight-GPU Bridge Stage-1 CTE training, with no policy/release CTE. The released effect-v3 objective is used, not an exact reproduction of the paper's MSE objective. See docs/bridge_cte_stage1.md for the full contract. """ from __future__ import annotations import argparse import hashlib import json import os import random import shutil import signal import sys import time import uuid from dataclasses import asdict from datetime import datetime, timedelta from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader from cosmos_framework.model.zeva.causal_transition_encoder import CausalTransitionEncoder, CausalTransitionEncoderConfig from cosmos_framework.model.zeva.cte_losses import CTELossConfig from scripts.bridge_cte_data import BridgeCTEDataset, PairedEpisodeSampler, ValidationSampler, collate_episodes, load_metadata from scripts.bridge_cte_loss import bridge_cte_loss from scripts.widowx_contract import digest_file, fingerprint def atomic_json(path, payload): path = Path(path) temporary = path.with_name(path.name + f".{os.getpid()}.tmp") temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n") temporary.replace(path) def model_digest(model): digest = hashlib.sha256() for name, tensor in model.state_dict().items(): digest.update(name.encode()) digest.update(tensor.detach().cpu().contiguous().numpy().tobytes()) return digest.hexdigest() class FrozenFrameEncoder: def __init__(self, path, device, microbatch=4): from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import WanVAE self.vae = WanVAE(vae_pth=str(path), device=device, dtype=torch.bfloat16, temporal_window=4, encode_exact_durations=[1]) self.device, self.microbatch = device, microbatch if any(p.requires_grad for p in self.vae.model.parameters()): raise AssertionError("VAE must be frozen") @torch.no_grad() def __call__(self, batch): valid = batch["valid"].to(self.device, non_blocking=True) latents = [] # Each single-frame encode resets the VAE's temporal cache. This is # intentionally the same operation online and offline, with no future # frames; resizing is not needed for Bridge's native 256x256 images. for chunk in batch["frames"].split(self.microbatch): pixels = chunk.to(self.device, dtype=torch.bfloat16, non_blocking=True) / 127.5 - 1 latent = self.vae.encode(pixels.unsqueeze(2)).squeeze(2) if tuple(latent.shape[1:]) != (48, 16, 16): raise ValueError(f"Unexpected VAE output: {latent.shape}") latents.append(latent) values = torch.cat(latents) if not torch.isfinite(values).all(): raise FloatingPointError("Nonfinite frozen VAE features") padded = values.new_zeros((*valid.shape, 48, 16, 16)) padded[valid] = values return padded def init_wandb(args, config): if args.wandb_mode == "disabled": return None import wandb id_path = args.output / "wandb_run_id.json" run_id = json.loads(id_path.read_text())["id"] if id_path.exists() else uuid.uuid4().hex[:8] atomic_json(id_path, {"id": run_id}) kwargs = dict(project=args.wandb_project, entity=os.environ.get("WANDB_ENTITY"), name=args.output.name, group="bridge_cte_stage1", id=run_id, dir=str(args.output), config=config, resume="allow", settings=wandb.Settings(init_timeout=20, base_url=os.environ.get("WANDB_BASE_URL", "https://api.wandb.ai"))) try: run = wandb.init(mode=args.wandb_mode, **kwargs) except wandb.errors.Error as error: if args.wandb_mode != "online": raise print(f"WANDB_OFFLINE_FALLBACK error_type={type(error).__name__}", flush=True) wandb.teardown() run = wandb.init(mode="offline", **kwargs) atomic_json(args.output / "wandb_status.json", dict(mode=run.settings.mode, project=run.project, entity=run.entity, id=run.id, directory=run.dir)) return run def build_loader(dataset, sampler, args, rank): # Private loader RNG: constructing/restarting a loader never consumes the # model's random stream. There are no stochastic data augmentations. generator = torch.Generator().manual_seed(args.seed + 10000 + rank) kwargs = dict(dataset=dataset, batch_sampler=sampler, collate_fn=collate_episodes, num_workers=args.workers, pin_memory=True, generator=generator) if args.workers: kwargs.update(prefetch_factor=2, multiprocessing_context="spawn") return DataLoader(**kwargs) def move_batch(batch, device): return {key: value.to(device, non_blocking=True) for key, value in batch.items() if key != "frames"} def reduce_metrics(metrics, device): keys = sorted(metrics) values = torch.stack([metrics[k].detach().float() for k in keys]).to(device) if dist.is_initialized(): dist.all_reduce(values) values /= dist.get_world_size() return dict(zip(keys, values.cpu().tolist(), strict=True)) @torch.no_grad() def validate(model, encoder, dataset, args, world, rank, device): model.eval() sampler = ValidationSampler(dataset.episodes, args.batch_size, world, rank, args.seed + 70000) loader = build_loader(dataset, sampler, args, rank) totals, batches = {}, 0 for raw in loader: inputs = encoder(raw) batch = move_batch(raw, device) with torch.autocast("cuda", dtype=torch.bfloat16): output = model(inputs, batch["actions"], batch["valid"], batch["transition_valid"]) losses = bridge_cte_loss(output, batch["actions"], batch["valid"], batch["semantic_ids"], batch["episode_ids"]) for key, value in losses.items(): totals[key] = totals.get(key, 0) + value.detach() batches += 1 metrics = reduce_metrics({key: value / batches for key, value in totals.items()}, device) model.train() return metrics def rng_state(): return dict(python=random.getstate(), numpy=np.random.get_state(), torch=torch.get_rng_state(), cuda=torch.cuda.get_rng_state()) def restore_rng(state): random.setstate(state["python"]) np.random.set_state(state["numpy"]) torch.set_rng_state(state["torch"]) torch.cuda.set_rng_state(state["cuda"]) def save_checkpoint(model, optimizer, args, contract, step, epoch, batch_cursor, world, rank): states = [None] * world if dist.is_initialized(): dist.all_gather_object(states, rng_state()) else: states[0] = rng_state() if rank == 0: directory = args.output / "checkpoints" directory.mkdir(exist_ok=True) path = directory / f"cte_step_{step:09d}.pt" if not path.exists(): if shutil.disk_usage(directory).free < args.min_free_gb * 2**30: raise OSError("Disk safety floor reached; no checkpoint is automatically deleted") payload = dict(format="bridge_cte_stage1_v1", model_config=model.cfg.to_dict(), model=model.state_dict(), optimizer=optimizer.state_dict(), step=step, epoch=epoch, batch_cursor=batch_cursor, rng_states=states, contract=contract, contract_fingerprint=fingerprint(contract), training_args=vars(args) | {"output": str(args.output)}) temporary = path.with_suffix(f".{os.getpid()}.tmp") torch.save(payload, temporary) with temporary.open("rb") as stream: os.fsync(stream.fileno()) temporary.replace(path) atomic_json(args.output / "latest_checkpoint.json", dict(path=str(path), step=step, epoch=epoch, batch_cursor=batch_cursor)) print(f"CHECKPOINT_SAVED step={step} path={path} bytes={path.stat().st_size}", flush=True) if dist.is_initialized(): dist.barrier() def parse_args(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset", default="/home/gpu4/tianyi/TTTVLA/datasets/bridge_orig_1.0.0_lerobot") parser.add_argument("--vae", default="/home/gpu4/tianyi/models/wan22_vae/Wan2.2_VAE.pth") parser.add_argument("--output", type=Path, required=True) parser.add_argument("--resume", help="Explicit Stage-1 checkpoint, never a RoboCasa CTE") parser.add_argument("--batch-size", type=int, default=8, help="Episodes per GPU, including paired partners") parser.add_argument("--workers", type=int, default=2) parser.add_argument("--vae-microbatch", type=int, default=4) parser.add_argument("--lr", type=float, default=0.0001) parser.add_argument("--warmup-steps", type=int, default=500) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--max-steps", type=int, default=2_147_483_647) parser.add_argument("--save-every", type=int, default=10000) parser.add_argument("--validate-every", type=int, default=2000) parser.add_argument("--log-every", type=int, default=10) parser.add_argument("--val-fraction", type=float, default=0.02) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--memory-fraction", type=float, default=0.125, help="PyTorch allocator budget as a fraction of total GPU RAM") parser.add_argument("--min-free-gb", type=float, default=2.0) parser.add_argument("--wandb-mode", choices=("online", "offline", "disabled"), default=os.environ.get("WANDB_MODE", "online")) parser.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "Zeva")) parser.add_argument("--skip-initial-validation", action="store_true", help="Smoke checks only") args = parser.parse_args() if min(args.batch_size, args.vae_microbatch, args.warmup_steps, args.max_steps, args.save_every, args.validate_every, args.log_every) < 1 or args.batch_size % 2: parser.error("Positive intervals and an even per-GPU batch size are required") if not 0 < args.memory_fraction <= 1 or not 0 < args.val_fraction < 1: parser.error("Invalid memory or validation fraction") return args def main(): args = parse_args() rank, world, local_rank = (int(os.environ.get(k, v)) for k, v in (("RANK", 0), ("WORLD_SIZE", 1), ("LOCAL_RANK", 0))) torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) torch.cuda.set_per_process_memory_fraction(args.memory_fraction, device) free, total = torch.cuda.mem_get_info(device) if free < total * args.memory_fraction + 2 * 2**30: raise RuntimeError(f"GPU {local_rank} has insufficient free VRAM for budget + 2 GiB headroom") if world > 1: dist.init_process_group("nccl", device_id=device, timeout=timedelta(minutes=10)) torch.set_num_threads(2) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.benchmark = False random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) info, train, val = load_metadata(args.dataset, args.seed, args.val_fraction) cfg = CausalTransitionEncoderConfig(action_dim=7, image_channels=48) model = CausalTransitionEncoder(cfg).to(device) # Unused cross_norm modules exist in the public architecture. DDP detects # those without silently freezing any actually used CTE parameters. optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr, betas=(0.9, 0.95), weight_decay=args.weight_decay) dataset = BridgeCTEDataset(args.dataset, info, train) val_dataset = BridgeCTEDataset(args.dataset, info, val) base_sampler = PairedEpisodeSampler(train, args.batch_size, world, rank, args.seed) files = [ROOT / "scripts" / n for n in ("train_bridge_cte.py", "bridge_cte_data.py", "bridge_cte_loss.py")] files += [ROOT / "cosmos_framework/model/zeva" / n for n in ("causal_transition_encoder.py", "cte_losses.py")] contract = dict(version="bridge_cte_stage1_v1", dataset=str(Path(args.dataset).resolve()), metadata={n: digest_file(Path(args.dataset) / "meta" / n) for n in ("info.json", "episodes.jsonl", "tasks.jsonl")}, vae_path=str(Path(args.vae).resolve()), vae_sha256=digest_file(args.vae), code_sha256={str(p.relative_to(ROOT)): digest_file(p) for p in files}, action_contract="bridge_q01_q99_xyzrpy_clipped_gripper_open_positive_v1", frame_encoding="single_frame_rgb_native256x256_wan22_bf16_v1", image_key="observation.images.image_0", boundary_stride=4, effect_raw_controls=16, fps=5, full_trajectories=True, cte_initialization="random_no_release_weights", loss="release_effect_v3_mse_diagnostic_only_bridge_masked_global_task_v1", model_config=cfg.to_dict(), loss_config=asdict(CTELossConfig()), seed=args.seed, val_fraction=args.val_fraction, world_size=world, per_rank_batch=args.batch_size, lr=args.lr, warmup_steps=args.warmup_steps, weight_decay=args.weight_decay, sampler="exhaustive_anchors_paired_partner_v1") args.output = args.output.resolve() if rank == 0: args.output.mkdir(parents=True, exist_ok=True) if (args.output / "manifest.json").exists() and not args.resume: raise FileExistsError("Run already exists; explicitly --resume or choose a new output directory") if shutil.disk_usage(args.output).free < args.min_free_gb * 2**30 + 1024**3: raise OSError("Not enough disk headroom for a new CTE run") if dist.is_initialized(): dist.barrier() step, epoch, cursor = 0, 0, 0 resume_rng = None initialization_hash = model_digest(model) if rank == 0 else None if args.resume: checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False) if checkpoint.get("format") != "bridge_cte_stage1_v1" or checkpoint.get("contract_fingerprint") != fingerprint(contract): raise ValueError("Checkpoint/data/code/batch/optimizer contract differs; refusing ambiguous resume") model.load_state_dict(checkpoint["model"], strict=True) optimizer.load_state_dict(checkpoint["optimizer"]) step, epoch, cursor = checkpoint["step"], checkpoint["epoch"], checkpoint["batch_cursor"] resume_rng = checkpoint["rng_states"][rank] del checkpoint config = dict(contract=contract, args=vars(args) | {"output": str(args.output)}, train_episodes=len(train), validation_episodes=len(val), train_unknown_instruction=sum(e["semantic_id"] < 0 for e in train), steps_per_anchor_epoch=base_sampler.total_batches, global_episode_batch=args.batch_size * world, trainable_parameters=sum(p.numel() for p in model.parameters() if p.requires_grad), initial_random_model_sha256=initialization_hash) if rank == 0: atomic_json(args.output / "manifest.json", config) atomic_json(args.output / "split.json", dict(train=[e["episode_index"] for e in train], validation=[e["episode_index"] for e in val])) logger = init_wandb(args, config) else: logger = None if dist.is_initialized(): dist.barrier() encoder = FrozenFrameEncoder(args.vae, device, args.vae_microbatch) ddp = DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True, broadcast_buffers=False) if world > 1 else model if resume_rng is not None: restore_rng(resume_rng) print(f"RANK_READY rank={rank} pid={os.getpid()} cuda={local_rank} trainable={config['trainable_parameters']} " f"resume_step={step} max_allocated_GiB={torch.cuda.max_memory_allocated()/2**30:.3f}", flush=True) def record(kind, values): if rank != 0: return record = dict(kind=kind, timestamp=datetime.now().astimezone().isoformat(), step=step, anchor_epoch=epoch + cursor / base_sampler.total_batches, **values) with (args.output / "metrics.jsonl").open("a") as stream: stream.write(json.dumps(record, allow_nan=False) + "\n") atomic_json(args.output / "status.json", record) if logger is not None: logger.log({f"{kind}/{k}": v for k, v in values.items()} | {"global_step": step}, step=step) print(json.dumps(record, allow_nan=False), flush=True) if not args.resume: save_checkpoint(model, optimizer, args, contract, step, epoch, cursor, world, rank) if not args.skip_initial_validation: record("validation", validate(model, encoder, val_dataset, args, world, rank, device)) requested_stop = False def handle_stop(signum, frame): nonlocal requested_stop requested_stop = True print(f"STOP_REQUESTED rank={rank} signal={signum}; finishing current update and saving", flush=True) signal.signal(signal.SIGTERM, handle_stop) signal.signal(signal.SIGINT, handle_stop) model.train() last_log_time, last_log_step = time.perf_counter(), step while step < args.max_steps: sampler = PairedEpisodeSampler(train, args.batch_size, world, rank, args.seed, epoch, cursor) loader = build_loader(dataset, sampler, args, rank) for raw in loader: optimizer.zero_grad(set_to_none=True) inputs = encoder(raw) batch = move_batch(raw, device) with torch.autocast("cuda", dtype=torch.bfloat16): output = ddp(inputs, batch["actions"], batch["valid"], batch["transition_valid"]) losses = bridge_cte_loss(output, batch["actions"], batch["valid"], batch["semantic_ids"], batch["episode_ids"]) if not torch.isfinite(losses["total"]): raise FloatingPointError(f"Nonfinite loss at step {step + 1}, rank {rank}") losses["total"].backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True) lr = args.lr * min(1.0, (step + 1) / args.warmup_steps) for group in optimizer.param_groups: group["lr"] = lr optimizer.step() model.update_ema_target() step, cursor = step + 1, cursor + 1 if step <= 3 or step % args.log_every == 0: values = reduce_metrics(losses | {"grad_norm": grad_norm}, device) now = time.perf_counter() values.update(lr=lr, seconds_per_step=(now-last_log_time)/(step-last_log_step), gpu_max_allocated_gb=torch.cuda.max_memory_allocated()/2**30, gpu_reserved_gb=torch.cuda.memory_reserved()/2**30) record("train", values) last_log_time, last_log_step = now, step stop_flag = torch.tensor(int(requested_stop), device=device) if dist.is_initialized(): dist.all_reduce(stop_flag, op=dist.ReduceOp.MAX) stop = bool(stop_flag.item()) or step >= args.max_steps if step % args.save_every == 0 or stop: save_checkpoint(model, optimizer, args, contract, step, epoch, cursor, world, rank) if step % args.validate_every == 0 and not stop: record("validation", validate(model, encoder, val_dataset, args, world, rank, device)) # Validation time is not reported as optimizer-step throughput. last_log_time, last_log_step = time.perf_counter(), step if stop: record("stopped", dict(checkpoint_step=step, reason="signal" if stop_flag.item() else "max_steps")) if logger is not None: logger.finish() if dist.is_initialized(): dist.destroy_process_group() return epoch, cursor = epoch + 1, 0 if rank == 0: print(f"ANCHOR_EPOCH_COMPLETE epoch={epoch} train_episodes={len(train)} all_anchors_seen=true", flush=True) if __name__ == "__main__": main()