import argparse import contextlib import logging import math import os import time import mup import numpy as np import torch import torchvision.transforms.functional as transforms_f from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import set_seed from einops import rearrange from lpips import lpips from torch.utils.data import DataLoader from tqdm.auto import tqdm import transformers import traceback from transformers import ( default_data_collator, get_scheduler, ) from collections import defaultdict from cont_data import RawFeatureDataset, get_maskgit_collator_feature from common.eval_utils import decode_tokens, compute_lpips, decode_features from genie.config import DiffusionGenieConfig from genie.st_mar import STMAR from visualize import decode_latents_wrapper from skimage import metrics as image_metrics from matplotlib import pyplot as plt from datetime import datetime from accelerate import DistributedDataParallelKwargs torch.autograd.set_detect_anomaly(True) # Get current date and time now = datetime.now() # Format the datetime object as a string formatted_date = now.strftime("%Y-%m-%d %H:%M:%S") torch.set_float32_matmul_precision("medium") logger = get_logger(__name__) SVD_SCALE = 0.18215 def parse_args(): # parser = argparse.ArgumentParser(description="Train a MaskGIT or Llama-style LLM on video generation.") parser = argparse.ArgumentParser(description="Train a spatial-temporal MaskGIT-style model on video generation.") # Data parser.add_argument( "--train_data_dir", type=str, default="data/1x_humanoid_magvit_traj1000_train", help="Directory containing tokenized data, should have a `video.bin`, `metadata.json` and `segment_ids.json`." ) parser.add_argument( "--val_data_dir", type=str, default="data/1x_humanoid_magvit_traj1000_val", help="Directory containing tokenized data, should have a `video.bin`, `metadata.json` and `segment_ids.json`." ) parser.add_argument( "--domain", type=str, default="1x_humanoid", help="The domain name for the dataset" ) parser.add_argument( "--window_size", type=int, default=12, help="Number of frames to in a sequence.", ) parser.add_argument( "--stride", type=int, default=None, help="Difference in frame count between consecutive frames in a sequence.", ) parser.add_argument( "--filter_overlaps", action="store_true", help=( "Whether to filter repeated frames in the train dataset (`filter_overlaps` always true for the val set). " "Filtering essentially makes the training dataset less correlated but ~16x smaller, " "see the `filter_overlaps` argument in `RawTokenDataset` for details."), default=True ) # Model parser.add_argument( "--llama_config", type=str, help="`transformers.LlamaConfig` json. " "E.g. https://huggingface.co/1x-technologies/Llama_1B_v0/blob/main/config.json", ) parser.add_argument( "--genie_config", type=str, help="GenieConfig json." ), parser.add_argument( "--warmstart_path", type=str, default=None, help="A path to a checkpoint to warmstart a model from, possibly not trained on the same dataset, " "will resize embeddings if needed.", ) parser.add_argument( "--resume_from_checkpoint", type=str, default=None, help="If the training should continue from a checkpoint folder.", ) # Training parser.add_argument( "--per_device_train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader.", ) parser.add_argument( "--per_device_eval_batch_size", type=int, default=2, help="Batch size (per device) for the evaluation dataloader.", ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.", ) parser.add_argument( "--gradient_checkpointing", default=False, action="store_true", ) parser.add_argument( "--learning_rate", type=float, default=2e-4, help="Initial learning rate (after the potential warmup period) to use.", ) parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay to use.") parser.add_argument("--num_train_epochs", type=int, default=2, help="Total number of training epochs to perform.") parser.add_argument( "--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.", ) parser.add_argument( "--max_eval_steps", type=int, default=int(1e10), help="Only evaluate on `max_eval_steps` batches of validation data per process, faster.", ) parser.add_argument( "--eval_every_n_steps", type=int, default=1000, help="Eval every N training steps.", ) parser.add_argument( "--vis_every_n_steps", type=int, default=1000, help="Visualize every N training steps.", ) parser.add_argument( "--lr_scheduler_type", type=str, default="constant_with_warmup", help="The scheduler type to use.", choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup", "custom_cosine"], ) parser.add_argument( "--num_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." ) parser.add_argument( "--max_grad_norm", type=float, default=10, help="Threshold to clip gradients.", ) parser.add_argument( "--attention_dropout", type=float, default=0.01, help="Attention dropout prob.", ) parser.add_argument( "--adam_beta_1", type=float, default=0.9, ) parser.add_argument( "--adam_beta_2", type=float, default=0.95, ) parser.add_argument( "--adam_eps", type=float, default=1e-8, ) # Misc parser.add_argument("--output_dir", type=str, required=True, help="Where to store the model checkpoints.") parser.add_argument( "--checkpointing_steps", type=str, default="10000", help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.", ) parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.") parser.add_argument( "--overfit_first_batch", action="store_true", help=( "Debug option that trains and validates on only the first batch of the training dataset." ), ) parser.add_argument( "--report_to", type=str, default="wandb", help="The integration to report the results and logs to.", ) parser.add_argument( "--mu_transfer", action="store_true", help="If specified, will train with mu transfer reparametrizations. Only supports Llama models.", default=False ) parser.add_argument( "--no_compile", action="store_true", help="If specified, will not compile the model.", default=True ) parser.add_argument( "--add_action_input", action="store_true", help=( "Whether to add action as input to the dynamics model. ") ) parser.add_argument( "--run_name", type=str, default="video_prediction", help="", ) parser.add_argument( "--cleanup_checkpoints", action="store_true", help=( "Whether to clean up checkpoints (to keep only the last 3) along the training. "), default=True ) parser.add_argument( "--use_raw_image_as_latent", action="store_true", help="If specified, will train with mu transfer reparametrizations. Only supports Llama models.", ) return parser def save_checkpoint(model, accelerator, args, filename): """ filename: `save_path = os.path.join(args.output_dir, filename)` """ unwrapped_model = accelerator.unwrap_model(model) save_path = os.path.join(args.output_dir, filename) if accelerator.is_main_process: unwrapped_model.save_pretrained( save_path, is_main_process=accelerator.is_main_process, save_function=accelerator.save ) accelerator.save_state(save_path) @torch.no_grad() def visualize(accelerator, model, dataloader, window_size, encoder_type, encoder_name_or_path, metrics_prefix="eval", max_steps=1, use_raw_image_as_latent=False): """ Visualizes model's autoregressive generation outputs, logged to wandb. It uses teacher-forcing (causal in time axis) """ # breakpoint() accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) if not unwrapped_model.config.jointly_predict_states: return metrics = defaultdict(list) if hasattr(dataloader.dataset, "metadata"): if accelerator.is_main_process: lpips_alex = lpips.LPIPS(net="alex") # Calculate LPIPS w/ AlexNet, the fastest option else: if accelerator.is_main_process: lpips_alex = lpips.LPIPS(net="alex") # Calculate LPIPS w/ AlexNet, the fastest option decode_latents = decode_latents_wrapper(encoder_type=encoder_type, encoder_name_or_path=encoder_name_or_path) # re-initializing every time to save memory unwrapped_model.eval() for step, batch in enumerate(dataloader): # Note: hardcoding 4 image cap for faster inference on small models TEST_NUM = 4 reshaped_labels = rearrange(batch["labels"][:TEST_NUM], "b (t s) c -> b t s c", t=window_size).to(accelerator.device) # `s` is really `(h, w)` domains = batch["domain"][:TEST_NUM] if 'action_ids' in batch: action_ids = batch["action_ids"][:TEST_NUM].to(accelerator.device) else: action_ids = None # hardcoding half of frames for context num_prompt_frames = unwrapped_model.config.num_prompt_frames num_new_tokens = batch["w"][0] * batch["h"][0] * (window_size - num_prompt_frames) prompt_input_ids = rearrange(reshaped_labels[:, :num_prompt_frames], "b t s c -> b (t s) c") outputs = unwrapped_model.generate(input_ids=prompt_input_ids, attention_mask=torch.ones_like(prompt_input_ids), max_new_tokens=num_new_tokens, min_new_tokens=num_new_tokens, action_ids=action_ids, domain=batch["domain"][:TEST_NUM], w=batch["w"][:TEST_NUM], h=batch["h"][:TEST_NUM]) output_tokens = rearrange(outputs, "b (t h w) c -> b t h w c", t=window_size, h=batch["h"][0], w=batch["w"][0]) gtruth_tokens = rearrange(reshaped_labels[:, num_prompt_frames:], "b t (h w) c -> b t h w c", h=batch["h"][0], w=batch["w"][0]) if not use_raw_image_as_latent: output_tokens = output_tokens / SVD_SCALE gtruth_tokens = gtruth_tokens / SVD_SCALE decoded_output = decode_features(output_tokens.cpu(), decode_latents) decoded_gtruth = decode_features(gtruth_tokens.cpu(), decode_latents) else: decoded_output = ((output_tokens + 0.5) * 255).long() decoded_gtruth = ((gtruth_tokens + 0.5) * 255).long() decoded_output = accelerator.gather(decoded_output.to(accelerator.device)).cpu() decoded_gtruth = accelerator.gather(decoded_gtruth.to(accelerator.device)).cpu() if accelerator.is_main_process and step < 2: exs_per_fig = 4 for j in range(0, len(decoded_output), exs_per_fig): # with 10 percent change we log some output. to save spaces fig, axs = plt.subplots(nrows=2 * exs_per_fig, ncols=window_size, figsize=(3 * window_size, 3 * 2 * exs_per_fig)) # If len(decoded_output) is not a multiple of 4, make sure to truncate properly for k in range(min(exs_per_fig, len(decoded_output) - j)): for i in range(num_prompt_frames): for ax in (axs[k * 2, i], axs[k * 2 + 1, i]): ax.imshow(transforms_f.to_pil_image(decoded_output[j + k, i])) ax.set_title("Context") ax.axis("off") for i in range(num_prompt_frames, window_size): axs[k * 2, i].imshow(transforms_f.to_pil_image(decoded_gtruth[j + k, i - num_prompt_frames])) axs[k * 2, i].set_title("Ground truth") axs[k * 2 + 1, i].imshow(transforms_f.to_pil_image(decoded_output[j + k, i])) axs[k * 2 + 1, i].set_title("Prediction") for ax in axs[:, i]: ax.axis("off") wandb_tracker = accelerator.get_tracker("wandb") wandb_tracker.log({f"vis_{metrics_prefix}_{j}": fig}, commit=False) wandb_tracker.log({f"{domains[0]}/vis_{metrics_prefix}_{j}": fig}, commit=False) plt.close(fig) metrics["ar_lpips"].extend(compute_lpips(decoded_gtruth, # Note: not parallelizing right now decoded_output[:, num_prompt_frames:], lpips_alex)) gt_frames_numpy = decoded_gtruth.detach().cpu().numpy() pred_frames_numpy = decoded_output[:, num_prompt_frames:].detach().cpu().numpy() psnr = [image_metrics.peak_signal_noise_ratio( gt_frames_numpy[i] / 255., pred_frames_numpy[i] / 255., data_range=1.0) for i in range(gt_frames_numpy.shape[0])] ssim = [np.mean([image_metrics.structural_similarity( gt_frames_numpy[i][j] / 255., pred_frames_numpy[i][j] / 255., data_range=1.0, channel_axis=0) \ for i in range(gt_frames_numpy.shape[0])]) for j in range(gt_frames_numpy.shape[1])] # compute some other metrics metrics[f"{metrics_prefix}/ar_psnr"].extend(psnr) metrics[f"{metrics_prefix}/ar_ssim"].extend(ssim) metrics[f"{batch['domain'][0]}/ar_lpips"].extend(compute_lpips(decoded_gtruth, # Note: not parallelizing right now decoded_output[:, num_prompt_frames:], lpips_alex)) if step + 1 >= max_steps: break unwrapped_model.train() if accelerator.is_main_process: metrics = {f"{metrics_prefix}_{key}": np.mean(val) for key, val in metrics.items() if len(val) > 0} print(f"{metrics=}") wandb_tracker = accelerator.get_tracker("wandb") wandb_tracker.log(metrics, commit=False) def train(accelerator, model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, experiment_config, config, args): total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps logger.info("***** Running training *****") logger.info(f" Num examples = {len(train_dataloader)}") logger.info(f" Num Epochs = {args.num_train_epochs}") logger.info(f" Instantaneous batch size per device = {args.per_device_train_batch_size}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") logger.info(f" Total optimization steps = {args.max_train_steps}") # Only show the progress bar once on each machine. progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process) completed_steps = 0 starting_epoch = 0 resume_step = None checkpoint_path = None # Potentially load in the weights and states from a previous save if args.resume_from_checkpoint: if os.path.exists(args.resume_from_checkpoint + "/pytorch_model.bin") or os.path.exists(args.resume_from_checkpoint + "/model.safetensors"): checkpoint_path = args.resume_from_checkpoint path = os.path.basename(args.resume_from_checkpoint.rstrip("/")) else: # Get the most recent checkpoint base_path = os.path.dirname(args.resume_from_checkpoint) dirs = [os.path.join(base_path, f.name) for f in os.scandir(base_path) if f.is_dir()] dirs.sort(key=os.path.getctime) # Sorts folders by date modified, most recent checkpoint is the last if len(dirs) > 0: path = dirs[-1] checkpoint_path = path path = os.path.basename(checkpoint_path) try: accelerator.print(f"Resumed from checkpoint: {checkpoint_path}") if os.path.exists(checkpoint_path): # for finetuning with a different structures print(f"loading checkpoint from {checkpoint_path}") accelerator.load_state(checkpoint_path, strict=False) # tied weights not saved so can't load strict, but also no need to tie again # Extract `epoch_{i}` or `step_{i}` training_difference = os.path.splitext(path)[0] else: print("No checkpoint found, training from scratch.") training_difference = "step_0" if "epoch" in training_difference: starting_epoch = int(training_difference.replace("epoch_", "")) + 1 resume_step = None num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) completed_steps = starting_epoch * num_update_steps_per_epoch else: # need to multiply `gradient_accumulation_steps` to reflect real steps resume_step = int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps starting_epoch = resume_step // len(train_dataloader) completed_steps = resume_step // args.gradient_accumulation_steps resume_step -= starting_epoch * len(train_dataloader) except Exception as e: print("load checkpoint incomplete", traceback.format_exc()) # update the progress_bar if load from checkpoint progress_bar.update(completed_steps) loss_info = torch.zeros(2, device=accelerator.device) # sum, count for epoch in range(starting_epoch, args.num_train_epochs): model.train() train_dataloader.set_epoch(epoch) # potenally cleanup the previous checkpoints if args.cleanup_checkpoints: if os.path.exists(args.output_dir): dirs = [os.path.join(args.output_dir, f.name) for f in os.scandir(args.output_dir) if f.is_dir()] if len(dirs) > 3: dirs.sort(key=os.path.getctime) paths = dirs[:-3] # only keep the last 3 for path in paths: print(f"remove rm -rf {path}") # os.system(f"rm -rf {path}") if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None: # We skip the first `n` batches in the dataloader when resuming from a checkpoint active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step) else: active_dataloader = train_dataloader _time = time.time() # accelerator.wait_for_everyone() dataloader_iter = iter(active_dataloader) for step in range(len(active_dataloader)): batch = next(dataloader_iter) batch_size = batch["input_ids"].size(0) # Manual gradient accumulation because accelerator somehow taking a lot of memory is_update_step = (step + 1) % args.gradient_accumulation_steps == 0 ctx_manager = contextlib.nullcontext() if is_update_step else accelerator.no_sync(model) train_action_loss = 0 with ctx_manager: outputs = model(**batch) loss = outputs.loss if not torch.isnan(loss).any(): loss_info[0] += loss.detach().mean() * batch_size if "action_loss" in outputs: train_action_loss = outputs.action_loss.item() loss += config.action_loss_weight * outputs.action_loss accelerator.backward(loss / args.gradient_accumulation_steps) else: print(f"Warning: NaN or Inf detected in loss for domain: {batch['domain']}. Skipping backward pass.") dummy_loss = torch.zeros_like(loss, requires_grad=True) accelerator.backward(dummy_loss) loss_info[1] += batch_size if not is_update_step: continue # Everything below only happens on update step if args.max_grad_norm is not None: accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) optimizer.step() lr_scheduler.step() optimizer.zero_grad() loss_info = accelerator.reduce(loss_info) train_action_loss = train_action_loss / loss_info[1] avg_train_loss = (loss_info[0] / loss_info[1]).item() # sum / count loss_info *= 0 # reset sum and count batch_time = time.time() - _time # accumulated batch _time = time.time() accelerator.log( { "train_loss": avg_train_loss, "train_action_loss": train_action_loss, f"stat/{batch['domain'][0]}_action_loss": train_action_loss, f"stat/{batch['domain'][0]}_train_loss": avg_train_loss, "epoch": epoch, "update_step": completed_steps, "examples_processed": completed_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * accelerator.num_processes, "learning_rate": lr_scheduler.get_last_lr()[0], "flops": (completed_steps + 1) * experiment_config["FLOPs_per_update_step"], "throughput_examples": experiment_config["effective_batch_size"] / batch_time, }, step=completed_steps) progress_bar.update(1) completed_steps += 1 # print(f"{completed_steps % args.checkpointing_steps=} {completed_steps=} {args.checkpointing_steps=}") if completed_steps % int(args.checkpointing_steps) == 0: print(f"Saving checkpoint at step {completed_steps}!") save_checkpoint(model, accelerator, args, f"step_{completed_steps}") if completed_steps % args.eval_every_n_steps == 0: time.sleep(1) # manual adding time sleep model.eval() eval_losses = [] # Compute token-level accuracy (w/ teacher forcing) num_correct = 0 num_total = 0 # to resolve the data collating issues eval_dataloader_iter = iter(eval_dataloader) for step in range(min(args.max_eval_steps, len(eval_dataloader))): batch = next(eval_dataloader_iter) batch_size = len(batch["input_ids"]) # Last batch might not be full with torch.no_grad(): outputs = model(**batch) loss = outputs.loss eval_losses.append(accelerator.gather_for_metrics(loss.repeat(batch_size))) if "acc" in outputs: # `num_correct` and `num_total` actually track mean accuracy in this case. num_correct_batch = accelerator.reduce(outputs.acc, reduction="mean").item() * batch_size num_total_batch = batch_size num_correct += num_correct_batch num_total += num_total_batch else: shifted_preds = torch.argmax(outputs.logits[:, :-1, :], dim=-1) shifted_labels = batch["labels"][:, 1:] num_correct_batch = accelerator.gather_for_metrics((shifted_preds == shifted_labels).sum()).sum().item() num_total_batch = accelerator.gather_for_metrics(torch.tensor(torch.numel(shifted_labels), device=accelerator.device)).sum().item() num_correct += num_correct_batch num_total += num_total_batch if step >= args.max_eval_steps * args.num_datasets: break try: accelerator.log( { f"stat/{batch['domain'][0]}_eval_loss": eval_losses[-1], # f"{batch['domain'][0]}_stat/eval_teacher_acc": num_correct_batch / num_total_batch }, step=completed_steps, ) except Exception as e: print("log failed", e) eval_losses = torch.cat(eval_losses) eval_loss = torch.mean(eval_losses).item() eval_teacher_acc = num_correct / num_total logger.info(f"{completed_steps=} {eval_loss=} {eval_teacher_acc=}") accelerator.log( { "eval_loss": eval_loss, "epoch": epoch, "update_step": completed_steps, "examples_processed": completed_steps * args.per_device_train_batch_size * args.gradient_accumulation_steps * accelerator.num_processes, "flops": completed_steps * experiment_config["FLOPs_per_update_step"], }, step=completed_steps, ) # Switch back to train mode model.train() if completed_steps % args.vis_every_n_steps == 0: if "encoder_type" not in experiment_config: experiment_config["encoder_name_or_path"] = "data/magvit2.ckpt" experiment_config["encoder_type"] = "magvit" if not args.overfit_first_batch: # val is same as train otherwise visualize(accelerator, model, eval_dataloader, args.window_size, experiment_config["encoder_type"], experiment_config["encoder_name_or_path"], "val", use_raw_image_as_latent=args.use_raw_image_as_latent) visualize(accelerator, model, train_dataloader, args.window_size, experiment_config["encoder_type"], experiment_config["encoder_name_or_path"], "train", use_raw_image_as_latent=args.use_raw_image_as_latent) if completed_steps >= args.max_train_steps: break if args.checkpointing_steps == "epoch": save_checkpoint(model, accelerator, args, f"epoch_{epoch}") accelerator.end_training() save_checkpoint(model, accelerator, args, f"final_checkpt") def main(): parser = parse_args() args = parser.parse_args() # Manual gradient accumulation ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator(gradient_accumulation_steps=1, log_with=args.report_to, project_dir=args.output_dir, kwargs_handlers=[ddp_kwargs]) accelerator.init_trackers("video") if accelerator.is_main_process: accelerator.trackers[0].run.name = formatted_date + "_" + args.run_name # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_info() print(f"Rank {accelerator.process_index} assigned to device {torch.cuda.current_device()}") else: transformers.utils.logging.set_verbosity_error() if args.seed is not None: set_seed(args.seed) if accelerator.is_main_process: os.makedirs(args.output_dir, exist_ok=True) accelerator.wait_for_everyone() config = DiffusionGenieConfig.from_pretrained(args.genie_config) train_dataset = RawFeatureDataset(args.train_data_dir, window_size=args.window_size, stride=args.stride, filter_overlaps=args.filter_overlaps, compute_stride_from_freq_table=(args.stride is None), use_actions=config.use_actions, use_raw_image_as_latent=args.use_raw_image_as_latent) if not args.overfit_first_batch: eval_dataset = RawFeatureDataset(args.val_data_dir, window_size=args.window_size, stride=args.stride, filter_overlaps=True, compute_stride_from_freq_table=(args.stride is None), use_actions=config.use_actions, use_raw_image_as_latent=args.use_raw_image_as_latent) else: train_dataset.valid_start_inds = train_dataset.valid_start_inds[:args.per_device_train_batch_size * args.gradient_accumulation_steps * accelerator.num_processes] eval_dataset = train_dataset shared_keys = ("s", "h", "w", "vocab_size", "latent_channels", "encoder_type", "encoder_name_or_path", "quantized") # TODO: check train/val hz per dataset? assert all(train_dataset.metadata.get(shared_key) == eval_dataset.metadata.get(shared_key) for shared_key in shared_keys) # if "encoder_type" not in train_dataset.metadata or "encoder_name_or_path" not in train_dataset.metadata: # accelerator.print("Assuming MAGVIT image encoder.") # train_dataset.metadata["encoder_type"] = "magvit" # train_dataset.metadata["encoder_name_or_path"] = "data/magvit2.ckpt" # Will not store key in metadata if it's missing, so that defaults can be filled by functions later? # TODO: don't think we are handling missing keys in function calls shared_metadata = {shared_key: train_dataset.metadata[shared_key] for shared_key in shared_keys if shared_key in train_dataset.metadata} if args.llama_config is not None: raise NotImplementedError("Have not factorized Llama vocabulary.") else: config.use_mup = args.mu_transfer # Note: changing this may affect pre-trained model due to attn scaling config.image_vocab_size = shared_metadata["vocab_size"] config.T = args.window_size config.S = shared_metadata["h"] * shared_metadata["w"] # TODO: make STMaskGIT use h and w instead of S config.vae_embed_dim = shared_metadata["latent_channels"] model = STMAR(config) if config.use_actions: print(f"Initializing action projectors with {train_dataset.n_action}d action") model.init_action_projectors([args.domain], [train_dataset.n_action], [train_dataset.action_stat], config.action_network) if args.mu_transfer: model.set_mup_shapes(rescale_params=True) model.init_weights() # might be unnecessary if `rescale_params` is True # Optimizer. Split weights in two groups, one with weight decay and the other not. no_decay = ["bias", "layer_norm.weight"] optimizer_grouped_parameters = [ { "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": args.weight_decay, }, { "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] opt_class = mup.MuAdamW if args.mu_transfer else torch.optim.AdamW # scale base learning rate effective_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps \ * accelerator.num_processes args.learning_rate = args.learning_rate * min(max(1, effective_batch_size / 64), 8) optimizer = opt_class(optimizer_grouped_parameters, lr=args.learning_rate, betas=(args.adam_beta_1, args.adam_beta_2), eps=args.adam_eps) # DataLoaders creation: collate_fn = default_data_collator if args.llama_config is not None else get_maskgit_collator_feature(config) train_dataloader = DataLoader( train_dataset, shuffle=True, collate_fn=collate_fn, batch_size=args.per_device_train_batch_size, num_workers=4, pin_memory=True, ) # Shuffle eval dataset and then set shuffle=False on the dataloader. # Shuffling in the dataloader results in reshuffling with each iteration. eval_dataset.valid_start_inds = torch.tensor(eval_dataset.valid_start_inds)[ torch.randperm(len(eval_dataset), generator=torch.Generator().manual_seed(0)) ].tolist() eval_dataloader = DataLoader( eval_dataset, shuffle=False, collate_fn=collate_fn, batch_size=args.per_device_eval_batch_size, pin_memory=True, num_workers=4, ) # Scheduler and math around the number of training steps. overrode_max_train_steps = False num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if args.max_train_steps is None: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch overrode_max_train_steps = True if args.lr_scheduler_type == "custom_cosine": # decay to `end_ratio` of the peak learning rate def get_lr_wrapper(warmup_steps, max_steps, end_ratio=0.1): def get_lr(step): if step < warmup_steps: return (step + 1) / warmup_steps remaining_steps = max_steps - warmup_steps return ((1 + math.cos(math.pi * (step - warmup_steps) / remaining_steps)) / 2) \ * (1 - end_ratio) + end_ratio return get_lr lr_scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, get_lr_wrapper(args.num_warmup_steps * accelerator.num_processes, args.max_train_steps if overrode_max_train_steps else args.max_train_steps * accelerator.num_processes) ) else: lr_scheduler = get_scheduler( name=args.lr_scheduler_type, optimizer=optimizer, num_warmup_steps=args.num_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps if overrode_max_train_steps else args.max_train_steps * accelerator.num_processes, ) # Enable gradient checkpointing to save memory if args.gradient_checkpointing: logger.info("Enabling gradient checkpointing") model.gradient_checkpointing_enable() model.config.use_cache = False # incompatible with grad checkpointing # Prepare everything with our `accelerator`. model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare( model, optimizer, train_dataloader, eval_dataloader, lr_scheduler ) if not args.no_compile: torch._dynamo.config.cache_size_limit = 256 torch._dynamo.config.optimize_ddp = False # https://github.com/pytorch/pytorch/issues/104674 # TODO: https://github.com/pytorch/pytorch/issues/109774#issuecomment-2046633776 model = torch.compile(model) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if overrode_max_train_steps: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch # Afterwards we recalculate our number of training epochs args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) # Figure out how many steps we should save the Accelerator states checkpointing_steps = args.checkpointing_steps if checkpointing_steps is not None and checkpointing_steps.isdigit(): checkpointing_steps = int(checkpointing_steps) # We need to initialize the trackers we use, and also store our configuration. # The trackers initialize automatically on the main process. experiment_config = vars(args) | vars(config) seq_len = shared_metadata["h"] * shared_metadata["w"] * args.window_size effective_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps \ * accelerator.num_processes args.num_datasets = 1 model_module = model.module if hasattr(model, "module") else model experiment_config.update(shared_metadata | { "model_parameters": sum(p.numel() for p in model.parameters()), "model_parameters_M": round(sum(p.numel() for p in model.parameters()) / 1e6), "trunk_parameters": sum(p.numel() for p in model_module.decoder.parameters()), "trunk_parameters_M": round(sum(p.numel() for p in model_module.decoder.parameters()) / 1e6), "seq_len": seq_len, # "hz": "HARDCODED_2", "train_data_tokens": len(train_dataset) * seq_len, # only one epoch "effective_batch_size": effective_batch_size, "effective_batch_size_tokens": effective_batch_size * seq_len, "mixed_precision": accelerator.mixed_precision, "num_datasets": 1 }) print("============================") print(f"model parameters: {experiment_config['model_parameters_M']}M") print("============================") experiment_config["FLOPs_per_update_step"] = 6 * experiment_config["model_parameters"] \ * experiment_config["effective_batch_size_tokens"] accelerator.init_trackers(project_name="video", config=experiment_config) train(accelerator, model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, experiment_config, config, args) if __name__ == "__main__": main()