[bugfix]: address Kandinsky5 QAD PR review findings

Fixes all P1/P2 items from the review of 7c7704da:

- E2E test's data_path pointed at a nonexistent combined_parquet_dataset/
  subdir; preprocess_kandinsky5_overfit.py writes directly to its output
  root.
- Stage 2 loaded the raw DCP checkpoint-N/ directly, which isn't a
  diffusers model dir. Added dcp_to_diffusers --verify (strict-reload
  the export immediately) and wired the conversion into the E2E test,
  launcher scripts, and README between stage 1 and stage 2.
- Kandinsky5DmdDenoisingStage called the transformer without
  set_forward_context, so LocalAttention's missing-context guard caught
  the AssertionError and silently fell back to plain SDPA -- Stage 2
  validation never actually exercised ATTN_QAT_TRAIN. Wrapped the call
  to match the parent stage's contract and added a scoped backend
  assertion (ATTN_QAT_TRAIN/ATTN_QAT_INFER only, since other backends
  can legitimately fall back to auto-selection).
- Exported DMD checkpoints kept the base T2V model_index.json, so
  VideoGenerator.from_pretrained resolved them to Kandinsky5T2VPipeline
  instead of the DMD sampler. Added Kandinsky5DMDConfig and documented
  the required override_pipeline_cls_name + pipeline_config combo.
- The FSDP-training dtype-detection special case (scan params, default
  to bf16) applied unconditionally, so an explicit fp32 pipeline was
  silently cast to bf16 for plain inference. Extracted
  _resolve_target_dtype and gated the scan on isinstance(transformer,
  FSDPModule).
- Preprocessing tokenized captions with a 512-token cap that doesn't
  account for the 129-token system template runtime strips off; now
  reads text_encoder_max_lengths from pipeline_config to match runtime.
  Also replaced unconditional item["cap"][0] indexing (silently takes
  the first character of a string caption) with a get_caption() helper
  that accepts/validates both the documented string schema and the
  list-of-variants schema some producers use.
- test_kandinsky5_qat_attention_engages.py set FASTVIDEO_ATTENTION_BACKEND
  and cleared the attention backend selector cache without teardown,
  leaking ATTN_QAT_TRAIN into later tests in the same pytest process.
  Now uses monkeypatch + try/finally.

Added regression tests: test_kandinsky5_target_dtype.py,
test_kandinsky5_dmd_pipeline_resolution.py,
test_kandinsky5_overfit_caption.py.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh

@@ -9,8 +9,20 @@
 # Kandinsky5-specific handling needed).
 #
 # Run examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh (stage 1)
-# first and point models.student/teacher/critic.init_from in
-# dmd2_t2v_480p_qat.yaml at that checkpoint before launching this script.
+# first. Its checkpoint-N/ is a raw DCP checkpoint (dcp/ + metadata/RNG
+# state only) -- models.student/teacher/critic.init_from need a diffusers
+# model directory (model_index.json + component subfolders), so convert it
+# first:
+#
+#   python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+#       --checkpoint outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N> \
+#       --output-dir outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N>-diffusers \
+#       --role student --verify
+#
+# then point models.student/teacher/critic.init_from in dmd2_t2v_480p_qat.yaml
+# at that exported checkpoint-<N>-diffusers directory (all three roles use
+# the same weights; teacher/critic full-precision masking happens at load
+# time, not by pointing at a different export) before launching this script.
 set -euo pipefail
 
 export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN

examples/train/configs/distribution_matching/kandinsky5/dmd2_t2v_480p_qat.yaml

@@ -1,6 +1,11 @@
 # DMD2 QAT-aware distillation: Kandinsky5 T2V Lite, 480p (teacher many-step
 # -> student 4-step), continuing from a stage-1 Attn-QAT finetune checkpoint.
 #
+# init_from below must be a diffusers model directory (model_index.json +
+# component subfolders), not the raw checkpoint-N/ stage 1 writes (dcp/ +
+# metadata/RNG state only). Convert with fastvideo.train.entrypoint.dcp_to_diffusers
+# first -- see distill_dmd_qat.sh's header for the exact command.
+#
 # - Teacher: frozen, full-precision Kandinsky5 (loaded from the stage-1 QAT
 #   checkpoint's base weights, or the original pretrained checkpoint --
 #   the _loading_teacher_critic_model gate in component_loader.py masks

examples/train/configs/fine_tuning/kandinsky5/README.md

@@ -62,9 +62,30 @@ during either training stage.
 
 ## Train stage 2 (QAT-aware DMD distillation)
 
+Stage 1 writes a raw DCP checkpoint (`checkpoint-<N>/dcp` + metadata/RNG
+state) under `training.checkpoint.output_dir`
+(`outputs/kandinsky5_t2v_qat_finetune/` by default) -- `models.*.init_from`
+needs a diffusers model directory instead (`model_index.json` + component
+subfolders), so convert it first:
+
+```bash
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+    --checkpoint outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N> \
+    --output-dir outputs/kandinsky5_t2v_qat_finetune/checkpoint-<N>-diffusers \
+    --role student --verify
+```
+
+`--role student` exports the trained transformer weights (the only role
+that matters here: student/teacher/critic in stage 2 all load from the same
+checkpoint, and `_loading_teacher_critic_model` handles the full-precision
+masking for teacher/critic at load time, not at export time). `--verify`
+strictly reloads the exported transformer immediately, so a key-mapping bug
+fails here instead of deep inside the stage-2 launch below. Then point
+`models.student/teacher/critic.init_from` in `dmd2_t2v_480p_qat.yaml` at
+that `checkpoint-<N>-diffusers` directory (already the placeholder value in
+the checked-in config) and launch:
+
 ```bash
-# point models.student/teacher/critic.init_from in dmd2_t2v_480p_qat.yaml
-# at the stage-1 checkpoint first
 bash examples/train/configs/distribution_matching/kandinsky5/distill_dmd_qat.sh
 ```
 
@@ -90,15 +111,35 @@ their default layer lists (`fastvideo/layers/quantization/nvfp4_qat_config.py`,
 `fastvideo/layers/quantization/fp8_config.py`), so no `target_layers`
 override is needed:
 
+`dcp_to_diffusers` copies the base T2V checkpoint's `model_index.json`
+unchanged into every export (see "Train stage 2" above), so `_class_name`
+still says the base T2V pipeline and the registry
+(`fastvideo/registry.py`) has no way to auto-detect that a given directory
+is actually a DMD (four-step re-noise sampler) export -- it will resolve to
+`Kandinsky5T2VPipeline` and run the full-length sampler on DMD-distilled
+weights. Pass `override_pipeline_cls_name` and a `Kandinsky5DMDConfig`
+explicitly to select the right pipeline/sampler and get
+`dmd_denoising_steps` set (`Kandinsky5T2VConfig`'s default of `None` makes
+`Kandinsky5DmdDenoisingStage` raise immediately):
+
 ```python
 from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig
 from fastvideo.layers.quantization import get_quantization_config
 
 gen = VideoGenerator.from_pretrained(
     "path/to/kandinsky5_dmd_checkpoint", num_gpus=1,
+    override_pipeline_cls_name="Kandinsky5DMDPipeline",
+    pipeline_config=Kandinsky5DMDConfig(),
     transformer_quant=get_quantization_config("nvfp4_qat")(),
     use_fsdp_inference=False,
 )
+# Kandinsky5DmdDenoisingStage ignores request.sampling.num_inference_steps
+# and guidance_scale -- it's a fixed 4-step, no-CFG sampler driven entirely
+# by pipeline_config.dmd_denoising_steps (set above). Adjust
+# Kandinsky5DMDConfig(dmd_denoising_steps=[...]) instead if the checkpoint
+# was distilled/validated with a different schedule than the default
+# [1000, 750, 500, 250].
 gen.generate(request={"prompt": "...", "output": {"save_video": True}})
 ```
 
@@ -114,7 +155,14 @@ to a supported dense backend while the FP4 linear layers still run.
 - `fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py` --
   confirms `ATTN_QAT_TRAIN` actually engages (not a silent SDPA fallback)
   and that gradients flow through a forward+backward pass.
+- `fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py` --
+  confirms an unmodified export resolves to `Kandinsky5T2VPipeline` (the bug)
+  and that `override_pipeline_cls_name="Kandinsky5DMDPipeline"` fixes it (the
+  documented workaround above), plus `Kandinsky5DMDConfig`'s
+  `dmd_denoising_steps` default.
 - `fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py` -- a few
-  steps of both training stages on a single synthetic sample. No golden
-  reference video exists yet (unlike the Wan e2e test); once a real training
-  run is available, snapshot one and add an SSIM assertion to match.
+  steps of both training stages on a single synthetic sample, including the
+  `dcp_to_diffusers --verify` conversion between them (the same command
+  documented above). No golden reference video exists yet (unlike the Wan
+  e2e test); once a real training run is available, snapshot one and add an
+  SSIM assertion to match.

examples/train/configs/fine_tuning/kandinsky5/finetune_qat.sh

@@ -15,6 +15,12 @@
 #
 # Data: preprocess with fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py
 # (or an equivalent parquet dataset matching pyarrow_schema_t2v) first.
+#
+# Next: this writes a raw DCP checkpoint (checkpoint-N/dcp + metadata/RNG
+# state), not something stage 2 can load directly. Convert it with
+# fastvideo.train.entrypoint.dcp_to_diffusers before running
+# ../../distribution_matching/kandinsky5/distill_dmd_qat.sh -- see that
+# script's header for the exact command.
 set -euo pipefail
 
 export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # <-- enables Attn-QAT training

fastvideo/configs/pipelines/__init__.py

@@ -6,7 +6,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
-from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
@@ -21,6 +21,6 @@ __all__ = [
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig",
-    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig",
+    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/kandinsky5.py

@@ -120,3 +120,21 @@ class Kandinsky5I2VConfig(Kandinsky5T2VConfig):
         super().__post_init__()
         # I2V needs the VAE encoder to encode the conditioning image.
         self.vae_config.load_encoder = True
+
+
+@dataclass
+class Kandinsky5DMDConfig(Kandinsky5T2VConfig):
+    """Kandinsky-5.0 DMD (few-step distilled) text-to-video pipeline configuration.
+
+    Checkpoints exported by ``fastvideo.train.entrypoint.dcp_to_diffusers``
+    copy their base T2V checkpoint's ``model_index.json`` unchanged, so
+    ``_class_name`` still says the base T2V pipeline and the registry
+    (``fastvideo/registry.py``) cannot auto-detect a DMD export -- pass this
+    config together with ``override_pipeline_cls_name="Kandinsky5DMDPipeline"``
+    explicitly to ``VideoGenerator.from_pretrained``/``from_config`` (see
+    ``examples/train/configs/fine_tuning/kandinsky5/README.md``). Without it,
+    ``Kandinsky5T2VConfig``'s ``dmd_denoising_steps=None`` makes
+    ``Kandinsky5DmdDenoisingStage`` raise immediately.
+    """
+
+    dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250])

fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py

@@ -100,6 +100,29 @@ def load_video(path: str, num_frames: int, height: int, width: int) -> torch.Ten
     return video
 
 
+def get_caption(item: dict) -> str:
+    """Extract a caption from one ``videos2caption.json`` entry.
+
+    The documented schema (``docs/training/data_preprocess.md``) stores
+    ``cap`` as a plain string; some producers (e.g. this repo's own e2e
+    test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) instead
+    store a non-empty list of caption variants and use the first one.
+    Indexing unconditionally with ``item["cap"][0]`` silently takes the
+    first *character* of a string caption instead of erroring, so accept
+    and validate both forms explicitly here.
+    """
+    cap = item.get("cap")
+    if isinstance(cap, str):
+        if not cap:
+            raise ValueError(f"Empty 'cap' string for entry: {item}")
+        return cap
+    if isinstance(cap, list):
+        if not cap or not isinstance(cap[0], str) or not cap[0]:
+            raise ValueError(f"'cap' list must be non-empty with a non-empty first string, got: {item}")
+        return cap[0]
+    raise ValueError(f"'cap' must be a string or a list of strings, got {type(cap).__name__}: {item}")
+
+
 def main() -> None:
     # FastVideo's native text-encoder layers (e.g. CLIPAttention's
     # QKVParallelLinear) are tensor-parallel-aware and assert the TP process
@@ -145,6 +168,12 @@ def main() -> None:
     ).to(device).eval()
     qwen_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
     qwen_tok_kwargs = dict(pipeline_config.text_encoder_configs[0].tokenizer_kwargs)
+    # TextEncodingStage overrides max_length from pipeline_config.text_encoder_max_lengths
+    # at runtime (fastvideo/pipelines/stages/text_encoding.py) -- Reason1Config's static
+    # tokenizer_kwargs default (text_len=512) doesn't account for the 129-token Kandinsky
+    # system template ENCODE_START_IDX strips off afterwards, so mirror the runtime value
+    # here or training conditions on fewer caption tokens than inference does.
+    qwen_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[0]
 
     print("Loading CLIP text encoder...")
     clip_enc = TextEncoderLoader().load(
@@ -153,12 +182,13 @@ def main() -> None:
     ).to(device).eval()
     clip_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2"))
     clip_tok_kwargs = dict(pipeline_config.text_encoder_configs[1].tokenizer_kwargs)
+    clip_tok_kwargs["max_length"] = pipeline_config.text_encoder_max_lengths[1]
 
     # --- Process each video ---
     records = []
     for item in caption_data:
         video_name = item["path"]
-        caption = item["cap"][0]
+        caption = get_caption(item)
         video_path = os.path.join(DATA_DIR, "videos", video_name)
 
         print(f"\nProcessing: {video_name}")
@@ -243,7 +273,7 @@ def main() -> None:
     # Write validation prompts for callback.
     # Wrap in "data" key -- ValidationDataset expects field="data".
     # Use "caption" field -- ValidationDataset aliases it to "prompt".
-    val_prompts = {"data": [{"caption": item["cap"][0]} for item in caption_data]}
+    val_prompts = {"data": [{"caption": get_caption(item)} for item in caption_data]}
     val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
     with open(val_path, "w") as f:
         json.dump(val_prompts, f, indent=2)

fastvideo/pipelines/stages/kandinsky5.py

@@ -10,7 +10,10 @@ import torch
 from diffusers.utils.torch_utils import randn_tensor
 from tqdm.auto import tqdm
 
+import fastvideo.envs as envs
+from fastvideo.attention import LocalAttention
 from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder
+from fastvideo.attention.selector import backend_name_to_enum
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
@@ -27,6 +30,7 @@ from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.encoding import EncodingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
@@ -169,6 +173,96 @@ class Kandinsky5DenoisingStage(PipelineStage):
         seq_len = int(mask.sum(1).max().item())
         return torch.arange(seq_len, device=device)
 
+    # Backends whose selection is all-or-nothing: fastvideo.platforms.cuda
+    # raises ImportError immediately if the kernel isn't available, rather
+    # than the generic "requested backend unsupported by this layer, falling
+    # back to automatic selection" path _cached_get_attn_backend takes for
+    # most other backend names (see fastvideo/attention/selector.py). Both
+    # values are also always in Kandinsky5ArchConfig's
+    # _supported_attention_backends, so that generic fallback never applies
+    # to them here. That combination makes an exact backend match a safe,
+    # false-positive-free signal *only* for these two -- any other
+    # FASTVIDEO_ATTENTION_BACKEND value (e.g. a backend meant for a
+    # different model family sharing the same process/env var) can
+    # legitimately resolve to something other than what was requested.
+    _STRICT_BACKENDS = frozenset({
+        AttentionBackendEnum.ATTN_QAT_TRAIN,
+        AttentionBackendEnum.ATTN_QAT_INFER,
+    })
+
+    def _assert_local_attention_backend_engaged(self) -> None:
+        """Guard against ``LocalAttention`` silently falling back to SDPA.
+
+        ``Kandinsky5Attention.forward`` catches the ``AssertionError``
+        ``LocalAttention`` raises when no pipeline forward context is set and
+        falls back to plain ``F.scaled_dot_product_attention`` -- silently
+        skipping whichever kernel ``FASTVIDEO_ATTENTION_BACKEND`` requested
+        (the fake-quantized ``ATTN_QAT_TRAIN``/``ATTN_QAT_INFER`` kernels).
+        ``LocalAttention.backend`` is resolved once at module construction,
+        independent of whether a forward context is later set, so this only
+        catches a backend that failed to resolve to what the env var
+        requested -- it does not prove forward context was present for a
+        given forward call. Pair it with always wrapping the actual
+        transformer call in ``set_forward_context`` (see the denoising loops
+        below), which is what prevents the runtime fallback.
+
+        Scoped to ``_STRICT_BACKENDS``: other backend names can legitimately
+        differ from what ``LocalAttention.backend`` resolved to (silent,
+        logged, intentional fallback when unsupported by this layer -- see
+        ``_STRICT_BACKENDS``'s comment), so asserting on those would flag
+        expected behavior as a bug.
+        """
+        backend_env = envs.FASTVIDEO_ATTENTION_BACKEND
+        if not backend_env:
+            return
+        expected = backend_name_to_enum(backend_env)
+        if expected is None or expected not in self._STRICT_BACKENDS:
+            return
+        for module in self.transformer.modules():
+            if isinstance(module, LocalAttention):
+                assert module.backend == expected, (
+                    f"Kandinsky5 local attention resolved to backend {module.backend}, expected "
+                    f"{expected} from FASTVIDEO_ATTENTION_BACKEND={backend_env}. This likely means "
+                    "LocalAttention's missing-forward-context guard silently fell back to SDPA.")
+                return
+
+    def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -> torch.dtype:
+        """Resolve the transformer's actual compute dtype.
+
+        Trust ``pipeline_config.dit_precision`` directly for a normal,
+        non-FSDP load -- ``TransformerLoader.load()`` asserts every
+        parameter matches it exactly (``fastvideo/models/loader/
+        component_loader.py``), so this holds for standalone T2V/I2V
+        inference regardless of which precision (including fp32) was
+        requested.
+
+        FSDP2-wrapped modules are the one case where that assertion doesn't
+        carry forward: ``maybe_load_fsdp_model`` hardcodes its
+        ``MixedPrecisionPolicy`` to ``param_dtype=torch.bfloat16`` for every
+        FSDP-wrapped load, independent of ``dit_precision`` -- this covers
+        both the live transformer ``ValidationCallback`` reuses from
+        training (whose ``pipeline_config.dit_precision`` reflects the
+        fp32 master-weight load dtype, not the actual bf16 compute dtype)
+        and multi-GPU ``use_fsdp_inference=True`` runs. Detect the real
+        compute dtype from the parameters only in that case, mirroring the
+        Cosmos 2.5 denoising stage.
+        """
+        declared_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
+        try:
+            from torch.distributed.fsdp import FSDPModule
+        except Exception:  # pragma: no cover - FSDP not always available
+            return declared_dtype
+
+        if not isinstance(self.transformer, FSDPModule):
+            return declared_dtype
+
+        target_dtype = torch.bfloat16  # safe default
+        for p in self.transformer.parameters():
+            if p.dtype != torch.float32:
+                target_dtype = p.dtype
+                break
+        return target_dtype
+
     @staticmethod
     def fast_sta_nabla(
         T: int,
@@ -270,20 +364,10 @@ class Kandinsky5DenoisingStage(PipelineStage):
             loader = TransformerLoader()
             self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
             fastvideo_args.model_loaded["transformer"] = True
+        self._assert_local_attention_backend_engaged()
 
         device = get_local_torch_device()
-        # FSDP-wrapped models (e.g. the live transformer reused by
-        # ValidationCallback during training) may have their true compute
-        # dtype disagree with pipeline_config.dit_precision -- the training
-        # config forces dit_precision to "fp32" (master-weight load dtype)
-        # even though the actual FSDP2 mixed-precision policy runs the model
-        # in bf16. Detect the real weight dtype instead, mirroring the
-        # Cosmos 2.5 denoising stage.
-        target_dtype = torch.bfloat16  # safe default
-        for p in self.transformer.parameters():
-            if p.dtype != torch.float32:
-                target_dtype = p.dtype
-                break
+        target_dtype = self._resolve_target_dtype(fastvideo_args)
         autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
         latents = batch.latents
         num_channels = getattr(
@@ -444,15 +528,10 @@ class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage):
             loader = TransformerLoader()
             self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
             fastvideo_args.model_loaded["transformer"] = True
+        self._assert_local_attention_backend_engaged()
 
         device = get_local_torch_device()
-        # See the parent stage's forward() for why this can't be
-        # PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision].
-        target_dtype = torch.bfloat16  # safe default
-        for p in self.transformer.parameters():
-            if p.dtype != torch.float32:
-                target_dtype = p.dtype
-                break
+        target_dtype = self._resolve_target_dtype(fastvideo_args)
         autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
         latents = batch.latents
         num_channels = getattr(
@@ -499,9 +578,18 @@ class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage):
                     continue
 
                 t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
+                attn_metadata = None
+                if sparse_params is not None:
+                    attn_metadata = NablaAttentionMetadataBuilder().build(
+                        current_timestep=i,
+                        sta_mask=sparse_params["sta_mask"],
+                        P=sparse_params["P"],
+                        visual_shape=sparse_params["visual_shape"],
+                    )
                 autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
                                 if device.type == "cuda" else contextlib.nullcontext())
-                with autocast_ctx:
+                with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
+                                         forward_batch=batch), autocast_ctx:
                     pred_velocity = self.transformer(
                         hidden_states=latents.to(dtype=target_dtype),
                         encoder_hidden_states=prompt_embeds,

fastvideo/tests/api/test_kandinsky5_dmd_pipeline_resolution.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression test for exported Kandinsky5 DMD checkpoints resolving to the
+wrong pipeline.
+
+``fastvideo.train.entrypoint.dcp_to_diffusers`` copies the base T2V
+checkpoint's ``model_index.json`` unchanged into every export (see
+``examples/train/configs/fine_tuning/kandinsky5/README.md``), so
+``_class_name`` still says the base T2V pipeline. Without an explicit
+override, ``fastvideo.registry.get_model_info`` resolves such a directory to
+``Kandinsky5T2VPipeline`` (via the path-based "t2v" fallback detector in
+``fastvideo/registry.py``) instead of ``Kandinsky5DMDPipeline`` -- silently
+running the full-length multi-step sampler on DMD-distilled weights instead
+of the four-step re-noise sampler. ``override_pipeline_cls_name`` bypasses
+that resolution entirely and must be paired with ``Kandinsky5DMDConfig`` (not
+the default ``Kandinsky5T2VConfig``, whose ``dmd_denoising_steps=None`` makes
+``Kandinsky5DmdDenoisingStage`` raise immediately).
+"""
+
+import json
+
+from fastvideo.configs.pipelines.kandinsky5 import (
+    Kandinsky5DMDConfig,
+    Kandinsky5T2VConfig,
+)
+from fastvideo.pipelines.basic.kandinsky5.kandinsky5_dmd_pipeline import Kandinsky5DMDPipeline
+from fastvideo.pipelines.basic.kandinsky5.kandinsky5_pipeline import Kandinsky5T2VPipeline
+from fastvideo.pipelines.pipeline_registry import PipelineType
+from fastvideo.registry import get_model_info
+
+
+def _write_fake_export(tmp_path, name: str):
+    """Minimal on-disk stand-in for a dcp_to_diffusers export: just enough
+    for verify_model_config_and_directory to accept it."""
+    model_dir = tmp_path / name
+    (model_dir / "transformer").mkdir(parents=True)
+    (model_dir / "model_index.json").write_text(
+        json.dumps({
+            "_class_name": "Kandinsky5T2VPipeline",
+            "_diffusers_version": "0.30.0",
+        }))
+    return model_dir
+
+
+def test_kandinsky5_dmd_config_sets_denoising_steps():
+    assert Kandinsky5T2VConfig().dmd_denoising_steps is None
+    assert Kandinsky5DMDConfig().dmd_denoising_steps == [1000, 750, 500, 250]
+
+
+def test_exported_dmd_checkpoint_resolves_to_t2v_pipeline_without_override(tmp_path):
+    """Documents the bug: an unmodified export resolves to the T2V pipeline."""
+    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat")
+    info = get_model_info(model_path=str(model_dir), pipeline_type=PipelineType.BASIC)
+    assert info.pipeline_cls is Kandinsky5T2VPipeline
+
+
+def test_override_pipeline_cls_name_resolves_dmd_pipeline(tmp_path):
+    model_dir = _write_fake_export(tmp_path, "kandinsky5_t2v_dmd2_4steps_qat_override")
+    info = get_model_info(
+        model_path=str(model_dir),
+        pipeline_type=PipelineType.BASIC,
+        override_pipeline_cls_name="Kandinsky5DMDPipeline",
+    )
+    assert info.pipeline_cls is Kandinsky5DMDPipeline

fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py

@@ -49,6 +49,14 @@ KANDINSKY5_OVERFIT_DATA_DIR = DATA_DIR / "kandinsky5_overfit"
 PREPROCESSED_DIR = DATA_DIR / "kandinsky5_overfit_preprocessed"
 STAGE1_OUTPUT_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage1"
 STAGE2_OUTPUT_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage2"
+# Kept out of STAGE1_OUTPUT_DIR: that directory is scanned with
+# glob("checkpoint-*") below to find the latest raw DCP checkpoint, and a
+# "checkpoint-<N>-diffusers" export dir sitting alongside "checkpoint-<N>"
+# would also match that pattern -- on a rerun that doesn't start from a
+# clean data/ dir, it would sort after "checkpoint-<N>" and get picked as
+# stage1_ckpt instead, feeding an already-exported diffusers directory back
+# into dcp_to_diffusers as if it were a DCP checkpoint.
+STAGE1_DIFFUSERS_DIR = DATA_DIR / "outputs" / "kandinsky5_e2e_stage1_diffusers"
 
 STAGE1_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "fine_tuning" / "kandinsky5" /
                  "t2v_480p_qat.yaml")
@@ -101,6 +109,29 @@ def _run_preprocessing() -> None:
     subprocess.run(cmd, cwd=str(REPO_ROOT), check=True)
 
 
+def _export_dcp_to_diffusers(checkpoint_dir: Path, output_dir: Path) -> None:
+    """Convert a stage-1 DCP checkpoint to a diffusers-style directory.
+
+    Stage 2's ``models.*.init_from`` (like real-training YAML configs, see
+    ``dmd2_t2v_480p_qat.yaml``) expects a diffusers model directory
+    (``model_index.json`` + component subfolders), not the raw
+    ``checkpoint-N`` DCP layout ``CheckpointManager`` writes (``dcp/`` +
+    metadata/RNG state only). ``--verify`` strictly reloads the exported
+    transformer immediately so a key-mapping bug fails here, at the export
+    boundary, rather than deep inside the stage-2 launch below.
+    """
+    cmd = [
+        sys.executable, "-m",
+        "fastvideo.train.entrypoint.dcp_to_diffusers",
+        "--checkpoint", str(checkpoint_dir),
+        "--output-dir", str(output_dir),
+        "--role", "student",
+        "--overwrite",
+        "--verify",
+    ]
+    subprocess.run(cmd, cwd=str(REPO_ROOT), check=True)
+
+
 def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_overrides: list[str],
               env_overrides: dict[str, str]) -> None:
     cmd = [
@@ -109,7 +140,7 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov
         "--nproc_per_node", NUM_GPUS,
         "-m", "fastvideo.train.entrypoint.train",
         "--config", str(config),
-        "--training.data.data_path", str(PREPROCESSED_DIR / "combined_parquet_dataset"),
+        "--training.data.data_path", str(PREPROCESSED_DIR),
         "--training.distributed.num_gpus", NUM_GPUS,
         "--training.distributed.hsdp_shard_dim", NUM_GPUS,
         "--training.loop.max_train_steps", str(max_train_steps),
@@ -145,16 +176,21 @@ def test_e2e_kandinsky5_dmd_overfit_single_sample():
         f"no stage-1 checkpoint produced under {STAGE1_OUTPUT_DIR}")
 
     stage1_checkpoints = sorted(STAGE1_OUTPUT_DIR.glob("checkpoint-*"))
-    stage1_ckpt = str(stage1_checkpoints[-1])
+    stage1_ckpt = stage1_checkpoints[-1]
+
+    stage1_diffusers_dir = STAGE1_DIFFUSERS_DIR / stage1_ckpt.name
+    _export_dcp_to_diffusers(stage1_ckpt, stage1_diffusers_dir)
+    assert (stage1_diffusers_dir / "model_index.json").exists(), (
+        f"dcp_to_diffusers export did not produce a diffusers model dir at {stage1_diffusers_dir}")
 
     _run_stage(
         STAGE2_CONFIG,
         STAGE2_OUTPUT_DIR,
         max_train_steps=3,
         extra_overrides=[
-            "--models.student.init_from", stage1_ckpt,
-            "--models.teacher.init_from", stage1_ckpt,
-            "--models.critic.init_from", stage1_ckpt,
+            "--models.student.init_from", str(stage1_diffusers_dir),
+            "--models.teacher.init_from", str(stage1_diffusers_dir),
+            "--models.critic.init_from", str(stage1_diffusers_dir),
         ],
         env_overrides={"FASTVIDEO_ATTENTION_BACKEND": "ATTN_QAT_TRAIN"},
     )

fastvideo/tests/stages/test_kandinsky5_target_dtype.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression test: Kandinsky5 denoising stages must honor
+``pipeline_config.dit_precision`` for a plain (non-FSDP) transformer.
+
+``Kandinsky5DenoisingStage._resolve_target_dtype`` scans parameters and
+defaults to bf16 when every one is fp32 -- a heuristic needed only for the
+FSDP2-wrapped transformer ``ValidationCallback`` reuses from training
+(``maybe_load_fsdp_model`` hardcodes its ``MixedPrecisionPolicy`` to
+``param_dtype=torch.bfloat16`` independent of ``dit_precision``, so the
+live training module's true compute dtype can disagree with what
+``dit_precision`` claims). Applying that same scan unconditionally broke
+plain inference: ``TransformerLoader.load()`` asserts every parameter
+matches ``dit_precision`` exactly for a non-FSDP load, so an explicit
+fp32 pipeline (all params fp32, loop never finds a non-fp32 param) fell
+through to the "safe default" of bf16 -- silently running fp32-requested
+inference in bf16 instead.
+
+This is a pure logic test on a plain ``torch.nn.Module`` stand-in (never an
+``FSDPModule``) -- no GPU, no model load, no distributed init needed.
+"""
+from __future__ import annotations
+
+import types
+
+import torch
+
+from fastvideo.pipelines.stages.kandinsky5 import Kandinsky5DenoisingStage
+
+
+def _make_stage(transformer: torch.nn.Module) -> Kandinsky5DenoisingStage:
+    """Bypass __init__: only set the field _resolve_target_dtype reads."""
+    stage = Kandinsky5DenoisingStage.__new__(Kandinsky5DenoisingStage)
+    stage.transformer = transformer
+    return stage
+
+
+def _fastvideo_args(dit_precision: str) -> types.SimpleNamespace:
+    return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(dit_precision=dit_precision))
+
+
+def test_resolve_target_dtype_honors_explicit_fp32_for_plain_transformer():
+    transformer = torch.nn.Linear(4, 4).to(torch.float32)
+    stage = _make_stage(transformer)
+
+    resolved = stage._resolve_target_dtype(_fastvideo_args("fp32"))
+
+    assert resolved == torch.float32, (
+        "an explicit fp32 pipeline_config must not be silently cast to bf16 for a "
+        "plain (non-FSDP) transformer")
+
+
+def test_resolve_target_dtype_honors_explicit_fp16_for_plain_transformer():
+    transformer = torch.nn.Linear(4, 4).to(torch.float16)
+    stage = _make_stage(transformer)
+
+    resolved = stage._resolve_target_dtype(_fastvideo_args("fp16"))
+
+    assert resolved == torch.float16
+
+
+def test_resolve_target_dtype_honors_explicit_bf16_for_plain_transformer():
+    transformer = torch.nn.Linear(4, 4).to(torch.bfloat16)
+    stage = _make_stage(transformer)
+
+    resolved = stage._resolve_target_dtype(_fastvideo_args("bf16"))
+
+    assert resolved == torch.bfloat16

fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py

@@ -16,6 +16,13 @@ Two silent-fallback traps this guards against:
     the defensive pattern in
     ``fastvideo/models/loader/component_loader.py``.
 
+This test itself must not become that stale-selection source for whichever
+model test runs next in the same pytest process: ``FASTVIDEO_ATTENTION_BACKEND``
+is set via ``monkeypatch`` (auto-restored at teardown, including on
+failure) and the cache is cleared again in a ``finally`` block, since
+``monkeypatch`` only restores the env var, not the separate process-wide
+cache keyed on it.
+
 Unlike ``fastvideo.platforms.cuda``'s other backends, ATTN_QAT_TRAIN has no
 silent-fallback path at all if the kernel isn't built -- backend selection
 raises ``ImportError`` loudly. This test skips (does not fail) when the
@@ -58,7 +65,7 @@ _FIXTURE = str(
 
 
 @pytest.mark.usefixtures("distributed_setup")
-def test_kandinsky5_attn_qat_train_engages_and_backprops():
+def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch):
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
 
@@ -67,82 +74,89 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops():
     if not is_attn_qat_train_available():
         pytest.skip("fastvideo_kernel ATTN_QAT_TRAIN kernel is not built")
 
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_TRAIN"
+    monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN")
     _cached_get_attn_backend.cache_clear()
+    try:
+        from fastvideo.train.models.kandinsky5 import Kandinsky5Model
+        from fastvideo.train.utils.config import load_run_config
 
-    from fastvideo.train.models.kandinsky5 import Kandinsky5Model
-    from fastvideo.train.utils.config import load_run_config
+        cfg = load_run_config(_FIXTURE)
+        model = Kandinsky5Model(
+            init_from=cfg.models["student"]["init_from"],
+            training_config=cfg.training,
+            trainable=True,
+        )
 
-    cfg = load_run_config(_FIXTURE)
-    model = Kandinsky5Model(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-    )
+        device = torch.device("cuda:0")
+        dtype = torch.bfloat16
+        transformer = model.transformer.to(device=device, dtype=dtype)
 
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-    transformer = model.transformer.to(device=device, dtype=dtype)
+        attn = transformer.visual_transformer_blocks[0].self_attention
+        assert isinstance(attn.local_attention, LocalAttention)
+        assert attn.local_attention.backend == AttentionBackendEnum.ATTN_QAT_TRAIN, (
+            f"expected ATTN_QAT_TRAIN, got {attn.local_attention.backend} -- "
+            "backend selection silently fell back")
 
-    attn = transformer.visual_transformer_blocks[0].self_attention
-    assert isinstance(attn.local_attention, LocalAttention)
-    assert attn.local_attention.backend == AttentionBackendEnum.ATTN_QAT_TRAIN, (
-        f"expected ATTN_QAT_TRAIN, got {attn.local_attention.backend} -- "
-        "backend selection silently fell back")
+        arch = transformer.config.arch_config
+        patch_size = arch.patch_size
+        in_visual_dim = arch.in_visual_dim
+        in_text_dim = arch.in_text_dim
+        in_text_dim2 = arch.in_text_dim2
 
-    arch = transformer.config.arch_config
-    patch_size = arch.patch_size
-    in_visual_dim = arch.in_visual_dim
-    in_text_dim = arch.in_text_dim
-    in_text_dim2 = arch.in_text_dim2
+        grid_t, grid_h, grid_w = 2, 4, 4
+        latent_t = grid_t * patch_size[0]
+        latent_h = grid_h * patch_size[1]
+        latent_w = grid_w * patch_size[2]
 
-    grid_t, grid_h, grid_w = 2, 4, 4
-    latent_t = grid_t * patch_size[0]
-    latent_h = grid_h * patch_size[1]
-    latent_w = grid_w * patch_size[2]
+        latents = torch.randn(
+            1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype,
+            requires_grad=True)
+        if bool(getattr(transformer, "visual_cond", False)):
+            # See Kandinsky5Model._build_distill_input_kwargs /
+            # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints
+            # always expect [real | zero_cond | zero_mask] concatenated on the
+            # channel dim.
+            cond = torch.zeros_like(latents)
+            mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype)
+            hidden_states = torch.cat([latents, cond, mask], dim=-1)
+        else:
+            hidden_states = latents
+        encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype)
+        pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype)
+        timestep = torch.tensor([500], device=device, dtype=dtype)
+        visual_rope_pos = [
+            torch.arange(grid_t, device=device),
+            torch.arange(grid_h, device=device),
+            torch.arange(grid_w, device=device),
+        ]
+        text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device)
 
-    latents = torch.randn(
-        1, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=dtype,
-        requires_grad=True)
-    if bool(getattr(transformer, "visual_cond", False)):
-        # See Kandinsky5Model._build_distill_input_kwargs /
-        # Kandinsky5LatentPreparationStage: visual_cond=True checkpoints
-        # always expect [real | zero_cond | zero_mask] concatenated on the
-        # channel dim.
-        cond = torch.zeros_like(latents)
-        mask = torch.zeros(*latents.shape[:-1], 1, device=device, dtype=dtype)
-        hidden_states = torch.cat([latents, cond, mask], dim=-1)
-    else:
-        hidden_states = latents
-    encoder_hidden_states = torch.randn(1, 8, in_text_dim, device=device, dtype=dtype)
-    pooled_projections = torch.randn(1, in_text_dim2, device=device, dtype=dtype)
-    timestep = torch.tensor([500], device=device, dtype=dtype)
-    visual_rope_pos = [
-        torch.arange(grid_t, device=device),
-        torch.arange(grid_h, device=device),
-        torch.arange(grid_w, device=device),
-    ]
-    text_rope_pos = torch.arange(encoder_hidden_states.shape[1], device=device)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            out = transformer(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                pooled_projections=pooled_projections,
+                timestep=timestep,
+                visual_rope_pos=visual_rope_pos,
+                text_rope_pos=text_rope_pos,
+                scale_factor=(1.0, 2.0, 2.0),
+                sparse_params=None,
+                return_dict=True,
+            ).sample
 
-    with set_forward_context(current_timestep=0, attn_metadata=None):
-        out = transformer(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            pooled_projections=pooled_projections,
-            timestep=timestep,
-            visual_rope_pos=visual_rope_pos,
-            text_rope_pos=text_rope_pos,
-            scale_factor=(1.0, 2.0, 2.0),
-            sparse_params=None,
-            return_dict=True,
-        ).sample
+        assert torch.isfinite(out).all().item(), "output contains NaN/Inf"
 
-    assert torch.isfinite(out).all().item(), "output contains NaN/Inf"
-
-    out.sum().backward()
-    assert latents.grad is not None
-    assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf"
-    assert attn.to_query.weight.grad is not None
-    assert torch.isfinite(attn.to_query.weight.grad.to_local()
-                          if hasattr(attn.to_query.weight.grad, "to_local") else
-                          attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf"
+        out.sum().backward()
+        assert latents.grad is not None
+        assert torch.isfinite(latents.grad).all().item(), "input grad contains NaN/Inf"
+        assert attn.to_query.weight.grad is not None
+        assert torch.isfinite(attn.to_query.weight.grad.to_local()
+                              if hasattr(attn.to_query.weight.grad, "to_local") else
+                              attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf"
+    finally:
+        # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the
+        # selector cache is a separate process-wide functools.cache keyed on
+        # (head_size, dtype, supported_backends) -- not on the env var -- so
+        # it must be cleared again here or a later model test in this same
+        # pytest process could reuse this test's ATTN_QAT_TRAIN selection.
+        _cached_get_attn_backend.cache_clear()

fastvideo/tests/workflow/test_kandinsky5_overfit_caption.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression test: preprocess_kandinsky5_overfit.py must not silently
+truncate a string caption to its first character.
+
+The documented ``videos2caption.json`` schema
+(``docs/training/data_preprocess.md``) stores ``cap`` as a plain string
+(e.g. ``"Ocean waves at sunset..."``), while some producers (this repo's
+own e2e test fixtures, mirroring ``preprocess_hunyuan_overfit.py``) store a
+non-empty list of caption variants instead. Indexing unconditionally with
+``item["cap"][0]`` silently takes the first *character* of a string caption
+("O") instead of erroring -- ``get_caption`` accepts and validates both
+forms once at ingestion.
+
+Pure logic test -- no GPU, no model load.
+"""
+from __future__ import annotations
+
+import pytest
+
+from fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit import get_caption
+
+
+def test_get_caption_accepts_plain_string():
+    assert get_caption({"path": "a.mp4", "cap": "Ocean waves at sunset..."}) == "Ocean waves at sunset..."
+
+
+def test_get_caption_accepts_nonempty_list():
+    assert get_caption({"path": "a.mp4", "cap": ["a synthetic test clip"]}) == "a synthetic test clip"
+
+
+def test_get_caption_uses_first_list_element():
+    assert get_caption({"path": "a.mp4", "cap": ["first", "second"]}) == "first"
+
+
+@pytest.mark.parametrize("bad_cap", [
+    "",
+    [],
+    [""],
+    None,
+    123,
+    {"nested": "dict"},
+])
+def test_get_caption_rejects_invalid_values(bad_cap):
+    with pytest.raises(ValueError):
+        get_caption({"path": "a.mp4", "cap": bad_cap})

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -20,6 +20,10 @@ The checkpoint must contain ``metadata.json`` (written by
 ``CheckpointManager``).  If the checkpoint predates metadata
 support, pass ``--config`` explicitly to provide the training
 YAML.
+
+Pass ``--verify`` to strictly reload the exported transformer immediately
+after writing it, so a key-mapping bug fails here instead of deep inside a
+later training/inference launch that loads the exported directory.
 """
 
 from __future__ import annotations
@@ -200,6 +204,27 @@ def _save_role_pretrained(
     return str(dst)
 
 
+def _strict_reload_verify(*, output_dir: str, training_config: Any) -> None:
+    """Reload the just-exported transformer from disk and fail loudly on
+    any key mismatch.
+
+    ``TransformerLoader.load()`` already loads strictly (``strict=True``)
+    for every non-Cosmos25 model, so this doesn't add leniency -- it moves
+    the failure to right after export, with a clear "the export is broken"
+    error, instead of surfacing deep inside a later training/inference
+    launch that happens to load this directory.
+    """
+    from fastvideo.train.utils.moduleloader import load_module_from_path
+
+    logger.info("Verifying export: strictly reloading transformer from %s", output_dir)
+    load_module_from_path(
+        model_path=output_dir,
+        module_type="transformer",
+        training_config=training_config,
+    )
+    logger.info("Strict reload verification passed.")
+
+
 def convert(
     *,
     checkpoint_dir: str,
@@ -207,6 +232,7 @@ def convert(
     config_path: str | None = None,
     role: str = "student",
     overwrite: bool = False,
+    verify: bool = False,
 ) -> str:
     """Load a DCP checkpoint and export as a diffusers model.
 
@@ -298,6 +324,10 @@ def convert(
         model=model,
     )
     logger.info("Export complete: %s", result)
+
+    if verify:
+        _strict_reload_verify(output_dir=result, training_config=tc)
+
     return result
 
 
@@ -406,6 +436,13 @@ def main() -> None:
         action="store_true",
         help="Overwrite output-dir if it exists.",
     )
+    parser.add_argument(
+        "--verify",
+        action="store_true",
+        help=("After exporting, strictly reload the transformer from "
+              "the exported directory to catch key-mapping bugs "
+              "immediately."),
+    )
     args = parser.parse_args(sys.argv[1:])
 
     convert(
@@ -414,6 +451,7 @@ def main() -> None:
         config_path=args.config,
         role=args.role,
         overwrite=args.overwrite,
+        verify=args.verify,
     )