[bugfix]: H3 encoder truncation review fixes: parity gate, override validation, key-filter bound

Review follow-ups for the layer-50 truncation of the Qwen3-VL conditioner:

- tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py compared all
  65 hidden states with equal lengths and zip(strict=True), so the gate
  failed deterministically once the production stack builds only 51. It now
  pins the production tuple to language_model.num_layers + 1 entries and
  compares every built state bit-exactly against the official value at the
  same index, which still catches a truncated stack that applies the final
  norm, and degrades to the original full comparison when the override is
  None.

- num_hidden_layers_override <= 0 was accepted: 0 built an empty decoder
  stack, and a negative value additionally left num_layers inconsistent with
  the built layers, making _is_above_the_tap drop every layer key so the
  conditioner loaded "successfully" with no transformer at all. The arch
  config __post_init__ now rejects non-positive overrides, both at
  construction and when update_model_arch re-validates a checkpoint config.

- _is_above_the_tap dropped any layer index >= num_layers, including indexes
  at or above the checkpoint's own num_hidden_layers that the full stack
  would reject as unexpected. The filter is now bounded to
  [num_layers, num_hidden_layers) so truncation only drops keys the full
  model would have accepted.

- The CPU truncation test now compares the whole shared hidden-state prefix,
  mirroring the production-loader parity gate, and covers the
  override == num_hidden_layers boundary. The conditioner's
  num_hidden_layers property documents that it reports the checkpoint's
  nominal depth, not the built depth.

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

examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh

@@ -19,16 +19,13 @@ export FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1
 export FASTVIDEO_TORCH_PROFILER_WITH_STACK=1  
 export FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1
 export FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1
-export FASTVIDEO_TORCH_PROFILER_WAIT_STEPS=2
-export FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS=1
-export FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS=1
 export FASTVIDEO_TORCH_PROFILER_DIR="../profiler_traces/wan_t2v_finetune/"
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_t2v_finetune"
   --output_dir "checkpoints/wan_t2v_finetune"
-  --max_train_steps 5000
+  --max_train_steps 1  # Profiler captures every selected region until shutdown.
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 8

fastvideo/configs/models/encoders/minimax_h3_qwen3_vl.py

@@ -62,6 +62,15 @@ class MiniMaxH3Qwen3VLArchConfig(TextEncoderArchConfig):
     hidden_size: int = 5120
     intermediate_size: int = 25600
     num_hidden_layers: int = 64
+    # H3 conditions on one intermediate hidden state and reads nothing above it,
+    # so the remaining layers are built, weight-loaded and then discarded: 14
+    # layers, 13.7 GB in bf16. Building exactly this many leaves that hidden
+    # state bit-identical, because the tuple records each layer's *input*, so
+    # entry N is the output of layer N-1. Set to None to keep the full stack.
+    # Must equal MINIMAX_H3_TEXT_ENCODER_LAYER in
+    # fastvideo/pipelines/basic/minimax_h3/packing.py; a test pins them together
+    # rather than importing across the models -> pipelines boundary.
+    num_hidden_layers_override: int | None = 50
     num_attention_heads: int = 64
     num_key_value_heads: int = 8
     head_dim: int = 128
@@ -118,6 +127,16 @@ class MiniMaxH3Qwen3VLArchConfig(TextEncoderArchConfig):
     ])
 
     def __post_init__(self) -> None:
+        # Runs both at construction and after ``update_model_arch`` merges the
+        # checkpoint's config.json, so it also guards config-file overrides. A
+        # non-positive override would build no decoder layers at all, and a
+        # negative one would additionally make the surplus-key filter drop
+        # every ``language_model.layers.*`` checkpoint key, so the conditioner
+        # would "load" with no transformer stack and only fail at generation.
+        if self.num_hidden_layers_override is not None and self.num_hidden_layers_override < 1:
+            raise ValueError("MiniMax H3 Qwen3-VL num_hidden_layers_override must be a positive layer count "
+                             f"or None for the full stack; got {self.num_hidden_layers_override}.")
+
         rope_scaling = dict(self.rope_scaling or {})
         self.mrope_interleaved = bool(rope_scaling.get("mrope_interleaved", self.mrope_interleaved))
         if not self.mrope_interleaved:

fastvideo/envs.py

@@ -30,11 +30,8 @@ if TYPE_CHECKING:
     FASTVIDEO_TORCH_PROFILER_DIR: str | None = None
     FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
-    FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = True
+    FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = False
     FASTVIDEO_TORCH_PROFILER_WITH_FLOPS: bool = False
-    FASTVIDEO_TORCH_PROFILER_WAIT_STEPS: int = 2
-    FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS: int = 1
-    FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS: int = 2
     FASTVIDEO_TORCH_PROFILE_REGIONS: str = ""
     FASTVIDEO_TRACE_ACTIVATIONS: bool = False
     FASTVIDEO_TRACE_LAYERS: str = ""
@@ -242,11 +239,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY":
     lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"),
 
-    # Enable torch profiler to profile stack if set
-    # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. If not set, torch profiler WILL
-    # profile stack by default.
+    # Enable torch profiler stack capture with
+    # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. Off by default: stack capture
+    # costs ~1.5x runtime overhead and ~1.4x trace size.
     "FASTVIDEO_TORCH_PROFILER_WITH_STACK":
-    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "1") != "0"),
+    lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "0") != "0"),
 
     # Enable torch profiler to profile flops if set
     # FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will
@@ -255,16 +252,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
     # Wait steps per profiling cycle (torch.profiler.schedule wait parameter)
     # Defaults to 2 if not set.
-    "FASTVIDEO_TORCH_PROFILER_WAIT_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WAIT_STEPS", "2")),
     # Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter)
     # Defaults to 1 if not set.
-    "FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS", "1")),
     # Active steps per profiling cycle (torch.profiler.schedule active parameter)
     # Defaults to 2 if not set.
-    "FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS":
-    lambda: int(os.getenv("FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS", "2")),
     "FASTVIDEO_TORCH_PROFILE_REGIONS":
     lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
 

fastvideo/models/encoders/minimax_h3_qwen3_vl.py

@@ -227,10 +227,21 @@ class MiniMaxH3Qwen3VLLanguageModel(nn.Module):
             org_num_embeddings=config.vocab_size,
             quant_config=quant_config,
         )
+        # Build only as far as the consumer reads. The hidden-state tuple records
+        # each layer's input, so stopping after N layers still yields entry N,
+        # the output of layer N-1, unchanged. Everything above it exists only to
+        # feed `last_hidden_state`, which nothing consumes.
+        override = config.num_hidden_layers_override
+        self.num_layers = (config.num_hidden_layers
+                           if override is None else min(config.num_hidden_layers, override))
         self.layers = nn.ModuleList(
             MiniMaxH3Qwen3VLTextDecoderLayer(config, prefix=f"{config.prefix}.language_model.layers.{index}")
-            for index in range(config.num_hidden_layers))
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            for index in range(self.num_layers))
+        # The final norm sits above the tapped layer, so a truncated stack drops
+        # it. Keeping it would overwrite the tapped entry with a normalised
+        # tensor and change conditioning without raising anything.
+        self.norm = (RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+                     if self.num_layers == config.num_hidden_layers else None)
         self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config)
 
     def forward(
@@ -258,7 +269,10 @@ class MiniMaxH3Qwen3VLLanguageModel(nn.Module):
                 visual = deepstack_visual_embeds[layer_index].to(hidden_states.device, hidden_states.dtype)
                 updated = hidden_states[mask].clone() + visual
                 hidden_states[mask] = updated
-        hidden_states = self.norm(hidden_states)
+        if self.norm is not None:
+            hidden_states = self.norm(hidden_states)
+        # Truncated or not, the last entry is appended here, so the tapped index
+        # lands in the same place either way.
         if all_hidden_states is not None:
             all_hidden_states += (hidden_states, )
         return BaseEncoderOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
@@ -516,6 +530,13 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder):
 
     @property
     def num_hidden_layers(self) -> int:
+        """The checkpoint architecture's nominal depth, matching its config.json.
+
+        When ``num_hidden_layers_override`` truncates the stack at the
+        conditioning tap, fewer layers exist; the built count is
+        ``self.language_model.num_layers``, and the hidden-state tuple has
+        ``num_layers + 1`` entries, not ``num_hidden_layers + 1``.
+        """
         return self.config.num_hidden_layers
 
     def _get_rope_index(
@@ -694,6 +715,31 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder):
         outputs.attention_mask = attention_mask
         return outputs
 
+    def _is_above_the_tap(self, name: str) -> bool:
+        """Whether this checkpoint key belongs to a layer we did not build.
+
+        A truncated language stack still ships every layer in the checkpoint, and
+        the unexpected-key check below is strict on purpose, so the surplus keys
+        have to be dropped here rather than by relaxing it.
+        """
+        language_model = self.language_model
+        # The final norm is dropped exactly when the stack is truncated, so its
+        # absence is the signal.
+        if language_model.norm is not None:
+            return False
+        if name == "language_model.norm.weight":
+            return True
+        prefix = "language_model.layers."
+        if not name.startswith(prefix):
+            return False
+        index = name[len(prefix):].split(".", 1)[0]
+        if not index.isdigit():
+            return False
+        # Only drop indexes the full stack would have built. Anything at or
+        # above the checkpoint's own num_hidden_layers is corrupt and must
+        # still raise below, exactly as it does without truncation.
+        return language_model.num_layers <= int(index) < self.config.num_hidden_layers
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         parameters = dict(self.named_parameters())
         loaded: set[str] = set()
@@ -702,6 +748,8 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder):
             if source_name == "lm_head.weight":
                 continue
             name = source_name[6:] if source_name.startswith("model.") else source_name
+            if self._is_above_the_tap(name):
+                continue
             if name not in parameters:
                 raise ValueError(f"Unexpected MiniMax-H3 Qwen3-VL checkpoint key: {source_name}")
             parameter = parameters[name]

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import contextlib
 from typing import Any
 
 import torch
@@ -12,7 +11,7 @@ from fastvideo.attention.selector import component_attention_backend, get_attn_b
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
-from fastvideo.profiler import get_global_controller
+from fastvideo.profiler import profiler_region
 from fastvideo.hooks.activation_trace import trace_step
 from fastvideo.pipelines.basic.minimax_h3.packing import (
     MINIMAX_H3_KEYFRAME_NOISE_AUG,
@@ -147,11 +146,8 @@ class MiniMaxH3DenoisingStage(PipelineStage):
             vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ()))
             vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0))
 
-        controller = get_global_controller()
-        denoise_region = (controller.region("profiler_region_inference_denoising")
-                          if controller is not None else contextlib.nullcontext())
         try:
-            with denoise_region:
+            with profiler_region("inference_denoising"):
                 for index, (video_timestep,
                             audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)):
                     unique_timesteps, timestep_indices = row_timestep_plan[index]

fastvideo/pipelines/composed_pipeline_base.py

@@ -94,7 +94,6 @@ class ComposedPipelineBase(ABC):
         # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace
         trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR
         self.profiler_controller = get_or_create_profiler(trace_dir)
-        self.profiler = self.profiler_controller.profiler
 
         self.local_rank = get_world_group().local_rank
 
@@ -498,17 +497,6 @@ class ComposedPipelineBase(ABC):
         self._stage_name_mapping[stage_name] = stage
         setattr(self, stage_name, stage)
 
-    def profile(self, is_start: bool = True):
-        if self.profiler is None:
-            raise RuntimeError("Profiler is not enabled.")
-        if is_start:
-            self.profiler.start()
-        else:
-            self.profiler.stop()
-            # only print profiler results on rank 0
-            if self.local_rank == 0:
-                print(self.profiler.key_averages().table(sort_by="self_cuda_time_total"))
-
     # TODO(will): don't hardcode no_grad
     @torch.no_grad()
     def forward(

fastvideo/profiler.py

@@ -5,8 +5,9 @@ The profiler is shared across the process; this module adds a light-weight
 controller that gates collection based on named *regions*. Regions may be
 enabled through dedicated environment variables (e.g.
 ``FASTVIDEO_TORCH_PROFILE_MODEL_LOADING=1``) or via the consolidated
-``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list (e.g.
-``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit``).
+``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list. Short names work
+(``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train`` resolves the
+``profiler_region_`` prefix automatically).
 
 Typical usage from client code::
 
@@ -34,7 +35,6 @@ from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
 
-_GLOBAL_PROFILER: torch.profiler.profile | None = None
 _GLOBAL_CONTROLLER: TorchProfilerController | None = None
 
 
@@ -44,7 +44,6 @@ class ProfilerRegion:
 
     name: str
     description: str
-    default_enabled: bool = False
 
     def __post_init__(self) -> None:
         if not self.name or self.name.strip() != self.name:
@@ -63,8 +62,6 @@ def _normalize_token(token: str) -> str:
 def register_profiler_region(
     name: str,
     description: str,
-    *,
-    default_enabled: bool = False,
 ) -> None:
     """Register a profiler region so configuration can validate inputs."""
 
@@ -75,16 +72,22 @@ def register_profiler_region(
     region = ProfilerRegion(
         name=canonical,
         description=description,
-        default_enabled=bool(default_enabled),
     )
     _REGISTERED_REGIONS[canonical] = region
 
 
 def resolve_profiler_region(name: str) -> ProfilerRegion | None:
-    """Return the registered region matching ``name`` or ``None`` if absent."""
+    """Return the registered region for ``name`` (long or short form).
+
+    Accepts both the canonical name and the short form without the
+    ``profiler_region_`` prefix, so ``REGIONS=inference_denoising`` works.
+    """
 
     canonical = _normalize_token(name)
-    return _REGISTERED_REGIONS.get(canonical)
+    region = _REGISTERED_REGIONS.get(canonical)
+    if region is None and not canonical.startswith("profiler_region_"):
+        region = _REGISTERED_REGIONS.get(f"profiler_region_{canonical}")
+    return region
 
 
 def list_profiler_regions() -> list[ProfilerRegion]:
@@ -99,17 +102,6 @@ _DEFAULT_ACTIVITIES: tuple[torch.profiler.ProfilerActivity, ...] = (
 )
 
 
-def get_global_profiler() -> torch.profiler.profile | None:
-    """Return the global profiler instance if one was created."""
-
-    return _GLOBAL_PROFILER
-
-
-def set_global_profiler(profiler: torch.profiler.profile | None) -> None:
-    global _GLOBAL_PROFILER
-    _GLOBAL_PROFILER = profiler
-
-
 def get_global_controller() -> TorchProfilerController | None:
     return _GLOBAL_CONTROLLER
 
@@ -122,7 +114,6 @@ def set_global_controller(controller: TorchProfilerController | None) -> None:
 register_profiler_region(
     name="profiler_region_model_loading",
     description="Module/model loading during pipeline initialization.",
-    default_enabled=False,
 )
 # register_profiler_region(
 #     name="profiler_region_inference_pre_denoising",
@@ -149,11 +140,11 @@ register_profiler_region(
 )
 register_profiler_region(
     name="profiler_region_training_train_one_step",
-    description="High-level step orchestration in the training loop.",
+    description="Single optimizer step including forward/backward passes.",
 )
 register_profiler_region(
     name="profiler_region_training_train",
-    description="Single optimizer step including forward/backward passes.",
+    description="High-level step orchestration in the training loop.",
 )
 
 # distillation specific regions
@@ -204,15 +195,18 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
         profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
         with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
         with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
-        schedule=torch.profiler.schedule(
-            wait=envs.FASTVIDEO_TORCH_PROFILER_WAIT_STEPS,
-            warmup=envs.FASTVIDEO_TORCH_PROFILER_WARMUP_STEPS,
-            active=envs.FASTVIDEO_TORCH_PROFILER_ACTIVE_STEPS,
-        ),
+        # No schedule: nothing in the codebase calls profiler.step(), so a
+        # wait/warmup schedule never advances and the profiler records nothing.
+        # Region toggling gates collection; the single trace exports at stop().
         on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True),
     )
     controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES)
+    controller._trace_dir = trace_dir
     controller.start()
+    # The trace only exports at stop(); inference paths have no shutdown hook
+    # that calls it, so register one. stop() is idempotent.
+    import atexit
+    atexit.register(controller.stop)
     logger.info("Torch profiler started")
     return controller
 
@@ -287,16 +281,16 @@ class TorchProfilerController:
         FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit \
         python fastvideo/training/wan_training_pipeline.py ...
 
-    Wrapping a code block in a custom region::
+    Wrapping a code block in a registered region::
 
-        controller = TorchProfilerController(profiler, activities)
-        with controller.region("training_validation"):
-            run_validation_epoch()
+        from fastvideo.profiler import profiler_region
 
-    Adding a new region requires three steps:
-      1. Define an env var in ``envs.py``.
-      2. Add a default entry to ``register_profiler_region`` in this module.
-      3. Wrap the target code in :meth:`region` using the new name.
+        with profiler_region("inference_denoising"):
+            run_denoising_loop()
+
+    Adding a new region requires two steps:
+      1. Register it with ``register_profiler_region`` in this module.
+      2. Wrap the target code in :func:`profiler_region` using the new name.
     """
 
     def __init__(
@@ -317,10 +311,13 @@ class TorchProfilerController:
         self._profiler = profiler
         self._activities = activities_tuple
         self._config = config or TorchProfilerConfig.from_env()
-        self._collection_enabled = False
+        # torch.profiler collects from start(); reflect that so the initial
+        # _set_collection(False) in start() actually toggles it off instead of
+        # short-circuiting (which captured everything before the first region).
+        self._collection_enabled = True
         self._active_region_depth = 0
+        self._trace_dir: str | None = None
         logger.info("PROFILER: TorchProfilerController initialized with config: %s", self._config)
-        set_global_profiler(self._profiler)
         set_global_controller(self)
 
     @property
@@ -336,7 +333,10 @@ class TorchProfilerController:
 
         if self._profiler is None:
             return False
-        return self._config.regions.get(region, False)
+        resolved = resolve_profiler_region(region)
+        if resolved is None:
+            return False
+        return self._config.regions.get(resolved.name, False)
 
     def _set_collection(self, enabled: bool) -> None:
         if self._profiler is None:
@@ -348,6 +348,8 @@ class TorchProfilerController:
             self._profiler.toggle_collection_dynamic(enabled, self._activities)
         self._collection_enabled = enabled
 
+    _warned_unregistered: set[str] = set()
+
     @contextlib.contextmanager
     def region(self, region: str):
         """Context manager that enables profiling for ``region`` if configured."""
@@ -356,24 +358,41 @@ class TorchProfilerController:
             yield
             return
 
+        if resolve_profiler_region(region) is None:
+            # a typo here would otherwise silently profile nothing, forever
+            if region not in self._warned_unregistered:
+                self._warned_unregistered.add(region)
+                logger.warning("PROFILER: region %r is not registered (typo?); available: %s", region,
+                               ", ".join(r.name for r in list_profiler_regions()))
+            yield
+            return
+
         if not self.is_region_enabled(region):
             yield
             return
 
-        with torch.profiler.record_function(f"fastvideo.region::{region}"):
-            self._active_region_depth += 1
-            if self._active_region_depth == 1:
-                logger.info("PROFILER: Setting collection to True (depth=%s) for region %s", self._active_region_depth,
-                            region)
-                self._set_collection(True)
-            try:
+        # NVTX range so the same region names are visible in nsys timelines
+        nvtx = torch.cuda.is_available()
+        if nvtx:
+            torch.cuda.nvtx.range_push(f"fastvideo.region::{region}")
+        self._active_region_depth += 1
+        if self._active_region_depth == 1:
+            logger.info("PROFILER: Setting collection to True (depth=%s) for region %s", self._active_region_depth,
+                        region)
+            self._set_collection(True)
+        try:
+            # record_function opens after collection is enabled so the region
+            # marker itself lands in the trace.
+            with torch.profiler.record_function(f"fastvideo.region::{region}"):
                 yield
-            finally:
-                self._active_region_depth -= 1
-                logger.info("PROFILER: Decreasing active region depth to %s", self._active_region_depth)
-                if self._active_region_depth == 0:
-                    logger.info("PROFILER: Setting collection to False upon exiting region %s", region)
-                    self._set_collection(False)
+        finally:
+            self._active_region_depth -= 1
+            logger.info("PROFILER: Decreasing active region depth to %s", self._active_region_depth)
+            if self._active_region_depth == 0:
+                logger.info("PROFILER: Setting collection to False upon exiting region %s", region)
+                self._set_collection(False)
+            if nvtx:
+                torch.cuda.nvtx.range_pop()
 
     def start(self) -> None:
         """Start the profiler and pause collection until a region is entered."""
@@ -388,6 +407,34 @@ class TorchProfilerController:
         self._set_collection(False)
         logger.info("PROFILER: Profiler started with collection disabled")
 
+    def _write_summary(self) -> None:
+        """Compact per-rank op summary next to the trace: a key_averages table
+        and a JSON with input shapes, so operator-split analysis does not
+        require parsing multi-GB chrome traces."""
+        if self._profiler is None or not self._trace_dir:
+            return
+        try:
+            import json as _json
+            import os as _os
+            rank = _os.environ.get("RANK", "0")
+            averages = self._profiler.key_averages(group_by_input_shape=True)
+            stem = _os.path.join(self._trace_dir, f"summary_rank{rank}")
+            with open(f"{stem}.txt", "w") as fh:
+                fh.write(averages.table(sort_by="self_device_time_total", row_limit=60))
+            rows = [{
+                "name": e.key,
+                "shapes": str(e.input_shapes),
+                "self_device_us": e.self_device_time_total,
+                "device_us": e.device_time_total,
+                "count": e.count,
+            } for e in averages]
+            with open(f"{stem}.json", "w") as fh:
+                _json.dump(rows, fh)
+            if rank == "0":
+                logger.info("PROFILER: summary written to %s.txt", stem)
+        except Exception:  # noqa: BLE001 -- summaries must never break shutdown
+            logger.exception("PROFILER: summary generation failed")
+
     def stop(self) -> None:
         """Stop the profiler after disabling collection and clearing state."""
 
@@ -396,9 +443,10 @@ class TorchProfilerController:
 
         logger.info("PROFILER: Stopping profiler...")
         self._profiler.stop()
+        self._write_summary()
+        self._profiler = None  # makes stop() idempotent (atexit may re-enter)
         logger.info("PROFILER: Profiler stopped")
         self._active_region_depth = 0
-        set_global_profiler(None)
         set_global_controller(None)
 
     @property
@@ -411,9 +459,20 @@ class TorchProfilerController:
     def activities(self) -> tuple[torch.profiler.ProfilerActivity, ...]:
         return tuple(self._activities)
 
-    @property
-    def profiler(self) -> torch.profiler.profile | None:
-        return self._profiler
+
+@contextlib.contextmanager
+def profiler_region(region: str):
+    """Module-level region context manager bound to the global controller.
+
+    A no-op when no profiler is configured — stages can use this without the
+    get_global_controller()/nullcontext dance.
+    """
+    controller = get_global_controller()
+    if controller is None:
+        yield
+        return
+    with controller.region(region):
+        yield
 
 
 def profile_region(region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:

fastvideo/tests/contract/test_profiler_regions.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contract tests for the profiler region system.
+
+The region system's failure mode is silent: a profiler that records nothing
+(the old wait/warmup schedule bug), a trace that never exports (no shutdown
+hook), or a typo'd region name that no-ops forever. Each test here runs a
+child process — the profiler is process-global and its export happens at
+atexit, so in-process tests can't observe the contract.
+"""
+from __future__ import annotations
+
+import glob
+import gzip
+import json
+import os
+import subprocess
+import sys
+
+# Five-window child: ops before any region, inside a region, between regions,
+# inside a second (short-named) region, after the last region. Exits without
+# calling stop() — export must happen via the atexit hook. Each window is
+# labeled with a record_function marker so the parent can check which windows
+# the trace actually captured.
+_CHILD = r"""
+import torch
+from fastvideo.profiler import get_or_create_profiler, profiler_region
+
+controller = get_or_create_profiler("{trace_dir}")
+
+def burn(tag):
+    with torch.profiler.record_function(tag):
+        torch.mm(torch.ones(8, 8), torch.ones(8, 8))
+
+burn("win_pre")
+with profiler_region("profiler_region_inference_denoising"):
+    burn("win_region1")
+burn("win_between")
+with profiler_region("model_loading"):  # short name must resolve
+    burn("win_region2")
+for _ in range(2):  # warn-once: second use must not log again
+    with profiler_region("definitely_not_a_region"):
+        burn("win_typo")
+burn("win_post")
+# no controller.stop(): atexit owns the export
+"""
+
+
+def _run_child(tmp_path):
+    trace_dir = str(tmp_path / "traces")
+    env = os.environ.copy()
+    env["FASTVIDEO_TORCH_PROFILER_DIR"] = trace_dir
+    env["FASTVIDEO_TORCH_PROFILE_REGIONS"] = "inference_denoising,model_loading"
+    proc = subprocess.run(
+        [sys.executable, "-c", _CHILD.format(trace_dir=trace_dir)],
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=300,
+    )
+    assert proc.returncode == 0, proc.stderr
+    return trace_dir, proc.stdout + proc.stderr
+
+
+def _trace_event_names(trace_dir):
+    traces = glob.glob(os.path.join(trace_dir, "**", "*.json*"), recursive=True)
+    traces = [t for t in traces if "summary" not in os.path.basename(t)]
+    assert traces, f"no trace exported to {trace_dir} (atexit hook missing?)"
+    opener = gzip.open if traces[0].endswith(".gz") else open
+    with opener(traces[0], "rt") as fh:
+        events = json.load(fh).get("traceEvents", [])
+    return {e.get("name", "") for e in events}
+
+
+def test_regions_gate_collection_and_atexit_exports(tmp_path):
+    trace_dir, output = _run_child(tmp_path)
+    names = _trace_event_names(trace_dir)
+
+    assert "win_region1" in names, "op inside an enabled region was not captured"
+    assert "win_region2" in names, "short region name did not resolve/capture"
+    assert "fastvideo.region::profiler_region_inference_denoising" in names
+
+    for leaked in ("win_pre", "win_between", "win_typo", "win_post"):
+        assert leaked not in names, f"op outside any region leaked into trace: {leaked}"
+
+    # unregistered region warns exactly once across repeated uses
+    assert output.count("definitely_not_a_region") >= 1
+    assert output.count("is not registered") == 1
+
+    # per-rank op summary written next to the trace
+    summaries = glob.glob(os.path.join(trace_dir, "summary_rank0.*"))
+    assert sorted(os.path.splitext(s)[1] for s in summaries) == [".json", ".txt"]
+
+
+def test_noop_without_profiler_dir(tmp_path):
+    # profiler_region must be a clean no-op when profiling is not configured
+    child = (
+        "from fastvideo.profiler import profiler_region\n"
+        "with profiler_region('inference_denoising'):\n"
+        "    x = 1\n"
+        "assert x == 1\n"
+    )
+    env = os.environ.copy()
+    env.pop("FASTVIDEO_TORCH_PROFILER_DIR", None)
+    env.pop("FASTVIDEO_TORCH_PROFILE_REGIONS", None)
+    proc = subprocess.run([sys.executable, "-c", child], env=env,
+                          capture_output=True, text=True, timeout=300)
+    assert proc.returncode == 0, proc.stderr

fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The Qwen3-VL stack is built only as far as MiniMax H3 reads.
+
+H3 conditions on one intermediate hidden state. The layers above it were built,
+weight-loaded and then discarded, which is 13.7 GB in bf16 and the difference
+between fitting and not fitting on a 121 GB unified-memory device.
+
+The dangerous part is not the truncation, it is getting the tuple index wrong.
+`hidden_states` records each layer's *input*, so entry N is the output of layer
+N-1, and the final entry comes from the norm that sits above the whole stack. A
+truncated stack that still applies that norm puts a normalised tensor where the
+raw one belongs: the length check in the conditioning stage still passes, and
+conditioning silently changes. These tests pin the index, the content, and the
+constant the two sides agree on.
+"""
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+# Matches the other encoder tests: the module registry these build against wants
+# a process group, and a single-rank one needs a rendezvous address.
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29513")
+
+from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (
+    MiniMaxH3Qwen3VLArchConfig,
+    MiniMaxH3Qwen3VLConfig,
+)
+from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLLanguageModel
+from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_ENCODER_LAYER
+
+
+def _small_arch(**overrides) -> MiniMaxH3Qwen3VLArchConfig:
+    """A stack small enough to run on CPU but shaped like the real one.
+
+    Everything goes through the constructor so ``__post_init__`` validates the
+    small shape the same way it validates the real one.
+    """
+    kwargs: dict = dict(
+        vocab_size=64,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=8,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        head_dim=8,
+        # __post_init__ reads the sections out of rope_scaling, and they must
+        # cover exactly half of each head.
+        rope_scaling={
+            "mrope_interleaved": True,
+            "mrope_section": [2, 1, 1],
+            "rope_type": "default",
+        },
+        vision_out_hidden_size=16,
+    )
+    kwargs.update(overrides)
+    return MiniMaxH3Qwen3VLArchConfig(**kwargs)
+
+
+def _small_config(**overrides) -> MiniMaxH3Qwen3VLConfig:
+    """The outer config, which is what the modules take.
+
+    ``ModelConfig.__getattr__`` forwards the architecture fields, so the modules
+    read ``prefix`` off this object and everything else off ``arch_config``.
+    """
+    config = MiniMaxH3Qwen3VLConfig()
+    config.arch_config = _small_arch(**overrides)
+    return config
+
+
+def test_default_matches_the_index_the_pipeline_reads() -> None:
+    """The two sides cannot import each other, so pin them here instead.
+
+    `fastvideo/models/` must not import from `fastvideo/pipelines/`, so the tap
+    is written down twice. If they drift, conditioning reads a hidden state that
+    was never built and the run dies with an index error at generation time,
+    after a full model load.
+    """
+    assert MiniMaxH3Qwen3VLArchConfig().num_hidden_layers_override == MINIMAX_H3_TEXT_ENCODER_LAYER
+
+
+def test_builds_only_up_to_the_override(distributed_setup) -> None:
+    model = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=5))
+
+    assert model.num_layers == 5
+    assert len(model.layers) == 5
+    # The norm sits above the tap, so a truncated stack must not keep it.
+    assert model.norm is None
+
+
+def test_override_none_keeps_the_full_stack(distributed_setup) -> None:
+    model = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=None))
+
+    assert model.num_layers == 8
+    assert model.norm is not None
+
+
+def test_override_above_the_stack_does_not_over_build(distributed_setup) -> None:
+    # num_hidden_layers comes from the checkpoint's config.json via
+    # update_model_arch, so a smaller variant must clamp rather than ask for
+    # layers that do not exist.
+    model = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=99))
+
+    assert model.num_layers == 8
+    assert model.norm is not None
+
+
+def test_override_equal_to_the_stack_keeps_the_norm(distributed_setup) -> None:
+    """The exact boundary of the clamp: a stack cut at its own depth is full.
+
+    A checkpoint with exactly ``override`` layers taps its final layer, whose
+    tuple entry sits after the norm in the full model, so the norm must stay
+    and nothing may be filtered from the checkpoint.
+    """
+    model = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=8))
+
+    assert model.num_layers == 8
+    assert model.norm is not None
+
+
+def test_non_positive_override_is_rejected() -> None:
+    """A non-positive override would build no decoder layers at all.
+
+    Worse, a negative one makes ``num_layers`` disagree with the built stack
+    and the surplus-key filter would then drop every layer key, so the
+    conditioner would load "successfully" with no transformer. Reject it at
+    config construction, and again when update_model_arch re-validates.
+    """
+    for override in (0, -1):
+        with pytest.raises(ValueError, match="num_hidden_layers_override"):
+            _small_arch(num_hidden_layers_override=override)
+
+    config = _small_config()
+    with pytest.raises(ValueError, match="num_hidden_layers_override"):
+        config.update_model_arch({"num_hidden_layers_override": 0})
+
+
+def test_tapped_hidden_state_is_unchanged_by_truncation(distributed_setup) -> None:
+    """The whole point: entry `tap` must be bit-identical either way."""
+    tap = 5
+    full = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=None))
+    cut = MiniMaxH3Qwen3VLLanguageModel(_small_config(num_hidden_layers_override=tap))
+
+    # These modules allocate uninitialised storage and expect a checkpoint, so
+    # give them finite weights before running anything through them.
+    torch.manual_seed(0)
+    for parameter in full.parameters():
+        parameter.data.normal_(std=0.02)
+    # Then make the shared prefix identical, which is the only part the tapped
+    # hidden state depends on.
+    for (_, a), (_, b) in zip(full.layers[:tap].named_parameters(),
+                              cut.layers[:tap].named_parameters(),
+                              strict=True):
+        b.data.copy_(a.data)
+    torch.manual_seed(1)
+    inputs_embeds = torch.randn(1, 6, 16)
+    # mRoPE indexes three axes (t, h, w); text tokens share the same position on
+    # all three.
+    position_ids = torch.arange(6).view(1, 1, 6).expand(3, 1, 6)
+    with torch.no_grad():
+        full_out = full(inputs_embeds, position_ids, None, True, None, None)
+        cut_out = cut(inputs_embeds, position_ids, None, True, None, None)
+
+    assert torch.equal(full_out.hidden_states[tap], cut_out.hidden_states[tap])
+    # And the truncated model must not offer states it never computed.
+    assert len(cut_out.hidden_states) == tap + 1
+    # The whole shared prefix must match, not just the tap: this is the same
+    # comparison the production-loader parity gate runs against the official
+    # model, and it is what catches a truncated stack that still applied the
+    # final norm to its last entry.
+    for index, (cut_state, full_state) in enumerate(zip(cut_out.hidden_states, full_out.hidden_states,
+                                                        strict=False)):
+        assert torch.equal(cut_state, full_state), f"hidden state {index} changed under truncation"
+
+
+def test_truncated_model_drops_the_surplus_checkpoint_keys(distributed_setup) -> None:
+    """The unexpected-key check is strict on purpose, so the surplus keys have
+    to be filtered rather than the check relaxed."""
+    from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
+
+    conditioner = MiniMaxH3Qwen3VLConditioner(_small_config(num_hidden_layers_override=5))
+
+    assert conditioner._is_above_the_tap("language_model.layers.5.mlp.gate_proj.weight")
+    assert conditioner._is_above_the_tap("language_model.layers.7.self_attn.q_proj.weight")
+    assert conditioner._is_above_the_tap("language_model.norm.weight")
+    # Kept: layers we built, the embeddings, and the vision tower.
+    assert not conditioner._is_above_the_tap("language_model.layers.4.mlp.gate_proj.weight")
+    assert not conditioner._is_above_the_tap("language_model.embed_tokens.weight")
+    assert not conditioner._is_above_the_tap("visual.blocks.0.attn.qkv.weight")
+    # The filter only drops indexes the full stack would have built. A key at
+    # or above the checkpoint's own num_hidden_layers is corrupt, and it must
+    # keep raising as unexpected exactly as it does without truncation.
+    assert not conditioner._is_above_the_tap("language_model.layers.8.mlp.gate_proj.weight")
+    with pytest.raises(ValueError, match="Unexpected"):
+        conditioner.load_weights([("model.language_model.layers.8.mlp.gate_proj.weight", torch.zeros(1))])
+
+
+def test_full_stack_filters_nothing(distributed_setup) -> None:
+    from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner
+
+    conditioner = MiniMaxH3Qwen3VLConditioner(_small_config(num_hidden_layers_override=None))
+
+    assert not conditioner._is_above_the_tap("language_model.layers.7.mlp.gate_proj.weight")
+    assert not conditioner._is_above_the_tap("language_model.norm.weight")

tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py

@@ -5,6 +5,12 @@ The test compares the exact Transformers base model used by the official
 pipeline with FastVideo's production ``TextEncoderLoader`` path.  It covers
 the three numerical branches the H3 pipelines exercise: text-only tokens,
 image features, and video features.
+
+The production encoder is built only as far as the layer-50 conditioning tap
+by default (``num_hidden_layers_override``), so it returns fewer hidden states
+than the official full stack.  Every state it does build is compared
+bit-exactly against the official value at the same index, which pins the tap
+and would catch a truncated stack that still applied the final norm.
 """
 
 from __future__ import annotations
@@ -220,9 +226,20 @@ def test_minimax_h3_qwen3_vl_parity() -> None:
     actual = _run_cases(production, cases, device)
 
     assert actual.keys() == expected.keys()
+    # The production stack is built only as far as the conditioning tap by
+    # default (``num_hidden_layers_override``), so it yields one hidden state
+    # per built layer plus the embeddings, while the official model always
+    # yields the full tuple. Every state the production model produces must be
+    # bit-identical to the official value at the same index; the shared-prefix
+    # comparison would in particular catch a truncated stack that still
+    # applied the final norm, which is the failure mode that silently changes
+    # conditioning. With the override set to None the lengths are equal and
+    # this remains the original full comparison, final normed state included.
+    built_layers = int(production.language_model.num_layers)
     for name in expected:
-        assert len(actual[name]) == len(expected[name])
-        for layer, (result, reference) in enumerate(zip(actual[name], expected[name], strict=True)):
+        assert len(actual[name]) == built_layers + 1
+        assert len(actual[name]) <= len(expected[name])
+        for layer, (result, reference) in enumerate(zip(actual[name], expected[name], strict=False)):
             assert_close(result, reference, atol=0.0, rtol=0.0, msg=lambda message: f"{name} layer {layer}: {message}")
         result = actual[name][MINIMAX_H3_TEXT_ENCODER_LAYER]
         reference = expected[name][MINIMAX_H3_TEXT_ENCODER_LAYER]

tests/local_tests/minimax_h3/README.md

@@ -49,7 +49,10 @@ pytest \
 ```
 
 With a gate enabled, missing CUDA, source, or weights is a failure. Recorded component evidence is exact for both DiT
-partitions, the video VAE, and all Qwen3-VL hidden states; audio decode has maximum absolute drift `2.4e-7`.
+partitions, the video VAE, and all Qwen3-VL hidden states; audio decode has maximum absolute drift `2.4e-7`. The
+production Qwen3-VL stack is now built only to the layer-50 conditioning tap by default
+(`num_hidden_layers_override`), so the encoder gate compares every hidden state the production model builds
+bit-exactly against the official full stack at the same index.
 
 FastVideo joint audio/video generation and SP=1/SP=4 latent consistency have been validated. T2VA, FL2VA, and
 Ref2VA video/audio latents match the pinned Diffusers pipeline exactly.