train.py

import argparse
import contextlib
import logging
import math
import os
import time

import matplotlib
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 data import RawTokenDataset, get_maskgit_collator

from common.eval_utils import decode_tokens, compute_lpips
from genie.st_mask_git import STMaskGIT
from genie.config import GenieConfig

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__)

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(
        "--diffusion",
        action="store_true",
        help="use diffusion model."
    ),
    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=1,
        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=1e-4,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument("--weight_decay", type=float, default=0.05, 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=20000,
        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=1.0,
        help="Threshold to clip gradients.",
    )
    parser.add_argument(
        "--attention_dropout",
        type=float,
        default=0.05,
        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=True
    )
    parser.add_argument(
        "--no_compile",
        action="store_true",
        help="If specified, will not compile the model.",
        default=True
    )

    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. "),
    )
    parser.add_argument(
        "--save_second_epoch",
        action="store_true",
        help="Whether to checkpoint at the end of the second epoch (1-indexing). This one will not be auto-deleted by cleanup.",
        default=True
    )
    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, metrics_prefix="train", max_steps=1):
    """
    Visualizes model's autoregressive generation outputs, logged to wandb.
    It uses teacher-forcing (causal in time axis)
    """

    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    if not unwrapped_model.config.jointly_predict_states:
        return
    metrics = defaultdict(list)
    if accelerator.is_main_process:
        lpips_alex = lpips.LPIPS(net="alex")  # Calculate LPIPS w/ AlexNet, the fastest option

    decode_latents = decode_latents_wrapper()  # re-initializing every time to save memory
    unwrapped_model.eval()
    rank = 0
    dataloader_iter = iter(dataloader)
    for step in range(len(dataloader)):
        try:
            batch = next(dataloader_iter)

            # Note: hardcoding 4 image cap for faster inference on small models
            TEST_NUM = 4
            reshaped_labels = rearrange(batch["labels"][:TEST_NUM], "b (t s) -> b t s", 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 -> b (t s)")
            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) -> b t h w", t=window_size,
                                    h=batch["h"][0], w=batch["w"][0])
            gtruth_tokens = rearrange(reshaped_labels[:, num_prompt_frames:], "b t (h w) -> b t h w",
                                    h=batch["h"][0], w=batch["w"][0])

            decoded_output = decode_tokens(output_tokens.cpu(), decode_latents)
            decoded_gtruth = decode_tokens(gtruth_tokens.cpu(), decode_latents)

            decoded_output = accelerator.gather(decoded_output.to(accelerator.device)).cpu()
            decoded_gtruth = accelerator.gather(decoded_gtruth.to(accelerator.device)).cpu()

            # As in Genie. we also compute psnr_delta = PSNR(x_t, x_t_hat) - PSNR(x_t, x_t_hatprime) where x_t_hatprime samples random actions
            # this difference in PSNR measures the controllability
            # actions need to be just uniform random actions
            if action_ids is not None:
                random_action_ids = torch.randn_like(action_ids)
                random_action_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=random_action_ids,
                                                    domain=batch["domain"][:TEST_NUM],
                                                    w=batch["w"][:TEST_NUM],
                                                    h=batch["h"][:TEST_NUM],
                                                    skip_normalization=True)

                random_output_tokens = rearrange(random_action_outputs, "b (t h w) -> b t h w", t=window_size,
                                        h=batch["h"][0], w=batch["w"][0])
                random_output_tokens = decode_tokens(random_output_tokens.cpu(), decode_latents)

                random_output_tokens = accelerator.gather(random_output_tokens.to(accelerator.device)).cpu()
                random_pred_frames_numpy = random_output_tokens[:, num_prompt_frames:].detach().cpu().numpy()


            if accelerator.is_main_process:
                exs_per_fig = 4

                for j in range(0, len(decoded_output), exs_per_fig):
                    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")

                    rank = accelerator.process_index
                    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 action_ids is not None:
                    # log controllability as random subtracts groundtruth
                    psnr_delta = [psnr[i] - image_metrics.peak_signal_noise_ratio(
                        gt_frames_numpy[i] / 255., random_pred_frames_numpy[i] / 255., data_range=1.0) for i in range(gt_frames_numpy.shape[0])]

                    metrics[f"{metrics_prefix}/ar_psnr_delta"].extend(psnr_delta)

        except Exception as e:
            print("batch failed", traceback.format_exc())

        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 = ""

    # Potentially load in the weights and states from a previous save
    if args.resume_from_checkpoint:
        try:
            if  os.path.exists(args.resume_from_checkpoint + "/pytorch_model.bin"):
                checkpoint_path = args.resume_from_checkpoint
                path = os.path.basename(args.resume_from_checkpoint.rstrip("/"))
            # else:
            #     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)

            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:
            training_difference = "step_0"
            starting_epoch = 0
            completed_steps = 0
            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)

        # potentially 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:  # must keep at least 2 checkpoints for second epoch and most recent one
                    if args.save_second_epoch and os.path.join(args.output_dir, "epoch_1") in dirs:  # never prune second epoch
                        dirs.remove(os.path.join(args.output_dir, "epoch_1"))

                    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()
        dataloader_iter = iter(active_dataloader)

        # Switch back to train mode
        model.train()
        num_iters_per_epoch = max(len(active_dataloader) - 8, 1) # avoid the last few iters

        for step in range(num_iters_per_epoch):
            try:
                train_action_loss = 0
                batch = next(dataloader_iter)
                # to reduce the numerical instability in the very beginning of training
                gradient_accumulation_steps = args.gradient_accumulation_steps
                batch_size = batch["input_ids"].size(0)
                # Manual gradient accumulation because accelerator somehow taking a lot of memory
                is_update_step = (step + 1) % gradient_accumulation_steps == 0
                ctx_manager = contextlib.nullcontext() if is_update_step else accelerator.no_sync(model)

                with ctx_manager:
                    accelerator.wait_for_everyone()
                    outputs = model(**batch)
                    loss = outputs.loss

                    if not torch.isnan(loss).any():
                        loss_info[0] += loss.detach().mean() * batch_size # only video loss
                        if "action_loss" in outputs:
                            train_action_loss = outputs.action_loss.item()
                            loss += config.action_loss_weight * outputs.action_loss

                        loss_info[1] += batch_size
                        accelerator.backward(loss / gradient_accumulation_steps)
                    else:
                        print("Warning: NaN or Inf detected in loss. Skipping backward pass.")
                        dummy_loss = torch.zeros_like(loss, requires_grad=True)
                        accelerator.backward(dummy_loss)

                if not is_update_step:
                    continue

            except Exception as e:
                # avoid final iteration batch concatenation problems
                print("batch failed",  traceback.format_exc())
                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)

            avg_train_loss = (loss_info[0] / loss_info[1]).item()  # sum / count
            loss_info *= 0  # reset sum and count
            try:
                perplexity = math.exp(avg_train_loss)
            except OverflowError:
                print("overflow error for perplexity")
                perplexity = float("inf")

            # print(f"{perplexity=} {avg_train_loss=}")
            batch_time = time.time() - _time  # accumulated batch
            rank = accelerator.process_index

            domain_iter = str(batch['domain'][0])
            _time = time.time()
            accelerator.log(
                {
                    "train_perplexity": perplexity,
                    "train_loss": avg_train_loss,
                    "train_action_loss": train_action_loss,
                    f"stat/{domain_iter}_action_loss": train_action_loss / loss_info[1],
                    f"stat/{domain_iter}_train_perplexity": perplexity,
                    f"stat/{domain_iter}_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

                # barrier

                # to resolve the data collating issues
                eval_dataloader_iter = iter(eval_dataloader)
                for step in range(args.max_eval_steps):
                    eval_action_loss = 0
                    try:
                        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
                        if "action_loss" in outputs:
                            eval_action_loss = outputs.action_loss.item()
                        eval_losses.append(accelerator.gather_for_metrics(loss.repeat(batch_size)))
                    except Exception as e:
                        print("error:", e)
                        continue

                    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/{domain_iter}_eval_teacher_acc': num_correct_batch / num_total_batch,
                            f'stat/{domain_iter}_eval_loss': (torch.mean(eval_losses[-1])).item(),
                            f'stat/{domain_iter}_eval_action_loss': eval_action_loss,

                        },
                        step=completed_steps,
                        )
                    except Exception as e:
                        print("log failed", e)
                        continue

                if len(eval_losses) > 0:
                    eval_losses = torch.cat(eval_losses)
                    eval_loss = torch.mean(eval_losses).item()
                    eval_teacher_acc = num_correct / num_total
                    try:
                        perplexity = math.exp(eval_loss)
                    except OverflowError:
                        print("overflow error for perplexity")
                        perplexity = float("inf")
                else:
                    continue

                logger.info(f"{completed_steps=} {perplexity=} {eval_loss=} {eval_teacher_acc=}")
                accelerator.log(
                    {
                        "eval_perplexity": perplexity,
                        "eval_loss": eval_loss,
                        "eval_action_loss": eval_action_loss,
                        "eval_teacher_acc": eval_teacher_acc,
                        "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,
                )

            if completed_steps % args.vis_every_n_steps == 0 or completed_steps >= args.max_train_steps:
                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, "val")

                visualize(accelerator, model, train_dataloader, args.window_size, "train")

            if completed_steps >= args.max_train_steps:
                break

        if args.checkpointing_steps == "epoch" or (args.save_second_epoch and epoch == 1):
            save_checkpoint(model, accelerator, args, f"epoch_{epoch}")

    save_checkpoint(model, accelerator, args, f"final_checkpt")
    accelerator.end_training()

def main():
    parser = parse_args()
    args = parser.parse_args()
    assert (args.llama_config is not None) ^ (args.genie_config is not None), \
        "Exactly one of `llama_config` and `genie_config` should be set."

    # 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 = GenieConfig.from_pretrained(args.genie_config)
    train_dataset = RawTokenDataset(args.train_data_dir, window_size=args.window_size, name=args.domain,
                                    stride=args.stride, filter_overlaps=args.filter_overlaps,
                                    compute_stride_from_freq_table=(args.stride is None),
                                    use_actions=config.use_actions)

    if not args.overfit_first_batch:
        eval_dataset = RawTokenDataset(args.val_data_dir, window_size=args.window_size, name=args.domain,
                                        stride=args.stride, filter_overlaps=True,
                                        compute_stride_from_freq_table=(args.stride is None),
                                        use_actions=config.use_actions)
    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

    assert all(train_dataset.metadata[shared_key] == eval_dataset.metadata[shared_key]
               for shared_key in ("s", "vocab_size", "hz"))

    latent_side_len, vocab_size, hz = [train_dataset.metadata[key] for key in ("s", "vocab_size", "hz")]

    # Note: changing this may affect pre-trained model due to attn scaling
    config.use_mup = args.mu_transfer
    config.image_vocab_size = vocab_size
    config.T = args.window_size
    model = STMaskGIT(config)

    if config.use_actions:
        print(f"Initializing action projectors with {train_dataset.n_action}d action")
        model.init_action_projectors([train_dataset.name], [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,
        },
    ]

    # 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)

    opt_class = mup.MuAdamW if args.mu_transfer else torch.optim.AdamW
    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(config)
    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=collate_fn,
        batch_size=args.per_device_train_batch_size, num_workers=8, 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=8,
    )

    # 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
        if args.max_train_steps < 2000 and args.resume_from_checkpoint is None: # minimal number of trainng steps
            args.max_train_steps = 2000
            args.num_train_epochs = math.ceil(args.max_train_steps / 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 = latent_side_len**2 * args.window_size
    args.num_datasets = 1

    model_module = model.module if hasattr(model, "module") else model
    experiment_config.update({
        "model_parameters": sum(p.numel() for p in model_module.parameters()),
        "model_parameters_M": round(sum(p.numel() for p in model_module.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": hz / train_dataset.stride,
        "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()