diff --git a/.github/workflows/test-checkpoint-selection.yml b/.github/workflows/test-checkpoint-selection.yml new file mode 100644 index 0000000000..1906cbf310 --- /dev/null +++ b/.github/workflows/test-checkpoint-selection.yml @@ -0,0 +1,65 @@ +name: Validate checkpoint selection + +on: + pull_request: + paths: + - "funasr/train_utils/**" + - "tests/test_checkpoint_metric_*.py" + - "tests/test_trainer_ds_unvalidated_ckpt.py" + - "tests/test_train_ds_distributed_config.py" + - "tests/test_sensevoice_continual_finetuning_docs.py" + - "examples/industrial_data_pretraining/sense_voice/**" + - ".github/workflows/test-checkpoint-selection.yml" + push: + branches: [main] + paths: + - "funasr/train_utils/**" + - "tests/test_checkpoint_metric_*.py" + - "tests/test_trainer_ds_unvalidated_ckpt.py" + - "tests/test_train_ds_distributed_config.py" + - "tests/test_sensevoice_continual_finetuning_docs.py" + - "examples/industrial_data_pretraining/sense_voice/**" + - ".github/workflows/test-checkpoint-selection.yml" + +permissions: + contents: read + +jobs: + checkpoint-selection: + runs-on: ubuntu-latest + timeout-minutes: 15 + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" + TORCH_FORCE_WEIGHTS_ONLY_LOAD: "1" + HF_HUB_OFFLINE: "1" + GLOO_SOCKET_IFNAME: lo + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + - name: Install CPU test dependencies + run: | + python -m pip install torch==2.10.0 torchaudio==2.10.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install "numpy<2" kaldiio librosa rapidfuzz six pytest omegaconf hydra-core pyyaml tqdm + - name: Check recipe syntax + run: bash -n examples/industrial_data_pretraining/sense_voice/finetune.sh + - name: Test validation, checkpoint lifecycle and two-rank collectives + run: | + python -m pytest -q \ + tests/test_checkpoint_metric_presence.py \ + tests/test_checkpoint_metric_distributed.py \ + tests/test_trainer_ds_unvalidated_ckpt.py \ + tests/test_train_ds_distributed_config.py \ + tests/test_sensevoice_continual_finetuning_docs.py \ + --junitxml=checkpoint-selection.xml + - name: Save test report + if: always() + uses: actions/upload-artifact@v4 + with: + name: checkpoint-selection-report + path: checkpoint-selection.xml + if-no-files-found: error diff --git a/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING.md b/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING.md index 86679756ab..6ca898eea5 100644 --- a/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING.md +++ b/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING.md @@ -24,16 +24,19 @@ relative regression = (candidate CER - baseline CER) / baseline CER The threshold is a product decision, not a FunASR guarantee. Select a checkpoint only when it meets every retention threshold and the new-domain target. -A real starting point from [issue #3388](https://github.com/modelscope/FunASR/issues/3388) -is: - -| Validation set | Utterances | Baseline CER | -| --- | ---: | ---: | -| Mandarin FLEURS test | 300 | 8.9% | -| Cantonese ASCEND dialogue | 314 | 9.1% | -| Teochew held-out, about 20 speakers in the source corpus | 300 | 93.3% | - -Keep the Teochew held-out speakers out of the 8.45-hour training set. +References must be independently audited for transcription policy and provenance. +Check whether new-domain labels are verbatim transcripts or written-language +paraphrases, and whether test references were generated by the same process as +training labels. Speaker separation alone does not make those references reliable. +An unaudited new-domain CER must not select checkpoints or tune replay ratios. + +The reporter in [issue #3388](https://github.com/modelscope/FunASR/issues/3388#issuecomment-5578763142) +later withdrew the earlier new-domain evaluation and rebuilt an independent test +set. The old baseline table is therefore not a valid selection reference. Recompute +the frozen baseline and all candidates with the same reviewed references and CER +normalization. For Mandarin and Cantonese, also separate read and spontaneous +speech when both matter in production. Keep held-out speakers out of the 8.45-hour +training example below; its replay budgets are illustrative, not validated optima. ## 2. Build a replay manifest @@ -72,15 +75,29 @@ pool totals disagree. ## 3. Use a supported language token -SenseVoiceSmall currently has a closed language-id set: `<|zh|>`, `<|en|>`, -`<|yue|>`, `<|ja|>`, and `<|ko|>`. For Teochew transcripts written as -standard Chinese, use `<|zh|>` for the first experiment. +The current model's named inference language mapping includes `zh`, `en`, `yue`, +`ja`, and `ko` (plus automatic/no-speech handling). This is not the complete +tokenizer vocabulary. Sharing `<|zh|>` with a new domain is an experiment, not an +isolation mechanism: it can interfere with Mandarin and must pass retention gates. + +The #3388 reporter found a useful `<|minnan|>` comparison in their setup. A token +existing in a vocabulary does not mean it is unused, freely repurposable, or +accepted by the stock inference language mapping. Pin the tokenizer/model revision +and inspect the token id, training target and query embedding path, and inference +mapping before adopting a language-tag change. Do not generalize that experiment +into a guarantee for another language or checkpoint. Do not add `<|teochew|>` to the JSONL and assume it becomes a new language id. The current tokenizer does not encode it as one supported special token. A real custom language id requires coordinated tokenizer, vocabulary, language embedding, checkpoint migration, export, and inference changes. +Check downstream language validation and every allowlist as well. Test an actual +result through export, service parsing and the final client: a valid model result +must not be silently discarded because its language tag is unfamiliar. Preserve +raw model output while diagnosing an empty response; do not simply disable all +validation or relabel the result to pass a closed list. + ## 4. Train in two stages Copy `finetune.sh` to an experiment directory, point it at the mixed training @@ -129,6 +146,31 @@ parameters trainable because no parameter name matches it. ## 5. Select the checkpoint by constraints +The maintained `finetune.sh` explicitly uses: + +```bash +++train_conf.avg_keep_nbest_models_type=loss +``` + +SenseVoice returns ASR CTC loss plus a rich-tag objective and `acc_rich`, not the +trainer's ASR `acc`. `acc_rich` is not CER or ASR accuracy. A missing metric must not +be treated as a real zero; genuine `acc=0` remains a valid score for models that +actually report it. The trainers exclude an unavailable/nonfinite selected metric +from best selection, keep-N ranking and metric-based averaging, with an explicit +warning. This can leave unranked checkpoints on disk; correct the criterion/data +instead of assuming the keep-N limit applies to them. + +Use a new output directory when changing from accuracy to loss ranking. Resuming +with a different criterion is rejected; old zero-valued histories cannot reliably +be identified as missing after the fact, and pruned weights cannot be recovered. +Earlier runs affected by this defect need independent reevaluation of retained +checkpoints, not just an upgrade. Revalidation cannot overwrite the current best +under the same checkpoint name if its selected metric becomes unavailable. + +`model.pt.best` now means best **total validation loss** for this recipe, not the +lowest CER or best retained-language behavior. Weight averaging is also only a +candidate: reevaluate the averaged model under all the same gates. + Evaluate every candidate checkpoint on Mandarin, Cantonese, and Teochew separately. @@ -156,6 +198,8 @@ Verify the trainable parameter list before relying on any adapter experiment. - Baseline CER/WER and acceptable relative regression are recorded per language. - Pseudo-labels are filtered and manually sampled for audit. - Replay is sampled to an explicit per-epoch budget with a Cantonese floor. -- Teochew uses the supported `<|zh|>` path unless the model is deliberately forked. +- References and normalization are independently audited before any CER comparison. +- Language-token mapping and downstream allowlists pass an end-to-end result check. +- The ranking metric is explicit; changed criteria use a new output directory. - Freeze logs and trainable parameters match the intended stage. - Every saved checkpoint is evaluated on all retained and new domains. diff --git a/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING_zh.md b/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING_zh.md index 6264d73602..5cf26b0587 100644 --- a/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING_zh.md +++ b/examples/industrial_data_pretraining/sense_voice/CONTINUAL_FINETUNING_zh.md @@ -20,15 +20,15 @@ replay 数据之间必须说话人不重叠。不要用一个混合 CER 选择 c 阈值应由业务要求决定,不是 FunASR 的保证值。候选 checkpoint 必须同时满足所有旧语种 阈值和新域目标。 -[issue #3388](https://github.com/modelscope/FunASR/issues/3388) 提供了一个真实起点: +参考文本的标注规范和来源必须经过独立核验:新域标签是逐字转写还是书面语意译?测试 +参考是否与训练标签来自同一套生成流程?仅做到说话人不重叠,并不能证明标注可靠。 +未经核验的新域 CER 不能用来选 checkpoint,也不能用来调整 replay 比例。 -| 验证集 | 句数 | 基线 CER | -| --- | ---: | ---: | -| 普通话 FLEURS test | 300 | 8.9% | -| 粤语 ASCEND 口语对话 | 314 | 9.1% | -| 潮汕话 held-out(源语料约 20 名说话人) | 300 | 93.3% | - -潮汕话 held-out 中的说话人不能出现在 8.45 小时训练数据中。 +[issue #3388 的后续反馈](https://github.com/modelscope/FunASR/issues/3388#issuecomment-5578763142) +已撤回早期新域评测,改用另建的独立测试集。因此旧基线表不能继续作为选点依据。 +应使用相同、经过人工复核的参考文本和 CER 归一化规则,重新评测原模型及所有候选。 +线上同时包含朗读和自发口语时,普通话、粤语也应分别检查这两类数据。下文 8.45 小时 +训练示例仍须排除 held-out 说话人;其 replay 小时数是预算演示,不是已验证的最优值。 ## 2. 构建 replay 清单 @@ -61,13 +61,23 @@ replay 数据之间必须说话人不重叠。不要用一个混合 CER 选择 c ## 3. 使用现有支持的语言 token -SenseVoiceSmall 当前语言 id 是闭集:`<|zh|>`、`<|en|>`、`<|yue|>`、 -`<|ja|>` 和 `<|ko|>`。潮汕话转写目标为标准中文时,第一轮实验使用 `<|zh|>`。 +当前模型具名推理语言映射包含 `zh`、`en`、`yue`、`ja`、`ko`,以及自动判断和无语音 +处理;这不是 tokenizer 完整词表。让新域共享 `<|zh|>` 只是实验方案,不是隔离机制, +可能干扰普通话,必须通过旧域保持评测。 + +#3388 报告者在自己的设置中获得了有价值的 `<|minnan|>` 对照结果,但“词表中存在” +不等于“空闲、可任意复用”,也不代表原版推理语言映射支持它。采用语言 tag 改动前, +固定 tokenizer/模型版本,核对 token id、训练目标和查询 embedding 路径以及推理映射。 +不能把该个例写成其他语种或 checkpoint 的保证。 不要直接在 JSONL 中添加 `<|teochew|>` 并假设它会成为新语言 id。当前 tokenizer 不会把它编码成一个受支持的特殊 token。真正新增自定义语言 id,需要同步修改 tokenizer、 词表、language embedding、旧权重迁移、导出和推理代码。 +还要检查下游语言校验与每一处白名单,用真实结果走通导出、服务端解析和最终客户端。 +不能因为语言 tag 不在旧枚举中,就静默丢弃有效的模型输出。排查空响应时保留原始模型 +结果,不要简单关闭全部校验或改写语言标签来绕过白名单。 + ## 4. 分两阶段训练 把 `finetune.sh` 复制到实验目录,指向混合训练 manifest 和固定验证集,然后把默认的 @@ -110,6 +120,26 @@ Setting encoder...requires_grad = False ## 5. 按约束选择 checkpoint +维护中的 `finetune.sh` 已明确指定: + +```bash +++train_conf.avg_keep_nbest_models_type=loss +``` + +SenseVoice 返回 ASR CTC loss 与 rich-tag 目标,以及 `acc_rich`,并不返回训练器所选的 +ASR `acc`。`acc_rich` 不是 CER 或 ASR 准确率。缺失指标不能伪装成真实零值;对确实 +返回准确率的模型,`acc=0` 仍是有效值。所选指标缺失、非有限值或只覆盖部分 batch/rank +时,训练器会明确告警,并将该点排除在 best、keep-N 排名和按指标平均之外。未参与排名 +的 checkpoint 可能继续留在磁盘上,应修正指标或数据,不要假设 keep-N 会删除它们。 + +从 accuracy 改为 loss 排名时,必须使用新的输出目录;恢复时混用两套指标会被拒绝。 +历史上的零值无法可靠追溯为“缺失”,已裁剪的权重也无法恢复。受影响的旧训练需要对 +仍保留的 checkpoint 独立复测,而不是只升级软件。同名重验若导致当前 best 的所选指标 +失效,也会被拒绝,避免覆盖现有 best。 + +这个配方的 `model.pt.best` 表示**总验证 loss 最优**,不是 CER 最低或旧域保持最好。 +权重平均结果也只是候选模型,仍须重新通过全部评测门槛。 + 每个候选 checkpoint 都要分别评估普通话、粤语和潮汕话。 1. 丢弃任何一个旧语种超过退化阈值的 checkpoint。 @@ -133,6 +163,8 @@ SenseVoice adapter;任何 adapter 实验都必须先核对实际可训练参 - 已记录每个语种的基线 CER/WER 和可接受相对退化。 - 已过滤伪标签,并按语种和说话人做人工抽检。 - replay 使用明确的每轮预算,并设置粤语最低占比。 -- 未 fork 模型时,潮汕话使用当前支持的 `<|zh|>` 路径。 +- 参考文本和归一化规则已经独立核验,再做 CER 对比。 +- 语言 token 映射和下游白名单已经完成真实结果的端到端检查。 +- 已明确排名指标;更换指标时使用新的输出目录。 - 冻结日志和可训练参数与当前阶段一致。 - 所有保存的 checkpoint 都分别评估旧语种和新域。 diff --git a/examples/industrial_data_pretraining/sense_voice/finetune.sh b/examples/industrial_data_pretraining/sense_voice/finetune.sh index 081b77b016..5867d93b70 100644 --- a/examples/industrial_data_pretraining/sense_voice/finetune.sh +++ b/examples/industrial_data_pretraining/sense_voice/finetune.sh @@ -63,6 +63,7 @@ ${train_tool} \ ++train_conf.save_checkpoint_interval=2000 \ ++train_conf.keep_nbest_models=20 \ ++train_conf.avg_nbest_model=10 \ +++train_conf.avg_keep_nbest_models_type=loss \ ++train_conf.use_deepspeed=false \ ++train_conf.deepspeed_config=${deepspeed_config} \ ++optim_conf.lr=0.0002 \ diff --git a/funasr/train_utils/average_nbest_models.py b/funasr/train_utils/average_nbest_models.py index 873f4192c1..5517903db3 100644 --- a/funasr/train_utils/average_nbest_models.py +++ b/funasr/train_utils/average_nbest_models.py @@ -14,6 +14,7 @@ import re from collections import OrderedDict from functools import cmp_to_key +from funasr.train_utils.checkpoint_metrics import finite_scalar def _get_checkpoint_paths(output_dir: str, last_n: int = 5, use_deepspeed=False, **kwargs): @@ -31,7 +32,19 @@ def _get_checkpoint_paths(output_dir: str, last_n: int = 5, use_deepspeed=False, ) avg_keep_nbest_models_type = checkpoint["avg_keep_nbest_models_type"] val_step_or_epoch = checkpoint[f"val_{avg_keep_nbest_models_type}_step_or_epoch"] - sorted_items = sorted(val_step_or_epoch.items(), key=lambda x: x[1], reverse=True) + eligible = [] + for key, value in val_step_or_epoch.items(): + value = finite_scalar(value) + ckpt = os.path.join(output_dir, key) + if use_deepspeed: + ckpt = os.path.join(ckpt, "mp_rank_00_model_states.pt") + if value is not None and os.path.isfile(ckpt): + eligible.append((key, value)) + if not eligible: + logging.warning("No checkpoints with an available finite %s metric for averaging.", + avg_keep_nbest_models_type) + return [] + sorted_items = sorted(eligible, key=lambda x: x[1], reverse=True) sorted_items = ( sorted_items[:last_n] if avg_keep_nbest_models_type == "acc" else sorted_items[-last_n:] ) @@ -43,8 +56,8 @@ def _get_checkpoint_paths(output_dir: str, last_n: int = 5, use_deepspeed=False, ckpt = os.path.join(output_dir, key, "mp_rank_00_model_states.pt") checkpoint_paths.append(ckpt) - except: - print(f"{checkpoint} does not exist, avg the lastet checkpoint.") + except (FileNotFoundError, KeyError): + logging.warning("Checkpoint ranking metadata unavailable; averaging latest checkpoints.") # List all files in the output directory files = os.listdir(output_dir) # Filter out checkpoint files and extract epoch numbers diff --git a/funasr/train_utils/checkpoint_metrics.py b/funasr/train_utils/checkpoint_metrics.py new file mode 100644 index 0000000000..a567e93a2f --- /dev/null +++ b/funasr/train_utils/checkpoint_metrics.py @@ -0,0 +1,99 @@ +"""Validation metric availability shared by the torch and DeepSpeed trainers.""" + +import logging +import math + +import torch +import torch.distributed as dist + + +def finite_scalar(value): + if isinstance(value, torch.Tensor): + if value.numel() != 1: + return None + value = value.detach().item() + try: + value = float(value) + except (TypeError, ValueError): + return None + return value if math.isfinite(value) else None + + +class ValidationMetrics: + def __init__(self): + self.count = 0 + self.sums = {"loss": 0.0, "acc": 0.0} + self.complete = {"loss": True, "acc": True} + + def update(self, loss, stats): + loss = finite_scalar(loss) + values = {"loss": loss, "acc": finite_scalar(stats.get("acc"))} + self.count += 1 + for name, value in values.items(): + if value is None or loss is None: + self.complete[name] = False + else: + self.sums[name] += value + + def compute(self, device, distributed=False): + names = ("loss", "acc") + means = [self.sums[name] / max(self.count, 1) for name in names] + available = [self.count > 0 and self.complete[name] for name in names] + if distributed: + # All ranks enter both collectives, including empty/nonfinite ranks. + values = torch.tensor(means, dtype=torch.float32, device=device) + flags = torch.tensor(available, dtype=torch.int32, device=device) + dist.all_reduce(values, op=dist.ReduceOp.SUM) + dist.all_reduce(flags, op=dist.ReduceOp.MIN) + means = (values / dist.get_world_size()).cpu().tolist() + available = flags.cpu().tolist() + return { + name: finite_scalar(value) if valid else None + for name, value, valid in zip(names, means, available) + } + + +def record_validation_metrics(trainer, ckpt_name, metrics): + selected = trainer.avg_keep_nbest_models_type + reject = metrics[selected] is None and trainer.best_step_or_epoch == ckpt_name + distributed = trainer.use_ddp or trainer.use_fsdp or getattr(trainer, "use_deepspeed", False) + if metrics[selected] is None and distributed: + # Only rank zero may have updated best during checkpoint saving. + flag = torch.tensor(int(reject), dtype=torch.int32, device=trainer.device) + dist.all_reduce(flag, op=dist.ReduceOp.MAX) + reject = bool(flag.item()) + if reject: + raise ValueError( + f"Cannot overwrite current best checkpoint {ckpt_name} with unavailable " + f"validation {selected}; use a new checkpoint name or output directory." + ) + for name, value in metrics.items(): + setattr(trainer, f"val_{name}_avg", value if value is not None else float("nan")) + history = getattr(trainer, f"val_{name}_step_or_epoch") + if value is None: + history.pop(ckpt_name, None) + else: + history[ckpt_name] = value + if metrics[selected] is None: + trainer.saved_ckpts.pop(ckpt_name, None) + if trainer.rank == 0 and not trainer._warned_validation_metric: + remedy = ( + " For models without ASR accuracy, start a new output directory with " + "++train_conf.avg_keep_nbest_models_type=loss; acc_rich is not ASR accuracy." + if selected == "acc" else " Check the validation data and model outputs." + ) + logging.warning( + "Validation has no complete, finite %s metric on every batch/rank; " + "this checkpoint is excluded from best, pruning and averaging.%s", + selected, remedy, + ) + trainer._warned_validation_metric = True + + +def check_resume_ranking(checkpoint, configured): + previous = checkpoint.get("avg_keep_nbest_models_type") + if previous is not None and previous != configured: + raise ValueError( + f"Cannot change checkpoint ranking metric from {previous} to {configured} " + "while resuming. Use a new output directory to avoid mixing metric histories." + ) diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py index b0493c03f3..38acf3fb75 100644 --- a/funasr/train_utils/trainer.py +++ b/funasr/train_utils/trainer.py @@ -13,6 +13,9 @@ from funasr.train_utils.device_funcs import to_device from funasr.train_utils.recursive_op import recursive_average from funasr.train_utils.average_nbest_models import average_checkpoints +from funasr.train_utils.checkpoint_metrics import ( + ValidationMetrics, check_resume_ranking, record_validation_metrics, +) from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler try: @@ -127,6 +130,7 @@ def __init__( self.best_step_or_epoch = "" self.val_acc_step_or_epoch = {} self.val_loss_step_or_epoch = {} + self._warned_validation_metric = False self.reset_gpu_cache = kwargs.get("reset_gpu_cache", False) self.start_data_split_i = 0 @@ -202,11 +206,9 @@ def save_checkpoint( else: ckpt_name = f"model.pt.ep{epoch}.{step}" filename = os.path.join(self.output_dir, ckpt_name) - torch.save(state, filename) - logging.info(f"Checkpoint saved to {filename}") - latest = Path(os.path.join(self.output_dir, f"model.pt")) - torch.save(state, latest) + write_best = False + prune_path = None if self.best_step_or_epoch == "" and ckpt_name in getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" @@ -219,7 +221,7 @@ def save_checkpoint( 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) + write_best = True logging.info( f"Update best acc: {cur_acc:.4f}, {best_ckpt}" ) @@ -237,7 +239,7 @@ def save_checkpoint( 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) + write_best = True logging.info( f"Update best loss: {cur_loss:.4f}, {best_ckpt}" ) @@ -276,10 +278,19 @@ def save_checkpoint( 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) + prune_path = os.path.join(self.output_dir, key) + + state["best_step_or_epoch"] = self.best_step_or_epoch + state["saved_ckpts"] = dict(self.saved_ckpts) + torch.save(state, filename) + torch.save(state, latest) + if write_best: + torch.save(state, best_ckpt) + logging.info(f"Checkpoint saved to {filename}") + # Do not delete a previous candidate before all new writes succeed. + if prune_path is not None and os.path.exists(prune_path): + logging.info(f"Delete: {prune_path}") + os.remove(prune_path) if self.use_ddp or self.use_fsdp: dist.barrier() @@ -302,6 +313,7 @@ def resume_checkpoint( ckpt = os.path.join(self.output_dir, "model.pt") if os.path.isfile(ckpt): checkpoint = torch.load(ckpt, map_location="cpu") + check_resume_ranking(checkpoint, self.avg_keep_nbest_models_type) self.start_epoch = checkpoint["epoch"] # self.model.load_state_dict(checkpoint['state_dict']) src_state = checkpoint["state_dict"] @@ -580,6 +592,7 @@ def validate_epoch( dist.barrier() logging.info(f"Validate epoch: {epoch}, rank: {self.rank}\n") model.eval() + metrics = ValidationMetrics() with torch.no_grad(): @@ -620,27 +633,10 @@ def validate_epoch( loss = loss time4 = time.perf_counter() - if torch.isfinite(loss): - self.val_loss_avg = ( - self.val_loss_avg * batch_idx + loss.detach().cpu().item() - ) / (batch_idx + 1) - - if "acc" in stats: - self.val_acc_avg = ( - self.val_acc_avg * batch_idx + stats["acc"].detach().cpu().item() - ) / (batch_idx + 1) - - if self.use_ddp or self.use_fsdp: - val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to( - self.device - ) - val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to( - self.device - ) - dist.all_reduce(val_loss_avg, op=dist.ReduceOp.SUM) - dist.all_reduce(val_acc_avg, op=dist.ReduceOp.SUM) - self.val_loss_avg = val_loss_avg.detach().cpu().item() / self.world_size - self.val_acc_avg = val_acc_avg.detach().cpu().item() / self.world_size + metrics.update(loss, stats) + running = metrics.compute(self.device) + self.val_loss_avg = running["loss"] if running["loss"] is not None else float("nan") + self.val_acc_avg = running["acc"] if running["acc"] is not None else float("nan") time5 = time.perf_counter() batch_num_epoch = 1 @@ -667,8 +663,10 @@ def validate_epoch( ckpt_name = f"model.pt.ep{epoch}" else: ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step_in_epoch")}' - self.val_acc_step_or_epoch[ckpt_name] = self.val_acc_avg - self.val_loss_step_or_epoch[ckpt_name] = self.val_loss_avg + record_validation_metrics( + self, ckpt_name, + metrics.compute(self.device, distributed=self.use_ddp or self.use_fsdp), + ) model.train() if self.use_ddp or self.use_fsdp: diff --git a/funasr/train_utils/trainer_ds.py b/funasr/train_utils/trainer_ds.py index 9bd7ca31ea..fc2f2b660f 100644 --- a/funasr/train_utils/trainer_ds.py +++ b/funasr/train_utils/trainer_ds.py @@ -14,6 +14,9 @@ from funasr.train_utils.device_funcs import to_device from funasr.train_utils.recursive_op import recursive_average from funasr.train_utils.average_nbest_models import average_checkpoints +from funasr.train_utils.checkpoint_metrics import ( + ValidationMetrics, check_resume_ranking, record_validation_metrics, +) from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler import funasr.utils.misc as misc_utils @@ -129,6 +132,7 @@ def __init__( self.best_step_or_epoch = "" self.val_acc_step_or_epoch = {} self.val_loss_step_or_epoch = {} + self._warned_validation_metric = False self.reset_gpu_cache = kwargs.get("reset_gpu_cache", False) self.start_data_split_i = 0 @@ -226,14 +230,8 @@ def save_checkpoint( ckpt_name = f"model.pt.ep{epoch}.{step}" filename = os.path.join(self.output_dir, ckpt_name) - # torch.save(state, filename) - with torch.no_grad(): - model.save_checkpoint(save_dir=self.output_dir, tag=ckpt_name, client_state=state) - logging.info(f"\nCheckpoint saved to {filename}\n") - latest = Path(os.path.join(self.output_dir, f"model.pt")) - # torch.save(state, latest) - with torch.no_grad(): - model.save_checkpoint(save_dir=self.output_dir, tag=f"model.pt", client_state=state) + write_best = False + prune_path = None if self.best_step_or_epoch == "" and ckpt_name in getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" ): @@ -245,11 +243,7 @@ def save_checkpoint( 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) - with torch.no_grad(): - model.save_checkpoint( - save_dir=self.output_dir, tag=f"model.pt.best", client_state=state - ) + write_best = True logging.info( f"Update best acc: {cur_acc:.4f}, {best_ckpt}" ) @@ -267,11 +261,7 @@ def save_checkpoint( 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) - with torch.no_grad(): - model.save_checkpoint( - save_dir=self.output_dir, tag=f"model.pt.best", client_state=state - ) + write_best = True logging.info( f"Update best loss: {cur_loss:.4f}, {best_ckpt}" ) @@ -311,11 +301,24 @@ def save_checkpoint( 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) + prune_path = os.path.join(self.output_dir, key) + + if dist.is_initialized(): + # Ranking is owned by rank zero; engine writes remain collective. + decision = [(self.best_step_or_epoch, self.saved_ckpts, write_best)] + dist.broadcast_object_list(decision, src=0) + self.best_step_or_epoch, self.saved_ckpts, write_best = decision[0] + state["best_step_or_epoch"] = self.best_step_or_epoch + state["saved_ckpts"] = dict(self.saved_ckpts) + with torch.no_grad(): + model.save_checkpoint(save_dir=self.output_dir, tag=ckpt_name, client_state=state) + model.save_checkpoint(save_dir=self.output_dir, tag="model.pt", client_state=state) + if write_best: + model.save_checkpoint(save_dir=self.output_dir, tag="model.pt.best", client_state=state) + logging.info(f"Checkpoint saved to {filename}") + if self.rank == 0 and prune_path is not None and os.path.exists(prune_path): + logging.info(f"Delete: {prune_path}") + misc_utils.smart_remove(prune_path) elif self.use_fsdp: raise NotImplementedError( @@ -373,11 +376,9 @@ def save_checkpoint( else: ckpt_name = f"model.pt.ep{epoch}.{step}" filename = os.path.join(self.output_dir, ckpt_name) - torch.save(state, filename) - - logging.info(f"\nCheckpoint saved to {filename}\n") latest = Path(os.path.join(self.output_dir, f"model.pt")) - torch.save(state, latest) + write_best = False + prune_path = None if self.best_step_or_epoch == "" and ckpt_name in getattr( self, f"val_{self.avg_keep_nbest_models_type}_step_or_epoch" ): @@ -389,7 +390,7 @@ def save_checkpoint( 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) + write_best = True logging.info( f"Update best acc: {cur_acc:.4f}, {best_ckpt}" ) @@ -407,7 +408,7 @@ def save_checkpoint( 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) + write_best = True logging.info( f"Update best loss: {cur_loss:.4f}, {best_ckpt}" ) @@ -446,11 +447,18 @@ def save_checkpoint( 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) + prune_path = os.path.join(self.output_dir, key) + + state["best_step_or_epoch"] = self.best_step_or_epoch + state["saved_ckpts"] = dict(self.saved_ckpts) + torch.save(state, filename) + torch.save(state, latest) + if write_best: + torch.save(state, best_ckpt) + logging.info(f"Checkpoint saved to {filename}") + if prune_path is not None and os.path.exists(prune_path): + logging.info(f"Delete: {prune_path}") + misc_utils.smart_remove(prune_path) if self.use_ddp or self.use_fsdp: dist.barrier() @@ -475,6 +483,7 @@ def resume_checkpoint( ckpt = os.path.join(self.output_dir, "model.pt") if os.path.exists(ckpt): _, checkpoint = model.load_checkpoint(self.output_dir, "model.pt") + check_resume_ranking(checkpoint, self.avg_keep_nbest_models_type) self.start_epoch = checkpoint["epoch"] self.saved_ckpts = checkpoint["saved_ckpts"] self.val_acc_step_or_epoch = ( @@ -520,6 +529,7 @@ def resume_checkpoint( ckpt = os.path.join(self.output_dir, "model.pt") if os.path.isfile(ckpt): checkpoint = torch.load(ckpt, map_location="cpu") + check_resume_ranking(checkpoint, self.avg_keep_nbest_models_type) self.start_epoch = checkpoint["epoch"] # self.model.load_state_dict(checkpoint['state_dict']) src_state = checkpoint["state_dict"] @@ -812,8 +822,7 @@ def validate_epoch( Args: epoch (int): The current epoch number. """ - self.val_loss_avg = 0.0 - self.val_acc_avg = 0.0 + metrics = ValidationMetrics() if self.use_ddp or self.use_fsdp or self.use_deepspeed: dist.barrier() @@ -860,31 +869,23 @@ def validate_epoch( loss_dict["batch_num_epoch"] = len(dataloader_val) + metrics.update(loss_dict["loss"], loss_dict["stats"]) + running = metrics.compute(self.device) + self.val_loss_avg = running["loss"] if running["loss"] is not None else float("nan") + self.val_acc_avg = running["acc"] if running["acc"] is not None else float("nan") self.log(loss_dict, tag="val") time_beg = time.perf_counter() - self.val_loss_avg = ( - self.val_loss_avg * batch_idx + loss_dict["loss"].detach().cpu().item() - ) / (batch_idx + 1) - if "acc" in loss_dict["stats"]: - self.val_acc_avg = ( - self.val_acc_avg * batch_idx - + loss_dict["stats"]["acc"].detach().cpu().item() - ) / (batch_idx + 1) - - if self.use_ddp or self.use_fsdp or self.use_deepspeed: - val_loss_avg = torch.tensor(self.val_loss_avg, dtype=torch.float32).to(self.device) - val_acc_avg = torch.tensor(self.val_acc_avg, dtype=torch.float32).to(self.device) - dist.all_reduce(val_loss_avg, op=dist.ReduceOp.SUM) - dist.all_reduce(val_acc_avg, op=dist.ReduceOp.SUM) - self.val_loss_avg = val_loss_avg.detach().cpu().item() / self.world_size - self.val_acc_avg = val_acc_avg.detach().cpu().item() / self.world_size if kwargs.get("step_in_epoch", None) is None: ckpt_name = f"model.pt.ep{epoch}" else: ckpt_name = f'model.pt.ep{epoch}.{kwargs.get("step_in_epoch")}' - self.val_acc_step_or_epoch[ckpt_name] = self.val_acc_avg - self.val_loss_step_or_epoch[ckpt_name] = self.val_loss_avg + record_validation_metrics( + self, ckpt_name, + metrics.compute( + self.device, distributed=self.use_ddp or self.use_fsdp or self.use_deepspeed + ), + ) if self.use_ddp or self.use_fsdp or self.use_deepspeed: dist.barrier() diff --git a/tests/test_checkpoint_metric_distributed.py b/tests/test_checkpoint_metric_distributed.py new file mode 100644 index 0000000000..cf6d76689e --- /dev/null +++ b/tests/test_checkpoint_metric_distributed.py @@ -0,0 +1,86 @@ +"""Actual two-rank CPU collectives; no distributed acoustic training claim.""" + +from datetime import timedelta +import json +from pathlib import Path +import time + +import pytest +import torch.distributed as dist +import torch.multiprocessing as mp + +from test_checkpoint_metric_presence import Run + + +def _validate_worker(rank, root): + root = Path(root) + dist.init_process_group("gloo", rank=rank, world_size=2, + init_method=(root / "gloo-init").as_uri(), + timeout=timedelta(seconds=30)) + results = {} + try: + for kind in ("torch", "ds_torch", "ds_engine"): + for case in ("valid", "missing", "nonfinite", "empty", "reject_best"): + run = Run(root / f"{kind}-{case}-{rank}", kind) + run.trainer.rank = rank + run.trainer.world_size = 2 + run.trainer.use_ddp = True + values = ([(1.0, 0.2), (3.0, 0.4)] if rank == 0 else + [(3.0, 0.6), (5.0, 0.8)]) + if rank == 1 and case == "missing": + values[-1] = (5.0, "missing") + if rank == 1 and case == "nonfinite": + values[-1] = (float("nan"), 0.8) + if rank == 1 and case == "empty": + values = [] + try: + run.validate(1, values) + if case == "reject_best": + run.save(1) + if kind == "ds_engine": + assert run.trainer.saved_ckpts == {"model.pt.ep1": pytest.approx(0.5)} + assert run.checkpoint("model.pt")["saved_ckpts"] == run.trainer.saved_ckpts + before = dict(run.trainer.val_acc_step_or_epoch) + rejected = False + try: + run.validate(1, [(1.0, "missing")]) + except ValueError as error: + rejected = "current best" in str(error) + results[f"{kind}-reject_best"] = { + "rejected": rejected, + "unchanged": before == run.trainer.val_acc_step_or_epoch, + } + continue + results[f"{kind}-{case}"] = { + "acc": run.trainer.val_acc_step_or_epoch, + "loss": run.trainer.val_loss_step_or_epoch, + } + finally: + run.close() + (root / f"rank-{rank}.json").write_text(json.dumps(results)) + finally: + dist.destroy_process_group() + + +def test_distributed_validation_agrees_on_availability_and_valid_means(tmp_path): + assert dist.is_gloo_available(), "CPU Gloo is required for this regression" + context = mp.spawn(_validate_worker, args=(str(tmp_path),), nprocs=2, join=False) + try: + deadline = time.monotonic() + 150 + while not context.join(timeout=5): + assert time.monotonic() < deadline, "Validation collectives did not terminate" + finally: + for process in context.processes: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + reports = [json.loads((tmp_path / f"rank-{rank}.json").read_text()) for rank in (0, 1)] + assert reports[0] == reports[1] + for kind in ("torch", "ds_torch", "ds_engine"): + assert reports[0][f"{kind}-reject_best"] == {"rejected": True, "unchanged": True} + assert reports[0][f"{kind}-valid"]["acc"]["model.pt.ep1"] == pytest.approx(0.5) + assert reports[0][f"{kind}-valid"]["loss"]["model.pt.ep1"] == pytest.approx(3.0) + assert reports[0][f"{kind}-missing"]["acc"] == {} + assert reports[0][f"{kind}-missing"]["loss"]["model.pt.ep1"] == pytest.approx(3.0) + for case in ("nonfinite", "empty"): + assert reports[0][f"{kind}-{case}"] == {"acc": {}, "loss": {}} diff --git a/tests/test_checkpoint_metric_presence.py b/tests/test_checkpoint_metric_presence.py new file mode 100644 index 0000000000..07cce56ecb --- /dev/null +++ b/tests/test_checkpoint_metric_presence.py @@ -0,0 +1,301 @@ +"""Exercise validation, checkpoint files and ranking without acoustic weights.""" + +import logging +from pathlib import Path + +import pytest +import torch + +from funasr.train_utils.average_nbest_models import _get_checkpoint_paths, average_checkpoints +from funasr.train_utils.trainer import Trainer as TrainerTorch +from funasr.train_utils.trainer_ds import Trainer as TrainerDs + + +SAVE_PATHS = ["torch", "ds_torch", "ds_engine"] + + +class MetricModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.marker = torch.nn.Parameter(torch.zeros(())) + + def forward(self, value, accuracy="missing"): + stats = {"loss": value, "acc_rich": value.new_tensor(0.8)} + if accuracy != "missing": + stats["acc"] = None if accuracy is None else value.new_tensor(accuracy) + return value, stats, value.new_tensor(1.0) + + def save_checkpoint(self, save_dir, tag, client_state): + # Model the engine's save boundary, not distributed DeepSpeed training. + path = Path(save_dir) / tag + path.mkdir(exist_ok=True) + torch.save({**client_state, "state_dict": self.state_dict()}, path / "mp_rank_00_model_states.pt") + + def load_checkpoint(self, save_dir, tag): + path = Path(save_dir) / tag / "mp_rank_00_model_states.pt" + state = torch.load(path, weights_only=True) + self.load_state_dict(state["state_dict"]) + return str(path), state + + +class Batches: + def __init__(self, values): + self.values = values + self.batch_sampler = self + + def set_epoch(self, epoch): + pass + + def __len__(self): + return len(self.values) + + def __iter__(self): + for loss, accuracy in self.values: + yield {"value": torch.tensor(loss), "accuracy": accuracy} + + +class Run: + def __init__(self, path, kind, ranking="acc", keep=2): + self.path = path + self.kind = kind + kwargs = dict(output_dir=str(path), device="cpu", log_interval=100, + avg_keep_nbest_models_type=ranking, keep_nbest_models=keep) + self.trainer = (TrainerTorch(local_rank=0, **kwargs) if kind == "torch" else + TrainerDs(rank=0, local_rank=0, world_size=1, use_deepspeed=False, **kwargs)) + self.model = MetricModel() + self.optim = torch.optim.SGD(self.model.parameters(), lr=0.01) + self.scheduler = torch.optim.lr_scheduler.StepLR(self.optim, step_size=1) + + def validate(self, epoch, values): + if self.kind == "ds_engine": + self.trainer.use_deepspeed = False + self.trainer.validate_epoch(model=self.model, dataloader_val=Batches(values), epoch=epoch) + + def save(self, epoch): + if self.kind == "ds_engine": + self.trainer.use_deepspeed = True + self.trainer.save_checkpoint(epoch=epoch, model=self.model, + optim=self.optim, scheduler=self.scheduler) + + def checkpoint(self, name): + path = self.path / name + if self.kind == "ds_engine": + path = path / "mp_rank_00_model_states.pt" + return torch.load(path, weights_only=True) + + def close(self): + writer = getattr(self.trainer, "writer", None) + if writer is not None: + writer.close() + + +@pytest.fixture(params=SAVE_PATHS) +def run(request, tmp_path): + case = Run(tmp_path, request.param) + yield case + case.close() + + +@pytest.mark.parametrize("accuracy", ["missing", None, float("nan"), float("inf")]) +def test_unavailable_accuracy_never_ranks_or_averages(run, accuracy, caplog): + with caplog.at_level(logging.WARNING): + for epoch in (1, 2): + run.validate(epoch, [(1.0, accuracy)]) + run.save(epoch) + assert run.trainer.val_acc_step_or_epoch == {} + assert run.trainer.saved_ckpts == {} + assert run.trainer.best_step_or_epoch == "" + assert not (run.path / "model.pt.best").exists() + assert run.checkpoint("model.pt")["val_acc_step_or_epoch"] == {} + assert _get_checkpoint_paths(str(run.path), use_deepspeed=run.kind == "ds_engine") == [] + messages = [r.message for r in caplog.records if "avg_keep_nbest_models_type=loss" in r.message] + assert len(messages) == 1 + + +@pytest.mark.parametrize("values", [[(1.0, 0.4), (1.0, "missing")], + [(1.0, "missing"), (1.0, 0.4)]]) +def test_partial_accuracy_is_not_a_validation_score(run, values): + run.validate(1, values) + run.save(1) + assert run.trainer.val_acc_step_or_epoch == {} + assert run.trainer.val_loss_step_or_epoch == {"model.pt.ep1": 1.0} + assert run.trainer.saved_ckpts == {} + + +def test_empty_validation_does_not_reuse_previous_averages(run): + run.validate(1, [(3.0, 0.5)]) + run.save(1) + run.validate(2, []) + run.save(2) + assert "model.pt.ep2" not in run.trainer.val_acc_step_or_epoch + assert "model.pt.ep2" not in run.trainer.val_loss_step_or_epoch + assert run.trainer.best_step_or_epoch == "model.pt.ep1" + + +@pytest.mark.parametrize("bad_loss", [float("nan"), float("inf")]) +def test_nonfinite_loss_does_not_publish_partial_metrics(run, bad_loss): + run.validate(1, [(1.0, 0.4), (bad_loss, 0.5)]) + run.save(1) + assert run.trainer.val_loss_step_or_epoch == {} + assert run.trainer.val_acc_step_or_epoch == {} + + +def test_real_zero_accuracy_preserves_existing_latest_tie(run): + for epoch in (1, 2, 3): + with torch.no_grad(): + run.model.marker.fill_(epoch) + run.validate(epoch, [(1.0, 0.0)]) + run.save(epoch) + assert run.trainer.best_step_or_epoch == "model.pt.ep3" + assert run.checkpoint("model.pt.best")["state_dict"]["marker"].item() == 3 + + +def test_validation_logs_current_running_means_not_previous_epoch(run, monkeypatch): + observed = [] + monkeypatch.setattr(run.trainer, "log", lambda *a, **kw: observed.append( + (run.trainer.val_loss_avg, run.trainer.val_acc_avg))) + run.validate(1, [(1.0, 0.2), (3.0, 0.4)]) + run.validate(2, [(5.0, 0.6)]) + assert [v[0] for v in observed] == pytest.approx([1.0, 2.0, 5.0]) + assert [v[1] for v in observed] == pytest.approx([0.2, 0.3, 0.6]) + + +def test_explicit_loss_best_pruning_and_average_agree(run, caplog): + run.trainer.avg_keep_nbest_models_type = "loss" + for epoch, loss in enumerate((3.0, 1.0, 2.0), 1): + with torch.no_grad(): + run.model.marker.fill_(epoch) + run.validate(epoch, [(loss, "missing")]) + run.save(epoch) + assert run.trainer.best_step_or_epoch == "model.pt.ep2" + assert run.checkpoint("model.pt.best")["state_dict"]["marker"].item() == 2 + assert set(run.trainer.saved_ckpts) == {"model.pt.ep2", "model.pt.ep3"} + path = average_checkpoints(str(run.path), last_n=2, use_deepspeed=run.kind == "ds_engine") + assert torch.load(path, weights_only=True)["state_dict"]["marker"].item() == 2.5 + assert not any("avg_keep_nbest_models_type=loss" in r.message for r in caplog.records) + + +def test_revalidation_removes_stale_nonbest_metric_before_serialization(run): + run.validate(1, [(1.0, 0.9)]) + run.save(1) + run.validate(2, [(1.0, 0.4)]) + run.save(2) + run.validate(2, [(1.0, "missing")]) + run.save(2) + state = run.checkpoint("model.pt") + assert "model.pt.ep2" not in state["val_acc_step_or_epoch"] + assert "model.pt.ep2" not in state["saved_ckpts"] + assert run.trainer.best_step_or_epoch == "model.pt.ep1" + + +def test_invalid_revalidation_cannot_overwrite_current_best(run): + run.validate(1, [(1.0, 0.5)]) + run.save(1) + before = run.checkpoint("model.pt.best") + with pytest.raises(ValueError, match="current best"): + run.validate(1, [(1.0, "missing")]) + assert run.trainer.val_acc_step_or_epoch == {"model.pt.ep1": 0.5} + assert torch.equal(before["state_dict"]["marker"], + run.checkpoint("model.pt.best")["state_dict"]["marker"]) + + +def test_resume_preserves_unavailable_metric_exclusion(run): + run.validate(1, [(1.0, "missing")]) + run.save(1) + resumed = Run(run.path, run.kind) + try: + resumed.trainer.use_deepspeed = run.kind == "ds_engine" if run.kind != "torch" else False + resumed.trainer.resume_checkpoint(model=resumed.model, optim=resumed.optim, scheduler=resumed.scheduler) + assert resumed.trainer.val_acc_step_or_epoch == {} + assert resumed.trainer.saved_ckpts == {} + finally: + resumed.close() + + +def test_resume_cannot_mix_accuracy_history_with_loss_ranking(run): + run.validate(1, [(1.0, 0.0)]) + run.save(1) + resumed = Run(run.path, run.kind, ranking="loss") + try: + if run.kind != "torch": + resumed.trainer.use_deepspeed = run.kind == "ds_engine" + with pytest.raises(ValueError, match="ranking metric"): + resumed.trainer.resume_checkpoint(model=resumed.model, optim=resumed.optim, scheduler=resumed.scheduler) + finally: + resumed.close() + + +def test_resume_retains_current_best_guard(run): + run.validate(1, [(1.0, 0.5)]) + run.save(1) + resumed = Run(run.path, run.kind) + try: + if run.kind != "torch": + resumed.trainer.use_deepspeed = run.kind == "ds_engine" + resumed.trainer.resume_checkpoint(model=resumed.model, optim=resumed.optim, scheduler=resumed.scheduler) + assert resumed.trainer.best_step_or_epoch == "model.pt.ep1" + with pytest.raises(ValueError, match="current best"): + resumed.validate(1, [(1.0, "missing")]) + finally: + resumed.close() + + +@pytest.mark.parametrize("ranking,values", [("acc", [0.3, 0.9, 0.5]), ("loss", [3.0, 1.0, 2.0])]) +def test_persisted_ranking_matches_best_and_pruned_files(run, ranking, values): + run.trainer.avg_keep_nbest_models_type = ranking + run.trainer.keep_nbest_models = 1 + for epoch, value in enumerate(values, 1): + loss, acc = (1.0, value) if ranking == "acc" else (value, "missing") + run.validate(epoch, [(loss, acc)]) + run.save(epoch) + state = run.checkpoint("model.pt") + assert state["best_step_or_epoch"] == run.trainer.best_step_or_epoch + assert state["saved_ckpts"] == run.trainer.saved_ckpts + best = run.checkpoint("model.pt.best") + assert best["best_step_or_epoch"] == run.trainer.best_step_or_epoch + assert all((run.path / name).exists() for name in state["saved_ckpts"]) + assert not (run.path / "model.pt.ep3").exists() + assert (run.path / "model.pt").exists() + + +@pytest.mark.parametrize("failed_tag", ["model.pt", "model.pt.best"]) +def test_write_failure_does_not_delete_previous_candidate(run, failed_tag, monkeypatch): + run.trainer.keep_nbest_models = 1 + with torch.no_grad(): + run.model.marker.fill_(1) + run.validate(1, [(1.0, 0.3)]) + run.save(1) + with torch.no_grad(): + run.model.marker.fill_(2) + run.validate(2, [(1.0, 0.9)]) + if run.kind == "ds_engine": + original = run.model.save_checkpoint + + def fail_engine(save_dir, tag, client_state): + if tag == failed_tag: + raise OSError("simulated checkpoint write failure") + return original(save_dir, tag, client_state) + + monkeypatch.setattr(run.model, "save_checkpoint", fail_engine) + else: + original = torch.save + + def fail_torch(state, path, *args, **kwargs): + if Path(path).name == failed_tag: + raise OSError("simulated checkpoint write failure") + return original(state, path, *args, **kwargs) + + monkeypatch.setattr(torch, "save", fail_torch) + with pytest.raises(OSError, match="simulated checkpoint"): + run.save(2) + assert (run.path / "model.pt.ep1").exists() + assert run.checkpoint("model.pt.best")["state_dict"]["marker"].item() == 1 + + +@pytest.mark.parametrize("ranking,scores", [("acc", {"model.pt.ep1": 0.9, "model.pt.ep2": 0.8, "model.pt.ep3": float("nan")}), + ("loss", {"model.pt.ep1": 0.1, "model.pt.ep2": 0.2, "model.pt.ep3": float("nan")})]) +def test_average_filters_nonfinite_and_missing_files_before_taking_n(tmp_path, ranking, scores): + torch.save({"avg_keep_nbest_models_type": ranking, f"val_{ranking}_step_or_epoch": scores}, tmp_path / "model.pt") + for epoch in (2, 3): + torch.save({"state_dict": {"x": torch.tensor(float(epoch))}}, tmp_path / f"model.pt.ep{epoch}") + assert _get_checkpoint_paths(str(tmp_path), last_n=1) == [str(tmp_path / "model.pt.ep2")] diff --git a/tests/test_sensevoice_continual_finetuning_docs.py b/tests/test_sensevoice_continual_finetuning_docs.py index 0c80dae3c1..299ce59107 100644 --- a/tests/test_sensevoice_continual_finetuning_docs.py +++ b/tests/test_sensevoice_continual_finetuning_docs.py @@ -38,6 +38,25 @@ def test_english_guide_covers_retention_controls_and_worked_example(): assert all(item in guide for item in required) +def test_checkpoint_selection_recipe_and_evaluation_boundaries(): + recipe = (SENSEVOICE_DIR / "finetune.sh").read_text(encoding="utf-8") + assert "++train_conf.avg_keep_nbest_models_type=loss" in recipe + for name in ("CONTINUAL_FINETUNING.md", "CONTINUAL_FINETUNING_zh.md"): + guide = (SENSEVOICE_DIR / name).read_text(encoding="utf-8") + assert "++train_conf.avg_keep_nbest_models_type=loss" in guide + assert "acc_rich" in guide + assert "<|minnan|>" in guide + assert "93.3%" not in guide + english = (SENSEVOICE_DIR / "CONTINUAL_FINETUNING.md").read_text(encoding="utf-8") + chinese = (SENSEVOICE_DIR / "CONTINUAL_FINETUNING_zh.md").read_text(encoding="utf-8") + assert "independently audited" in english + assert "new output directory" in english + assert "allowlist" in english + assert "独立核验" in chinese + assert "新的输出目录" in chinese + assert "白名单" in chinese + + def test_chinese_guide_covers_retention_controls_and_worked_example(): guide = (SENSEVOICE_DIR / "CONTINUAL_FINETUNING_zh.md").read_text(encoding="utf-8")