fix: disable timestamps when CTC weights are missing (#3211)

Fail closed for native, remote-code, and vLLM Fun-ASR-Nano timestamp inference when checkpoint CTC weights are missing, partial, shape-mismatched, DDP-prefixed, or wrapped.

Fixes #3208.

Author: LauraGPT

examples/industrial_data_pretraining/fun_asr_nano/README.md

@@ -64,6 +64,12 @@ pip install -r requirements.txt
 - [ ] Support speaker diarization
 - [x] Support model training
 
+> [!NOTE]
+> Character-level timestamps require a checkpoint with complete `ctc_decoder.*` and
+> `ctc.*` weights. The current `Fun-ASR-MLT-Nano-2512` checkpoint does not include
+> those weights, so FunASR logs a warning and returns text without `timestamps` or
+> `ctc_timestamps` instead of aligning with uninitialized layers.
+
 # Usage 🛠️
 
 ## Inference

examples/industrial_data_pretraining/fun_asr_nano/README_zh.md

@@ -64,6 +64,11 @@ pip install -r requirements.txt
 - [ ] 支持区分说话人识别
 - [x] 支持模型训练
 
+> [!NOTE]
+> 字符级时间戳要求 checkpoint 同时包含完整的 `ctc_decoder.*` 和 `ctc.*` 权重。
+> 当前 `Fun-ASR-MLT-Nano-2512` checkpoint 未包含这些权重，因此 FunASR 会记录警告并仅返回文本，
+> 不再使用未初始化的层生成 `timestamps` 或 `ctc_timestamps`。
+
 # 用法 🛠️
 
 ## 推理

examples/industrial_data_pretraining/fun_asr_nano/model.py

@@ -15,6 +15,10 @@
 from funasr.train_utils.device_funcs import force_gatherable, to_device
 from funasr.utils.datadir_writer import DatadirWriter
 from funasr.utils.load_utils import extract_fbank, load_audio_text_image_video
+from funasr.models.fun_asr_nano.checkpoint_utils import (
+    disable_incomplete_ctc,
+    normalize_checkpoint_state,
+)
 from funasr.models.fun_asr_nano.device_utils import resolve_autocast_device_type
 
 
@@ -109,6 +113,7 @@ def __init__(
 
         # ctc decoder
         self.ctc_decoder = None
+        self._externally_loaded_ctc_keys = set()
         # TODO: fix table name
         ctc_decoder_class = tables.adaptor_classes.get(kwargs.get("ctc_decoder", None))
         if ctc_decoder_class is not None:
@@ -135,8 +140,15 @@ def __init__(
             self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf)
             init_param_path = ctc_decoder_conf.get("init_param_path", None)
             if init_param_path is not None:
-                src_state = torch.load(init_param_path, map_location="cpu")
+                src_state = normalize_checkpoint_state(
+                    torch.load(init_param_path, map_location="cpu")
+                )
                 flag = self.ctc_decoder.load_state_dict(src_state, strict=False)
+                self._externally_loaded_ctc_keys.update(
+                    f"ctc_decoder.{key}"
+                    for key in self.ctc_decoder.state_dict()
+                    if key in src_state
+                )
                 logging.info(f"Loading ctc_decoder ckpt: {init_param_path}, status: {flag}")
             freeze = ctc_decoder_conf.get("freeze", False)
             if freeze:
@@ -160,6 +172,11 @@ def __init__(
         rank = int(os.environ.get("RANK", 0))
         logging.info(f"rank: {rank}, model is builded.")
 
+    def on_pretrained_model_loaded(self, loaded_keys):
+        """Fail closed when a checkpoint configures CTC without trained weights."""
+        loaded_keys = set(loaded_keys).union(getattr(self, "_externally_loaded_ctc_keys", ()))
+        disable_incomplete_ctc(self, loaded_keys, log=logging)
+
     def forward(
         self,
         speech: torch.Tensor = None,

funasr/models/fun_asr_nano/checkpoint_utils.py

@@ -0,0 +1,47 @@
+def normalize_checkpoint_state(state):
+    """Unwrap common checkpoint containers and remove one DDP key prefix."""
+    while isinstance(state, dict):
+        wrapped = next(
+            (
+                state[key]
+                for key in ("state_dict", "model_state_dict", "model")
+                if isinstance(state.get(key), dict)
+            ),
+            None,
+        )
+        if wrapped is None or wrapped is state:
+            break
+        state = wrapped
+
+    return {
+        key[len("module.") :] if key.startswith("module.") else key: value
+        for key, value in state.items()
+    }
+
+
+def disable_incomplete_ctc(model, loaded_keys, *, log):
+    """Disable timestamp inference unless every required CTC tensor was loaded."""
+    if model.ctc_decoder is None or model.ctc is None:
+        return []
+
+    expected = {f"ctc_decoder.{key}" for key in model.ctc_decoder.state_dict()}
+    expected.update(f"ctc.{key}" for key in model.ctc.state_dict())
+    missing = sorted(expected.difference(loaded_keys))
+    if not missing:
+        return []
+
+    preview = ", ".join(missing[:3])
+    suffix = "" if len(missing) <= 3 else ", ..."
+    log.warning(
+        "Disabling CTC timestamps because the checkpoint did not initialize "
+        "%d of %d required CTC tensors (%s%s). Text transcription remains available.",
+        len(missing),
+        len(expected),
+        preview,
+        suffix,
+    )
+    model.ctc_decoder = None
+    model.ctc = None
+    model.ctc_tokenizer = None
+    model.blank_id = None
+    return missing

funasr/models/fun_asr_nano/inference_vllm.py

@@ -33,6 +33,8 @@
 import torch
 import torch.nn as nn
 
+from .checkpoint_utils import disable_incomplete_ctc, normalize_checkpoint_state
+
 logger = logging.getLogger(__name__)
 
 dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
@@ -343,20 +345,7 @@ def _load_audio_components(self, model_dir: str, **kwargs):
 
             # CTC decoder
             if self.ctc_decoder is not None:
-                ctc_dec_state = {
-                    k[len("ctc_decoder."):]: v
-                    for k, v in state_dict.items()
-                    if k.startswith("ctc_decoder.")
-                }
-                if ctc_dec_state:
-                    self.ctc_decoder.load_state_dict(ctc_dec_state, strict=False)
-                ctc_state = {
-                    k[len("ctc."):]: v
-                    for k, v in state_dict.items()
-                    if k.startswith("ctc.") and not k.startswith("ctc_decoder.")
-                }
-                if ctc_state:
-                    self.ctc.load_state_dict(ctc_state, strict=False)
+                self._load_ctc_weights(state_dict)
 
         # Move to device
         self.audio_encoder = self.audio_encoder.to(self.device, dtype=torch.float32)
@@ -365,6 +354,35 @@ def _load_audio_components(self, model_dir: str, **kwargs):
             self.ctc_decoder = self.ctc_decoder.to(self.device, dtype=torch.float32)
             self.ctc = self.ctc.to(self.device, dtype=torch.float32)
 
+    def _load_ctc_weights(self, state_dict):
+        """Load timestamp modules and disable them when the checkpoint is incomplete."""
+        if self.ctc_decoder is None or self.ctc is None:
+            return
+
+        state_dict = normalize_checkpoint_state(state_dict)
+
+        def compatible_state(module, prefix):
+            expected = module.state_dict()
+            return {
+                key[len(prefix) :]: value
+                for key, value in state_dict.items()
+                if key.startswith(prefix)
+                and key[len(prefix) :] in expected
+                and value.size() == expected[key[len(prefix) :]].size()
+            }
+
+        ctc_dec_state = compatible_state(self.ctc_decoder, "ctc_decoder.")
+        if ctc_dec_state:
+            self.ctc_decoder.load_state_dict(ctc_dec_state, strict=False)
+
+        ctc_state = compatible_state(self.ctc, "ctc.")
+        if ctc_state:
+            self.ctc.load_state_dict(ctc_state, strict=False)
+
+        loaded_keys = {f"ctc_decoder.{key}" for key in ctc_dec_state}
+        loaded_keys.update(f"ctc.{key}" for key in ctc_state)
+        disable_incomplete_ctc(self, loaded_keys, log=logger)
+
     def _load_embedding_layer(self, model_dir: str):
         """Load the LLM embedding layer for text token embedding computation."""
         model_pt = os.path.join(model_dir, "model.pt")

funasr/models/fun_asr_nano/model.py

@@ -22,6 +22,7 @@
     AutoModelForCausalLM = None
 
 from .ctc import CTC
+from .checkpoint_utils import disable_incomplete_ctc, normalize_checkpoint_state
 from .device_utils import resolve_autocast_device_type
 from .tools.utils import forced_align
 
@@ -145,6 +146,7 @@ def __init__(
 
         # ctc decoder
         self.ctc_decoder = None
+        self._externally_loaded_ctc_keys = set()
         # TODO: fix table name
         ctc_decoder_class = tables.adaptor_classes.get(kwargs.get("ctc_decoder", None))
         if ctc_decoder_class is not None:
@@ -171,8 +173,15 @@ def __init__(
             self.ctc_decoder = ctc_decoder_class(**ctc_decoder_conf)
             init_param_path = ctc_decoder_conf.get("init_param_path", None)
             if init_param_path is not None:
-                src_state = torch.load(init_param_path, map_location="cpu")
+                src_state = normalize_checkpoint_state(
+                    torch.load(init_param_path, map_location="cpu")
+                )
                 flag = self.ctc_decoder.load_state_dict(src_state, strict=False)
+                self._externally_loaded_ctc_keys.update(
+                    f"ctc_decoder.{key}"
+                    for key in self.ctc_decoder.state_dict()
+                    if key in src_state
+                )
                 logging.info(f"Loading ctc_decoder ckpt: {init_param_path}, status: {flag}")
             freeze = ctc_decoder_conf.get("freeze", False)
             if freeze:
@@ -196,6 +205,11 @@ def __init__(
         rank = int(os.environ.get("RANK", 0))
         logging.info(f"rank: {rank}, model is builded.")
 
+    def on_pretrained_model_loaded(self, loaded_keys):
+        """Fail closed when a checkpoint configures CTC without trained weights."""
+        loaded_keys = set(loaded_keys).union(getattr(self, "_externally_loaded_ctc_keys", ()))
+        disable_incomplete_ctc(self, loaded_keys, log=logging)
+
     def forward(
         self,
         speech: torch.Tensor = None,

funasr/train_utils/load_pretrained_model.py

@@ -45,6 +45,7 @@ def load_pretrained_model(
     src_state = src_state["state_dict"] if "state_dict" in src_state else src_state
     src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state
     src_state = src_state["model"] if "model" in src_state else src_state
+    loaded_keys = set()
 
     if isinstance(scope_map, str):
         scope_map = scope_map.split(",")
@@ -96,8 +97,18 @@ def load_pretrained_model(
                 )
             else:
                 dst_state[k] = src_state[k_src]
+                loaded_keys.add(k)
         else:
             print(f"Warning, miss key in ckpt: {k}, {path}")
 
     flag = obj.load_state_dict(dst_state, strict=True)
+    on_loaded = getattr(obj, "on_pretrained_model_loaded", None)
+    if not callable(on_loaded) and hasattr(obj, "module"):
+        on_loaded = getattr(obj.module, "on_pretrained_model_loaded", None)
+        if callable(on_loaded):
+            loaded_keys = {
+                key[len("module.") :] if key.startswith("module.") else key for key in loaded_keys
+            }
+    if callable(on_loaded):
+        on_loaded(frozenset(loaded_keys))
     logging.info(f"Loading ckpt: {path}, status: {flag}")