init

steps/trainer.py

@@ -0,0 +1,467 @@
+import time
+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.tensorboard import SummaryWriter
+import numpy as np
+from torch.utils.data.distributed import DistributedSampler
+import logging
+from data import gigaspeech
+from models import voicecraft
+
+from .trainer_utils import DistributedDynamicBatchSampler, StatefulDistributedSampler, AverageMeter, print_model_info
+from .optim import ScaledAdam, Eden
+
+
+class Trainer:
+    
+    def __init__(self, args, world_size, rank):
+        self.start_time = time.time()
+        self.args = args
+        self.world_size, self.rank = world_size, rank
+        self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+        if self.rank == 0:
+            self.writer = SummaryWriter(args.exp_dir)
+        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._setup_models()
+
+        self.train_dataset_length, self.train_sampler, self.train_loader, self.valid_loader = self._setup_dataloader()
+        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.rank], find_unused_parameters=False)
+        
+        if self.rank == 0:
+            self.early_stop_accu_steps = 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 (summed over all GPUs): {self.args.batch_size}")
+
+    def train(self):
+        flag = True
+        skip_flag = False
+        data_start_time = time.time()
+        while flag:
+            self.train_sampler.set_epoch(self.progress['epoch'])
+            for i, batch in enumerate(self.train_loader):
+                data_end_time = time.time()
+                self.model.train()
+                if self.progress['step'] > self.total_step:
+                    flag = False
+                    self.validate_and_save()
+                    if self.rank == 0:
+                        self.writer.close()
+                    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)]
+                for j in range(self.args.gradient_accumulation_steps):
+                    cur_ind = all_inds[j::self.args.gradient_accumulation_steps]
+                    cur_batch = {key: batch[key][cur_ind] for key in batch}
+                    with torch.cuda.amp.autocast(dtype=torch.float16 if self.args.precision=="float16" else torch.float32):
+                        out = self.model(cur_batch)
+
+                    record_loss = out['loss'].detach().to(self.rank) 
+                    top10acc = out['top10acc'].to(self.rank)
+                    effective_ntoken = out['effective_ntoken'].to(self.rank)
+                    is_nan = torch.tensor(int(torch.isnan(record_loss).any()), dtype=torch.float32, device=self.rank)
+                    
+                    dist.all_reduce(record_loss, op=dist.ReduceOp.SUM)
+                    dist.all_reduce(top10acc, op=dist.ReduceOp.SUM)
+                    dist.all_reduce(effective_ntoken, op=dist.ReduceOp.SUM)
+                    dist.all_reduce(is_nan, op=dist.ReduceOp.SUM)
+                    
+                    # check if loss is nan
+                    if is_nan.item() > 0:
+                        logging.info(f"loss at step {self.progress['step']} is nan, therefore skip this batch")
+                        skip_flag = True
+                        continue
+
+                    sum_losses += record_loss.item()
+                    sum_top10acc += top10acc.item()
+                    sum_ntoken += effective_ntoken.item()
+
+                    if 'top10acc_by_codebook' in out:
+                        for cb in range(self.args.n_codebooks):
+                            top10acc_cbi = out['top10acc_by_codebook'][cb]
+                            dist.all_reduce(top10acc_cbi, op=dist.ReduceOp.SUM)
+                            sum_top10acc_cbi[cb] += top10acc_cbi.item()
+                        
+                    if self.rank == 0:
+                        average_loss = sum_losses / sum_ntoken
+                        average_top10acc = sum_top10acc / sum_ntoken
+                        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)
+                        average_top10acc_cbi = [sum_top10acc_cbi[cb] / sum_ntoken * self.args.n_codebooks for cb in range(self.args.n_codebooks)]
+                        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)
+
+                        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'])
+                            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'])
+
+                    if self.args.optimizer_name == "ScaledAdam":
+                        self.scaler.scale(out['loss']).backward() 
+                    else:
+                        self.scaler.scale(out['loss']/out['effective_ntoken']).backward()
+
+                if skip_flag:
+                    self.optimizer.zero_grad()
+                    skip_flag = False
+                    continue
+
+                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()
+
+                if self.rank == 0:
+                    self.meters['data_time'].update(data_end_time - data_start_time)
+                    self.meters['train_time'].update(time.time() - data_end_time)
+                    if self.progress['step'] % self.args.tb_write_every_n_steps == 0:
+                        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'])
+                        
+
+                    # 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...")
+
+                # validation and save models
+                if self.progress['step'] % self.args.val_every_n_steps == 0:
+                    dist.barrier()
+                    self.validate_and_save()
+
+                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.rank == 0:
+            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}")
+                        dist.destroy_process_group()
+                        raise RuntimeError("early stop")
+                else:
+                    self.early_stop_accu_steps = 0
+
+            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")
+                torch.save(
+                    {
+                        "model": self.model.module.state_dict(),
+                        "optimizer":  self.optimizer.state_dict(),
+                        "scheduler": self.scheduler.state_dict(),
+                        "config": 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']}")
+            self._save_progress()
+            save_path = os.path.join(self.args.exp_dir,"bundle.pth")
+            torch.save(
+                {
+                    "model": self.model.module.state_dict(),
+                    "optimizer":  self.optimizer.state_dict(),
+                    "scheduler": self.scheduler.state_dict(),
+                    "config": self.args,
+                    "phn2num": self.train_loader.dataset.phn2num
+                    },save_path
+            )
+            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']}")
+
+        dist.barrier()
+
+    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_top10acc_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)
+                sum_losses += out['loss']
+                sum_top10acc += out['top10acc']
+                sum_ntoken += out['effective_ntoken']
+                if 'top10acc_by_codebook' in out:
+                    for cb in range(self.args.n_codebooks):
+                        sum_top10acc_cbi[cb] += out['top10acc_by_codebook'][cb]
+                        
+        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 '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 self.rank == 0:
+            val_loss = sum_losses / sum_ntoken
+            val_top10acc = sum_top10acc / sum_ntoken
+            # 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.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'])
+            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'])
+
+            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()
+        else:
+            return None
+
+    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 += [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)]
+        for name in meter_names:
+            meters[name] = AverageMeter()
+        return meters
+    def _setup_progress(self):
+        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):
+        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])
+        with open("%s/progress.pkl" % self.args.exp_dir, "wb") as f:
+            pickle.dump(self.total_progress, f)
+
+    def _setup_dataloader(self):
+        assert self.args.dataset == 'gigaspeech', "only gigaspeech is supported for now"
+        train_dataset, val_dataset = gigaspeech.dataset(self.args, 'train'), gigaspeech.dataset(self.args, 'validation')
+        
+        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'])
+
+        if self.args.dynamic_batching:
+            train_loader = torch.utils.data.DataLoader(train_dataset, 
+                            batch_sampler=train_sampler, 
+                            num_workers=self.args.num_workers//self.world_size,
+                            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//self.world_size,
+                            collate_fn=val_dataset.collate, persistent_workers=True
+                            )
+        else:
+            train_loader = torch.utils.data.DataLoader(train_dataset, 
+                            batch_size=self.args.batch_size//self.world_size, sampler=train_sampler, num_workers=self.args.num_workers//self.world_size,
+                            collate_fn=train_dataset.collate, persistent_workers=True
+                            )
+            valid_loader = torch.utils.data.DataLoader(val_dataset, 
+                            batch_size=self.args.batch_size//self.world_size, sampler=valid_sampler,
+                            num_workers=self.args.num_workers//self.world_size,
+                            collate_fn=val_dataset.collate, persistent_workers=True
+                            )
+        return len(train_dataset), train_sampler, train_loader, valid_loader
+        
+
+        
+    def _setup_models(self):
+        model = voicecraft.VoiceCraft(self.args)
+
+        if self.rank == 0:
+            logging.info(model)
+            logging.info("model parameters")
+            print_model_info(model)
+
+        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']
+            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']
+        else:
+            optim_states = None
+            scheduler_states = None
+
+        if self.args.load_model_from != None and self.progress['step'] <= 1:
+            sd = torch.load(self.args.load_model_from, map_location="cpu")['model']
+            model.load_state_dict(sd)
+            del sd
+        
+        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
+        model.to(self.device)
+
+        return model, trainables, optim_states, scheduler_states
+
+    
+    def _setup_optimizer(self):
+        if self.args.optimizer_name == "ScaledAdam":
+            parameters_names = []
+            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)
+
+        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

steps/trainer_utils.py

@@ -0,0 +1,628 @@
+
+import torch
+import math
+import torch.distributed as dist
+from torch.utils.data.sampler import Sampler
+import copy
+import numpy as np
+from typing import List
+from scipy.stats import lognorm
+import logging
+
+class StatefulDistributedSampler(Sampler[int]):
+    def __init__(self, dataset, batch_size, num_replicas = None, rank = None, shuffle = True, seed = 0, drop_last = False):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.cur_epoch = 0
+        self.drop_last = drop_last
+        # If the dataset length is evenly divisible by # of replicas, then there
+        # is no need to drop any data, since the dataset will be split equally.
+        if self.drop_last and len(self.dataset) % self.num_replicas != 0:  # type: ignore[arg-type]
+            # Split to nearest available length that is evenly divisible.
+            # This is to ensure each rank receives the same amount of data when
+            # using this Sampler.
+            self.num_samples = math.ceil(
+                (len(self.dataset) - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil(len(self.dataset) / self.num_replicas)  # type: ignore[arg-type]
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.continue_flag = False
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        r"""
+        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
+        use a different random ordering for each epoch. Otherwise, the next iteration of this
+        sampler will yield the same ordering.
+
+        Args:
+            epoch (int): Epoch number.
+        """
+        self.epoch = epoch
+
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[:self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+        self.indices = indices
+
+        if self.continue_flag:
+            self.indices = self.indices[int(self.cur_step*self.batch_size):]
+            self.num_samples = len(self.indices)
+            self.continue_flag = False
+                
+    def __iter__(self):
+        for idx in self.indices:
+            yield idx  
+
+    def set_epoch_resume(self, epoch, cur_step):
+        self.epoch = epoch
+        self.cur_step = cur_step
+        self.continue_flag = True
+
+
+class StatefulSampler(Sampler):
+    def __init__(self, data_source_length, batch_size, use_random=True, seed=1, epoch=0):
+        self.use_random = use_random
+        self.data_source_length = data_source_length
+        self.num_samples = self.data_source_length
+        self.batch_size = batch_size
+        self.continue_flag = False
+        self.seed = seed
+        self.epoch = epoch
+        self.cur_step = 0
+
+    def __len__(self):
+        return self.num_samples
+
+    def __iter__(self):
+
+        for idx in self.indices:
+            yield idx
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+        if self.use_random:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            self.indices = torch.randperm(self.data_source_length, generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            self.indices = list(range(self.data_source_length))  # type: ignore[arg-type]
+        if self.continue_flag == True:
+            self.continue_flag = False
+            self.indices = self.indices[int(self.cur_step*self.batch_size):]
+        
+        self.num_samples = len(self.indices)
+    
+    def set_epoch_resume(self, epoch, cur_step):
+        self.epoch = epoch
+        self.cur_step = cur_step
+        self.continue_flag = True
+
+
+class AverageMeter:
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+def print_model_info(model, print_model = False, print_params = True):
+    if print_model:
+        logging.info(model)
+    if print_params:
+        all_params = {}
+        for name, p in model.named_parameters():
+            name = name.split(".")[0]
+            if name in all_params:
+                all_params[name] += p.numel()
+            else:
+                all_params[name] = p.numel()
+        logging.info("num of parameters of each components:")
+        for name in all_params:
+            logging.info(f"{name}: {all_params[name]/1000000.:.2f}m")
+
+
+class DistributedDynamicBatchSampler(Sampler):
+    """
+    modified from SpeechBrian, https://github.com/speechbrain/speechbrain/blob/develop/speechbrain/dataio/sampler.py#L307
+    This BatchSampler batches examples together by grouping them by their length.
+
+    Every example in the batch have approximately the same length and
+    thus padding is minimized.
+    This enables faster training on datasets
+    where length of examples can vary significantly (e.g Librispeech).
+    Inspired by: https://www.tensorflow.org/api_docs/python/tf/data/experimental/bucket_by_sequence_length
+
+    Dynamic batching is performed by specifying a max_batch_length which is the
+    upper limit for the sum of the length of examples in a batch:
+    e.g., if ex1 has length 4, ex2 length 5 and if max_batch_length is set to 6
+    ex1 and ex2 will be placed, alone, in two distinct batches.
+
+    Length for each example can be obtained in two manners.
+    If the input dataset is a DynamicItemDataset it can be obtained by specifying a
+    length_func. Default assumes a "duration" entry is in the annotation.
+    Length for each example can also be passed to this class upon instantiation
+    by specifying a list containing the length for each example and passing it to
+    lengths_list.
+
+    Examples are grouped together by defining a set of possible discrete intervals
+    (buckets). Examples whose length fall into these intervals can be batched together.
+
+    The number of buckets can be specified by using the arg num_buckets.
+    There is usually an optimal range for the value of this argument.
+
+    If num_buckets == 1, all examples can be batched together. You have maximum randomization
+    but your training speed will be slower due to the fact that a large amount of the values will be padding
+    as long and short examples can be batched together.
+    As the number of buckets grows only examples with similar
+    length can be grouped together.
+    This trades-off speed with randomization.
+    TLDR: Low number -> better randomization, High number -> faster training.
+    NOTE THAT: if set too high the training speed will decrease. If num_buckets -> number of examples in the dataset the batch size
+    will be small impacting training speed and possibly performance.
+
+    The buckets can also be specified by passing a list to the bucket_boundaries
+    argument instead of specifying a left_bucket_length and a bucket_length_multiplier.
+
+    Example
+    -------
+    >>> import torch
+    >>> import speechbrain as sb
+    >>> from speechbrain.dataio.sampler import DynamicBatchSampler
+    >>> from speechbrain.dataio.dataset import DynamicItemDataset
+    >>> from speechbrain.dataio.dataloader import SaveableDataLoader
+    >>> from speechbrain.dataio.batch import PaddedBatch
+    >>> import numpy as np
+    >>> item_lengths = sorted([np.random.randint(10, 100) for x in range(20)])
+    >>> dataset = {"ex_{}".format(x) : {"wav" :torch.randn(x)} for x in item_lengths}
+    >>> dataset = DynamicItemDataset(dataset)
+    >>> dataset.set_output_keys(["wav"])
+    >>> length_func = lambda x : len(x) # trivial in this example
+    >>> bsampler = DynamicBatchSampler(dataset, 20, 4, length_func, shuffle=False, batch_ordering='descending')
+    >>> dataloader = SaveableDataLoader(dataset, batch_sampler=bsampler, collate_fn=PaddedBatch)
+    >>> for i, b in enumerate(dataloader):
+    ...     data, length = b["wav"]
+    >>> assert data.shape[-1] == max(item_lengths)
+
+    Arguments
+    ---------
+    dataset : torch.utils.data.Dataset
+        Pytorch Dataset from which elements will be sampled.
+    max_batch_length : int
+        Upper limit for the sum of the length of examples in a batch.
+        Should be chosen based on your GPU memory.
+    num_buckets : int
+        Number of discrete buckets used to group examples together.
+        If num_buckets == 1, all examples can be batched together. As the number of buckets grows only examples with similar
+        length can be grouped together. This trades-off speed with randomization.
+        Low number -> better randomization, High number -> faster training.
+        However if set too high the training speed will decrease. If num_buckets -> number of examples in the dataset the batch size
+        will be small impacting training speed and possibly performance.
+        NOTE: you have either to specify manually the bucket_boundaries or the number of buckets.
+    length_func : callable
+        Function used to get length of each example from the dataset.
+        This argument can be used only when the dataset is a Speechbrain DynamicItemDataset object.
+        Can be anything: e.g. lambda x: x["duration"]*16000 returns number of samples
+        if duration key in the annotation is in seconds and the file has 16kHz sampling freq.
+    shuffle : bool
+        Whether or not shuffle examples between each epoch.
+    batch_ordering : string
+        If ``random``, batches are randomly permuted; otherwise ``ascending`` or ``descending`` sorted by length.
+    max_batch_ex: int
+        If set, it limits the maximum number of examples that can be in a batch superseeding max_batch_length
+        in instances where the amount of examples will exceeed the value specified here.
+        E.g. you have a lot of short examples and the batch size for those will be too high, you can use this argument
+        to limit the batch size for these short examples.
+    bucket_boundaries : list
+        Overrides bucket_length_multiplier and left_bucket_length by specifying manually
+        the buckets right boundaries.
+    lengths_list: list
+        Overrides length_func by passing a list containing the length of each example
+        in the dataset. This argument must be set when the dataset is a plain
+        Pytorch Dataset object and not a DynamicItemDataset object as length_func
+        cannot be used on Pytorch Datasets.
+    epoch : int
+        The epoch to start at.
+    drop_last : bool
+         If ``True``, the sampler will drop the last examples which
+         have not been grouped.
+    verbose: bool
+        If ``True``, log also the stats for each batch at the first epoch.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        args,
+        num_replicas = None, 
+        rank = None, 
+        shuffle = True, 
+        seed = 0, 
+        drop_last = False,
+        length_func=lambda x: x["duration"],
+        batch_ordering: str = "random",
+        max_batch_ex: int = None,
+        bucket_boundaries: List[int] = [],
+        lengths_list: List[int] = None,
+        epoch: int = 0,
+        verbose: bool = False,
+    ):
+        self.args = args
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1))
+        self.num_replicas = num_replicas
+        self.rank = rank
+        max_batch_length = self.args.max_num_tokens if dataset.split == "train" else self.args.val_max_num_tokens
+        logging.info(f"max_num_tokens per GPU for {dataset.split} split: {max_batch_length}")
+        num_buckets = self.args.num_buckets
+        #############
+        
+
+
+
+        self._dataset = dataset
+        self._ex_lengths = {}
+        # ex_ids = self._dataset.data_ids
+        self.verbose = verbose
+
+        # We do not put a default on num_buckets to encourage users to play with this parameter
+        if num_buckets is None and len(bucket_boundaries) == 0:
+            raise RuntimeError(
+                "Please specify either num_buckets or bucket boundaries."
+                "Check the docs, and/or the tutorial !"
+            )
+        assert lengths_list != None
+        max_len = int(self.args.audio_max_length * self.args.encodec_sr)
+        lengths_list = [min(l, max_len) for l in lengths_list] # replace all utt whose length is longer than max_len to max_len, will also do this in __getitem__ in dataset
+        for indx in range(len(lengths_list)):
+            self._ex_lengths[str(indx)] = lengths_list[indx]
+        # if lengths_list is not None:
+        #     # take length of examples from this argument and bypass length_key
+        #     for indx in range(len(lengths_list)):
+        #         self._ex_lengths[str(indx)] = lengths_list[indx]
+        # else:
+        #     # use length func
+        #     if not isinstance(dataset, DynamicItemDataset):
+        #         raise NotImplementedError(
+        #             "Dataset should be a Speechbrain DynamicItemDataset when using length function"
+        #         )
+        #     for indx in range(len(self._dataset)):
+        #         self._ex_lengths[str(indx)] = length_func(
+        #             self._dataset.data[ex_ids[indx]]
+        #         )
+
+        if len(bucket_boundaries) > 0:
+            if not all([x >= 0 for x in bucket_boundaries]):
+                raise ValueError(
+                    "All elements in bucket boundaries should be non-negative (>= 0)."
+                )
+            if not len(set(bucket_boundaries)) == len(bucket_boundaries):
+                raise ValueError(
+                    "Bucket_boundaries should not contain duplicates."
+                )
+            np.testing.assert_array_equal(
+                np.array(bucket_boundaries),
+                np.array(sorted(bucket_boundaries)),
+                err_msg="The arg bucket_boundaries should be an ascending sorted list of non negative values values!",
+            )
+            self._bucket_boundaries = np.array(sorted(bucket_boundaries))
+        else:
+            # use num_buckets
+            self._bucket_boundaries = np.array(
+                self._get_boundaries_through_warping(
+                    # max_batch_length=max_batch_length,
+                    max_batch_length=max(lengths_list),
+                    num_quantiles=num_buckets,
+                )
+            )
+
+        self._max_batch_length = max_batch_length
+        self._shuffle_ex = shuffle
+        self._batch_ordering = batch_ordering
+        self._seed = seed
+        self._drop_last = drop_last
+        if max_batch_ex is None:
+            max_batch_ex = np.inf
+        self._max_batch_ex = max_batch_ex
+        # Calculate bucket lengths - how often does one bucket boundary fit into max_batch_length?
+        self._bucket_lens = [
+            max(1, int(max_batch_length / self._bucket_boundaries[i]))
+            for i in range(len(self._bucket_boundaries))
+        ] + [1]
+        self._epoch = epoch
+        self._cur_step = 0
+        self.continue_flag = False
+        self._generate_batches()
+        self.num_samples = int(math.floor(len(self._batches) / self.num_replicas))
+        self.total_size = int(self.num_samples * self.num_replicas)
+        self._replica_batches = self._batches[self.rank:self.total_size:self.num_replicas]
+        assert len(self._replica_batches) == self.num_samples, f"len(self._batches): {len(self._batches)}, self.total_size: {self.total_size}, self.num_samples: {self.num_samples},len(self._replica_batches): {len(self._replica_batches)}"
+        logging.info(f"len(self._batches): {len(self._batches)}")
+        logging.info(f"self.num_replicas: {self.num_replicas}")
+        logging.info(f"num of batches on each replica: {self.num_samples}")
+
+    def get_durations(self, batch):
+        """Gets durations of the elements in the batch."""
+        return [self._ex_lengths[str(idx)] for idx in batch]
+
+    def _get_boundaries_through_warping(
+        self, max_batch_length: int, num_quantiles: int,
+    ) -> List[int]:
+
+        # NOTE: the following lines do not cover that there is only one example in the dataset
+        # warp frames (duration) distribution of train data
+        logging.info("Batch quantisation in latent space")
+        # linspace set-up
+        num_boundaries = num_quantiles + 1
+        # create latent linearly equal spaced buckets
+        latent_boundaries = np.linspace(
+            1 / num_boundaries, num_quantiles / num_boundaries, num_quantiles,
+        )
+        # get quantiles using lognormal distribution
+        quantiles = lognorm.ppf(latent_boundaries, 1)
+        # scale up to to max_batch_length
+        bucket_boundaries = quantiles * max_batch_length / quantiles[-1]
+        # compute resulting bucket length multipliers
+        length_multipliers = [
+            bucket_boundaries[x + 1] / bucket_boundaries[x]
+            for x in range(num_quantiles - 1)
+        ]
+        # logging
+        logging.debug(
+            "Latent bucket boundary - buckets: {} - length multipliers: {}".format(
+                list(map("{:.2f}".format, bucket_boundaries)),
+                list(map("{:.2f}".format, length_multipliers)),
+            )
+        )
+        return list(sorted(bucket_boundaries))
+
+    def _permute_batches(self):
+
+        if self._batch_ordering == "random":
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self._seed + self._epoch) # since the random seed is based on self._seed and self._epoch, it should be the same for different processes when using DDP, and therefore the generated order should be the same across different process, this is important, because each replica will only take a portion of it, we want to make sure they take a non-overlapping portion, and all of them constitute the entire dataset
+            sampler = torch.randperm(
+                len(self._batches), generator=g
+            ).tolist()  # type: ignore
+            tmp = []
+            for idx in sampler:
+                tmp.append(self._batches[idx])
+            self._batches = tmp
+
+        elif self._batch_ordering == "ascending":
+            self._batches = sorted(
+                self._batches,
+                key=lambda x: max([self._ex_lengths[str(idx)] for idx in x]),
+            )
+        elif self._batch_ordering == "descending":
+            self._batches = sorted(
+                self._batches,
+                key=lambda x: max([self._ex_lengths[str(idx)] for idx in x]),
+                reverse=True,
+            )
+        else:
+            raise NotImplementedError
+
+    def _generate_batches(self):
+        logging.info("DynamicBatchSampler: Generating dynamic batches")
+        if self._shuffle_ex:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self._seed + self._epoch) # since the random seed is based on self._seed and self._epoch, it should be the same for different processes when using DDP, and therefore the generated order should be the same across different process, this is important, because each replica will only take a portion of it, we want to make sure they take a non-overlapping portion, and all of them constitute the entire dataset
+            sampler = torch.randperm(len(self._dataset), generator=g).tolist()  # type: ignore 
+            # pyp note: this is actually randomly permoted indices
+        else:
+            # take examples as they are: e.g. they have been sorted
+            sampler = range(len(self._dataset))  # type: ignore
+
+        self._batches = []
+        bucket_batches = [[] for i in self._bucket_lens]
+
+        stats_tracker = [
+            {"min": np.inf, "max": -np.inf, "tot": 0, "n_ex": 0}
+            for i in self._bucket_lens
+        ]
+
+        for idx in sampler:
+            # length of pre-sampled audio
+            item_len = self._ex_lengths[str(idx)]
+            # bucket to fill up most padding
+            bucket_id = np.searchsorted(self._bucket_boundaries, item_len)
+            # fill audio's duration into that bucket
+            bucket_batches[bucket_id].append(idx)
+
+            stats_tracker[bucket_id]["min"] = min(
+                stats_tracker[bucket_id]["min"], item_len
+            )
+            stats_tracker[bucket_id]["max"] = max(
+                stats_tracker[bucket_id]["max"], item_len
+            )
+            stats_tracker[bucket_id]["tot"] += item_len
+            stats_tracker[bucket_id]["n_ex"] += 1
+            # track #samples - why not duration/#frames; rounded up?
+            # keep track of durations, if necessary
+
+            if (
+                len(bucket_batches[bucket_id]) >= self._bucket_lens[bucket_id]
+                or len(bucket_batches[bucket_id]) >= self._max_batch_ex
+            ):
+                self._batches.append(bucket_batches[bucket_id])
+                bucket_batches[bucket_id] = []
+                # keep track of durations
+
+        # Dump remaining batches
+        if not self._drop_last:
+            for batch in bucket_batches:
+                if batch:
+                    self._batches.append(batch)
+
+        self._permute_batches()  # possibly reorder batches
+
+        if self._epoch == 0:  # only log at first epoch
+            # frames per batch & their padding remaining
+            boundaries = [0] + self._bucket_boundaries.tolist()
+
+            for bucket_indx in range(len(self._bucket_boundaries)):
+                try:
+                    num_batches = stats_tracker[bucket_indx]["tot"] // (
+                        self._max_batch_length
+                    )
+                    pad_factor = (
+                        stats_tracker[bucket_indx]["max"]
+                        - stats_tracker[bucket_indx]["min"]
+                    ) / (
+                        stats_tracker[bucket_indx]["tot"]
+                        / stats_tracker[bucket_indx]["n_ex"]
+                    )
+                except ZeroDivisionError:
+                    num_batches = 0
+                    pad_factor = 0
+
+                logging.debug(
+                    (
+                        "DynamicBatchSampler: Bucket {} with boundary {:.1f}-{:.1f} and "
+                        + "batch_size {}: Num Examples {:.1f}, Num Full Batches {:.3f}, Pad Factor {:.3f}."
+                    ).format(
+                        bucket_indx,
+                        boundaries[bucket_indx],
+                        boundaries[bucket_indx + 1],
+                        self._bucket_lens[bucket_indx],
+                        stats_tracker[bucket_indx]["n_ex"],
+                        num_batches,
+                        pad_factor * 100,
+                    )
+                )
+
+            if self.verbose:
+                batch_stats = {
+                    "tot_frames": [],
+                    "tot_pad_frames": [],
+                    "pad_%": [],
+                }
+                for batch in self._batches:
+                    tot_frames = sum(
+                        [self._ex_lengths[str(idx)] for idx in batch]
+                    )
+                    batch_stats["tot_frames"].append(tot_frames)
+                    max_frames = max(
+                        [self._ex_lengths[str(idx)] for idx in batch]
+                    )
+                    tot_pad = sum(
+                        [
+                            max_frames - self._ex_lengths[str(idx)]
+                            for idx in batch
+                        ]
+                    )
+                    batch_stats["tot_pad_frames"].append(tot_pad)
+                    batch_stats["pad_%"].append(tot_pad / tot_frames * 100)
+
+                padding_details = "Batch {} with {:.1f} frames with {} files - {:.1f} padding, {:.2f} (%) of total."
+                padding_details = "DynamicBatchSampler: " + padding_details
+                for i in range(len(self._batches)):
+                    logging.debug(
+                        padding_details.format(
+                            i,
+                            batch_stats["tot_frames"][i],
+                            len(self._batches[i]),
+                            batch_stats["tot_pad_frames"][i],
+                            batch_stats["pad_%"][i],
+                        )
+                    )
+
+    def __iter__(self):
+
+        for batch in self._replica_batches:
+            yield batch
+
+
+        # if self._shuffle_ex:  # re-generate examples if ex_ordering == "random"
+        #     self._generate_batches()
+        # if self._batch_ordering == "random":
+        #     # we randomly permute the batches only --> faster
+        #     self._permute_batches()
+
+    def set_epoch(self, epoch):
+        """
+        You can also just access self.epoch, but we maintain this interface
+        to mirror torch.utils.data.distributed.DistributedSampler
+        """
+        self._epoch = epoch
+        self._generate_batches()
+        self._replica_batches = self._batches[self.rank:self.total_size:self.num_replicas]
+        self.num_samples = int(math.floor(len(self._batches) / self.num_replicas))
+        assert len(self._replica_batches) == self.num_samples, f"len(self._batches): {len(self._batches)}, self.total_size: {self.total_size}, self.num_samples: {self.num_samples},len(self._replica_batches): {len(self._replica_batches)}"
+
+        if self.continue_flag:
+            self.continue_flag = False
+            self._replica_batches = self._replica_batches[self._cur_step:]
+            self.num_samples = len(self._replica_batches)
+ 
+
+    def __len__(self):
+        return self.num_samples
+    
+    def set_epoch_resume(self, epoch, cur_step):
+        self.continue_flag = True
+        self._epoch = epoch
+        self._cur_step = cur_step