diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py index f78729b71..15654cd47 100644 --- a/funasr/train_utils/trainer.py +++ b/funasr/train_utils/trainer.py @@ -208,56 +208,78 @@ class Trainer: latest = Path(os.path.join(self.output_dir, f"model.pt")) torch.save(state, latest) - if self.best_step_or_epoch == "": + if self.best_step_or_epoch == "" and ckpt_name in getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + ): self.best_step_or_epoch = ckpt_name if self.avg_keep_nbest_models_type == "acc": - if ( - self.val_acc_step_or_epoch[ckpt_name] - >= self.val_acc_step_or_epoch[self.best_step_or_epoch] - ): + cur_acc = self.val_acc_step_or_epoch.get(ckpt_name) + best_acc = self.val_acc_step_or_epoch.get(self.best_step_or_epoch) + if cur_acc is not None and (best_acc is None or cur_acc >= best_acc): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) logging.info( - f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best acc: {cur_acc:.4f}, {best_ckpt}" + ) + elif cur_acc is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation acc yet; not considered for best." ) else: logging.info( - f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in acc: {cur_acc:.4f} < {best_acc:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) elif self.avg_keep_nbest_models_type == "loss": - if ( - self.val_loss_step_or_epoch[ckpt_name] - <= self.val_loss_step_or_epoch[self.best_step_or_epoch] - ): + cur_loss = self.val_loss_step_or_epoch.get(ckpt_name) + best_loss = self.val_loss_step_or_epoch.get(self.best_step_or_epoch) + if cur_loss is not None and (best_loss is None or cur_loss <= best_loss): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) logging.info( - f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best loss: {cur_loss:.4f}, {best_ckpt}" + ) + elif cur_loss is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation loss yet; not considered for best." ) else: logging.info( - f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in loss: {cur_loss:.4f} > {best_loss:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) else: print("Undo") - self.saved_ckpts[ckpt_name] = getattr( + # Only checkpoints carrying the configured validation metric + # (acc or loss) are ranked. A checkpoint saved at an unvalidated + # step (e.g. save_checkpoint_interval is not a multiple of + # validate_interval) is kept on disk but excluded from saved_ckpts, + # so it never competes in best-model ranking or keep_nbest_models + # pruning with a fabricated score and cannot evict a validated + # best checkpoint. + metric_value = getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" - )[ckpt_name] - if self.keep_nbest_models > 0: - if len(self.saved_ckpts) > self.keep_nbest_models: - if self.avg_keep_nbest_models_type == "acc": - key = min(self.saved_ckpts, key=self.saved_ckpts.get) - else: - key = max(self.saved_ckpts, key=self.saved_ckpts.get) - if key in self.saved_ckpts: - del self.saved_ckpts[key] - filename = os.path.join(self.output_dir, key) - logging.info(f"Delete: {filename}") - if os.path.exists(filename): - os.remove(filename) + ).get(ckpt_name) + if metric_value is None: + logging.info( + f"Checkpoint {ckpt_name} has no {self.avg_keep_nbest_models_type} metric; " + "kept on disk but excluded from keep_nbest_models ranking." + ) + else: + self.saved_ckpts[ckpt_name] = metric_value + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + os.remove(filename) if self.use_ddp or self.use_fsdp: dist.barrier() diff --git a/funasr/train_utils/trainer_ds.py b/funasr/train_utils/trainer_ds.py index ff649e30e..6a4aaefba 100644 --- a/funasr/train_utils/trainer_ds.py +++ b/funasr/train_utils/trainer_ds.py @@ -234,14 +234,15 @@ class Trainer: # torch.save(state, latest) with torch.no_grad(): model.save_checkpoint(save_dir=self.output_dir, tag=f"model.pt", client_state=state) - if self.best_step_or_epoch == "": + if self.best_step_or_epoch == "" and ckpt_name in getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + ): self.best_step_or_epoch = ckpt_name if self.avg_keep_nbest_models_type == "acc": - if ( - self.val_acc_step_or_epoch[ckpt_name] - >= self.val_acc_step_or_epoch[self.best_step_or_epoch] - ): + cur_acc = self.val_acc_step_or_epoch.get(ckpt_name) + best_acc = self.val_acc_step_or_epoch.get(self.best_step_or_epoch) + if cur_acc is not None and (best_acc is None or cur_acc >= best_acc): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) # torch.save(state, best_ckpt) @@ -250,17 +251,20 @@ class Trainer: save_dir=self.output_dir, tag=f"model.pt.best", client_state=state ) logging.info( - f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best acc: {cur_acc:.4f}, {best_ckpt}" + ) + elif cur_acc is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation acc yet; not considered for best." ) else: logging.info( - f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in acc: {cur_acc:.4f} < {best_acc:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) elif self.avg_keep_nbest_models_type == "loss": - if ( - self.val_loss_step_or_epoch[ckpt_name] - <= self.val_loss_step_or_epoch[self.best_step_or_epoch] - ): + cur_loss = self.val_loss_step_or_epoch.get(ckpt_name) + best_loss = self.val_loss_step_or_epoch.get(self.best_step_or_epoch) + if cur_loss is not None and (best_loss is None or cur_loss <= best_loss): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) # torch.save(state, best_ckpt) @@ -269,31 +273,49 @@ class Trainer: save_dir=self.output_dir, tag=f"model.pt.best", client_state=state ) logging.info( - f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best loss: {cur_loss:.4f}, {best_ckpt}" + ) + elif cur_loss is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation loss yet; not considered for best." ) else: logging.info( - f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in loss: {cur_loss:.4f} > {best_loss:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) else: print("Undo") if self.rank == 0: - self.saved_ckpts[ckpt_name] = getattr( + # Only checkpoints carrying the configured validation metric + # (acc or loss) are ranked. A checkpoint saved at an unvalidated + # step (e.g. save_checkpoint_interval is not a multiple of + # validate_interval) is kept on disk but excluded from + # saved_ckpts, so it never competes in best-model ranking or + # keep_nbest_models pruning with a fabricated score and cannot + # evict a validated best checkpoint. + metric_value = getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" - )[ckpt_name] - if self.keep_nbest_models > 0: - if len(self.saved_ckpts) > self.keep_nbest_models: - if self.avg_keep_nbest_models_type == "acc": - key = min(self.saved_ckpts, key=self.saved_ckpts.get) - else: - key = max(self.saved_ckpts, key=self.saved_ckpts.get) - if key in self.saved_ckpts: - del self.saved_ckpts[key] - filename = os.path.join(self.output_dir, key) - logging.info(f"Delete: {filename}") - if os.path.exists(filename): - # os.remove(filename) - misc_utils.smart_remove(filename) + ).get(ckpt_name) + if metric_value is None: + logging.info( + f"Checkpoint {ckpt_name} has no {self.avg_keep_nbest_models_type} metric; " + "kept on disk but excluded from keep_nbest_models ranking." + ) + else: + self.saved_ckpts[ckpt_name] = metric_value + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + # os.remove(filename) + misc_utils.smart_remove(filename) elif self.use_fsdp: raise NotImplementedError( @@ -356,57 +378,79 @@ class Trainer: logging.info(f"\nCheckpoint saved to {filename}\n") latest = Path(os.path.join(self.output_dir, f"model.pt")) torch.save(state, latest) - if self.best_step_or_epoch == "": + if self.best_step_or_epoch == "" and ckpt_name in getattr( + self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" + ): self.best_step_or_epoch = ckpt_name if self.avg_keep_nbest_models_type == "acc": - if ( - self.val_acc_step_or_epoch[ckpt_name] - >= self.val_acc_step_or_epoch[self.best_step_or_epoch] - ): + cur_acc = self.val_acc_step_or_epoch.get(ckpt_name) + best_acc = self.val_acc_step_or_epoch.get(self.best_step_or_epoch) + if cur_acc is not None and (best_acc is None or cur_acc >= best_acc): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) logging.info( - f"Update best acc: {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best acc: {cur_acc:.4f}, {best_ckpt}" + ) + elif cur_acc is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation acc yet; not considered for best." ) else: logging.info( - f"No improvement in acc: {self.val_acc_step_or_epoch[ckpt_name]:.4f} < {self.val_acc_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in acc: {cur_acc:.4f} < {best_acc:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) elif self.avg_keep_nbest_models_type == "loss": - if ( - self.val_loss_step_or_epoch[ckpt_name] - <= self.val_loss_step_or_epoch[self.best_step_or_epoch] - ): + cur_loss = self.val_loss_step_or_epoch.get(ckpt_name) + best_loss = self.val_loss_step_or_epoch.get(self.best_step_or_epoch) + if cur_loss is not None and (best_loss is None or cur_loss <= best_loss): self.best_step_or_epoch = ckpt_name best_ckpt = Path(os.path.join(self.output_dir, f"model.pt.best")) torch.save(state, best_ckpt) logging.info( - f"Update best loss: {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {best_ckpt}" + f"Update best loss: {cur_loss:.4f}, {best_ckpt}" + ) + elif cur_loss is None: + logging.info( + f"Checkpoint {ckpt_name} saved at a step with no validation loss yet; not considered for best." ) else: logging.info( - f"No improvement in loss: {self.val_loss_step_or_epoch[ckpt_name]:.4f} > {self.val_loss_step_or_epoch[self.best_step_or_epoch]:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" + f"No improvement in loss: {cur_loss:.4f} > {best_loss:.4f}, {os.path.join(self.output_dir, self.best_step_or_epoch)}" ) else: print("Undo") - self.saved_ckpts[ckpt_name] = getattr( + # Only checkpoints carrying the configured validation metric + # (acc or loss) are ranked. A checkpoint saved at an unvalidated + # step (e.g. save_checkpoint_interval is not a multiple of + # validate_interval) is kept on disk but excluded from saved_ckpts, + # so it never competes in best-model ranking or keep_nbest_models + # pruning with a fabricated score and cannot evict a validated + # best checkpoint. + metric_value = getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" - )[ckpt_name] - if self.keep_nbest_models > 0: - if len(self.saved_ckpts) > self.keep_nbest_models: - if self.avg_keep_nbest_models_type == "acc": - key = min(self.saved_ckpts, key=self.saved_ckpts.get) - else: - key = max(self.saved_ckpts, key=self.saved_ckpts.get) - if key in self.saved_ckpts: - del self.saved_ckpts[key] - filename = os.path.join(self.output_dir, key) - logging.info(f"Delete: {filename}") - if os.path.exists(filename): - # os.remove(filename) - misc_utils.smart_remove(filename) + ).get(ckpt_name) + if metric_value is None: + logging.info( + f"Checkpoint {ckpt_name} has no {self.avg_keep_nbest_models_type} metric; " + "kept on disk but excluded from keep_nbest_models ranking." + ) + else: + self.saved_ckpts[ckpt_name] = metric_value + if self.keep_nbest_models > 0: + if len(self.saved_ckpts) > self.keep_nbest_models: + if self.avg_keep_nbest_models_type == "acc": + key = min(self.saved_ckpts, key=self.saved_ckpts.get) + else: + key = max(self.saved_ckpts, key=self.saved_ckpts.get) + if key in self.saved_ckpts: + del self.saved_ckpts[key] + filename = os.path.join(self.output_dir, key) + logging.info(f"Delete: {filename}") + if os.path.exists(filename): + # os.remove(filename) + misc_utils.smart_remove(filename) if self.use_ddp or self.use_fsdp: dist.barrier() diff --git a/tests/test_trainer_ds_unvalidated_ckpt.py b/tests/test_trainer_ds_unvalidated_ckpt.py new file mode 100644 index 000000000..4c08fd390 --- /dev/null +++ b/tests/test_trainer_ds_unvalidated_ckpt.py @@ -0,0 +1,147 @@ +import os + +import pytest +import torch + +from funasr.train_utils.trainer import Trainer as TrainerTorch +from funasr.train_utils.trainer_ds import Trainer as TrainerDs + + +class _DummyModule(torch.nn.Module): + """A tiny real model for the torch.save checkpoint path.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + + +class _FakeDeepSpeedEngine: + """Minimal stand-in for a DeepSpeed engine (only save_checkpoint is used). + + DeepSpeed writes one checkpoint directory per tag; this fake writes a + placeholder *file* per tag so that keep_nbest_models pruning's + smart_remove() actually deletes something observable on disk. + """ + + def __init__(self): + self.tags = [] + self.client_states = [] + + def save_checkpoint(self, save_dir=None, tag=None, client_state=None): + self.tags.append(tag) + self.client_states.append(client_state) + with open(os.path.join(save_dir, tag), "w") as f: + f.write("placeholder") + return True + + +# The same ranking logic exists in three checkpointing paths: +# trainer_ds_torch -- TrainerDs, use_deepspeed=False (torch.save path) +# trainer_ds_deepspeed -- TrainerDs, use_deepspeed=True (DeepSpeed save path) +# trainer_torch -- non-DeepSpeed Trainer (funasr-train), torch.save path +SAVE_PATHS = ["trainer_ds_torch", "trainer_ds_deepspeed", "trainer_torch"] + + +def _make_trainer(tmp_path, ranking, save_path, keep_nbest_models=1): + kwargs = dict( + output_dir=str(tmp_path), + device="cpu", + keep_nbest_models=keep_nbest_models, + avg_keep_nbest_models_type=ranking, + ) + if save_path == "trainer_torch": + return TrainerTorch(local_rank=0, **kwargs) + return TrainerDs( + rank=0, + local_rank=0, + world_size=1, + use_ddp=False, + use_fsdp=False, + use_fp16=False, + use_bf16=False, + use_deepspeed=(save_path == "trainer_ds_deepspeed"), + **kwargs, + ) + + +def _make_fixtures(save_path): + if save_path == "trainer_ds_deepspeed": + return _FakeDeepSpeedEngine(), None, None, None + model = _DummyModule() + optim = torch.optim.SGD(model.parameters(), lr=0.01) + scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=1) + return model, optim, scheduler, None + + +def _good_metric(ranking): + # a loss is good when it is low; an acc is good when it is high + return 0.5 if ranking == "loss" else 0.9 + + +def _bad_metric(ranking): + return 1.0 if ranking == "loss" else 0.5 + + +@pytest.mark.parametrize("save_path", SAVE_PATHS) +@pytest.mark.parametrize("ranking", ["loss", "acc"]) +def test_unvalidated_checkpoint_cannot_evict_validated_best(tmp_path, ranking, save_path): + """A checkpoint saved at an unvalidated step must not enter ranking. + + Regression for the KeyError fix: it used to fall back to a fabricated + score of 0.0, which under loss-based ranking pruned the validated best + checkpoint and left only the unvalidated one. + """ + trainer = _make_trainer(tmp_path, ranking, save_path, keep_nbest_models=1) + metric_dict = getattr(trainer, f"val_{ranking}_step_or_epoch") + model, optim, scheduler, scaler = _make_fixtures(save_path) + + # step 1 was validated; register its metric so it is a ranked best. + metric_dict["model.pt.ep1.1"] = _bad_metric(ranking) + trainer.save_checkpoint( + epoch=1, step=1, step_in_epoch=1, + model=model, optim=optim, scheduler=scheduler, scaler=scaler, + ) + assert trainer.saved_ckpts == {"model.pt.ep1.1": _bad_metric(ranking)} + assert trainer.best_step_or_epoch == "model.pt.ep1.1" + + # step 2 was never validated (e.g. save_checkpoint_interval is not a + # multiple of validate_interval) -> save_checkpoint used to KeyError here. + trainer.save_checkpoint( + epoch=1, step=2, step_in_epoch=2, + model=model, optim=optim, scheduler=scheduler, scaler=scaler, + ) + + # The unvalidated checkpoint is kept on disk but excluded from ranking: + # no synthetic score, no best change, and the validated best is never pruned. + # Both the validated best and the unvalidated checkpoint files survive. + assert trainer.saved_ckpts == {"model.pt.ep1.1": _bad_metric(ranking)} + assert trainer.best_step_or_epoch == "model.pt.ep1.1" + assert os.path.exists(os.path.join(tmp_path, "model.pt.ep1.1")) + assert os.path.exists(os.path.join(tmp_path, "model.pt.ep1.2")) + + +@pytest.mark.parametrize("save_path", SAVE_PATHS) +@pytest.mark.parametrize("ranking", ["loss", "acc"]) +def test_validated_checkpoints_still_rank_and_prune(tmp_path, ranking, save_path): + """Validated checkpoints must keep ranking and keep_nbest_models pruning.""" + trainer = _make_trainer(tmp_path, ranking, save_path, keep_nbest_models=1) + metric_dict = getattr(trainer, f"val_{ranking}_step_or_epoch") + model, optim, scheduler, scaler = _make_fixtures(save_path) + + metric_dict["model.pt.ep1.1"] = _bad_metric(ranking) + metric_dict["model.pt.ep1.2"] = _good_metric(ranking) + + trainer.save_checkpoint( + epoch=1, step=1, step_in_epoch=1, + model=model, optim=optim, scheduler=scheduler, scaler=scaler, + ) + trainer.save_checkpoint( + epoch=1, step=2, step_in_epoch=2, + model=model, optim=optim, scheduler=scheduler, scaler=scaler, + ) + + assert trainer.saved_ckpts == {"model.pt.ep1.2": _good_metric(ranking)} + assert trainer.best_step_or_epoch == "model.pt.ep1.2" + # the worse validated checkpoint was pruned, the better one remains + assert not os.path.exists(os.path.join(tmp_path, "model.pt.ep1.1")) + assert os.path.exists(os.path.join(tmp_path, "model.pt.ep1.2"))