chore: import upstream snapshot with attribution

2026-07-13 12:35:17 +08:00
commit 344816a5d8
136 changed files with 25044 additions and 0 deletions
@@ -0,0 +1,514 @@
+import logging
+import os
+import pathlib
+import shutil
+import sys
+from typing import Dict
+
+import matplotlib
+
+import utils
+
+matplotlib.use('Agg')
+
+import torch.utils.data
+from torchmetrics import Metric, MeanMetric
+import lightning.pytorch as pl
+from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_only
+
+from basics.base_module import CategorizedModule
+from utils.hparams import hparams
+from utils.training_utils import (
+    DsModelCheckpoint, DsTQDMProgressBar,
+    DsBatchSampler, DsTensorBoardLogger,
+    get_latest_checkpoint_path, get_strategy
+)
+from utils.phoneme_utils import load_phoneme_dictionary
+
+torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system'))
+
+log_format = '%(asctime)s %(message)s'
+logging.basicConfig(stream=sys.stdout, level=logging.INFO,
+                    format=log_format, datefmt='%m/%d %I:%M:%S %p')
+
+
+class BaseTask(pl.LightningModule):
+    """
+        Base class for training tasks.
+        1. *load_ckpt*:
+            load checkpoint;
+        2. *training_step*:
+            record and log the loss;
+        3. *optimizer_step*:
+            run backwards step;
+        4. *start*:
+            load training configs, backup code, log to tensorboard, start training;
+        5. *configure_ddp* and *init_ddp_connection*:
+            start parallel training.
+
+        Subclasses should define:
+        1. *build_model*, *build_optimizer*, *build_scheduler*:
+            how to build the model, the optimizer and the training scheduler;
+        2. *_training_step*:
+            one training step of the model;
+        3. *on_validation_end* and *_on_validation_end*:
+            postprocess the validation output.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_batch_frames = hparams['max_batch_frames']
+        self.max_batch_size = hparams['max_batch_size']
+        self.max_val_batch_frames = hparams['max_val_batch_frames']
+        if self.max_val_batch_frames == -1:
+            hparams['max_val_batch_frames'] = self.max_val_batch_frames = self.max_batch_frames
+        self.max_val_batch_size = hparams['max_val_batch_size']
+        if self.max_val_batch_size == -1:
+            hparams['max_val_batch_size'] = self.max_val_batch_size = self.max_batch_size
+
+        self.training_sampler = None
+        self.skip_immediate_validation = False
+        self.skip_immediate_ckpt_save = False
+
+        self.phoneme_dictionary = load_phoneme_dictionary()
+        self.build_model()
+
+        self.valid_losses: Dict[str, Metric] = {}
+        self.valid_metrics: Dict[str, Metric] = {}
+
+    def _finish_init(self):
+        self.register_validation_loss('total_loss')
+        self.build_losses_and_metrics()
+        assert len(self.valid_losses) > 0, "No validation loss registered. Please check your configuration file."
+
+    ###########
+    # Training, validation and testing
+    ###########
+    def setup(self, stage):
+        self.train_dataset = self.dataset_cls('train')
+        self.valid_dataset = self.dataset_cls('valid')
+        self.num_replicas = (self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1)
+
+    def get_need_freeze_state_dict_key(self, model_state_dict) -> list:
+        key_list = []
+        for i in hparams['frozen_params']:
+            for j in model_state_dict:
+                if j.startswith(i):
+                    key_list.append(j)
+        return list(set(key_list))
+
+    def freeze_params(self) -> None:
+        model_state_dict = self.state_dict().keys()
+        freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict)
+
+        for i in freeze_key:
+            params=self.get_parameter(i)
+
+            params.requires_grad = False
+
+    def unfreeze_all_params(self) -> None:
+        for i in self.model.parameters():
+            i.requires_grad = True
+
+    def load_finetune_ckpt(
+            self, state_dict
+    ) -> None:
+        adapt_shapes = hparams['finetune_strict_shapes']
+        if not adapt_shapes:
+            cur_model_state_dict = self.state_dict()
+            unmatched_keys = []
+            for key, param in state_dict.items():
+                if key in cur_model_state_dict:
+                    new_param = cur_model_state_dict[key]
+                    if new_param.shape != param.shape:
+                        unmatched_keys.append(key)
+                        print('| Unmatched keys: ', key, new_param.shape, param.shape)
+            for key in unmatched_keys:
+                del state_dict[key]
+        self.load_state_dict(state_dict, strict=False)
+
+    def load_pre_train_model(self):
+        pre_train_ckpt_path = hparams['finetune_ckpt_path']
+        blacklist = hparams['finetune_ignored_params']
+        # whitelist=hparams['pre_train_whitelist']
+        if blacklist is None:
+            blacklist = []
+        # if whitelist is  None:
+        #     raise RuntimeError("")
+
+        if pre_train_ckpt_path is not None:
+            ckpt = torch.load(pre_train_ckpt_path)
+            # if ckpt.get('category') is None:
+            #     raise RuntimeError("")
+
+            if isinstance(self.model, CategorizedModule):
+                self.model.check_category(ckpt.get('category'))
+
+            state_dict = {}
+            for i in ckpt['state_dict']:
+                # if 'diffusion' in i:
+                # if i in rrrr:
+                #     continue
+                skip = False
+                for b in blacklist:
+                    if i.startswith(b):
+                        skip = True
+                        break
+
+                if skip:
+                    continue
+
+                state_dict[i] = ckpt['state_dict'][i]
+                print(i)
+            return state_dict
+        else:
+            raise RuntimeError("")
+
+    def _build_model(self):
+        raise NotImplementedError()
+
+    def build_model(self):
+        self.model = self._build_model()
+        # utils.load_warp(self)
+        self.unfreeze_all_params()
+        if hparams['freezing_enabled']:
+            self.freeze_params()
+        if hparams['finetune_enabled'] and get_latest_checkpoint_path(pathlib.Path(hparams['work_dir'])) is None:
+            self.load_finetune_ckpt(self.load_pre_train_model())
+        self.print_arch()
+
+    @rank_zero_only
+    def print_arch(self):
+        utils.print_arch(self.model)
+
+    def build_losses_and_metrics(self):
+        raise NotImplementedError()
+
+    def register_validation_metric(self, name: str, metric: Metric):
+        assert isinstance(metric, Metric)
+        self.valid_metrics[name] = metric
+
+    def register_validation_loss(self, name: str, Aggregator: Metric = MeanMetric):
+        assert issubclass(Aggregator, Metric)
+        self.valid_losses[name] = Aggregator()
+
+    def run_model(self, sample, infer=False):
+        """
+        steps:
+            1. run the full model
+            2. calculate losses if not infer
+        """
+        raise NotImplementedError()
+
+    def on_train_epoch_start(self):
+        if self.training_sampler is not None:
+            self.training_sampler.set_epoch(self.current_epoch)
+
+    def _training_step(self, sample):
+        """
+        :return: total loss: torch.Tensor, loss_log: dict, other_log: dict
+        """
+        losses = self.run_model(sample)
+        total_loss = sum(losses.values())
+        return total_loss, {**losses, 'batch_size': float(sample['size'])}
+
+    def training_step(self, sample, batch_idx):
+        total_loss, log_outputs = self._training_step(sample)
+
+        # logs to progress bar
+        self.log_dict(log_outputs, prog_bar=True, logger=False, on_step=True, on_epoch=False)
+        self.log('lr', self.lr_schedulers().get_last_lr()[0], prog_bar=True, logger=False, on_step=True, on_epoch=False)
+        # logs to tensorboard
+        if self.global_step % hparams['log_interval'] == 0:
+            tb_log = {f'training/{k}': v for k, v in log_outputs.items()}
+            tb_log['training/lr'] = self.lr_schedulers().get_last_lr()[0]
+            self.logger.log_metrics(tb_log, step=self.global_step)
+
+        return total_loss
+
+    # def on_before_optimizer_step(self, *args, **kwargs):
+    #     self.log_dict(grad_norm(self, norm_type=2))
+
+    def _on_validation_start(self):
+        pass
+
+    def on_validation_start(self):
+        if self.skip_immediate_validation:
+            rank_zero_debug("Skip validation")
+            return
+        self._on_validation_start()
+        for metric in self.valid_losses.values():
+            metric.to(self.device)
+            metric.reset()
+        for metric in self.valid_metrics.values():
+            metric.to(self.device)
+            metric.reset()
+
+    def _validation_step(self, sample, batch_idx):
+        """
+
+        :param sample:
+        :param batch_idx:
+        :return: loss_log: dict, weight: int
+        """
+        raise NotImplementedError()
+
+    def validation_step(self, sample, batch_idx):
+        """
+
+        :param sample:
+        :param batch_idx:
+        """
+        if self.skip_immediate_validation:
+            rank_zero_debug("Skip validation")
+            return
+        if sample['size'] > 0:
+            with torch.autocast(self.device.type, enabled=False):
+                losses, weight = self._validation_step(sample, batch_idx)
+            losses = {
+                'total_loss': sum(losses.values()),
+                **losses
+            }
+            for k, v in losses.items():
+                self.valid_losses[k].update(v, weight=weight)
+
+    def _on_validation_epoch_end(self):
+        pass
+
+    def on_validation_epoch_end(self):
+        if self.skip_immediate_validation:
+            self.skip_immediate_validation = False
+            self.skip_immediate_ckpt_save = True
+            return
+        self._on_validation_epoch_end()
+        loss_vals = {k: v.compute() for k, v in self.valid_losses.items()}
+        metric_vals = {k: v.compute() for k, v in self.valid_metrics.items()}
+        self.log('val_loss', loss_vals['total_loss'], on_epoch=True, prog_bar=True, logger=False, sync_dist=True)
+        self.logger.log_metrics({f'validation/{k}': v for k, v in loss_vals.items()}, step=self.global_step)
+        self.logger.log_metrics({f'metrics/{k}': v for k, v in metric_vals.items()}, step=self.global_step)
+
+    # noinspection PyMethodMayBeStatic
+    def build_scheduler(self, optimizer):
+        from utils import build_lr_scheduler_from_config
+
+        scheduler_args = hparams['lr_scheduler_args']
+        assert scheduler_args['scheduler_cls'] != ''
+        scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args)
+        return scheduler
+
+    # noinspection PyMethodMayBeStatic
+    def build_optimizer(self, model):
+        from utils import build_object_from_class_name
+
+        optimizer_args = hparams['optimizer_args']
+        assert optimizer_args['optimizer_cls'] != ''
+        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
+            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
+        optimizer = build_object_from_class_name(
+            optimizer_args['optimizer_cls'],
+            torch.optim.Optimizer,
+            model if optimizer_args['optimizer_cls'] == 'modules.optimizer.muon.Muon_AdamW' else model.parameters(),
+            **optimizer_args
+        )
+        return optimizer
+
+    def configure_optimizers(self):
+        optm = self.build_optimizer(self.model)
+        scheduler = self.build_scheduler(optm)
+        if scheduler is None:
+            return optm
+        return {
+            "optimizer": optm,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1
+            }
+        }
+
+    def train_dataloader(self):
+        self.training_sampler = DsBatchSampler(
+            self.train_dataset,
+            max_batch_frames=self.max_batch_frames,
+            max_batch_size=self.max_batch_size,
+            num_replicas=self.num_replicas,
+            rank=self.global_rank,
+            sort_by_similar_size=hparams['sort_by_len'],
+            size_reversed=True,
+            required_batch_count_multiple=hparams['accumulate_grad_batches'],
+            shuffle_sample=True,
+            shuffle_batch=True
+        )
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            collate_fn=self.train_dataset.collater,
+            batch_sampler=self.training_sampler,
+            num_workers=hparams['ds_workers'],
+            prefetch_factor=hparams['dataloader_prefetch_factor'],
+            pin_memory=True,
+            persistent_workers=True
+        )
+
+    def val_dataloader(self):
+        sampler = DsBatchSampler(
+            self.valid_dataset,
+            max_batch_frames=self.max_val_batch_frames,
+            max_batch_size=self.max_val_batch_size,
+            num_replicas=self.num_replicas,
+            rank=self.global_rank,
+            shuffle_sample=False,
+            shuffle_batch=False,
+            disallow_empty_batch=False,
+            pad_batch_assignment=False
+        )
+        return torch.utils.data.DataLoader(
+            self.valid_dataset,
+            collate_fn=self.valid_dataset.collater,
+            batch_sampler=sampler,
+            num_workers=hparams['ds_workers'],
+            prefetch_factor=hparams['dataloader_prefetch_factor'],
+            persistent_workers=True
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()
+
+    def on_test_start(self):
+        self.on_validation_start()
+
+    def test_step(self, sample, batch_idx):
+        return self.validation_step(sample, batch_idx)
+
+    def on_test_end(self):
+        return self.on_validation_end()
+
+    ###########
+    # Running configuration
+    ###########
+
+    @classmethod
+    def start(cls):
+        task = cls()
+
+        # if pre_train is not None:
+        #     task.load_state_dict(pre_train,strict=False)
+        #     print("load success-------------------------------------------------------------------")
+
+        work_dir = pathlib.Path(hparams['work_dir'])
+        trainer = pl.Trainer(
+            accelerator=hparams['pl_trainer_accelerator'],
+            devices=hparams['pl_trainer_devices'],
+            num_nodes=hparams['pl_trainer_num_nodes'],
+            strategy=get_strategy(
+                hparams['pl_trainer_devices'],
+                hparams['pl_trainer_num_nodes'],
+                hparams['pl_trainer_accelerator'],
+                hparams['pl_trainer_strategy'],
+                hparams['pl_trainer_precision'],
+            ),
+            precision=hparams['pl_trainer_precision'],
+            callbacks=[
+                DsModelCheckpoint(
+                    dirpath=work_dir,
+                    filename='model_ckpt_steps_{step}',
+                    auto_insert_metric_name=False,
+                    monitor='step',
+                    mode='max',
+                    save_last=False,
+                    # every_n_train_steps=hparams['val_check_interval'],
+                    save_top_k=hparams['num_ckpt_keep'],
+                    permanent_ckpt_start=hparams['permanent_ckpt_start'],
+                    permanent_ckpt_interval=hparams['permanent_ckpt_interval'],
+                    verbose=True
+                ),
+                # LearningRateMonitor(logging_interval='step'),
+                DsTQDMProgressBar(),
+            ],
+            logger=DsTensorBoardLogger(
+                save_dir=str(work_dir),
+                name='lightning_logs',
+                version='latest'
+            ),
+            gradient_clip_val=hparams['clip_grad_norm'],
+            val_check_interval=hparams['val_check_interval'] * hparams['accumulate_grad_batches'],
+            # so this is global_steps
+            check_val_every_n_epoch=None,
+            log_every_n_steps=1,
+            max_steps=hparams['max_updates'],
+            use_distributed_sampler=False,
+            num_sanity_val_steps=hparams['num_sanity_val_steps'],
+            accumulate_grad_batches=hparams['accumulate_grad_batches']
+        )
+        if not hparams['infer']:  # train
+            @rank_zero_only
+            def train_payload_copy():
+                # Copy files to work_dir
+                binary_dir = pathlib.Path(hparams['binary_data_dir'])
+                spk_map_dst = work_dir / 'spk_map.json'
+                spk_map_src = binary_dir / 'spk_map.json'
+                shutil.copy(spk_map_src, spk_map_dst)
+                print(f'| Copied spk map to {spk_map_dst}.')
+                lang_map_dst = work_dir / 'lang_map.json'
+                lang_map_src = binary_dir / 'lang_map.json'
+                shutil.copy(lang_map_src, lang_map_dst)
+                print(f'| Copied lang map to {lang_map_dst}.')
+                for lang in hparams['dictionaries'].keys():
+                    dict_dst = work_dir / f'dictionary-{lang}.txt'
+                    dict_src = binary_dir / f'dictionary-{lang}.txt'
+                    shutil.copy(dict_src, dict_dst)
+                    print(f'| Copied dictionary for language \'{lang}\' to {dict_dst}.')
+
+            train_payload_copy()
+            trainer.fit(task, ckpt_path=get_latest_checkpoint_path(work_dir))
+        else:
+            trainer.test(task)
+
+    def on_save_checkpoint(self, checkpoint):
+        if isinstance(self.model, CategorizedModule):
+            checkpoint['category'] = self.model.category
+        checkpoint['trainer_stage'] = self.trainer.state.stage.value
+
+    def on_load_checkpoint(self, checkpoint):
+        from lightning.pytorch.trainer.states import RunningStage
+        from utils import simulate_lr_scheduler
+        if checkpoint.get('trainer_stage', '') == RunningStage.VALIDATING.value:
+            self.skip_immediate_validation = True
+
+        optimizer_args = hparams['optimizer_args']
+        scheduler_args = hparams['lr_scheduler_args']
+
+        if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args:
+            optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2'])
+
+        if checkpoint.get('optimizer_states', None):
+            opt_states = checkpoint['optimizer_states']
+            assert len(opt_states) == 1  # only support one optimizer
+            opt_state = opt_states[0]
+            for param_group in opt_state['param_groups']:
+                for k, v in optimizer_args.items():
+                    if k in param_group and param_group[k] != v:
+                        if 'lr_schedulers' in checkpoint and checkpoint['lr_schedulers'] and k == 'lr':
+                            continue
+                        rank_zero_info(f'| Overriding optimizer parameter {k} from checkpoint: {param_group[k]} -> {v}')
+                        param_group[k] = v
+                if 'initial_lr' in param_group and param_group['initial_lr'] != optimizer_args['lr']:
+                    rank_zero_info(
+                        f'| Overriding optimizer parameter initial_lr from checkpoint: {param_group["initial_lr"]} -> {optimizer_args["lr"]}'
+                    )
+                    param_group['initial_lr'] = optimizer_args['lr']
+
+        if checkpoint.get('lr_schedulers', None):
+            assert checkpoint.get('optimizer_states', False)
+            assert len(checkpoint['lr_schedulers']) == 1  # only support one scheduler
+            checkpoint['lr_schedulers'][0] = simulate_lr_scheduler(
+                optimizer_args, scheduler_args,
+                step_count=checkpoint['global_step'],
+                num_param_groups=len(checkpoint['optimizer_states'][0]['param_groups'])
+            )
+            for param_group, new_lr in zip(
+                checkpoint['optimizer_states'][0]['param_groups'],
+                checkpoint['lr_schedulers'][0]['_last_lr'],
+            ):
+                if param_group['lr'] != new_lr:
+                    rank_zero_info(f'| Overriding optimizer parameter lr from checkpoint: {param_group["lr"]} -> {new_lr}')
+                    param_group['lr'] = new_lr