train_reader.py

# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import time
import sys
import torch
import transformers
import numpy as np
import logging
import wandb
import os
import re
import warnings
#import converter

from datetime import datetime
from tqdm import tqdm
from pathlib import Path
from torch.utils.data import DataLoader, RandomSampler, DistributedSampler, SequentialSampler
from src.options import Options, Config

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

import src.slurm
import src.util
import src.evaluation
import src.data
import src.model


warnings.filterwarnings("ignore")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S %Z")

def main(opt):
    wandb.require("service")  # service improves wandb's handling of multiprocessing
    wandb.setup()
    src.util.variable_sanity_check(opt)
    mp.spawn(
        train,
        nprocs=opt.world_size,
        join=True,
        args=(
            opt,
        ),
    )


def train(rank, opt):

    torch.cuda.set_device(rank)
    device = f"cuda:{rank}"
    opt.device = device
    
    src.util.setup(rank, opt.world_size, opt.port, opt.master_addr)
    opt.group = dist.new_group(list(range(opt.world_size)))
    transformers.set_seed(opt.seed)
    best_dev_em = 0.0

    checkpoint_path = Path(opt.checkpoint_dir)/opt.name
    checkpoint_exists = checkpoint_path.exists()
    if opt.is_distributed:
        dist.barrier()
    checkpoint_path.mkdir(parents=True, exist_ok=True)
    #if not checkpoint_exists and opt.is_main:
    #    options.print_options(opt)
    #checkpoint_path, checkpoint_exists = util.get_checkpoint_path(opt)
    opt.is_main = (rank == 0)
    if opt.is_main:
        wandb.init(project="ATLAS")
    logger = src.util.init_logger(
        opt.is_main,
        opt.is_distributed,
        checkpoint_path / 'run.log'
    )

    model_name = 'google/t5-' + opt.model_size + '-lm-adapt'
    model_class = src.model.FiDT5

    #load data
    tokenizer = transformers.T5Tokenizer.from_pretrained(model_name)
    #Set bos token for tokenizer
    tokenizer.bos_token = '<pad>'
    collator = src.data.Collator(opt.text_maxlength, tokenizer, answer_maxlength=opt.answer_maxlength)

    # use golbal rank and world size to split the eval set on multiple gpus
    train_examples = src.data.load_data(
        opt.train_data, 
        rank=rank, 
        world_size=opt.world_size,
    )
    if opt.toy is True:
        train_examples = train_examples[:64]
    train_dataset = src.data.Dataset(train_examples, opt.n_context, finetune=not opt.is_pretrain)

    # use golbal rank and world size to split the eval set on multiple gpus
    eval_examples = src.data.load_data(
        opt.eval_data,
        rank=rank,
        world_size=opt.world_size,
    )
    if opt.toy is True:
        eval_examples = eval_examples[:4]
    eval_dataset = src.data.Dataset(eval_examples, opt.n_context, finetune=not opt.is_pretrain)

    # For toy setting


    # Load model
    # Note: Loading from model_path is NOT for continuing the paused train loop. This is for starting a new iteration.
    if not checkpoint_exists and opt.model_path == "none":
        logger.info("Initialize with HF pretrained model...")
        t5 = transformers.T5ForConditionalGeneration.from_pretrained(model_name)
        model = src.model.FiDT5(t5.config)
        model.load_t5(t5.state_dict())
        #model = model.to(opt.local_rank)
        #model = model.to(os.environ['LOCAL_RANK'])
        optimizer, scheduler = src.util.set_optim(opt, model)
        step, best_dev_em = 0, 0.0
        logger.info(f"Model loaded from {model_name}")
    elif opt.model_path == "none":
        load_path = checkpoint_path / 'checkpoint' / 'latest'
        model = model_class.from_pretrained(opt.model_path)
        step, best_dev_em = 0, 0.0
        optimizer, scheduler = src.util.set_optim(opt, model)
        #model, optimizer, scheduler, opt_checkpoint, step, best_dev_em = \
        #    src.util.load(model_class, load_path, opt, reset_params=False)
        logger.info(f"Model loaded from {load_path}")
    else:
        model = model_class.from_pretrained(opt.model_path)
        step, best_dev_em = 0, 0.0
        optimizer, scheduler = src.util.set_optim(opt, model)
        #model, optimizer, scheduler, opt_checkpoint, step, best_dev_em = \
        #    src.util.load(model_class, opt.model_path, opt, reset_params=True)
        logger.info(f"Model loaded from {opt.model_path}")

    model = model.to(device)
    model.set_checkpoint(opt.use_checkpoint)

    # Setup codes for distributed learning
    logger.info(f"Distrubited learning: {opt.is_distributed}")
    if opt.is_distributed:
        model = DDP(
            model, device_ids=[rank], output_device=rank, find_unused_parameters=False,
        )
        logger.info(f"Rank {rank}: Model configured for distributed learning")

    # Setup dataloader
    torch.manual_seed(rank + opt.seed) #different seed for different sampling depending on global_rank
    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(
        train_dataset,
        sampler=train_sampler,
        batch_size=opt.per_gpu_batch_size,
        drop_last=True,
        num_workers=0,
        collate_fn=collator
    )

    loss, curr_loss = 0.0, 0.0
    epoch = 1

    logger.info(f"Rank {rank}: Start training")

    # Training loop
    model.train()
    while step < opt.total_steps:
        epoch += 1
        for i, batch in enumerate(tqdm(train_dataloader)):
            step += 1
            (idx, labels, _, context_ids, context_mask) = batch
            train_loss = model(
                input_ids=context_ids.cuda(),
                attention_mask=context_mask.cuda(),
                labels=labels.cuda(),
                return_dict=False
            )[0]
            train_loss.backward()

            if step % opt.accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), opt.clip)
                optimizer.step()
                scheduler.step()
                model.zero_grad()

            train_loss = src.util.average_main(train_loss, opt)
            curr_loss += train_loss.item()

            if opt.is_main:
                log_dict = {"train": {"loss": train_loss}, "lr": scheduler.get_last_lr()[0]}
                wandb.log(log_dict)
            dist.barrier()

            if step % opt.eval_freq == 0:
                logging.info('Start evaluation')
                dev_em = evaluate(model, eval_dataset, tokenizer, collator, opt, rank)
                logging.info('Finished evaluation')
                model.train()
                if opt.is_main:
                    if dev_em > best_dev_em:
                        best_dev_em = dev_em
                        print(f"saving checkpoint to {checkpoint_path}...")
                        src.util.save(model, optimizer, scheduler, step, best_dev_em,
                                opt, checkpoint_path, 'best_dev')
                    log = f"train: {curr_loss/opt.eval_freq:.3f} |"
                    log += f"evaluation: {100*dev_em:.2f}EM |"
                    log += f"lr: {scheduler.get_last_lr()[0]:.5f}"
                    logger.info(log)

                    eval_log_dict = {"val": {"EM acc.": dev_em}}
                    wandb.log(eval_log_dict)
                    curr_loss = 0.
                dist.barrier()

            if opt.is_main and step % opt.save_freq == 0:
                src.util.save(model, optimizer, scheduler, step, best_dev_em,
                        opt, checkpoint_path, f"step-{step}")
            dist.barrier()

            if step > opt.total_steps:
                break

    dist.barrier()
    src.util.cleanup()


def evaluate(model, dataset, tokenizer, collator, opt, rank):
    sampler = SequentialSampler(dataset)
    dataloader = DataLoader(dataset,
        sampler=sampler,
        batch_size=opt.per_gpu_batch_size,
        drop_last=False,
        num_workers=10,
        collate_fn=collator
    )
    model.eval()
    total = 0
    exactmatch = []
    model = model.module if hasattr(model, "module") else model
    with torch.no_grad():
        count = 0
        flag = False
        for i, batch in enumerate(tqdm(dataloader)):
            # For pretraining, evaluate fn only provides one example
            if opt.is_pretrain and flag is True:
                break
            (idx, _, _, context_ids, context_mask) = batch
            outputs = model.generate(
                input_ids=context_ids.cuda(),
                attention_mask=context_mask.cuda(),
                max_length=50
            )

            for k, o in enumerate(outputs):
                skip_special_tokens = not opt.is_pretrain
                ans = tokenizer.decode(o, skip_special_tokens=skip_special_tokens)
                if opt.is_pretrain:
                    ans = re.sub(r"<pad>", '', ans)
                ex = dataset.get_example(idx[k])
                gold = [ex['target']] if opt.is_pretrain else ex['answers']
                score = src.evaluation.ems(ans, gold)
                total += 1
                exactmatch.append(score)
                if flag is False:
                    print("="*50)
                    print("Eval example")
                    print(f"questinon: {ex['question']}")
                    print(f"answers: {gold}")
                    print(f"prediction: {ans}")
                    print(f"score: {score}")
                    print("="*50)
                    count += 1
                if count > 4:
                    flag = True

    if opt.is_pretrain == False:
        exactmatch, total = src.util.weighted_average(np.mean(exactmatch), total, opt, rank)

    else:
        exactmatch = 0.0

    return exactmatch


if __name__ == "__main__":
    mp.set_start_method('spawn')
    options = Options()
    options.add_reader_options()
    options.add_optim_options()
    opt = options.parse()
    if opt.config == 'none':
        opt = options.get_options(use_reader=True, use_optim=True)
    else:
        opt = Config.from_yaml(opt.config)

    opt.is_distributed = torch.cuda.device_count() > 1
    if opt.is_distributed:
        logging.info("Initialize distributed training")

    opt.world_size = torch.cuda.device_count()

    print('='*50)
    print(f"config: {opt.__dict__}")
    print(f"\nworld size: {torch.cuda.device_count()}")
    print('='*50)

    main(opt)