steps/trainer.py

import time, sys, subprocess, json, re
from pathlib import Path
import os, random
import torch
import math, pickle
from tqdm import tqdm
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data.sampler import Sampler
import copy
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from torch.utils.data.distributed import DistributedSampler
import logging
# from data import librilight, gigaspeech, gigaspeech_waveform
from data import combined_dataset
from models import voice_star

from .trainer_utils import DistributedDynamicBatchSampler, StatefulDistributedSampler, StatefulSampler, AverageMeter, print_model_info
from .optim import ScaledAdam, Eden
import run_gen
import wandb, socket

class Trainer:
    
    def __init__(self, args, world_size, rank, local_rank):
        self.start_time = time.time()
        self.args = args
        if self.args.val_max_num_tokens == None:
            self.args.val_max_num_tokens = self.args.max_num_tokens 
        self.world_size, self.rank, self.local_rank = world_size, rank, local_rank
        self.device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
        if self.rank == 0:
            self.writer = SummaryWriter(args.exp_dir)
            self.wandb = wandb.init(project="voice_editor", name=args.exp_dir.split("/")[-1], config=args, dir=args.exp_dir, entity=self.args.wandb_entity)
        self.seed_everything(seed=self.args.seed)
        self.meters = self._setup_meters()

        self.progress, self.total_progress = self._setup_progress()

        self.model, self.trainables, self.optim_states, self.scheduler_states, self.phn2num = self._setup_models()

        self.train_dataset_length, self.train_sampler, self.train_loader, self.valid_loader = self._setup_dataloader() # both are use DistributedSampler, train sampler is stateful
        if self.args.num_steps != None:
            self.total_step = self.args.num_steps
            self.args.num_epochs = math.ceil(self.total_step / math.floor(self.train_dataset_length / self.args.batch_size)) if not self.args.dynamic_batching else None
        else:
            self.total_step = int(math.floor(self.train_dataset_length / self.args.batch_size))*self.args.num_epochs

        self.optimizer, self.scheduler = self._setup_optimizer()
        self.scaler = torch.cuda.amp.GradScaler()
        self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[self.local_rank], find_unused_parameters=False)
        self.early_stop_accu_steps = 0
        if self.rank == 0:
            if self.args.dynamic_batching:
                logging.info(f"max number of tokens per GPU in a training batch: {self.args.max_num_tokens}, max number of tokens per GPU in a inference batch: {self.args.val_max_num_tokens}")
            else:
                logging.info(f"batch size (per gpu): {self.args.batch_size}")
        
        self.args.inference_every_n_steps = getattr(self.args, "inference_every_n_steps", self.args.val_every_n_steps*5)
        assert self.args.inference_every_n_steps > self.args.val_every_n_steps and self.args.inference_every_n_steps % self.args.val_every_n_steps == 0, "inference_every_n_steps should be divisible by val_every_n_steps, otherwise the code will not get a chance to run inference"

    def train(self):
        flag = True
        skip_flag = False
        data_start_time = time.time()
        if self.progress['step'] >= self.total_step:
            if self.rank == 0:
                self.writer.close()
                self.wandb.finish()
            return
        while flag:
            self.train_sampler.set_epoch(self.progress['epoch'])
            for i, batch in enumerate(self.train_loader):
                if len(batch['y_lens']) < self.args.gradient_accumulation_steps:
                    continue
                data_end_time = time.time()
                self.model.train()
                if self.progress['step'] >= getattr(self.args, "uniform_weight_start_step", 1e50):
                    if self.progress['step'] == getattr(self.args, "uniform_weight_start_step", 1e50) and self.rank == 0:
                        logging.info("NOTE: start using uniform weight from step: {}".format(self.progress['step']))
                    self.args.codebook_weight = [2.5,2,1.5,0.6]
                    self.model.module.args.codebook_weight = [2.5,2,1.5,0.6]

                if self.progress['step'] >= self.total_step:
                    dist.barrier()
                    flag = False
                    self.validate_and_save()
                    if self.rank == 0:
                        self.writer.close()
                        self.wandb.finish()
                    break
                if isinstance(self.scheduler, Eden):
                    self.scheduler.step_epoch(self.progress['step']//self.args.pseudo_epoch_size + 1)
                if self.args.optimizer_name == "ScaledAdam":
                    cur_lr = self.scheduler.get_last_lr()[0]
                else:
                    lrs = [param_group['lr'] for param_group in self.optimizer.param_groups]
                    assert lrs[0] == lrs[1]
                    cur_lr = lrs[0]

                if self.rank == 0 and self.progress['step'] % self.args.tb_write_every_n_steps == 0:
                    self.writer.add_scalar("train/lr", cur_lr, self.progress['step'])
                    self.wandb.log({"train/lr": cur_lr}, step=self.progress['step'])

                all_inds = list(range(len(batch['y'])))
                sum_losses = 0
                sum_top10acc = 0
                sum_ntoken = 0
                sum_top10acc_cbi = [0 for _ in range(self.args.n_codebooks)]
                # extra losses
                sum_extra_losses = {}
                # when using prompt-based training, it's likely that due to prompt, the total length gets much longer, which make effective batch size in each accumulation step much bigger and then lead to OOM.
                # therefore we re-calculate graduent_accumulation_steps based on the effective batch size

                if self.args.neighbor_prompt_prob > 0:
                    effective_batch_size = self.args.max_num_tokens // self.args.gradient_accumulation_steps
                    total_batch_size = sum(batch['y_lens']).item()
                    cur_gradient_accumulation_steps = max(self.args.gradient_accumulation_steps, total_batch_size // effective_batch_size)
                    gas = torch.tensor(cur_gradient_accumulation_steps, dtype=torch.int, device=self.local_rank)
                    dist.all_reduce(gas, op=dist.ReduceOp.MAX)
                    cur_gradient_accumulation_steps = gas.item()
                    len_batch = torch.tensor(len(batch['y']), dtype=torch.int, device=self.local_rank)
                    dist.all_reduce(len_batch, op=dist.ReduceOp.MIN)
                    len_batch = len_batch.item()
                    cur_gradient_accumulation_steps = min(cur_gradient_accumulation_steps, len_batch)
                    # for those that cur_gradient_accumulation_steps * effective_batch_size < total_batch_size, we only use the first cur_gradient_accumulation_steps * effective_batch_size samples
                    cur_len = 0
                    final_all_inds = []
                    pointer = 0
                    while cur_len < self.args.max_num_tokens and pointer < len(all_inds):
                        cur_len += batch['y_lens'][pointer]
                        final_all_inds.append(all_inds[pointer])
                        pointer += 1
                    all_inds = final_all_inds
                else:
                    cur_gradient_accumulation_steps = self.args.gradient_accumulation_steps
                

                sum_losses_local = 0.0
                sum_top10acc_local = 0.0
                sum_entropy_loss_local = 0.0
                sum_ctc_loss_local = 0.0
                sum_ntoken_local = 0.0
                sum_top10acc_cbi_local = [0.0 for _ in range(self.args.n_codebooks)]

                global_nan_flag = 0
                for j in range(cur_gradient_accumulation_steps):
                    cur_ind = all_inds[j::cur_gradient_accumulation_steps]
                    cur_batch = {key: batch[key][cur_ind] for key in batch}

                    # Automatic casting
                    if self.args.precision == "float16":
                        precision_used = torch.float16
                    elif self.args.precision in ["bf16", "bfloat16"]:
                        precision_used = torch.bfloat16
                    else:
                        precision_used = torch.float32

                    with torch.amp.autocast('cuda', dtype=precision_used):
                        out = self.model(cur_batch, calc_loss=True)
                    if out is None:
                        continue

                    if torch.isnan(out['loss']).any():
                        local_nan_flag = torch.tensor(1, device=self.local_rank)
                    else:
                        local_nan_flag = torch.tensor(0, device=self.local_rank)

                    # All ranks check if *any* rank got a NaN
                    dist.all_reduce(local_nan_flag, op=dist.ReduceOp.SUM)
                    global_nan_flag = local_nan_flag.item()
                    if global_nan_flag > 0:
                        # Now *all* ranks break at the same j
                        logging.info(f"rank: {self.rank}. Loss at micro-batch {j} in step {self.progress['step']} was NaN on at least one rank; skipping.")
                        break

                    # Accumulate local values
                    record_loss = out['loss'].detach()
                    top10acc = out['top10acc'].detach()
                    effective_ntoken = out['effective_ntoken'].detach()

                    sum_losses_local += record_loss.item()
                    sum_top10acc_local += top10acc.item()
                    sum_ntoken_local += effective_ntoken.item()

                    # Optional losses
                    if 'entropy_loss' in out:
                        sum_entropy_loss_local += out['entropy_loss'].detach().item()
                    if 'ctc_loss' in out:
                        sum_ctc_loss_local += out['ctc_loss'].detach().item()

                    # Codebook accuracy
                    if 'top10acc_by_codebook' in out:
                        for cb in range(self.args.n_codebooks):
                            sum_top10acc_cbi_local[cb] += out['top10acc_by_codebook'][cb].detach().item()

                    # Backprop on this micro-batch
                    if self.args.optimizer_name == "ScaledAdam":
                        self.scaler.scale(out['loss']).backward()
                    else:
                        self.scaler.scale(out['loss'] / out['effective_ntoken']).backward()

                if global_nan_flag > 0:
                    # If *any* rank had NaN, skip this step
                    logging.info(f"rank: {self.rank}. Loss at one micro-batch in step {self.progress['step']} was NaN on at least one rank; skipping.")
                    self.progress['step'] += 1
                    self.progress['cur_step'] += 1
                    self.optimizer.zero_grad()
                    continue

                # Otherwise, do one big reduce for the summed metrics
                metrics_tensor = torch.tensor([
                    sum_losses_local,
                    sum_top10acc_local,
                    sum_entropy_loss_local,
                    sum_ctc_loss_local,
                    sum_ntoken_local
                ], device=self.local_rank, dtype=torch.float32)

                dist.all_reduce(metrics_tensor, op=dist.ReduceOp.SUM)

                # Also reduce the codebook array in one shot if needed
                codebook_tensor = torch.tensor(sum_top10acc_cbi_local, device=self.local_rank, dtype=torch.float32)
                dist.all_reduce(codebook_tensor, op=dist.ReduceOp.SUM)

                # Convert them back to Python scalars
                sum_losses = metrics_tensor[0].item()
                sum_top10acc = metrics_tensor[1].item()
                sum_entropy_loss = metrics_tensor[2].item()
                sum_ctc_loss = metrics_tensor[3].item()
                sum_ntoken = metrics_tensor[4].item()

                sum_top10acc_cbi = codebook_tensor.tolist()

                if self.args.optimizer_name != "ScaledAdam":
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.gradient_clip_val)

                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()

                if self.args.optimizer_name == "ScaledAdam":
                    self.scheduler.step_batch(self.progress['step'])
                else:
                    self.scheduler.step()
                
                # logging
                if self.rank == 0:
                    average_loss = sum_losses / sum_ntoken
                    average_top10acc = sum_top10acc / sum_ntoken
                    average_top10acc_cbi = [sum_top10acc_cbi[cb] / sum_ntoken * self.args.n_codebooks for cb in range(self.args.n_codebooks)]
                    self.meters['train_loss'].update(average_loss, batch['x'].shape[0]*self.world_size)
                    self.meters['train_top10acc'].update(average_top10acc, batch['x'].shape[0]*self.world_size)
                    self.meters['train_top10acc'].update(average_top10acc, batch['x'].shape[0]*self.world_size)
                    for cb in range(self.args.n_codebooks):
                        self.meters[f'train_top10acc_cb{cb+1}'].update(average_top10acc_cbi[cb], batch['x'].shape[0]*self.world_size)
                    self.meters['data_time'].update(data_end_time - data_start_time)
                    self.meters['train_time'].update(time.time() - data_end_time)

                    # log extra losses
                    for key in sum_extra_losses:
                        if "train_"+key not in self.meters:
                            self.meters["train_"+key] = AverageMeter()
                        self.meters["train_"+key].update(sum(sum_extra_losses[key])/len(sum_extra_losses[key]), batch['x'].shape[0]*self.world_size)

                    if self.progress['step'] % self.args.tb_write_every_n_steps == 0:
                        self.writer.add_scalar('train/loss', average_loss, self.progress['step'])
                        self.writer.add_scalar('train/top10acc', average_top10acc, self.progress['step'])
                        self.writer.add_scalar("train/ntokens", sum_ntoken, self.progress['step'])
                        self.wandb.log({"train/loss": average_loss, "train/top10acc": average_top10acc, "train/ntokens": sum_ntoken, "train/data_time": data_end_time - data_start_time, "train/train_time": time.time() - data_end_time}, step=self.progress['step'])

                        for cb in range(self.args.n_codebooks):
                            self.writer.add_scalar(f'train/top10acc_cb{cb+1}', average_top10acc_cbi[cb], self.progress['step'])
                            self.wandb.log({f'train/top10acc_cb{cb+1}': average_top10acc_cbi[cb]}, step=self.progress['step'])
                        self.writer.add_scalar("train/data_time", data_end_time - data_start_time, self.progress['step'])
                        self.writer.add_scalar("train/train_time", time.time() - data_end_time, self.progress['step'])
                        # write extra losses
                        for key in sum_extra_losses:
                            self.writer.add_scalar(f"train/{key}", sum(sum_extra_losses[key])/len(sum_extra_losses[key]), self.progress['step'])
                            self.wandb.log({f"train/{key}": sum(sum_extra_losses[key])/len(sum_extra_losses[key])}, step=self.progress['step'])
                    # logging.info(f"ntoken: {sum_ntoken}")
                        
                    # logging
                    if self.progress['step'] % self.args.print_every_n_steps == 0:
                        log_out = {}
                        log_out['cur_epoch'] = f"{self.progress['epoch']}/{self.args.num_epochs}" if self.args.num_epochs is not None else f"{self.progress['epoch']}"
                        log_out['cur_step'] = f"{int(self.progress['cur_step']+1)}"
                        log_out['total_step'] = f"{self.progress['step']}/{self.args.num_steps}"
                        log_out['lr'] = f"{cur_lr:.7f}"
                        log_out['ntokens'] = f"{sum_ntoken}"
                        for key in self.meters:
                            if self.meters[key].val != 0 or self.meters[key].avg != 0:
                                log_out[key] = f"{self.meters[key].val:.4f} ({self.meters[key].avg:.4f})" if isinstance(self.meters[key].val, float) else f"{self.meters[key].val}"
                        logging.info(log_out)
                        if np.isnan(self.meters['train_loss'].avg):
                            logging.warning("training diverged...")
                            raise RuntimeError("training diverged...")

                # save the model only
                if self.progress['step'] % self.args.save_every_n_steps == 0:
                    dist.barrier()
                    if self.rank == 0:
                        save_path = os.path.join(self.args.exp_dir,f"bundle_step{self.progress['step']}.pth")
                        self.save_progress(name=f"step{self.progress['step']}")
                        torch.save(
                            {
                                "model": self.model.module.state_dict(),
                                "args": self.args,
                                "phn2num": self.train_loader.dataset.phn2num,
                                "optimizer":  self.optimizer.state_dict(),
                                "scheduler": self.scheduler.state_dict(),
                            },save_path
                        )
                        logging.info(f"save model, optimizer, scheduler and progress at {save_path} at global step {self.progress['step']}")
                    dist.barrier()
                
                # validation and save models
                if self.progress['step'] % self.args.val_every_n_steps == 0:
                    dist.barrier()
                    continue_training = self.validate_and_save()
                    # broadcast continue_training to all processes, so that all processes gets into generation stage
                    continue_training = torch.tensor(int(continue_training), dtype=torch.int, device=self.local_rank)
                    dist.broadcast(continue_training, src=0)
                    continue_training = bool(continue_training.item())
                    dist.barrier() # need this to ensure all processes get to the next line?
                    logging.info(f"rank: {self.rank}, continue_training: {continue_training}")
                    if not continue_training: 
                        if self.rank == 0:
                            self.writer.close()
                            self.wandb.finish()
                        flag = False
                        break

                self.progress['step'] += 1
                self.progress['cur_step'] += 1

                data_start_time = time.time()
            self.progress['epoch'] += 1
            self.progress['cur_step'] = 0 # reset cur_step to be 0
        dist.destroy_process_group()

    def validate_and_save(self):
        self.model.eval()
        
        score = self.validate(self.valid_loader)

        if self.args.early_stop_threshold > 0:
            if self.progress['best_score'] - score < self.args.early_stop_threshold:
                self.early_stop_accu_steps += self.args.val_every_n_steps
                if self.early_stop_accu_steps >= self.args.early_stop_step-1:
                    logging.info(f"early stop based on self.args.early_stop_threshold: {self.args.early_stop_threshold}, and self.args.early_stop_step: {self.args.early_stop_step}")
                    logging.info(f"best validation score at step: {self.progress['best_step']}, and the score is {self.progress['best_score']:.4f}")
                    return False
            else:
                self.early_stop_accu_steps = 0

        if self.rank == 0:
            save_path = os.path.join(self.args.exp_dir,"bundle.pth")
            if os.path.isfile(save_path):
                os.system(f"mv {save_path} {save_path.replace('.pth', '_prev.pth')}")
            torch.save(
                {
                    "model": self.model.module.state_dict(),
                    "optimizer":  self.optimizer.state_dict(),
                    "scheduler": self.scheduler.state_dict(),
                    "args": self.args,
                    "phn2num": self.train_loader.dataset.phn2num
                    },save_path
            )
            self.save_progress()
            logging.info(f"save models, indices, acc and other statistics at {save_path} and {self.args.exp_dir}/progress.pkl at global step {self.progress['step']}")
            if (score < self.progress['best_score']):
                self.progress['best_step'] = self.progress['step']
                self.progress['best_score'] = score
                save_path = os.path.join(self.args.exp_dir,"best_bundle.pth")
                if os.path.isfile(save_path):
                    os.system(f"mv {save_path} {save_path.replace('.pth', '_prev.pth')}")
                torch.save(
                    {
                        "model": self.model.module.state_dict(),
                        "optimizer":  self.optimizer.state_dict(),
                        "scheduler": self.scheduler.state_dict(),
                        "args": self.args,
                        "phn2num": self.train_loader.dataset.phn2num
                    },save_path
                )
                logging.info(f"save *best* models at {save_path} at global step {self.progress['step']}")
        
        # sync best score and best step, so that all processes early stop at the same time
        best_score_tensor = torch.tensor(self.progress['best_score'], device=self.local_rank)
        dist.broadcast(best_score_tensor, src=0)
        self.progress['best_score'] = float(best_score_tensor.item())
        best_step_tensor = torch.tensor(self.progress['best_step'], device=self.local_rank)
        dist.broadcast(best_step_tensor, src=0)
        self.progress['best_step'] = int(best_step_tensor.item())
        dist.barrier()
        return True

    def validate(self, valid_loader=None, hide_progress=True):
        if valid_loader == None:
            valid_loader = self.valid_loader
        self.model.eval()

        start_val_time = time.time()
        sum_losses = 0
        sum_top10acc = 0
        sum_ntoken = 0
        sum_dur_loss = 0
        sum_dur_acc = 0
        sum_entropy_loss = 0
        sum_ctc_loss = 0

        sum_top10acc_cbi = [0 for _ in range(self.args.n_codebooks)]
        mean_perplexity_cbi = [0 for _ in range(self.args.n_codebooks)]

        with torch.no_grad():
            for i, batch in enumerate(tqdm(valid_loader, disable=hide_progress)):
                out = self.model(batch, calc_loss=True) # no reduction is applied to loss
                sum_losses += out['loss']
                sum_top10acc += out['top10acc']
                sum_ntoken += out['effective_ntoken']
                if "dur_loss" in out:
                    sum_dur_loss += out['dur_loss']
                    sum_dur_acc += out['dur_acc']
                if "entropy_loss" in out:
                    sum_entropy_loss += out['entropy_loss']
                if "ctc_loss" in out:
                    sum_ctc_loss += out['ctc_loss']
                # logging.info(f"iter {i}::: {sum_losses}, {sum_top10acc}, {sum_ntoken}")
                if 'top10acc_by_codebook' in out:
                    for cb in range(self.args.n_codebooks):
                        sum_top10acc_cbi[cb] += out['top10acc_by_codebook'][cb]
                        
                if 'perplexity_by_codebook' in out:
                    for cb in range(self.args.n_codebooks):
                        mean_perplexity_cbi[cb] += out['perplexity_by_codebook'][cb]
                # if i > 10:
                #     break

        
        dist.all_reduce(sum_losses, op=dist.ReduceOp.SUM)
        dist.all_reduce(sum_top10acc, op=dist.ReduceOp.SUM)
        dist.all_reduce(sum_ntoken, op=dist.ReduceOp.SUM)
        if "dur_loss" in out:
            dist.all_reduce(sum_dur_loss, op=dist.ReduceOp.SUM)
            dist.all_reduce(sum_dur_acc, op=dist.ReduceOp.SUM)
        if "entropy_loss" in out:
            dist.all_reduce(sum_entropy_loss, op=dist.ReduceOp.SUM)
        if "ctc_loss" in out:
            dist.all_reduce(sum_ctc_loss, op=dist.ReduceOp.SUM)

        if 'top10acc_by_codebook' in out:
            for cb in range(self.args.n_codebooks):
                dist.all_reduce(sum_top10acc_cbi[cb], op=dist.ReduceOp.SUM)
        
        if 'perplexity_by_codebook' in out:
            for cb in range(self.args.n_codebooks):
                dist.all_reduce(mean_perplexity_cbi[cb], op=dist.ReduceOp.SUM)
        
        val_loss = sum_losses / sum_ntoken
        val_top10acc = sum_top10acc / sum_ntoken
        
        if self.rank == 0:
            if "dur_loss" in out:
                val_dur_loss = sum_dur_loss / sum_ntoken
                val_dur_acc = sum_dur_acc / sum_ntoken
                self.meters['val_dur_loss'].update(val_dur_loss)
                logging.info(f"val dur_loss: {val_dur_loss:.5f}")
                self.meters['val_dur_acc'].update(val_dur_acc)
                logging.info(f"val dur_acc: {val_dur_acc:.5f}")
                self.writer.add_scalar("val/dur_loss", val_dur_loss, self.progress['step'])
                self.writer.add_scalar("val/dur_acc", val_dur_acc, self.progress['step'])
                self.wandb.log({"val/dur_loss": val_dur_loss, "val/dur_acc": val_dur_acc}, step=self.progress['step'])
            # logging
            self.meters['val_loss'].update(val_loss)
            logging.info(f"val loss: {val_loss:.5f}")
            self.writer.add_scalar("val/loss", val_loss, self.progress['step'])
            self.wandb.log({"val/loss": val_loss}, step=self.progress['step'])

            self.meters['val_top10acc'].update(val_top10acc)
            logging.info(f"val top10acc: {val_top10acc:.5f}")
            self.writer.add_scalar("val/top10acc", val_top10acc, self.progress['step'])
            self.wandb.log({"val/top10acc": val_top10acc}, step=self.progress['step'])
            for cb in range(self.args.n_codebooks):
                average_top10acc_cbi = sum_top10acc_cbi[cb] / sum_ntoken * self.args.n_codebooks
                self.meters[f'val_top10acc_cb{cb+1}'].update(average_top10acc_cbi)
                self.writer.add_scalar(f'val/top10acc_cb{cb+1}', average_top10acc_cbi, self.progress['step'])
                self.wandb.log({f'val/top10acc_cb{cb+1}': average_top10acc_cbi}, step=self.progress['step'])

                temp = mean_perplexity_cbi[cb]/len(valid_loader)
                self.writer.add_scalar(f'val/perplexity_cb{cb+1}', temp, self.progress['step'])
                self.wandb.log({f'val/perplexity_cb{cb+1}': temp}, step=self.progress['step'])
                
            average_perplexity = sum(mean_perplexity_cbi)/(self.args.n_codebooks*len(valid_loader))
            self.wandb.log({"val/average_perplexity": average_perplexity}, step=self.progress['step'])
            self.writer.add_scalar('val/average_perplexity', average_perplexity, self.progress['step'])
            
            # log entropy and ctc loss
            if "entropy_loss" in out:
                val_entropy_loss = sum_entropy_loss / ((i+1) * self.world_size)
                self.meters['val_entropy_loss'].update(val_entropy_loss)
                logging.info(f"val entropy_loss: {val_entropy_loss:.5f}")
                self.writer.add_scalar("val/entropy_loss", val_entropy_loss, self.progress['step'])
                self.wandb.log({"val/entropy_loss": val_entropy_loss}, step=self.progress['step'])
            if "ctc_loss" in out:
                val_ctc_loss = sum_ctc_loss / ((i+1) * self.world_size)
                self.meters['val_ctc_loss'].update(val_ctc_loss)
                logging.info(f"val ctc_loss: {val_ctc_loss:.5f}")
                self.writer.add_scalar("val/ctc_loss", val_ctc_loss, self.progress['step'])
                self.wandb.log({"val/ctc_loss": val_ctc_loss}, step=self.progress['step'])

            logging.info(f"validation takes: {time.time() - start_val_time:.2f}s")
            logging.info(f"Step [{self.progress['step']}/{self.total_step}]\t Time elapsed {(time.time() - self.start_time)/3600.:.2f}h, Val Loss: {val_loss:.4f}, Val Top10Acc: {val_top10acc:.4f}")

        return val_loss.item()

    def _setup_meters(self):
        meters = {}
        meter_names = ['train_loss', 'val_loss', 'train_top10acc', 'val_top10acc', 'data_time', 'train_time']
        meter_names += ['train_dur_loss', 'train_dur_acc', 'val_dur_loss', 'val_dur_acc']
        meter_names += ['val_perplexity']
        meter_names += [f'train_top10acc_cb{cb+1}' for cb in range(self.args.n_codebooks)]
        meter_names += [f'val_top10acc_cb{cb+1}' for cb in range(self.args.n_codebooks)]
        meter_names += [f'val_perplexity_cb{cb+1}' for cb in range(self.args.n_codebooks)]
        for name in meter_names:
            meters[name] = AverageMeter()
        return meters
    def _setup_progress(self):
        """
        Need to customize it
        """
        progress = {}
        progress['best_step'] = 1
        progress['best_score'] = np.inf # this records loss value
        progress['step'] = 1
        progress['epoch'] = 1
        progress['cur_step'] = 0 # step in the current epoch, for resuming the sampler
        total_progress = []
        # if self.args.resume or self.args.validate:
        if self.args.resume:
            progress_pkl = "%s/progress.pkl" % self.args.exp_dir
            with open(progress_pkl, "rb") as f:
                total_progress = pickle.load(f)
                progress['best_step'], progress['best_score'], progress['step'], progress['epoch'], progress['cur_step'], _ = total_progress[-1]
            if self.rank == 0:
                logging.info("\nResume training from:")
                logging.info("  epoch = %s" % progress['epoch'])
                logging.info("  cur_step = %s" % progress['cur_step'])
                logging.info("  step = %s" % progress['step'])
                logging.info("  best_step = %s" % progress['best_step'])
                logging.info("  best_score = %s" % progress['best_score'])
        return progress, total_progress
    
    def save_progress(self, name=None):
        self.total_progress.append([self.progress['best_step'], self.progress['best_score'], int(self.progress['step']+1), self.progress['epoch'], int(self.progress['cur_step']+1), time.time() - self.start_time])
        if name is not None:
            progress_fn = f"{self.args.exp_dir}/progress_{name}.pkl"
        else:
            progress_fn = f"{self.args.exp_dir}/progress.pkl"
        with open(progress_fn, "wb") as f:
            pickle.dump(self.total_progress, f)

    def _setup_dataloader(self):
        train_dataset, val_dataset = combined_dataset.dataset(self.args, 'train'), combined_dataset.dataset(self.args, 'valid') # need to change 'train' to 'valid' in actual training

        if self.args.dynamic_batching:
            train_sampler = DistributedDynamicBatchSampler(train_dataset, self.args, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True, lengths_list=train_dataset.lengths_list, verbose=True, epoch=0)
            valid_sampler = DistributedDynamicBatchSampler(val_dataset, self.args, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True, lengths_list=val_dataset.lengths_list, verbose=True, epoch=0)
        else:
            train_sampler = StatefulDistributedSampler(train_dataset, self.args.batch_size//self.world_size, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True)
            valid_sampler = DistributedSampler(val_dataset, num_replicas=self.world_size, rank=self.rank, shuffle=False, seed=self.args.seed, drop_last=False)
            
        if self.progress['step'] > 1:
            train_sampler.set_epoch_resume(self.progress['epoch'], self.progress['cur_step'])
            assert self.phn2num != None

        if self.phn2num != None:
            train_dataset.phn2num = self.phn2num
            val_dataset.phn2num = self.phn2num

        if self.args.dynamic_batching:
            train_loader = torch.utils.data.DataLoader(train_dataset, 
                            batch_sampler=train_sampler, 
                            num_workers=self.args.num_workers,
                            collate_fn=train_dataset.collate, persistent_workers=True
                            )
            valid_loader = torch.utils.data.DataLoader(val_dataset, 
                            batch_sampler=valid_sampler, 
                            num_workers=self.args.num_workers,
                            collate_fn=val_dataset.collate, persistent_workers=True
                            )
        else:
            train_loader = torch.utils.data.DataLoader(train_dataset, 
                            batch_size=self.args.batch_size, sampler=train_sampler, num_workers=self.args.num_workers,
                            collate_fn=train_dataset.collate, persistent_workers=True
                            )
            valid_loader = torch.utils.data.DataLoader(val_dataset, 
                            batch_size=self.args.batch_size, sampler=valid_sampler,
                            num_workers=self.args.num_workers,
                            collate_fn=val_dataset.collate, persistent_workers=True
                            )
        return len(train_dataset), train_sampler, train_loader, valid_loader
        

        
    def _setup_models(self):
        model = voice_star.VoiceStar(self.args)

        if self.rank == 0:
            logging.info(model)
            logging.info("model parameters")
            print_model_info(model)
        
        phn2num = None
        optim_states = None
        scheduler_states = None
        if self.progress['step'] > 1:
            bundle = torch.load(os.path.join(self.args.exp_dir, "bundle.pth"), map_location="cpu")
            model.load_state_dict(bundle['model'])
            optim_states = bundle['optimizer']
            scheduler_states = bundle['scheduler']
            phn2num = bundle['phn2num']
            if self.rank == 0:
                logging.info("loaded parameters and data indices from epoch %d, global step %d" % (self.progress['epoch'], self.progress['step']))
            del bundle['model']

        if self.args.load_model_from != None and self.progress['step'] <= 1:
            logging.info(f"load weights from {self.args.load_model_from}")
            sd = torch.load(self.args.load_model_from, map_location="cpu")
            if hasattr(model, "carefully_load_state_dict"):
                model.carefully_load_state_dict(sd['model'])
            else:
                model.load_state_dict(sd['model'])
            phn2num = sd['phn2num']
            del sd
        

        #### below operations is for getting params for optimizer, which is at wrapper level ###
        if self.args.optimizer_name == "ScaledAdam":
            trainables = [p for p in model.parameters() if p.requires_grad]
        else:
            no_decay = [".bias", ".audio_embeddings.weight", ".text_embeddings.weight", ".norm.weight", ".norm1.weight", ".norm2.weight"]
            optimizer_grouped_parameters = [
                {
                    "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad],
                    "weight_decay": 0.0,
                },
            ]
            if len(optimizer_grouped_parameters[1]['params']) == 0:
                logging.info("there is no embedding weights, bias, and layernorm parameters in the model, which should be True, check model parameter names")
                trainables = optimizer_grouped_parameters[0]
            else:
                trainables = optimizer_grouped_parameters
        #### below operations is for getting params for optimizer, which is at wrapper level ###
        model.to(self.device)

        return model, trainables, optim_states, scheduler_states, phn2num

    
    def _setup_optimizer(self):
        if self.args.optimizer_name == "ScaledAdam":
            parameters_names = []
            _model = self.model.module if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) else self.model
            parameters_names.append([n for n,p in self.model.named_parameters() if p.requires_grad])
            optimizer = ScaledAdam(
                self.trainables,
                lr=self.args.lr,
                betas=(0.9, 0.95),
                clipping_scale=2.0,
                parameters_names=parameters_names,
                show_dominant_parameters=False,
                clipping_update_period=self.args.clipping_update_period,
            )
            scheduler = Eden(optimizer, self.args.reduce_lr_start_step, self.args.reduce_lr_start_epoch, warmup_batches=self.total_step * self.args.warmup_fraction) # NOTE: if using ScaledAdam, we will use the Eden scheduler!

        else:
            optimizer = AdamW(self.trainables, lr=self.args.lr)
            warmup_steps = self.total_step * self.args.warmup_fraction
            def lr_lambda(current_step: int):
                if current_step < warmup_steps:
                    return float(current_step) / float(max(1, warmup_steps))
                return max(
                    0.0, float(self.total_step - current_step) / float(max(1, self.total_step - warmup_steps))
                )

            scheduler = LambdaLR(optimizer, lr_lambda, last_epoch=-1)
            
        # if resume
        if self.progress['step'] > 1:
            optimizer.load_state_dict(self.optim_states)
            for state in optimizer.state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        state[k] = v.cuda()
            del self.optim_states

            scheduler.load_state_dict(self.scheduler_states)

        optimizer.zero_grad()
        return optimizer, scheduler
    
    def seed_everything(self, seed=1):
        os.environ['PYTHONHASHSEED'] = str(seed)
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True