train_multi_diffusion.py

import logging
import math
import os

import mup
import numpy as np
import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed

from torch.utils.data import DataLoader
from tqdm.auto import tqdm
import transformers
from transformers import (
    default_data_collator,
    get_scheduler,
)
import wandb

from cont_data import RawFeatureDataset, get_maskgit_collator_feature
from genie.config import DiffusionGenieConfig
from genie.st_mar import STMAR

from datetime import datetime
from accelerate import DistributedDataParallelKwargs
from common import data_sampler
import yaml
from train_diffusion import parse_args, train

# 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__)
torch.autograd.set_detect_anomaly(True)

def parse_args_multi():
    # parser = argparse.ArgumentParser(description="Train a MaskGIT or Llama-style LLM on video generation.")
    parser = parse_args()

    # Data
    parser.add_argument(
        "--train_split", type=str, default="experiments/datasplit/dataset2.yaml",
        help="Config files for using multiple datasets."
    )

    parser.add_argument(
        "--num_episodes_per_dataset",
        type=int,
        default=1000000,
        help="Maximum number of trajectories per dataset",
    )
    parser.add_argument(
        "--image_maskgit_path",
        type=str,
        default=None,
        help="Optional path to the official MaskGIT checkpoint. "
             "If specified, will copy relevant weights from the checkpoint. "
             "These weights will have a different (hard-coded) warmup schedule.",
    )
    parser.add_argument(
        "--action_network",
        type=str,
        default=None,
        choices=["concat", "cross_attention"],  # TODO: add other methods (resampler_concat, modulate, etc)
        help="If specified, will override the action in the config. Helps reduce the number of config jsons."
    )
    args = parser.parse_args()
    return args


def main():
    args = parse_args_multi()
    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,
                                even_batches=False, 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()
    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()

    # create multiple datasets
    with open(args.train_split, 'r') as file:
        datasplit = yaml.safe_load(file)

    config = DiffusionGenieConfig.from_pretrained(args.genie_config)

    # Extract the 'domains' value and split it into a list
    domains_list = [domain.strip() for domain in datasplit['domains'].split(',')]
    train_datasets = []
    val_datasets = []
    dataset_num_samples = []
    val_dataset_num_samples = []

    action_dimensions = []
    action_stats = []

    shared_keys = ("s", "h", "w", "vocab_size", "latent_channels",
                   "encoder_type", "encoder_name_or_path", "quantized")  # TODO: check train/val hz per dataset?
    for domain in domains_list:
        try:
 
            # train_data_dir = f"data/{domain}_vae_traj500_train" # {args.num_episodes_per_dataset}
            # val_data_dir = f"data/{domain}_vae_traj500_val"
            train_data_dir = f"data/{domain}_noquant_temporalvae_shard0_of_1_train"  # {args.num_episodes_per_dataset}
            val_data_dir = f"data/{domain}_noquant_temporalvae_shard0_of_1_val"

            # train_data_dir = f"data/{domain}_vae_traj{args.num_episodes_per_dataset}_train" # {args.num_episodes_per_dataset}
            # val_data_dir = f"data/{domain}_vae_traj{args.num_episodes_per_dataset}_val"

            if config.drop_action_ratio > 0:
                raise NotImplementedError

            train_dataset = RawFeatureDataset(train_data_dir, window_size=args.window_size,
                                            stride=args.stride, filter_overlaps=args.filter_overlaps,
                                            max_traj_num=args.num_episodes_per_dataset,
                                            use_actions=config.use_actions, domain=domain)
            dataset_num_samples.append(len(train_dataset))
            action_dimensions.append(train_dataset.n_action)
            if config.use_actions:
                action_stats.append(train_dataset.action_stat)

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

            # 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()
            val_dataset_num_samples.append(len(eval_dataset))
        except Exception as e:
            import traceback
            print(traceback.format_exc())

        train_datasets.append(train_dataset)
        val_datasets.append(eval_dataset)
        assert all(train_dataset.metadata.get(shared_key) == eval_dataset.metadata.get(shared_key)
                   for shared_key in shared_keys)  # TODO: check this across all datasets

    print("dataset_num_samples:", dataset_num_samples)

    # Will not store key in metadata if it's missing, so that defaults can be filled by functions later?  # TODO: handle missing keys
    shared_metadata = {shared_key: train_dataset.metadata[shared_key]
                       for shared_key in shared_keys if shared_key in train_dataset.metadata}

    config.use_mup = args.mu_transfer  # Note: changing this may affect pre-trained model due to attn scaling
    config.image_vocab_size = None
    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"]
    if args.action_network is not None:
        print("Using action network", args.action_network)
        config.action_network = args.action_network

    model = STMAR(config)

    if config.use_actions:
        # TODO: use new list instead of domains_list, in case domain fails
        model.init_action_projectors(domains_list, action_dimensions, action_stats, config.action_network)

    if args.image_maskgit_path is not None:
        model.init_weights()
        model.load_pretrained_image_weights(args.image_maskgit_path)
        if args.mu_transfer:
            model.set_mup_shapes(rescale_params=False)
    elif 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.
    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)

    no_decay = ["bias", "layer_norm.weight"]

    pretrained_params = {  # more accurately the params we want lower lr for, some weights like pos_embed_TSC are pre-trained but not treated as lower lr
        param_name
        for param_name, _ in model.named_parameters()
        if any(term in param_name for term in ("spatial_attn.qkv", "spatial_attn.proj", "mlp"))
    } if args.image_maskgit_path is not None else set()
    # Give pre-trained weights 10x lower learning rate
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters()
                       if not any(nd in n for nd in no_decay) and n not in pretrained_params],
            "weight_decay": args.weight_decay,
            "lr": args.learning_rate,
        },
        {
            "params": [p for n, p in model.named_parameters()
                       if any(nd in n for nd in no_decay) and n not in pretrained_params],
            "weight_decay": 0.0,
            "lr": args.learning_rate,
        },
        {
            "params": [p for n, p in model.named_parameters()
                       if not any(nd in n for nd in no_decay) and n in pretrained_params],
            "weight_decay": args.weight_decay,
            "lr": args.learning_rate * 0.1,
        },
        {
            "params": [p for n, p in model.named_parameters()
                       if any(nd in n for nd in no_decay) and n in pretrained_params],
            "weight_decay": 0.0,
            "lr": args.learning_rate * 0.1,
        },
    ]

    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)
    combined_dataset = torch.utils.data.ConcatDataset(train_datasets)

    batch_sampler = data_sampler.MultiTaskBatchSampler(
        dataset_num_samples,
        batch_size=args.per_device_train_batch_size,
        temperature=3. # the higher the more flat the distribution
    )
    dataset_traj_image = data_sampler.make_dataset_pie_plot(domains_list, dataset_num_samples)
    accelerator.log(({"dataset_mixture": wandb.Image(dataset_traj_image)}), log_kwargs={"wandb": {"commit": False}})
    dataset_weights = batch_sampler.generate_tasks_distribution().cpu().numpy()
    dataset_weight_image = data_sampler.make_dataset_pie_plot(domains_list, dataset_weights)
    accelerator.log(({"dataset_mixture_weight": wandb.Image(dataset_weight_image)}), log_kwargs={"wandb": {"commit": False}})

    train_dataloader = DataLoader(combined_dataset, batch_sampler=batch_sampler, collate_fn=collate_fn,
                                   num_workers=24, pin_memory=False)

    batch_val_sampler = data_sampler.MultiTaskBatchSampler(
        val_dataset_num_samples,
        batch_size=args.per_device_train_batch_size,
        temperature=4. # the higher the more flat the distribution
    )

    combined_val_dataset = torch.utils.data.ConcatDataset(val_datasets)
    eval_dataloader = DataLoader(combined_val_dataset, batch_sampler=batch_val_sampler, collate_fn=collate_fn,
                                    num_workers=24, pin_memory=False)

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

    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`.
    accelerator.wait_for_everyone()
    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 = len(train_datasets)
    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,
        "train_data_tokens": len(train_dataset) * seq_len,
        "effective_batch_size": effective_batch_size,
        "effective_batch_size_tokens": effective_batch_size * seq_len,
        "mixed_precision": accelerator.mixed_precision,
        "num_datasets": args.num_datasets
    })

    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!
    train(accelerator, model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, experiment_config, config, args)


if __name__ == "__main__":
    main()