Spaces:
Running on Zero
Running on Zero
| from pathlib import Path | |
| import time | |
| import glob | |
| import subprocess | |
| import os | |
| import re | |
| from typing import Optional | |
| from torch.distributed.checkpoint.state_dict import ( | |
| get_model_state_dict, | |
| get_optimizer_state_dict, | |
| set_model_state_dict, | |
| StateDictOptions, | |
| ) | |
| from torch.distributed.fsdp import FSDPModule | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| MODEL_CHECKPOINT = "model_state_dict.pt" | |
| OPTIMIZER_CHECKPOINT = "optim_state_dict.pt" | |
| STATE = "state.pt" | |
| def readable_timestamp(): | |
| """Generate a sortable timestamp for filenames (no weekday).""" | |
| return time.strftime("%Y_%m_%d_%H_%M_%S") | |
| def find_latest_checkpoint(base_dir, model_name, run_root_dir: Optional[str] = None, stage_name: Optional[str] = None): | |
| """Find latest checkpoint. | |
| If run_root_dir (and optional stage_name) are provided, search only under | |
| <run_root_dir>/<stage_name>/checkpoints (or <run_root_dir>/**/checkpoints if stage_name None). | |
| Otherwise, fall back to project-wide model type-based search. | |
| Newest run dir first, then highest step within it. | |
| If the newest run root has none for this model, keep searching older runs. | |
| """ | |
| def collect_checkpoint_paths(roots, model_name): | |
| alias_map = { | |
| 'video_tokenizer': ['video_tokenizer'], | |
| 'latent_actions': ['latent_actions', 'lam', 'actions', 'action_tokenizer'], | |
| 'dynamics': ['dynamics'], | |
| } | |
| aliases = alias_map.get(model_name, [model_name]) | |
| candidates = [] | |
| seen = set() | |
| def add_candidate(path: str) -> None: | |
| norm = os.path.normpath(path) | |
| if norm in seen: | |
| return | |
| if os.path.isdir(norm): | |
| # require at least state or model file | |
| state_file = os.path.join(norm, STATE) | |
| model_file = os.path.join(norm, MODEL_CHECKPOINT) | |
| if os.path.isfile(state_file) or os.path.isfile(model_file): | |
| candidates.append(norm) | |
| seen.add(norm) | |
| else: | |
| _, ext = os.path.splitext(norm) | |
| if ext in ('.pt', '.pth'): | |
| candidates.append(norm) | |
| seen.add(norm) | |
| patterns = [] | |
| for alias in aliases: | |
| patterns.append(f"*{alias}_step_*") | |
| patterns.append(f"*{alias}_checkpoint_*") | |
| for root in roots: | |
| for pat in patterns: | |
| search_pattern = os.path.join(root, "**", pat) | |
| for match in glob.glob(search_pattern, recursive=True): | |
| add_candidate(match) | |
| return candidates | |
| def run_dir_of(path: str) -> str: | |
| # Walk up until we reach the 'checkpoints' directory, then return its parent (the stage dir) | |
| d = path if os.path.isdir(path) else os.path.dirname(path) | |
| while d and os.path.basename(d) != 'checkpoints': | |
| parent = os.path.dirname(d) | |
| if parent == d: | |
| break | |
| d = parent | |
| # If we found 'checkpoints', return its parent; otherwise, fallback to two-level up | |
| if d and os.path.basename(d) == 'checkpoints': | |
| return os.path.dirname(d) | |
| # Fallback: use the directory containing the checkpoint path (or its parent) | |
| candidate_dir = path if os.path.isdir(path) else os.path.dirname(path) | |
| return candidate_dir | |
| def project_wide_search(): | |
| # Fallback to model_type/results search in repository | |
| # Allow multiple possible directory names per model | |
| model_type_dirs = { | |
| 'video_tokenizer': ['video_tokenizer'], | |
| 'latent_actions': ['latent_actions', 'lam', 'actions', 'action_tokenizer'], | |
| 'dynamics': ['dynamics',], | |
| } | |
| dirs = model_type_dirs.get(model_name) | |
| if not dirs: | |
| roots = [os.path.join(base_dir, 'results', '**', 'checkpoints')] | |
| else: | |
| roots = [os.path.join(base_dir, 'results', '**', d, 'checkpoints') for d in dirs] | |
| files = collect_checkpoint_paths(roots, model_name) | |
| if not files: | |
| # Generic fallback: search all checkpoints regardless of stage dir name | |
| generic_roots = [os.path.join(base_dir, 'results', '**', 'checkpoints')] | |
| files = collect_checkpoint_paths(generic_roots, model_name) | |
| if not files: | |
| raise Exception(f"No checkpoint files found for {model_name}") | |
| run_dir_to_files = {} | |
| for p in files: | |
| rd = run_dir_of(p) | |
| run_dir_to_files.setdefault(rd, []).append(p) | |
| newest_run_dir = max(run_dir_to_files.keys(), key=lambda d: os.path.getctime(d)) | |
| candidate_files = run_dir_to_files[newest_run_dir] | |
| return candidate_files | |
| if run_root_dir is not None: | |
| if stage_name: | |
| roots = [os.path.join(run_root_dir, stage_name, 'checkpoints')] | |
| else: | |
| roots = [os.path.join(run_root_dir, '**', 'checkpoints')] | |
| files = collect_checkpoint_paths(roots, model_name) | |
| if not files: | |
| # Fallback: search project-wide older runs until found | |
| candidate_files = project_wide_search() | |
| else: | |
| # Group by run dir within the provided root and choose newest run dir | |
| run_dir_to_files = {} | |
| for p in files: | |
| rd = run_dir_of(p) | |
| run_dir_to_files.setdefault(rd, []).append(p) | |
| newest_run_dir = max(run_dir_to_files.keys(), key=lambda d: os.path.getctime(d)) | |
| candidate_files = run_dir_to_files[newest_run_dir] | |
| else: | |
| candidate_files = project_wide_search() | |
| def extract_step(path: str) -> int: | |
| fname = os.path.basename(path) | |
| m = re.search(r"_step_(\d+)", fname) | |
| return int(m.group(1)) if m else -1 | |
| candidate_files.sort(key=lambda p: (extract_step(p), os.path.getctime(p))) | |
| return candidate_files[-1] | |
| def run_command(cmd, description): | |
| # empirical max dataLoader throughput settings (I used 1-6 H100s) | |
| env = os.environ.copy() | |
| env.setdefault("NG_NUM_WORKERS", str(max(2, (os.cpu_count() or 4) - 2))) | |
| env.setdefault("NG_PREFETCH_FACTOR", "4") | |
| env.setdefault("NG_PIN_MEMORY", "1") | |
| env["NG_PERSISTENT_WORKERS"] = "0" | |
| env.setdefault("TORCH_CUDNN_V8_API_ENABLED", "1") | |
| try: | |
| result = subprocess.run(cmd, check=True, capture_output=False, env=env) | |
| return True | |
| except subprocess.CalledProcessError as e: | |
| print(f"Error: {e.stderr}") | |
| return False | |
| except KeyboardInterrupt: | |
| return False | |
| def save_training_state(model, optimizer, scheduler, config, checkpoints_dir, prefix, step): | |
| """Save a checkpoint with model/optimizer/scheduler and the exact config. | |
| The filename includes the global step and a timestamp for uniqueness. | |
| """ | |
| import torch | |
| ts = readable_timestamp() | |
| if isinstance(model, (FSDPModule, DDP)): | |
| state_dict = get_model_state_dict( | |
| model=model, | |
| options=StateDictOptions( | |
| full_state_dict=True, | |
| cpu_offload=True, | |
| ), | |
| ) | |
| optimizer_state_dict = get_optimizer_state_dict( | |
| model=model, | |
| optimizers=optimizer, | |
| options=StateDictOptions( | |
| full_state_dict=True, | |
| cpu_offload=True, | |
| ), | |
| ) | |
| else: | |
| # Avoid saving the model with _orig_mod prefix if it's compiled | |
| state_dict = getattr(model, '_orig_mod', model).state_dict() | |
| optimizer_state_dict = optimizer.state_dict() | |
| state = { | |
| 'scheduler_state_dict': scheduler.state_dict() if scheduler is not None else None, | |
| 'config': config, | |
| 'step': int(step) if step is not None else None, | |
| 'timestamp': ts, | |
| } | |
| os.makedirs(checkpoints_dir, exist_ok=True) | |
| ckpt_path = os.path.join(checkpoints_dir, f"{prefix}_step_{int(step) if step is not None else 0}") | |
| os.makedirs(ckpt_path, exist_ok=True) | |
| torch.save(state_dict, Path(ckpt_path) / MODEL_CHECKPOINT) | |
| torch.save(optimizer_state_dict, Path(ckpt_path) / OPTIMIZER_CHECKPOINT) | |
| torch.save(state, Path(ckpt_path) / STATE) | |
| return ckpt_path | |
| def load_videotokenizer_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False): | |
| """Instantiate VideoTokenizer from a checkpoint's saved config and load weights.""" | |
| import torch | |
| from models.video_tokenizer import VideoTokenizer | |
| model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True) | |
| state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False) | |
| cfg = state_cfg.get('config', {}) or {} | |
| frame_size = cfg.get('frame_size', 128) | |
| kwargs = { | |
| 'frame_size': (frame_size, frame_size), | |
| 'patch_size': cfg.get('patch_size', 8), | |
| 'embed_dim': cfg.get('embed_dim', 128), | |
| 'num_heads': cfg.get('num_heads', 8), | |
| 'hidden_dim': cfg.get('hidden_dim', 256), | |
| 'num_blocks': cfg.get('num_blocks', 4), | |
| 'latent_dim': cfg.get('latent_dim', 6), | |
| 'num_bins': cfg.get('num_bins', 4), | |
| } | |
| if model is None: | |
| model = VideoTokenizer(**kwargs) | |
| set_model_state_dict( | |
| model=model, | |
| model_state_dict=model_sd, | |
| options=StateDictOptions( | |
| full_state_dict=True, | |
| broadcast_from_rank0=is_distributed, | |
| ), | |
| ) | |
| model = model.to(device) | |
| return model, state_cfg | |
| def load_latent_actions_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False): | |
| """Instantiate LatentActionModel from a checkpoint's saved config and load weights.""" | |
| import torch | |
| from models.latent_actions import LatentActionModel | |
| model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True) | |
| state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False) | |
| cfg = state_cfg.get('config', {}) or {} | |
| frame_size = cfg.get('frame_size', 128) | |
| kwargs = { | |
| 'frame_size': (frame_size, frame_size), | |
| 'n_actions': cfg.get('n_actions', 8), | |
| 'patch_size': cfg.get('patch_size', 8), | |
| 'embed_dim': cfg.get('embed_dim', 128), | |
| 'num_heads': cfg.get('num_heads', 8), | |
| 'hidden_dim': cfg.get('hidden_dim', 256), | |
| 'num_blocks': cfg.get('num_blocks', 4), | |
| } | |
| if model is None: | |
| model = LatentActionModel(**kwargs) | |
| set_model_state_dict( | |
| model=model, | |
| model_state_dict=model_sd, | |
| options=StateDictOptions( | |
| full_state_dict=True, | |
| broadcast_from_rank0=is_distributed, | |
| ), | |
| ) | |
| model = model.to(device) | |
| return model, state_cfg | |
| def load_dynamics_from_checkpoint(checkpoint_path, device, model = None, is_distributed = False): | |
| """Instantiate DynamicsModel from a checkpoint's saved config and load weights.""" | |
| import torch | |
| from models.dynamics import DynamicsModel | |
| model_sd = torch.load(Path(checkpoint_path) / MODEL_CHECKPOINT, map_location='cpu', weights_only=True) | |
| state_cfg = torch.load(Path(checkpoint_path) / STATE, map_location='cpu', weights_only=False) | |
| cfg = state_cfg.get('config', {}) or {} | |
| frame_size = cfg.get('frame_size', 128) | |
| # Infer conditioning_dim from checkpoint if missing | |
| conditioning_dim = cfg.get('conditioning_dim', None) | |
| if conditioning_dim is None: | |
| cond_inferred = None | |
| for k, v in model_sd.items(): | |
| # Linear weight shape: [out_features, in_features]; in_features is conditioning dim | |
| if k.endswith('to_gamma_beta.1.weight'): | |
| cond_inferred = int(v.shape[1]) | |
| break | |
| conditioning_dim = cond_inferred if cond_inferred is not None else 3 | |
| kwargs = { | |
| 'frame_size': (frame_size, frame_size), | |
| 'patch_size': cfg.get('patch_size', 8), | |
| 'embed_dim': cfg.get('embed_dim', 128), | |
| 'num_heads': cfg.get('num_heads', 8), | |
| 'hidden_dim': cfg.get('hidden_dim', 256), | |
| 'num_blocks': cfg.get('num_blocks', 4), | |
| 'conditioning_dim': conditioning_dim, | |
| 'latent_dim': cfg.get('latent_dim', 6), | |
| 'num_bins': cfg.get('num_bins', 4), | |
| 'use_moe': cfg.get('use_moe', False), | |
| 'num_experts': cfg.get('num_experts', 4), | |
| 'top_k_experts': cfg.get('top_k_experts', 2), | |
| 'moe_aux_loss_coeff': cfg.get('moe_aux_loss_coeff', 0.01), | |
| } | |
| if model is None: | |
| model = DynamicsModel(**kwargs) | |
| set_model_state_dict( | |
| model=model, | |
| model_state_dict=model_sd, | |
| options=StateDictOptions( | |
| full_state_dict=True, | |
| broadcast_from_rank0=is_distributed, | |
| ) | |
| ) | |
| model = model.to(device) | |
| return model, state_cfg | |
| def prepare_pipeline_run_root(run_name: Optional[str] = None, base_cwd: Optional[str] = None): | |
| """Create a top-level run root directory results/<timestamp_or_name>""" | |
| cwd = base_cwd or os.getcwd() | |
| ts = readable_timestamp() | |
| name = run_name or ts | |
| run_root = os.path.join(cwd, 'results', name) | |
| os.makedirs(run_root, exist_ok=True) | |
| return run_root, name | |
| def prepare_stage_dirs(run_root_dir: str, stage_name: str): | |
| """Create stage subdirectories under the given run root. | |
| Structure: | |
| <run_root_dir>/<stage_name>/checkpoints | |
| <run_root_dir>/<stage_name>/visualizations | |
| Returns (stage_dir, checkpoints_dir, visualizations_dir). | |
| """ | |
| stage_dir = os.path.join(run_root_dir, stage_name) | |
| checkpoints_dir = os.path.join(stage_dir, 'checkpoints') | |
| visualizations_dir = os.path.join(stage_dir, 'visualizations') | |
| os.makedirs(checkpoints_dir, exist_ok=True) | |
| os.makedirs(visualizations_dir, exist_ok=True) | |
| return stage_dir, checkpoints_dir, visualizations_dir | |