debug training code

speech/cosyvoice/dataset/processor.py

@@ -313,9 +313,23 @@ def compute_fbank(data,
         feat = feat_extractor(waveform).squeeze(dim=0).transpose(0, 1)
         if token_mel_ratio != 0:
             # trim to align speech_token and speech_feat
-            token_len = int(min(feat.shape[0] / token_mel_ratio, sample["speech_token"].shape[0]))
+            # token_len = int(min(feat.shape[0] / token_mel_ratio, sample["speech_token"].shape[0]))
+            # feat = feat[:token_mel_ratio * token_len]
+            # sample["speech_token"] = sample["speech_token"][:token_len]
+
+            # Convert speech_token to tensor if it's a list
+            if isinstance(sample["speech_token"], list):
+                speech_token_tensor = torch.tensor(sample["speech_token"])
+            else:
+                speech_token_tensor = sample["speech_token"]
+            
+            # trim to align speech_token and speech_feat
+            token_len = int(min(feat.shape[0] / token_mel_ratio, speech_token_tensor.shape[0]))
             feat = feat[:token_mel_ratio * token_len]
-            sample["speech_token"] = sample["speech_token"][:token_len]
+            
+            # Update speech_token - keep as tensor for consistency
+            sample["speech_token"] = speech_token_tensor[:token_len]
+
         sample['speech_feat'] = feat
         yield sample
 

speech/cosyvoice/flow/decoder.py

@@ -421,6 +421,8 @@ class CausalConditionalDecoder(ConditionalDecoder):
         """
         t = self.time_embeddings(t).to(t.dtype)
         t = self.time_mlp(t)
+        # print('x shape', x.shape)
+        # print('mu shape:', mu.shape)
 
         x = pack([x, mu], "b * t")[0]
 

speech/cosyvoice/flow/flow.py

@@ -73,6 +73,7 @@ class MaskedDiffWithXvec(torch.nn.Module):
         embedding = self.spk_embed_affine_layer(embedding)
 
         # concat text and prompt_text
+        print('token_len values: ', token_len)
         mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
 
@@ -197,13 +198,19 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         feat_len = batch['speech_feat_len'].to(device)
         embedding = batch['embedding'].to(device)
 
+        # print('token: ', token.shape)
+        # print('token_len: ', token_len.shape)
+        # print('feat: ', feat.shape)
+        # print('feat_len: ', feat_len.shape)
+        # print('embedding: ', embedding.shape)
+
         # NOTE unified training, static_chunk_size > 0 or = 0
-        streaming = True if random.random() < 0.5 else False
+        streaming = False# if random.random() < 0.5 else False
 
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
         embedding = self.spk_embed_affine_layer(embedding)
-
+        # print('token_len values: ', token_len)
         # concat text and prompt_text
         mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
@@ -222,6 +229,14 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         conds = conds.transpose(1, 2)
 
         mask = (~make_pad_mask(h_lengths.sum(dim=-1).squeeze(dim=1))).to(h)
+
+        # print('feat shape: ', feat.shape)
+        # print('mask shape: ', mask.shape)
+        # print('h shape: ', h.shape)
+        # print('embedding shape: ', embedding.shape)
+        # print('conds shape: ', conds.shape)
+        # print('streaming: ', streaming)
+
         loss, _ = self.decoder.compute_loss(
             feat.transpose(1, 2).contiguous(),
             mask.unsqueeze(1),

speech/cosyvoice/flow/flow_matching.py

@@ -187,7 +187,13 @@ class ConditionalCFM(BASECFM):
             mu = mu * cfg_mask.view(-1, 1, 1)
             spks = spks * cfg_mask.view(-1, 1)
             cond = cond * cfg_mask.view(-1, 1, 1)
-
+        # print('y shape: ', y.shape)
+        # print('mask shape: ', mask.shape)
+        # print('mu shape: ', mu.shape)
+        # print('t shape: ', t.shape)
+        # print('spks shape: ', spks.shape)
+        # print('cond shape: ', cond.shape)
+        # print('streaming: ', streaming)
         pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond, streaming=streaming)
         loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
         return loss, y

speech/cosyvoice/transformer/upsample_encoder.py

@@ -272,6 +272,14 @@ class UpsampleConformerEncoder(torch.nn.Module):
             checkpointing API because `__call__` attaches all the hooks of the module.
             https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2
         """
+
+        # print('xs shape: ', xs.shape)
+        # print('xs_lens shape: ', xs_lens.shape)
+        # print('context shape: ', context.shape)
+        # print('decoding_chunk_size: ', decoding_chunk_size)
+        # print('num_decoding_left_chunks: ', num_decoding_left_chunks)
+        # print('streaming: ', streaming)
+
         T = xs.size(1)
         masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
         if self.global_cmvn is not None:
@@ -303,6 +311,8 @@ class UpsampleConformerEncoder(torch.nn.Module):
         # Here we assume the mask is not changed in encoder layers, so just
         # return the masks before encoder layers, and the masks will be used
         # for cross attention with decoder later
+        # print('output xs shape: ', xs.shape)
+        # print('output masks shape: ', masks.shape)
         return xs, masks
 
     def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,

speech/cosyvoice/utils/executor.py

@@ -48,11 +48,9 @@ class Executor:
         optimizer,
         scheduler,
         train_data_loader,
-        cv_data_loader,
         experiment,
         info_dict,
         scaler,
-        group_join,
         model_type
     ):
         """Train one epoch"""
@@ -68,58 +66,51 @@ class Executor:
         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
-        )
-        with model_context():
-            for batch_idx, batch_dict in enumerate(train_data_loader):
-                info_dict["tag"] = "TRAIN"
-                info_dict["step"] = self.step
-                info_dict["epoch"] = self.epoch
-                info_dict["batch_idx"] = batch_idx
-                if cosyvoice_join(group_join, info_dict):
-                    break
-
-               
-                if (
-                    info_dict["train_engine"] == "torch_ddp"
-                    and (batch_idx + 1) % info_dict["accum_grad"] != 0
-                ):
-                    context = model.no_sync
-            
-                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_backward(model, scaler, info_dict)
-
-                info_dict = update_parameter_and_lr(
-                    model, optimizer, scheduler, scaler, info_dict, model_type=model_type
+
+        use_ddp = info_dict["train_engine"] == "torch_ddp"
+
+        for batch_idx, batch_dict in enumerate(train_data_loader):
+            info_dict["tag"] = "TRAIN"
+            info_dict["step"] = self.step
+            info_dict["epoch"] = self.epoch
+            info_dict["batch_idx"] = batch_idx
+
+            if use_ddp and (batch_idx + 1) % info_dict["accum_grad"] != 0:
+                context = model.no_sync
+            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,
                 )
-                log_per_step(experiment, info_dict)
-                
-                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
-                ):
+                info_dict = batch_backward(model, scaler, info_dict)
+
+            info_dict = update_parameter_and_lr(
+                model, optimizer, scheduler, scaler, info_dict, model_type=model_type
+            )
+            log_per_step(experiment, info_dict)
+
+            if (
+                info_dict.get("save_per_step", -1) > 0
+                and (self.step + 1) % info_dict["save_per_step"] == 0
+                and (batch_idx + 1) % info_dict["accum_grad"] == 0
+            ):
+                if dist.is_initialized():
                     dist.barrier()
-                    self.cv(
-                        model, cv_data_loader, experiment, info_dict, on_batch_end=False
-                    )
-                    model.train()
-                if (batch_idx + 1) % info_dict["accum_grad"] == 0:
-                    self.step += 1
+                model_name = (
+                    f"epoch_{self.epoch}_step_{self.step + 1}"
+                )
+                save_model(model, model_name, info_dict)
+                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)
 
     @torch.inference_mode()
     def cv(self, model, cv_data_loader, experiment, info_dict, on_batch_end=True):

speech/cosyvoice/utils/train_utils.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import logging
 import os
 import torch
 import json
@@ -33,7 +32,49 @@ from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem
 
 from cosyvoice.dataset.dataset import Dataset
 
-from torch.optim.lr_scheduler import LinearLR, ConstantLR, SequentialLR
+from torch.optim.lr_scheduler import LinearLR, ConstantLR, SequentialLR, _LRScheduler
+
+from loguru import logger
+
+class ResumableSequentialLR(_LRScheduler):
+    """A resumable version of SequentialLR that supports set_step"""
+    def __init__(self, optimizer, schedulers, milestones, last_epoch=-1):
+        self.schedulers = schedulers
+        self.milestones = milestones
+        self._last_lr = [group['lr'] for group in optimizer.param_groups]
+        super().__init__(optimizer, last_epoch)
+        
+    def get_lr(self):
+        # Find which scheduler to use based on last_epoch
+        idx = 0
+        for i, milestone in enumerate(self.milestones):
+            if self.last_epoch >= milestone:
+                idx = i + 1
+        
+        if idx >= len(self.schedulers):
+            idx = len(self.schedulers) - 1
+            
+        # Get lr from the appropriate scheduler
+        scheduler = self.schedulers[idx]
+        if hasattr(scheduler, '_get_closed_form_lr'):
+            return scheduler._get_closed_form_lr()
+        else:
+            return scheduler.get_lr()
+    
+    def step(self, epoch=None):
+        if epoch is None:
+            self.last_epoch += 1
+        else:
+            self.last_epoch = epoch
+            
+        # Update learning rates
+        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
+            param_group['lr'] = lr
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+    
+    def set_step(self, step):
+        """Set the current step for resuming training"""
+        self.last_epoch = step - 1  # -1 because step() will increment it
 
 def init_distributed(args):
     world_size = int(os.environ.get('WORLD_SIZE', 1))
@@ -100,7 +141,7 @@ def wrap_cuda_model(args, model):
         model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)
     else:
         if int(os.environ.get('RANK', 0)) == 0:
-            logging.info("Estimating model states memory needs (zero2)...")
+            logger.info("Estimating model states memory needs (zero2)...")
             estimate_zero2_model_states_mem_needs_all_live(
                 model,
                 num_gpus_per_node=local_world_size,
@@ -132,12 +173,13 @@ def init_optimizer_and_scheduler(configs, model):
     )
     
     # Combine schedulers: warmup for 5k steps, then constant
-    scheduler = SequentialLR(
+    scheduler = ResumableSequentialLR(
         optimizer,
         schedulers=[warmup_scheduler, constant_scheduler],
-        milestones=[5000]  # Switch after 5k steps
+        milestones=[5000]
     )
 
+
     return model, optimizer, scheduler
 
 
@@ -146,7 +188,9 @@ def save_model(model, model_name, info_dict):
     """Save model"""
     rank = int(os.environ.get('RANK', 0))
     model_dir = info_dict["model_dir"]
+    # os.makedirs(model_dir, exist_ok=True)
     save_model_path = os.path.join(model_dir, '{}.pt'.format(model_name))
+    
 
     if info_dict["train_engine"] == "torch_ddp":
         if rank == 0:
@@ -162,7 +206,7 @@ def save_model(model, model_name, info_dict):
         with open(info_path, 'w') as fout:
             data = yaml.dump(info_dict)
             fout.write(data)
-        logging.info('[Rank {}] Checkpoint: save to checkpoint {}'.format(rank, save_model_path))
+        logger.info('[Rank {}] Checkpoint: save to checkpoint {}'.format(rank, save_model_path))
 
 
 def cosyvoice_join(group_join, info_dict):
@@ -178,7 +222,7 @@ def cosyvoice_join(group_join, info_dict):
                                    timeout=group_join.options._timeout)
             return False
         except RuntimeError as e:
-            logging.info("Detected uneven workload distribution: {}\n".format(e) +
+            logger.info("Detected uneven workload distribution: {}\n".format(e) +
                          "Break current worker to manually join all workers, " +
                          "world_size {}, current rank {}, current local_rank {}\n".
                          format(world_size, rank, local_rank))
@@ -326,14 +370,14 @@ def update_parameter_and_lr(model, optimizer, scheduler, scaler, info_dict, mode
             if torch.isfinite(grad_norm):
                 scaler.step(optimizer)
             else:
-                logging.warning('get infinite grad_norm, check your code/data if it appears frequently')
+                logger.warning('get infinite grad_norm, check your code/data if it appears frequently')
             scaler.update()
         else:
             grad_norm = clip_grad_norm_(model.parameters(), info_dict['grad_clip'])
             if torch.isfinite(grad_norm):
                 optimizer.step()
             else:
-                logging.warning('get infinite grad_norm, check your code/data if it appears frequently')
+                logger.warning('get infinite grad_norm, check your code/data if it appears frequently')
         optimizer.zero_grad()
         scheduler.step()
     info_dict["lr"] = optimizer.param_groups[0]['lr']
@@ -376,7 +420,7 @@ def log_per_step(experiment, info_dict):
         if tag == "TRAIN":
             log_str += f'lr {info_dict["lr"]:.8f} grad_norm {info_dict["grad_norm"]:.6f}'
         log_str += f' rank {rank}'
-        logging.debug(log_str)
+        logger.info(log_str)
 
 def log_per_save(experiment, info_dict):
     """Log per save using Comet ML"""
@@ -387,7 +431,7 @@ def log_per_save(experiment, info_dict):
     lr = info_dict['lr']
     rank = int(os.environ.get('RANK', 0))
     
-    # Create loss string for logging
+    # Create loss string for logger
     loss_str = ' '.join([f"{k} {v.item() if isinstance(v, torch.Tensor) else v}" for k, v in loss_dict.items()])
     logger.info(f'Epoch {epoch} Step {step + 1} CV info lr {lr} {rank} {loss_str}')
 

speech/train.py

@@ -22,11 +22,11 @@ from copy import deepcopy
 import deepspeed
 import torch
 import torch.distributed as dist
+from comet_ml import Experiment
 from hyperpyyaml import load_hyperpyyaml
 from loguru import logger
 from torch.distributed.elastic.multiprocessing.errors import record
 
-from comet_ml import Experiment
 from cosyvoice.utils.executor import Executor
 from cosyvoice.utils.losses import DPOLoss
 from cosyvoice.utils.train_utils import (check_modify_and_save_config,
@@ -35,6 +35,8 @@ from cosyvoice.utils.train_utils import (check_modify_and_save_config,
                                          save_model)
 
 os.environ["COMET_LOGGING_CONSOLE"] = "ERROR"  # Only show errors
+
+
 def get_args():
     parser = argparse.ArgumentParser(description="training your network")
     parser.add_argument(
@@ -107,14 +109,8 @@ def get_args():
         default=False,
         help="Disable comet ml experiment",
     )
-    parser.add_argument(
-        "--comet_project",
-        default="speech"
-    )
-    parser.add_argument(
-        "--comet_experiment_name",
-        default="test"
-    )
+    parser.add_argument("--comet_project", default="speech")
+    parser.add_argument("--comet_experiment_name", default="test")
     parser = deepspeed.add_config_arguments(parser)
     args = parser.parse_args()
     return args
@@ -122,8 +118,8 @@ def get_args():
 
 def init_comet_experiment(args, configs):
     """Initialize Comet ML experiment"""
-    rank = int(os.environ.get('RANK', 0))
-    
+    rank = int(os.environ.get("RANK", 0))
+
     # Only create experiment on rank 0 to avoid duplicates
     if rank == 0 and not args.comet_disabled:
         # Set up Comet ML experiment
@@ -131,7 +127,7 @@ def init_comet_experiment(args, configs):
             project_name=args.comet_project,
             experiment_name=args.comet_experiment_name,
         )
-        
+
         # Log hyperparameters
         experiment.log_parameters(configs["train_conf"])
         experiment.log_parameter("model_type", args.model)
@@ -141,24 +137,29 @@ def init_comet_experiment(args, configs):
         experiment.log_parameter("dpo", args.dpo)
         experiment.log_parameter("num_workers", args.num_workers)
         experiment.log_parameter("prefetch", args.prefetch)
-        
+
         # Log model architecture if available
         if args.model in configs:
-            model_config = configs[args.model].__dict__ if hasattr(configs[args.model], '__dict__') else {}
+            model_config = (
+                configs[args.model].__dict__
+                if hasattr(configs[args.model], "__dict__")
+                else {}
+            )
             experiment.log_parameters(model_config, prefix=f"{args.model}/")
-        
+
         # Add tags
         experiment.add_tag(args.model)
         if args.dpo:
             experiment.add_tag("dpo")
         if args.use_amp:
             experiment.add_tag("amp")
-            
+
         logger.info(f"Comet ML experiment initialized: {experiment.get_name()}")
         return experiment
     else:
         return None
 
+
 @record
 def main():
     args = get_args()
@@ -182,12 +183,14 @@ def main():
 
     configs["train_conf"].update(vars(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))
-    logger.info(f'training on multiple gpus, this gpu {local_rank}, rank {rank}, world_size {world_size}')
+    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))
+    logger.info(
+        f"training on multiple gpus, this gpu {local_rank}, rank {rank}, world_size {world_size}"
+    )
     torch.cuda.set_device(local_rank)
-    dist.init_process_group(args.dist_backend)
+    dist.init_process_group("nccl")
 
     # Get dataset & dataloader
     train_dataset, _, train_data_loader, cv_data_loader = init_dataset_and_dataloader(
@@ -200,7 +203,6 @@ def main():
     # Tensorboard summary
     experiment = init_comet_experiment(args, configs)
 
-
     # load checkpoint
     if args.dpo is True:
         configs[args.model].forward = configs[args.model].forward_dpo
@@ -230,9 +232,7 @@ def main():
     )
 
     # Get optimizer & scheduler
-    model, optimizer, scheduler = (
-        init_optimizer_and_scheduler(configs, model)
-    )
+    model, optimizer, scheduler = init_optimizer_and_scheduler(configs, model)
     scheduler.set_step(start_step)
 
     # Save init checkpoints
@@ -246,7 +246,7 @@ def main():
         experiment.log_model(
             name=f"{args.model}_init",
             file_or_folder=os.path.join(args.model_dir, "init.pt"),
-            metadata=info_dict
+            metadata=info_dict,
         )
 
     # DPO related
@@ -279,26 +279,23 @@ def main():
     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)
-        )
-        
         executor.train_one_epoc(
             model,
             optimizer,
             scheduler,
             train_data_loader,
-            cv_data_loader,
             experiment,
             info_dict,
             scaler,
-            group_join,
-            model_type=args.model
+            model_type=args.model,
         )
-        dist.destroy_process_group(group_join)
+
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
     if experiment:
         experiment.end()
 
+
 if __name__ == "__main__":
     main()