mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
fix(trainer): avoid KeyError when a checkpoint is saved at an unvalidated step (#3455)
* fix(trainer): avoid KeyError when a checkpoint is saved at an unvalidated step When save_checkpoint_interval is not a multiple of validate_interval, a checkpoint can be saved at a step that was never validated. The best-ckpt comparison then indexes val_acc/val_loss_step_or_epoch[ckpt_name], raising KeyError right after the checkpoint was already written and aborting the run. Guard both the DeepSpeed and the torch.save checkpoint paths with .get() so an unvalidated checkpoint is skipped for the "best" bookkeeping (and logged) instead of crashing; fall back to 0.0 when recording saved_ckpts. Co-Authored-By: Claude <[email protected]> * fix(trainer): exclude unvalidated-step checkpoints from best-model ranking Address review feedback on PR #3455. A checkpoint saved at a step that was never validated (save_checkpoint_interval not a multiple of validate_interval) used to be inserted into `saved_ckpts` with a fabricated score of 0.0. Under loss-based ranking that synthetic "perfect" loss could evict a real validated best checkpoint via keep_nbest_models pruning. Now such checkpoints are kept on disk but excluded from `saved_ckpts` entirely, so they never compete in best-model ranking or pruning, and `best_step_or_epoch` only ever references validated checkpoints. Adds regression coverage for both loss and acc ranking on the DeepSpeed and torch.save paths, asserting an unvalidated checkpoint cannot evict a validated best checkpoint. * test(trainer): make unvalidated-ckpt regression catch the DeepSpeed+acc case The fake DeepSpeed engine now writes a placeholder file per tag (mirroring DeepSpeed's per-tag checkpoint dirs), so keep_nbest_models pruning's smart_remove() is observable. This closes a gap where the [DeepSpeed, acc] combination passed on the buggy code: under acc ranking the fabricated 0.0 is the minimum, so pruning deletes the unvalidated checkpoint itself rather than the validated best, and without an on-disk artifact that deletion was invisible. With the fix, all 8 cases pass; on the previous code all 4 unvalidated-step cases now fail (loss: fabricated 0.0 evicts the validated best; acc: the unvalidated checkpoint file is pruned). * fix(trainer): apply the same unvalidated-step guard to the non-DeepSpeed Trainer `funasr-train` (funasr/train_utils/trainer.py) shares the identical checkpoint-ranking code as trainer_ds.py and had the same bug: direct indexing of `val_{acc,loss}_step_or_epoch[ckpt_name]` raised KeyError when a checkpoint was saved at a step that was never validated, and the keep_nbest_models bookkeeping inserted a fabricated score for such checkpoints. Apply the same fix: checkpoints saved at an unvalidated step are kept on disk but excluded from saved_ckpts (no synthetic score), best_step_or_epoch only ever references validated checkpoints, and the metric dicts are read via .get(). Tests now parametrize over all three checkpointing paths (trainer_ds torch.save, trainer_ds DeepSpeed, and trainer torch.save) x loss/acc. On upstream main all 6 unvalidated-step cases fail with `KeyError: 'model.pt.ep1.2'`; with the fix all 12 cases pass. --------- Co-authored-by: Claude <[email protected]>
This commit is contained in:
张思远 (ZhangSiYuan)
co-authored by
Claude
Claude
parent
cb0eb3bb40
commit
036ed6936b
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"))
|
||||
Reference in New Issue
Block a user