update train

speech/cosyvoice/utils/executor.py

@@ -13,42 +13,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
-from contextlib import nullcontext
 import os
+from contextlib import nullcontext
 
 import torch
 import torch.distributed as dist
+from cosyvoice.utils.train_utils import (batch_backward, batch_forward,
+                                         cosyvoice_join, log_per_save,
+                                         log_per_step, save_model,
+                                         update_parameter_and_lr)
 
-from cosyvoice.utils.train_utils import update_parameter_and_lr, log_per_step, log_per_save, batch_forward, batch_backward, save_model, cosyvoice_join
+from loguru import logger
 
 
 class Executor:
-
-    def __init__(self, gan: bool = False, ref_model: torch.nn.Module = None, dpo_loss: torch.nn.Module = None):
+    """Executor for training and cross validation"""
+    def __init__(
+        self,
+        gan: bool = False,
+        ref_model: torch.nn.Module = None,
+        dpo_loss: torch.nn.Module = None,
+    ):
         self.gan = gan
         self.ref_model = ref_model
         self.dpo_loss = dpo_loss
         self.step = 0
         self.epoch = 0
-        self.rank = int(os.environ.get('RANK', 0))
-        self.device = torch.device('cuda:{}'.format(self.rank))
-
-    def train_one_epoc(self, model, optimizer, scheduler, train_data_loader, cv_data_loader, writer, info_dict, scaler, group_join, ref_model=None):
-        ''' Train one epoch
-        '''
-
-        lr = optimizer.param_groups[0]['lr']
-        logging.info('Epoch {} TRAIN info lr {} rank {}'.format(self.epoch, lr, self.rank))
-        logging.info('using accumulate grad, new batch size is {} times'
-                     ' larger than before'.format(info_dict['accum_grad']))
-        # A context manager to be used in conjunction with an instance of
-        # torch.nn.parallel.DistributedDataParallel to be able to train
-        # with uneven inputs across participating processes.
+        self.rank = int(os.environ.get("RANK", 0))
+        self.device = torch.device(f"cuda:{self.rank}")
+
+    def train_one_epoc(
+        self,
+        model,
+        optimizer,
+        scheduler,
+        train_data_loader,
+        cv_data_loader,
+        writer,
+        info_dict,
+        scaler,
+        group_join,
+    ):
+        """Train one epoch"""
+
+        lr = optimizer.param_groups[0]["lr"]
+        logger.info(
+            f"Epoch {self.epoch} TRAIN info lr {lr} rank {self.rank}"
+        )
+        logger.info(
+            f"using accumulate grad, new batch size is {info_dict['accum_grad']} times larger than before"
+        )
+       
         model.train()
         if self.ref_model is not None:
             self.ref_model.eval()
-        model_context = model.join if info_dict['train_engine'] == 'torch_ddp' else nullcontext
+        model_context = (
+            model.join if info_dict["train_engine"] == "torch_ddp" else nullcontext
+        )
         with model_context():
             for batch_idx, batch_dict in enumerate(train_data_loader):
                 info_dict["tag"] = "TRAIN"
@@ -58,47 +79,77 @@ class Executor:
                 if cosyvoice_join(group_join, info_dict):
                     break
 
-                # Disable gradient synchronizations across DDP processes.
-                # Within this context, gradients will be accumulated on module
-                # variables, which will later be synchronized.
-                if info_dict['train_engine'] == 'torch_ddp' and (batch_idx + 1) % info_dict["accum_grad"] != 0:
+               
+                if (
+                    info_dict["train_engine"] == "torch_ddp"
+                    and (batch_idx + 1) % info_dict["accum_grad"] != 0
+                ):
                     context = model.no_sync
-                # Used for single gpu training and DDP gradient synchronization
-                # processes.
+            
                 else:
                     context = nullcontext
 
                 with context():
-                    info_dict = batch_forward(model, batch_dict, scaler, info_dict, ref_model=self.ref_model, dpo_loss=self.dpo_loss)
+                    info_dict = batch_forward(
+                        model,
+                        batch_dict,
+                        scaler,
+                        info_dict,
+                        ref_model=self.ref_model,
+                        dpo_loss=self.dpo_loss,
+                    )
                     info_dict = batch_backward(model, scaler, info_dict)
 
-                info_dict = update_parameter_and_lr(model, optimizer, scheduler, scaler, info_dict)
+                info_dict = update_parameter_and_lr(
+                    model, optimizer, scheduler, scaler, info_dict
+                )
                 log_per_step(writer, info_dict)
                 # NOTE specify save_per_step in cosyvoice.yaml if you want to enable step save
-                if info_dict['save_per_step'] > 0 and (self.step + 1) % info_dict['save_per_step'] == 0 and \
-                   (batch_idx + 1) % info_dict["accum_grad"] == 0:
+                if (
+                    info_dict["save_per_step"] > 0
+                    and (self.step + 1) % info_dict["save_per_step"] == 0
+                    and (batch_idx + 1) % info_dict["accum_grad"] == 0
+                ):
                     dist.barrier()
-                    self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=False)
+                    self.cv(
+                        model, cv_data_loader, writer, info_dict, on_batch_end=False
+                    )
                     model.train()
                 if (batch_idx + 1) % info_dict["accum_grad"] == 0:
                     self.step += 1
         dist.barrier()
         self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=True)
 
-    def train_one_epoc_gan(self, model, optimizer, scheduler, optimizer_d, scheduler_d, train_data_loader, cv_data_loader,
-                           writer, info_dict, scaler, group_join):
-        ''' Train one epoch
-        '''
-
-        lr = optimizer.param_groups[0]['lr']
-        logging.info('Epoch {} TRAIN info lr {} rank {}'.format(self.epoch, lr, self.rank))
-        logging.info('using accumulate grad, new batch size is {} times'
-                     ' larger than before'.format(info_dict['accum_grad']))
+    def train_one_epoc_gan(
+        self,
+        model,
+        optimizer,
+        scheduler,
+        optimizer_d,
+        scheduler_d,
+        train_data_loader,
+        cv_data_loader,
+        writer,
+        info_dict,
+        scaler,
+        group_join,
+    ):
+        """Train one epoch"""
+
+        lr = optimizer.param_groups[0]["lr"]
+        logger.info(
+            f"Epoch {self.epoch} TRAIN info lr {lr} rank {self.rank}"
+        )
+        logger.info(
+            f"using accumulate grad, new batch size is {info_dict['accum_grad']} times larger than before"
+        )
         # A context manager to be used in conjunction with an instance of
         # torch.nn.parallel.DistributedDataParallel to be able to train
         # with uneven inputs across participating processes.
         model.train()
-        model_context = model.join if info_dict['train_engine'] == 'torch_ddp' else nullcontext
+        model_context = (
+            model.join if info_dict["train_engine"] == "torch_ddp" else nullcontext
+        )
         with model_context():
             for batch_idx, batch_dict in enumerate(train_data_loader):
                 info_dict["tag"] = "TRAIN"
@@ -111,7 +162,10 @@ class Executor:
                 # Disable gradient synchronizations across DDP processes.
                 # Within this context, gradients will be accumulated on module
                 # variables, which will later be synchronized.
-                if info_dict['train_engine'] == 'torch_ddp' and (batch_idx + 1) % info_dict["accum_grad"] != 0:
+                if (
+                    info_dict["train_engine"] == "torch_ddp"
+                    and (batch_idx + 1) % info_dict["accum_grad"] != 0
+                ):
                     context = model.no_sync
                 # Used for single gpu training and DDP gradient synchronization
                 # processes.
@@ -119,35 +173,43 @@ class Executor:
                     context = nullcontext
 
                 with context():
-                    batch_dict['turn'] = 'discriminator'
+                    batch_dict["turn"] = "discriminator"
                     info_dict = batch_forward(model, batch_dict, scaler, info_dict)
                     info_dict = batch_backward(model, scaler, info_dict)
-                info_dict = update_parameter_and_lr(model, optimizer_d, scheduler_d, scaler, info_dict)
+                info_dict = update_parameter_and_lr(
+                    model, optimizer_d, scheduler_d, scaler, info_dict
+                )
                 optimizer.zero_grad()
                 log_per_step(writer, info_dict)
                 with context():
-                    batch_dict['turn'] = 'generator'
+                    batch_dict["turn"] = "generator"
                     info_dict = batch_forward(model, batch_dict, scaler, info_dict)
                     info_dict = batch_backward(model, scaler, info_dict)
-                info_dict = update_parameter_and_lr(model, optimizer, scheduler, scaler, info_dict)
+                info_dict = update_parameter_and_lr(
+                    model, optimizer, scheduler, scaler, info_dict
+                )
                 optimizer_d.zero_grad()
                 log_per_step(writer, info_dict)
                 # NOTE specify save_per_step in cosyvoice.yaml if you want to enable step save
-                if info_dict['save_per_step'] > 0 and (self.step + 1) % info_dict['save_per_step'] == 0 and \
-                   (batch_idx + 1) % info_dict["accum_grad"] == 0:
+                if (
+                    info_dict["save_per_step"] > 0
+                    and (self.step + 1) % info_dict["save_per_step"] == 0
+                    and (batch_idx + 1) % info_dict["accum_grad"] == 0
+                ):
                     dist.barrier()
-                    self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=False)
+                    self.cv(
+                        model, cv_data_loader, writer, info_dict, on_batch_end=False
+                    )
                     model.train()
                 if (batch_idx + 1) % info_dict["accum_grad"] == 0:
                     self.step += 1
         dist.barrier()
-        self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=True)
+        # self.cv(model, cv_data_loader, writer, info_dict, on_batch_end=True)
 
     @torch.inference_mode()
     def cv(self, model, cv_data_loader, writer, info_dict, on_batch_end=True):
-        ''' Cross validation on
-        '''
-        logging.info('Epoch {} Step {} on_batch_end {} CV rank {}'.format(self.epoch, self.step + 1, on_batch_end, self.rank))
+        """Cross validation on"""
+        logger.info(f"Epoch {self.epoch} Step {self.step + 1} on_batch_end {on_batch_end} CV rank {self.rank}")
         model.eval()
         total_num_utts, total_loss_dict = 0, {}  # avoid division by 0
         for batch_idx, batch_dict in enumerate(cv_data_loader):
@@ -160,17 +222,21 @@ class Executor:
             total_num_utts += num_utts
 
             if self.gan is True:
-                batch_dict['turn'] = 'generator'
+                batch_dict["turn"] = "generator"
             info_dict = batch_forward(model, batch_dict, None, info_dict)
 
-            for k, v in info_dict['loss_dict'].items():
+            for k, v in info_dict["loss_dict"].items():
                 if k not in total_loss_dict:
                     total_loss_dict[k] = []
                 total_loss_dict[k].append(v.item() * num_utts)
             log_per_step(None, info_dict)
         for k, v in total_loss_dict.items():
             total_loss_dict[k] = sum(v) / total_num_utts
-        info_dict['loss_dict'] = total_loss_dict
+        info_dict["loss_dict"] = total_loss_dict
         log_per_save(writer, info_dict)
-        model_name = 'epoch_{}_whole'.format(self.epoch) if on_batch_end else 'epoch_{}_step_{}'.format(self.epoch, self.step + 1)
+        model_name = (
+            "epoch_{}_whole".format(self.epoch)
+            if on_batch_end
+            else "epoch_{}_step_{}".format(self.epoch, self.step + 1)
+        )
         save_model(model, model_name, info_dict)

speech/cosyvoice/utils/train_utils.py

@@ -29,7 +29,7 @@ import torch.distributed as dist
 from torch.utils.tensorboard import SummaryWriter
 from torch.utils.data import DataLoader
 from torch.nn.utils import clip_grad_norm_
-
+from loguru import logger
 from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_live
 
 from cosyvoice.dataset.dataset import Dataset
@@ -40,8 +40,7 @@ def init_distributed(args):
     world_size = int(os.environ.get('WORLD_SIZE', 1))
     local_rank = int(os.environ.get('LOCAL_RANK', 0))
     rank = int(os.environ.get('RANK', 0))
-    logging.info('training on multiple gpus, this gpu {}'.format(local_rank) +
-                 ', rank {}, world_size {}'.format(rank, world_size))
+    logger.info(f'training on multiple gpus, this gpu {local_rank}, rank {rank}, world_size {world_size}')
     if args.train_engine == 'torch_ddp':
         torch.cuda.set_device(local_rank)
         dist.init_process_group(args.dist_backend)
@@ -70,6 +69,7 @@ def init_dataset_and_dataloader(args, configs, gan, dpo):
 
 
 def check_modify_and_save_config(args, configs):
+    """Check and modify config"""
     if args.train_engine == "torch_ddp":
         configs['train_conf']["dtype"] = 'fp32'
     else:
@@ -92,6 +92,7 @@ def check_modify_and_save_config(args, configs):
 
 
 def wrap_cuda_model(args, model):
+    """Wrap model to cuda"""
     local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', 1))
     world_size = int(os.environ.get('WORLD_SIZE', 1))
     if args.train_engine == "torch_ddp":  # native pytorch ddp
@@ -109,6 +110,7 @@ def wrap_cuda_model(args, model):
 
 
 def init_optimizer_and_scheduler(args, configs, model, gan):
+    """Init optimizer and scheduler"""
     if gan is False:
         if configs['train_conf']['optim'] == 'adam':
             optimizer = optim.Adam(model.parameters(), **configs['train_conf']['optim_conf'])
@@ -185,6 +187,7 @@ def init_optimizer_and_scheduler(args, configs, model, gan):
 
 
 def init_summarywriter(args):
+    
     writer = None
     if int(os.environ.get('RANK', 0)) == 0:
         os.makedirs(args.model_dir, exist_ok=True)
@@ -215,6 +218,7 @@ def save_model(model, model_name, info_dict):
 
 
 def cosyvoice_join(group_join, info_dict):
+    """Join all ranks"""
     world_size = int(os.environ.get('WORLD_SIZE', 1))
     local_rank = int(os.environ.get('LOCAL_RANK', 0))
     rank = int(os.environ.get('RANK', 0))
@@ -236,6 +240,7 @@ def cosyvoice_join(group_join, info_dict):
 
 
 def batch_forward(model, batch, scaler, info_dict, ref_model=None, dpo_loss=None):
+    """ Forward batch and compute loss"""
     device = int(os.environ.get('LOCAL_RANK', 0))
 
     dtype = info_dict["dtype"]
@@ -276,7 +281,7 @@ def batch_forward(model, batch, scaler, info_dict, ref_model=None, dpo_loss=None
 
 def batch_backward(model, scaler, info_dict):
     if info_dict["train_engine"] == "deepspeed":
-        scaled_loss = model.backward(info_dict['loss_dict']['loss'])
+        scaled_loss = model.backward(info_dict['loss_dict']['loss'])    
     else:
         scaled_loss = info_dict['loss_dict']['loss'] / info_dict['accum_grad']
         if scaler is not None:
@@ -356,9 +361,8 @@ def log_per_save(writer, info_dict):
     loss_dict = info_dict["loss_dict"]
     lr = info_dict['lr']
     rank = int(os.environ.get('RANK', 0))
-    logging.info(
-        'Epoch {} Step {} CV info lr {} {} rank {}'.format(
-            epoch, step + 1, lr, rank, ' '.join(['{} {}'.format(k, v) for k, v in loss_dict.items()])))
+    logger.info(
+        f'Epoch {epoch} Step {step + 1} CV info lr {lr} {rank} {''.join([f"{k} {v}" for k, v in loss_dict.items()])}')
 
     if writer is not None:
         for k in ['epoch', 'lr']:

speech/train.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+import argparse
+import datetime
+import logging
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+from copy import deepcopy
+import os
+import torch
+import torch.distributed as dist
+import deepspeed
+from loguru import logger
+
+from hyperpyyaml import load_hyperpyyaml
+
+from torch.distributed.elastic.multiprocessing.errors import record
+
+from cosyvoice.utils.losses import DPOLoss
+from cosyvoice.utils.executor import Executor
+from cosyvoice.utils.train_utils import (
+    init_distributed,
+    init_dataset_and_dataloader,
+    init_optimizer_and_scheduler,
+    init_summarywriter, save_model,
+    check_modify_and_save_config)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='training your network')
+    parser.add_argument('--train_engine',
+                        default='torch_ddp',
+                        choices=['torch_ddp', 'deepspeed'],
+                        help='Engine for paralleled training')
+    parser.add_argument('--model', required=True, help='model which will be trained')
+    parser.add_argument('--ref_model', required=False, help='ref model used in dpo')
+    parser.add_argument('--config', required=True, help='config file')
+    parser.add_argument('--train_data', required=True, help='train data file')
+    parser.add_argument('--cv_data', required=True, help='cv data file')
+    parser.add_argument('--qwen_pretrain_path', required=False, help='qwen pretrain path')
+    parser.add_argument('--checkpoint', help='checkpoint model')
+    parser.add_argument('--model_dir', required=True, help='save model dir')
+    parser.add_argument('--tensorboard_dir',
+                        default='tensorboard',
+                        help='tensorboard log dir')
+    parser.add_argument('--ddp.dist_backend',
+                        dest='dist_backend',
+                        default='nccl',
+                        choices=['nccl', 'gloo'],
+                        help='distributed backend')
+    parser.add_argument('--num_workers',
+                        default=0,
+                        type=int,
+                        help='num of subprocess workers for reading')
+    parser.add_argument('--prefetch',
+                        default=100,
+                        type=int,
+                        help='prefetch number')
+    parser.add_argument('--pin_memory',
+                        action='store_true',
+                        default=False,
+                        help='Use pinned memory buffers used for reading')
+    parser.add_argument('--use_amp',
+                        action='store_true',
+                        default=False,
+                        help='Use automatic mixed precision training')
+    parser.add_argument('--dpo',
+                        action='store_true',
+                        default=False,
+                        help='Use Direct Preference Optimization')
+    parser.add_argument('--deepspeed.save_states',
+                        dest='save_states',
+                        default='model_only',
+                        choices=['model_only', 'model+optimizer'],
+                        help='save model/optimizer states')
+    parser.add_argument('--timeout',
+                        default=60,
+                        type=int,
+                        help='timeout (in seconds) of cosyvoice_join.')
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args()
+    return args
+
+
+@record
+def main():
+    args = get_args()
+    logging.basicConfig(level=logging.DEBUG,
+                        format='%(asctime)s %(levelname)s %(message)s')
+    # gan train has some special initialization logic
+    gan = True if args.model == 'hifigan' else False
+
+    override_dict = {k: None for k in ['llm', 'flow', 'hift', 'hifigan'] if k != args.model}
+    if gan is True:
+        override_dict.pop('hift')
+    try:
+        with open(args.config, 'r') as f:
+            configs = load_hyperpyyaml(f, overrides={**override_dict, 'qwen_pretrain_path': args.qwen_pretrain_path})
+    except Exception:
+        with open(args.config, 'r') as f:
+            configs = load_hyperpyyaml(f, overrides=override_dict)
+    if gan is True:
+        configs['train_conf'] = configs['train_conf_gan']
+    configs['train_conf'].update(vars(args))
+
+    # Init env for ddp
+    init_distributed(args)
+
+    # Get dataset & dataloader
+    train_dataset, _, train_data_loader, cv_data_loader = \
+        init_dataset_and_dataloader(args, configs, gan, args.dpo)
+
+    # Do some sanity checks and save config to arsg.model_dir
+    configs = check_modify_and_save_config(args, configs)
+
+    # Tensorboard summary
+    writer = init_summarywriter(args)
+
+    # load checkpoint
+    if args.dpo is True:
+        configs[args.model].forward = configs[args.model].forward_dpo
+
+    model = configs[args.model]
+    start_step, start_epoch = 0, -1
+    if args.checkpoint is not None:
+        if os.path.exists(args.checkpoint):
+            state_dict = torch.load(args.checkpoint, map_location='cpu')
+            model.load_state_dict(state_dict, strict=False)
+            if 'step' in state_dict:
+                start_step = state_dict['step']
+            if 'epoch' in state_dict:
+                start_epoch = state_dict['epoch']
+        else:
+            logging.warning('checkpoint {} do not exsist!'.format(args.checkpoint))
+
+    # Dispatch model from cpu to gpu
+    model = model.cuda()
+    model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)
+   
+
+    # Get optimizer & scheduler
+    model, optimizer, scheduler, optimizer_d, scheduler_d = init_optimizer_and_scheduler(args, configs, model, gan)
+    scheduler.set_step(start_step)
+    if scheduler_d is not None:
+        scheduler_d.set_step(start_step)
+
+    # Save init checkpoints
+    info_dict = deepcopy(configs['train_conf'])
+    info_dict['step'] = start_step
+    info_dict['epoch'] = start_epoch
+    save_model(model, 'init', info_dict)
+
+    # DPO related
+    if args.dpo is True:
+        ref_model = deepcopy(configs[args.model])
+        state_dict = torch.load(args.ref_model, map_location='cpu')
+        ref_model.load_state_dict(state_dict, strict=False)
+        dpo_loss = DPOLoss(beta=0.01, label_smoothing=0.0, ipo=False)
+        ref_model = ref_model.cuda()
+        ref_model = torch.nn.parallel.DistributedDataParallel(ref_model, find_unused_parameters=True)
+    else:
+        ref_model, dpo_loss = None, None
+
+    # Get executor
+    executor = Executor(gan=gan, ref_model=ref_model, dpo_loss=dpo_loss)
+    executor.step = start_step
+
+    # Init scaler, used for pytorch amp mixed precision training
+    scaler = torch.amp.GradScaler() if args.use_amp else None
+    logger.info(f'start step {start_step} start epoch {start_epoch}')
+
+    # Start training loop
+    for epoch in range(start_epoch + 1, info_dict['max_epoch']):
+        executor.epoch = epoch
+        train_dataset.set_epoch(epoch)
+        dist.barrier()
+        group_join = dist.new_group(backend="nccl", timeout=datetime.timedelta(seconds=args.timeout))
+        if gan is True:
+            executor.train_one_epoc_gan(model, optimizer, scheduler, optimizer_d, scheduler_d, train_data_loader, cv_data_loader,
+                                        writer, info_dict, scaler, group_join)
+        else:
+            executor.train_one_epoc(model, optimizer, scheduler, train_data_loader, cv_data_loader, writer, info_dict, scaler, group_join)
+        dist.destroy_process_group(group_join)
+
+
+if __name__ == '__main__':
+    main()