third_party/InternVideo/InternVideo2/multi_modality/tasks/pretrain.py

import datetime
import logging
import time
from os.path import join

import pandas as pd
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import wandb

from dataset.serialize import local_broadcast_process_authkey
from dataset import MetaLoader, MetaLoader_rs2, create_dataset, create_loader, create_sampler, create_stateful_sampler
from models import *
from tasks.retrieval_utils import evaluation_wrapper
from tasks.shared_utils import get_media_types, setup_model
from utils.basic_utils import (MetricLogger, SmoothedValue,
                               remove_files_if_exist, setup_seed)
from utils.config_utils import setup_main
from utils.distributed import get_rank, get_world_size, is_main_process
from utils.logger import log_dict_to_wandb, setup_wandb
try:
    from petrel_client.client import Client
except:
    Client = None
import io
import os
import shutil

logger = logging.getLogger(__name__)

ceph_ckpt_bucket = "shdd:s3://avp_ckpt"


def train(
    model,
    train_loaders,
    optimizer,
    tokenizer,
    epoch,
    global_step,
    device,
    scheduler,
    scaler,
    config,
    skip_num=0
):

    try:
        ceph_ckpt_path = f"{ceph_ckpt_bucket}/{config.output_dir.split('/')[-3]}/{config.output_dir.split('/')[-2]}/{config.output_dir.split('/')[-1]}"
        client_ckpt = Client(conf_path='~/petreloss.conf')
    except Exception as e:
        print(e)
        logger.info("Ceph is not working!!!")


    if config.use_half_precision: 
        if config.get('use_bf16', False):
            cast_dtype = torch.bfloat16
        else:
            cast_dtype = torch.float16
    else:
        cast_dtype = None

    model.train()

    metric_logger = MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", SmoothedValue(window=100, fmt="{value:.6f}"))
    metric_logger.add_meter("temperature", SmoothedValue(window=100, fmt="{value:.4f}"))
    loss_names = ["loss_" + k for k, v in config.criterion.loss_weight.items() if v != 0]
    
    if config.get("use_raw_text", False): # for cosa
        loss_names = loss_names + ["c_loss_" + k for k, v in config.criterion.loss_weight.items() if v != 0]
    uta_all = config.criterion.get('uta_all', False)

    media_types = get_media_types(train_loaders)

    for name in loss_names:
        for m in media_types:
            metric_logger.add_meter(
                f"{m}-{name}", SmoothedValue(window=100, fmt="{value:.4f}")
            )

    header = f"Train Epoch: [{epoch}]"
    log_freq = config.log_freq

    if config.distributed:
        for d in train_loaders:
            d.sampler.set_epoch(epoch)

    if config.get('use_iter_train', False):
        train_loader = MetaLoader_rs2(name2loader=dict(list(zip(media_types, train_loaders))), skip_num=skip_num)
    else:
        train_loader = MetaLoader(name2loader=dict(list(zip(media_types, train_loaders))))

    model_without_ddp = model.module if config.distributed else model
    iterator = metric_logger.log_every(train_loader, log_freq, header)
    
    begin_step = global_step % len(train_loader)
    logger.info(f"Epoch={epoch}, begin_step={begin_step} save_ckpt_iter={config.get('save_ckpt_iter', None)}")

    for local_step, (media_type, (media, text, idx)) in enumerate(iterator):
        if local_step < begin_step: 
            logger.warn(f"Jump local_step: {local_step} (begin_step={begin_step})!!!")
            metric_logger.update(lr=optimizer.param_groups[0]["lr"])
            continue

        if config.get("save_ckpt_iter", None) is not None: # and not is_iter_resume:
            if local_step != 0 and local_step % config.get("save_ckpt_iter") == 0:
                if hasattr(config, "deepspeed") and config.deepspeed.enable: 
                    tag = f"ckpt_e{epoch:02d}_local{local_step}_global{global_step}"
                    client_state = {'epoch': epoch, 'global_step': global_step, 'local_step': local_step}
                    model.save_checkpoint(config.output_dir, tag=tag, save_latest=False, client_state=client_state)
                    logger.info(f"save ckpt file to local ({config.output_dir}/{tag})!!!")
                elif is_main_process():
                    state_dict = model_without_ddp.state_dict()
                    for k in config.get("no_save_params_prefix", []):
                        kk = [x for x in state_dict.keys() if x.startswith(k)]
                        logger.info(f"Not saving {len(kk)} params with prefix {k}")
                        for kkk in kk:
                            state_dict.pop(kkk)

                    save_obj = {
                        "model": state_dict,
                        "optimizer": optimizer.state_dict(),
                        "scheduler": scheduler.state_dict(),
                        "scaler": scaler.state_dict(),
                        "config": config,
                        "epoch": epoch,
                        "local_step": local_step,
                        "global_step": global_step,
                    }
                    try:
                        with io.BytesIO() as buffer:
                            torch.save(save_obj, buffer)
                            client_ckpt.put(f"{ceph_ckpt_path}/ckpt_{epoch:02d}_local{local_step}_global{global_step}.pth", buffer.getvalue())
                            logger.info(f"Save to ceph ({ceph_ckpt_path}/ckpt_{epoch:02d}_local{local_step}_global{global_step}.pth)!!!")
                    except Exception as e:
                        print(e)
                        torch.save(save_obj, join(config.output_dir, f"ckpt_{epoch:02d}_local{local_step}_global{global_step}.pth"))
                        logger.warn(f"Ceph is not working, save to local ({join(config.output_dir, f'ckpt_{epoch:02d}_local{local_step}_global{global_step}.pth')})!!!")

        if media_type == 'audio_video':
            if type(media[0]) is list:
                assert len(media[0]) == 2
                audio = [media[0][0].to(device, dtype=cast_dtype, non_blocking=True), media[0][1].to(device, non_blocking=True)]
            else:
                audio = media[0].to(device, dtype=cast_dtype, non_blocking=True)
            video = media[1].to(device, dtype=cast_dtype, non_blocking=True)
            media = [audio, video]
        else:
            media = media.to(device, dtype=cast_dtype, non_blocking=True)
        idx = idx.to(device, non_blocking=True)
        if config.get("use_raw_text", False) or config.get("use_cosa", False):
            max_length = config.inputs.max_txt_l[media_type]
        else:
            if type(text) is dict:
                text_input = {}
                for k in text.keys():
                    text_input[k] = tokenizer(
                    text[k],
                    padding="max_length",
                    truncation=True,
                    max_length=config.inputs.max_txt_l[media_type],
                    return_tensors="pt",
                ).to(
                    device) # change from "longest" to "max_length"
            else:
                text_input = tokenizer(
                    text,
                    padding="max_length",
                    truncation=True,
                    max_length=config.inputs.max_txt_l[media_type],
                    return_tensors="pt",
                ).to(
                    device) # change from "longest" to "max_length"
            

        if hasattr(config, "deepspeed") and config.deepspeed.enable:
            loss_dict = model(media, text_input, idx=idx, media_type=media_type)
            loss = sum(loss_dict.values())

            model.backward(loss)
            model.step()

        else: # NOTE We shouldn't use scaler if we only involve bf16, check this!
            with torch.cuda.amp.autocast(enabled=config.use_half_precision, dtype=cast_dtype):
                loss_dict = model(media, text_input, idx=idx, media_type=media_type)
                loss = sum(loss_dict.values())

            if not config.use_half_precision or config.get('use_bf16', False):
                optimizer.zero_grad()
                loss.backward()
                if config.optimizer.max_grad_norm > 0:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.optimizer.max_grad_norm)
                    
                optimizer.step()
                scheduler.step()
            else:
                optimizer.zero_grad()
                scaler.scale(loss).backward()
                if config.optimizer.max_grad_norm > 0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.optimizer.max_grad_norm)
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()

        # logging
        for name in loss_names:
            if name in loss_dict.keys():
                value = loss_dict[name]
                value = value if isinstance(value, float) else value.item()
                metric_logger.update(**{f"{media_type}-{name}": value})
        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
        metric_logger.update(temperature=model_without_ddp.temp.item())

        if is_main_process() and config.wandb.enable and global_step % log_freq == 0:
            try:
                logs = metric_logger.get_global_avg_dict()
                log_dict_to_wandb(logs, step=global_step, prefix="train/")
            except Exception as e:
                logger.warn("Wandb is not working!!!")
                print(e)
            
        global_step += 1

        if config.debug and global_step % 20 == 0:
            logger.info("debug mode, break training loop")
            break

        if config.debug and global_step % (2 * log_freq + 3) == 0:
            logger.info("debug mode, break training loop")
            break

    # gather the stats from all processes
    metric_logger.synchronize_between_processes()
    logger.info(f"Averaged stats: {metric_logger.global_avg()}")
    return global_step


def setup_dataloaders(config, mode="pt"):
    # train datasets, create a list of data loaders
    logger.info(f"Creating dataset for {mode} use_iter_train={config.get('use_iter_train', False)}")
    train_datasets = create_dataset(f"{mode}_train", config)
    media_types = get_media_types(train_datasets)

    if config.get('use_iter_train', False):
        if config.distributed:
            batch_size = [config.inputs.batch_size[k] for k in media_types] # batch_size for each GPU
            samplers = create_stateful_sampler(train_datasets, batch_size)
        else:
            raise NotImplementedError
    else:
        if config.distributed:
            num_tasks = get_world_size()
            global_rank = get_rank()
            samplers = create_sampler(
                train_datasets, [True] * len(media_types), num_tasks, global_rank
            )
        else:
            samplers = [None] * len(media_types)

    train_loaders = create_loader(
        train_datasets,
        samplers,
        batch_size=[config.inputs.batch_size[k] for k in media_types],
        num_workers=[config.num_workers] * len(media_types),
        is_trains=[True] * len(media_types),
        collate_fns=[None] * len(media_types),
    )  # [0]

    # test datasets, a mapping from dataset name to data loader
    test_datasets, test_dataset_names = create_dataset(f"{mode}_eval", config)
    test_loaders = create_loader(
        test_datasets,
        [None] * len(test_datasets),
        batch_size=[config.inputs.batch_size_test[d.media_type] for d in test_datasets],
        num_workers=[config.num_workers] * len(test_datasets),
        is_trains=[False] * len(test_datasets),
        collate_fns=[None] * len(test_datasets),
    )
    test_name2loaders = {k: v for k, v in zip(test_dataset_names, test_loaders)}
    return train_loaders, test_name2loaders, media_types


def main(config):
    if config.get('use_flash_sdp', False):
        torch.backends.cuda.enable_flash_sdp(enabled=True)
    elif config.get('use_mem_efficient_sdp', False):
        torch.backends.cuda.enable_mem_efficient_sdp(enabled=True)

    try:
        ceph_ckpt_path = f"{ceph_ckpt_bucket}/{config.output_dir.split('/')[-3]}/{config.output_dir.split('/')[-2]}/{config.output_dir.split('/')[-1]}"
        client_ckpt = Client(conf_path='~/petreloss.conf')
    except Exception as e:
        print(e)
        logger.info("Ceph is not working!!!")
        
    if is_main_process() and config.wandb.enable:
        try:
            run = setup_wandb(config)
            logger.info("Wandb is working!!!")
        except Exception as e:
            logger.warn("Wandb is not working!!!")
            print(e)

    is_pretrain = config.mode == "pt"

    logger.info(f"train_file: {config.train_file}")

    setup_seed(config.seed + get_rank())
    device = torch.device(config.device)

    train_loaders, test_name2loaders, train_media_types = setup_dataloaders(
        config, mode=config.mode
    )
    num_steps_per_epoch = sum(len(d) for d in train_loaders)
    config.scheduler.num_training_steps = num_steps_per_epoch * config.scheduler.epochs
    config.scheduler.num_warmup_steps = num_steps_per_epoch * config.scheduler.warmup_epochs
    # set cudnn.benchmark=True only when input size is fixed
    # https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
    cudnn.benchmark = len(train_media_types) == 1

    print(f"\033[31m CURRENT NODE NAME: {os.environ['SLURMD_NODENAME']} dataloader is OK {datetime.datetime.now().strftime('%Y-%m-%d-%H_%M_%S')}!!! \033[0m")

    find_unused_parameters = config.model.get('find_unused_parameters', False)
    logger.info(f"find_unused_parameters={find_unused_parameters}")

    model_cls = eval(config.model.get('model_cls'))
    (
        model,
        model_without_ddp,
        optimizer,
        scheduler,
        scaler,
        tokenizer,
        start_epoch,
        global_step,
    ) = setup_model(
        config,
        model_cls=model_cls,
        add_decoder=False,
        pretrain=is_pretrain,
        find_unused_parameters=find_unused_parameters,
    )

    if is_main_process() and config.wandb.enable:
        try:
            wandb.watch(model)
        except Exception as e:
            logger.warn("Wandb is not working!!!")
            print(e)

    best = 0
    best_epoch = 0
    if type(config.best_key) is str:
        best_key = [config.best_key, "t2v_r1"]
    elif type(config.best_key) is list and len(config.best_key) == 2:
        best_key = config.best_key
    else:
        raise NotImplementedError(config.best_key)

    best_ckpt_id = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    logger.info(f"Start training, start_epoch={start_epoch}")
    start_time = time.time()
    start_step = start_epoch * num_steps_per_epoch
    for epoch in range(start_epoch, config.scheduler.epochs):
        if not config.evaluate:
            global_step = train(
                model,
                train_loaders,
                optimizer,
                tokenizer,
                epoch,
                global_step,
                device,
                scheduler,
                scaler,
                config,
                skip_num = global_step - start_step
            )

            if hasattr(config, "deepspeed") and config.deepspeed.enable:
                if config.get("save_latest", False):
                    tag = "ckpt_latest"
                else:
                    tag = f"ckpt_{epoch:02d}"

                client_state = {'epoch': epoch, 'global_step': global_step}
                model.save_checkpoint(config.output_dir, tag=tag, save_latest=False, client_state=client_state)

                logger.info(f"save ckpt file to local ({config.output_dir}/{tag})!!!") 
                if is_main_process() and config.get("delete_ds_optim_states", False):
                    if config.get("save_latest", False):
                        if epoch == (config.scheduler.epochs - 1): # last epoch
                            last_tag = "ckpt_latest"
                            last_ckpt_path = f"{config.output_dir}/{last_tag}"
                            if os.path.exists(last_ckpt_path):
                                logger.info(f"remove optim states in ({config.output_dir}/{last_tag})!!!")
                                for file in os.listdir(last_ckpt_path):
                                    if file.endswith('optim_states.pt'):
                                        os.remove(os.path.join(last_ckpt_path, file))
                    else:
                        last_tag = f"ckpt_{epoch-1:02d}"
                        last_ckpt_path = f"{config.output_dir}/{last_tag}"
                        if os.path.exists(last_ckpt_path):
                            logger.info(f"remove optim states in ({config.output_dir}/{last_tag})!!!")
                            for file in os.listdir(last_ckpt_path):
                                if file.endswith('optim_states.pt'):
                                    os.remove(os.path.join(last_ckpt_path, file))
                        
                        if epoch == (config.scheduler.epochs - 1): # last epoch
                            last_tag = f"ckpt_{epoch:02d}"
                            last_ckpt_path = f"{config.output_dir}/{last_tag}"
                            if os.path.exists(last_ckpt_path):
                                logger.info(f"remove optim states in ({config.output_dir}/{last_tag})!!!")
                                for file in os.listdir(last_ckpt_path):
                                    if file.endswith('optim_states.pt'):
                                        os.remove(os.path.join(last_ckpt_path, file))

            if is_main_process():
                if not (hasattr(config, "deepspeed") and config.deepspeed.enable):
                    state_dict = model_without_ddp.state_dict()
                    for k in config.get("no_save_params_prefix", []):
                        kk = [x for x in state_dict.keys() if x.startswith(k)]
                        logger.info(f"Not saving {len(kk)} params with prefix {k}")
                        for kkk in kk:
                            state_dict.pop(kkk)

                    save_obj = {
                        "model": state_dict,
                        "optimizer": optimizer.state_dict(),
                        "scheduler": scheduler.state_dict(),
                        "scaler": scaler.state_dict(),
                        "config": config,
                        "epoch": epoch,
                        "global_step": global_step,
                    }
                    try:
                        with io.BytesIO() as buffer:
                            torch.save(save_obj, buffer)
                            if config.get("save_latest", False):
                                client_ckpt.put(f"{ceph_ckpt_path}/ckpt_latest.pth", buffer.getvalue())
                                logger.info(f"Save to ceph ({ceph_ckpt_path}/ckpt_latest.pth)!!!")
                            else:
                                client_ckpt.put(f"{ceph_ckpt_path}/ckpt_{epoch:02d}.pth", buffer.getvalue())
                                logger.info(f"Save to ceph ({ceph_ckpt_path}/ckpt_{epoch:02d}.pth)!!!")
                    except Exception as e:
                        print(e)
                        if config.get("save_latest", False):
                            torch.save(save_obj, join(config.output_dir, "ckpt_latest.pth"))
                            logger.warn(f"Ceph is not working, save to local ({join(config.output_dir, 'ckpt_latest.pth')})!!!")
                        else:
                            torch.save(save_obj, join(config.output_dir, f"ckpt_{epoch:02d}.pth"))
                            logger.warn(f"Ceph is not working, save to local ({join(config.output_dir, f'ckpt_{epoch:02d}.pth')})!!!")



        if config.get("jump_evaluate", False) and not config.evaluate:
            logger.warn(f"Jump the evaluation'))!!!")
        else:
            try:
                eval_res = {}
                for test_name, test_loader in test_name2loaders.items():
                    if test_name not in config.test_types:
                        logger.info(
                            f"Skip eval {test_name} split. All test_types {config.test_types}"
                        )
                        continue
                    res = evaluation_wrapper(
                        model_without_ddp, test_loader, tokenizer, device, config, prefix=test_name
                    )
                    eval_res.update(res)

                if is_main_process():
                    # log to wandb
                    if config.wandb.enable:
                        try:
                            for p, v in eval_res.items():
                                log_dict_to_wandb(v, step=global_step, prefix=p)
                        except Exception as e:
                            logger.warn("Wandb is not working!!!")
                            print(e)

                    try:
                        cur_recall = eval_res[best_key[0]][best_key[1]]
                    except Exception as e:
                        logger.warn(e)
                        print(e)
                        # print(eval_res)
                        cur_recall = best - 1

                    eval_res = pd.DataFrame(eval_res)
                    logger.info(f"Epoch {epoch}")
                    logger.info(f"\n{eval_res.transpose().to_string(max_cols=30)}")

                    eval_res.to_json(join(config.output_dir, "eval_res_latest.json"))

                    state_dict = model_without_ddp.state_dict()

                    for k in config.get("no_save_params_prefix", []):
                        kk = [x for x in state_dict.keys() if x.startswith(k)]
                        logger.info(f"Not saving {len(kk)} params with prefix {k}")
                        for kkk in kk:
                            state_dict.pop(kkk)

                    if not config.evaluate and cur_recall > best:
                        if not (hasattr(config, "deepspeed") and config.deepspeed.enable):
                            try:
                                with io.BytesIO() as buffer:
                                    torch.save(save_obj, buffer)
                                    client_ckpt.put(f"{ceph_ckpt_path}/ckpt_best_{best_ckpt_id}.pth", buffer.getvalue())
                                    logger.info(f"Save to ceph ({f'{ceph_ckpt_path}/ckpt_best_{best_ckpt_id}.pth'})!!!")
                            except Exception as e:
                                print(e)
                                torch.save(save_obj, join(config.output_dir, f"ckpt_best_{best_ckpt_id}.pth"))
                                logger.warn(f"Ceph is not working, save to local ({join(config.output_dir, f'ckpt_best_{best_ckpt_id}.pth')})!!!")
                        else:
                            
                            now_ckpt_path = f"{config.output_dir}/{tag}/mp_rank_00_model_states.pt"
                            best_ckpt_path = f"{config.output_dir}/best_mp_rank_00_model_states.pt"

                            if os.path.exists(now_ckpt_path):
                                shutil.copy(now_ckpt_path, best_ckpt_path)
                                logger.info(f"Copy {now_ckpt_path} to {best_ckpt_path}!!!")
                            else:
                                logger.warn(f"Can't find {now_ckpt_path}, there's some wrong!!!")

                        eval_file = "eval_res_best.json"
                        eval_res.to_json(join(config.output_dir, eval_file))
                        best = cur_recall
                        best_epoch = epoch
            except Exception as e:
                logger.warn("Something wrong when eval or save!!!")
                print(e)
                if config.evaluate:
                    raise e


        if config.evaluate:
            break
        
        start_step = global_step

        dist.barrier()

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    logger.info(f"Training time {total_time_str}")
    logger.info(f"best epoch {best_epoch} [best_key {best_key}]")
    logger.info(f"Checkpoints and Logs saved at {config.output_dir}")

    if is_main_process() and config.wandb.enable:
        try:
            run.finish()
        except Exception as e:
            logger.warn("Wandb is not working!!!")
            print(e)


if __name__ == "__main__":
    print(f"\033[31m NODE LIST: {os.environ['SLURM_NODELIST']} \033[0m")
    logger.info(f"NODE LIST: {os.environ['SLURM_NODELIST']}")
    cfg = setup_main()
    local_broadcast_process_authkey()
    main(cfg)