diff --git a/.file_mapping.json b/.file_mapping.json index 4ff7c2e0a..60d1258cc 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "845dbfbd3aa3913f3d031aa1f09fb16af8b37aa1-dirty", - "_dest_commit": "cf5d68c00d97ccd2480a2320ed652b92dec63102", - "_generated_at": "2026-10-07T07:59:16Z", + "_source_commit": "c956695327e59d1b759868e8e0a36cca5733cdb3-dirty", + "_dest_commit": "8aca062dd4d93111d8afc4c6e7c7c592a81b81bc", + "_generated_at": "2026-10-08T08:05:55Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -199,6 +199,9 @@ "projects/cosmos3/cosmos3/algorithm/loss/flow_matching_test.py": "cosmos_framework/model/generator/algorithm/loss/flow_matching_test.py", "projects/cosmos3/cosmos3/algorithm/loss/load_balancing.py": "cosmos_framework/model/generator/algorithm/loss/load_balancing.py", "projects/cosmos3/cosmos3/algorithm/loss/load_balancing_test.py": "cosmos_framework/model/generator/algorithm/loss/load_balancing_test.py", + "projects/cosmos3/cosmos3/algorithm/loss/modality_reduction.py": "cosmos_framework/model/generator/algorithm/loss/modality_reduction.py", + "projects/cosmos3/cosmos3/algorithm/loss/modality_reduction_distributed_test.py": "cosmos_framework/model/generator/algorithm/loss/modality_reduction_distributed_test.py", + "projects/cosmos3/cosmos3/algorithm/loss/modality_reduction_test.py": "cosmos_framework/model/generator/algorithm/loss/modality_reduction_test.py", "projects/cosmos3/cosmos3/algorithm/loss/time_weight.py": "cosmos_framework/model/generator/algorithm/loss/time_weight.py", "projects/cosmos3/cosmos3/callbacks/compile_tokenizer.py": "cosmos_framework/callbacks/compile_tokenizer.py", "projects/cosmos3/cosmos3/callbacks/data_stats.py": "cosmos_framework/callbacks/data_stats.py", @@ -489,6 +492,7 @@ "projects/cosmos3/cosmos3/models/reasoner/qwen3_vl_moe/shared_expert_test.py": "cosmos_framework/model/generator/reasoner/qwen3_vl_moe/shared_expert_test.py", "projects/cosmos3/cosmos3/models/sensor_encoder.py": "cosmos_framework/model/generator/sensor_encoder.py", "projects/cosmos3/cosmos3/models/utils/__init__.py": "cosmos_framework/model/generator/utils/__init__.py", + "projects/cosmos3/cosmos3/models/utils/batch_normalization.py": "cosmos_framework/model/generator/utils/batch_normalization.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition.py": "cosmos_framework/model/generator/utils/data_and_condition.py", "projects/cosmos3/cosmos3/models/utils/data_and_condition_test.py": "cosmos_framework/model/generator/utils/data_and_condition_test.py", "projects/cosmos3/cosmos3/models/utils/load_balancing_stats.py": "cosmos_framework/model/generator/utils/load_balancing_stats.py", @@ -604,7 +608,6 @@ "projects/cosmos3/interactive/configs/defaults/flex_attention.py": "cosmos_framework/configs/base/defaults/causal_flex_attention.py", "projects/cosmos3/interactive/configs/defaults/replay_attention.py": "cosmos_framework/configs/base/defaults/replay_attention.py", "projects/cosmos3/interactive/models/attention_io_layout.py": "cosmos_framework/model/generator/attention_io_layout.py", - "projects/cosmos3/interactive/models/joint_transfer_ar.py": "cosmos_framework/model/generator/joint_transfer_ar.py", "projects/cosmos3/interactive/models/mot/attention.py": "cosmos_framework/model/generator/mot/causal_attention.py", "projects/cosmos3/interactive/models/mot/context_parallel_test.py": "cosmos_framework/model/generator/mot/causal_context_parallel_test.py", "projects/cosmos3/interactive/models/mot/cosmos3_vfm_network.py": "cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py", @@ -619,7 +622,6 @@ "projects/cosmos3/interactive/models/mot/post_saturation/static_compile.py": "cosmos_framework/model/generator/mot/post_saturation/static_compile.py", "projects/cosmos3/interactive/models/mot/replicated_io.py": "cosmos_framework/model/generator/mot/replicated_io.py", "projects/cosmos3/interactive/models/mot/three_way_attention_test.py": "cosmos_framework/model/generator/mot/three_way_attention_test.py", - "projects/cosmos3/interactive/models/multiview_transfer_ar.py": "cosmos_framework/model/generator/multiview_transfer_ar.py", "projects/cosmos3/interactive/models/omni_mot_causal_model.py": "cosmos_framework/model/generator/omni_mot_causal_model.py", "projects/cosmos3/interactive/models/omni_mot_causal_model_test.py": "cosmos_framework/model/generator/omni_mot_causal_model_test.py", "projects/cosmos3/interactive/models/rolling_prompt.py": "cosmos_framework/model/generator/rolling_prompt.py", @@ -632,8 +634,14 @@ "projects/cosmos3/interactive/models/utils/kv_cache_test.py": "cosmos_framework/model/generator/utils/kv_cache_test.py", "projects/cosmos3/interactive/models/utils/kv_storage_backend.py": "cosmos_framework/model/generator/utils/kv_storage_backend.py", "projects/cosmos3/interactive/models/utils/kv_storage_backend_test.py": "cosmos_framework/model/generator/utils/kv_storage_backend_test.py", + "projects/cosmos3/interactive/models/utils/multiview_ar.py": "cosmos_framework/model/generator/utils/multiview_ar.py", "projects/cosmos3/interactive/models/utils/nvfp4.py": "cosmos_framework/model/generator/utils/nvfp4.py", "projects/cosmos3/interactive/models/utils/nvfp4_test.py": "cosmos_framework/model/generator/utils/nvfp4_test.py", + "projects/cosmos3/interactive/models/utils/rolling_kv/__init__.py": "cosmos_framework/model/generator/utils/rolling_kv/__init__.py", + "projects/cosmos3/interactive/models/utils/rolling_kv/rolling_prompt.py": "cosmos_framework/model/generator/utils/rolling_kv/rolling_prompt.py", + "projects/cosmos3/interactive/models/utils/rolling_kv/rolling_replay.py": "cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py", + "projects/cosmos3/interactive/models/utils/rolling_kv/rolling_sink_rope.py": "cosmos_framework/model/generator/utils/rolling_kv/rolling_sink_rope.py", + "projects/cosmos3/interactive/models/utils/rolling_kv/rolling_transfer_ar.py": "cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py", "projects/cosmos3/interactive/sequence_packing.py": "cosmos_framework/data/generator/sequence_packing/autoregressive.py", "projects/cosmos3/interactive/utils/data_batch.py": "cosmos_framework/utils/generator/data_batch.py", "projects/cosmos3/tokenizer/lidar_tokenizer/__init__.py": "cosmos_framework/model/generator/tokenizers/lidar/__init__.py", diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 6cb012deb..23420203e 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -361,12 +361,14 @@ jobs: # and inference2/, are intentionally omitted to keep this single-runner job within # its time budget. Add an L0 marker when one of those tests should join this gate. # L1, L2, and unmarked cases remain available for explicit or scheduled runs. + # Non-GPU tests generated from i4 are already covered by i4 CI and are deselected here; + # CF-owned tests plus mapped tests explicitly marked GPU or gpus(N>0) remain in this job. - name: Unit tests run: | export LD_LIBRARY_PATH= export COSMOS_DOWNLOAD_CACHE_DIR="$RUNNER_WORKSPACE/cosmos_input_cache" uv run --all-extras --group=cu128-train python -m pytest -v -s \ - cosmos_framework/ -o addopts= -m L0 \ + cosmos_framework/ -o addopts= -m L0 --skip-mapped-non-gpu-tests \ --deselect cosmos_framework/model/generator/mot/attention_test.py::test_two_way_attention_flex_matches_dense_across_batch_shapes # Deselected above and run here instead, in a process of its own. diff --git a/conftest.py b/conftest.py index 362f77472..c29985bb5 100644 --- a/conftest.py +++ b/conftest.py @@ -7,13 +7,21 @@ lazy_call._CONVERT_TARGET_TO_STRING = True import gc +import json import os from functools import cache from pathlib import Path import pytest -from cosmos_framework.inference.fixtures.args import ALL_LEVELS, ALL_NUM_GPUS, ALLOWED_GPUS_BY_LEVEL, Args, get_args, init_args +from cosmos_framework.inference.fixtures.args import ( + ALL_LEVELS, + ALL_NUM_GPUS, + ALLOWED_GPUS_BY_LEVEL, + Args, + get_args, + init_args, +) @pytest.fixture(scope="module") @@ -39,6 +47,12 @@ def _get_available_gpus() -> int: def pytest_addoption(parser: pytest.Parser): parser.addoption("--manual", action="store_true", default=False, help="Run manual tests") + parser.addoption( + "--skip-mapped-non-gpu-tests", + action="store_true", + default=False, + help="Skip tests generated from imaginaire4 unless they explicitly require a GPU.", + ) parser.addoption( "--num-gpus", default=None, @@ -129,13 +143,26 @@ def _parse_gpus_marker(mark: pytest.Mark) -> int: return required_gpus +@cache +def _mapped_destinations(root_dir: Path) -> frozenset[str]: + """Return repository-relative files generated from imaginaire4 sources.""" + payload = json.loads((root_dir / ".file_mapping.json").read_text()) + files = payload.get("files", payload) + if not isinstance(files, dict) or any(not isinstance(path, str) for path in files.values()): + raise TypeError(".file_mapping.json must contain a string-to-string files mapping.") + return frozenset(files.values()) + + def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]): args = get_args() + skip_mapped_non_gpu_tests = bool(config.getoption("--skip-mapped-non-gpu-tests")) + mapped_destinations = _mapped_destinations(config.rootpath) if skip_mapped_non_gpu_tests else frozenset() for item in items: manual_mark = _get_marker(item, "manual") level_mark = _get_marker(item, "level") gpus_mark = _get_marker(item, "gpus") + gpu_mark = _get_marker(item, "GPU") try: level = _parse_level_marker(level_mark) if level_mark else 0 gpus = _parse_gpus_marker(gpus_mark) if gpus_mark else 0 @@ -159,6 +186,10 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item item.add_marker( pytest.mark.skip(reason=f"test requires {gpus} GPUs, but only {available_gpus} are available") ) + if skip_mapped_non_gpu_tests and gpu_mark is None and gpus == 0: + item_path = Path(str(item.path)).relative_to(config.rootpath).as_posix() + if item_path in mapped_destinations: + item.add_marker(pytest.mark.skip(reason="mapped non-GPU test is covered by imaginaire4 CI")) # Exclude skipped tests selected_items = [] @@ -257,7 +288,6 @@ def init_torch_test(): torch.cuda.empty_cache() - _WHITELIST_ENV_VARS = { # Both set by sklearn/__init__.py at import time (OpenMP workarounds: allow a # duplicate OpenMP runtime, and skip the intel-openmp 2019.5 fork bug). Reached diff --git a/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py b/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py index 89d63df6b..b165db153 100644 --- a/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py +++ b/cosmos_framework/auxiliary/guardrail/qwen3guard/qwen3guard.py @@ -3,9 +3,10 @@ import argparse import re +from dataclasses import dataclass +from typing import Any import torch -from transformers import AutoModelForCausalLM, AutoTokenizer from cosmos_framework.auxiliary.guardrail.common.core import ContentSafetyGuardrail, GuardrailRunner from cosmos_framework.auxiliary.guardrail.qwen3guard.categories import UNSAFE_CATEGORIES @@ -15,51 +16,79 @@ UNSAFE = misc.Color.red("UNSAFE") +@dataclass(frozen=True) +class Qwen3GuardClassification: + """Structured Qwen3Guard generation result.""" + + label: str + categories: tuple[str, ...] + raw_output: str + + class Qwen3Guard(ContentSafetyGuardrail): + offload_model: bool + dtype: torch.dtype + model_id: str + model: Any + tokenizer: Any + def __init__( self, offload_model_to_cpu: bool = True, ) -> None: - """Llama Guard 3 model for text filtering safety check. + """Load Qwen3Guard for text filtering safety checks. Args: - checkpoint_dir (str): Path to the checkpoint directory. offload_model_to_cpu (bool, optional): Whether to offload the model to CPU. Defaults to True. """ + from transformers import AutoModelForCausalLM, AutoTokenizer + self.offload_model = offload_model_to_cpu self.dtype = torch.bfloat16 - model_id = "Qwen/Qwen3Guard-Gen-0.6B" + self.model_id = "Qwen/Qwen3Guard-Gen-0.6B" - self.model = AutoModelForCausalLM.from_pretrained(model_id) - self.tokenizer = AutoTokenizer.from_pretrained(model_id) + self.model = AutoModelForCausalLM.from_pretrained(self.model_id) + self.tokenizer = AutoTokenizer.from_pretrained(self.model_id) # Move model to GPU unless offload_model_to_cpu is True if not offload_model_to_cpu: self.model = self.model.to("cuda", dtype=self.dtype).eval() - log.debug("Moved llamaGuard3 model to GPU") + log.debug("Moved Qwen3Guard model to GPU") else: self.model = self.model.to("cpu", dtype=self.dtype).eval() log.debug("Moved Qwen3Guard model to CPU") - def extract_label_and_categories(self, prompt): + @staticmethod + def parse_output(content: str) -> Qwen3GuardClassification: + """Parse Qwen3Guard's generated safety label without collapsing its middle class.""" safe_pattern = r"Safety: (Safe|Unsafe|Controversial)" - category_pattern = r"(" + "|".join(UNSAFE_CATEGORIES.values()) + ")" + category_pattern = r"(" + "|".join(re.escape(value) for value in UNSAFE_CATEGORIES.values()) + ")" + safe_label_match = re.search(safe_pattern, content) + if safe_label_match is None: + raise ValueError(f"Qwen3Guard output did not contain a safety label: {content!r}") + label = safe_label_match.group(1) + categories = tuple(dict.fromkeys(re.findall(category_pattern, content))) + return Qwen3GuardClassification(label=label, categories=categories, raw_output=content) + + @torch.inference_mode() + def classify(self, prompt: str) -> Qwen3GuardClassification: + """Generate and parse one native Safe/Controversial/Unsafe decision.""" messages = [{"role": "user", "content": prompt}] text = self.tokenizer.apply_chat_template(messages, tokenize=False) - model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) - generated_ids = self.model.generate(**model_inputs, max_new_tokens=128) - output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist() + model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) # [1, input_tokens] + generated_ids = self.model.generate(**model_inputs, max_new_tokens=128) # [1, input_tokens + output_tokens] + output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist() # [output_tokens] content = self.tokenizer.decode(output_ids, skip_special_tokens=True) + return self.parse_output(content) - safe_label_match = re.search(safe_pattern, content) - label = safe_label_match.group(1) if safe_label_match else None - categories = re.findall(category_pattern, content) - if label.lower() == "unsafe": - return False, f"Prompt blocked by Qwen3Guard. Safety: {label}, Categories: {categories}" - else: - return True, "" + def extract_label_and_categories(self, prompt: str) -> tuple[bool, str]: + """Return the legacy binary result while retaining native classification through classify().""" + result = self.classify(prompt) + if result.label.lower() == "unsafe": + return False, f"Prompt blocked by Qwen3Guard. Safety: {result.label}, Categories: {result.categories}" + return True, "" def is_safe(self, prompt: str) -> tuple[bool, str]: """Check if the input prompt is safe according to the Qwen3Guard model.""" @@ -70,13 +99,13 @@ def is_safe(self, prompt: str) -> tuple[bool, str]: return True, "Unexpected error occurred when running Qwen3Guard guardrail." -def parse_args(): +def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--prompt", type=str, required=True, help="Input prompt") return parser.parse_args() -def main(args): +def main(args: argparse.Namespace) -> None: qwen3guard = Qwen3Guard() runner = GuardrailRunner(safety_models=[qwen3guard]) with misc.timer("Qwen3Guard safety check"): diff --git a/cosmos_framework/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py b/cosmos_framework/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py index b75f7b35b..65ccab55e 100644 --- a/cosmos_framework/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py +++ b/cosmos_framework/auxiliary/guardrail/video_content_safety_filter/vision_encoder.py @@ -25,6 +25,9 @@ def encode_image(self, input_img: Image.Image) -> torch.Tensor: """Encode an image into a feature vector.""" with torch.no_grad(): inputs = self.processor(images=input_img, return_tensors="pt").to(self.device, dtype=self.dtype) - image_features = self.model.get_image_features(**inputs) - image_features /= image_features.norm(dim=-1, keepdim=True) + model_output = self.model.get_image_features(**inputs) + image_features = getattr(model_output, "pooler_output", model_output) # [B,D] + if not isinstance(image_features, torch.Tensor): + raise TypeError("SigLIP image encoder did not return pooled tensor features") + image_features = torch.nn.functional.normalize(image_features, dim=-1) # [B,D] return image_features diff --git a/cosmos_framework/callbacks/dataloader_state.py b/cosmos_framework/callbacks/dataloader_state.py index 22a518903..c3659e236 100644 --- a/cosmos_framework/callbacks/dataloader_state.py +++ b/cosmos_framework/callbacks/dataloader_state.py @@ -11,12 +11,21 @@ import torch from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.trainer import persistent_validation_enabled from cosmos_framework.utils import log from cosmos_framework.utils.callback import Callback LANCE_VLM_RESUME_FORMAT = "lance_vlm_cursor_v1" LANCE_VLM_RESUME_STATE_KEY = "_lance_vlm_resume_state" LANCE_VLM_RESUME_WORKER_ENV_PREFIX = "LANCE_VLM_RESUME_STATE_WORKER_" +VAL_STATE_KEY = "val" +_STATE_LABELS = {"train": "dataloader", "val": "validation dataloader"} + + +def no_replace_resume_env_names(split: str, worker_id: int) -> tuple[str, str]: + """Return the env vars that carry one worker's last consumed (epoch, index) for a split.""" + prefix = "NSL_STATE_WORKER" if split == "train" else "NSL_STATE_VAL_WORKER" + return f"{prefix}_{worker_id}_EPOCH", f"{prefix}_{worker_id}_INDEX" @dataclass @@ -36,10 +45,15 @@ def __init__( self.distributor_type = distributor_type self.config: Any = None self.state: dict[int, NoReplaceShardlistState] = {} + self.val_state: dict[int, NoReplaceShardlistState] = {} self.lance_state: dict[int, dict[str, Any]] = {} self._pending_lance_state: dict[str, Any] | None = None self.verbose = True + def _tracks_validation(self) -> bool: + """Validation cursors only carry over when the trainer keeps one validation iterator alive.""" + return self.distributor_type == "no_replace" and persistent_validation_enabled(self.config.trainer) + def on_training_step_batch_start( self, model: ImaginaireModel, @@ -109,18 +123,19 @@ def _update_lance_state(self, update: dict[str, Any]) -> None: } self.lance_state[worker_id] = {"previous": previous_current, "current": current} - def _update_state_from_batch(self, data_batch: dict[str, torch.Tensor]) -> None: + @staticmethod + def _update_state_from_batch( + state: dict[int, NoReplaceShardlistState], data_batch: dict[str, torch.Tensor] + ) -> None: worker_ids = data_batch["sample_worker_id"].tolist() # [B] epochs = data_batch["sample_epoch"].tolist() # [B] indices = data_batch["sample_index"].tolist() # [B] for worker_id, epoch, index in zip(worker_ids, epochs, indices, strict=True): - if worker_id not in self.state: - self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) + if worker_id not in state: + state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) - elif self.state[worker_id].epoch < epoch or ( - self.state[worker_id].index < index and self.state[worker_id].epoch == epoch - ): - self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) + elif state[worker_id].epoch < epoch or (state[worker_id].index < index and state[worker_id].epoch == epoch): + state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) def on_training_step_batch_end( self, @@ -135,7 +150,7 @@ def on_training_step_batch_end( self._update_lance_state(self._pending_lance_state) self._pending_lance_state = None else: - self._update_state_from_batch(data_batch) + self._update_state_from_batch(self.state, data_batch) def on_training_step_end( self, @@ -157,6 +172,18 @@ def on_training_step_end( msg += f"worker {wid}: epoch={state.epoch}, index={state.index}\n" log.info(msg) + def on_validation_step_end( + self, + model: ImaginaireModel, + data_batch: dict[str, torch.Tensor], + output_batch: dict[str, torch.Tensor], + loss: torch.Tensor, + iteration: int = 0, + ) -> None: + # Lance batches carry no shardlist cursor. + if "sample_worker_id" in data_batch and self._tracks_validation(): + self._update_state_from_batch(self.val_state, data_batch) + def has_checkpoint_state(self) -> bool: return self.distributor_type == "no_replace" @@ -165,15 +192,11 @@ def state_dict(self) -> dict[Any, Any]: return {} if self.lance_state: - return {"format": LANCE_VLM_RESUME_FORMAT, "workers": self.lance_state} - - state_dict: dict[int, dict[str, int]] = {} - for worker_id, per_worker_state in self.state.items(): - state_dict[worker_id] = {"epoch": per_worker_state.epoch, "index": per_worker_state.index} - log.info( - f"Saved dataloader state for worker {worker_id}: " - f"epoch={per_worker_state.epoch}, index={per_worker_state.index}" - ) + state_dict: dict[Any, Any] = {"format": LANCE_VLM_RESUME_FORMAT, "workers": self.lance_state} + else: + state_dict = self._save_shardlist_state(self.state, "train") + if self.val_state: + state_dict[VAL_STATE_KEY] = self._save_shardlist_state(self.val_state, "val") return state_dict def load_state_dict(self, state_dict: dict[Any, Any]) -> None: @@ -184,6 +207,14 @@ def load_state_dict(self, state_dict: dict[Any, Any]) -> None: log.info("No dataloader state found in checkpoint") return + state_dict = dict(state_dict) + val_state_dict = state_dict.pop(VAL_STATE_KEY, None) + if val_state_dict is not None: + if self._tracks_validation(): + self.val_state = self._load_shardlist_state(val_state_dict, "val") + else: + log.info("Ignoring validation dataloader state because the validation iterator is not persistent") + if state_dict.get("format") == LANCE_VLM_RESUME_FORMAT: workers = state_dict.get("workers") if not isinstance(workers, dict): @@ -203,11 +234,31 @@ def load_state_dict(self, state_dict: dict[Any, Any]) -> None: ) return - self.state = {} + self.state = self._load_shardlist_state(state_dict, "train") + + @staticmethod + def _save_shardlist_state(state: dict[int, NoReplaceShardlistState], split: str) -> dict[int, dict[str, int]]: + state_dict: dict[int, dict[str, int]] = {} + for worker_id, per_worker_state in state.items(): + state_dict[worker_id] = {"epoch": per_worker_state.epoch, "index": per_worker_state.index} + log.info( + f"Saved {_STATE_LABELS[split]} state for worker {worker_id}: " + f"epoch={per_worker_state.epoch}, index={per_worker_state.index}" + ) + return state_dict + + @staticmethod + def _load_shardlist_state(state_dict: dict[int, dict[str, int]], split: str) -> dict[int, NoReplaceShardlistState]: + """Restore per-worker cursors and export them to the split's dataloader workers, which start after loading.""" + state: dict[int, NoReplaceShardlistState] = {} for worker_id, per_worker_state in state_dict.items(): epoch = per_worker_state["epoch"] index = per_worker_state["index"] - self.state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) - os.environ[f"NSL_STATE_WORKER_{worker_id}_EPOCH"] = str(epoch) - os.environ[f"NSL_STATE_WORKER_{worker_id}_INDEX"] = str(index) - log.info(f"Loaded no replace dataloader state for worker {worker_id}: epoch={epoch}, index={index}") + state[worker_id] = NoReplaceShardlistState(epoch=epoch, index=index) + epoch_env, index_env = no_replace_resume_env_names(split, worker_id) + os.environ[epoch_env] = str(epoch) + os.environ[index_env] = str(index) + log.info( + f"Loaded no replace {_STATE_LABELS[split]} state for worker {worker_id}: epoch={epoch}, index={index}" + ) + return state diff --git a/cosmos_framework/callbacks/every_n_draw_sample.py b/cosmos_framework/callbacks/every_n_draw_sample.py index 539830456..0ad30edc7 100644 --- a/cosmos_framework/callbacks/every_n_draw_sample.py +++ b/cosmos_framework/callbacks/every_n_draw_sample.py @@ -969,18 +969,25 @@ def every_n_impl( if distributed.is_rank0() and wandb.run: sample_counter = getattr(trainer, "sample_counter", iteration) data_type = "image" if model.is_image_batch(data_batch) else "video" + image_flags = data_batch.get("sample_is_image") + if image_flags and any(image_flags) and not all(image_flags): + data_type = "mixed" + media_caption = f"{sample_counter}" + if image_flags is not None: + sources = data_batch.get("dataset_name", []) + media_caption += " | sources: " + ", ".join(str(source) for source in sources) tag += f"_{data_type}" info = { "trainer/global_step": iteration, "sample_counter": sample_counter, } if self.do_x0_prediction: - _add_wandb_media(info, f"{self.name}/{tag}_x0", x0_img_fp, f"{sample_counter}") + _add_wandb_media(info, f"{self.name}/{tag}_x0", x0_img_fp, media_caption) # convert mse_loss to a dict mse_loss = mse_loss.tolist() info.update({f"x0_pred_mse_{tag}/Sigma{sigmas[i]:0.5f}": mse_loss[i] for i in range(len(mse_loss))}) - _add_wandb_media(info, f"{self.name}/{tag}_sample", sample_img_fp, f"{sample_counter}") + _add_wandb_media(info, f"{self.name}/{tag}_sample", sample_img_fp, media_caption) wandb.log( info, step=iteration, diff --git a/cosmos_framework/callbacks/iter_speed.py b/cosmos_framework/callbacks/iter_speed.py index 2ef103c87..596fe25f8 100644 --- a/cosmos_framework/callbacks/iter_speed.py +++ b/cosmos_framework/callbacks/iter_speed.py @@ -370,7 +370,14 @@ def every_n_impl( per_sample_batch_counter = dict() # for VFM - if hasattr(model, "is_image_batch") and hasattr(model, "input_image_key") and hasattr(model, "input_video_key"): + if "image_batch_size" in output_batch and "video_batch_size" in output_batch: + # The model counts samples in the processed CP payload. Always log + # both modalities, including zero when one is absent from this batch. + per_sample_batch_counter["image_batch_size"] = output_batch["image_batch_size"] + per_sample_batch_counter["video_batch_size"] = output_batch["video_batch_size"] + elif ( + hasattr(model, "is_image_batch") and hasattr(model, "input_image_key") and hasattr(model, "input_video_key") + ): is_image_batch = model.is_image_batch(data_batch) if is_image_batch: image_batch_size = len(data_batch[model.input_image_key]) diff --git a/cosmos_framework/callbacks/sigma_loss_analysis.py b/cosmos_framework/callbacks/sigma_loss_analysis.py index 439956011..6301951b8 100644 --- a/cosmos_framework/callbacks/sigma_loss_analysis.py +++ b/cosmos_framework/callbacks/sigma_loss_analysis.py @@ -388,7 +388,12 @@ def on_training_step_end( if sigma.ndim == 2: sigma = sigma.mean(dim=-1) # [B] (reduced from [B,T_max] or [B,1]) - if cast(_SupportsIsImageBatch, model).is_image_batch(data_batch): + vision_flags = output_batch.get("vision_is_image") + if vision_flags is not None: + is_image = torch.tensor(vision_flags, dtype=torch.bool, device=sigma.device) # [B_items] + self.image_cache.add(sigma[is_image], fm_loss_vision_per_instance[is_image]) + self.video_cache.add(sigma[~is_image], fm_loss_vision_per_instance[~is_image]) + elif cast(_SupportsIsImageBatch, model).is_image_batch(data_batch): self.image_cache.add(sigma, fm_loss_vision_per_instance) else: self.video_cache.add(sigma, fm_loss_vision_per_instance) diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 0be08be20..5ad64112b 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -133,22 +133,28 @@ class RectifiedFlowTrainingConfig: # control items; this is independent of per-item active-token normalization. exclude_fully_conditioned_items: bool | None = None - # Sample-level (vs rank-level) loss averaging for the vision modality. + # Loss reduction mode for all supervised modalities. # - # By default the vision loss on each rank is a mean over that rank's samples, and + # With "local_item_mean" (default), vision loss is a mean over each rank's items, and # FSDP/DDP averages gradients across ranks — so every *rank* contributes equally - # regardless of how many image/video samples it holds ("rank-level averaging"). - # When ranks carry different sample counts this over-weights samples on sparse ranks. + # regardless of how many image/video items it holds ("rank-level averaging"). + # When ranks carry different item counts this over-weights items on sparse ranks. # - # When True, the loss is renormalized so every *sample* contributes equally across - # the whole data-parallel group: each iteration all-reduces the total number of image - # and video samples over the DP group, and image and video losses are each normalized - # by their own global sample count and then summed. This counteracts the framework's - # rank-level gradient averaging so the effective objective is a true per-sample mean. - # The two independently normalized terms are weighted by the existing `image_loss_scale` - # (images; falls back to `loss_scale` when None) and `loss_scale` (videos), so their - # balance is tuned with the same knobs as the legacy rank-level path. - sample_level_loss_averaging: bool = False + # With "global_sample_mean", image/video/action/sound/LiDAR/radar losses are independently normalized + # by each modality's global sample count, then summed with their configured weights. + # This counteracts FSDP/DDP rank-level gradient averaging so every *sample* + # contributes equally within its modality, for homogeneous and mixed batches. + # Images use `image_loss_scale` (falling back to `loss_scale`), videos use + # `loss_scale`, and action/sound/LiDAR/radar retain their own configured loss weights. + # Denominators count logical samples contributing supervision, after averaging + # their timestep-weighted supervised item losses; + # neither tokens, clean controls, nor missing modalities add to those counts. + # Weights are never renormalized, including when a modality is globally absent. + # This intentional objective change applies per microbatch; gradient accumulation + # and independently weighted MoE auxiliary losses keep their existing semantics. + # Select this behavior independently of the dataloader class. "local_item_mean" preserves + # local rank-level averaging, including mixed image/video item weights. + loss_reduction_mode: Literal["local_item_mean", "global_sample_mean"] = "local_item_mean" @attrs.define(slots=False) @@ -176,6 +182,18 @@ class FixedStepSamplerConfig: sample_type: str = "sde" +@attrs.define(slots=False) +class MultiviewActionConditioningConfig: + """Opt-in synchronized RGB targets with full frame-rate actions replicated per view. + + Fixed view codes are written into padded action channels after normalization; + the first (reference) view must use zeros to retain single-view compatibility. + """ + + view_code_start: int + view_codes: list[list[float]] + + # Don't have any defaults and init only in config file. @attrs.define(slots=False) class OmniMoTModelConfig: @@ -311,6 +329,11 @@ class OmniMoTModelConfig: # (the maskless decomposition or a masked FlexAttention call), and under what mask. multiview_attention: MultiviewAttentionConfig = MultiviewAttentionConfig() + # 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. + multiview_action_conditioning: MultiviewActionConditioningConfig | None = None + # Per-layer NATTEN parameters # Must use "three_way" attention if used. # If None, all attention layers remain dense. diff --git a/cosmos_framework/configs/base/defaults/replay_attention.py b/cosmos_framework/configs/base/defaults/replay_attention.py index fb894e9d2..3fac24ec1 100644 --- a/cosmos_framework/configs/base/defaults/replay_attention.py +++ b/cosmos_framework/configs/base/defaults/replay_attention.py @@ -58,6 +58,10 @@ class TeacherForcingReplayPolicyConfig: validator=attrs.validators.optional(attrs.validators.ge(0)), ) + # Maskless decomposed replay merges same-view and cross-view partitions. + # Exclude their overlapping keys when reproducing deduplicated attention. + deduplicate_cross_view: bool = attrs.field(default=False, validator=attrs.validators.instance_of(bool)) + def __attrs_post_init__(self) -> None: """Reject replay graphs that can relay future RGB through global control K/V.""" self.validate() diff --git a/cosmos_framework/data/generator/joint_dataloader.py b/cosmos_framework/data/generator/joint_dataloader.py index 2a09808ca..d016c14d4 100644 --- a/cosmos_framework/data/generator/joint_dataloader.py +++ b/cosmos_framework/data/generator/joint_dataloader.py @@ -1232,10 +1232,8 @@ def __init__( raise ValueError(f"Invalid packing limits for {name}: {limits}") if any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in limits.values()): raise ValueError(f"Packing limits for {name} must be positive integers: {limits}") - # Keep PyTorch/WebDataset/Lance worker creation on the main thread. - # The packer thread starts only after child iterator prewarm has completed. - if enable_async_batch_building and not prewarm: - raise ValueError("enable_async_batch_building=True requires prewarm=True.") + # Set this before the base constructor initializes child iterators. + self.enable_async_batch_building: bool = enable_async_batch_building # A time ceiling skews the token mix here exactly as it does on the random # loader, but the correction cannot be applied here: this loader seeds the draw @@ -1291,10 +1289,17 @@ def __init__( total_ratio = sum(self.data_ratios) self.data_probs = np.array([ratio / total_ratio for ratio in self.data_ratios]) # Async batch builder state - self.enable_async_batch_building: bool = enable_async_batch_building self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s) self._async_batch_builder: _AsyncBatchBuilder | None = None + def _initialize_child_iterators_once(self) -> None: + # Keep PyTorch/WebDataset/Lance worker creation outside the packer thread. + # Workerless loaders can initialize lazily without fetching every stream first. + if self.enable_async_batch_building and not self.prewarm: + if any(getattr(loader, "num_workers", None) != 0 for loader in self.dataloader_list): + raise ValueError("Async batch building without prewarm requires num_workers=0 on every child loader.") + super()._initialize_child_iterators_once() + def __iter__(self) -> Iterator[dict[str, Any]]: self._initialize_child_iterators_once() if not self.enable_async_batch_building: @@ -1892,3 +1897,236 @@ def __iter__(self): # exactly the quantity whose split across sources is being held. self.token_mix.observe(index_id, metrics.current_sequence_length) yield output_batch + + +@dataclass(frozen=True) +class SampleTokenMixConfig: + """Estimate source token lengths per fresh sample, including deferred draws. + + Ratios become target token shares when enabled. The estimate is updated once + per fresh, packable candidate; rechecking a deferred sample never counts it a + second time. Both loaders use the shared ``TokenMixController`` estimator. + """ + + enabled: bool = True + half_life_samples: float = 50.0 + warmup_samples_per_source: int = 8 + + def __post_init__(self) -> None: + if not math.isfinite(self.half_life_samples) or self.half_life_samples <= 0: + raise ValueError("half_life_samples must be finite and positive.") + if self.warmup_samples_per_source < 1: + raise ValueError("warmup_samples_per_source must be positive.") + + +@dataclass(frozen=True) +class _Candidate: + source: int + sample: dict[str, Any] + tokens: int + seconds: float + + +def collate_mixed_samples(samples: list[dict[str, Any]]) -> dict[str, Any]: + """Keep every column aligned and expose ragged vision items under one key. + + Pixel normalization belongs to the model input boundary. Here image and video + tensors retain their native shapes/dtypes, with a per-sample image flag. Empty + vision lists represent samples with another modality, such as LiDAR. Optional + metadata stays aligned even when only one source supplies it. + """ + if not samples: + raise ValueError("Cannot collate an empty mixed batch.") + excluded = {"images", "video", *_BATCH_TIMING_KEYS} + keys = set().union(*(sample.keys() for sample in samples)) - excluded + batch = {key: [sample.get(key) for sample in samples] for key in keys} + batch["vision"] = [sample.get("images", sample.get("video", [])) for sample in samples] + batch["sample_is_image"] = ["images" in sample for sample in samples] + return batch + + +class RandomMixedJointDataLoader(JointDataLoader): + """Mix sources within each packed batch without altering the legacy loaders. + + Each rank owns an independent RNG. Under a time budget, configured ratios + target token shares; without it, ratios are sample draw probabilities unless + sample token-mix control is explicitly enabled. Random draws allow, but do not + require, several sources in a batch. + + Deferred candidates are kept in draw order and offered first in the next + batch. In particular its oldest candidate starts that batch, even if its cost + exceeds the time target, so lookahead cannot starve expensive samples. A token + limit still rejects individually oversized samples using the base contract. + Finite sources are retired on exhaustion; the remaining sources continue + with renormalized probabilities until all sources and pending samples finish. + """ + + def __init__( + self, + dataloaders: dict[str, Any], + tokenizer_spatial_compression_factor: int, + tokenizer_temporal_compression_factor: int, + patch_spatial: int, + max_sequence_length: int | None = None, + max_samples_per_batch: int | None = None, + lidar_spatial_compression: Sequence[int] | None = None, + lidar_temporal_compression_factor: int | None = None, + sound_latent_fps: float = 0, + audio_sample_rate: int = 48000, + default_lookahead_limit: int = 10, + lookahead_limits: dict[str, int] | None = None, + uniae_chunk_frames: int | Mapping[str, int] | None = None, + uniae_pad_frames: int | None = None, + iteration_time_budget: IterationTimeBudgetConfig | None = None, + token_mix_control: SampleTokenMixConfig | None = None, + seed: int = 42, + prewarm: bool = True, + prewarm_concurrency: int = 1, + lazy_initialize_child_iterators: bool = False, + forkserver_preload_modules: list[str] | None = None, + ) -> None: + if default_lookahead_limit < 1 or any(value < 1 for value in (lookahead_limits or {}).values()): + raise ValueError("Lookahead limits must be positive.") + ratios = [entry["ratio"] for entry in dataloaders.values() if entry is not None] + if not ratios or not all(math.isfinite(ratio) and ratio >= 0 for ratio in ratios) or not any(ratios): + raise ValueError("Source ratios must be finite and nonnegative, with at least one positive ratio.") + super().__init__( + dataloaders=dataloaders, + tokenizer_spatial_compression_factor=tokenizer_spatial_compression_factor, + tokenizer_temporal_compression_factor=tokenizer_temporal_compression_factor, + patch_spatial=patch_spatial, + max_sequence_length=max_sequence_length, + max_samples_per_batch=max_samples_per_batch, + lidar_spatial_compression=lidar_spatial_compression, + lidar_temporal_compression_factor=lidar_temporal_compression_factor, + sound_latent_fps=sound_latent_fps, + audio_sample_rate=audio_sample_rate, + prewarm=prewarm, + prewarm_concurrency=prewarm_concurrency, + default_lookahead_limit=default_lookahead_limit, + lookahead_limits=lookahead_limits, + uniae_chunk_frames=uniae_chunk_frames, + uniae_pad_frames=uniae_pad_frames, + lazy_initialize_child_iterators=lazy_initialize_child_iterators, + iteration_time_budget=iteration_time_budget, + forkserver_preload_modules=forkserver_preload_modules, + ) + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + self.rng: np.random.Generator = np.random.default_rng(np.random.SeedSequence([seed, rank])) + config = token_mix_control or SampleTokenMixConfig(enabled=self.iteration_time_budget is not None) + # Legacy batch-named fields count observations; here each is one fresh sample. + # Sample lengths can differ far more than packed batches, so keep this draw unclipped. + self.token_mix: TokenMixController = TokenMixController( + self.dataset_name_list, + self.data_ratios, + TokenMixControlConfig( + enabled=config.enabled, + half_life_batches=config.half_life_samples, + warmup_batches_per_source=config.warmup_samples_per_source, + max_weight_ratio=None, + ), + ) + self.accepted_tokens: np.ndarray = np.zeros(len(self.dataset_name_list), dtype=np.int64) + self._active: np.ndarray = np.ones(len(self.dataset_name_list), dtype=bool) + self._pending: deque[_Candidate] = deque() + + def _fresh_candidate(self, metrics: _PackingMetrics, worker_timings: list[dict[str, Any]]) -> _Candidate | None: + while np.any(self._active): + source = int(self.rng.choice(len(self.dataloader_list), p=self.token_mix.probabilities(self._active))) + from_buffer = bool(self.buffers[source]) + try: + sample = self._get_next_sample(source) + except StopIteration: + self._active[source] = False + continue + metrics.from_buffer += int(from_buffer) + metrics.from_workers += int(not from_buffer) + if not from_buffer: + worker_timings.append( + { + "dataset_name": self.dataset_name_list[source], + **{key: sample[key] for key in _BATCH_TIMING_KEYS if key in sample}, + } + ) + if _sample_should_drop(sample): + self._log_drop_sample(sample, self.dataset_name_list[source]) + metrics.dropped_count += 1 + _extend_action_sampler_draw_counts(metrics.dropped_action_sampler_draw_counts, sample) + continue + tokens, seconds = self._compute_sample_cost(sample) + if not self._sample_fits( + num_tokens=tokens, + sample_seconds=seconds, + packed_tokens=0, + packed_sample_seconds=0.0, + batch_started=False, + ): + metrics.dropped_count += 1 + _extend_action_sampler_draw_counts(metrics.dropped_action_sampler_draw_counts, sample) + continue + if tokens <= 0: + raise ValueError("A packable sample must contain at least one token.") + self.token_mix.observe(source, tokens, record_realized=False) + return _Candidate(source, sample, tokens, seconds) + return None + + def __iter__(self) -> Iterator[dict[str, Any]]: + self._initialize_child_iterators_once() + while self._pending or np.any(self._active): + metrics = _PackingMetrics() + samples: list[dict[str, Any]] = [] + deferred: deque[_Candidate] = deque() + source_stats: dict[str, dict[str, int | float]] = {} + worker_timings: list[dict[str, Any]] = [] + misses = 0 + while self.max_samples_per_batch is None or len(samples) < self.max_samples_per_batch: + if self._pending: + candidate = self._pending.popleft() + metrics.from_buffer += 1 + else: + candidate = self._fresh_candidate(metrics, worker_timings) + if candidate is None: + break + if not self._sample_fits( + num_tokens=candidate.tokens, + sample_seconds=candidate.seconds, + packed_tokens=metrics.current_sequence_length, + packed_sample_seconds=metrics.sample_seconds, + batch_started=bool(samples), + ): + deferred.append(candidate) + misses += 1 + if misses >= self.lookahead_limits[candidate.source]: + break + continue + + name = self.dataset_name_list[candidate.source] + sample = dict(candidate.sample) + sample["dataset_name"] = name + samples.append(sample) + metrics.num_samples += 1 + metrics.current_sequence_length += candidate.tokens + metrics.sample_seconds += candidate.seconds + self.accepted_tokens[candidate.source] += candidate.tokens + stats = source_stats.setdefault(name, {"_num_tokens": 0, "_num_samples": 0, "_sample_seconds": 0.0}) + stats["_num_tokens"] += candidate.tokens + stats["_num_samples"] += 1 + stats["_sample_seconds"] += candidate.seconds + + # Existing pending candidates preceded all fresh draws. Put candidates + # skipped this batch back ahead of those we have not visited yet. + deferred.extend(self._pending) + self._pending = deferred + if not samples: + return + batch = collate_mixed_samples(samples) + batch["_source_packing_metrics"] = source_stats + # One timing record per child batch fetch, including dropped draws. + # Buffered/deferred samples do not charge their worker a second time. + batch["_source_worker_timings"] = worker_timings + metrics.attach_to( + batch, + buffer_size=len(self._pending) + sum(len(buffer) for buffer in self.buffers), + projected_iteration_sec=self._projected_iteration_sec(metrics.sample_seconds), + ) + yield batch diff --git a/cosmos_framework/data/generator/sequence_packing/packers.py b/cosmos_framework/data/generator/sequence_packing/packers.py index 0c16e2f85..37b79666a 100644 --- a/cosmos_framework/data/generator/sequence_packing/packers.py +++ b/cosmos_framework/data/generator/sequence_packing/packers.py @@ -5,6 +5,7 @@ from __future__ import annotations +from dataclasses import replace from typing import TYPE_CHECKING import torch @@ -13,6 +14,7 @@ from cosmos_framework.data.generator.sequence_packing.temporal_causal import pack_supertokens_temporal_causal if TYPE_CHECKING: + from cosmos_framework.configs.base.defaults.model_config import MultiviewActionConditioningConfig from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean @@ -148,6 +150,117 @@ def uses_single_timestep(input_timesteps: torch.Tensor) -> bool: return bool((flat == flat[0]).all()) +def replicate_multiview_actions( + data: GenerationDataClean, *, num_views: int, temporal_factor: int, action_dim: int +) -> GenerationDataClean: + """Replicate all frame-rate actions camera-major; leave single-view data untouched.""" + if data.num_vision_items_per_sample is None: + return data + if data.num_vision_items_per_sample != [num_views] * data.batch_size: + raise ValueError("Multiview action conditioning requires one RGB item per configured view and no controls.") + assert data.x0_tokens_vision is not None and data.x0_tokens_action is not None + actions: list[torch.Tensor] = [] + domains: list[torch.Tensor] = [] + for sample, raw_action in enumerate(data.x0_tokens_action): + latent_t = data.x0_tokens_vision[num_views * sample].shape[2] + if ( + raw_action.ndim != 2 + or raw_action.shape[1] != action_dim + or raw_action.shape[0] != (latent_t - 1) * temporal_factor + ): + raise ValueError(f"Multiview requires all {action_dim}D frame-rate actions aligned to T+1 observations.") + actions.append(raw_action.repeat(num_views, 1)) + if data.action_domain_id is not None: + domain = data.action_domain_id[sample].reshape(-1) + if domain.numel() not in (1, raw_action.shape[0]): + raise ValueError("Action domains must be scalar or have one ID per frame-rate action.") + domains.append(domain.repeat(num_views) if domain.numel() > 1 else domain) + return replace( + data, + x0_tokens_action=actions, + raw_state_action=actions, + action_domain_id=domains or None, + num_views_per_action_item=[num_views] * data.batch_size, + ) + + +def pack_multiview_action_conditioning( + plans: list[SequencePlan], + text_tokens: list[list[int]], + data: GenerationDataClean, + timesteps: torch.Tensor, + *, + config: MultiviewActionConditioningConfig, + special_tokens: dict[str, int], + patch_spatial: int, + temporal_factor: int, + temporal_margin: int, + base_fps: float, +) -> PackedSequence: + """Pack separate RGB and full frame-rate action blocks with synchronized ragged views.""" + num_views = len(config.view_codes) + if data.num_vision_items_per_sample != [num_views] * data.batch_size: + raise ValueError("Multiview action batches must match the configured view count.") + assert data.x0_tokens_vision is not None and data.x0_tokens_action is not None + builder = PackedSequenceBuilder(uses_single_timestep=uses_single_timestep(timesteps)) + for sample, plan in enumerate(plans): + if not plan.has_vision or not plan.has_action or plan.has_sound or plan.has_lidar or plan.has_radar: + raise ValueError("Multiview action conditioning supports RGB + conditioned action forward dynamics only.") + sample_start = builder.current_seq_index + builder.set_mrope_temporal_offset(0) + builder.pack_text_tokens(text_tokens[sample], special_tokens, has_generation=True, use_float_positions=True) + builder.advance_mrope_temporal_offset(temporal_margin) + temporal_offset = builder.mrope_temporal_offset + fps = float(data.fps_vision[sample]) if data.fps_vision is not None else base_fps + views = data.x0_tokens_vision[num_views * sample : num_views * (sample + 1)] + latent_t = views[0].shape[2] + if any(view.shape[2] != latent_t for view in views): + raise ValueError("All views must share latent time, even when their spatial grids differ.") + action = data.x0_tokens_action[sample] + actions_per_view = (latent_t - 1) * temporal_factor + if action.ndim != 2 or action.shape[0] != num_views * actions_per_view: + raise ValueError("All frame-rate actions must be replicated camera-major without temporal subsampling.") + # Fixed codes occupy configured padded slots after normalization; + # fully conditioned frame-rate actions bypass RF masking. + action = action.clone() + # Torch probing OmegaConf containers retains caller-frame tensors until + # cyclic GC; materialize plain lists before passing codes to Torch. + codes = action.new_tensor([list(code) for code in config.view_codes]) + start = config.view_code_start + action.reshape(num_views, actions_per_view, action.shape[1])[..., start : start + codes.shape[1]] = codes[ + :, None, : + ] + generation_start = builder.current_seq_index + for view in views: + builder.set_mrope_temporal_offset(temporal_offset) + builder.pack_vision_tokens( + view, + plan.condition_frame_indexes_vision, + float(timesteps[sample].flatten()[0]), + patch_spatial, + fps, + True, + base_fps, + temporal_factor, + None, + None, + ) + builder.pack_action_tokens( + action, + list(range(action.shape[0])), + 0.0, + action_temporal_offset=temporal_offset, + enable_fps_modulation=True, + base_fps=base_fps, + action_fps=fps, + base_temporal_compression_factor=temporal_factor, + action_start_frame_offset=1, + num_views=num_views, + ) + builder.finish_sample(builder.current_seq_index - generation_start, builder.current_seq_index - sample_start) + return builder.finalize(gen_data_clean=data) + + def pack_input_sequence( sequence_plans: list[SequencePlan], input_text_indexes: list[list[int]], @@ -876,6 +989,17 @@ def pack_input_sequence( seq_builder.null_action_supertokens = null_action_flags[0] # Finalize and return packed data - return seq_builder.finalize( + packed = seq_builder.finalize( gen_data_clean=gen_data_clean, ) + # Loss normalization counts logical owners, not flattened control/target items. + # Keep this metadata outside the token layout so attention packing is unchanged. + for name in ("vision", "lidar", "radar", "action", "sound"): + counts = getattr(gen_data_clean, f"num_{name}_items_per_sample", None) + packed.modality_sample_ids[name] = [ + sample_index + for sample_index, plan in enumerate(sequence_plans) + if getattr(plan, f"has_{name}") + for _ in range(counts[sample_index] if counts is not None else 1) + ] + return packed diff --git a/cosmos_framework/data/generator/sequence_packing/sequence.py b/cosmos_framework/data/generator/sequence_packing/sequence.py index e9407b70c..cba781e00 100644 --- a/cosmos_framework/data/generator/sequence_packing/sequence.py +++ b/cosmos_framework/data/generator/sequence_packing/sequence.py @@ -1191,6 +1191,7 @@ def finalize( split_lens=split_lens, attn_modes=attn_modes, is_image_batch=gen_data_clean.is_image_batch, + vision_is_image=gen_data_clean.vision_image_flags(), uses_single_timestep=self.uses_single_timestep, # Text modality (converted to tensors) text_ids=torch.tensor(self.text_ids, dtype=torch.long), # [N_text_tokens] @@ -1354,6 +1355,10 @@ class PackedSequence: # FlexAttention mask as its vision items followed by its LiDAR items. num_lidar_items_per_sample: list[int] | None = None + # Modality identity and ownership are data properties, independent of the + # loader. Sample IDs map flattened supervised items back to logical samples. + vision_is_image: list[bool] | None = None + modality_sample_ids: dict[str, list[int]] = field(default_factory=dict) # Same bookkeeping for the radar stream, whose items follow the sample's LiDAR items. # None when the batch carries no radar. num_radar_items_per_sample: list[int] | None = None diff --git a/cosmos_framework/data/generator/token_mix_control.py b/cosmos_framework/data/generator/token_mix_control.py index f8ec878b1..06d1828f6 100644 --- a/cosmos_framework/data/generator/token_mix_control.py +++ b/cosmos_framework/data/generator/token_mix_control.py @@ -21,6 +21,11 @@ cooperation from the packer: it prices what the packer did, drops, look-ahead misses, token ceiling and all. +The same estimator supports sample-level mixing: each observation is one fresh, +packable candidate and ``T_i`` is mean tokens per sample. Deferred candidates must +be observed only once. Their eventual acceptance is accounted for by the loader, +separately from the candidate estimate, using ``record_realized=False`` here. + This is a plug-in rather than a servo. ``T_i`` is a property of a source's data under the packer's ceilings and does not depend on how often that source is drawn, so the loop has no feedback path into the quantity it measures: no gain to tune, nothing to @@ -29,8 +34,8 @@ Every rank runs its own controller over its own batches, which keeps it out of the collective path -- the packer's daemon thread makes any collective there a hazard. -That is also why only ``RandomJointDataLoader`` may use it: ranks there already draw -independently and converge to the same weights up to noise, whereas +Both ``RandomJointDataLoader`` and ``RandomMixedJointDataLoader`` draw independently +across ranks and converge to the same weights up to noise, whereas ``IterativeJointDataLoader`` seeds its draw so that every rank picks the same source at a given iteration, and rank-local weights would break that. @@ -52,6 +57,10 @@ class TokenMixControlConfig: """Config for the token-mix controller. + The historical ``*_batches`` names remain compatible with batch-loader configs. + EMA half-life and warmup count observations: the mixed loader maps its sample + settings to these fields and disables correction clipping explicitly. + Attributes: enabled: Whether to adapt the draw at all. When false the loader draws on the configured ratios and the token mix skews by whatever the packer's @@ -70,7 +79,8 @@ class TokenMixControlConfig: against a bad estimate starving a source rather than a tuning knob, but one with a cost: a mixture whose tokens per batch spread wider than this is only partly corrected, and the residual shows up as a realized share - that sits off target with the source pinned at the bound. + that sits off target with the source pinned at the bound. ``None`` + disables clipping, allowing arbitrary sample-length differences. log_every_n_batches: How often to log target against realized share. The realized figure is measured over that window, so it is also how long the diagnostic takes to reflect a change. @@ -79,7 +89,7 @@ class TokenMixControlConfig: enabled: bool = True half_life_batches: float = 50.0 warmup_batches_per_source: int = 8 - max_weight_ratio: float = 8.0 + max_weight_ratio: float | None = 8.0 log_every_n_batches: int = 500 def build(self, dataset_names: Sequence[str], target_shares: Sequence[float]) -> TokenMixController | None: @@ -102,7 +112,7 @@ def build(self, dataset_names: Sequence[str], target_shares: Sequence[float]) -> raise ValueError(f"half_life_batches must be positive, got {self.half_life_batches}.") if self.warmup_batches_per_source < 1: raise ValueError(f"warmup_batches_per_source must be at least 1, got {self.warmup_batches_per_source}.") - if self.max_weight_ratio < 1.0: + if self.max_weight_ratio is not None and self.max_weight_ratio < 1.0: raise ValueError(f"max_weight_ratio must be at least 1, got {self.max_weight_ratio}.") if self.log_every_n_batches < 1: raise ValueError(f"log_every_n_batches must be at least 1, got {self.log_every_n_batches}.") @@ -118,8 +128,9 @@ class TokenMixController: """Draw probabilities that keep each source's token share on its configured one. Attributes: - sampling_probs: The distribution to draw the next batch's source from. Sums - to one, and equals the configured shares until warmup completes. + sampling_probs: The distribution over all sources after the last + observation. Use ``probabilities(active)`` when sources can exhaust. + draw_counts: Number of positive-token observations from each source. """ def __init__( @@ -130,44 +141,52 @@ def __init__( ) -> None: """See :meth:`TokenMixControlConfig.build`, which validates and calls this.""" shares = np.asarray(target_shares, dtype=np.float64) - if shares.min() <= 0.0 or shares.sum() <= 0.0: - raise ValueError(f"Every target share must be positive, got {target_shares}.") + if shares.ndim != 1 or not shares.size or not np.isfinite(shares).all() or np.any(shares <= 0.0): + raise ValueError(f"Every target share must be positive and finite, got {target_shares}.") + if len(dataset_names) != len(shares): + raise ValueError(f"Got {len(dataset_names)} dataset names and {len(shares)} target shares.") self.names: list[str] = list(dataset_names) self.target_shares: np.ndarray = shares / shares.sum() - self.config = config + self.config: TokenMixControlConfig = config self.sampling_probs: np.ndarray = self.target_shares.copy() # Per observation rather than per iteration, matching half_life_batches. - self._decay = 0.5 ** (1.0 / config.half_life_batches) - self._mean_tokens = np.zeros(len(self.names), dtype=np.float64) - self._batches = np.zeros(len(self.names), dtype=np.int64) - self._window_tokens = np.zeros(len(self.names), dtype=np.float64) - self._window_batches = 0 - self._at_band = np.zeros(len(self.names), dtype=bool) + self._decay: float = 0.5 ** (1.0 / config.half_life_batches) + self._mean_tokens: np.ndarray = np.zeros(len(self.names), dtype=np.float64) + self.draw_counts: np.ndarray = np.zeros(len(self.names), dtype=np.int64) + self._window_tokens: np.ndarray = np.zeros(len(self.names), dtype=np.float64) + self._window_batches: int = 0 + self._at_band: np.ndarray = np.zeros(len(self.names), dtype=bool) - def observe(self, index: int, tokens: int) -> None: - """Fold one packed batch into the estimate and refresh the draw. + def observe(self, index: int, tokens: int, *, record_realized: bool = True) -> None: + """Fold one source draw into the estimate and refresh the probabilities. Args: - index: Source the batch was packed from, in the loader's index order. - tokens: Tokens in that batch. Non-positive counts are ignored, an empty - batch being a fact about the input stream rather than about cost. + index: Source of this observation, in the loader's index order. + tokens: Tokens in the batch or candidate. Non-positive counts are + ignored, an empty draw being a fact about the input stream. + record_realized: Record this observation in the accepted-token window. + Batch draws are already packed. Sample candidates may be deferred, + so their loader passes false and tracks accepted tokens separately. """ if tokens <= 0: return - self._batches[index] += 1 - if self._batches[index] == 1: + self.draw_counts[index] += 1 + if self.draw_counts[index] == 1: # Seed rather than decay from zero, which would read a source's first - # batches as cheap and draw it even harder while the error was largest. + # observations as cheap and draw it even harder while the error was largest. self._mean_tokens[index] = float(tokens) else: self._mean_tokens[index] *= self._decay self._mean_tokens[index] += (1.0 - self._decay) * tokens + self.sampling_probs = self.probabilities() + if not record_realized: + return + self._window_tokens[index] += tokens self._window_batches += 1 - self.sampling_probs = self._draw_probabilities() if self._window_batches >= self.config.log_every_n_batches: # Rank 0 only, for volume: every rank holds its own controller, but they see # the same data and differ by estimator noise, so one rank's view is the mix. @@ -185,10 +204,10 @@ def realized_shares(self) -> np.ndarray: def describe(self) -> str: """Target against realized share and the draw in force, one line per source.""" realized = self.realized_shares() - warming = int((self._batches < self.config.warmup_batches_per_source).sum()) + warming = int((self.draw_counts < self.config.warmup_batches_per_source).sum()) header = f"token mix over {self._window_batches} batches" if warming: - header += f" (warming up: {warming} source(s) below {self.config.warmup_batches_per_source} batches)" + header += f" (warming up: {warming} source(s) below {self.config.warmup_batches_per_source} observations)" lines = [header] for index, name in enumerate(self.names): # Naming the sources at the band is the whole diagnostic for a mix that is @@ -197,14 +216,18 @@ def describe(self) -> str: at_band = " (at the band, so only partly corrected)" if self._at_band[index] else "" lines.append( f" {name}: target {self.target_shares[index]:.1%}, realized {realized[index]:.1%}, " - f"drawn at {self.sampling_probs[index]:.1%}, {self._mean_tokens[index]:,.0f} tokens/batch{at_band}" + f"drawn at {self.sampling_probs[index]:.1%}, {self._mean_tokens[index]:,.0f} tokens/observation{at_band}" ) return "\n".join(lines) - def _draw_probabilities(self) -> np.ndarray: - """``ratio_i / T_i``, normalised, once every source has been measured. + def probabilities(self, active: np.ndarray | None = None) -> np.ndarray: + """``ratio_i / T_i``, normalised over the active sources after warmup. + + Disabled controllers return the configured ratios over active sources. + An all-false mask returns zeros; the caller must stop drawing in that case. + Inactive sources affect neither warmup nor correction clipping. - The correction ``1 / T_i`` is clipped in log space around its geometric mean + When configured, ``1 / T_i`` is clipped in log space around its geometric mean rather than after normalising. Ratios between corrections survive normalisation where absolute values do not, so bounding their spread there is what makes ``max_weight_ratio`` an exact bound on the departure of every @@ -212,12 +235,21 @@ def _draw_probabilities(self) -> np.ndarray: Also records which sources the clip caught, for :meth:`describe`. """ - if (self._batches < self.config.warmup_batches_per_source).any(): - return self.target_shares.copy() - - correction = -np.log(self._mean_tokens) - correction -= correction.mean() - half_band = 0.5 * np.log(self.config.max_weight_ratio) - self._at_band = np.abs(correction) > half_band - weights = self.target_shares * np.exp(np.clip(correction, -half_band, half_band)) + if active is None: + active = np.ones(len(self.names), dtype=bool) + weights = np.where(active, self.target_shares, 0.0) + self._at_band[:] = False + if not active.any(): + return weights + if not self.config.enabled or (self.draw_counts[active] < self.config.warmup_batches_per_source).any(): + return weights / weights.sum() + + if self.config.max_weight_ratio is None: + weights[active] /= self._mean_tokens[active] + else: + correction = -np.log(self._mean_tokens[active]) + correction -= correction.mean() + half_band = 0.5 * np.log(self.config.max_weight_ratio) + self._at_band[active] = np.abs(correction) > half_band + weights[active] *= np.exp(np.clip(correction, -half_band, half_band)) return weights / weights.sum() diff --git a/cosmos_framework/model/generator/algorithm/loss/flow_matching.py b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py index 7f8e227cc..a2264b8ad 100644 --- a/cosmos_framework/model/generator/algorithm/loss/flow_matching.py +++ b/cosmos_framework/model/generator/algorithm/loss/flow_matching.py @@ -32,19 +32,35 @@ class ActionSlotLossStats: sample_count: torch.Tensor +@dataclass +class FlowMatchingLossItems: + """Optional differentiable output for sample-aware supervised reductions. + + Logging uses the helper's existing unweighted return value. Training reducers + must use ``weighted_losses``, which includes per-frame timestep weights. + ``valid`` excludes items with no supervised elements, including fully masked + action channels. The legacy scalar reduction and logging outputs are unchanged. + """ + + weighted_losses: torch.Tensor | None = None # [N_items] + valid: torch.Tensor | None = None # [N_items] + + def compute_flow_matching_loss( pred: list[torch.Tensor], target: list[torch.Tensor], condition_mask: list[torch.Tensor], timesteps: torch.Tensor, has_valid_tokens: bool, - rectified_flow: RectifiedFlow, + rectified_flow: RectifiedFlow | list[RectifiedFlow], tensor_kwargs_fp32: dict, raw_action_dim: list[torch.Tensor] | None = None, action_valid_mask: list[torch.Tensor] | None = None, normalize_by_active: bool = False, exclude_fully_conditioned_items: bool = False, action_slot_stats: ActionSlotLossStats | None = None, + item_losses: FlowMatchingLossItems | None = None, + time_axis: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute flow matching loss for a modality. @@ -76,6 +92,10 @@ def compute_flow_matching_loss( samples from diluting the generated target loss. action_slot_stats: Optional collector for detached normalized per-sample losses over the canonical unified Action slots. + item_losses: Optional collector for differentiable timestep-weighted item + losses and supervision validity. Does not change the legacy reduction. + time_axis: Temporal axis of the conditioning mask. The new reduction uses + axis 1 for [C,T] sound predictions; default 0 preserves legacy weighting. Returns: tuple: A tuple containing two elements: @@ -85,6 +105,9 @@ def compute_flow_matching_loss( if not has_valid_tokens: # Dummy loss to maintain backward graph consistency across ranks dummy_loss = 0.0 * sum(p.sum() for p in pred) + if item_losses is not None: + item_losses.weighted_losses = torch.stack([p.sum() * 0.0 for p in pred]) # [N_items] + item_losses.valid = torch.zeros(len(pred), dtype=torch.bool, device=dummy_loss.device) # [N_items] return dummy_loss, dummy_loss.unsqueeze(0) # make per-instance loss 1-D # condition_mask[i] is T-first with trailing singletons: [T,1,1] vision, [T,1] action. @@ -92,8 +115,9 @@ def compute_flow_matching_loss( per_instance_losses = [] per_instance_weighted_losses = [] has_noisy_items: list[torch.Tensor] = [] + has_supervised_items: list[torch.Tensor] = [] for i in range(len(pred)): - T_i = condition_mask[i].shape[0] + T_i = condition_mask[i].shape[time_axis] sqerr_i = (pred[i] - target[i]) ** 2 # vision:[C,T,H,W] action/sound:[T,D] noisy_mask_i = 1.0 - condition_mask[i] # vision:[T,1,1] action/sound:[T,1] has_noisy_items.append(torch.any(noisy_mask_i != 0)) # [] @@ -106,6 +130,11 @@ def compute_flow_matching_loss( raise ValueError(f"action_valid_mask width {slot_mask_i.shape[-1]} < action width {sqerr_i.shape[-1]}") slot_mask_i = slot_mask_i[:, : sqerr_i.shape[-1]] sqerr_i = sqerr_i * slot_mask_i + if item_losses is not None: + valid_i = torch.any(noisy_mask_i != 0) & (sqerr_i.numel() > 0) # [] + if slot_mask_i is not None: + valid_i = valid_i & torch.any(slot_mask_i != 0) # [] + has_supervised_items.append(valid_i) if normalize_by_active: active_channels = slot_mask_i.sum() if slot_mask_i is not None else sqerr_i.numel() // noisy_mask_i.numel() active_count = (noisy_mask_i.sum() * active_channels).clamp(min=1) @@ -120,8 +149,11 @@ def compute_flow_matching_loss( per_instance_losses.append((sqerr_i * noisy_mask_i).mean()) # [] ts_i = timesteps[i, :T_i] if timesteps.dim() > 1 else timesteps[i] # DF:[T_i] TF:[1] - tw_i = rectified_flow.train_time_weight(ts_i, tensor_kwargs_fp32) # DF:[T_i] TF:[1] - tw_i = tw_i.reshape(-1, *([1] * (condition_mask[i].ndim - 1))) # vision:[T_i,1,1] action/sound:[T_i,1] + flow_i = rectified_flow[i] if isinstance(rectified_flow, list) else rectified_flow + tw_i = flow_i.train_time_weight(ts_i, tensor_kwargs_fp32) # DF:[T_i] TF:[1] + weight_shape = [1] * condition_mask[i].ndim + weight_shape[time_axis] = -1 + tw_i = tw_i.reshape(weight_shape) # vision:[T_i,1,1] action:[T_i,1] sound:[1,T_i] weighted_sqerr_i = sqerr_i * tw_i * noisy_mask_i if normalize_by_active or slot_mask_i is not None: per_instance_weighted_losses.append(weighted_sqerr_i.sum() / active_count) @@ -155,6 +187,9 @@ def compute_flow_matching_loss( per_instance_loss = torch.stack(per_instance_losses) # [B] per_instance_weighted_loss = torch.stack(per_instance_weighted_losses) # [B] + if item_losses is not None: + item_losses.weighted_losses = per_instance_weighted_loss # [B] + item_losses.valid = torch.stack(has_supervised_items) # [B] if exclude_fully_conditioned_items: active_item_mask = torch.stack(has_noisy_items).to( device=per_instance_weighted_loss.device, diff --git a/cosmos_framework/model/generator/algorithm/loss/load_balancing.py b/cosmos_framework/model/generator/algorithm/loss/load_balancing.py index 94bd1fbaa..bd6f4203b 100644 --- a/cosmos_framework/model/generator/algorithm/loss/load_balancing.py +++ b/cosmos_framework/model/generator/algorithm/loss/load_balancing.py @@ -31,7 +31,7 @@ def _scale_rank_sample_mean( device_mesh, ) # [num_layers] # FSDP averages gradients over the mesh, so pre-multiply by the mesh size. - # This is analogous to OmniMoTModel._sample_level_loss_scale. + # This also counteracts gradient averaging in the supervised per-modality sample means. pre_fsdp_sample_scale = ( device_mesh.size() * local_sample_count_per_layer / global_sample_count_per_layer.clamp_min(1) ) # [num_layers] diff --git a/cosmos_framework/model/generator/algorithm/loss/modality_reduction.py b/cosmos_framework/model/generator/algorithm/loss/modality_reduction.py new file mode 100644 index 000000000..8f45b8688 --- /dev/null +++ b/cosmos_framework/model/generator/algorithm/loss/modality_reduction.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Global sample means for the supervised terms of one training microbatch. + +The objective is a fixed weighted sum of image, video, action, sound, LiDAR +and radar sample means. A sample can supervise several modalities, but owns at most one +entry in each modality's denominator. Multiple supervised items in that entry +are averaged first. This prevents extra camera/control items, missing audio, or +the number of tokens in a sample from implicitly changing its loss weight. + +These are timestep-weighted training losses. The unweighted per-instance values +returned for logging by ``compute_flow_matching_loss`` are not suitable inputs. +The trainer retains its existing gradient-accumulation averaging: normalization +here is per global microbatch, not across an optimizer's accumulation window. +MoE auxiliary losses are added by the caller after this supervised reduction. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from cosmos_framework.model.generator.algorithm.loss.flow_matching import FlowMatchingLossItems + +MODALITIES: tuple[str, ...] = ("image", "video", "action", "sound", "lidar", "radar") + + +@dataclass(frozen=True) +class ModalityLossItems: + """Differentiable item losses and their ownership in a single modality. + + ``valid`` excludes fully conditioned items and items with no valid supervised + elements. An enabled but locally absent modality should supply a zero loss + connected to its network predictions with ``valid=False``; its parameters + then participate in backward without adding a sample to the denominator. + Sample IDs only need to be unique within the local batch, and may be sparse. + """ + + weighted_losses: torch.Tensor # [N_items] + sample_ids: torch.Tensor # [N_items] + valid: torch.Tensor # [N_items] + + def sample_sum_and_count(self) -> tuple[torch.Tensor, torch.Tensor]: + """Average supervised items per sample, then return a sum and sample count.""" + losses = self.weighted_losses # [N_items] + if losses.ndim != 1 or self.sample_ids.shape != losses.shape or self.valid.shape != losses.shape: + raise ValueError("Item losses, sample IDs and validity must be aligned one-dimensional tensors.") + if self.sample_ids.dtype != torch.long or self.valid.dtype != torch.bool: + raise ValueError("Sample IDs must be int64 and item validity must be boolean.") + if self.sample_ids.device != losses.device or self.valid.device != losses.device: + raise ValueError("Losses, sample IDs and validity must share a device.") + if losses.numel() == 0: + return losses.sum(), losses.new_zeros(()) # [], [] + + sample_ids, inverse = torch.unique(self.sample_ids, sorted=True, return_inverse=True) # [B], [N_items] + valid_losses = torch.where(self.valid, losses, torch.zeros_like(losses)) # [N_items] + sample_sums = losses.new_zeros(sample_ids.numel()).scatter_add(0, inverse, valid_losses) # [B] + item_counts = losses.new_zeros(sample_ids.numel()).scatter_add(0, inverse, self.valid.to(losses.dtype)) # [B] + sample_means = sample_sums / item_counts.clamp(min=1) # [B] + sample_count = item_counts.gt(0).sum().to(losses.dtype) # [] + return sample_means.sum(), sample_count # [], [] + + +def build_modality_loss_items( + collectors: Mapping[str, FlowMatchingLossItems], + predictions: Mapping[str, list[torch.Tensor]], + sample_ids: Mapping[str, list[int]], + *, + vision_is_image: list[bool] | None, + is_image_batch: bool, +) -> dict[str, ModalityLossItems]: + """Associate weighted flow losses with logical owners and image/video identity. + + Absent enabled streams retain the network's zero-weighted probes, so all + enabled heads remain connected to backward even on a rank without that data. + """ + result: dict[str, ModalityLossItems] = {} + for name, collector in collectors.items(): + losses, valid = collector.weighted_losses, collector.valid # [N_items] or None + if losses is None or valid is None: + losses = torch.stack([prediction.sum() * 0.0 for prediction in predictions[f"preds_{name}"]]) # [N_probe] + valid = torch.zeros_like(losses, dtype=torch.bool) # [N_probe] + if name not in sample_ids: + raise ValueError(f"Missing logical sample ownership for {name}.") + owners = sample_ids[name] + # An absent stream has no owners but carries one graph-connected probe. + owner_ids = owners if owners else [0] * losses.numel() + if len(owner_ids) != losses.numel(): + raise ValueError(f"{name} has {losses.numel()} item losses but {len(owner_ids)} owners.") + owner_tensor = torch.tensor(owner_ids, dtype=torch.long, device=losses.device) # [N_items] + if name == "vision": + flags = vision_is_image if vision_is_image else [is_image_batch] * losses.numel() + if len(flags) != losses.numel(): + raise ValueError("Vision image identities must align with item losses.") + image_mask = torch.tensor(flags, dtype=torch.bool, device=losses.device) # [N_items] + for modality, mask in (("image", image_mask), ("video", ~image_mask)): + result[modality] = ModalityLossItems(losses[mask], owner_tensor[mask], valid[mask]) + else: + result[name] = ModalityLossItems(losses, owner_tensor, valid) + return result + + +def reduce_global_modality_means( + items: Mapping[str, ModalityLossItems], + weights: Mapping[str, float], + *, + group: dist.ProcessGroup | None = None, + gradient_average_size: int = 1, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Return the weighted global supervised objective and unscaled modality means. + + Every rank reduces the same six sums/counts, even when it has no local data + for a modality. Differentiable local sums are multiplied by the actual + FSDP/DDP gradient-averaging group size and divided by global sample counts. + Detached global sums supply consistent logged values without changing that + backward computation. CP copies must participate through the same averaging + group as the model: equal replication of sums and counts cancels, rather than + giving each copy a new share of the objective. + Sharded CP hidden states additionally require summed gather gradients + (``correct_cp_gradients=True``); counting replicas alone cannot correct the + backbone gradient without incorrectly rescaling replicated output heads. + Disabling correction is allowed for ablations and retains legacy backbone + gradient scaling while preserving the forward loss values. + + Fixed configured weights are summed, never normalized over present modalities. + A globally absent modality contributes a graph-connected zero and keeps its + weight; it does not transfer that weight to other modalities. This deliberately + differs from both legacy rank-level and homogeneous sample-level averaging. + """ + unknown = (set(items) | set(weights)) - set(MODALITIES) + if unknown: + raise ValueError(f"Unknown supervised modalities: {sorted(unknown)}") + if not items: + raise ValueError("At least one modality must provide losses or graph-connected dummy predictions.") + if gradient_average_size < 1: + raise ValueError("gradient_average_size must be positive.") + if gradient_average_size > 1: + if not dist.is_initialized() or dist.get_world_size(group) != gradient_average_size: + raise ValueError("The reduction group must match the model's gradient-averaging group size.") + + reference = next(iter(items.values())).weighted_losses # [N_items] + local_sums: list[torch.Tensor] = [] + local_counts: list[torch.Tensor] = [] + for name in MODALITIES: + if name in items: + local_sum, local_count = items[name].sample_sum_and_count() # [], [] + else: + local_sum = reference.sum() * 0.0 # [] + local_count = reference.new_zeros(()) # [] + local_sums.append(local_sum) + local_counts.append(local_count) + + sums = torch.stack(local_sums) # [6] + counts = torch.stack(local_counts) # [6] + # Counts are logical contributing samples, not tokens or flattened item counts. + # FP64 keeps the small collective stable for large distributed sample totals. + global_stats = torch.stack((sums.detach().double(), counts.detach().double())) # [2,6] + if gradient_average_size > 1: + dist.all_reduce(global_stats, group=group) + + total_loss = sums.sum() * 0.0 # [] + means: dict[str, torch.Tensor] = {} + for index, name in enumerate(MODALITIES): + denominator = global_stats[1, index].clamp(min=1).to(sums.dtype) # [] + backward_mean = sums[index] * gradient_average_size / denominator # [] + logged_mean = (global_stats[0, index] / global_stats[1, index].clamp(min=1)).to(sums.dtype) # [] + mean = logged_mean + (backward_mean - backward_mean.detach()) # [] + means[name] = mean + total_loss = total_loss + float(weights.get(name, 0.0)) * mean # [] + return total_loss, means diff --git a/cosmos_framework/model/generator/algorithm/loss/modality_reduction_distributed_test.py b/cosmos_framework/model/generator/algorithm/loss/modality_reduction_distributed_test.py new file mode 100644 index 000000000..e65087557 --- /dev/null +++ b/cosmos_framework/model/generator/algorithm/loss/modality_reduction_distributed_test.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Real Gloo/DDP checks of global sample means with sharded CP intermediates.""" + +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.nn.parallel import DistributedDataParallel + +from cosmos_framework.model.generator.algorithm.loss.modality_reduction import ModalityLossItems, reduce_global_modality_means +from cosmos_framework.model.generator.mot.context_parallel_utils import all_gather_tensor + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] +_WEIGHTS = {"image": 2.0, "video": 3.0, "action": 5.0, "sound": 7.0} + + +class _Network(torch.nn.Module): + def __init__(self, cp_mesh: DeviceMesh | None) -> None: + super().__init__() + self.backbone: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(1.5)) # [] + self.heads: torch.nn.Parameter = torch.nn.Parameter(torch.tensor([1.0, 2.0, 3.0, 4.0])) # [4] + self.cp_mesh: DeviceMesh | None = cp_mesh + + def forward(self, features: torch.Tensor) -> torch.Tensor: # [N] -> [4,4] + hidden = self.backbone * features # [N] + if self.cp_mesh is not None: + hidden = all_gather_tensor(hidden, 0, self.cp_mesh, correct_cp_gradients=True) # [4] + return hidden[:, None] * self.heads[None, :] # [4,4] + + +def _items(predictions: torch.Tensor, owner: int) -> dict[str, ModalityLossItems]: + # Group 0 has two controls/targets in one image sample, another image, and + # one video. Group 1 has four videos, only the last of which has actions. + weights = 1 + (torch.arange(4, dtype=predictions.dtype) + 1 + owner * 4) / 10 # [4] + losses = predictions.square() * weights[:, None] # [4,4] + ids = torch.tensor([0, 0, 1, 2] if owner == 0 else [0, 1, 2, 3]) + owner * 10 # [4] + masks = { + "image": [True, True, True, False] if owner == 0 else [False] * 4, + "video": [False, False, False, True] if owner == 0 else [True] * 4, + "action": [False, False, False, owner == 1], + "sound": [False] * 4, + } + return { + name: ModalityLossItems(losses[:, index], ids, torch.tensor(masks[name])) for index, name in enumerate(_WEIGHTS) + } + + +def _worker(rank: int, cp_size: int, port: int) -> None: + torch.set_num_threads(1) + world_size = 2 * cp_size + dist.init_process_group( + "gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world_size, timeout=timedelta(seconds=90) + ) + mesh = init_device_mesh("cpu", (2, cp_size), mesh_dim_names=("dp", "cp")) + owner = rank // cp_size + network = _Network(mesh["cp"] if cp_size > 1 else None) + ddp = DistributedDataParallel(network) + for reverse_owners in (False, True): + batch_owner = 1 - owner if reverse_owners else owner + features = torch.arange(1 + batch_owner * 4, 5 + batch_owner * 4, dtype=torch.float32) # [4] + local_features = features.chunk(cp_size)[rank % cp_size] # [4/CP] + ddp.zero_grad(set_to_none=True) + predictions = ddp(local_features) # [4,4] + loss, _ = reduce_global_modality_means( + _items(predictions, batch_owner), _WEIGHTS, gradient_average_size=world_size + ) # [] + loss.backward() + + reference = _Network(None) + groups = [ + _items(reference(torch.arange(1 + index * 4, 5 + index * 4, dtype=torch.float32)), index) + for index in range(2) + ] + combined = { + name: ModalityLossItems( + torch.cat([group[name].weighted_losses for group in groups]), # [8] + torch.cat([group[name].sample_ids for group in groups]), # [8] + torch.cat([group[name].valid for group in groups]), # [8] + ) + for name in _WEIGHTS + } + # The global reference is local: avoid entering an extra all-reduce here. + reference_loss = sum( + _WEIGHTS[name] * item.sample_sum_and_count()[0] / item.sample_sum_and_count()[1].clamp(min=1) + for name, item in combined.items() + ) # [] + reference_loss.backward() + torch.testing.assert_close(loss, reference_loss) + torch.testing.assert_close(network.backbone.grad, reference.backbone.grad) + torch.testing.assert_close(network.heads.grad, reference.heads.grad) + assert network.heads.grad[3].item() == 0 + dist.destroy_process_group() + + +@pytest.mark.parametrize("cp_size", [1, 2]) +def test_ddp_global_means_and_cp_backbone_head_gradients(cp_size: int) -> None: + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + mp.spawn(_worker, args=(cp_size, port), nprocs=2 * cp_size, join=True) diff --git a/cosmos_framework/model/generator/algorithm/loss/modality_reduction_test.py b/cosmos_framework/model/generator/algorithm/loss/modality_reduction_test.py new file mode 100644 index 000000000..6b9846840 --- /dev/null +++ b/cosmos_framework/model/generator/algorithm/loss/modality_reduction_test.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Reference objectives for sparse modalities and uneven distributed batches.""" + +from unittest.mock import patch + +import pytest +import torch + +from cosmos_framework.model.generator.algorithm.loss.modality_reduction import ( + MODALITIES, + ModalityLossItems, + reduce_global_modality_means, +) + +pytestmark = [pytest.mark.L0, pytest.mark.CPU] + + +def _items(losses: torch.Tensor, owners: list[int], valid: list[bool] | None = None) -> ModalityLossItems: + return ModalityLossItems( + weighted_losses=losses, + sample_ids=torch.tensor(owners, dtype=torch.long), # [N_items] + valid=torch.tensor(valid if valid is not None else [True] * len(owners)), # [N_items] + ) + + +def test_sample_means_exclude_invalid_items_and_preserve_fixed_weights() -> None: + image = torch.tensor([2.0, 6.0, 1000.0, 10.0], requires_grad=True) # [4] + video = torch.tensor([7.0], requires_grad=True) # [1] + action = torch.tensor([3.0, 900.0], requires_grad=True) # [2] + sound_probe = torch.tensor([8.0], requires_grad=True) # [1] + loss, means = reduce_global_modality_means( + { + "image": _items(image, [0, 0, 0, 1], [True, True, False, True]), + "video": _items(video, [2]), + "action": _items(action, [2, 3], [True, False]), + "sound": _items(sound_probe * 0, [0], [False]), + }, + {"image": 2.0, "video": 3.0, "action": 10.0, "sound": 100.0}, + ) + # Image sample means are 4 and 10. Neither invalid items nor the absent audio + # modality dilute the sample denominators or redistribute configured weights. + assert loss.item() == pytest.approx(2 * 7 + 3 * 7 + 10 * 3) + assert means["sound"].item() == 0 + loss.backward() + torch.testing.assert_close(image.grad, torch.tensor([0.5, 0.5, 0.0, 1.0])) + torch.testing.assert_close(video.grad, torch.tensor([3.0])) + torch.testing.assert_close(action.grad, torch.tensor([10.0, 0.0])) + torch.testing.assert_close(sound_probe.grad, torch.zeros(1)) + + +def test_reordering_items_keeps_objective_and_gradients() -> None: + values = torch.tensor([1.0, 5.0, 9.0, 11.0], requires_grad=True) # [4] + permutation = torch.tensor([3, 1, 0, 2]) # [4] + baseline, _ = reduce_global_modality_means({"video": _items(values, [0, 1, 1, 2])}, {"video": 2.0}) + reordered, _ = reduce_global_modality_means({"video": _items(values[permutation], [2, 1, 0, 1])}, {"video": 2.0}) + baseline_grad = torch.autograd.grad(baseline, values, retain_graph=True)[0] # [4] + reordered_grad = torch.autograd.grad(reordered, values)[0] # [4] + torch.testing.assert_close(baseline, reordered) + torch.testing.assert_close(baseline_grad, reordered_grad) + + +@pytest.mark.parametrize("cp_copies", [1, 2]) +def test_uneven_rank_partitions_and_cp_copies_match_global_reference(cp_copies: int) -> None: + # Two independent batches: the first rank has two images and one video, the + # second only a video with action. CP replicas process the same owner batch. + parameter = torch.tensor(2.0, requires_grad=True) # [] + image_losses = parameter.square() * torch.tensor([1.0, 3.0]) # [2] + video_losses = parameter.square() * torch.tensor([5.0, 7.0]) # [2] + action_losses = parameter.square() * torch.tensor([11.0]) # [1] + global_items = { + "image": _items(image_losses, [0, 1]), + "video": _items(video_losses, [2, 3]), + "action": _items(action_losses, [3]), + } + weights = {"image": 2.0, "video": 3.0, "action": 0.5, "sound": 20.0} + reference, _ = reduce_global_modality_means(global_items, weights) + reference_gradient = torch.autograd.grad(reference, parameter, retain_graph=True)[0] # [] + sums = torch.tensor([16.0, 48.0, 44.0, 0.0, 0.0, 0.0], dtype=torch.float64) # [6] + counts = torch.tensor([2.0, 2.0, 1.0, 0.0, 0.0, 0.0], dtype=torch.float64) # [6] + global_stats = torch.stack((sums, counts)) * cp_copies # [2,6] + world_size = 2 * cp_copies + local_items = [ + {"image": global_items["image"], "video": _items(video_losses[:1], [2])}, + {"video": _items(video_losses[1:], [3]), "action": global_items["action"]}, + ] + + def all_reduce(tensor: torch.Tensor, **kwargs: object) -> None: + tensor.copy_(global_stats) # [2,6] + + gradients: list[torch.Tensor] = [] + with ( + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_world_size", return_value=world_size), + patch("torch.distributed.all_reduce", side_effect=all_reduce) as collective, + ): + for rank in range(world_size): + local_loss, means = reduce_global_modality_means( + local_items[rank // cp_copies], weights, gradient_average_size=world_size + ) + torch.testing.assert_close(local_loss, reference) + assert set(means) == set(MODALITIES) + gradients.append(torch.autograd.grad(local_loss, parameter, retain_graph=True)[0]) # [] + assert collective.call_count == world_size + averaged_gradient = torch.stack(gradients).mean() # [] + torch.testing.assert_close(averaged_gradient, reference_gradient) + + +def test_invalid_metadata_fails_before_reduction() -> None: + values = torch.ones(2, requires_grad=True) # [2] + with pytest.raises(ValueError, match="aligned"): + reduce_global_modality_means({"image": _items(values, [0])}, {"image": 1.0}) + with pytest.raises(ValueError, match="Unknown"): + reduce_global_modality_means({"unknown": _items(values, [0, 1])}, {}) + + +def test_radar_sample_means_retain_supervision_and_absent_head_gradients() -> None: + radar = torch.tensor([2.0, 6.0, 10.0, 99.0], requires_grad=True) # [4] + probe = torch.tensor([7.0], requires_grad=True) # [1] + loss, means = reduce_global_modality_means( + {"radar": _items(radar, [0, 0, 2, 2], [True, True, True, False])}, + {"radar": 3.0}, + ) # [], dict of [] + assert means["radar"].item() == pytest.approx(7.0) + assert loss.item() == pytest.approx(21.0) + loss.backward() + torch.testing.assert_close(radar.grad, torch.tensor([0.75, 0.75, 1.5, 0.0])) + absent_loss, absent_means = reduce_global_modality_means( + {"radar": _items(probe * 0, [0], [False])}, {"radar": 3.0} + ) # [], dict of [] + absent_loss.backward() + assert absent_means["radar"].item() == 0 + torch.testing.assert_close(probe.grad, torch.zeros(1)) diff --git a/cosmos_framework/model/generator/distillation/dmd2_rf_test.py b/cosmos_framework/model/generator/distillation/dmd2_rf_test.py index 835157a1b..c88e5a39d 100644 --- a/cosmos_framework/model/generator/distillation/dmd2_rf_test.py +++ b/cosmos_framework/model/generator/distillation/dmd2_rf_test.py @@ -343,6 +343,7 @@ def test_build_net_lora_override_controls_injection_and_initialization( model.config.lora_rank = 8 model.config.lora_alpha = 16 model.config.lora_target_modules = "q_proj_moe_gen" + model.config.multiview_action_conditioning = None model.vlm_config = MagicMock() model.tokenizer_vision_gen = MagicMock() model.parallel_dims = MagicMock() diff --git a/cosmos_framework/model/generator/mot/causal_attention.py b/cosmos_framework/model/generator/mot/causal_attention.py index 4279e33ad..56ec7797f 100644 --- a/cosmos_framework/model/generator/mot/causal_attention.py +++ b/cosmos_framework/model/generator/mot/causal_attention.py @@ -42,7 +42,6 @@ replay_maskless_attention, ) from cosmos_framework.model.generator.mot.merge_attention import merge_attentions_ac_safe -from cosmos_framework.model.generator.rolling_prompt import RollingPromptLayout, RollingTextSinkLayer from cosmos_framework.model.generator.utils.kv_cache import ( ARMemoryValue, JointARMemoryValue, @@ -51,6 +50,10 @@ TFNoisyMemoryValue, TFReplayCleanMemoryValue, ) +from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( + RollingPromptLayout, + RollingTextSinkLayer, +) class ConcatenateAttentionsBridge(torch.autograd.Function): diff --git a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py index 943ff136f..e98c40a82 100644 --- a/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/causal_cosmos3_vfm_network.py @@ -36,7 +36,10 @@ build_teacher_forcing_multiview_flex_metadata, ) from cosmos_framework.model.generator.mot.maskless_attention import build_replay_maskless_plan -from cosmos_framework.model.generator.rolling_replay import RollingReplayRequest, build_rolling_replay_metadata +from cosmos_framework.model.generator.utils.rolling_kv.rolling_replay import ( + RollingReplayRequest, + build_rolling_replay_metadata, +) def build_interactive_multiview_mask_items( diff --git a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py index 381a757c5..fc8e7c4d1 100644 --- a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py @@ -13,6 +13,7 @@ from cosmos_framework.utils import log from cosmos_framework.configs.base.defaults.joint_attention import packing_layout from cosmos_framework.configs.base.defaults.multiview_attention import ( + CaptionAccess, MultiviewAttentionConfig, ResolvedBackend, TemporalWindow, @@ -90,6 +91,7 @@ def __init__( predict_text_tokens=False, joint_attn_implementation="two_way", multiview_attention_config: MultiviewAttentionConfig | None = None, + multiview_action_conditioning: bool = False, action_dim=32, num_embodiment_domains=32, action_io_projector_type: str = ACTION_IO_PROJECTOR_DOMAIN_AWARE, @@ -148,6 +150,7 @@ def __init__( # window, the folds read all of it through ``maskless_unavailable_reason``, and a copy of # each on this config could disagree with the other. self.multiview_attention_config = multiview_attention_config or MultiviewAttentionConfig() + self.multiview_action_conditioning = multiview_action_conditioning # Exported checkpoint configs can arrive as plain DictConfig objects, bypassing # the attrs converter. Restore the legacy counting default and normalize durations # here, before building attention. @@ -849,6 +852,14 @@ def _encode_vision( """ if self.config.enable_vision_modality_embeddings: modality_embed = self.image_modality_embed if packed_seq.is_image_batch else self.video_modality_embed + if packed_seq.vision_is_image is not None and packed_seq.vision is not None: + shapes = packed_seq.vision.token_shapes + if len(packed_seq.vision_is_image) != len(shapes): + raise ValueError("Image identities must match the packed vision items.") + flags = torch.tensor(packed_seq.vision_is_image, device=packed_sequence.device, dtype=torch.long) # [I] + lengths = torch.tensor([t * h * w for t, h, w in shapes], device=flags.device) # [I] + embeddings = torch.stack((self.video_modality_embed, self.image_modality_embed)) # [2,H] + modality_embed = embeddings[flags].repeat_interleave(lengths, dim=0) # [N_vision,H] elif self.config.enable_media_modality_embedding: modality_embed = self.media_modality_embed else: @@ -959,7 +970,12 @@ def _encode_grid_stream( packed_tokens, modality.token_shapes, view_ids, self.rig_view_embed ) # [total_patches,hidden_size] if modality_embed is not None: - packed_tokens = packed_tokens + modality_embed.view(1, -1) # [total_patches,hidden_size] + if modality_embed.ndim == 2: + # The linear output is not saved for backward. Reuse its buffer + # so per-token identities do not add another full hidden-state allocation. + packed_tokens.add_(modality_embed) # [total_patches,hidden_size] + else: + packed_tokens = packed_tokens + modality_embed # [total_patches,hidden_size] if modality.mse_loss_indexes.numel() > 0: timesteps = modality.timesteps.to(dtype=torch.float32) * self.timestep_scale # [N_noisy_frames] @@ -1006,6 +1022,12 @@ def _decode_vision( ) ) + if getattr(packed_seq, "vision_is_image", None) is not None and self.config.enable_vision_modality_embeddings: + # A vision-less rank must still reach both enabled image/video + # embeddings when another rank has camera supervision. + embedding_probe = 0.0 * (self.image_modality_embed.sum() + self.video_modality_embed.sum()) # [] + output_dict["preds_vision"][0] = output_dict["preds_vision"][0] + embedding_probe # [1,C,T,H,W] + def _decode_lidar( self, packed_seq: PackedSequence, @@ -1339,6 +1361,78 @@ def _prepare_multiview_attention( if self.multiview_backend is None: return + if getattr(self.config, "multiview_action_conditioning", False): + counts = packed_seq.num_vision_items_per_sample + if counts is None: + # Preserve dense attention for the unchanged single-view streams + # mixed into an opt-in run; their existing action block is unchanged. + return + if ( + self.multiview_backend != "maskless" + or packed_seq.num_action_tokens_per_supertoken != 0 + or packed_seq.vision is None + or packed_seq.action is None + or packed_seq.sound is not None + or packed_seq.lidar is not None + or packed_seq.radar is not None + ): + raise ValueError("Multiview action conditioning requires RGB/action blocks and maskless attention.") + config = self.config.multiview_attention_config + full_seq, _ = get_full_only_seq(input_pack) + shapes: list[tuple[int, int, int]] = [] + rates: list[float] = [] + view_axis: list[int] = [] + frame_offsets: list[int] = [] + items_per_sample: list[int] = [] + factor = self.config.temporal_compression_factor_vision + if len(packed_seq.action.token_shapes) != len(counts): + raise ValueError("Multiview packs require one camera-major action payload per sample.") + vision_cursor = 0 + for sample, num_views in enumerate(counts): + view_shapes = packed_seq.vision.token_shapes[vision_cursor : vision_cursor + num_views] + view_rates = packed_seq.vision.seconds_per_frame[vision_cursor : vision_cursor + num_views] + latent_t = view_shapes[0][0] + if any(shape[0] != latent_t for shape in view_shapes) or packed_seq.action.token_shapes[sample] != ( + num_views * (latent_t - 1) * factor, + ): + raise ValueError("Multiview action blocks must retain every frame-rate action for every view.") + shapes.extend((t, h, w) for t, h, w in view_shapes) + rates.extend(view_rates) + view_axis.extend(range(num_views)) + frame_offsets.extend([0] * num_views) + # The action block stays separate. Describe its full frame-rate + # rows in latent-sized groups, sharing each RGB view's attention + # identity. Match pack_supertokens_temporal_causal's [null, V0], + # [actions[0:factor], V1], ... alignment, but omit its null prefix. + # The first real action group therefore starts at latent frame 1. + if latent_t > 1: + shapes.extend([(latent_t - 1, 1, factor)] * num_views) + rates.extend(view_rates) + view_axis.extend(range(num_views)) + frame_offsets.extend([1] * num_views) + items_per_sample.append(num_views * (2 if latent_t > 1 else 1)) + vision_cursor += num_views + caption_access: list[CaptionAccess] = ["all_captions"] * len(shapes) + # All RGB items are synchronized targets, not Auto's prefix controls/ + # last target. Conditioned actions join the same view/time groups; + # their clean inputs do not make them separate control-image streams. + attention_meta.multiview_maskless = build_multiview_maskless_plan( + [1] * len(shapes), + shapes, + device=full_seq.device, + seconds_per_frame=rates, + frame_offsets=frame_offsets, + items_per_sample=items_per_sample, + is_control=[False] * len(shapes), + view_axis=view_axis, + caption_access=caption_access, + padded_gen_tokens=full_seq.shape[0], + attention_scope=config.mask.attention_scope, + deduplicate_cross_view=config.deduplicate_cross_view, + decomposed_temporal_window_seconds=config.mask.decomposed_temporal_window_seconds, + ) + return + # No pathway check here: ``multiview_backend`` is non-None exactly when # ``joint_attn_implementation == "multiview"``, so the two cannot disagree. # natten_metadata_list is always None here (only the three-way packer builds it). diff --git a/cosmos_framework/model/generator/mot/cosmos3_vfm_network_multiview_forward_test.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network_multiview_forward_test.py index cf67ba380..d6d8ddde7 100644 --- a/cosmos_framework/model/generator/mot/cosmos3_vfm_network_multiview_forward_test.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network_multiview_forward_test.py @@ -24,21 +24,31 @@ import random from types import SimpleNamespace +from typing import TYPE_CHECKING import attrs import pytest import torch from omegaconf import OmegaConf +from cosmos_framework.configs.base.defaults.model_config import MultiviewActionConditioningConfig from cosmos_framework.configs.base.defaults.multiview_attention import MultiviewAttentionConfig +from cosmos_framework.data.generator.action.utils.transforms import build_sequence_plan_from_mode from cosmos_framework.model.generator.mot.attention import SplitInfo, dispatch_attention from cosmos_framework.model.generator.mot.multiview_attention import resolve_multiview_backend from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.data.generator.sequence_packing import PackedSequence -from cosmos_framework.data.generator.sequence_packing.packers import pack_input_sequence +from cosmos_framework.data.generator.sequence_packing.packers import ( + pack_input_sequence, + pack_multiview_action_conditioning, + replicate_multiview_actions, +) from cosmos_framework.data.generator.sequence_packing.runtime import SequencePack from cosmos_framework.data.generator.sequence_packing.sequence import SequencePlan +if TYPE_CHECKING: + from cosmos_framework.model.generator.mot.cosmos3_vfm_network import Cosmos3VFMNetwork + # The rig this harness packs: two cameras, two latent frames each, on a 4x4 latent grid that # one patch step halves to 2x2. Sixteen GEN tokens in all -- small enough that the flex mask's # 128-token block padding is almost the whole stream, which is the case the decomposition has @@ -191,7 +201,8 @@ def _multiview_network( device: torch.device, temporal_window_seconds: tuple[float, float] | None = None, deduplicate_cross_view: bool = False, -) -> torch.nn.Module: + action_conditioning: bool = False, +) -> Cosmos3VFMNetwork: """The network under test, on the stub reasoner, with the multiview mask configured.""" from cosmos_framework.configs.base.defaults.multiview_attention import ( MultiviewAttentionConfig, @@ -206,6 +217,9 @@ def _multiview_network( config = Cosmos3VFMNetworkConfig( vlm_config=language_model.config, vision_gen=True, + action_gen=action_conditioning, + action_dim=64 if action_conditioning else 32, + multiview_action_conditioning=action_conditioning, latent_channel_size=LATENT_CHANNELS, latent_patch_size=PATCH_SPATIAL, latent_downsample_factor=16, @@ -375,3 +389,51 @@ def test_forward_trains_through_the_decomposition() -> None: assert network.vae2llm.weight.grad is not None assert torch.isfinite(network.vae2llm.weight.grad).all() assert float(network.vae2llm.weight.grad.abs().sum()) > 0.0 + + +@pytest.mark.L0 +@pytest.mark.GPU +@pytest.mark.skipif(not torch.cuda.is_available(), reason="The attention kernels require a GPU.") +def test_opt_in_action_conditioning_train_all_ragged_views_and_action_encoder() -> None: + network = _multiview_network(maskless_attention=True, device=torch.device("cuda"), action_conditioning=True) + views = [torch.randn(1, LATENT_CHANNELS, 3, h, w) for h, w in [(4, 6), (2, 2), (2, 4)]] + data = GenerationDataClean( + batch_size=1, + is_image_batch=False, + x0_tokens_vision=views, + x0_tokens_action=[torch.randn(8, 64)], + action_domain_id=[torch.tensor([2])], + fps_vision=torch.tensor([30.0]), + num_vision_items_per_sample=[3], + ) + data = replicate_multiview_actions( + data, + num_views=3, + temporal_factor=4, + action_dim=64, + ) + packed = pack_multiview_action_conditioning( + [build_sequence_plan_from_mode("forward_dynamics", video_length=9, action_length=8)], + [[3, 4]], + data, + torch.tensor([0.5]), + config=MultiviewActionConditioningConfig(view_code_start=59, view_codes=[[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]), + special_tokens=SPECIAL_TOKENS, + patch_spatial=PATCH_SPATIAL, + temporal_factor=4, + temporal_margin=15_000, + base_fps=24.0, + ) + packed.to_cuda() + output = network(packed) + assert [pred.shape for pred in output["preds_vision"]] == [view.shape for view in views] + torch.stack([pred.square().mean() for pred in output["preds_vision"]]).sum().backward() + attention = network.language_model.seen_attention_mask + assert isinstance(attention, SplitInfo) + assert attention.flex_block_mask is None + assert attention.multiview_maskless is not None + assert attention.multiview_maskless.is_control == (False,) * 6 + assert network.vae2llm.weight.grad is not None and torch.isfinite(network.vae2llm.weight.grad).all() + grads = [parameter.grad for parameter in network.action2llm.parameters() if parameter.grad is not None] + assert grads and all(torch.isfinite(grad).all() for grad in grads) + assert sum(grad.abs().sum() for grad in grads) > 0 diff --git a/cosmos_framework/model/generator/mot/maskless_attention.py b/cosmos_framework/model/generator/mot/maskless_attention.py index 9ebc42b59..b40343803 100644 --- a/cosmos_framework/model/generator/mot/maskless_attention.py +++ b/cosmos_framework/model/generator/mot/maskless_attention.py @@ -379,8 +379,11 @@ def build_replay_maskless_plan( & cross_enabled[:, None] & cross_time ) # [R_Q,R_KV] - # Preserve production maskless weighting: a same-view key present in - # both partitions contributes twice to the merged softmax denominator. + # Legacy maskless weighting counts a same-view key in both partitions. + # Deduplicated recipes retain it only in the same-view partition so the + # merged softmax denominator counts every permitted key once. + if policy.deduplicate_cross_view: + cross = cross & ~view_match # [R_Q,R_KV] partitions.append(("cross_instant" if window is None else "cross_window", cross)) partitions.append(("caption", allowed & caption)) # [R_Q,R_KV] passes = tuple( diff --git a/cosmos_framework/model/generator/mot/multiview_maskless_attention.py b/cosmos_framework/model/generator/mot/multiview_maskless_attention.py index 89e3c7fd2..8f6ac004d 100644 --- a/cosmos_framework/model/generator/mot/multiview_maskless_attention.py +++ b/cosmos_framework/model/generator/mot/multiview_maskless_attention.py @@ -636,6 +636,7 @@ def build_multiview_maskless_plan( *, device: torch.device, seconds_per_frame: Sequence[float] | None = None, + frame_offsets: Sequence[int] | None = None, items_per_sample: Sequence[int] | None = None, is_control: Sequence[bool] | None = None, control_attends_sensor: bool | None = None, @@ -690,7 +691,7 @@ def build_multiview_maskless_plan( Without a temporal window, the instant of a token is its frame's **midpoint**, quantised by the anchor item's frame period:: - instant = floor((frame_id + 0.5) * seconds_per_frame / anchor_seconds_per_frame) + instant = floor((frame_id + frame_offset + 0.5) * seconds_per_frame / anchor_seconds_per_frame) A latent frame is a span rather than a point -- temporal compression folds several pixel frames into one -- so the midpoint is what says which anchor frame a frame mostly happened @@ -717,6 +718,9 @@ def build_multiview_maskless_plan( seconds_per_frame: real time between two latent frames of each item. ``None`` gives every item 1.0, which makes an instant the frame index -- correct whenever the batch carries one rate. + frame_offsets: first frame of each item on its own clock; ``None`` keeps every + item starting at zero. Frame-rate robot actions grouped by the VAE factor + start at latent frame 1, after the initial observation, without null padding. items_per_sample: how many items each sample owns. ``None`` means one each. is_control: whether each item conditions the target that follows it. ``None`` means none of them do. @@ -786,6 +790,7 @@ def build_multiview_maskless_plan( if not num_items: raise ValueError("build_multiview_maskless_plan needs at least one item.") rates = [1.0] * num_items if seconds_per_frame is None else list(seconds_per_frame) + offsets = [0] * num_items if frame_offsets is None else list(frame_offsets) control = [False] * num_items if is_control is None else list(is_control) axes = [0] * num_items if view_axis is None else list(view_axis) # ``None`` reproduces what the axis used to decide on its own: the cameras' axis is the one @@ -797,6 +802,7 @@ def build_multiview_maskless_plan( ) for name, values in ( ("seconds_per_frame", rates), + ("frame_offsets", offsets), ("is_control", control), ("view_axis", axes), ("caption_access", accesses), @@ -805,6 +811,8 @@ def build_multiview_maskless_plan( raise ValueError(f"num_views describes {num_items} items but {name} describes {len(values)}.") if any(rate <= 0 or not math.isfinite(rate) for rate in rates): raise ValueError(f"seconds_per_frame must be positive and finite, got {rates}.") + if any(not isinstance(offset, int) or offset < 0 for offset in offsets): + raise ValueError(f"frame_offsets must be nonnegative integers, got {offsets}.") counts = [1] * num_items if items_per_sample is None else list(items_per_sample) if sum(counts) != num_items: raise ValueError(f"items_per_sample sums to {sum(counts)} but the batch holds {num_items} items.") @@ -881,6 +889,8 @@ def build_multiview_maskless_plan( control_to_sensor_window, ) uses_frame_windows = any(window is not None for window in frame_windows) + if uses_frame_windows and any(offsets): + raise ValueError("Nonzero item frame offsets are not supported by same-view neighborhood windows.") view_group: dict[tuple[int, int, int], int] = {} group_sample: dict[int, int] = {} @@ -969,7 +979,9 @@ def build_multiview_maskless_plan( # Match Flex's float32 frame-start clock exactly for windows; # leave the original midpoint quantisation intact otherwise. frame_times = ( - (torch.arange(frames, dtype=torch.float32, device="cpu") * rates[item]).tolist() # [F] -> list + ( + (torch.arange(frames, dtype=torch.float32, device="cpu") + offsets[item]) * rates[item] + ).tolist() # [F] -> list if decomposed_temporal_window_seconds is not None else None ) @@ -979,7 +991,7 @@ def build_multiview_maskless_plan( instant = ( frame_times[frame] if frame_times is not None - else math.floor((frame + 0.5) * (rates[item] / anchor_rate) + 1e-6) + else math.floor((frame + offsets[item] + 0.5) * (rates[item] / anchor_rate) + 1e-6) ) start = position + (view * frames + frame) * spatial instant_cells.setdefault((sample, instant), {}).setdefault(group, []).append( @@ -987,7 +999,9 @@ def build_multiview_maskless_plan( ) else: frame_ids = torch.arange(frames, device=device, dtype=torch.float64) # [F] - instants = torch.floor((frame_ids + 0.5) * (rates[item] / anchor_rate) + 1e-6).long() # [F] + instants = torch.floor( + (frame_ids + offsets[item] + 0.5) * (rates[item] / anchor_rate) + 1e-6 + ).long() # [F] instant_ids.append(instants.repeat_interleave(spatial).repeat(views) + (sample << 32)) # [V*F*S] sensor_positions.append( torch.arange(position, position + item_lens[item], device=device) diff --git a/cosmos_framework/model/generator/omni_mot_causal_model.py b/cosmos_framework/model/generator/omni_mot_causal_model.py index 2df2d7349..a9a2bd377 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model.py @@ -43,7 +43,6 @@ TeacherForcingReplayPolicyConfig, ) from cosmos_framework.model.generator.attention_io_layout import AttentionIOLayout -from cosmos_framework.model.generator.joint_transfer_ar import sample_joint_transfer_ar from cosmos_framework.model.generator.mot.causal_attention import dispatch_attention_with_memory from cosmos_framework.model.generator.mot.causal_cosmos3_vfm_network import ( InteractiveCosmos3VFMNetwork, @@ -63,7 +62,6 @@ from cosmos_framework.model.generator.mot.post_saturation.static_compile import ( validate_ar_static_und_cache_lengths, ) -from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend from cosmos_framework.model.generator.teacher_forcing import ( make_teacher_forcing_clean_pack, ) @@ -76,6 +74,7 @@ TeacherForcingMemoryState, ) from cosmos_framework.model.generator.utils.kv_storage_backend import validate_kv_cache_dtype +from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend, sample_joint_transfer_ar from cosmos_framework.model.generator.utils.nvfp4 import resolve_legacy_nvfp4_mode from cosmos_framework.data.generator.sequence_packing.autoregressive import ( pack_input_sequence_autoregressive, @@ -2941,7 +2940,9 @@ def _iter_samples_multiview_transfer_autoregressive( seed = _broadcast_seed([seed], self.parallel_dims.cfgp_mesh.get_group(), self.parallel_dims.cfgp_rank)[0] cfg_active = guidance != 1.0 or cfgp_enabled if getattr(self.config, "rolling_kv_cache_chunks", None) is not None: - from cosmos_framework.model.generator.rolling_transfer_ar import iter_rolling_transfer_ar + from cosmos_framework.model.generator.utils.rolling_kv.rolling_transfer_ar import ( + iter_rolling_transfer_ar, + ) if normalize_cfg and cfg_active: raise ValueError("Rolling replay requires normalize_cfg=False when guidance is active.") diff --git a/cosmos_framework/model/generator/omni_mot_causal_model_test.py b/cosmos_framework/model/generator/omni_mot_causal_model_test.py index ddaaaa6c3..66101cc24 100644 --- a/cosmos_framework/model/generator/omni_mot_causal_model_test.py +++ b/cosmos_framework/model/generator/omni_mot_causal_model_test.py @@ -1209,8 +1209,8 @@ def test_multiview_transfer_ar_yields_logical_frames_as_chunks_finish( ) -> None: """Multi-chunk transfer exposes progress and refreshes RGB-aware control K/V in order.""" from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel - from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend num_views = 2 frames_per_view = 5 @@ -1322,7 +1322,7 @@ def generate_chunk(**kwargs: object) -> torch.Tensor: # returns [B,C,V*T_chunk, return_value=[sequence_plan], ), patch( - "cosmos_framework.model.generator.multiview_transfer_ar.build_multiview_transfer_ar_memory_layout", + "cosmos_framework.model.generator.utils.multiview_ar.build_multiview_transfer_ar_memory_layout", return_value=memory_layout, ), ): @@ -1403,7 +1403,7 @@ def generate_chunk(**kwargs: object) -> torch.Tensor: # returns [B,C,V*T_chunk, @pytest.mark.CPU def test_multiview_transfer_ar_pack_sets_metadata_and_aligned_view_positions() -> None: """Synchronized chunks align camera-local mRoPE positions and retain Flex metadata.""" - from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend + from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend model = MagicMock() model.config.diffusion_expert_config.patch_spatial = 1 @@ -1419,7 +1419,7 @@ def test_multiview_transfer_ar_pack_sets_metadata_and_aligned_view_positions() - vision_latent = torch.zeros(1, 4, 4, 2, 2) # [B,C,V*chunk_len,H,W] with patch( - "cosmos_framework.model.generator.multiview_transfer_ar.pack_input_sequence_autoregressive", + "cosmos_framework.model.generator.utils.multiview_ar.pack_input_sequence_autoregressive", return_value=packed_seq, ) as pack_input: result = MultiviewTransferARBackend(model).build_current_pack( @@ -1459,8 +1459,8 @@ def test_multiview_transfer_ar_pack_sets_metadata_and_aligned_view_positions() - def test_multiview_transfer_prefill_is_clean_without_extending_conditioned_prefix(condition_count: int) -> None: """Recomputed RGB history omits diffusion embeddings without changing the rollout prefix.""" from cosmos_framework.data.generator.sequence_packing import ModalityData - from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend from cosmos_framework.model.generator.omni_mot_causal_model import _multiview_conditioned_prefix_length + from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend target_mask = (torch.arange(5) < condition_count).repeat(2).reshape(10, 1, 1) # [V*T,1,1] control_mask = torch.ones_like(target_mask) # [V*T,1,1] @@ -1515,7 +1515,7 @@ def test_multiview_transfer_prefill_is_clean_without_extending_conditioned_prefi @pytest.mark.CPU def test_multiview_transfer_clean_pack_uses_teacher_forcing_condition_semantics() -> None: """Clean target history omits timestep embeddings while retaining its explicit Flex role.""" - from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend + from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend model = MagicMock() model.config.diffusion_expert_config.patch_spatial = 1 @@ -1537,7 +1537,7 @@ def test_multiview_transfer_clean_pack_uses_teacher_forcing_condition_semantics( memory_layout = MagicMock() with patch( - "cosmos_framework.model.generator.multiview_transfer_ar.pack_input_sequence_autoregressive", + "cosmos_framework.model.generator.utils.multiview_ar.pack_input_sequence_autoregressive", return_value=packed_seq, ): result = MultiviewTransferARBackend(model).build_current_pack( @@ -1568,7 +1568,7 @@ def test_multiview_transfer_clean_pack_uses_teacher_forcing_condition_semantics( @pytest.mark.CPU def test_multiview_transfer_ar_memory_merge_updates_only_selected_slots() -> None: """A clean recomputation refreshes control slots without erasing RGB history.""" - from cosmos_framework.model.generator.multiview_transfer_ar import MultiviewTransferARBackend + from cosmos_framework.model.generator.utils.multiview_ar import MultiviewTransferARBackend destination_k = torch.full((1, 5, 1, 1), -1.0) # [1,S_memory,H_kv,D] destination_v = torch.full((1, 5, 1, 1), -2.0) # [1,S_memory,H_kv,D] @@ -1603,7 +1603,7 @@ def test_multiview_transfer_backend_assigns_prefill_cache_by_cfg_branch( expected_pack_name: str | None, ) -> None: """Sequential CFG keeps two caches while CFGP stores only the rank-local branch.""" - from cosmos_framework.model.generator.multiview_transfer_ar import ( + from cosmos_framework.model.generator.utils.multiview_ar import ( MultiviewTransferARBackend, MultiviewTransferARSession, ) @@ -1679,7 +1679,7 @@ def test_multiview_transfer_backend_resolves_cfgp_rank_local_memory( remote_branch: Literal["conditional", "unconditional"], ) -> None: """CFGP exposes the single allocated cache through the logical branch owned by each rank.""" - from cosmos_framework.model.generator.multiview_transfer_ar import ( + from cosmos_framework.model.generator.utils.multiview_ar import ( MultiviewTransferARBackend, MultiviewTransferARSession, ) @@ -1720,11 +1720,11 @@ def test_multiview_transfer_backend_resolves_cfgp_rank_local_memory( @pytest.mark.CPU def test_multiview_transfer_backend_prefills_and_commits_fixed_cache_slots() -> None: """The real backend preserves prefill slots while committing camera-major clean history.""" - from cosmos_framework.model.generator.multiview_transfer_ar import ( + from cosmos_framework.model.generator.utils.kv_cache import MultiviewARMemoryState + from cosmos_framework.model.generator.utils.multiview_ar import ( MultiviewTransferARBackend, MultiviewTransferARSession, ) - from cosmos_framework.model.generator.utils.kv_cache import MultiviewARMemoryState denoise_call_count = 0 diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 5f7b7f6c1..81c71ec45 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -31,9 +31,14 @@ ACTION_SLOT_SAMPLE_COUNT_KEY, ACTION_SLOT_SAMPLE_LOSS_KEY, ActionSlotLossStats, + FlowMatchingLossItems, compute_flow_matching_loss, ) from cosmos_framework.model.generator.algorithm.loss.load_balancing import compute_load_balancing_loss +from cosmos_framework.model.generator.algorithm.loss.modality_reduction import ( + build_modality_loss_items, + reduce_global_modality_means, +) from cosmos_framework.configs.base.defaults.joint_attention import JointAttnImplementation from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig from cosmos_framework.configs.base.defaults.parallelism import PRECISION_TO_TORCH_DTYPE @@ -70,6 +75,7 @@ get_vae_pixel_shapes, normalize_uint8_item, ) +from cosmos_framework.model.generator.utils.batch_normalization import normalize_vision_batch_inplace from cosmos_framework.model.generator.utils.data_and_condition import ( GenerationDataClean, GenerationDataNoised, @@ -103,7 +109,12 @@ pack_input_sequence, ) from cosmos_framework.data.generator.sequence_packing.modality import add_special_tokens -from cosmos_framework.data.generator.sequence_packing.packers import is_item_generated, uses_single_timestep +from cosmos_framework.data.generator.sequence_packing.packers import ( + is_item_generated, + pack_multiview_action_conditioning, + replicate_multiview_actions, + uses_single_timestep, +) from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface from cosmos_framework.model.generator.upsampler.prompts import build_messages, clean_response from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution, read_positive_int_metadata @@ -477,6 +488,29 @@ def build_net( lora_enabled: Override for ``config.lora_enabled``. """ # Build model network and parallelize it. + multiview_action_conditioning = getattr(self.config, "multiview_action_conditioning", None) + if multiview_action_conditioning is not None: + if ( + not self.config.action_gen + or self.config.video_temporal_causal + or self.config.causal_training_strategy != "none" + or self.config.joint_attn_implementation != "multiview" + or self.config.multiview_attention.backend != "maskless" + ): + raise ValueError("Multiview action conditioning requires bidirectional RGB/action maskless attention.") + if self.config.diffusion_expert_config.num_view_embeddings != 0: + raise ValueError("num_view_embeddings must be 0 with fixed action view codes (no learned embeddings).") + codes = torch.tensor([list(code) for code in multiview_action_conditioning.view_codes]) + if ( + codes.ndim != 2 + or codes.shape[0] < 2 + or codes.shape[1] == 0 + or multiview_action_conditioning.view_code_start < 0 + or multiview_action_conditioning.view_code_start + codes.shape[1] > self.config.max_action_dim + or not torch.isfinite(codes).all() + or codes[0].any() + ): + raise ValueError("Fixed view codes must fit padded action slots and keep the reference view zero.") lora_enabled = self.config.lora_enabled if lora_enabled is None else lora_enabled with torch.device("meta"): assert self.vlm_config.model_instance is not None, "Model instance should be specified" @@ -528,6 +562,7 @@ def build_net( sound_gen=self.config.sound_gen, joint_attn_implementation=self.config.joint_attn_implementation, multiview_attention_config=self.config.multiview_attention, + multiview_action_conditioning=multiview_action_conditioning is not None, timestep_scale=1.0 / float(num_train_timesteps) * self.config.diffusion_expert_config.timestep_range, timestep_range=self.config.diffusion_expert_config.timestep_range, action_dim=self.config.max_action_dim, @@ -958,6 +993,22 @@ def _pack_input_sequence( plus three optional flags. """ assert self.tokenizer_vision_gen is not None + multiview_action_conditioning = getattr(self.config, "multiview_action_conditioning", None) + if multiview_action_conditioning is not None and gen_data_clean.num_vision_items_per_sample is not None: + if skip_text_tokens or include_end_of_generation_token or initial_mrope_temporal_offset: + raise ValueError("Multiview action conditioning keeps sample-level text and no generation markers.") + return pack_multiview_action_conditioning( + sequence_plans, + input_text_indexes, + gen_data_clean, + input_timesteps, + config=multiview_action_conditioning, + special_tokens=self.llm_special_tokens, + patch_spatial=self.config.diffusion_expert_config.patch_spatial, + temporal_factor=self.tokenizer_vision_gen.temporal_compression_factor, + temporal_margin=self.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin, + base_fps=float(self.config.diffusion_expert_config.base_fps), + ) packed = pack_input_sequence( sequence_plans=sequence_plans, input_text_indexes=input_text_indexes, @@ -1187,6 +1238,9 @@ def _prepare_training_data( leave it off deliberately: balancing runs collectives over the whole ``lb`` group, and those paths are not guaranteed to run on every rank of it. """ + normalize_vision_batch_inplace( + data_batch, video_key=self.input_video_key, tensor_kwargs_fp32=getattr(self, "tensor_kwargs_fp32", {}) + ) input_text_indexes = self._load_and_tokenize_text_data(data_batch, iteration) sequence_plans = build_sequence_plans_from_data_batch( data_batch=data_batch, @@ -1206,15 +1260,36 @@ def _prepare_training_data( gen_data_clean, memory_info = self.memory_init_training(gen_data_clean, data_batch, input_text_indexes) # image_size[i] may be (1, 4) from IterativeJointDataLoader or (4,) from custom_collate_fn. - if "image_size" in data_batch: - data_resolutions: list[str] | str | None = [] + if ( + getattr(self.config, "multiview_action_conditioning", None) is not None + and gen_data_clean.num_vision_items_per_sample is not None + ): + # There is no last-item transfer target. Use the configured reference + # resolution once per synchronized sample, not the last smaller view. + data_resolutions: list[str] | str | None = [self.config.resolution] * gen_data_clean.batch_size + elif "image_size" in data_batch: + data_resolutions = [] # Multi-item samples (transfer, SR) carry one image_size per vision item; the noise # schedule must follow the generated (last) item of each sample. - target_image_sizes = select_target_image_sizes( - data_batch["image_size"], gen_data_clean.num_vision_items_per_sample, gen_data_clean.batch_size - ) + vision_counts = gen_data_clean.num_vision_items_per_sample + target_image_sizes: list[torch.Tensor | None] + if vision_counts is not None and any(count == 0 for count in vision_counts): + # Mixed batches can contain samples without vision. Select camera-bearing + # targets first, then restore the slots that use the configured resolution. + camera_counts = [count for count in vision_counts if count > 0] + camera_target_sizes = iter( + select_target_image_sizes(data_batch["image_size"], camera_counts, len(camera_counts)) + ) + target_image_sizes = [next(camera_target_sizes) if count else None for count in vision_counts] + else: + target_image_sizes = select_target_image_sizes( + data_batch["image_size"], vision_counts, gen_data_clean.batch_size + ) for i in range(gen_data_clean.batch_size): img_size = target_image_sizes[i] + if img_size is None: + data_resolutions.append(self.config.resolution) + continue if img_size.dim() == 2: img_size = img_size[0] target_h = int(img_size[0].item()) @@ -1386,6 +1461,19 @@ def training_step( # Sample a random noise level (sigma) and corresponding interpolation coefficient ("timesteps" in RF) # Apply shift per sample based on each sample's resolution num_vision_latent_frames = [x.shape[2] for x in gen_data_clean.x0_tokens_vision] + sample_is_image = getattr(gen_data_clean, "sample_is_image", None) + if sample_is_image is not None: + # Schedules belong to logical samples; multi-item controls/targets use + # the same sample schedule, expanded to items later in this method. + counts = gen_data_clean.num_vision_items_per_sample or [1] * gen_data_clean.batch_size + offset = 0 + sample_tokens: list[int] = [] + sample_frames: list[int] = [] + for count in counts: + sample_tokens.append(sum(num_tokens_per_sample[offset : offset + count])) + sample_frames.append(max(num_vision_latent_frames[offset : offset + count], default=1)) + offset += count + num_tokens_per_sample, num_vision_latent_frames = sample_tokens, sample_frames timesteps_vision, sigmas_vision = self._get_train_noise_level_vision( batch_size=gen_data_clean.batch_size, is_image_batch=gen_data_clean.is_image_batch, @@ -1393,6 +1481,7 @@ def training_step( num_vision_latent_frames=num_vision_latent_frames, num_tokens=num_tokens_per_sample, iteration=iteration, + **({"sample_is_image": sample_is_image} if sample_is_image is not None else {}), sequence_plans=sequence_plans, ) # [B, T_vis] each @@ -1418,6 +1507,12 @@ def training_step( else: timesteps_action, sigmas_action = (None, None) + if timesteps_action is None and sample_is_image is not None and action_sample_indices: + # The shared schedule remains indexed by logical sample before the + # vision-item expansion. Sparse actions need their own dense view. + timesteps_action = timesteps_vision[action_sample_indices] # [N_action,T_vis] + sigmas_action = sigmas_vision[action_sample_indices] # [N_action,T_vis] + # Optional independent sound schedule: sample a scalar sound sigma per batch # slot, then reindex to the dense audio-bearing subset. sound_sample_indices = [i for i, plan in enumerate(sequence_plans) if getattr(plan, "has_sound", False)] @@ -1477,7 +1572,7 @@ def training_step( # Under independent_action_schedule, overwrite the vision-based action timestep the # packer injected with the action timestep, so the denoiser's action timestep embedding # matches the sigma used to noise action tokens. - if timesteps_action is not None and packed_sequence.action is not None: + if rf_cfg.independent_action_schedule and timesteps_action is not None and packed_sequence.action is not None: action_has_noisy_tokens = any(nfi.numel() > 0 for nfi in packed_sequence.action.noisy_frame_indexes) if action_has_noisy_tokens: sample_ts = timesteps_action.squeeze(1).cpu() # [n_action] @@ -1626,6 +1721,17 @@ def training_step( } if sigmas_action is not None: output_batch["sigma_action"] = sigmas_action # [n_action, 1] — dense over action-bearing samples + if sample_is_image is not None: + # Callback identity follows the CP-rotated payload, not this rank's + # raw loader batch, and is aligned with per-vision-item losses. + output_batch["vision_is_image"] = packed_sequence.vision_is_image + # Batch sizes count logical camera-bearing samples: multiple views or + # editing controls count once, and vision-less samples count as neither. + vision_sample_flags = [ + is_image for is_image, plan in zip(sample_is_image, sequence_plans, strict=True) if plan.has_vision + ] + output_batch["image_batch_size"] = sum(vision_sample_flags) + output_batch["video_batch_size"] = len(vision_sample_flags) - output_batch["image_batch_size"] if getattr(rf_cfg, "independent_sound_schedule", False) and sigmas_sound is not None: output_batch["sigma_sound"] = sigmas_sound # [n_sound, 1] — dense over sound-bearing samples @@ -1724,41 +1830,6 @@ def _loss_averaging_group(self) -> tuple[torch.distributed.ProcessGroup | None, return None, torch.distributed.get_world_size() return None, 1 - def _sample_level_loss_scale( - self, - is_image_batch: bool, - num_samples: int, - device: torch.device, - ) -> torch.Tensor: - """Multiplier that converts rank-level to sample-level loss averaging. - - A packed batch is homogeneous — either all image or all video (video may carry - accompanying action/audio) — with ``num_samples`` samples on this rank. Each - iteration the per-modality sample counts are all-reduced over the data-parallel - group to get the global totals ``N_image`` / ``N_video``, and this rank's batch is - scaled by ``group_size * num_samples / N_modality``. - - Applied to the whole ``total_loss`` (video + action + audio for video batches), - this scales every term by the same factor, so the video/action/audio balance set - by ``loss_scale`` / ``action_loss_weight`` / ``sound_loss_scale`` is preserved, - while the effective objective becomes a per-sample mean. The base per-modality - weight already baked into ``total_loss`` (``loss_scale`` for video, ``image_loss_scale`` - for image) cancels out of this multiplier, so it does not appear here. The - ``group_size`` factor cancels the framework's ``1/group_size`` gradient averaging; - for balanced batches (``num_samples ≈ N/group_size``) the scale is ≈ 1, so logged - loss magnitudes stay comparable to rank-level averaging. - """ - dp_group, group_size = self._loss_averaging_group() - - # counts = [num_local_image_samples, num_local_video_samples]; exactly one is non-zero. - counts = torch.zeros(2, dtype=torch.float64, device=device) - counts[0 if is_image_batch else 1] = float(num_samples) - if group_size > 1: - torch.distributed.all_reduce(counts, op=torch.distributed.ReduceOp.SUM, group=dp_group) - - global_num_samples = (counts[0] if is_image_batch else counts[1]).clamp(min=1.0) - return (group_size * num_samples) / global_num_samples - def _compute_losses( self, out_net: dict, @@ -1792,7 +1863,28 @@ def _compute_losses( normalize_by_active = rf_cfg.normalize_loss_by_active exclude_fully_conditioned_items = rf_cfg.exclude_fully_conditioned_items if exclude_fully_conditioned_items is None: - exclude_fully_conditioned_items = self.config.causal_training_strategy == "teacher_forcing" + exclude_fully_conditioned_items = ( + getattr(self.config, "causal_training_strategy", None) == "teacher_forcing" + ) + # CP gather gradient correction is independent of loss reduction; + # disabling it preserves legacy gradient scaling for ablations. + use_global_sample_mean = rf_cfg.loss_reduction_mode == "global_sample_mean" + vision_flags = getattr(data_batch_packed, "vision_is_image", None) + use_vision_item_weights = not use_global_sample_mean and bool(vision_flags) + enabled_streams = { + "vision": self.config.vision_gen, + "action": self.config.action_gen, + "sound": self.config.sound_gen, + "lidar": self.config.vision_gen and self.config.lidar_state_ch is not None, + "radar": self.config.vision_gen and self.config.radar_state_ch is not None, + } + item_losses = ( + {name: FlowMatchingLossItems() for name, enabled in enabled_streams.items() if enabled} + if use_global_sample_mean + else {} + ) + if use_vision_item_weights: + item_losses["vision"] = FlowMatchingLossItems() if self.config.vision_gen: # Only a batch that generates no camera stream, as the LiDAR-only recipe does, may # arrive with vision unpacked; for anything else that would silently sink the vision @@ -1801,6 +1893,10 @@ def _compute_losses( "Vision packed data required when the batch carries camera items" ) rectified_flow_vision = self.rectified_flow_image if is_image_batch else self.rectified_flow_video + if vision_flags: + rectified_flow_vision = [ + self.rectified_flow_image if flag else self.rectified_flow_video for flag in vision_flags + ] # With no camera stream there are no noisy vision tokens, so this returns its dummy # loss over the network's zero-weighted vision predictions. That is what keeps @@ -1819,11 +1915,33 @@ def _compute_losses( tensor_kwargs_fp32=self.tensor_kwargs_fp32, normalize_by_active=normalize_by_active, exclude_fully_conditioned_items=exclude_fully_conditioned_items, + item_losses=item_losses.get("vision"), ) - loss_scale = ( - rf_cfg.image_loss_scale if is_image_batch and rf_cfg.image_loss_scale is not None else rf_cfg.loss_scale - ) - total_loss += fm_loss_vision * loss_scale # [] + if use_vision_item_weights: + # Keep legacy rank-level averaging across all local vision items, + # applying each item's image/video scale before taking the mean. + # Vision-less samples must not change the weight of neighboring images. + vision_losses = item_losses["vision"].weighted_losses # [N_vision] or None + valid_vision = item_losses["vision"].valid # [N_vision] or None + assert vision_losses is not None and valid_vision is not None + vision_scales = vision_losses.new_tensor( + [ + rf_cfg.image_loss_scale if flag and rf_cfg.image_loss_scale is not None else rf_cfg.loss_scale + for flag in vision_flags + ] + ) # [N_vision] + scaled_vision_losses = vision_losses * vision_scales # [N_vision] + if exclude_fully_conditioned_items: + total_loss += (scaled_vision_losses * valid_vision).sum() / valid_vision.sum().clamp(min=1) # [] + else: + total_loss += scaled_vision_losses.mean() # [] + else: + loss_scale = ( + rf_cfg.image_loss_scale + if is_image_batch and rf_cfg.image_loss_scale is not None + else rf_cfg.loss_scale + ) + total_loss += fm_loss_vision * loss_scale # [] losses_dict["flow_matching_loss_vision"] = fm_loss_vision # [] losses_dict["flow_matching_loss_vision_per_instance"] = fm_loss_vision_per_instance # [N] else: @@ -1849,6 +1967,7 @@ def _compute_losses( tensor_kwargs_fp32=self.tensor_kwargs_fp32, normalize_by_active=normalize_by_active, exclude_fully_conditioned_items=exclude_fully_conditioned_items, + item_losses=item_losses.get("lidar"), ) lidar_loss_scale = rf_cfg.lidar_loss_scale if rf_cfg.lidar_loss_scale is not None else rf_cfg.loss_scale total_loss += fm_loss_lidar * lidar_loss_scale # [] @@ -1890,6 +2009,7 @@ def _compute_losses( rectified_flow=self.rectified_flow_video, tensor_kwargs_fp32=self.tensor_kwargs_fp32, normalize_by_active=normalize_by_active, + item_losses=item_losses.get("radar"), ) radar_loss_scale = rf_cfg.radar_loss_scale if rf_cfg.radar_loss_scale is not None else rf_cfg.loss_scale total_loss += fm_loss_radar * radar_loss_scale # [] @@ -1934,6 +2054,7 @@ def _compute_losses( action_valid_mask=data_batch_packed.action.action_valid_mask, normalize_by_active=normalize_by_active, action_slot_stats=action_slot_stats, + item_losses=item_losses.get("action"), ) # Yihuai: In case the video loss is too large (1.5) and covers the action loss (0.05), we scale up the action loss to match the video loss to improve action precision. @@ -1969,6 +2090,8 @@ def _compute_losses( rectified_flow=self.rectified_flow_sound, tensor_kwargs_fp32=self.tensor_kwargs_fp32, normalize_by_active=normalize_by_active, + item_losses=item_losses.get("sound"), + time_axis=1 if use_global_sample_mean else 0, ) loss_scale = rf_cfg.sound_loss_scale if rf_cfg.sound_loss_scale is not None else rf_cfg.loss_scale total_loss += fm_loss_sound * loss_scale @@ -1984,19 +2107,38 @@ def _compute_losses( else: losses_dict["flow_matching_loss_sound"] = torch.tensor(0.0, **self.tensor_kwargs_fp32) - # Sample-level (vs rank-level) loss averaging. Scale the whole batch loss - # (vision + accompanying action/audio for video batches) by one factor so every - # image / video sample contributes equally across the data-parallel group. Applied - # here — after all flow-matching terms, before load balancing — so the video/action/ - # audio balance is preserved and the auxiliary load-balancing losses stay unscaled. - if rf_cfg.sample_level_loss_averaging and self.config.vision_gen: - num_samples = len(out_net["preds_vision"]) - sample_level_scale = self._sample_level_loss_scale( + # Sample-level (vs rank-level) loss averaging. Normalize each supervised + # modality by its own global sample count so sparse action/audio samples + # do not inherit the image/video denominator. Apply this after all flow-matching + # terms, before load balancing, preserving configured modality weights and + # leaving auxiliary load-balancing losses unscaled. + if use_global_sample_mean: + # This opt-in policy normalizes contributing logical samples per modality + # over the gradient-averaging group, within this microbatch only. The + # collector exposes weighted training losses, not unweighted logging + # values. Clean/invalid items and absent modalities add no sample count. + # Fixed modality weights are summed without renormalizing; the trainer's + # accumulation averaging and the auxiliary terms below are unchanged. + group, group_size = self._loss_averaging_group() + modality_items = build_modality_loss_items( + item_losses, + out_net, + data_batch_packed.modality_sample_ids, + vision_is_image=vision_flags, is_image_batch=is_image_batch, - num_samples=num_samples, - device=self.tensor_kwargs_fp32["device"], ) - total_loss = total_loss * sample_level_scale.to(dtype=total_loss.dtype) + weights = { + "image": rf_cfg.image_loss_scale if rf_cfg.image_loss_scale is not None else rf_cfg.loss_scale, + "video": rf_cfg.loss_scale, + "action": rf_cfg.action_loss_weight, + "sound": rf_cfg.sound_loss_scale if rf_cfg.sound_loss_scale is not None else rf_cfg.loss_scale, + "lidar": rf_cfg.lidar_loss_scale if rf_cfg.lidar_loss_scale is not None else rf_cfg.loss_scale, + "radar": rf_cfg.radar_loss_scale if rf_cfg.radar_loss_scale is not None else rf_cfg.loss_scale, + } + total_loss, modality_means = reduce_global_modality_means( + modality_items, weights, group=group, gradient_average_size=group_size + ) + losses_dict.update({f"global_sample_loss_{name}": value for name, value in modality_means.items()}) # 2. Load balancing auxiliary losses device_mesh, context_parallel_mesh = self._get_load_balancing_loss_meshes() @@ -2021,7 +2163,11 @@ def _update_train_stats(self, gen_data_clean: GenerationDataClean) -> None: """Update sample counters from a processed ``GenerationDataClean`` payload.""" if not isinstance(self.net, WeightTrainingStat): return - if gen_data_clean.is_image_batch: + if gen_data_clean.sample_is_image is not None: + image_count = sum(gen_data_clean.sample_is_image) + self.net.accum_image_sample_counter += image_count + self.net.accum_video_sample_counter += gen_data_clean.batch_size - image_count + elif gen_data_clean.is_image_batch: self.net.accum_image_sample_counter += gen_data_clean.batch_size else: self.net.accum_video_sample_counter += gen_data_clean.batch_size @@ -2086,6 +2232,8 @@ def _apply_per_view_caption_plan( f"has {len(sequence_plans)} sequence plans." ) for sample_idx, (caption_lengths, plan) in enumerate(zip(caption_lengths_per_sample, sequence_plans)): + if caption_lengths is None: + continue num_captions = len(caption_lengths) if num_captions < 1: raise ValueError(f"Sample {sample_idx} carries no per-view caption.") @@ -2120,6 +2268,7 @@ def _get_train_noise_level_vision( resolutions: list[str] | str | None = None, num_tokens: list[int] | None = None, iteration: int | None = None, + sample_is_image: list[bool] | None = None, sequence_plans: list[SequencePlan] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ @@ -2137,6 +2286,7 @@ def _get_train_noise_level_vision( Can be a single string (applied to all samples) or a list of strings (one per sample). If None, defaults to self.config.resolution (can be used for other modalities). num_tokens: Number of tokens for each sample (before 2x2 merge). Needed for dynamic shift. + sample_is_image: Per-sample image identities for mixed image/video schedules. sequence_plans: Effective per-sample modality metadata after CP input selection. Batches with LiDAR and no RGB use shift_lidar when configured. RGB and joint RGB/LiDAR batches share the vision schedule. None retains the vision schedule. @@ -2145,6 +2295,35 @@ def _get_train_noise_level_vision( (timesteps, sigmas): Both [B,1] for TF/base, or [B,T_max] for diffusion_forcing. """ + if sample_is_image is not None: + if len(sample_is_image) != batch_size: + raise ValueError("Image identities must match the logical batch size.") + if all(sample_is_image) or not any(sample_is_image): + is_image_batch = all(sample_is_image) + else: + # Homogeneous inputs retain the original sampling call and RNG order. + # A mixed input draws each modality's own schedule and restores sample order. + schedules: list[tuple[list[int], torch.Tensor, torch.Tensor]] = [] + for image_flag in (True, False): + indexes = [i for i, flag in enumerate(sample_is_image) if flag == image_flag] + times, noise = self._get_train_noise_level_vision( + len(indexes), + image_flag, + [num_vision_latent_frames[i] for i in indexes], + resolutions=[resolutions[i] for i in indexes] if isinstance(resolutions, list) else resolutions, + num_tokens=[num_tokens[i] for i in indexes] if num_tokens is not None else None, + iteration=iteration, + sequence_plans=[sequence_plans[i] for i in indexes] if sequence_plans is not None else None, + ) # [N,T_modality] each + schedules.append((indexes, times, noise)) + width = max(times.shape[1] for _, times, _ in schedules) + timesteps = schedules[0][1].new_zeros((batch_size, width)) # [B,T_max] + sigmas = torch.zeros_like(timesteps) # [B,T_max] + for indexes, times, noise in schedules: + timesteps[indexes, : times.shape[1]] = times # [N,T_modality] + sigmas[indexes, : noise.shape[1]] = noise # [N,T_modality] + return timesteps, sigmas + rectified_flow = self.rectified_flow_image if is_image_batch else self.rectified_flow_video assert not self.config.rectified_flow_training_config.use_discrete_rf, ( @@ -2394,6 +2573,8 @@ def _add_noise_to_input( sigmas_vision = [ sigmas[i].view(-1, 1, 1)[: x0_vision[i].shape[2]] * noisy_mask_vision[i] for i in range(num_vision_items) ] + # Image/video interpolation has the same linear equation; modality-specific + # distributions and weights are selected in scheduling and loss computation. rectified_flow_vision = ( self.rectified_flow_image if gen_data_clean.is_image_batch else self.rectified_flow_video ) @@ -2731,6 +2912,9 @@ def _prepare_inference_data( callers must thread it to _get_velocity so both sides split the flat state identically. """ + normalize_vision_batch_inplace( + data_batch, video_key=self.input_video_key, tensor_kwargs_fp32=getattr(self, "tensor_kwargs_fp32", {}) + ) # 1. Build sequence plans (same as training) sequence_plans = build_sequence_plans_from_data_batch( data_batch=data_batch, @@ -2803,6 +2987,10 @@ def _prepare_inference_data( seed_dict["vision"].extend([seed[sample_idx]] * num_vision_items) seed_dict["action"].append(seed[sample_idx]) seed_dict["sound"].append(seed[sample_idx]) + if gen_data_clean.action_sample_ids is not None: + seed_dict["action"] = [seed[index] for index in gen_data_clean.action_sample_ids] + if gen_data_clean.sound_sample_ids is not None: + seed_dict["sound"] = [seed[index] for index in gen_data_clean.sound_sample_ids] # Generate noise and apply conditioning per vision item (supports variable shapes). # Noise and the conditioning blend are kept in fp32 so the sampler accumulates @@ -3078,7 +3266,7 @@ def _release_inference_raw_vision( ): raise TypeError(f"{INFERENCE_RAW_VISION_RETAINED_ITEMS_KEY} must be a list of CPU vision tensors.") - media_key = self.input_image_key if gen_data_clean.is_image_batch else self.input_video_key + media_key = self.input_image_key if self.input_image_key in data_batch else self.input_video_key raw_items = data_batch.get(media_key) if not isinstance(raw_items, list): raise TypeError(f"Expected data_batch[{media_key!r}] to be a list before releasing raw vision tensors.") @@ -3217,6 +3405,8 @@ def _dup(value: Any) -> Any: kwargs[field.name] = val * 2 elif field.name == "is_image_batch": kwargs[field.name] = val + elif field.name in {"action_sample_ids", "sound_sample_ids"} and val is not None: + kwargs[field.name] = val + [index + gd.batch_size for index in val] else: kwargs[field.name] = _dup(val) return GenerationDataClean(**kwargs) @@ -3386,6 +3576,10 @@ def _get_velocity( # sequence. Without this, multi_control_two_way_attention never runs # and all controls are blended equally (weights ignored). control_weights=gen_data_clean.control_weights, + sample_is_image=gen_data_clean.sample_is_image, + action_sample_ids=gen_data_clean.action_sample_ids, + sound_sample_ids=gen_data_clean.sound_sample_ids, + action_family=gen_data_clean.action_family, ) if packed_sequence_template is None: @@ -4407,7 +4601,9 @@ def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: if n_vis == 1 or is_item_generated( sequence_plans[i].condition_frame_indexes_vision, item_idx=j, - num_items=n_vis, + num_items=( + 1 if getattr(self.config, "multiview_action_conditioning", None) is not None else n_vis + ), latent_t=int(vision_shape[2]), ): result_vision.append(latents[i][offset : offset + vision_dim].reshape(vision_shape)) @@ -4632,19 +4828,29 @@ def _slice_gen_data_clean(self, gen_data_clean: GenerationDataClean, start: int, subset_num_radar_items = None if num_radar_items is None else num_radar_items[start:limit] fps_radar = gen_data_clean.fps_radar[start:limit] if gen_data_clean.fps_radar is not None else None + def dense_slice(owners: list[int] | None) -> slice: + if owners is None: + return slice(start, limit) + return slice(sum(owner < start for owner in owners), sum(owner < limit for owner in owners)) + + action_slice = dense_slice(gen_data_clean.action_sample_ids) + sound_slice = dense_slice(gen_data_clean.sound_sample_ids) + if has_action: subset_raw_action = ( - gen_data_clean.raw_state_action[start:limit] if gen_data_clean.raw_state_action else None + gen_data_clean.raw_state_action[action_slice] if gen_data_clean.raw_state_action else None + ) + x0_tokens_action = gen_data_clean.x0_tokens_action[action_slice] + fps_action = gen_data_clean.fps_action[action_slice] if gen_data_clean.fps_action is not None else None + action_domain_id = ( + gen_data_clean.action_domain_id[action_slice] if gen_data_clean.action_domain_id else None ) - x0_tokens_action = gen_data_clean.x0_tokens_action[start:limit] - fps_action = gen_data_clean.fps_action[start:limit] if gen_data_clean.fps_action is not None else None - action_domain_id = gen_data_clean.action_domain_id[start:limit] if gen_data_clean.action_domain_id else None - raw_action_dim = gen_data_clean.raw_action_dim[start:limit] if gen_data_clean.raw_action_dim else None + raw_action_dim = gen_data_clean.raw_action_dim[action_slice] if gen_data_clean.raw_action_dim else None action_valid_mask = ( - gen_data_clean.action_valid_mask[start:limit] if gen_data_clean.action_valid_mask else None + gen_data_clean.action_valid_mask[action_slice] if gen_data_clean.action_valid_mask else None ) num_views_per_action_item = ( - gen_data_clean.num_views_per_action_item[start:limit] + gen_data_clean.num_views_per_action_item[action_slice] if gen_data_clean.num_views_per_action_item else None ) @@ -4658,17 +4864,20 @@ def _slice_gen_data_clean(self, gen_data_clean: GenerationDataClean, start: int, num_views_per_action_item = None if has_sound: - subset_raw_sound = gen_data_clean.raw_state_sound[start:limit] if gen_data_clean.raw_state_sound else None - x0_tokens_sound = gen_data_clean.x0_tokens_sound[start:limit] - fps_sound = gen_data_clean.fps_sound[start:limit] if gen_data_clean.fps_sound is not None else None + subset_raw_sound = gen_data_clean.raw_state_sound[sound_slice] if gen_data_clean.raw_state_sound else None + x0_tokens_sound = gen_data_clean.x0_tokens_sound[sound_slice] + fps_sound = gen_data_clean.fps_sound[sound_slice] if gen_data_clean.fps_sound is not None else None else: subset_raw_sound = None x0_tokens_sound = None fps_sound = None + image_flags = ( + gen_data_clean.sample_is_image[start:limit] if gen_data_clean.sample_is_image is not None else None + ) return GenerationDataClean( batch_size=limit - start, - is_image_batch=gen_data_clean.is_image_batch, + is_image_batch=all(image_flags) if image_flags else gen_data_clean.is_image_batch, raw_state_vision=subset_raw_vision, raw_state_action=subset_raw_action, raw_state_sound=subset_raw_sound, @@ -4694,6 +4903,19 @@ def _slice_gen_data_clean(self, gen_data_clean: GenerationDataClean, start: int, x0_tokens_lidar=subset_x0_lidar, fps_lidar=fps_lidar, num_lidar_items_per_sample=subset_num_lidar_items, + sample_is_image=image_flags, + action_sample_ids=( + [owner - start for owner in gen_data_clean.action_sample_ids[action_slice]] + if gen_data_clean.action_sample_ids is not None + else None + ), + sound_sample_ids=( + [owner - start for owner in gen_data_clean.sound_sample_ids[sound_slice]] + if gen_data_clean.sound_sample_ids is not None + else None + ), + action_family=gen_data_clean.action_family[action_slice] if gen_data_clean.action_family else None, + control_weights=gen_data_clean.control_weights[start:limit] if gen_data_clean.control_weights else None, raw_state_radar=subset_raw_radar, x0_tokens_radar=subset_x0_radar, fps_radar=fps_radar, @@ -4786,8 +5008,8 @@ def _validate_and_get_num_views( num_views_per_vision_item: list[int] = [] frames_per_vision_item: list[int] = [] for sample_idx, num_vision_items in enumerate(vision_item_counts): - if num_vision_items <= 0: - raise ValueError(f"Vision item counts must be positive, got {vision_item_counts}.") + if num_vision_items < 0: + raise ValueError(f"Vision item counts must be nonnegative, got {vision_item_counts}.") num_views_per_vision_item.extend([sample_n_views[sample_idx]] * num_vision_items) frames_per_vision_item.extend([sample_frames_per_view[sample_idx]] * num_vision_items) @@ -5173,7 +5395,7 @@ def _encode_sound_x0_tokens( def get_data_and_condition( self, - data_batch: dict[str, torch.Tensor], + data_batch: dict[str, Any], vision_condition_indexes: list[list[int]] | None = None, retain_raw_state_vision: bool = True, balance_vae_encode: bool = False, @@ -5203,13 +5425,22 @@ def get_data_and_condition( over the whole group: only training, which every rank enters in lockstep, may turn it on. """ + normalize_vision_batch_inplace( + data_batch, video_key=self.input_video_key, tensor_kwargs_fp32=getattr(self, "tensor_kwargs_fp32", {}) + ) if not self._has_vision_stream(data_batch): return self._get_lidar_only_data_and_condition(data_batch) + multiview_action_conditioning = getattr(self.config, "multiview_action_conditioning", None) + if multiview_action_conditioning is not None: + sizes = data_batch.get("image_size") + if isinstance(sizes, list) and sizes and isinstance(sizes[0], list): + data_batch["image_size"] = [size for sample_sizes in sizes for size in sample_sizes] + # Detect whether any sample has multiple vision items (e.g. image editing). # If so, track the count per sample before all vision items from this batch are flattened into a list. is_image_batch = self.is_image_batch(data_batch) - media_key = self.input_video_key if not is_image_batch else self.input_image_key + media_key = self.input_image_key if self.input_image_key in data_batch else self.input_video_key sample_vision_list = data_batch[media_key] @@ -5336,8 +5567,28 @@ def get_data_and_condition( x0_tokens_action = raw_state_action raw_action_dim = data_batch.get("raw_action_dim", None) action_valid_mask = data_batch.get("action_valid_mask", None) + raw_action_view_counts = data_batch.get("num_views_per_action_item") + sample_is_image = data_batch.get("sample_is_image") + action_sample_ids = None + sound_sample_ids = None + if sample_is_image is not None: + action_sample_ids = [i for i, plan in enumerate(data_batch["sequence_plan"]) if plan.has_action] + raw_action_dim = [raw_action_dim[i] for i in action_sample_ids] if raw_action_dim is not None else None + action_valid_mask = ( + [action_valid_mask[i] for i in action_sample_ids] if action_valid_mask is not None else None + ) + if raw_action_view_counts is not None: + # Keep batch metadata sample-aligned for callbacks. Dense Action rows + # without explicit multiview metadata use the standard single-view layout. + raw_action_view_counts = [ + 1 if raw_action_view_counts[i] is None else raw_action_view_counts[i] for i in action_sample_ids + ] num_views_per_action_item = ( - read_positive_int_metadata(data_batch, "num_views_per_action_item", expected_count=len(raw_state_action)) + read_positive_int_metadata( + {"num_views_per_action_item": raw_action_view_counts}, + "num_views_per_action_item", + expected_count=len(raw_state_action), + ) if raw_state_action is not None else None ) @@ -5352,6 +5603,9 @@ def get_data_and_condition( else: x0_tokens_sound = None + if sample_is_image is not None: + sound_sample_ids = [i for i, plan in enumerate(data_batch["sequence_plan"]) if plan.has_sound] + # FPS metadata is used by the sequence packer for mRoPE temporal IDs when # FPS modulation is enabled in the training config. fps_raw = data_batch.get("conditioning_fps", None) @@ -5367,6 +5621,9 @@ def get_data_and_condition( fps_action_raw = torch.stack(fps_action_raw).flatten() fps_action = fps_action_raw.to(**self.tensor_kwargs) + if action_sample_ids is not None and fps_action is not None: + fps_action = fps_action[action_sample_ids] # [N_action] + # LiDAR sweep rate for mRoPE, a property of the recording rather than of the clip, so # it comes from the config the way the LiDAR VAE's compression does. fps_lidar = None @@ -5403,7 +5660,7 @@ def get_data_and_condition( num_embeddings=self.config.diffusion_expert_config.num_view_embeddings, ) # list[[V]] control_weights: list[list[float]] | None = data_batch.get("control_weights", None) - return GenerationDataClean( + data = GenerationDataClean( batch_size=batch_size, is_image_batch=is_image_batch, raw_state_vision=output_raw_state_vision, @@ -5433,7 +5690,19 @@ def get_data_and_condition( action_valid_mask=action_valid_mask, num_views_per_action_item=num_views_per_action_item, control_weights=control_weights, + sample_is_image=sample_is_image, + action_sample_ids=action_sample_ids, + sound_sample_ids=sound_sample_ids if x0_tokens_sound is not None or sample_is_image is None else [], ) + if multiview_action_conditioning is not None and data.num_vision_items_per_sample is not None: + assert self.tokenizer_vision_gen is not None + return replicate_multiview_actions( + data, + num_views=len(multiview_action_conditioning.view_codes), + temporal_factor=self.tokenizer_vision_gen.temporal_compression_factor, + action_dim=self.config.max_action_dim, + ) + return data def _has_vision_stream(self, data_batch: dict[str, Any]) -> bool: """Whether the batch carries a camera stream to tokenize. @@ -6028,6 +6297,9 @@ def is_image_batch(self, data_batch: dict[str, torch.Tensor]) -> bool: differentiate image_batch and video_batch, another from a single dataloader which we assume as video_data by default. """ + if "sample_is_image" in data_batch: + flags = data_batch["sample_is_image"] + return bool(flags) and all(flags) is_image = self.input_image_key in data_batch is_video = self.input_video_key in data_batch if not (is_image or is_video): diff --git a/cosmos_framework/model/generator/utils/batch_normalization.py b/cosmos_framework/model/generator/utils/batch_normalization.py new file mode 100644 index 000000000..1179299a5 --- /dev/null +++ b/cosmos_framework/model/generator/utils/batch_normalization.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Normalize ragged raw vision batches at the model input boundary. + +Legacy image/video dictionaries already implement the model's input contract and +are left intact. A unified ``vision`` column carries ragged per-sample items and +explicit image identities; this adapter flattens it into the same canonical +vision layout. Model stages consume sample metadata, never a loader-class flag. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from cosmos_framework.data.generator.sequence_packing.sequence import SequencePlan + + +def _items(value: Any) -> list[torch.Tensor]: + if value is None: + return [] + values = list(value) if isinstance(value, (list, tuple)) else [value] + return [item for item in values if item is not None] + + +def _scalar(value: Any, default: float) -> float: + if value is None: + return default + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("Expected scalar sample metadata.") + return float(value.item()) + return float(value) + + +def normalize_vision_batch_inplace( + batch: dict[str, Any], + *, + video_key: str, + tensor_kwargs_fp32: dict[str, Any], +) -> None: + """Convert unified raw vision items once, preserving sample ownership. + + Each canonical item has shape [1,C,T,H,W]. uint8 pixels use the existing + [-1,1] map; floating-point pixels retain their values. Per-camera uint8 items + remain unnormalized until the VAE encodes each camera, avoiding a full-size + normalized multiview allocation. Empty vision lists are valid for samples + with LiDAR. Optional sample columns keep their original aligned positions. + """ + if "vision" not in batch: + return + grouped = [_items(value) for value in batch["vision"]] + image_flags = batch.get("sample_is_image") + if not isinstance(image_flags, list) or len(image_flags) != len(grouped): + raise ValueError("Unified vision requires one sample_is_image flag per sample.") + if not all(isinstance(flag, bool) for flag in image_flags): + raise ValueError("sample_is_image entries must be boolean.") + batch_size = len(grouped) + plans = batch.get("sequence_plan", [None] * batch_size) + if len(plans) != batch_size: + raise ValueError("Sequence plans must align with the sample list.") + frames_per_view = batch.get("num_video_frames_per_view", [None] * batch_size) + views_per_sample = batch.get("sample_n_views", [None] * batch_size) + per_camera = "enable_per_camera_vae_encoding" in batch + sizes = batch.get("image_size", [None] * batch_size) + for column in (frames_per_view, views_per_sample, sizes): + if len(column) != batch_size: + raise ValueError("Vision metadata must align with the logical sample list.") + flat_sizes: list[torch.Tensor] = [] + sample_sizes: list[torch.Tensor | None] = [] + flat_vision: list[torch.Tensor] = [] + normalized_plans: list[SequencePlan] = [] + view_counts: list[torch.Tensor] = [] + frame_counts: list[torch.Tensor] = [] + + for sample_index, (items, is_image, plan) in enumerate(zip(grouped, image_flags, plans, strict=True)): + if plan is None: + for key in ("action", "sound", "lidar"): + if _items(batch.get(key, [None] * batch_size)[sample_index]): + raise ValueError(f"Samples with {key} must supply an explicit SequencePlan.") + plan = SequencePlan(has_text=True, has_vision=bool(items)) + if bool(plan.has_vision) != bool(items): + raise ValueError("SequencePlan.has_vision must match its sample's vision items.") + normalized_plans.append(plan) + sample_size = None + for item_index, item in enumerate(items): + if not isinstance(item, torch.Tensor): + raise TypeError("Vision items must be tensors.") + if is_image and item.ndim == 3: + item = item.unsqueeze(0).unsqueeze(2) # [1,C,1,H,W] + elif item.ndim == 4: + item = item.unsqueeze(0) # [1,C,T,H,W] + if item.ndim != 5 or item.shape[0] != 1 or (is_image and item.shape[2] != 1): + raise ValueError(f"Invalid vision item shape {tuple(item.shape)} for is_image={is_image}.") + if item.dtype == torch.uint8: + if not per_camera: + item = item.to(**tensor_kwargs_fp32) / 127.5 - 1.0 # [1,C,T,H,W] + elif torch.is_floating_point(item): + item = item.to(**tensor_kwargs_fp32) # [1,C,T,H,W] + else: + raise TypeError(f"Vision pixels must be uint8 or normalized floats, got {item.dtype}.") + flat_vision.append(item) + supplied_sizes = _items(sizes[sample_index]) + if supplied_sizes: + size = supplied_sizes[min(item_index, len(supplied_sizes) - 1)] # [4] or [1,4] + else: + height, width = item.shape[-2:] + size = torch.tensor([height, width, height, width]) # [4] + flat_sizes.append(size) + if sample_size is None: + sample_size = size # [4] or [1,4] + sample_sizes.append(sample_size) + if per_camera: + views = int(_scalar(views_per_sample[sample_index], 1)) + if views < 1: + raise ValueError("Camera counts must be positive.") + default_frames = items[0].shape[-3] // views if items and not is_image else 1 + frames = int(_scalar(frames_per_view[sample_index], default_frames)) + view_counts.append(torch.tensor(views)) # [] + frame_counts.append(torch.tensor(frames)) # [] + + batch[video_key] = flat_vision + batch["num_vision_items_per_sample"] = [len(items) for items in grouped] + batch["sequence_plan"] = normalized_plans + batch["image_size"] = flat_sizes + batch["sample_image_size"] = sample_sizes + batch["is_preprocessed"] = True + if per_camera: + batch["sample_n_views"] = view_counts + batch["num_video_frames_per_view"] = frame_counts + if "conditioning_fps" in batch: + batch["conditioning_fps"] = [ + torch.tensor(_scalar(value, 1.0 if is_image else 24.0)) # [] + for value, is_image in zip(batch["conditioning_fps"], image_flags, strict=True) + ] + if "conditioning_fps_action" in batch: + fallback = batch.get("conditioning_fps", [24.0] * batch_size) + batch["conditioning_fps_action"] = [ + torch.tensor(_scalar(value, _scalar(fps, 24.0))) # [] + for value, fps in zip(batch["conditioning_fps_action"], fallback, strict=True) + ] + if "lidar" in batch: + batch["lidar"] = [_items(value) for value in batch["lidar"]] + for name in ("action", "sound", "domain_id", "raw_action_dim", "action_valid_mask"): + if name in batch: + values = [_items(value) for value in batch[name]] + if any(len(items) > 1 for items in values): + raise ValueError(f"Each logical sample supports at most one {name} item.") + batch[name] = [items[0] if items else None for items in values] + if "control_weights" in batch: + batch["control_weights"] = [value if value is not None else [] for value in batch["control_weights"]] + # Mark completion by consuming the raw-format key. Calling the boundary again + # (e.g. an online sampling callback) must not normalize pixels a second time. + del batch["vision"] diff --git a/cosmos_framework/model/generator/utils/data_and_condition.py b/cosmos_framework/model/generator/utils/data_and_condition.py index 115ffddbe..c6987352f 100644 --- a/cosmos_framework/model/generator/utils/data_and_condition.py +++ b/cosmos_framework/model/generator/utils/data_and_condition.py @@ -85,6 +85,24 @@ class GenerationDataClean: # None for non-transfer or single-control samples. control_weights: list[list[float]] | None = None + # Explicit image identity per logical sample. None preserves the legacy + # homogeneous flag and its scheduling/RNG behavior. Multiple vision items + # belonging to a sample inherit its identity, including image-editing controls. + sample_is_image: list[bool] | None = None + # Sparse action/audio rows retain their owning logical sample for inference + # slicing and batched CFG. None keeps the existing homogeneous layout. + action_sample_ids: list[int] | None = None + sound_sample_ids: list[int] | None = None + + def vision_image_flags(self) -> list[bool] | None: + """Expand sample identities to the flattened vision-item order.""" + if self.sample_is_image is None: + return None + counts = self.num_vision_items_per_sample or [1] * self.batch_size + if len(counts) != self.batch_size or len(self.sample_is_image) != self.batch_size: + raise ValueError("Vision ownership and image identity must match the logical batch size.") + return [flag for flag, count in zip(self.sample_is_image, counts, strict=True) for _ in range(count)] + @dataclass(slots=True) class GenerationDataNoised: diff --git a/cosmos_framework/model/generator/utils/kv_cache.py b/cosmos_framework/model/generator/utils/kv_cache.py index 365429569..67d581d15 100644 --- a/cosmos_framework/model/generator/utils/kv_cache.py +++ b/cosmos_framework/model/generator/utils/kv_cache.py @@ -24,12 +24,15 @@ from cosmos_framework.model.generator.utils.memory import KVToStore, MemoryState, MemoryValue from cosmos_framework.data.generator.sequence_packing.runtime import get_num_real_samples, to_device_nonblocking from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig -from cosmos_framework.model.generator.rolling_prompt import RollingPromptLayout, RollingTextSinkLayer from cosmos_framework.model.generator.utils.kv_storage_backend import ( BF16StorageBackend, FP8StorageBackend, KVStorageBackend, ) +from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( + RollingPromptLayout, + RollingTextSinkLayer, +) # See cosmos_framework/data/generator/sequence_packing/ for the canonical definition. SequencePack = dict[str, Any] diff --git a/cosmos_framework/model/generator/utils/multiview_ar.py b/cosmos_framework/model/generator/utils/multiview_ar.py new file mode 100644 index 000000000..21fb1761c --- /dev/null +++ b/cosmos_framework/model/generator/utils/multiview_ar.py @@ -0,0 +1,1423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared RGB and joint RGB/LiDAR multiview Transfer autoregressive backends and replay context.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass, fields, replace +from typing import TYPE_CHECKING, Any, Literal, Protocol + +import torch + +from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_broadcast_tensor_list +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.model.generator.utils.memory import MemoryState +from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan +from cosmos_framework.model.generator.mot.causal_flex_attention import ( + _ROLE_CLEAN_TARGET, + _ROLE_CURRENT_TARGET, + FlexQueryMetadata, + MultiviewTransferARCurrentRole, + MultiviewTransferARMemoryLayout, + TeacherForcingFlexMetadata, + build_multiview_transfer_ar_memory_layout, +) +from cosmos_framework.model.generator.teacher_forcing import mark_modality_as_clean_condition +from cosmos_framework.model.generator.utils.kv_cache import ( + JointChunkMemory, + MultiviewARMemoryState, + TeacherForcingMemoryState, +) +from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import RollingPromptSchedule +from cosmos_framework.data.generator.sequence_packing.autoregressive import pack_input_sequence_autoregressive + +if TYPE_CHECKING: + from cosmos_framework.model.generator.mot.flex_attention import SensorMaskItem + from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + +MultiviewTransferARKVCache = list[tuple[torch.Tensor, torch.Tensor] | None] +_ARBranch = Literal["conditional", "unconditional"] + + +class MultiviewTransferARHost(Protocol): + """Model surface required by the shared multiview Transfer AR backend.""" + + config: Any + llm_special_tokens: dict[str, int] + net: torch.nn.Module + parallel_dims: Any + tensor_kwargs: dict[str, Any] + tokenizer_vision_gen: Any + + def _pack_input_sequence( + self, + sequence_plans: list[SequencePlan], + input_text_indexes: list[list[int]], + gen_data_clean: GenerationDataClean, + input_timesteps: torch.Tensor, + include_end_of_generation_token: bool = False, + skip_text_tokens: bool = False, + initial_mrope_temporal_offset: int | float = 0, + ) -> PackedSequence: ... + + def _cast_generated_tokens_to_precision(self, packed_sequence: PackedSequence) -> None: ... + + def denoise( + self, + net: torch.nn.Module | None = None, + data_batch_packed: PackedSequence | None = None, + memory: MemoryState | None = None, + video_temporal_causal: bool | None = None, + ) -> dict[str, Any]: ... + + +@dataclass +class MultiviewTransferARSession: + """Mutable backend-owned state for one multiview Transfer AR rollout.""" + + token_shapes: tuple[tuple[int, int, int], ...] + target_condition_mask: torch.Tensor # [V*T,1,1] + num_views: int + frames_per_view: int + frames_per_chunk: int + condition_count: int + memory_seq_len: int + control_frame_ranges: list[tuple[int, int]] + target_condition_frame_ranges: list[tuple[int, int]] + history_frame_ranges: list[tuple[int, int]] + conditional_cache: MultiviewTransferARKVCache # Sequential conditional or CFGP rank-local K/V: [1,M,H_kv,D] + unconditional_cache: MultiviewTransferARKVCache | None # K/V: [1,M,H_kv,D] or None + cfg_active: bool + cfgp_enabled: bool + text_view_ids: tuple[int, ...] | None = None + target_view_ids: torch.Tensor | None = None # [V] physical rig IDs + + +@dataclass(frozen=True) +class MultiviewTransferARReplayContext: + """Immutable layout snapshot plus shared detached cache for one replayed chunk.""" + + cache: MultiviewTransferARKVCache # K/V: [1,M,H_kv,D] + token_shapes: tuple[tuple[int, int, int], ...] + target_condition_mask: torch.Tensor # [V*T,1,1] + num_views: int + frames_per_view: int + frames_per_chunk: int + condition_count: int + history_frame_ranges: tuple[tuple[int, int], ...] + memory_seq_len: int + text_tokens: tuple[tuple[int, ...], ...] + fps_vision: tuple[float, ...] + text_view_ids: tuple[int, ...] | None = None + target_view_ids: torch.Tensor | None = None # [V] physical rig IDs + + +class MultiviewTransferARBackend: + """Own multiview Transfer packing, Flex KV layout, cache writes, and replay state.""" + + def __init__(self, host: MultiviewTransferARHost) -> None: + self.host = host + + def build_prefill_pack( + self, + *, + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + text_tokens: list[list[int]], + materialized_target_frame_ranges: Sequence[tuple[int, int]] | None = None, + ) -> PackedSequence: + """Build one full-geometry clean prefill pack for cache capture.""" + pack = self.host._pack_input_sequence( + sequence_plans, + text_tokens, + gen_data_clean, + torch.zeros(1, dtype=torch.float32), # [1] + ) + if pack.vision is None: + raise ValueError("Multiview transfer AR prefill requires packed vision data.") + original_masks = [mask.clone() for mask in pack.vision.condition_mask] # list[[V*T,1,1]] + # Clean prefill must not add a diffusion embedding to materialized RGB history. + # Keep the original conditioning mask for prefix validation and cache geometry. + mark_modality_as_clean_condition(pack.vision) + pack.vision.condition_mask = original_masks + pack.teacher_forcing_pass = "clean" + pack.teacher_forcing_original_condition_masks_vision = original_masks + if materialized_target_frame_ranges is not None: + # Preserve the full two-item geometry while preventing real queries + # from reading ungenerated target suffix values. + pack.teacher_forcing_materialized_target_frame_ranges = tuple(materialized_target_frame_ranges) + pack.to_cuda() + self.host._cast_generated_tokens_to_precision(pack) + return pack + + def create_session( + self, + *, + prefill_pack: PackedSequence, + num_views: int, + frames_per_view: int, + condition_count: int, + cfg_active: bool, + cfgp_enabled: bool, + text_view_ids: list[int] | None = None, + ) -> MultiviewTransferARSession: + """Allocate backend state from a validated prefill pack.""" + if prefill_pack.vision is None or len(prefill_pack.vision.token_shapes) != 2: + raise ValueError("Multiview transfer AR requires packed [control, target] vision metadata.") + flex_backend = getattr(self.host.net, "flex_backend", None) + maskless_replay = getattr(self.host.net, "teacher_forcing_maskless", False) + if flex_backend is None and not maskless_replay: + raise ValueError("Multiview transfer AR requires an initialized FlexAttention backend.") + if text_view_ids is not None and text_view_ids != list(range(num_views)): + raise ValueError( + f"Multiview transfer AR text must cover camera-major views {list(range(num_views))}, " + f"got {text_view_ids}." + ) + control_shape, target_shape = prefill_pack.vision.token_shapes + total_memory_tokens = control_shape[0] * control_shape[1] * control_shape[2] + total_memory_tokens += target_shape[0] * target_shape[1] * target_shape[2] + kv_alignment = 1 if maskless_replay else int(flex_backend.block_size[1]) + memory_seq_len = ((total_memory_tokens + kv_alignment - 1) // kv_alignment) * kv_alignment + target_condition_ranges = [(0, condition_count)] if condition_count else [] + num_layers = int(self.host.net.num_hidden_layers) # type: ignore[attr-defined] + return MultiviewTransferARSession( + token_shapes=tuple(prefill_pack.vision.token_shapes), + target_condition_mask=prefill_pack.vision.condition_mask[1], # [V*T,1,1] + num_views=num_views, + frames_per_view=frames_per_view, + frames_per_chunk=int(self.host.config.teacher_forcing_frames_per_chunk), + condition_count=condition_count, + memory_seq_len=memory_seq_len, + control_frame_ranges=[], + target_condition_frame_ranges=target_condition_ranges, + history_frame_ranges=[], + conditional_cache=[None] * num_layers, + unconditional_cache=[None] * num_layers if cfg_active and not cfgp_enabled else None, + cfg_active=cfg_active, + cfgp_enabled=cfgp_enabled, + text_view_ids=tuple(text_view_ids) if text_view_ids is not None else None, + target_view_ids=prefill_pack.vision_view_ids[1] + if prefill_pack.vision_view_ids is not None + else None, # [V] + ) + + @staticmethod + def _current_pack_text_tokens( + text_tokens: list[list[int]], + text_view_ids: tuple[int, ...] | None, + ) -> list[int] | list[list[int]]: + """Restore the sample-level or per-view AR text payload shape.""" + if text_view_ids is None: + if len(text_tokens) != 1: + raise ValueError(f"Sample-level multiview AR text requires one caption, got {len(text_tokens)}.") + return text_tokens[0] + if len(text_tokens) != len(text_view_ids): + raise ValueError( + f"Per-view multiview AR text carries {len(text_tokens)} captions but {len(text_view_ids)} view IDs." + ) + return text_tokens + + @staticmethod + def build_memory_layout(session: MultiviewTransferARSession) -> MultiviewTransferARMemoryLayout: + """Resolve the current fixed-slot Flex memory layout.""" + return build_multiview_transfer_ar_memory_layout( + token_shapes=list(session.token_shapes), + target_condition_mask=session.target_condition_mask, + num_views=session.num_views, + frames_per_chunk=session.frames_per_chunk, + control_frame_ranges=session.control_frame_ranges, + target_condition_frame_ranges=session.target_condition_frame_ranges, + history_frame_ranges=session.history_frame_ranges, + memory_seq_len=session.memory_seq_len, + device=session.target_condition_mask.device, + ) + + def build_current_pack( + self, + *, + vision_latent: torch.Tensor, # [1,C,V*chunk_len,H,W] + text_tokens: list[list[int]], + text_view_ids: tuple[int, ...] | None, + fps_vision: list[float], + num_views: int, + frames_per_view: int, + chunk_start: int, + memory_layout: MultiviewTransferARMemoryLayout, + current_role: MultiviewTransferARCurrentRole, + vision_view_ids: torch.Tensor | None = None, # [V] physical rig IDs, distinct from caption-local text_view_ids + ) -> PackedSequence: + """Pack one synchronized chunk at its shared camera-local mRoPE positions.""" + if vision_latent.shape[2] % num_views != 0: + raise ValueError( + f"Multiview transfer chunk latent_t={vision_latent.shape[2]} must be divisible by num_views={num_views}." + ) + chunk_len = vision_latent.shape[2] // num_views + temporal_positions = torch.arange( + chunk_start, + chunk_start + chunk_len, + dtype=torch.float32, + ).repeat(num_views) # [V*chunk_len] + ar_text_tokens = self._current_pack_text_tokens(text_tokens, text_view_ids) + pack = pack_input_sequence_autoregressive( + vision_latent=vision_latent, + action_latent=None, + text_tokens=ar_text_tokens, + timestep=0.0, + fps_vision=fps_vision, + fps_action=None, + special_tokens=self.host.llm_special_tokens, + latent_patch_size=self.host.config.diffusion_expert_config.patch_spatial, + condition_frame_indexes_vision=[], + frame_idx=0, + temporal_compression_factor=self.host.tokenizer_vision_gen.temporal_compression_factor or 4, + video_temporal_causal=False, + action_dim=self.host.config.max_action_dim, + enable_fps_modulation=self.host.config.diffusion_expert_config.enable_fps_modulation, + base_fps=float(self.host.config.diffusion_expert_config.base_fps), + unified_3d_mrope_temporal_modality_margin=( + self.host.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin + ), + vision_temporal_positions=temporal_positions, + num_views=num_views, + text_view_ids=list(text_view_ids) if text_view_ids is not None else None, + ) + pack.vision_view_ids = [vision_view_ids] if vision_view_ids is not None else None # list[[V]] or None + pack.to_cuda() + pack.multiview_transfer_ar_metadata = { + "current_frame_start": chunk_start, + "frames_per_view": frames_per_view, + "frames_per_chunk": self.host.config.teacher_forcing_frames_per_chunk, + "current_role": current_role, + "memory_layout": memory_layout, + } + if current_role == "clean_target": + if pack.vision is None: + raise ValueError("Multiview transfer clean-history packing requires vision tokens.") + mark_modality_as_clean_condition(pack.vision) + self.host._cast_generated_tokens_to_precision(pack) + return pack + + def capture_memory( + self, + *, + pack: PackedSequence, + cache: MultiviewTransferARKVCache, + memory_seq_len: int, + write_indexes: torch.Tensor, # [S_write] + write_offset: int, + cache_write_indexes: torch.Tensor | None = None, # [S_write] or None + ) -> None: + """Run a clean pass and commit selected generated-token K/V.""" + memory = MultiviewARMemoryState( + num_layers=int(self.host.net.num_hidden_layers), # type: ignore[attr-defined] + memory_seq_len=memory_seq_len, + cache=cache, + write_indexes=write_indexes, + write_offset=write_offset, + cache_write_indexes=cache_write_indexes, + ) + self.host.denoise(data_batch_packed=pack, memory=memory) + + @staticmethod + def merge_memory( + *, + destination: MultiviewTransferARKVCache, + source: MultiviewTransferARKVCache, + cache_indexes: torch.Tensor, # [S_write] + ) -> None: + """Copy selected fixed-slot K/V from a scratch no-memory clean pass.""" + if len(destination) != len(source): + raise ValueError(f"Expected matching cache layers, got {len(destination)} and {len(source)}.") + for layer_idx, source_kv in enumerate(source): + if source_kv is None: + raise ValueError(f"Clean replay did not capture K/V for layer {layer_idx}.") + source_k, source_v = source_kv + destination_kv = destination[layer_idx] + if destination_kv is None: + destination_k = torch.zeros_like(source_k) # [1,S_memory,H_kv,D] + destination_v = torch.zeros_like(source_v) # [1,S_memory,H_kv,D] + destination[layer_idx] = (destination_k, destination_v) + else: + destination_k, destination_v = destination_kv + layer_cache_indexes = cache_indexes.to(device=source_k.device, dtype=torch.long) # [S_write] + selected_k = torch.index_select(source_k, 1, layer_cache_indexes) # [1,S_write,H_kv,D] + selected_v = torch.index_select(source_v, 1, layer_cache_indexes) # [1,S_write,H_kv,D] + destination_k.index_copy_(1, layer_cache_indexes, selected_k) # [1,S_memory,H_kv,D] + destination_v.index_copy_(1, layer_cache_indexes, selected_v) # [1,S_memory,H_kv,D] + + def capture_prefill( + self, + *, + session: MultiviewTransferARSession, + pack: PackedSequence, + destination: MultiviewTransferARKVCache, + memory_layout: MultiviewTransferARMemoryLayout, + ) -> None: + """Capture prefill without allowing it to read partially built AR memory.""" + scratch_cache: MultiviewTransferARKVCache = [None] * len(destination) + self.capture_memory( + pack=pack, + cache=scratch_cache, + memory_seq_len=session.memory_seq_len, + write_indexes=memory_layout.prefill_source_token_indexes, + write_offset=0, + cache_write_indexes=memory_layout.prefill_cache_token_indexes, + ) + self.merge_memory( + destination=destination, + source=scratch_cache, + cache_indexes=memory_layout.prefill_cache_token_indexes, + ) + + def prime_control_cache( + self, + *, + session: MultiviewTransferARSession, + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + conditional_text_tokens: list[list[int]], + unconditional_text_tokens: list[list[int]] | None, + ) -> None: + """Capture control/condition K/V against the currently materialized target history.""" + session.control_frame_ranges[:] = [(0, session.frames_per_view)] + conditional_pack = self.build_prefill_pack( + sequence_plans=sequence_plans, + gen_data_clean=gen_data_clean, + text_tokens=conditional_text_tokens, + materialized_target_frame_ranges=session.history_frame_ranges, + ) + unconditional_pack = None + if session.cfg_active: + if unconditional_text_tokens is None: + raise ValueError("CFG multiview transfer AR requires unconditional text tokens.") + unconditional_pack = self.build_prefill_pack( + sequence_plans=sequence_plans, + gen_data_clean=gen_data_clean, + text_tokens=unconditional_text_tokens, + materialized_target_frame_ranges=session.history_frame_ranges, + ) + self.capture_control_cache( + session=session, + conditional_pack=conditional_pack, + unconditional_pack=unconditional_pack, + ) + + def capture_control_cache( + self, + *, + session: MultiviewTransferARSession, + conditional_pack: PackedSequence, + unconditional_pack: PackedSequence | None, + ) -> None: + """Capture already-built control prefills into every active branch cache.""" + session.control_frame_ranges[:] = [(0, session.frames_per_view)] + memory_layout = self.build_memory_layout(session) + if session.cfgp_enabled: + local_pack = conditional_pack if int(self.host.parallel_dims.cfgp_rank) == 0 else unconditional_pack + if local_pack is None: + raise RuntimeError("CFGP multiview transfer AR is missing its rank-local prefill pack.") + self.capture_prefill( + session=session, + pack=local_pack, + destination=session.conditional_cache, + memory_layout=memory_layout, + ) + return + self.capture_prefill( + session=session, + pack=conditional_pack, + destination=session.conditional_cache, + memory_layout=memory_layout, + ) + if session.unconditional_cache is not None: + if unconditional_pack is None: + raise RuntimeError("CFG multiview transfer AR is missing its unconditional prefill pack.") + self.capture_prefill( + session=session, + pack=unconditional_pack, + destination=session.unconditional_cache, + memory_layout=memory_layout, + ) + + def make_memory( + self, + session: MultiviewTransferARSession, + *, + branch: _ARBranch = "conditional", + ) -> MultiviewARMemoryState: + """Create the read-only multiview memory view for a denoising branch.""" + if session.cfgp_enabled: + if self.host.parallel_dims is None: + raise ValueError("CFGP multiview transfer AR requires initialized parallel dimensions.") + cfgp_rank = int(self.host.parallel_dims.cfgp_rank) + if cfgp_rank not in (0, 1): + raise ValueError(f"CFGP multiview transfer AR requires rank 0 or 1, got {cfgp_rank}.") + local_branch: _ARBranch = "conditional" if cfgp_rank == 0 else "unconditional" + if branch != local_branch: + raise ValueError(f"CFGP rank {cfgp_rank} owns the {local_branch} branch, not {branch}.") + # CFGP stores the rank-local branch in the only allocated cache. + cache = session.conditional_cache + else: + cache = session.conditional_cache if branch == "conditional" else session.unconditional_cache + if cache is None: + raise ValueError(f"Multiview transfer AR has no {branch} cache.") + return MultiviewARMemoryState( + num_layers=int(self.host.net.num_hidden_layers), # type: ignore[attr-defined] + memory_seq_len=session.memory_seq_len, + cache=cache, + ) + + def commit_clean_chunk( + self, + *, + session: MultiviewTransferARSession, + denoised_chunk: torch.Tensor, # [1,C,V*chunk_len,H,W] + chunk_start: int, + chunk_end: int, + conditional_text_tokens: list[list[int]], + unconditional_text_tokens: list[list[int]] | None, + fps_vision: list[float], + ) -> None: + """Write one finalized target chunk into branch caches and advance history.""" + memory_layout = self.build_memory_layout(session) + conditional_pack = self.build_current_pack( + vision_latent=denoised_chunk, + text_tokens=conditional_text_tokens, + text_view_ids=session.text_view_ids, + vision_view_ids=session.target_view_ids, + fps_vision=fps_vision, + num_views=session.num_views, + frames_per_view=session.frames_per_view, + chunk_start=chunk_start, + memory_layout=memory_layout, + current_role="clean_target", + ) + unconditional_pack = None + if session.cfg_active: + if unconditional_text_tokens is None: + raise ValueError("CFG multiview transfer AR requires unconditional text tokens.") + unconditional_pack = self.build_current_pack( + vision_latent=denoised_chunk, + text_tokens=unconditional_text_tokens, + text_view_ids=session.text_view_ids, + vision_view_ids=session.target_view_ids, + fps_vision=fps_vision, + num_views=session.num_views, + frames_per_view=session.frames_per_view, + chunk_start=chunk_start, + memory_layout=memory_layout, + current_role="clean_target", + ) + target_shape = session.token_shapes[1] + chunk_len = chunk_end - chunk_start + spatial_tokens = target_shape[1] * target_shape[2] + chunk_token_count = session.num_views * chunk_len * spatial_tokens + write_indexes = torch.arange( + chunk_token_count, + device=session.target_condition_mask.device, + dtype=torch.long, + ) # [chunk_tokens] + cache_write_indexes = memory_layout.target_cache_token_indexes((chunk_start, chunk_end)) # [chunk_tokens] + if session.cfgp_enabled: + local_pack = conditional_pack if int(self.host.parallel_dims.cfgp_rank) == 0 else unconditional_pack + if local_pack is None: + raise RuntimeError("CFGP multiview transfer AR is missing its rank-local clean chunk pack.") + self.capture_memory( + pack=local_pack, + cache=session.conditional_cache, + memory_seq_len=session.memory_seq_len, + write_indexes=write_indexes, + write_offset=0, + cache_write_indexes=cache_write_indexes, + ) + else: + self.capture_memory( + pack=conditional_pack, + cache=session.conditional_cache, + memory_seq_len=session.memory_seq_len, + write_indexes=write_indexes, + write_offset=0, + cache_write_indexes=cache_write_indexes, + ) + if session.unconditional_cache is not None: + if unconditional_pack is None: + raise RuntimeError("CFG multiview transfer AR is missing its unconditional clean chunk pack.") + self.capture_memory( + pack=unconditional_pack, + cache=session.unconditional_cache, + memory_seq_len=session.memory_seq_len, + write_indexes=write_indexes, + write_offset=0, + cache_write_indexes=cache_write_indexes, + ) + session.history_frame_ranges.append((chunk_start, chunk_end)) + + @staticmethod + def slice_chunk( + vision_tokens: torch.Tensor, # [1,C,V*T,H,W] + *, + num_views: int, + frames_per_view: int, + chunk_start: int, + chunk_end: int, + ) -> torch.Tensor: # [1,C,V*chunk_len,H,W] + """Slice one synchronized camera-major chunk from every view.""" + if vision_tokens.ndim != 5 or vision_tokens.shape[0] != 1: + raise ValueError(f"Expected [1,C,V*T,H,W] latents, got shape {tuple(vision_tokens.shape)}.") + if num_views < 1 or vision_tokens.shape[2] != num_views * frames_per_view: + raise ValueError( + f"Expected camera-major latent_t={num_views * frames_per_view}, got {vision_tokens.shape[2]}." + ) + if not 0 <= chunk_start < chunk_end <= frames_per_view: + raise ValueError(f"Invalid synchronized chunk [{chunk_start},{chunk_end}) for T={frames_per_view}.") + return torch.cat( + [ + vision_tokens[ + :, + :, + view_idx * frames_per_view + chunk_start : view_idx * frames_per_view + chunk_end, + ] + for view_idx in range(num_views) + ], + dim=2, + ) # [1,C,V*chunk_len,H,W] + + @staticmethod + def scatter_chunk( + destination: torch.Tensor, # [1,C,V*T,H,W] + chunk: torch.Tensor, # [1,C,V*chunk_len,H,W] + *, + num_views: int, + frames_per_view: int, + chunk_start: int, + chunk_end: int, + ) -> None: + """Write one synchronized camera-major chunk into its per-view ranges.""" + chunk_len = chunk_end - chunk_start + if destination.ndim != 5 or destination.shape[2] != num_views * frames_per_view: + raise ValueError(f"Invalid multiview destination shape {tuple(destination.shape)}.") + if chunk.ndim != 5 or chunk.shape[2] != num_views * chunk_len: + raise ValueError(f"Invalid multiview chunk shape {tuple(chunk.shape)}.") + for view_idx in range(num_views): + source_start = view_idx * chunk_len + target_start = view_idx * frames_per_view + chunk_start + destination[:, :, target_start : target_start + chunk_len].copy_( + chunk[:, :, source_start : source_start + chunk_len] + ) # [1,C,chunk_len,H,W] + + @staticmethod + def snapshot_replay_context( + session: MultiviewTransferARSession, + *, + text_tokens: list[list[int]], + fps_vision: list[float], + ) -> MultiviewTransferARReplayContext: + """Snapshot layout metadata while sharing the detached cache storage.""" + return MultiviewTransferARReplayContext( + cache=session.conditional_cache, + token_shapes=session.token_shapes, + target_condition_mask=session.target_condition_mask, + num_views=session.num_views, + frames_per_view=session.frames_per_view, + frames_per_chunk=session.frames_per_chunk, + condition_count=session.condition_count, + history_frame_ranges=tuple(session.history_frame_ranges), + memory_seq_len=session.memory_seq_len, + text_tokens=tuple(tuple(tokens) for tokens in text_tokens), + fps_vision=tuple(fps_vision), + text_view_ids=session.text_view_ids, + target_view_ids=session.target_view_ids, + ) + + def build_replay_pack_and_memory( + self, + *, + context: MultiviewTransferARReplayContext, + vision_latent: torch.Tensor, # [1,C,V*chunk_len,H,W] + chunk_start: int, + timestep: float, + ) -> tuple[PackedSequence, MultiviewARMemoryState]: + """Rebuild one differentiable current-chunk input from a backend replay context.""" + target_condition_ranges = [(0, context.condition_count)] if context.condition_count else [] + memory_layout = build_multiview_transfer_ar_memory_layout( + token_shapes=list(context.token_shapes), + target_condition_mask=context.target_condition_mask, + num_views=context.num_views, + frames_per_chunk=context.frames_per_chunk, + control_frame_ranges=[(0, context.frames_per_view)], + target_condition_frame_ranges=target_condition_ranges, + history_frame_ranges=list(context.history_frame_ranges), + memory_seq_len=context.memory_seq_len, + device=context.target_condition_mask.device, + ) + pack = self.build_current_pack( + vision_latent=vision_latent, + text_tokens=[list(tokens) for tokens in context.text_tokens], + text_view_ids=context.text_view_ids, + vision_view_ids=context.target_view_ids, + fps_vision=list(context.fps_vision), + num_views=context.num_views, + frames_per_view=context.frames_per_view, + chunk_start=chunk_start, + memory_layout=memory_layout, + current_role="current_target", + ) + if pack.vision is None: + raise ValueError("Multiview transfer replay requires packed vision tokens.") + num_vision_patches = len(pack.vision.mse_loss_indexes) + pack.vision.timesteps = torch.full( + (num_vision_patches,), + timestep, + device=self.host.tensor_kwargs["device"], + dtype=torch.float32, + ) # [N_noisy_vision] + memory = MultiviewARMemoryState( + num_layers=int(self.host.net.num_hidden_layers), # type: ignore[attr-defined] + memory_seq_len=context.memory_seq_len, + cache=context.cache, + ) + return pack, memory + + +class JointARCache: + """Handle persistent-K/V layout and attention metadata for a joint AR pipeline. + + This coordinator maps current RGB/LiDAR chunks to cache slots, restores their + original position IDs, and builds token labels for maskless attention. + JointChunkMemory owns the actual per-layer K/V buffers. + + Each layer's K and V buffers have shape [1,S,H_kv/CP,D], where S is the allocated + RGB+LiDAR token capacity, including unwritten slots. + Storage along S and its matching metadata are ordered as follows:: + + K/V slots: [ RGB control | RGB target | LiDAR control | LiDAR target ] + Metadata slots: [ RGB control | RGB target | LiDAR control | LiDAR target ] + same slot index identifies the same token + + Within each RGB item: [view 0: frames_0,frame_1..][view 1: ...]... + Within each LiDAR item: [frames 0,1...] + + Each frame contains contiguous spatial tokens. + Unwritten slots have padding roles and are excluded from attention groups. + """ + + reference: PackedSequence + memory: JointChunkMemory + metadata: dict[str, torch.Tensor] | None + current_metadata: dict[str, torch.Tensor] | None + item_layouts: list[tuple[int, int, int, int]] + rgb_frames: tuple[int, ...] + lidar_frames: tuple[int, ...] + spans: list[tuple[int, int, int]] + + def __init__(self, reference: PackedSequence, num_layers: int) -> None: + self.reference = reference + assert reference.vision is not None and reference.lidar is not None + views = reference.num_views_per_vision_item + assert views is not None + if len(reference.vision.token_shapes) != 2 or len(reference.lidar.token_shapes) != 2 or len(views) != 2: + raise ValueError("Joint AR cache requires one control and one target item for each of RGB and LiDAR.") + rgb_control_shape, rgb_target_shape = reference.vision.token_shapes + lidar_control_shape, lidar_target_shape = reference.lidar.token_shapes + # Cache order is RGB control, RGB target, LiDAR control, LiDAR target. + # RGB items have per-item camera counts; LiDAR has no camera-view axis. + items = ( + (rgb_control_shape, views[0]), + (rgb_target_shape, views[1]), + (lidar_control_shape, 1), + (lidar_target_shape, 1), + ) + self.item_layouts = [] + capacity = 0 + for (total_frames, height, width), num_views in items: + slots = total_frames // num_views + self.item_layouts.append((capacity, num_views, slots, height * width)) + capacity += num_views * slots * height * width + self.memory = JointChunkMemory(num_layers, capacity) + self.metadata = None + self.current_metadata = None + self.rgb_frames = () + self.lidar_frames = () + self.spans = [] + + def prepare(self, pack: PackedSequence, rgb_frames: tuple[int, ...], lidar_frames: tuple[int, ...]) -> None: + """Prepare current-chunk cache writes, positional encoding, and attention visibility. + + 1. Compute spans mapping current-chunk K/V into persistent history for + the branch's clean forward. + 2. Restore current tokens' original position IDs from the full reference + pack before RoPE runs. Packing only a chunk can give its tokens local + positions. Restoring the original IDs keeps current Q/K in the same + positional coordinate system as the cached historical keys. + 3. Record original RGB/LiDAR frame IDs and attach this coordinator to the + pack. The network hook uses those IDs and stored history labels to + build attention visibility for current tokens and cached history in + both noisy denoising and clean forwards. + + Position IDs preserve positional encoding; metadata preserves attention + visibility. This method does not write K/V or enable cache writes. + """ + self.rgb_frames, self.lidar_frames = rgb_frames, lidar_frames + self.spans = [] + source = 0 + for modality_name, frame_ids in (("vision", rgb_frames), ("lidar", lidar_frames)): + modality = getattr(pack, modality_name) + original = getattr(self.reference, modality_name) + assert modality is not None and original is not None + current_offset = original_offset = 0 + for local_item, (total_frames, height, width) in enumerate(original.token_shapes): + item = local_item + (0 if modality_name == "vision" else 2) + base, views, slots, spatial = self.item_layouts[item] + original_frames = total_frames // views + selected = [view * original_frames + frame for view in range(views) for frame in frame_ids] + frame_tensor = torch.tensor( + selected, dtype=torch.long, device=original.sequence_indexes.device + ) # [V*C] + original_indexes = original.sequence_indexes[ + original_offset : original_offset + total_frames * spatial + ] # [V*T*S] + selected_indexes = original_indexes.reshape(total_frames, spatial)[frame_tensor].flatten() # [V*C*S] + length = len(selected) * spatial + current_indexes = modality.sequence_indexes[current_offset : current_offset + length] # [V*C*S] + pack.position_ids[:, current_indexes] = self.reference.position_ids[:, selected_indexes].to( + pack.position_ids.device + ) # [3,V*C*S] + for view in range(views): + for frame in frame_ids: + destination = base + (view * slots + frame) * spatial + if ( + self.spans + and self.spans[-1][0] + self.spans[-1][2] == source + and self.spans[-1][1] + self.spans[-1][2] == destination + ): + previous = self.spans[-1] + self.spans[-1] = (previous[0], previous[1], previous[2] + spatial) + else: + self.spans.append((source, destination, spatial)) + source += spatial + current_offset += length + original_offset += total_frames * spatial + setattr(pack, "joint_ar_cache", self) + + def augment_metadata( + self, + metadata: TeacherForcingFlexMetadata, + sensor_items: list[list[SensorMaskItem]], + chunk_size: int, + *, + noisy: bool, + ) -> TeacherForcingFlexMetadata: + """Build attention-visibility metadata for current tokens and persistent history. + + The returned metadata describes current sensor queries and keys ordered as follows:: + + Query labels: [current sensor chunk] + Key labels: [text | current sensor chunk | persistent history] + + The network uses this metadata to build the maskless attention plan. + """ + names = [field.name for field in fields(FlexQueryMetadata)] + current = {name: getattr(metadata.query, name).clone() for name in names} # each [Q] + if noisy: + current["token_role_id"] = torch.where( + current["token_role_id"] == _ROLE_CLEAN_TARGET, + _ROLE_CURRENT_TARGET, + current["token_role_id"], + ) # [Q] + offset = 0 + for item_index, item in enumerate(sensor_items[0]): + frame_ids = self.rgb_frames if item_index < 2 else self.lidar_frames + spatial = item.token_shape[1] * item.token_shape[2] + logical = ( + torch.tensor(frame_ids, dtype=torch.long, device=current["frame_id"].device) + .repeat(item.num_views) + .repeat_interleave(spatial) + ) # [V*C*S] + length = logical.numel() + current["frame_id"][offset : offset + length] = logical # [V*C*S] + current["timestamp"][offset : offset + length] = logical * item.seconds_per_frame # [V*C*S] + current["causal_step_id"][offset : offset + length] = torch.ceil( + logical.double() * item.seconds_per_frame / (chunk_size * sensor_items[0][0].seconds_per_frame) - 1e-5 + ).long() # [V*C*S] + offset += length + self.current_metadata = current + values = {} + for name in names: + parts = [getattr(metadata, name)[: metadata.num_und], current[name]] # [UND], [Q] + if self.metadata is not None: + parts.append(self.metadata[name]) # [M] + values[name] = torch.cat(parts) # [UND+Q+M] + return replace(metadata, **values, query=FlexQueryMetadata(**current)) + + def commit_metadata(self) -> None: + """Publish occupied slots only after every layer has written its clean K/V.""" + assert self.current_metadata is not None + # Stored targets are clean memory, never noisy keys in later chunks. + self.current_metadata["is_noisy"] = torch.zeros_like(self.current_metadata["is_noisy"]) # [Q] + if self.metadata is None: + self.metadata = { + name: torch.full( + (self.memory.memory_seq_len,), + False if value.dtype == torch.bool else -1, + device=value.device, + dtype=value.dtype, + ) + for name, value in self.current_metadata.items() + } # each [M] + for name, value in self.current_metadata.items(): # value: [Q] + for source, destination, length in self.spans: + self.metadata[name][destination : destination + length] = value[source : source + length] # [L] + self.memory.spans = None + + +@dataclass(frozen=True) +class JointARChunk: + """A complete training-time causal chunk on the shared RGB/LiDAR clock.""" + + step: int + vision_frames: tuple[int, ...] + lidar_frames: tuple[int, ...] + vision_prefix_end: int + lidar_prefix_end: int + + +def joint_ar_chunks( + *, + vision_frames: int, + lidar_frames: int, + vision_seconds_per_frame: float, + lidar_seconds_per_frame: float, + frames_per_chunk: int, +) -> list[JointARChunk]: + """Match the teacher-forcing singleton-zero and shared-time chunk boundaries.""" + if min(vision_frames, lidar_frames, frames_per_chunk) < 1: + raise ValueError("Joint AR requires positive frame counts and chunk size") + if not all(math.isfinite(value) and value > 0 for value in (vision_seconds_per_frame, lidar_seconds_per_frame)): + raise ValueError("Joint AR requires finite positive sensor periods") + chunk_seconds = frames_per_chunk * vision_seconds_per_frame + vision_steps = [math.ceil(frame / frames_per_chunk - 1e-5) for frame in range(vision_frames)] + lidar_steps = [math.ceil(frame * lidar_seconds_per_frame / chunk_seconds - 1e-5) for frame in range(lidar_frames)] + chunks: list[JointARChunk] = [] + for step in sorted(set(vision_steps) | set(lidar_steps)): + chunks.append( + JointARChunk( + step=step, + vision_frames=tuple(frame for frame, value in enumerate(vision_steps) if value == step), + lidar_frames=tuple(frame for frame, value in enumerate(lidar_steps) if value == step), + vision_prefix_end=sum(value <= step for value in vision_steps), + lidar_prefix_end=sum(value <= step for value in lidar_steps), + ) + ) + return chunks + + +def _camera_prefix( + latent: torch.Tensor, # [1,C,V*T,H,W] + num_views: int, + end: int, +) -> torch.Tensor: # [1,C,V*end,H,W] + batch, channels, total_frames, height, width = latent.shape + grid = latent.reshape(batch, channels, num_views, total_frames // num_views, height, width) # [1,C,V,T,H,W] + return grid[:, :, :, :end].reshape(batch, channels, num_views * end, height, width) # [1,C,V*end,H,W] + + +def _conditioned_target( + latent: torch.Tensor, # [1,C,T,H,W] + mask: torch.Tensor, # [T,1,1] +) -> torch.Tensor: # [1,C,T,H,W] + """Discard every unconditioned target value before any model forward.""" + condition = mask.to(device=latent.device, dtype=torch.bool).reshape(1, 1, -1, 1, 1) # [1,1,T,1,1] + return torch.where(condition, latent, torch.zeros_like(latent)) # [1,C,T,H,W] + + +def _prefix_condition_count(mask: torch.Tensor, num_views: int) -> int: # mask: [V*T,1,1] + """Reject unsynchronized or non-prefix conditioning in this joint sampler.""" + grid = mask.to(dtype=torch.bool).reshape(num_views, -1) # [V,T] + if not torch.equal(grid, grid[:1].expand_as(grid)): + raise ValueError("Joint AR currently requires the same conditioned prefix for every RGB view") + count = int(grid[0].sum().item()) + expected = torch.arange(grid.shape[1], device=grid.device) < count # [T] + if not torch.equal(grid[0], expected): + raise ValueError("Joint AR requires a contiguous conditioned target prefix") + return count + + +def _validate_persistent_chunks(chunks: list[JointARChunk], vision_condition: int, lidar_condition: int) -> None: + """Check whether persistent K/V can replace replay for the complete chunk schedule. + + Persistent K/V reuses stored clean history instead of rebuilding it through replay, + but cannot currently handle: + + 1. A causal chunk missing RGB or LiDAR frames; both sensors must be present. + 2. User-provided RGB/LiDAR target conditioning extending beyond the first + causal chunk. Later conditioned queries may read earlier noisy target + representations in replay; substituting stored clean K/V is not equivalent. + + For either case, use the replay K/V mechanism by setting joint_ar_use_persistent_kv=False. + This check raises rather than silently falling back to replay. + """ + if any(not chunk.vision_frames or not chunk.lidar_frames for chunk in chunks): + raise ValueError( + "Persistent joint K/V currently requires both sensors (RGB and LiDAR) in every causal chunk. " + "Use the replay K/V mechanism: set joint_ar_use_persistent_kv=False." + ) + if vision_condition > chunks[0].vision_prefix_end or lidar_condition > chunks[0].lidar_prefix_end: + raise ValueError( + "Persistent joint K/V currently supports target conditioning only in the first causal chunk; " + "later conditioning can read noisy replay history. Use the replay K/V mechanism: " + "set joint_ar_use_persistent_kv=False." + ) + + +def _prefix_data( + data: GenerationDataClean, + *, + vision_target: torch.Tensor, # [1,Cv,V*Tv,Hv,Wv] + lidar_target: torch.Tensor, # [1,Cl,Tl,Hl,Wl] + num_views: int, + chunk: JointARChunk, +) -> GenerationDataClean: + """Keep complete sensor chunks, with no future target or control positions.""" + assert data.x0_tokens_vision is not None and data.x0_tokens_lidar is not None + if data.temporal_positions_vision is not None: + raise ValueError("Joint AR currently requires the training latent-index temporal positions") + return replace( + data, + raw_state_vision=None, + raw_state_lidar=None, + x0_tokens_vision=[ + _camera_prefix(data.x0_tokens_vision[0], num_views, chunk.vision_prefix_end), # [1,Cv,V*Tv_prefix,Hv,Wv] + _camera_prefix(vision_target, num_views, chunk.vision_prefix_end), # [1,Cv,V*Tv_prefix,Hv,Wv] + ], + x0_tokens_lidar=[ + data.x0_tokens_lidar[0][:, :, : chunk.lidar_prefix_end], # [1,Cl,Tl_prefix,Hl,Wl] + lidar_target[:, :, : chunk.lidar_prefix_end], # [1,Cl,Tl_prefix,Hl,Wl] + ], + ) + + +def _build_replay_branch( + host: OmniMoTCausalModel, + plans: list[SequencePlan], + data: GenerationDataClean, + text_tokens: list[list[int]], +) -> tuple[PackedSequence, TeacherForcingMemoryState]: + """Use the actual joint training clean/noisy passes, including caption isolation.""" + pack = host._pack_input_sequence(plans, text_tokens, data, torch.zeros(1, dtype=torch.float32)) # timestep: [1] + pack.to_cuda() + host._cast_generated_tokens_to_precision(pack) + host._validate_teacher_forcing_pack(pack) + # Replay recomputes clean K/V for the complete causal prefix. + memory = host._build_clean_tf_cache( + net=host.net, + packed_sequence=pack, + gen_data_clean=data, + memory_info={}, + detach_clean_kv=True, + ) + return pack, memory + + +def _current_chunk_data( + data: GenerationDataClean, + vision_target: torch.Tensor, # [1,Cv,V*Tv,Hv,Wv] + lidar_target: torch.Tensor, # [1,Cl,Tl,Hl,Wl] + num_views: int, + chunk: JointARChunk, +) -> GenerationDataClean: + """Materialize only this shared-clock chunk, keeping control/target pairing.""" + assert data.x0_tokens_vision is not None and data.x0_tokens_lidar is not None + frames = vision_target.shape[2] // num_views + if data.temporal_positions_vision is not None: + raise ValueError("Joint AR currently requires the training latent-index temporal positions") + indexes = [view * frames + frame for view in range(num_views) for frame in chunk.vision_frames] + return replace( + data, + raw_state_vision=None, + raw_state_lidar=None, + x0_tokens_vision=[ + data.x0_tokens_vision[0][:, :, indexes], + vision_target[:, :, indexes], + ], # each [1,Cv,V*C_rgb,Hv,Wv] + x0_tokens_lidar=[ + data.x0_tokens_lidar[0][:, :, list(chunk.lidar_frames)], + lidar_target[:, :, list(chunk.lidar_frames)], + ], # each [1,Cl,C_lidar,Hl,Wl] + ) + + +def _build_persistent_branch( + host: OmniMoTCausalModel, + plans: list[SequencePlan], + data: GenerationDataClean, + text: list[list[int]], + cache: JointARCache, + chunk: JointARChunk, + *, + clean: bool = False, +) -> tuple[PackedSequence, JointChunkMemory]: + """Prepare current-chunk inputs for attention against persistent history K/V.""" + local_plans = [ + replace( + plans[0], + condition_frame_indexes_vision=[ + i for i, frame in enumerate(chunk.vision_frames) if frame in plans[0].condition_frame_indexes_vision + ], + condition_frame_indexes_lidar=[ + i for i, frame in enumerate(chunk.lidar_frames) if frame in plans[0].condition_frame_indexes_lidar + ], + ) + ] + pack = host._pack_input_sequence(local_plans, text, data, torch.zeros(1)) # timestep: [1] + host._validate_teacher_forcing_pack(pack) + assert pack.vision is not None and pack.lidar is not None + original_masks = [ + mask.clone() for modality in (pack.vision, pack.lidar) for mask in modality.condition_mask + ] # list of [V*C,1,1] + cache.prepare(pack, chunk.vision_frames, chunk.lidar_frames) + if clean: + for modality in (pack.vision, pack.lidar): + mark_modality_as_clean_condition(modality) + pack.teacher_forcing_pass = "clean" if clean else "noisy" + pack.teacher_forcing_original_condition_masks_sensors = original_masks + pack.to_cuda() + host._cast_generated_tokens_to_precision(pack) + return pack, cache.memory + + +def _commit_persistent_joint_chunk( + host: OmniMoTCausalModel, + plans: list[SequencePlan], + data: GenerationDataClean, + branches: list[tuple[list[list[int]], JointARCache]], + chunk: JointARChunk, +) -> None: + """Store a finalized chunk's clean K/V and attention metadata in persistent history. + + Prior persistent history metadata stays unchanged. + """ + for text, cache in branches: + pack, memory = _build_persistent_branch(host, plans, data, text, cache, chunk, clean=True) + cache.memory.spans = cache.spans + try: + host.denoise(data_batch_packed=pack, memory=memory) + cache.commit_metadata() + finally: + cache.memory.spans = None + + +def sample_joint_transfer_ar( + host: OmniMoTCausalModel, + *, + plans: list[SequencePlan], + data: GenerationDataClean, + conditional_text: list[list[int]], + unconditional_text: list[list[int]] | None, + guidance: float, + seed: int, + num_steps: int, + shift: float, + sampler_mode: Literal["rf", "distilled"] | None = None, + distilled_num_steps: int | None = None, + prompt_schedule: RollingPromptSchedule | None = None, + joint_ar_use_persistent_kv: bool = False, +) -> dict[str, list[torch.Tensor]]: # returns vision: [1,Cv,V*Tv,Hv,Wv], lidar: [1,Cl,Tl,Hl,Wl] + """Joint AR sampling with reference replay or opt-in persistent clean history. + + Both targets are denoised together. Each clean replay contains only conditions + and previously generated values; ungenerated targets are zero placeholders. + The existing noisy-pass mask reads only strictly earlier clean chunks, and its + current-target edges stay inside the current chunk. A complete prefix is used + so controls can recompute their training-time history-dependent K/V. This is + deliberately a separate inference path: it does not alter training or the + optimized RGB-only AR cache. + + ``joint_ar_use_persistent_kv=True`` selects current-chunk-only maskless + execution. Each finalized chunk is committed once; later denoising reads + those immutable keys/values instead of recomputing historical queries. + This temporary opt-in keeps replay available during numerical qualification. + """ + fixed_sampler = getattr(host.config, "fixed_step_sampler_config", None) + if sampler_mode is None: + sampler_mode = "distilled" if fixed_sampler is not None else "rf" + if sampler_mode not in ("rf", "distilled"): + raise ValueError("Joint AR sampler_mode must be 'rf' or 'distilled'.") + if sampler_mode == "rf" and fixed_sampler is not None: + raise ValueError("A joint student checkpoint requires sampler_mode='distilled'.") + if sampler_mode == "distilled": + if fixed_sampler is None: + raise ValueError("Joint distilled sampling requires a fixed-step student configuration.") + if len(plans) != 1 or not plans[0].has_lidar or not plans[0].has_vision: + raise ValueError("Joint AR requires exactly one RGB+LiDAR sequence") + if plans[0].has_action or plans[0].has_sound: + raise ValueError("Joint AR supports only RGB and LiDAR targets") + if not host._uses_multiview_replay_kv() or host.config.compile.enabled: + raise ValueError("Joint AR requires eager multiview teacher forcing") + if host.parallel_dims is not None and host.parallel_dims.cfgp_enabled: + raise ValueError("Joint AR currently uses serial CFG branches; set cfg_parallel_shard_degree=1") + policy = host._get_teacher_forcing_replay_policy() + if policy.control_visibility != "causal": + raise ValueError("Joint AR prefix replay requires causal controls") + if data.num_vision_items_per_sample != [2] or data.num_lidar_items_per_sample != [2]: + raise ValueError("Joint AR requires control and target items for each sensor modality") + if data.x0_tokens_vision is None or data.x0_tokens_lidar is None: + raise ValueError("Joint AR requires encoded camera and LiDAR streams") + views = data.num_views_per_vision_item + if views is None or len(views) != 2 or views[0] < 1 or views[0] != views[1]: + raise ValueError("Joint AR requires matching per-camera VAE metadata") + num_views = views[0] + if len(data.x0_tokens_vision) != 2 or len(data.x0_tokens_lidar) != 2: + raise ValueError("Joint AR requires exactly two encoded items per modality") + for items in (data.x0_tokens_vision, data.x0_tokens_lidar): + if items[0].ndim != 5 or items[0].shape[0] != 1 or items[0].shape != items[1].shape: + raise ValueError("Joint AR requires aligned [1,C,T,H,W] control and target latents") + if data.x0_tokens_vision[1].shape[2] % num_views: + raise ValueError("Joint AR camera-major latent count must be divisible by its view count") + if guidance != 1.0 and unconditional_text is None: + raise ValueError("Joint AR guidance requires unconditional captions") + if seed < 0 or num_steps < 1 or not math.isfinite(guidance) or not math.isfinite(shift) or shift <= 0: + raise ValueError("Joint AR requires a nonnegative seed and finite valid sampling settings") + + # Rank-local VAE encoding can differ even with identical pixels and seeds. + # CP shards must read one canonical set of controls and conditioned targets + # before the initial pack and every subsequent clean/noisy replay. + # Keep the existing Flex inference input behavior unchanged. + if host._get_teacher_forcing_kv_implementation() == "multiview_maskless_kv": + context_parallel_broadcast_tensor_list(data.x0_tokens_vision, host.parallel_dims) # each [1,Cv,V*Tv,Hv,Wv] + context_parallel_broadcast_tensor_list(data.x0_tokens_lidar, host.parallel_dims) # each [1,Cl,Tl,Hl,Wl] + + if getattr(host.config, "rolling_kv_cache_chunks", None) is not None: + if joint_ar_use_persistent_kv: + raise ValueError( + "rolling_kv_cache_chunks uses a separate cache path; disable --joint-ar-use-persistent-kv." + ) + from cosmos_framework.model.generator.utils.rolling_kv.rolling_transfer_ar import ( + sample_rolling_joint_transfer_ar, + ) + + return sample_rolling_joint_transfer_ar( + host, + plans=plans, + data=data, + conditional_text=conditional_text, + unconditional_text=unconditional_text, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + prompt_schedule=prompt_schedule, + ) + + if prompt_schedule is not None: + raise ValueError("Segmented prompts require rolling KV inference.") + + initial_pack = host._pack_input_sequence(plans, conditional_text, data, torch.zeros(1)) # timestep: [1] + if initial_pack.vision is None or initial_pack.lidar is None: + raise ValueError("The joint inference pack omitted a required modality") + # Keep clock and prefix helpers independent of full transformer construction. + from cosmos_framework.model.generator.mot.causal_cosmos3_vfm_network import build_interactive_multiview_mask_items + + sensor_items = build_interactive_multiview_mask_items(initial_pack) + if len(sensor_items) != 1 or len(sensor_items[0]) != 4: + raise ValueError("Joint AR requires the training [RGB control, RGB target, LiDAR control, LiDAR target] layout") + rgb_control, rgb_item, lidar_control, lidar_item = sensor_items[0] + if not (rgb_control.is_control and lidar_control.is_control) or rgb_item.is_control or lidar_item.is_control: + raise ValueError("Joint AR sensor control/target roles differ from the training layout") + vision_condition = _prefix_condition_count(initial_pack.vision.condition_mask[1], num_views) + lidar_condition = _prefix_condition_count(initial_pack.lidar.condition_mask[1], 1) + vision_target = _conditioned_target(data.x0_tokens_vision[1], initial_pack.vision.condition_mask[1]).to( + device=host.tensor_kwargs["device"], dtype=torch.float32 + ) # [1,Cv,V*Tv,Hv,Wv] + lidar_target = _conditioned_target(data.x0_tokens_lidar[1], initial_pack.lidar.condition_mask[1]).to( + device=host.tensor_kwargs["device"], dtype=torch.float32 + ) # [1,Cl,Tl,Hl,Wl] + frames_per_view = vision_target.shape[2] // num_views + chunks = joint_ar_chunks( + vision_frames=frames_per_view, + lidar_frames=lidar_target.shape[2], + vision_seconds_per_frame=rgb_item.seconds_per_frame, + lidar_seconds_per_frame=lidar_item.seconds_per_frame, + frames_per_chunk=int(host.config.teacher_forcing_frames_per_chunk), + ) + persistent_branches = None + if joint_ar_use_persistent_kv: + if host._get_teacher_forcing_kv_implementation() != "multiview_maskless_kv": + raise ValueError("Persistent joint K/V requires maskless attention") + if host.config.kv_cache_inference_size is not None or host.config.attention_sink_size != 0: + raise ValueError("Persistent joint K/V is unbounded in this phase; leave the window unset and sink at zero") + _validate_persistent_chunks(chunks, vision_condition, lidar_condition) + persistent_branches = [(conditional_text, JointARCache(initial_pack, host.net.num_hidden_layers))] + if guidance != 1.0 and unconditional_text is not None: + reference = host._pack_input_sequence(plans, unconditional_text, data, torch.zeros(1)) # timestep: [1] + persistent_branches.append((unconditional_text, JointARCache(reference, host.net.num_hidden_layers))) + del initial_pack + for chunk in chunks: + _sample_joint_chunk( + host, + plans=plans, + data=data, + conditional_text=conditional_text, + unconditional_text=unconditional_text, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + vision_target=vision_target, + lidar_target=lidar_target, + num_views=num_views, + frames_per_view=frames_per_view, + vision_condition=vision_condition, + lidar_condition=lidar_condition, + chunk=chunk, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + persistent_branches=persistent_branches, + # The last chunk need not be retained: no later chunk reads its K/V. + retain_history=persistent_branches is not None and chunk is not chunks[-1], + ) + return {"vision": [vision_target], "lidar": [lidar_target]} + + +def _sample_joint_chunk( + host: OmniMoTCausalModel, + *, + plans: list[SequencePlan], + data: GenerationDataClean, + conditional_text: list[list[int]], + unconditional_text: list[list[int]] | None, + guidance: float, + seed: int, + num_steps: int, + shift: float, + vision_target: torch.Tensor, # [1,Cv,V*Tv,Hv,Wv] + lidar_target: torch.Tensor, # [1,Cl,Tl,Hl,Wl] + num_views: int, + frames_per_view: int, + vision_condition: int, + lidar_condition: int, + chunk: JointARChunk, + sampler_mode: Literal["rf", "distilled"] = "rf", + distilled_num_steps: int | None = None, + persistent_branches: list[tuple[list[list[int]], JointARCache]] | None = None, + retain_history: bool = True, +) -> None: + """Denoise the current RGB/LiDAR chunk and cache its clean K/V if persistent history is needed.""" + vision_frames = [frame for frame in chunk.vision_frames if frame >= vision_condition] + lidar_frames = [frame for frame in chunk.lidar_frames if frame >= lidar_condition] + if not vision_frames and not lidar_frames: + if persistent_branches is not None and retain_history: + current = _current_chunk_data(data, vision_target, lidar_target, num_views, chunk) + _commit_persistent_joint_chunk(host, plans, current, persistent_branches, chunk) + return + forward_data = ( + _prefix_data(data, vision_target=vision_target, lidar_target=lidar_target, num_views=num_views, chunk=chunk) + if persistent_branches is None + else _current_chunk_data(data, vision_target, lidar_target, num_views, chunk) + ) + assert forward_data.x0_tokens_vision is not None and forward_data.x0_tokens_lidar is not None + vision_indexes = [view * chunk.vision_prefix_end + frame for view in range(num_views) for frame in vision_frames] + final_vision_indexes = [view * frames_per_view + frame for view in range(num_views) for frame in vision_frames] + local_lidar_frames = lidar_frames + if persistent_branches is not None: + vision_indexes = [ + view * len(chunk.vision_frames) + i + for view in range(num_views) + for i, frame in enumerate(chunk.vision_frames) + if frame >= vision_condition + ] + local_lidar_frames = [i for i, frame in enumerate(chunk.lidar_frames) if frame >= lidar_condition] + vision_shape = (1, vision_target.shape[1], len(vision_indexes), *vision_target.shape[-2:]) + lidar_shape = (1, lidar_target.shape[1], len(lidar_frames), *lidar_target.shape[-2:]) + vision_size = math.prod(vision_shape) + lidar_size = math.prod(lidar_shape) + generator = torch.Generator(device=vision_target.device).manual_seed(seed + chunk.step) + initial_noise = torch.randn( + (1, vision_size + lidar_size), device=vision_target.device, dtype=torch.float32, generator=generator + ) # [1,D_vision+D_lidar] + if persistent_branches is None: + cond_pack, cond_memory = _build_replay_branch(host, plans, forward_data, conditional_text) + uncond_branch = ( + _build_replay_branch(host, plans, forward_data, unconditional_text) + if guidance != 1.0 and unconditional_text is not None + else None + ) + else: + branches = [ + _build_persistent_branch(host, plans, forward_data, text, cache, chunk) + for text, cache in persistent_branches + ] + cond_pack, cond_memory = branches[0] + uncond_branch = branches[1] if len(branches) > 1 else None + + def branch_velocity( + pack: PackedSequence, + memory: TeacherForcingMemoryState | JointChunkMemory, + noise: torch.Tensor, # [1,D_vision+D_lidar] + timestep: torch.Tensor, # [1,1] + ) -> torch.Tensor: # [1,D_vision+D_lidar] + assert pack.vision is not None and pack.lidar is not None + noisy_vision = forward_data.x0_tokens_vision[1].clone() # [1,Cv,V*Tv_forward,Hv,Wv] + noisy_lidar = forward_data.x0_tokens_lidar[1].clone() # [1,Cl,Tl_forward,Hl,Wl] + noisy_vision[:, :, vision_indexes] = noise[:, :vision_size].reshape(vision_shape) # [1,Cv,V*Tv_chunk,Hv,Wv] + noisy_lidar[:, :, local_lidar_frames] = noise[:, vision_size:].reshape(lidar_shape) # [1,Cl,Tl_chunk,Hl,Wl] + host._update_inference_pack_template( + packed_sequence=pack, + noise_x_vision=[forward_data.x0_tokens_vision[0], noisy_vision], # list of [1,Cv,V*Tv_forward,Hv,Wv] + noise_x_lidar=[forward_data.x0_tokens_lidar[0], noisy_lidar], # list of [1,Cl,Tl_forward,Hl,Wl] + noise_x_radar=None, + noise_x_action=None, + noise_x_sound=None, + timestep=timestep, # [1,1] + ) + prediction = host.denoise(data_batch_packed=pack, memory=memory) + vision_velocity = prediction["preds_vision"][1] # [1,Cv,V*Tv_forward,Hv,Wv] + lidar_velocity = prediction["preds_lidar"][1] # [1,Cl,Tl_forward,Hl,Wl] + if vision_velocity.shape != noisy_vision.shape or lidar_velocity.shape != noisy_lidar.shape: + raise ValueError("Joint model predictions must retain each input item's [1,C,T,H,W] shape") + # Grid-stream decoding retains its singleton batch dimension. Select time, + # preserving every channel before flattening for the shared diffusion solver. + return torch.cat( + [ + vision_velocity[:, :, vision_indexes].reshape(1, -1), + lidar_velocity[:, :, local_lidar_frames].reshape(1, -1), + ], + dim=1, + ) # [1,D_vision+D_lidar] + + def velocity(noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: # [1,D], [1,1] -> [1,D] + conditional = branch_velocity(cond_pack, cond_memory, noise, timestep) # [1,D] + if uncond_branch is None: + return conditional # [1,D] + unconditional = branch_velocity(*uncond_branch, noise, timestep) # [1,D] + return unconditional + guidance * (conditional - unconditional) # [1,D] + + # Match the joint SF-DMD RGB clock, including a final LiDAR-only chunk. + # Keep chunk-based noise seeds independent of the schedule's frame index. + schedule_frame_idx = chunk.vision_frames[0] if chunk.vision_frames else frames_per_view - 1 + denoised = host._run_ar_sampler( + velocity, + initial_noise, + sampler_mode=sampler_mode, + num_steps=num_steps, + shift=shift, + seed=seed, + sample_idx=chunk.step, + num_frames=frames_per_view, + distilled_num_steps=distilled_num_steps, + schedule_frame_idx=schedule_frame_idx, + ) # [1,D_vision+D_lidar] + if not torch.isfinite(denoised).all(): + raise FloatingPointError(f"Joint AR produced nonfinite latents in causal chunk {chunk.step}") + vision_target[:, :, final_vision_indexes] = denoised[:, :vision_size].reshape( + vision_shape + ) # [1,Cv,V*Tv_chunk,Hv,Wv] + lidar_target[:, :, lidar_frames] = denoised[:, vision_size:].reshape(lidar_shape) # [1,Cl,Tl_chunk,Hl,Wl] + if persistent_branches is not None and retain_history: + current = _current_chunk_data(data, vision_target, lidar_target, num_views, chunk) + # No later query reads the last chunk; do not run an unused clean forward. + _commit_persistent_joint_chunk(host, plans, current, persistent_branches, chunk) diff --git a/cosmos_framework/model/generator/utils/rolling_kv/__init__.py b/cosmos_framework/model/generator/utils/rolling_kv/__init__.py new file mode 100644 index 000000000..d854f8f9e --- /dev/null +++ b/cosmos_framework/model/generator/utils/rolling_kv/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bounded rolling KV, prompt scheduling, and sink-position helpers.""" diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_prompt.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_prompt.py new file mode 100644 index 000000000..0303c3051 --- /dev/null +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_prompt.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bounded segmented prompt visibility and immutable initial text-prefix K/V.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from itertools import accumulate + +import torch + +from cosmos_framework.data.generator.sequence_packing import PackedSequence +from cosmos_framework.data.generator.sequence_packing.runtime import SequencePack + + +@dataclass(frozen=True) +class RollingPromptSegment: + """A half-open source-frame interval and its unframed tokenizer output.""" + + start_frame: int + end_frame: int + token_ids: tuple[int, ...] + + +@dataclass(frozen=True) +class RollingPromptDocument: + """One independently encoded caption in the current UND stream.""" + + view_id: int + start_frame: int + end_frame: int + token_ids: tuple[int, ...] + is_sink: bool = False + # Caption identity keeps its original interval. Only direct media visibility + # expands when neighboring caption context is requested. + visible_start_frame: int | None = None + visible_end_frame: int | None = None + + +@dataclass(frozen=True) +class RollingPromptLayout: + """Global post-Ulysses text coordinates, independent of local head shards.""" + + documents: tuple[RollingPromptDocument, ...] + fps: float + sink_tokens: int + num_views: int + # These are global UND token indices; Ulysses shards heads at attention time. + sink_capture_indexes: torch.Tensor # [V*K] + sink_read_indexes: torch.Tensor # [V*K] + sink_cache_indexes: torch.Tensor # [V*K] + + @property + def lengths(self) -> tuple[int, ...]: + return tuple(len(doc.token_ids) for doc in self.documents) + + def install_document_offsets(self, input_pack: SequencePack, num_und: int) -> None: + """Keep time segments and cameras isolated, including CP's pad document.""" + offsets = list(accumulate(self.lengths, initial=0)) + if offsets[-1] > num_und: + raise ValueError("Segmented prompt document lengths exceed the global UND stream.") + if offsets[-1] < num_und: + offsets.append(num_und) + input_pack["_caption_seq_offsets"] = torch.tensor( + offsets, dtype=torch.int32, device=self.sink_read_indexes.device + ) # [N_docs+1 or N_docs+2] + input_pack["max_caption_len"] = max(b - a for a, b in zip(offsets, offsets[1:])) + + def time_bounds(self, kv_length: int) -> tuple[torch.Tensor, torch.Tensor]: # returns 2*[KV] + starts = torch.full((kv_length,), -torch.inf, device=self.sink_read_indexes.device) # [KV] + ends = torch.full((kv_length,), torch.inf, device=starts.device) # [KV] + offset = 0 + for doc in self.documents: + stop = offset + len(doc.token_ids) + if not doc.is_sink: + first = doc.start_frame if doc.visible_start_frame is None else doc.visible_start_frame + last = doc.end_frame if doc.visible_end_frame is None else doc.visible_end_frame + starts[offset:stop] = first / self.fps # [N_doc] + ends[offset:stop] = last / self.fps # [N_doc] + offset = stop + return starts, ends # 2*[KV] + + +@dataclass(frozen=True) +class RollingPromptPack: + """Current documents, grouped by camera for the existing base mask builder.""" + + documents: tuple[RollingPromptDocument, ...] + text_tokens: list[list[int]] + fps: float + sink_tokens: int + num_views: int + fixed_text_extent: int + + def prepare( + self, + packed: PackedSequence, + *, + position_mode: str = "fixed", + temporal_units_per_second: float = 1.0, + chunk_start_seconds: float = 0.0, + ) -> RollingPromptLayout: + """Position independent documents without changing their visibility or sinks.""" + if position_mode not in ("fixed", "segment_start", "chunk_start"): + raise ValueError("Rolling text position mode must be fixed, segment_start or chunk_start.") + if ( + not math.isfinite(temporal_units_per_second) + or temporal_units_per_second <= 0 + or not math.isfinite(chunk_start_seconds) + or chunk_start_seconds < 0 + ): + raise ValueError("Rolling text positioning requires a positive temporal scale and nonnegative chunk time.") + lengths = [sum(len(d.token_ids) for d in self.documents if d.view_id == view) for view in range(self.num_views)] + if packed.text_caption_lens != [lengths]: + raise ValueError("Prompt framing disagrees with native caption packing.") + if position_mode != "fixed" and not packed.position_ids.is_floating_point(): + packed.position_ids = packed.position_ids.float() # [3,S] + offset = 0 + for doc in self.documents: + length = len(doc.token_ids) + # The native per-view pack rewinds each view to zero. Rewind each + # independent time segment as well. Optional temporal translation + # follows the absolute rollout clock; sink K/V stays at its origin. + positions = torch.arange( + length, dtype=packed.position_ids.dtype, device=packed.position_ids.device + ) # [N_doc] + packed.position_ids[:, offset : offset + length] = positions[None, :] # [3,N_doc] + if not doc.is_sink and position_mode != "fixed": + seconds = doc.start_frame / self.fps if position_mode == "segment_start" else chunk_start_seconds + packed.position_ids[0, offset : offset + length] += seconds * temporal_units_per_second # [N_doc] + offset += length + # Length changes must never move a media key relative to previously + # committed keys. The native text cursor advances past the longest view. + delta = self.fixed_text_extent - max(lengths) + for modality in (packed.vision, packed.lidar): + if modality is not None: + packed.position_ids[0, modality.sequence_indexes] += delta # [S_modality] + return self.layout(packed.position_ids.device) + + def layout(self, device: torch.device) -> RollingPromptLayout: + """Build global attention indices once per current pack, not once per layer.""" + first_by_view: dict[int, int] = {} + read_indexes: list[int] = [] + cache_indexes: list[int] = [] + offset = 0 + for doc in self.documents: + if doc.is_sink: + if doc.view_id in first_by_view or len(doc.token_ids) != self.sink_tokens: + raise ValueError("Segmented prompting requires exactly one fixed-size sink document per camera.") + first_by_view[doc.view_id] = offset + read_indexes.extend(range(offset, offset + self.sink_tokens)) + cache_indexes.extend(range(doc.view_id * self.sink_tokens, (doc.view_id + 1) * self.sink_tokens)) + offset += len(doc.token_ids) + if self.sink_tokens and set(first_by_view) != set(range(self.num_views)): + raise ValueError("Segmented prompting requires a retained sink for every camera.") + capture_indexes = [ + index + for view in range(self.num_views) + for index in range(first_by_view.get(view, 0), first_by_view.get(view, 0) + self.sink_tokens) + ] + return RollingPromptLayout( + documents=self.documents, + fps=self.fps, + sink_tokens=self.sink_tokens, + num_views=self.num_views, + sink_capture_indexes=torch.tensor(capture_indexes, device=device, dtype=torch.long), # [V*K] + sink_read_indexes=torch.tensor(read_indexes, device=device, dtype=torch.long), # [V*K] + sink_cache_indexes=torch.tensor(cache_indexes, device=device, dtype=torch.long), # [V*K] + ) + + +class RollingPromptSchedule: + """Read captions on the original finite video clock, without repeating them.""" + + fps: float + num_frames: int + views: tuple[tuple[RollingPromptSegment, ...], ...] + sink_tokens: int + special_tokens: dict[str, int] + prefixes: tuple[tuple[int, ...], ...] + fixed_text_extent: int + context_before: int + context_after: int + + def __init__( + self, + *, + fps: float, + num_frames: int, + views: tuple[tuple[RollingPromptSegment, ...], ...], + sink_tokens: int, + special_tokens: dict[str, int], + context_before: int = 0, + context_after: int = 0, + media_text_extent: int | None = None, + ) -> None: + if ( + not math.isfinite(fps) + or fps <= 0 + or isinstance(num_frames, bool) + or not isinstance(num_frames, int) + or num_frames < 1 + or not views + ): + raise ValueError("Segmented prompting requires a positive FPS, frame count and camera count.") + if isinstance(sink_tokens, bool) or not isinstance(sink_tokens, int) or sink_tokens < 0: + raise ValueError("Rolling text sink size must be a nonnegative integer.") + if any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in (context_before, context_after) + ): + raise ValueError("Neighboring caption counts must be nonnegative integers.") + self.fps = fps + self.num_frames = num_frames + self.views = views + self.sink_tokens = sink_tokens + self.special_tokens = special_tokens + self.context_before = context_before + self.context_after = context_after + for segments in views: + cursor = 0 + for segment in segments: + if segment.start_frame != cursor or not cursor < segment.end_frame <= num_frames: + raise ValueError("Rolling captions must cover the video without gaps or overlaps.") + cursor = segment.end_frame + if cursor != num_frames: + raise ValueError("Rolling captions do not cover the video.") + # Empty captions are valid for unconditional sampling. A short prompt + # retains only the available prefix; the full current document remains. + self.sink_tokens = min(sink_tokens, min(len(self._frame(segments[0].token_ids)) for segments in views)) + self.prefixes = tuple(self._frame(segments[0].token_ids)[: self.sink_tokens] for segments in views) + longest = max(len(self._frame(segment.token_ids)) for segments in views for segment in segments) + if media_text_extent is not None and ( + isinstance(media_text_extent, bool) or not isinstance(media_text_extent, int) or media_text_extent < longest + ): + raise ValueError("Media text extent must cover the longest independently framed caption.") + self.fixed_text_extent = longest if media_text_extent is None else media_text_extent + + def _frame(self, tokens: tuple[int, ...]) -> tuple[int, ...]: + prefix = (self.special_tokens["bos_token_id"],) if "bos_token_id" in self.special_tokens else () + return prefix + tokens + (self.special_tokens["eos_token_id"], self.special_tokens["start_of_generation"]) + + def for_times(self, timestamps: tuple[float, ...]) -> RollingPromptPack: + if not timestamps or any(not math.isfinite(t) or t < 0 for t in timestamps): + raise ValueError("Segmented prompting requires finite nonnegative query timestamps.") + frames = sorted({math.floor(t * self.fps + 1e-5) for t in timestamps}) + if frames[-1] >= self.num_frames: + raise ValueError("Rolling query exceeds the video horizon; prompts are never repeated.") + documents: list[RollingPromptDocument] = [] + packed_tokens: list[list[int]] = [] + for view, segments in enumerate(self.views): + selected: set[int] = set() + for frame in frames: + for index, segment in enumerate(segments): + if segment.start_frame <= frame < segment.end_frame: + selected.update( + range( + max(0, index - self.context_before), min(len(segments), index + self.context_after + 1) + ) + ) + break + view_tokens: list[int] = [] + if self.sink_tokens: + # Keep the original prefix in its own causal text document. + # Current captions retain every token, including a different + # opening word after a switch. Media can always read this sink. + documents.append(RollingPromptDocument(view, 0, self.num_frames, self.prefixes[view], is_sink=True)) + view_tokens.extend(self.prefixes[view]) + for index in sorted(selected): + segment = segments[index] + tokens = self._frame(segment.token_ids) + # A document d is visible to query-segment q exactly when + # q-before <= d <= q+after. The inverse bounds are therefore + # [d-after, d+before], clipped to this video's actual captions. + first_index = max(0, index - self.context_after) + last_index = min(len(segments) - 1, index + self.context_before) + documents.append( + RollingPromptDocument( + view_id=view, + start_frame=segment.start_frame, + end_frame=segment.end_frame, + token_ids=tokens, + visible_start_frame=segments[first_index].start_frame, + visible_end_frame=segments[last_index].end_frame, + ) + ) + view_tokens.extend(tokens) + # Native packing adds precisely these outside tokens. Internal + # document delimiters stay explicit; the replay hook isolates them. + prefix_count = int("bos_token_id" in self.special_tokens) + packed_tokens.append(view_tokens[prefix_count:-2]) + return RollingPromptPack( + documents=tuple(documents), + text_tokens=packed_tokens, + fps=self.fps, + sink_tokens=self.sink_tokens, + num_views=len(self.views), + fixed_text_extent=self.fixed_text_extent, + ) + + +class RollingTextSinkLayer: + """One layer's original RoPE K, normalized K and V; never overwritten.""" + + cached: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None # each [V*K,H_KV_local,D] + reads: int + + def __init__(self) -> None: + self.cached = None + self.reads = 0 + + def apply( + self, + k_rope: torch.Tensor, # [UND,H_KV_local,D] + k_normalized: torch.Tensor, # [UND,H_KV_local,D] + value: torch.Tensor, # [UND,H_KV_local,D] + layout: RollingPromptLayout, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # each [UND,H_KV_local,D] + if layout.sink_tokens == 0: + return k_rope, k_normalized, value + if torch.is_grad_enabled(): + raise ValueError("Retained text K/V is inference-only.") + if self.cached is None: + initial_views = {doc.view_id for doc in layout.documents if not doc.is_sink and doc.start_frame == 0} + if initial_views != set(range(layout.num_views)): + raise ValueError("Segmented prompting must capture its text sink at the start of the rollout.") + self.cached = tuple( + tensor.index_select(0, layout.sink_capture_indexes).detach().clone() + for tensor in (k_rope, k_normalized, value) + ) # 3*[V*K,H_KV_local,D] + result: list[torch.Tensor] = [] # each [UND,H_KV_local,D] + for tensor, saved in zip((k_rope, k_normalized, value), self.cached, strict=True): + if saved.shape != (layout.num_views * layout.sink_tokens, *tensor.shape[1:]): + raise ValueError("Rolling text sink camera/head geometry changed.") + restored = tensor.clone() # [UND,H_KV_local,D] + restored.index_copy_( + 0, layout.sink_read_indexes, saved.index_select(0, layout.sink_cache_indexes) + ) # [UND,H_KV_local,D] + result.append(restored) + self.reads += 1 + return result[0], result[1], result[2] diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py new file mode 100644 index 000000000..cea406326 --- /dev/null +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_replay.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Bounded clean replay K/V with pinned sink chunks and absolute sensor clocks.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, fields, replace +from typing import Literal + +import torch + +from cosmos_framework.model.generator.mot.flex_attention import CaptionMaskItem, SensorMaskItem +from cosmos_framework.configs.base.defaults.replay_attention import TeacherForcingReplayPolicyConfig +from cosmos_framework.model.generator.mot.causal_flex_attention import ( + _ROLE_CLEAN_TARGET, + _ROLE_CURRENT_TARGET, + FlexQueryMetadata, + TeacherForcingFlexMetadata, + build_teacher_forcing_multiview_flex_metadata, +) +from cosmos_framework.model.generator.utils.kv_cache import KVToStore, MultiviewARMemoryState +from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( + RollingPromptLayout, + RollingTextSinkLayer, +) +from cosmos_framework.model.generator.utils.rolling_kv.rolling_sink_rope import ( + MediaSinkKeyRotation, + MediaSinkMemoryState, +) + + +@dataclass +class RollingReplayRequest: + """One current chunk plus the branch's previously committed clean memory. + + Original condition masks survive a clean refresh, whose embedding inputs + mark all tokens as clean. The network records the current query metadata so + the cache can commit exactly the same roles and absolute positions as K/V. + """ + + step: int + frame_starts: tuple[int, ...] + frames_per_chunk: int + pass_kind: Literal["clean", "noisy"] + condition_masks: list[torch.Tensor] # each [V*T,1,1] + history: FlexQueryMetadata | None = None + current: FlexQueryMetadata | None = None + prompt_layout: RollingPromptLayout | None = None + + +def build_rolling_replay_metadata( + *, + request: RollingReplayRequest, + seq_len: int, + full_q_offsets: torch.Tensor, # [B+1 or B+2] + sensor_mask_items: Sequence[Sequence[SensorMaskItem]], + caption_mask_items: Sequence[Sequence[CaptionMaskItem]] | None, + device: torch.device, + num_und: int, + causal_offsets: torch.Tensor, # [B+1 or B+2] + policy: TeacherForcingReplayPolicyConfig, +) -> TeacherForcingFlexMetadata: + """Retain replay visibility while moving current and cached frames globally.""" + if len(sensor_mask_items) != 1 or len(sensor_mask_items[0]) != len(request.frame_starts): + raise ValueError("Rolling replay requires one sample and one absolute start per sensor item.") + if request.step < 0 or request.frames_per_chunk < 1 or any(start < 0 for start in request.frame_starts): + raise ValueError("Rolling replay requires nonnegative positions and a positive chunk size.") + if request.pass_kind not in ("clean", "noisy"): + raise ValueError("Rolling replay pass_kind must be clean or noisy.") + # Build only the current pack. The legacy noisy builder expects an entire + # prefix when constructing its suffix; rolling memory already owns that + # metadata, so assign the current noisy role after this clean-only build. + current = build_teacher_forcing_multiview_flex_metadata( + seq_len=seq_len, + full_q_offsets=full_q_offsets, + sensor_mask_items=sensor_mask_items, + caption_mask_items=caption_mask_items, + device=device, + num_und=num_und, + causal_offsets=causal_offsets, + frames_per_chunk=request.frames_per_chunk, + pass_kind="clean", + teacher_forcing_replay_policy=policy, + ) + frame_id = current.frame_id.clone() # [UND+GEN] + timestamp = current.timestamp.clone() # [UND+GEN] + causal_step_id = current.causal_step_id.clone() # [UND+GEN] + offset = num_und + for item, start in zip(sensor_mask_items[0], request.frame_starts, strict=True): + count = item.latent_t * item.spatial_tokens + frame_id[offset : offset + count] += start # [S_item] + timestamp[offset : offset + count] = ( + frame_id[offset : offset + count].to(torch.float32) * item.seconds_per_frame + ) # [S_item] + # Do not divide large float32 timestamps to recover a chunk index. + # The shared clock assigned this complete chunk before token packing. + causal_step_id[offset : offset + count] = request.step # [S_item] + offset += count + roles = current.token_role_id # [UND+GEN] + if request.pass_kind == "noisy": + roles = torch.where(roles == _ROLE_CLEAN_TARGET, _ROLE_CURRENT_TARGET, roles) # [UND+GEN] + current = replace( + current, frame_id=frame_id, timestamp=timestamp, causal_step_id=causal_step_id, token_role_id=roles + ) + query = FlexQueryMetadata( + **{field.name: getattr(current, field.name)[num_und:] for field in fields(FlexQueryMetadata)} # each [GEN] + ) + request.current = query + if request.history is None: + result = replace(current, query=query) + else: + combined = { + field.name: torch.cat((getattr(current, field.name), getattr(request.history, field.name))) + for field in fields(FlexQueryMetadata) + } # each [UND+GEN+M] + result = TeacherForcingFlexMetadata( + **combined, num_und=num_und, query=query, teacher_forcing_replay_policy=policy + ) + if request.prompt_layout is not None: + result = replace(result, caption_time_bounds=request.prompt_layout.time_bounds(result.seq_len)) + return result + + +class _RollingReplayWriter(MediaSinkMemoryState): + """Read committed history while staging one clean chunk outside the live ring.""" + + owner: RollingReplayCache + step: int + token_count: int + staged: MultiviewARMemoryState + written_layers: set[int] + + def __init__(self, owner: RollingReplayCache, *, step: int, token_count: int, device: torch.device) -> None: + indexes = torch.arange(token_count, device=device, dtype=torch.long) # [S_current] + super().__init__( + num_layers=len(owner.cache), + memory_seq_len=owner.capacity, + cache=owner.cache, + write_indexes=indexes, + text_sink_layers=owner.text_sink_layers, + prompt_layout=owner.prompt_layout, + ) + self.owner, self.step, self.token_count = owner, step, token_count + self.sink_key_rotation = owner.sink_key_rotation + self.staged = MultiviewARMemoryState( + num_layers=len(owner.cache), + memory_seq_len=owner.tokens_per_chunk, + cache=owner._staging_cache, + write_indexes=indexes, + ) + self.written_layers = set() + + def write_for_layer(self, layer_idx: int, kv_to_store: KVToStore) -> None: + if self.owner._pending is not self: + raise ValueError("Only the active rolling refresh can stage clean K/V.") + self.staged.write_for_layer(layer_idx, kv_to_store) + self.written_layers.add(layer_idx) + + +class RollingReplayCache: + """One CFG branch's fixed-capacity ring of complete clean sensor chunks. + + ``cache_chunks`` counts completed history chunks, including ``sink_chunks``; + the current query pack is separate. Chunk zero is the training singleton. + Every slot reserves the maximum synchronized chunk token count, so short + chunks need only mark their unused slots as padding. Cache tensors remain + allocated and their addresses remain stable after the first clean refresh. + Refreshes use one reusable staging chunk per layer; publication temporarily + backs up the overwritten region so a failed commit can restore the ring. + """ + + cache_chunks: int + sink_chunks: int + tokens_per_chunk: int + capacity: int + cache: list[tuple[torch.Tensor, torch.Tensor] | None] # per layer, each [1,M,H_kv,D] + metadata: FlexQueryMetadata | None + slot_steps: list[int | None] + next_step: int + _staging_cache: list[tuple[torch.Tensor, torch.Tensor] | None] # per layer, each [1,S_slot,H_kv,D] + _pending: _RollingReplayWriter | None + text_sink_layers: list[RollingTextSinkLayer] | None + prompt_layout: RollingPromptLayout | None + initial_sink_token_count: int + sink_key_rotation: MediaSinkKeyRotation | None + + def __init__(self, *, num_layers: int, cache_chunks: int, sink_chunks: int, tokens_per_chunk: int) -> None: + if any( + isinstance(value, bool) or not isinstance(value, int) + for value in (num_layers, cache_chunks, sink_chunks, tokens_per_chunk) + ): + raise ValueError("Rolling replay cache sizes must be integers.") + if num_layers < 1 or tokens_per_chunk < 1 or not 0 <= sink_chunks < cache_chunks: + raise ValueError("Rolling replay requires positive capacity and 0 <= sink_chunks < cache_chunks.") + self.cache_chunks = cache_chunks + self.sink_chunks = sink_chunks + self.tokens_per_chunk = tokens_per_chunk + self.capacity = cache_chunks * tokens_per_chunk + self.cache = [None] * num_layers + self.metadata = None + self.slot_steps = [None] * cache_chunks + self.next_step = 0 + self._staging_cache = [None] * num_layers + self._pending = None + self.text_sink_layers = None + self.prompt_layout = None + self.initial_sink_token_count = 0 + self.sink_key_rotation = None + + def _slot(self, step: int) -> int: + if step != self.next_step: + raise ValueError(f"Expected the next rolling chunk {self.next_step}, got {step}.") + if step < self.sink_chunks: + return step + return self.sink_chunks + (step - self.sink_chunks) % (self.cache_chunks - self.sink_chunks) + + def read(self) -> MultiviewARMemoryState: + """Denoising cannot write K/V, even if the network returns intermediate K/V.""" + state_type = MultiviewARMemoryState if self.sink_key_rotation is None else MediaSinkMemoryState + state = state_type( + num_layers=len(self.cache), + memory_seq_len=self.capacity, + cache=self.cache, + text_sink_layers=self.text_sink_layers, + prompt_layout=self.prompt_layout, + ) + if isinstance(state, MediaSinkMemoryState): + state.sink_key_rotation = self.sink_key_rotation + return state + + def writer(self, *, step: int, token_count: int, device: torch.device) -> MultiviewARMemoryState: + """Stage real current GEN tokens after each layer has read committed history.""" + self._slot(step) + if not 0 < token_count <= self.tokens_per_chunk: + raise ValueError(f"Chunk has {token_count} tokens, exceeding slot capacity {self.tokens_per_chunk}.") + self._pending = _RollingReplayWriter(self, step=step, token_count=token_count, device=device) + return self._pending + + def commit(self, *, step: int, token_count: int, request: RollingReplayRequest) -> None: + """Publish clean K/V and metadata only after every layer finished this refresh.""" + slot = self._slot(step) + if request.pass_kind != "clean" or request.step != step or request.current is None: + raise ValueError("Only a completed clean replay can be committed to rolling memory.") + if not 0 < token_count <= self.tokens_per_chunk or token_count > request.current.sample_id.numel(): + raise ValueError("Rolling replay commit has invalid current token geometry.") + pending = self._pending + if pending is None or pending.step != step or pending.token_count != token_count: + raise ValueError("Rolling replay commit must match the active refresh.") + if pending.written_layers != set(range(len(self.cache))): + raise ValueError("Every transformer layer must write clean K/V before committing a chunk.") + if self.metadata is None: + metadata = FlexQueryMetadata( + **{ + field.name: getattr(request.current, field.name).new_full( + (self.capacity,), False if getattr(request.current, field.name).dtype == torch.bool else -1 + ) + for field in fields(FlexQueryMetadata) + } # each [M] + ) + else: + metadata = FlexQueryMetadata( + **{field.name: getattr(self.metadata, field.name).clone() for field in fields(FlexQueryMetadata)} + ) # each [M] + start = slot * self.tokens_per_chunk + for field in fields(FlexQueryMetadata): + destination = getattr(metadata, field.name) # [M] + destination[start : start + self.tokens_per_chunk].fill_( + False if destination.dtype == torch.bool else -1 + ) # [S_slot] + source = getattr(request.current, field.name)[:token_count] # [S_current] + destination[start : start + token_count].copy_(source) # [S_current] + # Clean history is never a second current-noise stream in a later pass. + metadata.is_noisy[start : start + token_count] = False # [S_current] + + # Finish allocations and validation before touching any committed K/V. + targets: list[tuple[torch.Tensor, torch.Tensor]] = [] # per layer, each [1,M,H_kv,D] + backups: list[tuple[torch.Tensor, torch.Tensor] | None] = [] # each [1,S_current,H_kv,D] + for cached, staged in zip(self.cache, pending.staged.cache, strict=True): + assert staged is not None + staged_k, staged_v = staged # each [1,S_slot,H_kv,D] + if cached is None: + shape = (staged_k.shape[0], self.capacity, *staged_k.shape[2:]) + targets.append((staged_k.new_zeros(shape), staged_v.new_zeros(shape))) # each [1,M,H_kv,D] + backups.append(None) + else: + for target, source in zip(cached, staged, strict=True): + if ( + target.shape[:1] + target.shape[2:] != source.shape[:1] + source.shape[2:] + or target.dtype != source.dtype + or target.device != source.device + ): + raise ValueError("Rolling clean K/V geometry, dtype and device must remain unchanged.") + targets.append(cached) + backups.append( + ( + cached[0][:, start : start + token_count].clone(), + cached[1][:, start : start + token_count].clone(), + ) + ) # each [1,S_current,H_kv,D] + try: + for target, staged in zip(targets, pending.staged.cache, strict=True): + assert staged is not None + target[0][:, start : start + token_count].copy_(staged[0][:, :token_count]) # [1,S_current,H_kv,D] + target[1][:, start : start + token_count].copy_(staged[1][:, :token_count]) # [1,S_current,H_kv,D] + except Exception: + for target, backup in zip(targets, backups, strict=True): + if backup is not None: + target[0][:, start : start + token_count].copy_(backup[0]) # [1,S_current,H_kv,D] + target[1][:, start : start + token_count].copy_(backup[1]) # [1,S_current,H_kv,D] + raise + self.cache[:] = targets + self.metadata = metadata + self.slot_steps[slot] = step + if step == 0 and self.sink_chunks: + self.initial_sink_token_count = token_count + self.next_step += 1 + self._pending = None + + @property + def retained_steps(self) -> tuple[int, ...]: + """Return retained absolute chunk IDs without recording the entire rollout.""" + return tuple(sorted(step for step in self.slot_steps if step is not None)) diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_sink_rope.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_sink_rope.py new file mode 100644 index 000000000..b26c138b2 --- /dev/null +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_sink_rope.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Optional read-time temporal rephasing of an immutable frame-zero media sink.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace + +import torch + +from cosmos_framework.model.generator.reasoner.qwen3_vl.qwen3_vl import Qwen3VLTextRotaryEmbedding +from cosmos_framework.model.generator.utils.kv_cache import MultiviewARMemoryState, MultiviewARMemoryValue + + +@dataclass(frozen=True) +class MediaSinkKeyRotation: + """Unit rotation shared by layers; original cache tensors are never changed.""" + + token_count: int + time_offset_seconds: float + temporal_offset: float + cos: torch.Tensor # [1,1,1,D] + sin: torch.Tensor # [1,1,1,D] + + def apply(self, keys: torch.Tensor) -> torch.Tensor: # keys, returns: [1,M,H_kv,D] + if keys.ndim != 4 or keys.shape[-1] != self.cos.shape[-1] or not 0 < self.token_count <= keys.shape[1]: + raise ValueError("Media sink rotation does not match the cached key geometry.") + original = keys[:, : self.token_count].float() # [1,S_sink,H_kv,D] + half = original.shape[-1] // 2 + rotated = torch.cat((-original[..., half:], original[..., :half]), dim=-1) # [1,S_sink,H_kv,D] + shifted_ = original * self.cos + rotated * self.sin # [1,S_sink,H_kv,D] + result = keys.clone() # [1,M,H_kv,D] + result[:, : self.token_count] = shifted_.to(dtype=keys.dtype) # [1,S_sink,H_kv,D] + return result # [1,M,H_kv,D] + + +@torch.no_grad() +def prepare_media_sink_key_rotation( + rotary_embedding: Qwen3VLTextRotaryEmbedding, + *, + token_count: int, + chunk_start_seconds: float, + age_cap_seconds: float, + temporal_units_per_second: float, +) -> MediaSinkKeyRotation | None: + """Translate only temporal phase using the loaded native interleaved MRoPE. + + The cap is relative to the earliest query timestamp in a synchronized + chunk. A later query in that same chunk can be older by its time offset. + Static frequencies are required: dynamic RoPE could mutate the frequency + basis of the already-rotated stored keys. Attention scaling is removed + from the delta factors because the keys already contain that scaling. + """ + if not isinstance(rotary_embedding, Qwen3VLTextRotaryEmbedding) or rotary_embedding.rope_type != "default": + raise ValueError("Media sink age capping requires static Qwen3-VL rotary embeddings.") + if isinstance(token_count, bool) or not isinstance(token_count, int) or token_count < 0: + raise ValueError("Media sink token count must be a nonnegative integer.") + if ( + not math.isfinite(chunk_start_seconds) + or chunk_start_seconds < 0 + or isinstance(age_cap_seconds, bool) + or not math.isfinite(age_cap_seconds) + or age_cap_seconds <= 0 + or not math.isfinite(temporal_units_per_second) + or temporal_units_per_second <= 0 + ): + raise ValueError("Media sink age capping requires finite positive age and temporal scale.") + scaling = float(rotary_embedding.attention_scaling) + if not math.isfinite(scaling) or scaling <= 0: + raise ValueError("The native rotary attention scaling must be finite and positive.") + time_offset = max(0.0, chunk_start_seconds - age_cap_seconds) + if token_count == 0 or time_offset == 0: + return None + temporal_offset = time_offset * temporal_units_per_second + device = rotary_embedding.inv_freq.device + delta = torch.tensor([[[temporal_offset]], [[0.0]], [[0.0]]], device=device, dtype=torch.float32) # [3,1,1] + sentinel = torch.empty((), device=device, dtype=torch.float32) # [] + cos, sin = rotary_embedding(sentinel, delta) # each [1,1,D] + unit_cos = (cos / scaling).unsqueeze(2) # [1,1,1,D] + unit_sin = (sin / scaling).unsqueeze(2) # [1,1,1,D] + return MediaSinkKeyRotation(token_count, time_offset, temporal_offset, unit_cos, unit_sin) + + +class MediaSinkMemoryState(MultiviewARMemoryState): + """Serve a temporary sink-key phase while inheriting the clean write path.""" + + sink_key_rotation: MediaSinkKeyRotation | None = None + + def read_for_layer(self, layer_idx: int) -> MultiviewARMemoryValue: + value = super().read_for_layer(layer_idx) + if self.sink_key_rotation is None or value.cached_gen_k is None: + return value + keys_ = self.sink_key_rotation.apply(value.cached_gen_k) # [1,M,H_kv,D] + return replace(value, cached_gen_k=keys_) diff --git a/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py new file mode 100644 index 000000000..0eee69b0f --- /dev/null +++ b/cosmos_framework/model/generator/utils/rolling_kv/rolling_transfer_ar.py @@ -0,0 +1,635 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Opt-in incremental RGB and joint RGB/LiDAR inference with bounded replay KV.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterator +from dataclasses import replace +from typing import TYPE_CHECKING, Literal + +import torch + +from cosmos_framework.utils import log +from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean +from cosmos_framework.data.generator.sequence_packing import PackedSequence, SequencePlan +from cosmos_framework.model.generator.teacher_forcing import mark_modality_as_clean_condition +from cosmos_framework.model.generator.utils.multiview_ar import JointARChunk, joint_ar_chunks +from cosmos_framework.model.generator.utils.rolling_kv.rolling_prompt import ( + RollingPromptPack, + RollingPromptSchedule, + RollingPromptSegment, + RollingTextSinkLayer, +) +from cosmos_framework.model.generator.utils.rolling_kv.rolling_replay import RollingReplayCache, RollingReplayRequest +from cosmos_framework.model.generator.utils.rolling_kv.rolling_sink_rope import prepare_media_sink_key_rotation + +if TYPE_CHECKING: + from cosmos_framework.model.generator.omni_mot_causal_model import OmniMoTCausalModel + + +def _camera_slice( + latent: torch.Tensor, + num_views: int, + start: int, + stop: int, # latent: [1,C,V*T,H,W] +) -> torch.Tensor: # [1,C,V*(stop-start),H,W] + batch, channels, total, height, width = latent.shape + grid = latent.reshape(batch, channels, num_views, total // num_views, height, width) # [1,C,V,T,H,W] + return grid[:, :, :, start:stop].reshape( + batch, channels, num_views * (stop - start), height, width + ) # [1,C,V*Tc,H,W] + + +def _initial_target( + target: torch.Tensor, + num_views: int, + condition_count: int, + device: torch.device, # target: [1,C,V*T,H,W] +) -> torch.Tensor: # [1,C,V*T,H,W] + """Read only explicitly observed targets, never future ground truth.""" + result = torch.zeros(target.shape, device=device, dtype=torch.float32) # [1,C,V*T,H,W] + if condition_count: + frames = target.shape[2] // num_views + indexes = [view * frames for view in range(num_views)] + result[:, :, indexes] = target[:, :, indexes].to(device=device, dtype=torch.float32) # [1,C,V,H,W] + return result + + +def _chunk_data( + data: GenerationDataClean, + vision: torch.Tensor, # [1,Cv,V*Tv,Hv,Wv] + lidar: torch.Tensor | None, # [1,Cl,Tl,Hl,Wl] or None + num_views: int, + chunk: JointARChunk, +) -> GenerationDataClean: + assert data.x0_tokens_vision is not None + start, stop = chunk.vision_frames[0], chunk.vision_frames[-1] + 1 + lidar_items = None + if lidar is not None: + assert data.x0_tokens_lidar is not None + lidar_start, lidar_stop = chunk.lidar_frames[0], chunk.lidar_frames[-1] + 1 + lidar_items = [ + data.x0_tokens_lidar[0][:, :, lidar_start:lidar_stop], + lidar[:, :, lidar_start:lidar_stop], + ] # each [1,Cl,Tlc,Hl,Wl] + return replace( + data, + raw_state_vision=None, + raw_state_lidar=None, + x0_tokens_vision=[ + _camera_slice(data.x0_tokens_vision[0], num_views, start, stop), + _camera_slice(vision, num_views, start, stop), + ], # each [1,Cv,V*Tvc,Hv,Wv] + x0_tokens_lidar=lidar_items, + ) + + +def shift_rolling_positions( + pack: PackedSequence, + *, + vision_start: int, + lidar_start: int, + enable_fps_modulation: bool, + base_fps: float, + vision_temporal_compression: int, + temporal_scale: float = 1.0, +) -> None: + """Offset sensor RoPE only; caption positions and spatial grids stay fixed. + + The caller requires shared control/target vision positions, so each sensor + starts at the caption boundary regardless of the local chunk length. Both + sensor RoPE grids use the vision base temporal compression, as in training. + """ + for modality, start in ((pack.vision, vision_start), (pack.lidar, lidar_start)): + if modality is None: + continue + if not modality.seconds_per_frame or len(set(modality.seconds_per_frame)) != 1: + raise ValueError("Rolling replay requires one sensor period per modality.") + scale = modality.seconds_per_frame[0] * base_fps / vision_temporal_compression if enable_fps_modulation else 1 + if temporal_scale == 1.0: + pack.position_ids[0, modality.sequence_indexes] += start * scale # [S_modality] + continue + if not pack.position_ids.is_floating_point(): + pack.position_ids = pack.position_ids.float() # [3,S] + local_time = pack.position_ids[0, modality.sequence_indexes] # [S_modality] + origin = local_time.min() # [] + pack.position_ids[0, modality.sequence_indexes] = ( + origin + (local_time - origin + start * scale) * temporal_scale + ) # [S_modality] + + +def _build_branch_pack( + host: OmniMoTCausalModel, + *, + plan: SequencePlan, + data: GenerationDataClean, + text: list[list[int]], + chunk: JointARChunk, + cache: RollingReplayCache, + clean: bool, + prompt: RollingPromptPack | None = None, +) -> tuple[PackedSequence, RollingReplayRequest]: + pack = host._pack_input_sequence([plan], text, data, torch.zeros(1)) # timestep: [1] + expert = host.config.diffusion_expert_config + shift_rolling_positions( + pack, + vision_start=chunk.vision_frames[0], + lidar_start=chunk.lidar_frames[0] if plan.has_lidar else 0, + enable_fps_modulation=expert.enable_fps_modulation, + base_fps=float(expert.base_fps), + vision_temporal_compression=host.tokenizer_vision_gen.temporal_compression_factor, + temporal_scale=getattr(host.config, "rolling_media_rope_scale", 1.0), + ) + pack.to_cuda() + age_cap = getattr(host.config, "rolling_media_sink_age_cap_seconds", None) + cache.sink_key_rotation = None + if age_cap is not None: + assert pack.vision is not None + starts_seconds = [chunk.vision_frames[0] * pack.vision.seconds_per_frame[0]] + if pack.lidar is not None: + starts_seconds.append(chunk.lidar_frames[0] * pack.lidar.seconds_per_frame[0]) + cache.sink_key_rotation = prepare_media_sink_key_rotation( + host.net.language_model.model.rotary_emb, + token_count=cache.initial_sink_token_count, + chunk_start_seconds=min(starts_seconds), + age_cap_seconds=age_cap, + temporal_units_per_second=( + float(expert.base_fps) + / host.tokenizer_vision_gen.temporal_compression_factor + * getattr(host.config, "rolling_media_rope_scale", 1.0) + ), + ) + cache.prompt_layout = None + if prompt is not None: + assert pack.vision is not None + vision_period = pack.vision.seconds_per_frame[0] + temporal_units = ( + float(expert.base_fps) / host.tokenizer_vision_gen.temporal_compression_factor + if expert.enable_fps_modulation + else 1.0 / vision_period + ) + cache.prompt_layout = prompt.prepare( + pack, + position_mode=getattr(host.config, "rolling_text_position_mode", "fixed"), + temporal_units_per_second=temporal_units * getattr(host.config, "rolling_media_rope_scale", 1.0), + chunk_start_seconds=chunk.vision_frames[0] * vision_period, + ) + host._cast_generated_tokens_to_precision(pack) + host._validate_teacher_forcing_pack(pack) + condition_masks = [ + mask.clone() + for modality in (pack.vision, pack.lidar) + if modality is not None + for mask in modality.condition_mask + ] # each [V*T,1,1] + starts = (chunk.vision_frames[0],) * 2 + if plan.has_lidar: + starts += (chunk.lidar_frames[0],) * 2 + request = RollingReplayRequest( + step=chunk.step, + frame_starts=starts, + frames_per_chunk=int(host.config.teacher_forcing_frames_per_chunk), + pass_kind="clean" if clean else "noisy", + condition_masks=condition_masks, + history=cache.metadata, + prompt_layout=cache.prompt_layout, + ) + if clean: + for modality in (pack.vision, pack.lidar): + if modality is not None: + mark_modality_as_clean_condition(modality) + setattr(pack, "rolling_replay_request", request) + return pack, request + + +def _validate_rolling_transfer( + host: OmniMoTCausalModel, + plans: list[SequencePlan], + data: GenerationDataClean, + sampler_mode: str, +) -> None: + if len(plans) != 1 or not plans[0].has_vision or any((plans[0].has_action, plans[0].has_sound, plans[0].has_radar)): + raise ValueError("Rolling replay requires one RGB or RGB/LiDAR transfer sample.") + if not host._uses_multiview_replay_kv() or host.config.compile.enabled: + raise ValueError("Rolling replay requires eager multiview teacher forcing.") + if host.parallel_dims is not None and host.parallel_dims.cfgp_enabled: + raise ValueError("Rolling replay uses serial CFG; set cfg_parallel_shard_degree=1.") + if host._get_teacher_forcing_replay_policy().control_visibility != "causal": + raise ValueError("Rolling replay requires causal control visibility.") + if ( + host.config.kv_cache_dtype is not None + or host.config.kv_cache_inference_size is not None + or host.config.attention_sink_size + ): + raise ValueError( + "Use rolling_kv_cache_chunks/rolling_kv_sink_chunks without framewise or quantized KV options." + ) + groups = plans[0].vision_temporal_position_groups + shared_grid = plans[0].share_vision_temporal_positions and groups is None + # The native joint input builder expresses the same shared RGB clock with + # one temporal group instead of the legacy all-items shared-grid flag. + shared_group = ( + not plans[0].share_vision_temporal_positions + and groups is not None + and len(groups) == 2 + and groups[0] is not None + and groups[0] == groups[1] + ) + if data.temporal_positions_vision is not None or not (shared_grid or shared_group): + raise ValueError("Rolling transfer requires shared latent-index control/target positions.") + if plans[0].condition_view_indexes_vision or any( + indexes not in ([], [0]) + for indexes in (plans[0].condition_frame_indexes_vision, plans[0].condition_frame_indexes_lidar) + ): + raise ValueError( + "Rolling replay currently supports zero or one initial target frame per sensor, in every view." + ) + if sampler_mode not in ("rf", "distilled"): + raise ValueError("Rolling replay sampler_mode must be rf or distilled.") + student = host.config.fixed_step_sampler_config is not None + if student != (sampler_mode == "distilled"): + raise ValueError("Rolling replay sampler_mode must match the checkpoint's fixed-step sampler configuration.") + if data.num_vision_items_per_sample != [2] or data.x0_tokens_vision is None or len(data.x0_tokens_vision) != 2: + raise ValueError("Rolling replay requires RGB control and target items.") + views = data.num_views_per_vision_item + if views is None or len(views) != 2 or views[0] < 1 or views[0] != views[1]: + raise ValueError("Rolling replay requires matching camera-major view counts.") + if plans[0].has_lidar and ( + data.num_lidar_items_per_sample != [2] or data.x0_tokens_lidar is None or len(data.x0_tokens_lidar) != 2 + ): + raise ValueError("Rolling joint replay requires LiDAR control and target items.") + for items in (data.x0_tokens_vision, data.x0_tokens_lidar if plans[0].has_lidar else None): + if items is not None and (items[0].ndim != 5 or items[0].shape[0] != 1 or items[0].shape != items[1].shape): + raise ValueError("Rolling replay requires aligned [1,C,T,H,W] control/target latents.") + if data.x0_tokens_vision[0].shape[2] % views[0]: + raise ValueError("Rolling replay camera latents must be divisible by the view count.") + + +@torch.no_grad() +def iter_rolling_transfer_ar( + host: OmniMoTCausalModel, + *, + plans: list[SequencePlan], + data: GenerationDataClean, + conditional_text: list[list[int]], + unconditional_text: list[list[int]] | None, + guidance: float, + seed: int, + num_steps: int, + shift: float, + sampler_mode: Literal["rf", "distilled"], + distilled_num_steps: int | None = None, + output_frames: int | None = None, + on_cache_update: Callable[[RollingReplayCache], None] | None = None, + on_clean_vision_chunk: Callable[[torch.Tensor], None] | None = None, + prompt_schedule: RollingPromptSchedule | None = None, +) -> Iterator[dict[str, torch.Tensor]]: # each vision [1,Cv,V*Tvc,Hv,Wv], optional lidar [1,Cl,Tlc,Hl,Wl] + """Notify the decoder, refresh KV for later chunks, then yield the current chunk. + + Only the current chunk is packed or forwarded. Encoded input controls and + the output backing latents still scale with the input clip; this bounds the + transformer cache and query work, not VAE/media memory or rollout quality. + Input canonicalization and seed broadcast belong to the public AR entry. + """ + _validate_rolling_transfer(host, plans, data, sampler_mode) + policy = getattr(host.config, "rolling_prompt_policy", None) + if policy not in (None, "segmented") or (policy is None and prompt_schedule is not None): + raise ValueError("Segmented prompts require rolling_prompt_policy='segmented'.") + rope_scale = getattr(host.config, "rolling_media_rope_scale", 1.0) + if not math.isfinite(rope_scale) or rope_scale <= 0 or (policy is None and rope_scale != 1.0): + raise ValueError("A positive temporal RoPE scale is supported only with segmented prompts.") + age_cap = getattr(host.config, "rolling_media_sink_age_cap_seconds", None) + if age_cap is not None and ( + isinstance(age_cap, bool) + or not math.isfinite(age_cap) + or age_cap <= 0 + or policy != "segmented" + or host.config.rolling_kv_sink_chunks != 1 + or not host.config.diffusion_expert_config.enable_fps_modulation + ): + raise ValueError( + "Media sink age capping requires segmented prompts, one initial sink, FPS modulation and a positive age." + ) + text_position_mode = getattr(host.config, "rolling_text_position_mode", "fixed") + if text_position_mode not in ("fixed", "segment_start", "chunk_start") or ( + policy is None and text_position_mode != "fixed" + ): + raise ValueError("Nonfixed text temporal positioning requires segmented prompts and a supported position mode.") + if seed < 0 or num_steps < 1 or not math.isfinite(guidance) or not math.isfinite(shift) or shift <= 0: + raise ValueError("Rolling replay requires finite valid sampling settings.") + if guidance != 1.0 and unconditional_text is None: + raise ValueError("Rolling replay CFG requires unconditional captions.") + assert data.x0_tokens_vision is not None and data.num_views_per_vision_item is not None + plan, num_views = plans[0], data.num_views_per_vision_item[0] + if policy == "segmented": + if plan.text_view_ids is None: + if len(conditional_text) != 1 or (unconditional_text is not None and len(unconditional_text) != 1): + raise ValueError("A shared prompt must have one text stream.") + conditional_text = conditional_text * num_views + unconditional_text = unconditional_text * num_views if unconditional_text is not None else None + plan = replace(plan, text_view_ids=list(range(num_views))) + plans = [plan] + if plan.text_view_ids != list(range(num_views)) or len(conditional_text) != num_views: + raise ValueError("Segmented prompting requires camera-ordered text streams.") + if prompt_schedule is not None and ( + len(prompt_schedule.views) != num_views or plan.text_view_ids != list(range(num_views)) + ): + raise ValueError("Segmented prompting requires one scheduled caption stream per camera, in native view order.") + device = torch.device(host.tensor_kwargs["device"]) + vision_condition, lidar_condition = ( + len(plan.condition_frame_indexes_vision), + len(plan.condition_frame_indexes_lidar), + ) + vision = _initial_target(data.x0_tokens_vision[1], num_views, vision_condition, device) # [1,Cv,V*Tv,Hv,Wv] + lidar = ( + _initial_target(data.x0_tokens_lidar[1], 1, lidar_condition, device) + if plan.has_lidar and data.x0_tokens_lidar is not None + else None + ) # [1,Cl,Tl,Hl,Wl] or None + total_frames = vision.shape[2] // num_views + frame_count = total_frames if output_frames is None else min(output_frames, total_frames) + # A singleton probe resolves the same sensor periods and patch geometry as + # production packing, without materializing an all-horizon token pack. + first = JointARChunk(0, (0,), (0,), 1, 1) + probe = host._pack_input_sequence( + plans, conditional_text, _chunk_data(data, vision, lidar, num_views, first), torch.zeros(1) + ) # timestep: [1] + assert probe.vision is not None + vision_period = probe.vision.seconds_per_frame[0] + lidar_period = probe.lidar.seconds_per_frame[0] if probe.lidar is not None else vision_period + source_fps = host.tokenizer_vision_gen.temporal_compression_factor / vision_period + if prompt_schedule is not None and not math.isclose(prompt_schedule.fps, source_fps, rel_tol=1e-6, abs_tol=1e-6): + raise ValueError(f"Prompt schedule FPS {prompt_schedule.fps} does not match the source clock {source_fps}.") + negative_schedule = None + if policy == "segmented": + if prompt_schedule is None: + # Native evaluation callbacks may supply one static caption per + # camera. Give it the actual generated horizon, without requiring + # a separate schedule file or changing the input media length. + fps = source_fps + num_frames = round((frame_count - 1) * vision_period * fps) + 1 + prompt_schedule = RollingPromptSchedule( + fps=fps, + num_frames=num_frames, + views=tuple((RollingPromptSegment(0, num_frames, tuple(tokens)),) for tokens in conditional_text), + sink_tokens=host.config.rolling_text_sink_tokens, + special_tokens=host.llm_special_tokens, + context_before=host.config.rolling_prompt_context_before, + context_after=host.config.rolling_prompt_context_after, + media_text_extent=host.config.rolling_media_text_extent, + ) + if guidance != 1.0: + assert unconditional_text is not None + if len(unconditional_text) != num_views: + raise ValueError("Segmented prompting requires one unconditional text stream per camera.") + # RF and distilled samplers share serial CFG. The negative branch owns separate + # KV and a static negative caption; it never reads positive sinks. + negative_schedule = RollingPromptSchedule( + fps=prompt_schedule.fps, + num_frames=prompt_schedule.num_frames, + views=tuple( + (RollingPromptSegment(0, prompt_schedule.num_frames, tuple(tokens)),) + for tokens in unconditional_text + ), + sink_tokens=0, + special_tokens=host.llm_special_tokens, + ) + chunk_size = int(host.config.teacher_forcing_frames_per_chunk) + chunks = joint_ar_chunks( + vision_frames=frame_count, + lidar_frames=lidar.shape[2] if lidar is not None else 1, + vision_seconds_per_frame=vision_period, + lidar_seconds_per_frame=lidar_period, + frames_per_chunk=chunk_size, + ) + if any(not chunk.vision_frames or (lidar is not None and not chunk.lidar_frames) for chunk in chunks): + raise ValueError("Rolling joint replay requires an aligned horizon with both sensors in every chunk.") + if prompt_schedule is not None: + last_timestamp = max( + (frame_count - 1) * vision_period, + chunks[-1].lidar_frames[-1] * lidar_period if lidar is not None else 0.0, + ) + # Validate the complete finite horizon before forwarding its first chunk. + # Longer schedules may still serve a requested prefix of the input clip. + prompt_schedule.for_times((last_timestamp,)) + max_vision_frames = max(len(chunk.vision_frames) for chunk in chunks) + token_capacity = sum(math.prod(shape[1:]) * num_views * max_vision_frames for shape in probe.vision.token_shapes) + if probe.lidar is not None: + token_capacity += sum( + math.prod(shape[1:]) * max(len(chunk.lidar_frames) for chunk in chunks) + for shape in probe.lidar.token_shapes + ) + flex_backend = getattr(host.net, "flex_backend", None) + if flex_backend is not None and not getattr(host.net, "teacher_forcing_maskless", False): + # Flex appends every reserved history slot to the aligned current stream. + # Pad storage only; cache writes and metadata retain the real token counts. + kv_alignment = int(flex_backend.block_size[1]) + token_capacity = ((token_capacity + kv_alignment - 1) // kv_alignment) * kv_alignment + del probe + cache_chunks = host.config.rolling_kv_cache_chunks + assert cache_chunks is not None + texts = [conditional_text] if guidance == 1.0 else [conditional_text, unconditional_text] + caches = [ + RollingReplayCache( + num_layers=int(host.net.num_hidden_layers), + cache_chunks=cache_chunks, + sink_chunks=host.config.rolling_kv_sink_chunks, + tokens_per_chunk=token_capacity, + ) + for _ in texts + ] + if prompt_schedule is not None: + for cache in caches: + cache.text_sink_layers = [RollingTextSinkLayer() for _ in cache.cache] + log.info( + f"Rolling replay: {cache_chunks} history chunks, {host.config.rolling_kv_sink_chunks} sinks, {caches[0].capacity} KV tokens/layer/branch." + ) + for chunk in chunks: + prompts: list[RollingPromptPack | None] = [None] * len(caches) + if prompt_schedule is not None: + timestamps = tuple(frame * vision_period for frame in chunk.vision_frames) + if lidar is not None: + timestamps += tuple(frame * lidar_period for frame in chunk.lidar_frames) + prompts = [prompt_schedule.for_times(timestamps)] + if negative_schedule is not None: + prompts.append(negative_schedule.for_times(timestamps)) + texts = [prompt.text_tokens for prompt in prompts if prompt is not None] + chunk_plan = replace( + plan, + condition_frame_indexes_vision=[0] if vision_condition and chunk.step == 0 else [], + condition_frame_indexes_lidar=[0] if lidar_condition and chunk.step == 0 else [], + ) + current = _chunk_data(data, vision, lidar, num_views, chunk) + assert current.x0_tokens_vision is not None + vision_indexes = [ + view * len(chunk.vision_frames) + index + for view in range(num_views) + for index, frame in enumerate(chunk.vision_frames) + if frame >= vision_condition + ] + lidar_indexes = ( + [index for index, frame in enumerate(chunk.lidar_frames) if frame >= lidar_condition] + if lidar is not None + else [] + ) + vision_shape = (1, vision.shape[1], len(vision_indexes), *vision.shape[-2:]) + lidar_shape = (1, lidar.shape[1], len(lidar_indexes), *lidar.shape[-2:]) if lidar is not None else (1, 0) + vision_size, lidar_size = math.prod(vision_shape), math.prod(lidar_shape) + if vision_size + lidar_size: + # Match the reference joint seed clock; preserve RGB's frame-start + # seed convention for its separate pre-existing public entry. + sample_index = chunk.step if lidar is not None else chunk.vision_frames[0] + generator = torch.Generator(device=device).manual_seed(seed + sample_index) + noise = torch.randn( + (1, vision_size + lidar_size), device=device, dtype=torch.float32, generator=generator + ) # [1,Dv+Dl] + branches = [ + _build_branch_pack( + host, plan=chunk_plan, data=current, text=text, chunk=chunk, cache=cache, clean=False, prompt=prompt + )[0] + for text, cache, prompt in zip(texts, caches, prompts, strict=True) + if text is not None + ] + + def velocity(noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: # [1,Dv+Dl], [1,1] -> [1,Dv+Dl] + predictions: list[torch.Tensor] = [] # each [1,Dv+Dl] + for pack, cache in zip(branches, caches, strict=True): + assert current.x0_tokens_vision is not None + noisy_vision = current.x0_tokens_vision[1].clone() # [1,Cv,V*Tvc,Hv,Wv] + noisy_vision[:, :, vision_indexes] = noise[:, :vision_size].reshape( + vision_shape + ) # [1,Cv,V*Tvc_noisy,Hv,Wv] + lidar_items = None + if current.x0_tokens_lidar is not None: + noisy_lidar = current.x0_tokens_lidar[1].clone() # [1,Cl,Tlc,Hl,Wl] + noisy_lidar[:, :, lidar_indexes] = noise[:, vision_size:].reshape( + lidar_shape + ) # [1,Cl,Tlc_noisy,Hl,Wl] + lidar_items = [current.x0_tokens_lidar[0], noisy_lidar] + host._update_inference_pack_template( + packed_sequence=pack, + noise_x_vision=[current.x0_tokens_vision[0], noisy_vision], + noise_x_lidar=lidar_items, + noise_x_radar=None, + noise_x_action=None, + noise_x_sound=None, + timestep=timestep, # [1,1] + ) + output = host.denoise(data_batch_packed=pack, memory=cache.read()) + values = [output["preds_vision"][1][:, :, vision_indexes].reshape(1, -1)] # each [1,Dv] + if lidar is not None: + values.append(output["preds_lidar"][1][:, :, lidar_indexes].reshape(1, -1)) # [1,Dl] + predictions.append(torch.cat(values, dim=1)) # [1,Dv+Dl] + if len(predictions) == 1: + return predictions[0] # [1,Dv+Dl] + return predictions[1] + guidance * (predictions[0] - predictions[1]) # [1,Dv+Dl] + + denoised = host._run_ar_sampler( + velocity, + noise, + sampler_mode=sampler_mode, + num_steps=num_steps, + shift=shift, + seed=seed, + sample_idx=sample_index, + num_frames=frame_count, + distilled_num_steps=distilled_num_steps, + schedule_frame_idx=chunk.vision_frames[0], + ) # [1,Dv+Dl] + if not torch.isfinite(denoised).all(): + raise FloatingPointError(f"Rolling replay produced nonfinite latents in chunk {chunk.step}.") + final_vision_indexes = [ + view * total_frames + frame + for view in range(num_views) + for frame in chunk.vision_frames + if frame >= vision_condition + ] + vision[:, :, final_vision_indexes] = denoised[:, :vision_size].reshape( + vision_shape + ) # [1,Cv,V*Tvc_noisy,Hv,Wv] + if lidar is not None: + final_lidar_indexes = [frame for frame in chunk.lidar_frames if frame >= lidar_condition] + lidar[:, :, final_lidar_indexes] = denoised[:, vision_size:].reshape( + lidar_shape + ) # [1,Cl,Tlc_noisy,Hl,Wl] + del branches + clean_data = _chunk_data(data, vision, lidar, num_views, chunk) + assert clean_data.x0_tokens_vision is not None + if on_clean_vision_chunk is not None: + on_clean_vision_chunk(clean_data.x0_tokens_vision[1]) + if chunk.step != chunks[-1].step: + for text, cache, prompt in zip(texts, caches, prompts, strict=True): + assert text is not None + clean_pack, request = _build_branch_pack( + host, + plan=chunk_plan, + data=clean_data, + text=text, + chunk=chunk, + cache=cache, + clean=True, + prompt=prompt, + ) + token_count = sum( + math.prod(shape) + for modality in (clean_pack.vision, clean_pack.lidar) + if modality is not None + for shape in modality.token_shapes + ) + host.denoise( + data_batch_packed=clean_pack, + memory=cache.writer(step=chunk.step, token_count=token_count, device=device), + ) + cache.commit(step=chunk.step, token_count=token_count, request=request) + del clean_pack, request + if on_cache_update is not None: + on_cache_update(caches[0]) + assert clean_data.x0_tokens_vision is not None + result = {"vision": clean_data.x0_tokens_vision[1]} # [1,Cv,V*Tvc,Hv,Wv] + if clean_data.x0_tokens_lidar is not None: + result["lidar"] = clean_data.x0_tokens_lidar[1] # [1,Cl,Tlc,Hl,Wl] + yield result + + +def sample_rolling_joint_transfer_ar( + host: OmniMoTCausalModel, + *, + plans: list[SequencePlan], + data: GenerationDataClean, + conditional_text: list[list[int]], + unconditional_text: list[list[int]] | None, + guidance: float, + seed: int, + num_steps: int, + shift: float, + sampler_mode: Literal["rf", "distilled"], + distilled_num_steps: int | None, + prompt_schedule: RollingPromptSchedule | None = None, +) -> dict[str, list[torch.Tensor]]: # vision [1,Cv,V*Tv,Hv,Wv], lidar [1,Cl,Tl,Hl,Wl] + """Collect chunks in the existing joint API's camera-major output layout.""" + assert data.num_views_per_vision_item is not None + views = data.num_views_per_vision_item[0] + rgb_parts: list[torch.Tensor] = [] # each [1,Cv,V,Tvc,Hv,Wv] + lidar_parts: list[torch.Tensor] = [] # each [1,Cl,Tlc,Hl,Wl] + for part in iter_rolling_transfer_ar( + host, + plans=plans, + data=data, + conditional_text=conditional_text, + unconditional_text=unconditional_text, + guidance=guidance, + seed=seed, + num_steps=num_steps, + shift=shift, + sampler_mode=sampler_mode, + distilled_num_steps=distilled_num_steps, + prompt_schedule=prompt_schedule, + ): + rgb = part["vision"] # [1,Cv,V*Tvc,Hv,Wv] + rgb_parts.append(rgb.reshape(1, rgb.shape[1], views, -1, *rgb.shape[-2:])) # [1,Cv,V,Tvc,Hv,Wv] + lidar_parts.append(part["lidar"]) # [1,Cl,Tlc,Hl,Wl] + vision = torch.cat(rgb_parts, dim=3).flatten(2, 3) # [1,Cv,V*Tv,Hv,Wv] + return {"vision": [vision], "lidar": [torch.cat(lidar_parts, dim=2)]} # lidar [1,Cl,Tl,Hl,Wl] diff --git a/cosmos_framework/trainer/__init__.py b/cosmos_framework/trainer/__init__.py index 94aea9449..e9d1fd728 100644 --- a/cosmos_framework/trainer/__init__.py +++ b/cosmos_framework/trainer/__init__.py @@ -76,6 +76,15 @@ def assert_synced_with_model(self, model: Any) -> None: ) +def persistent_validation_enabled(config_trainer: Any) -> bool: + """Whether training keeps one validation iterator alive, so validation data advances across calls.""" + return bool( + config_trainer.run_validation + and getattr(config_trainer, "prefetch_validation", False) + and config_trainer.max_val_iter is not None + ) + + class ImaginaireTrainer: """The base trainer class of Imaginaire. @@ -317,12 +326,7 @@ def train( # callbacks) so data-augmentation randomness starts from a deterministic state # regardless of how much RNG state init consumed. misc.set_random_seed(seed=self.config.trainer.seed, by_rank=True) - if ( - self.config.trainer.run_validation - and getattr(self.config.trainer, "prefetch_validation", False) - and self.config.trainer.max_val_iter is not None - and self._validation_iterator is None - ): + if persistent_validation_enabled(self.config.trainer) and self._validation_iterator is None: self._validation_iterator = iter(dataloader_val) with ( maybe_enable_profiling(self.config, global_step=iteration) as torch_profiler, diff --git a/cosmos_framework/utils/generator/model_loader.py b/cosmos_framework/utils/generator/model_loader.py index 33b25ae95..37df7a5e6 100644 --- a/cosmos_framework/utils/generator/model_loader.py +++ b/cosmos_framework/utils/generator/model_loader.py @@ -53,13 +53,16 @@ def write_lock(self) -> FileLock: # Thus these load_model functions are designed with less dependency. -def checkpoint_path_to_cached_path(path: str, cache_rootdir: Optional[str] = None) -> str: +def checkpoint_path_to_cached_path( + path: str, cache_rootdir: str | None = None, *, load_ema_to_reg: bool = False +) -> str: + """Cache S3 checkpoints by bucket, object key and the selected weight set.""" if cache_rootdir is None: homedir = os.getenv("HOME") or "" cache_rootdir = osp.join(homedir, ".cache/imaginaire4/checkpoints/") if path.startswith("s3://"): - return osp.join(cache_rootdir, path.removeprefix("s3://").split("/", maxsplit=1)[1]) + return osp.join(cache_rootdir, "ema" if load_ema_to_reg else "regular", path.removeprefix("s3://")) else: return path @@ -359,7 +362,8 @@ def load_model_from_checkpoint( seed: Random seed used for initialization (if applicable). experiment_opts: Extra experiment/config override options. use_cache_checkpoint: If True, locally save & read remote checkpoints to speed up repeated loads. - Be aware, the default cache path is $HOME/.cache/imaginaire4/checkpoints/. + The default cache path is $HOME/.cache/imaginaire4/checkpoints///. + Legacy caches without weight-set and bucket identity are left intact but not reused. Applies to the DCP path only; for safetensors checkpoints a warning is logged and caching is skipped (the DCP write cache is not a meaningful round-trip for safetensors sources). @@ -514,7 +518,9 @@ def load_model(checkpoint_load_path: str) -> None: checkpoint_cache_path = None if use_cache_checkpoint: - checkpoint_cache_path = checkpoint_path_to_cached_path(checkpoint_path, cache_checkpoint_rootdir) + checkpoint_cache_path = checkpoint_path_to_cached_path( + checkpoint_path, cache_checkpoint_rootdir, load_ema_to_reg=load_ema_to_reg + ) if checkpoint_cache_path is None: load_model(checkpoint_path)