Release 2026-10-24 (i4 b3866eb3)

cosmos_framework/auxiliary/guardrail/common/presets.py

@@ -9,12 +9,15 @@
 from cosmos_framework.utils import log
 
 
-def create_text_guardrail_runner(offload_model_to_cpu: bool = False) -> GuardrailRunner:
+def create_text_guardrail_runner(
+    offload_model_to_cpu: bool = False, *, raise_on_error: bool = False
+) -> GuardrailRunner:
     """Create the text guardrail runner."""
+    strict_options: dict[str, bool] = {"raise_on_error": True} if raise_on_error else {}
     return GuardrailRunner(
         safety_models=[
             Blocklist(),
-            Qwen3Guard(offload_model_to_cpu=offload_model_to_cpu),
+            Qwen3Guard(offload_model_to_cpu=offload_model_to_cpu, **strict_options),
         ]
     )
 

cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py

@@ -31,19 +31,23 @@ class Qwen3Guard(ContentSafetyGuardrail):
     model_id: str
     model: Any
     tokenizer: Any
+    raise_on_error: bool
 
     def __init__(
         self,
         offload_model_to_cpu: bool = True,
+        raise_on_error: bool = False,
     ) -> None:
         """Load Qwen3Guard for text filtering safety checks.
 
         Args:
             offload_model_to_cpu (bool, optional): Whether to offload the model to CPU. Defaults to True.
+            raise_on_error (bool): Propagate runtime errors for strict evaluation. Defaults to False.
         """
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
         self.offload_model = offload_model_to_cpu
+        self.raise_on_error = raise_on_error
         self.dtype = torch.bfloat16
 
         self.model_id = "Qwen/Qwen3Guard-Gen-0.6B"
@@ -95,6 +99,8 @@ def is_safe(self, prompt: str) -> tuple[bool, str]:
         try:
             return self.extract_label_and_categories(prompt)
         except Exception as e:
+            if self.raise_on_error:
+                raise
             log.error(f"Unexpected error occurred when running Qwen3Guard guardrail: {e}")
             return True, "Unexpected error occurred when running Qwen3Guard guardrail."
 

cosmos_framework/callbacks/data_stats_writer.py

@@ -52,7 +52,7 @@ def arrow(self) -> pa.Field:
 
 _FIELDS = (
     _Field("phase", pa.string(), 'Data split that produced the row: "train", "val", or "data_only".'),
-    _Field("step", pa.int64(), "Completed optimizer step associated with the sample."),
+    _Field("step", pa.int64(), "Completed training iteration or data-loading step associated with the sample."),
     _Field("rank", pa.int32(), "Distributed process rank that consumed the sample."),
     _Field("sample_rank", pa.int32(), "Rank recorded by the dataset distributor."),
     _Field("worker_id", pa.int32(), "Dataloader worker that selected the sample."),
@@ -111,7 +111,7 @@ def arrow(self) -> pa.Field:
     _Field("context_length", pa.int32(), "Configured model context length."),
     _Field("dataset_weight", pa.float64(), "Effective normalized sampling share of the dataset."),
     _Field("dataset_num_urls", pa.int64(), "Dataset URL count after recipe expansion."),
-    _Field("total_training_steps", pa.int64(), "Configured maximum optimizer iterations."),
+    _Field("total_training_steps", pa.int64(), "Configured maximum training or data-loading iterations."),
 )
 
 DATA_STATS_SCHEMA = pa.schema(
@@ -290,6 +290,9 @@ class DataStatsWriterCallback(Callback):
     training continues. The next save, the end of training, and resume each wait for that merge and stop the
     job if it cannot succeed, so rank files never accumulate unnoticed. Rows recorded after the final save,
     which can only come from the last validation, are not saved.
+
+    With ``trainer.skip_model_update``, train rows retain sample metadata and trainable-token counts;
+    CE and objective fields are null because no model forward pass ran.
     """
 
     def __init__(
@@ -349,6 +352,8 @@ def _record(
             return
         if not self._ready:
             raise RuntimeError("DataStatsWriterCallback was not initialized")
+        if int(data_batch.get("context_parallel_size", 1)) > 1:
+            raise NotImplementedError("Disable per-sample data-stats writing for CP; shard aggregation is required")
         metadata = data_batch.get("data_stats")
         if not isinstance(metadata, list) or not all(isinstance(sample_metadata, dict) for sample_metadata in metadata):
             raise TypeError(
@@ -357,6 +362,8 @@ def _record(
             )
         if data_batch.get("true_packing") or "cu_seq_lens_q" in data_batch:
             raise NotImplementedError("Per-sample data-stats logging currently requires padded batches")
+        if phase == "train" and getattr(self.config.trainer, "skip_model_update", False):
+            output_batch = None
         missing = [key for key in _LOSS_KEYS if output_batch is not None and key not in output_batch]
         if missing:
             raise KeyError(f"Model output is missing data-stats loss fields: {missing}")
@@ -374,7 +381,7 @@ def _record(
                 batch_values.pop("dataset_num_urls")
             _check_sample(keys[index], sample_metadata, batch_values)
         if output_batch is None:
-            labels = data_batch["labels"]
+            labels = data_batch["labels"]  # [B,T]
             ignore_indices = _values(data_batch["ignore_index"], size)
             valid_counts = [int((labels[index, 1:] != ignore_indices[index]).sum()) for index in range(size)]
             token_sums = numerators = weights = [None] * size

cosmos_framework/callbacks/dataloader_state.py

@@ -55,7 +55,7 @@ def __init__(
         self.lance_state: dict[int, dict[str, Any]] = {}
         self.val_lance_state: dict[int, dict[str, Any]] = {}
         self.val_lance_batches: int = 0
-        self._pending_lance_state: dict[str, Any] | None = None
+        self._pending_lance_state: dict[str, Any] | list[dict[str, Any]] | None = None
         self._pending_val_lance_state: dict[str, Any] | None = None
         self.verbose = True
 
@@ -75,7 +75,12 @@ def on_training_step_batch_start(
             return
         lance_state = data_batch.pop(LANCE_VLM_RESUME_STATE_KEY, None)
         if lance_state is not None and not isinstance(lance_state, dict):
-            raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict, got {type(lance_state).__name__}.")
+            if (
+                not isinstance(lance_state, list)
+                or not lance_state
+                or not all(isinstance(x, dict) for x in lance_state)
+            ):
+                raise TypeError(f"{LANCE_VLM_RESUME_STATE_KEY} must be a dict or a non-empty list of cursor dicts.")
         self._pending_lance_state = lance_state
 
     def _update_lance_state(self, update: dict[str, Any], split: str = "train") -> None:
@@ -168,7 +173,9 @@ def on_training_step_batch_end(
     ) -> None:
         if self.distributor_type == "no_replace":
             if self._pending_lance_state is not None:
-                self._update_lance_state(self._pending_lance_state)
+                updates = self._pending_lance_state
+                for update in updates if isinstance(updates, list) else [updates]:
+                    self._update_lance_state(update)
                 self._pending_lance_state = None
             else:
                 self._update_state_from_batch(self.state, data_batch)

cosmos_framework/callbacks/every_n_draw_sample.py

@@ -3,8 +3,8 @@
 
 import math
 import os
-from collections.abc import Callable
-from contextlib import nullcontext
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from functools import partial
 from typing import Any, Literal, NamedTuple
@@ -30,6 +30,7 @@
     TEXT_SYSTEM_PROMPT_KEY,
 )
 from cosmos_framework.model.generator.mot.context_parallel_utils import broadcast_context_parallel_object
+from cosmos_framework.model.generator.utils.moe_utils import preserve_moe_learning_accumulators
 from cosmos_framework.utils.generator.data_utils import slice_data_batch
 from cosmos_framework.utils.generator.multiview import (
     decode_multiview_latent_per_view,
@@ -50,6 +51,17 @@ class WandbAnimation(NamedTuple):
 WandbMedia = str | WandbAnimation | list[wandb.Object3D] | dict[str, "WandbMedia"]
 
 
+@contextmanager
+def _sampling_observation_scope(model: Any) -> Iterator[None]:
+    """Keep diagnostic sampling from contributing to the next router update."""
+    net = getattr(model, "net", None)
+    if not isinstance(net, torch.nn.Module):
+        yield
+        return
+    with preserve_moe_learning_accumulators(net):
+        yield
+
+
 def resize_image(image: torch.Tensor, size: int = 1024) -> torch.Tensor:
     """
     Resize the image to the given size. This is done so that wandb can display the image correctly.
@@ -947,7 +959,9 @@ def every_n_impl(
             )
 
         log.debug("entering, every_n_impl", rank0_only=False)
-        with context():
+        # Keep this scope outside ``ema_scope`` so its ``finally`` restores learning
+        # accumulators after EMA weights have been swapped back, including on exceptions.
+        with _sampling_observation_scope(model), context():
             if self.do_x0_prediction:
                 log.debug("entering, x0_pred", rank0_only=False)
                 x0_img_fp, mse_loss, sigmas = self.x0_pred(

cosmos_framework/callbacks/sampled_media_recorder.py

@@ -181,7 +181,7 @@ def _extract_records(self, data_batch: dict[str, Any], iteration: int, rank: int
         count = len(sample_ids)
         media_urls = _as_list(data_batch.get("__url__"), count)
         dataset_names = _as_list(data_batch.get("dataset_name"), count)
-        source_streams = _as_list(data_batch.get("transfer_source_stream"), count)
+        source_streams = _as_list(data_batch.get("task_family"), count)
         source_names = _as_list(data_batch.get("source_dataset_name"), count)
         captions = _as_optional_string_list(data_batch.get("ai_caption"), count) if self.record_caption else []
         source_ids = _as_list(data_batch.get("source_id"), count)

cosmos_framework/callbacks/sampled_media_recorder_test.py

@@ -59,12 +59,12 @@ def test_extract_records_from_consumed_image_batch(monkeypatch: pytest.MonkeyPat
     assert all(record["rank"] == 3 for record in records)
 
 
-def test_extract_records_preserves_mixed_transfer_source_stream() -> None:
+def test_extract_records_preserves_token_mixed_source_stream() -> None:
     callback = SampledMediaRecorder(enabled=True, output_uri="/tmp/samples.lance")
     batch = {
         "__key__": ["sample-a"],
         "dataset_name": ["transfer_rgb"],
-        "transfer_source_stream": ["transfer_depth_single_view_720"],
+        "task_family": ["transfer_depth_single_view_720"],
     }
 
     records = callback._extract_records(batch, iteration=7, rank=3)

cosmos_framework/callbacks/tokens_per_sec.py

@@ -114,6 +114,7 @@ def __init__(
         self._fwd_bwd_flop_coeff: float | None = fwd_bwd_flop_coeff
         self._flop_coeff_resolved: float | None = None  # cached auto-resolved coeff
         self._num_params: int | None = None  # lazily filled from model.parameters() (global/DTensor)
+        self._context_parallel_size: int = 1
         self._hit_counter: int = 0
         self._tokens_in_window: int = 0
         self._useful_tokens_in_window: int = 0
@@ -179,6 +180,7 @@ def on_training_step_batch_end(
         if self._hit_counter < self.hit_thres:
             return
 
+        self._context_parallel_size = int(data_batch.get("context_parallel_size", 1))
         ids = data_batch.get(self.length_key)
         if ids is not None:
             # numel() is the physical token count. shape[0] counts rows, while a true-packed
@@ -338,7 +340,8 @@ def _synchronize_window(self, elapsed: float) -> tuple[float, int, float, float,
             if dist.is_available() and dist.is_initialized() and dist.get_backend() == "nccl"
             else torch.device("cpu")
         )
-        local_useful_tps = self._useful_tokens_in_window / elapsed if elapsed > 0 else 0.0
+        cp_size = self._context_parallel_size
+        local_useful_tps = self._useful_tokens_in_window / cp_size / elapsed if elapsed > 0 else 0.0
         local_predicted_microbatch_ms = (
             self._predicted_runtime_ms_in_window / self._predicted_steps_in_window
             if self._predicted_steps_in_window
@@ -371,6 +374,10 @@ def _synchronize_window(self, elapsed: float) -> tuple[float, int, float, float,
             dtype=torch.float64,
             device=reduce_device,
         )
+        # The callback sees complete inputs replicated across CP ranks. Deduplicate
+        # token/sample/work totals, retaining physical-rank timing and step counts.
+        # Divide in FP64 before the global sum so small supervised counts stay exact.
+        sums[[0, 1, 2, 6, 7, 8, 14, 15]] /= cp_size
         peak_allocated_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
         peak_reserved_gb = torch.cuda.max_memory_reserved() / 1e9 if torch.cuda.is_available() else 0.0
         maxima = torch.tensor(  # [N_MAX]

cosmos_framework/callbacks/wall_clock_checkpoint.py

@@ -263,7 +263,7 @@ def __init__(
         # hook has since described different optimizer/scheduler components.
         self._pending_dcp_components: dict[int, tuple[str, ...]] = {}
 
-        # Captured from on_before_optimizer_step so we can call checkpointer.save().
+        # Captured before optimizer steps, or at train start for dataloading-only runs, to call checkpointer.save().
         self._optimizer: torch.optim.Optimizer | None = None
         self._scheduler: torch.optim.lr_scheduler.LRScheduler | None = None
         self._grad_scaler: torch.amp.GradScaler | None = None
@@ -309,6 +309,10 @@ def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
         if not self.enabled:
             log.info(f"[WallClockCheckpoint] Disabled; set {INTERVAL_MINUTES_ENV_VAR} or interval_minutes to enable.")
             return
+        if getattr(self.config.trainer, "skip_model_update", False):
+            self._optimizer = self.trainer.optimizer
+            self._scheduler = self.trainer.scheduler
+            self._grad_scaler = self.trainer.grad_scaler
         log.info(
             f"[WallClockCheckpoint] Enabled: saving at least every {self._interval_seconds / 60:.1f} minutes, "
             f"retaining the last {self._keep_last} wall-clock checkpoints."

cosmos_framework/checkpoint/base.py

@@ -8,11 +8,11 @@
 
 import torch
 
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig, ObjectStoreConfig
+from cosmos_framework.utils.flags import INTERNAL
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig, ObjectStoreConfig
 from cosmos_framework.utils.easy_io import easy_io
-from cosmos_framework.utils.flags import INTERNAL
 
 
 @dataclass(frozen=True, slots=True)

cosmos_framework/checkpoint/dcp.py

@@ -80,9 +80,9 @@
 )
 from cosmos_framework.checkpoint.base import AbstractCheckpointer, CheckpointLoadSource
 from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback, distributed, log, misc
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
 

cosmos_framework/checkpoint/dcp_distill.py

@@ -27,6 +27,9 @@
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.nn.modules.module import _IncompatibleKeys
 
+from cosmos_framework.model._base import ImaginaireModel
+from cosmos_framework.utils import log, misc
+from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.checkpoint.dcp import (
     AsyncMode,
     CustomLoadPlanner,
@@ -36,10 +39,7 @@
     DistributedCheckpointer as _DistributedCheckpointer,
 )
 from cosmos_framework.checkpoint.dcp import ModelWrapper as VFMModelWrapper
-from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.model.generator.distillation.optimizer import OptimizerContainerLike, is_optimizer_container
-from cosmos_framework.utils import log, misc
-from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
 
 __all__: tuple[str, ...] = (

cosmos_framework/configs/base/defaults/model_config.py

@@ -340,6 +340,8 @@ class OmniMoTModelConfig:
     # None preserves Auto's separate action controls and control/target item roles.
     # Opt-in batches use separate ragged RGB and full frame-rate action blocks;
     # single-view batches in the same run keep their original dense packing.
+    # This config controls packing/view codes; attention uses the resulting item
+    # layout and the shared multiview_attention settings above.
     multiview_action_conditioning: MultiviewActionConditioningConfig | None = None
 
     # Per-layer NATTEN parameters

cosmos_framework/configs/base/defaults/multiview_attention.py

@@ -538,6 +538,21 @@ class MultiviewAttentionConfig:
         default="other_views", validator=attrs.validators.in_(("all_views", "own_view", "other_views"))
     )
 
+    # Let a camera query's cross-view pass also key the other views' camera-pose action
+    # tokens, beside their vision tokens, under the same capture-time rule. Action tokens stay
+    # keys only: they never query the cross-view pass, and their own view's actions remain in
+    # the same-view pass alone. Requires maskless decomposed attention, no cross-view band
+    # expansion, no same-view temporal window and baseline geometry. False preserves existing
+    # recipes.
+    maskless_cross_view_attends_action: bool = False
+
+    # Let camera-pose action tokens query the cross-view pass too, keyed only by the other
+    # views' action tokens under the same capture-time rule; they still key no vision token
+    # there, and their own view's actions stay in the same-view pass alone. Independent of
+    # ``maskless_cross_view_attends_action``, with the same requirements. False preserves
+    # existing recipes.
+    maskless_cross_view_action_attends_action: bool = False
+
     # What the multiview backend lets the noisy tokens attend to.
     mask: MultiviewAttentionMaskConfig = attrs.field(factory=MultiviewAttentionMaskConfig)
 

cosmos_framework/configs/base/defaults/reasoner.py

@@ -183,7 +183,7 @@ class VLMConfig:
 class SoundUnderstandingConfig:
     """Optional audio-understanding pathway for standalone Reasoner/VLM models.
 
-    The encoder artifact initializes the frozen modality encoder independently
+    The encoder artifact initializes the optional trainable modality encoder independently
     of the Reasoner checkpoint. The audio projector is trainable by default and
     may be frozen explicitly after alignment.
     """
@@ -207,6 +207,9 @@ class SoundUnderstandingConfig:
     # ``separate_no_timestamps`` omits only audio-side timestamps, while
     # ``interleaved_av`` inserts each audio segment after its native video chunk.
     audio_layout: str = "separate_with_timestamps"
+    # Standalone audio remains enabled independently of video soundtrack extraction.
+    extract_audio_from_video: bool = False
+    freeze_encoder: bool = True
     freeze_projector: bool = False
     # Keep the full input embedding table as an optimizer parameter while
     # allowing gradients only through the audio start/end token positions.
@@ -247,6 +250,8 @@ def validate_sound_understanding_config(config: Any, *, sound_und: bool) -> None
     if config.audio_layout not in _AUDIO_LAYOUTS:
         raise ValueError(f"sound_und_config.audio_layout must be one of {_AUDIO_LAYOUTS}, got {config.audio_layout!r}")
     for name, value in (
+        ("extract_audio_from_video", config.extract_audio_from_video),
+        ("freeze_encoder", config.freeze_encoder),
         ("freeze_projector", config.freeze_projector),
         ("train_boundary_token_embeddings_only", config.train_boundary_token_embeddings_only),
         ("encoder_checkpoint_enabled", config.encoder_checkpoint_enabled),
@@ -260,6 +265,8 @@ def validate_sound_understanding_config(config: Any, *, sound_und: bool) -> None
         raise ValueError("An enabled audio encoder checkpoint requires a non-empty path")
     if config.exclude_frozen_encoder_from_training_checkpoint and not config.encoder_checkpoint_enabled:
         raise ValueError("Excluding the frozen audio encoder from training checkpoints requires checkpoint loading")
+    if config.exclude_frozen_encoder_from_training_checkpoint and not config.freeze_encoder:
+        raise ValueError("A trainable audio encoder must be included in training checkpoints")
 
 
 # Configs for LLM models