diff --git a/README.md b/README.md index b9585825..0bf7965c 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,6 @@ first; the skills above are referenced from there. | [Code Structure](./docs/code_structure.md) | Repository layout and a per-subpackage tour of `cosmos_framework/` — where each concern lives and where to add new code. | | [Training](./docs/training.md) | Launching multi-GPU and multi-node runs; parallelism strategies; mixed precision; resuming. | | [Inference (from a trained checkpoint)](./docs/inference.md) | Loading a trained checkpoint into one of the inference backends. | -| [Policy Server](./docs/action_policy_droid_server.md) | Running the server-client pipeline for Cosmos3-Nano-Policy-DROID. | +| [Policy Server](./docs/action_policy_droid_server.md) | Running the server-client pipeline for Cosmos3-Policy-DROID. | | [FAQ](./docs/faq.md) | Troubleshooting (OOM, NCCL hangs, slow training), environment variables, and common pitfalls. | | [AGENTS.md](./AGENTS.md) + [Agent Skills](#agent-skills) | Repo map and task-specific `SKILL.md` files loaded automatically by AGENTS.md-aware coding agents (Claude Code, Codex CLI, Cursor, etc.). | diff --git a/cosmos_framework/callbacks/dmd2_metrics.py b/cosmos_framework/callbacks/dmd2_metrics.py new file mode 100644 index 00000000..4af5bab1 --- /dev/null +++ b/cosmos_framework/callbacks/dmd2_metrics.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from __future__ import annotations + +import hashlib +import json +import pickle +import random +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils import distributed +from cosmos_framework.utils.callback import Callback + +__all__: tuple[str, ...] = ("DMD2Metrics", "DMD2ParityLedger", "PARITY_KEYS") + +PARITY_KEYS: tuple[str, ...] = ( + "vsd_loss", + "total_generator_loss", + "dmd_loss_generator", + "dmd_vsd_grad_norm", + "fake_score_loss", + "total_critic_loss", + "dmd_loss_critic", + "clip_grad_norm/net_selected_preclip", + "clip_grad_norm/net_selected_clip_scale", + "clip_grad_norm/net_selected_clip_norm", + "clip_grad_norm/fake_score_selected_preclip", + "clip_grad_norm/fake_score_selected_clip_scale", + "clip_grad_norm/fake_score_selected_clip_norm", +) + +_SAMPLE_KEY_FIELDS: tuple[str, ...] = ("__key__", "sample_id", "uuid") +_VOLATILE_INPUT_FIELDS: frozenset[str] = frozenset( + { + "_aug_step_times", + "_aug_time", + "_pre_aug_time", + "_sample_time", + "_worker_aug_step_times", + "_worker_aug_time", + "_worker_batch_time", + "_worker_id", + "_worker_io_time", + } +) + + +def _local_rng_checksum() -> str: + """Hash local RNG states without advancing any generator.""" + digest = hashlib.sha256() + torch_state = torch.get_rng_state() # [N_cpu_rng] + digest.update(torch_state.numpy().tobytes()) + if torch.cuda.is_available(): + cuda_state = torch.cuda.get_rng_state() # [N_cuda_rng] + digest.update(cuda_state.cpu().numpy().tobytes()) + digest.update(pickle.dumps(np.random.get_state(), protocol=pickle.HIGHEST_PROTOCOL)) + digest.update(pickle.dumps(random.getstate(), protocol=pickle.HIGHEST_PROTOCOL)) + return digest.hexdigest() + + +def _gather_objects(payload: object) -> list[object]: + if not dist.is_available() or not dist.is_initialized(): + return [payload] + return distributed.all_gather_object(payload) + + +def _flatten_sample_keys(value: object) -> list[str]: + if value is None: + return [] + if isinstance(value, torch.Tensor): + local_value = value.detach().cpu() # [...] + return [str(item) for item in local_value.reshape(-1).tolist()] + if isinstance(value, bytes): + return [value.decode(errors="replace")] + if isinstance(value, str): + return [value] + if isinstance(value, Iterable): + flattened: list[str] = [] + for item in value: + flattened.extend(_flatten_sample_keys(item)) + return flattened + return [str(value)] + + +def _sample_keys(data_batch: dict[str, object]) -> list[str]: + local_keys: list[str] = [] + for field in _SAMPLE_KEY_FIELDS: + if field in data_batch: + local_keys = _flatten_sample_keys(data_batch[field]) + break + gathered_keys = _gather_objects(local_keys) + return sorted(key for rank_keys in gathered_keys for key in _flatten_sample_keys(rank_keys)) + + +def _rng_checksum() -> str: + gathered_checksums = [str(checksum) for checksum in _gather_objects(_local_rng_checksum())] + return hashlib.sha256("\n".join(gathered_checksums).encode()).hexdigest() + + +def _update_input_digest(digest: Any, value: object) -> None: + if isinstance(value, torch.Tensor): + to_local = getattr(value, "to_local", None) + local_value = to_local() if callable(to_local) else value # [...] + cpu_value = local_value.detach().cpu().contiguous() # [...] + byte_view = cpu_value.view(torch.uint8) # [N_bytes] + digest.update(f"tensor:{cpu_value.dtype}:{tuple(cpu_value.shape)}:".encode()) + digest.update(byte_view.numpy().tobytes()) + return + if isinstance(value, np.ndarray): + contiguous_value = np.ascontiguousarray(value) + digest.update(f"ndarray:{contiguous_value.dtype}:{contiguous_value.shape}:".encode()) + digest.update(contiguous_value.tobytes()) + return + if isinstance(value, dict): + digest.update(b"dict:") + for key in sorted(value, key=lambda item: (type(item).__qualname__, str(item))): + _update_input_digest(digest, key) + _update_input_digest(digest, value[key]) + return + if isinstance(value, (list, tuple)): + digest.update(f"{type(value).__qualname__}:{len(value)}:".encode()) + for item in value: + _update_input_digest(digest, item) + return + if isinstance(value, bytes): + digest.update(f"bytes:{len(value)}:".encode()) + digest.update(value) + return + if isinstance(value, str): + encoded_value = value.encode() + digest.update(f"str:{len(encoded_value)}:".encode()) + digest.update(encoded_value) + return + digest.update( + pickle.dumps((type(value).__module__, type(value).__qualname__, value), protocol=pickle.HIGHEST_PROTOCOL) + ) + + +def _input_digest(data_batch: dict[str, object]) -> str: + stable_batch = {key: value for key, value in data_batch.items() if key not in _VOLATILE_INPUT_FIELDS} + local_hasher = hashlib.sha256() + _update_input_digest(local_hasher, stable_batch) + local_digest = local_hasher.hexdigest() + gathered_digests = [str(digest) for digest in _gather_objects(local_digest)] + return hashlib.sha256("\n".join(gathered_digests).encode()).hexdigest() + + +def _average_metric(value: object, device: torch.device) -> float: + if isinstance(value, torch.Tensor): + to_local = getattr(value, "to_local", None) + local_value = to_local() if callable(to_local) else value # [...] + metric = local_value.detach().float().clone() # [...] + else: + metric = torch.tensor(value, device=device, dtype=torch.float32) # [] + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(metric, op=dist.ReduceOp.AVG) + return metric.mean().item() + + +def _phase(model: torch.nn.Module, iteration: int) -> str: + get_phase = getattr(model, "get_phase", None) + if callable(get_phase): + return str(get_phase(iteration)) + get_optimizer_key = getattr(model, "get_optimizer_key", None) + if callable(get_optimizer_key): + return "student" if str(get_optimizer_key(iteration)) == "net" else "critic" + raise AttributeError(f"{type(model).__name__} must define get_phase() or get_optimizer_key()") + + +class DMD2ParityLedger(Callback): + """Append deterministic DMD2 parity evidence to rank-0 JSONL.""" + + def __init__(self, output_path: str = "") -> None: + super().__init__() + self.output_path: str = output_path + + def on_training_step_end( + self, + model: ImaginaireModel, + data_batch: dict[str, torch.Tensor], + output_batch: dict[str, torch.Tensor], + loss: torch.Tensor, + iteration: int = 0, + ) -> None: + if not self.output_path: + return + + grad_metrics = getattr(model, "_distillation_parity_grad_metrics", {}) + metric_values = {**output_batch, **grad_metrics} + record: dict[str, object] = { + "iteration": iteration, + "phase": _phase(model, max(iteration - 1, 0)), + "sample_keys": _sample_keys(data_batch), + "input_digest": _input_digest(data_batch), + "rng_checksum": _rng_checksum(), + } + for key in PARITY_KEYS: + if key in metric_values: + record[key] = _average_metric(metric_values[key], loss.device) + + if not distributed.is_rank0(): + return + output_path = Path(self.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("a", encoding="utf-8") as output_file: + output_file.write(json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n") + + +class DMD2Metrics(DMD2ParityLedger): + """Public OSS name for the DMD2 parity and diagnostic metrics callback.""" diff --git a/cosmos_framework/callbacks/dmd2_metrics_test.py b/cosmos_framework/callbacks/dmd2_metrics_test.py new file mode 100644 index 00000000..89d48ee0 --- /dev/null +++ b/cosmos_framework/callbacks/dmd2_metrics_test.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest +import torch + +if "megatron.core" not in sys.modules: + megatron_module = ModuleType("megatron") + megatron_core_module = ModuleType("megatron.core") + megatron_core_module.parallel_state = SimpleNamespace() # type: ignore[attr-defined] + megatron_module.core = megatron_core_module # type: ignore[attr-defined] + sys.modules["megatron"] = megatron_module + sys.modules["megatron.core"] = megatron_core_module + +import cosmos_framework.callbacks.dmd2_metrics as ledger_module +from cosmos_framework.callbacks.dmd2_metrics import DMD2ParityLedger +from cosmos_framework.callbacks.grad_clip_distillation import GradClip + + +class _PhaseModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(grad_clip=True) + self._distillation_parity_grad_metrics: dict[str, float] = {} + + def get_phase(self, iteration: int) -> str: + return "student" if iteration % 5 == 0 else "critic" + + def get_optimizer_key(self, iteration: int) -> str: + return "net" if self.get_phase(iteration) == "student" else "fake_score" + + def is_image_batch(self, data_batch: dict[str, Any]) -> bool: + return bool(data_batch.get("is_image", False)) + + +class _PhaseOptimizer: + def __init__(self, parameter: torch.nn.Parameter) -> None: + self.parameter = parameter + + def parameters_for_key(self, key: str) -> list[torch.nn.Parameter]: + assert key == "net" + return [self.parameter] + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_disabled_ledger_does_not_write_or_run_collectives(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + output_path = tmp_path / "disabled.jsonl" + callback = DMD2ParityLedger(output_path="") + model = _PhaseModel() + + def _unexpected_collective(payload: object) -> list[object]: + raise AssertionError(f"disabled ledger gathered {payload!r}") + + monkeypatch.setattr(ledger_module.distributed, "all_gather_object", _unexpected_collective) + callback.on_training_step_end( + model, + {"sample_id": ["sample-a"]}, + {"fake_score_loss": torch.tensor(1.0)}, # [] + torch.tensor(1.0), # [] + iteration=1, + ) + + assert not output_path.exists() + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_ledger_writes_sorted_sample_keys_phase_rng_and_present_metrics( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + output_path = tmp_path / "parity.jsonl" + callback = DMD2ParityLedger(output_path=str(output_path)) + model = _PhaseModel() + model._distillation_parity_grad_metrics = { + "clip_grad_norm/net_selected_preclip": 3.0, + "clip_grad_norm/net_selected_clip_scale": 0.5, + "clip_grad_norm/net_selected_clip_norm": 1.5, + } + monkeypatch.setattr(ledger_module, "_local_rng_checksum", lambda: "rank-0-rng") + reduced_metrics: list[torch.Tensor] = [] + + def reduce_metric(metric: torch.Tensor, op: object) -> None: # metric: [] + assert op == torch.distributed.ReduceOp.AVG + reduced_metrics.append(metric.clone()) # [] + + monkeypatch.setattr(ledger_module.dist, "is_available", lambda: True) + monkeypatch.setattr(ledger_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(ledger_module.dist, "all_reduce", reduce_metric) + monkeypatch.setattr(ledger_module.distributed, "all_gather_object", lambda payload: [payload]) + monkeypatch.setattr(ledger_module.distributed, "is_rank0", lambda: True) + + callback.on_training_step_end( + model, + { + "sample_id": ["sample-b", "sample-a"], + "images": torch.tensor([[1.0], [2.0]]), # [B,C] + }, + { + "vsd_loss": torch.tensor(2.0), # [] + "total_generator_loss": torch.tensor(4.0), # [] + "ignored_metric": torch.tensor(9.0), # [] + }, + torch.tensor(4.0), # [] + iteration=5, + ) + + record = json.loads(output_path.read_text().strip()) + assert record["iteration"] == 5 + assert record["phase"] == "critic" + assert record["sample_keys"] == ["sample-a", "sample-b"] + assert len(record["input_digest"]) == 64 + assert len(record["rng_checksum"]) == 64 + assert record["vsd_loss"] == 2.0 + assert record["total_generator_loss"] == 4.0 + assert record["clip_grad_norm/net_selected_preclip"] == 3.0 + assert record["clip_grad_norm/net_selected_clip_scale"] == 0.5 + assert record["clip_grad_norm/net_selected_clip_norm"] == 1.5 + assert "ignored_metric" not in record + assert len(reduced_metrics) == 5 + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_input_digest_is_independent_of_mapping_order(monkeypatch: pytest.MonkeyPatch) -> None: + first_batch = { + "sample_id": ["sample-a"], + "images": torch.tensor([[1.0, 2.0]]), # [B,C] + "metadata": {"height": 1, "width": 2}, + } + reordered_batch = { + "metadata": {"width": 2, "height": 1}, + "images": torch.tensor([[1.0, 2.0]]), # [B,C] + "sample_id": ["sample-a"], + } + monkeypatch.setattr(ledger_module.distributed, "all_gather_object", lambda payload: [payload]) + + assert ledger_module._input_digest(first_batch) == ledger_module._input_digest(reordered_batch) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_input_digest_ignores_worker_timing_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + first_batch = { + "sample_id": ["sample-a"], + "images": torch.tensor([[1.0, 2.0]]), # [B,C] + "_worker_io_time": torch.tensor([0.1]), # [B] + "_worker_aug_step_times": {"resize": 0.2}, + } + later_batch = { + "sample_id": ["sample-a"], + "images": torch.tensor([[1.0, 2.0]]), # [B,C] + "_worker_io_time": torch.tensor([9.9]), # [B] + "_worker_aug_step_times": {"resize": 8.8}, + } + monkeypatch.setattr(ledger_module.distributed, "all_gather_object", lambda payload: [payload]) + + assert ledger_module._input_digest(first_batch) == ledger_module._input_digest(later_batch) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + ("iteration", "expected_phase"), + [ + (1, "student"), + (5, "critic"), + (6, "student"), + ], +) +def test_ledger_phase_describes_completed_step( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + iteration: int, + expected_phase: str, +) -> None: + output_path = tmp_path / "parity.jsonl" + callback = DMD2ParityLedger(output_path=str(output_path)) + model = _PhaseModel() + monkeypatch.setattr(ledger_module, "_local_rng_checksum", lambda: "rank-0-rng") + monkeypatch.setattr(ledger_module.distributed, "all_gather_object", lambda payload: [payload]) + monkeypatch.setattr(ledger_module.distributed, "is_rank0", lambda: True) + + callback.on_training_step_end( + model, + {"sample_id": ["sample-a"]}, + {}, + torch.tensor(0.0), # [] + iteration=iteration, + ) + + record = json.loads(output_path.read_text().strip()) + assert record["iteration"] == iteration + assert record["phase"] == expected_phase + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_grad_clip_exposes_computed_phase_metrics() -> None: + parameter = torch.nn.Parameter(torch.tensor([0.0, 0.0])) # [2] + parameter.grad = torch.tensor([3.0, 4.0]) # [2] + model = _PhaseModel() + optimizer = _PhaseOptimizer(parameter) + callback = GradClip(clip_norm=10.0) + callback.config = SimpleNamespace(trainer=SimpleNamespace(logging_iter=10)) + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + + callback.on_before_optimizer_step(model, optimizer, None, None, iteration=5) # type: ignore[arg-type] + + assert model._distillation_parity_grad_metrics == { + "clip_grad_norm/net_selected_preclip": 5.0, + "clip_grad_norm/net_selected_clip_scale": 1.0, + "clip_grad_norm/net_selected_clip_norm": 10.0, + } + + + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_metrics_name_preserves_internal_ledger_alias() -> None: + public_metrics = getattr(ledger_module, "DMD2Metrics", None) + + assert public_metrics is not None + assert issubclass(public_metrics, ledger_module.DMD2ParityLedger) + assert ledger_module.__all__ == ("DMD2Metrics", "DMD2ParityLedger", "PARITY_KEYS") diff --git a/cosmos_framework/callbacks/grad_clip_distillation.py b/cosmos_framework/callbacks/grad_clip_distillation.py new file mode 100644 index 00000000..1ca7cc89 --- /dev/null +++ b/cosmos_framework/callbacks/grad_clip_distillation.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from dataclasses import dataclass +from typing import cast + +import torch +import wandb +from torch.nn.utils.clip_grad import clip_grad_norm_ + +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils.callback import Callback + +__all__: tuple[str, ...] = ("GradClip",) + + +@torch.jit.script +def _fused_nan_to_num(params: list[torch.Tensor]) -> None: + for param in params: + torch.nan_to_num(param, nan=0.0, posinf=0.0, neginf=0.0, out=param) + + +def _to_scalar(value: torch.Tensor | float) -> float: + """Convert tensor or DTensor to Python float for logging; avoids DTensor collectives.""" + if not isinstance(value, torch.Tensor): + return float(value) + to_local = getattr(value, "to_local", None) + t = cast(torch.Tensor, to_local() if callable(to_local) else value) + return t.detach().float().item() + + +def _replicate_norm(value: torch.Tensor | float) -> torch.Tensor | float: + """Materialize a DTensor norm partial before logging it as a global scalar.""" + full_tensor = getattr(value, "full_tensor", None) + if callable(full_tensor): + return cast(torch.Tensor, full_tensor()) # [] + return value + + +def _clip_scale(total_norm: torch.Tensor | float, clip_norm: float) -> float: + """Return the coefficient used by PyTorch gradient clipping.""" + if isinstance(total_norm, torch.Tensor): + scale = torch.clamp(clip_norm / (total_norm + 1e-6), max=1.0) # [] + return _to_scalar(scale) + return min(1.0, clip_norm / (float(total_norm) + 1e-6)) + + +def _phase_optimizer_parameters( + model: torch.nn.Module, optimizer: object, iteration: int +) -> tuple[str, list[torch.nn.Parameter]]: + assert hasattr(model, "get_optimizer_key"), ( + f"{type(model).__name__} missing get_optimizer_key — GradClipCosmos3 requires a PhaseOptimizer-compatible model" + ) + assert hasattr(optimizer, "parameters_for_key"), ( + f"{type(optimizer).__name__} missing parameters_for_key — GradClipCosmos3 requires PhaseOptimizer" + ) + key = str(model.get_optimizer_key(iteration)) + return key, list(optimizer.parameters_for_key(key)) + + +@dataclass +class _MagnitudeRecord: + state: float = 0 + iter_count: int = 0 + + def reset(self) -> None: + self.state = 0 + self.iter_count = 0 + + def update(self, cur_state: torch.Tensor | float) -> None: + self.state += _to_scalar(cur_state) + self.iter_count += 1 + + def get_stat(self) -> float: + if self.iter_count > 0: + avg_state = self.state / self.iter_count + else: + avg_state = 0.0 + self.reset() + return avg_state + + +class GradClip(Callback): + """ + This callback is used to clip the gradient norm of the model. + It also logs the average gradient norm of the model to wandb. + """ + + def __init__(self, clip_norm: float = 1.0, force_finite: bool = True) -> None: + self.clip_norm: float = clip_norm + self.force_finite: bool = force_finite + + self.img_mag_log: _MagnitudeRecord = _MagnitudeRecord() + self.video_mag_log: _MagnitudeRecord = _MagnitudeRecord() + self.phase_grad_norm_logs: dict[str, _MagnitudeRecord] = {} + self.phase_clip_scale_logs: dict[str, _MagnitudeRecord] = {} + self._cur_state: _MagnitudeRecord | None = None + + def on_training_step_start( + self, model: torch.nn.Module, data_batch: dict[str, torch.Tensor], iteration: int = 0 + ) -> None: + model._distillation_parity_grad_metrics = {} + if model.is_image_batch(data_batch): + self._cur_state = self.img_mag_log + else: + self._cur_state = self.video_mag_log + + def on_before_optimizer_step( + self, + model: ImaginaireModel, + optimizer: object, + scheduler: object, + grad_scaler: object, + iteration: int = 0, + ) -> None: + del scheduler + optimizer_key, clip_params = _phase_optimizer_parameters(model, optimizer, iteration) + if self.force_finite: + params = [param.grad for param in clip_params if param.grad is not None] + _fused_nan_to_num(params) + + clip_norm = 1e12 if not getattr(model.config, "grad_clip", True) else self.clip_norm + total_norm = _replicate_norm(clip_grad_norm_(clip_params, clip_norm)) # [] + grad_norm = _to_scalar(total_norm) + clip_scale = _clip_scale(total_norm, clip_norm) + grad_metrics: dict[str, float] | None = getattr(model, "_distillation_parity_grad_metrics", None) + if grad_metrics is None: + grad_metrics = {} + model._distillation_parity_grad_metrics = grad_metrics + grad_metrics.update( + { + f"clip_grad_norm/{optimizer_key}_selected_preclip": grad_norm, + f"clip_grad_norm/{optimizer_key}_selected_clip_scale": clip_scale, + f"clip_grad_norm/{optimizer_key}_selected_clip_norm": clip_norm, + } + ) + self.phase_grad_norm_logs.setdefault(optimizer_key, _MagnitudeRecord()).update(grad_norm) + self.phase_clip_scale_logs.setdefault(optimizer_key, _MagnitudeRecord()).update(clip_scale) + + self._cur_state.update(total_norm) # type: ignore[union-attr] + if iteration % self.config.trainer.logging_iter == 0: + avg_img_mag, avg_video_mag = self.img_mag_log.get_stat(), self.video_mag_log.get_stat() + metrics = { + "clip_grad_norm/image": avg_img_mag, + "clip_grad_norm/video": avg_video_mag, + "iteration": iteration, + } + for key, record in self.phase_grad_norm_logs.items(): + if record.iter_count > 0: + metrics[f"clip_grad_norm/{key}_selected_preclip"] = record.get_stat() + for key, record in self.phase_clip_scale_logs.items(): + if record.iter_count > 0: + metrics[f"clip_grad_norm/{key}_selected_clip_scale"] = record.get_stat() + if wandb.run: + wandb.log(metrics, step=iteration) diff --git a/cosmos_framework/callbacks/grad_clip_distillation_test.py b/cosmos_framework/callbacks/grad_clip_distillation_test.py new file mode 100644 index 00000000..b258bb78 --- /dev/null +++ b/cosmos_framework/callbacks/grad_clip_distillation_test.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import cosmos_framework.callbacks.grad_clip_distillation as grad_clip_module +from cosmos_framework.callbacks.grad_clip_distillation import GradClip +from cosmos_framework.model.generator.distillation.optimizer import PhaseOptimizer + + +class _PhaseModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.net: torch.nn.Linear = torch.nn.Linear(2, 1, bias=False) + self.net_fake_score: torch.nn.Linear = torch.nn.Linear(2, 1, bias=False) + self.config: SimpleNamespace = SimpleNamespace(grad_clip=True) + self._distillation_parity_grad_metrics: dict[str, float] = {} + self.optimizer_key_override: str | None = None + + def get_optimizer_key(self, iteration: int) -> str: + if self.optimizer_key_override is not None: + return self.optimizer_key_override + return "net" if iteration % 5 == 0 else "fake_score" + + def is_image_batch(self, data_batch: dict[str, object]) -> bool: + return bool(data_batch["is_image"]) + + +class _PartialNorm: + def __init__(self, local_value: float, global_value: float) -> None: + self.local_value: float = local_value + self.global_value: float = global_value + self.full_tensor_calls: int = 0 + + def __float__(self) -> float: + return self.local_value + + def full_tensor(self) -> torch.Tensor: # returns: [] + self.full_tensor_calls += 1 + return torch.tensor(self.global_value) # [] + + +def _callback() -> GradClip: + callback = GradClip(clip_norm=1.0) + callback.config = SimpleNamespace(trainer=SimpleNamespace(logging_iter=100)) + return callback + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_grad_clip_exports_are_explicit() -> None: + assert grad_clip_module.__all__ == ("GradClip",) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_student_phase_clips_only_student_parameters() -> None: + model = _PhaseModel() + optimizer = PhaseOptimizer( + { + "net": torch.optim.SGD(model.net.parameters(), lr=0.1), + "fake_score": torch.optim.SGD(model.net_fake_score.parameters(), lr=0.1), + } + ) + model.net.weight.grad = torch.tensor([[3.0, 4.0]]) # [1,2] + model.net_fake_score.weight.grad = torch.tensor([[0.0, 12.0]]) # [1,2] + critic_grad_before = model.net_fake_score.weight.grad.clone() # [1,2] + callback = _callback() + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=5) + + torch.testing.assert_close(model.net.weight.grad.norm(), torch.tensor(1.0)) + torch.testing.assert_close(model.net_fake_score.weight.grad, critic_grad_before) + assert model._distillation_parity_grad_metrics == { + "clip_grad_norm/net_selected_preclip": 5.0, + "clip_grad_norm/net_selected_clip_scale": pytest.approx(0.2), + "clip_grad_norm/net_selected_clip_norm": 1.0, + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_critic_phase_clips_only_fake_score_parameters() -> None: + model = _PhaseModel() + optimizer = PhaseOptimizer( + { + "net": torch.optim.SGD(model.net.parameters(), lr=0.1), + "fake_score": torch.optim.SGD(model.net_fake_score.parameters(), lr=0.1), + } + ) + model.net.weight.grad = torch.tensor([[6.0, 8.0]]) # [1,2] + model.net_fake_score.weight.grad = torch.tensor([[5.0, 12.0]]) # [1,2] + student_grad_before = model.net.weight.grad.clone() # [1,2] + callback = _callback() + callback.on_training_step_start(model, {"is_image": False}, iteration=6) + + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=6) + + torch.testing.assert_close(model.net.weight.grad, student_grad_before) + torch.testing.assert_close(model.net_fake_score.weight.grad.norm(), torch.tensor(1.0)) + assert model._distillation_parity_grad_metrics == { + "clip_grad_norm/fake_score_selected_preclip": 13.0, + "clip_grad_norm/fake_score_selected_clip_scale": pytest.approx(1.0 / 13.0), + "clip_grad_norm/fake_score_selected_clip_norm": 1.0, + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_partial_dtensor_norm_is_reduced_before_recording_parity_metrics() -> None: + model = _PhaseModel() + optimizer = PhaseOptimizer( + { + "net": torch.optim.SGD(model.net.parameters(), lr=0.1), + "fake_score": torch.optim.SGD(model.net_fake_score.parameters(), lr=0.1), + } + ) + model.net.weight.grad = torch.tensor([[3.0, 4.0]]) # [1,2] + partial_norm = _PartialNorm(local_value=2.0, global_value=5.0) + callback = _callback() + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + + with patch.object(grad_clip_module, "clip_grad_norm_", return_value=partial_norm): + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=5) + + assert partial_norm.full_tensor_calls == 1 + assert model._distillation_parity_grad_metrics == { + "clip_grad_norm/net_selected_preclip": 5.0, + "clip_grad_norm/net_selected_clip_scale": pytest.approx(1.0 / (5.0 + 1e-6)), + "clip_grad_norm/net_selected_clip_norm": 1.0, + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_logging_with_wandb_disabled_does_not_call_wandb_log(monkeypatch: pytest.MonkeyPatch) -> None: + model = _PhaseModel() + optimizer = PhaseOptimizer( + { + "net": torch.optim.SGD(model.net.parameters(), lr=0.1), + "fake_score": torch.optim.SGD(model.net_fake_score.parameters(), lr=0.1), + } + ) + model.net.weight.grad = torch.tensor([[3.0, 4.0]]) # [1,2] + callback = _callback() + callback.config.trainer.logging_iter = 1 + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + monkeypatch.setattr(grad_clip_module.wandb, "run", None) + wandb_log = MagicMock() + monkeypatch.setattr(grad_clip_module.wandb, "log", wandb_log) + + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=5) + + wandb_log.assert_not_called() + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_multiple_optimizer_phases_accumulate_parity_metrics_within_one_training_step() -> None: + model = _PhaseModel() + optimizer = PhaseOptimizer( + { + "net": torch.optim.SGD(model.net.parameters(), lr=0.1), + "fake_score": torch.optim.SGD(model.net_fake_score.parameters(), lr=0.1), + } + ) + model.net.weight.grad = torch.tensor([[3.0, 4.0]]) # [1,2] + model.net_fake_score.weight.grad = torch.tensor([[5.0, 12.0]]) # [1,2] + callback = _callback() + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=5) + model.optimizer_key_override = "fake_score" + callback.on_before_optimizer_step(model, optimizer, MagicMock(), MagicMock(), iteration=5) + + assert model._distillation_parity_grad_metrics == { + "clip_grad_norm/net_selected_preclip": 5.0, + "clip_grad_norm/net_selected_clip_scale": pytest.approx(0.2), + "clip_grad_norm/net_selected_clip_norm": 1.0, + "clip_grad_norm/fake_score_selected_preclip": 13.0, + "clip_grad_norm/fake_score_selected_clip_scale": pytest.approx(1.0 / 13.0), + "clip_grad_norm/fake_score_selected_clip_norm": 1.0, + } + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_new_training_step_clears_stale_parity_grad_metrics() -> None: + model = _PhaseModel() + model._distillation_parity_grad_metrics = {"stale": 1.0} + callback = _callback() + + callback.on_training_step_start(model, {"is_image": True}, iteration=5) + + assert model._distillation_parity_grad_metrics == {} diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py index 97de1fd8..88c9069e 100644 --- a/cosmos_framework/callbacks/hf_export.py +++ b/cosmos_framework/callbacks/hf_export.py @@ -356,6 +356,9 @@ def _do_save_and_upload( # hf_config.save_pretrained) are never overwritten. # HARD failure: a broken copy leaves the checkpoint unloadable, so any I/O error # propagates to the background-worker wrapper (same as shard writes). + # Native-format snapshots (e.g. nvidia/Cosmos3-Edge, model_type + # "cosmos3_edge") ship no remote-code .py files — nothing to copy is + # normal; the export loads through the framework-registered classes. if model_name_or_path and os.path.isdir(model_name_or_path): real_src = os.path.realpath(model_name_or_path) real_out = os.path.realpath(output_dir) @@ -389,6 +392,11 @@ def _do_save_and_upload( copied.append(os.path.join(rel_dir, fname) if rel_dir != "." else fname) if copied: log.info(f"[HFExportCallback] Copied missing files from source model: {copied}") + else: + log.info( + "[HFExportCallback] No missing .py/.json files to copy from source model " + "(expected for native-format snapshots, which ship no remote code)." + ) # 5. Tokenizer (best-effort — may fail for custom / gated models). if AutoTokenizer is not None and model_name_or_path: diff --git a/cosmos_framework/checkpoint/base.py b/cosmos_framework/checkpoint/base.py index 0c116017..f17ebd2d 100644 --- a/cosmos_framework/checkpoint/base.py +++ b/cosmos_framework/checkpoint/base.py @@ -43,7 +43,7 @@ def __init__( if not INTERNAL: from cosmos_framework.utils.checkpoint_db import download_checkpoint_v2 - if load_path: + if load_path and not self.load_from_object_store: load_path = download_checkpoint_v2(load_path) self.load_path = load_path self.load_training_state = config_checkpoint.load_training_state diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py index 43145c88..9f095ae9 100644 --- a/cosmos_framework/checkpoint/dcp.py +++ b/cosmos_framework/checkpoint/dcp.py @@ -71,9 +71,9 @@ from cosmos_framework.checkpoint.base import AbstractCheckpointer from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter -from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.model._base import ImaginaireModel from cosmos_framework.utils import callback, distributed, log, misc +from cosmos_framework.utils.config import CheckpointConfig, JobConfig from cosmos_framework.utils.easy_io import easy_io from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict diff --git a/cosmos_framework/checkpoint/dcp_distill.py b/cosmos_framework/checkpoint/dcp_distill.py new file mode 100644 index 00000000..da03e197 --- /dev/null +++ b/cosmos_framework/checkpoint/dcp_distill.py @@ -0,0 +1,389 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +""" +Distillation-tailored DCP for Cosmos3. + +Extends the Cosmos3 VFM DistributedCheckpointer to support multi-model / +multi-optimizer training (e.g., student + fake-score + discriminator). +""" + +import functools +import os +import time +from typing import Any + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) +from torch.distributed.checkpoint.stateful import Stateful +from torch.nn.modules.module import _IncompatibleKeys + +from cosmos_framework.model._base import ImaginaireModel +from cosmos_framework.utils import log, misc +from cosmos_framework.utils.easy_io import easy_io +from cosmos_framework.checkpoint.dcp import ( + AsyncMode, + CustomLoadPlanner, + _DataloaderWrapper, +) +from cosmos_framework.checkpoint.dcp import ( + DistributedCheckpointer as _DistributedCheckpointer, +) +from cosmos_framework.checkpoint.dcp import ModelWrapper as VFMModelWrapper +from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict +from cosmos_framework.model.generator.distillation.optimizer import OptimizerContainerLike, is_optimizer_container + +__all__: tuple[str, ...] = ( + "DistributedCheckpointer", + "ModelWrapper", + "OptimizerWrapper", +) + + +class OptimizerWrapper(Stateful): + """FSDP/DCP-aware optimizer state wrapper for per-phase distillation optimizers.""" + + def __init__( + self, + model: torch.nn.Module | list[torch.nn.Module], + optim: torch.optim.Optimizer | OptimizerContainerLike | list[torch.optim.Optimizer], + ) -> None: + self.optim_container: OptimizerContainerLike | None = optim if is_optimizer_container(optim) else None + if self.optim_container is not None: + self.model: list[torch.nn.Module] = [] + self.optim: list[torch.optim.Optimizer] = [] + return + + self.model = [model] if isinstance(model, torch.nn.Module) else model + self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim + if len(self.model) != len(self.optim): + raise ValueError(f"Expected matched model/optimizer lists, got {len(self.model)} and {len(self.optim)}") + + def state_dict(self) -> dict[str, Any]: + if self.optim_container is not None: + return self.optim_container.state_dict() + + func = functools.partial( + get_optimizer_state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + return {k: v for sd in map(func, self.model, self.optim) for k, v in sd.items()} + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if self.optim_container is not None: + self.optim_container.load_state_dict(state_dict) + return + + func = functools.partial( + set_optimizer_state_dict, + optim_state_dict=state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + list(map(func, self.model, self.optim)) + + +class ModelWrapper(VFMModelWrapper): + """ + Modify the Cosmos3 VFM model wrapper to exclude + the teacher weights from the model state dict. + """ + + def __init__( + self, + model: ImaginaireModel, + exclude_teacher_weights: bool = True, + strict_resume: bool = True, + ) -> None: + super().__init__(model) + self.exclude_teacher_weights: bool = exclude_teacher_weights + self.strict_resume: bool = strict_resume + + def _unregistered_fake_score(self) -> torch.nn.Module | None: + """Return ``net_fake_score`` only when it is held outside the module registry. + + Self-forcing DMD and dense DMD2 hide ``net_fake_score`` from the module + registry (so inference ``get_model_state_dict`` stays student-only); it + must therefore be serialized explicitly. Older layouts may have registered + ``net_fake_score``, so those must NOT be double-handled. + """ + if "net_fake_score" in dict(self.model.named_children()): + return None + fake_score = getattr(self.model, "net_fake_score", None) + return fake_score if isinstance(fake_score, torch.nn.Module) else None + + def state_dict(self) -> dict[str, Any]: + sd = super().state_dict() + if self.exclude_teacher_weights: + sd = {k: v for k, v in sd.items() if not k.startswith("net_teacher.")} + fake_score = self._unregistered_fake_score() + if fake_score is not None: + for k, v in get_model_state_dict(fake_score).items(): + sd[f"net_fake_score.{k}"] = v + return sd + + def load_state_dict(self, state_dict: dict[str, Any]) -> _IncompatibleKeys: + # DCP/FSDP-aware load path + fake_score = self._unregistered_fake_score() + fake_score_results: _IncompatibleKeys | None = None + main_state_dict = state_dict + if fake_score is not None: + prefix = "net_fake_score." + fake_score_state_dict = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} + main_state_dict = {k: v for k, v in state_dict.items() if not k.startswith(prefix)} + fake_score_results = set_model_state_dict( + model=fake_score, + model_state_dict=fake_score_state_dict, + options=StateDictOptions(strict=False, ignore_frozen_params=False), + ) + + results = set_model_state_dict( + model=self.model, + model_state_dict=main_state_dict, + options=StateDictOptions(strict=False, ignore_frozen_params=False), + ) + + if self.strict_resume: + bad_missing = [k for k in results.missing_keys if not k.startswith("net_teacher.")] + bad_unexpected = [k for k in results.unexpected_keys if not k.startswith("net_teacher.")] + if fake_score_results is not None: + bad_missing += list(fake_score_results.missing_keys) + bad_unexpected += list(fake_score_results.unexpected_keys) + if bad_missing or bad_unexpected: + raise ValueError( + f"Strict resume failed. Missing(non-teacher)={bad_missing[:20]}, " + f"Unexpected(non-teacher)={bad_unexpected[:20]}" + ) + return results + + +class DistributedCheckpointer(_DistributedCheckpointer): + @misc.timer("checkpoint loading") + def load( + self, + model: ImaginaireModel, + optimizer: Any = None, + scheduler: Any = None, + grad_scaler: torch.amp.GradScaler | None = None, + ) -> int: + if self.callbacks is not None: + self.callbacks.on_load_checkpoint_start(model) + + model_dict = model.model_dict() + + resume_keys, checkpoint_path, warm_start = self.keys_to_resume_during_load() + resume_keys = sorted(resume_keys) + log.critical(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}") + + iteration = 0 + + if checkpoint_path is not None: + self._check_checkpoint_exists(checkpoint_path) + for key in resume_keys: + dist.barrier() + + cur_key_ckpt_full_path = os.path.join(checkpoint_path, key) + log.critical(f"Start loading checkpoint from {cur_key_ckpt_full_path}") + + strict_resume = self.config_checkpoint.strict_resume + keys_to_skip_loading = self.config_checkpoint.keys_to_skip_loading if warm_start else [] + load_planner = CustomLoadPlanner( + allow_partial_load=not strict_resume, + keys_to_skip_loading=keys_to_skip_loading, + ) + + if key == "model": + storage_reader = self.get_storage_reader(cur_key_ckpt_full_path) + log.info("- Loading the model...") + _model_wrapper = ModelWrapper( + model, + exclude_teacher_weights=True, + strict_resume=strict_resume, + ) + _state_dict = _model_wrapper.state_dict() + dcp.load(_state_dict, storage_reader=storage_reader, planner=load_planner) + _model_wrapper.load_state_dict(_state_dict) + + elif key == "optim": + if optimizer is None: + raise ValueError("optimizer must be provided when loading optim state.") + for optim_key, optim in optimizer.items(): + storage_reader = self.get_storage_reader(f"{cur_key_ckpt_full_path}_{optim_key}") + log.info(f"- Loading the optimizer ({optim_key})...") + _optim_wrapper = OptimizerWrapper(model_dict[optim_key], optim) + _state_dict = _optim_wrapper.state_dict() + optim_load_planner = CustomLoadPlanner( + allow_partial_load=not strict_resume, + keys_to_skip_loading=[], + ) + dcp.load( + _state_dict, + storage_reader=storage_reader, + planner=optim_load_planner, + ) + _optim_wrapper.load_state_dict(_state_dict) + + elif key == "scheduler": + if scheduler is None: + raise ValueError("scheduler must be provided when loading scheduler state.") + for sched_key, sched in scheduler.items(): + storage_reader = self.get_storage_reader(f"{cur_key_ckpt_full_path}_{sched_key}") + log.info(f"- Loading the scheduler ({sched_key})...") + _state_dict = sched.state_dict() + dcp.load( + _state_dict, + storage_reader=storage_reader, + planner=load_planner, + ) + sched.load_state_dict(_state_dict) + + elif key == "trainer": + if grad_scaler is None: + raise ValueError("grad_scaler must be provided when loading trainer state.") + storage_reader = self.get_storage_reader(cur_key_ckpt_full_path) + log.info("- Loading the trainer...") + + rng_key = f"rng_state_{dist.get_rank()}" + current_rng_state = get_rand_state_dict() + _state_dict = { + "grad_scaler": grad_scaler.state_dict(), + "iteration": iteration, + } + metadata = storage_reader.read_metadata() + rng_key_exists = any( + k.startswith(f"{rng_key}.") or k == rng_key for k in metadata.state_dict_metadata.keys() + ) + if rng_key_exists: + _state_dict[rng_key] = current_rng_state + + dcp.load( + _state_dict, + storage_reader=storage_reader, + planner=load_planner, + ) + grad_scaler.load_state_dict(_state_dict["grad_scaler"]) + iteration = _state_dict["iteration"] + set_rand_state_dict(_state_dict.get(rng_key, current_rng_state)) + + elif key == "dataloader": + # Per-rank dataloader-state resume (main's base checkpointer). Skips + # gracefully when no dataloader-state callback is registered or no + # dataloader checkpoint was written. + if not easy_io.exists(cur_key_ckpt_full_path, backend_key=self.load_s3_backend_key): + log.info( + f"Checkpoint {cur_key_ckpt_full_path} does not exist, skip loading dataloader.", + rank0_only=False, + ) + continue + rank = dist.get_rank() + dataloader_pkl_path = os.path.join(cur_key_ckpt_full_path, f"rank_{rank}.pkl") + if not easy_io.exists(dataloader_pkl_path, backend_key=self.load_s3_backend_key): + log.info(f"No dataloader checkpoint found at {dataloader_pkl_path}", rank0_only=False) + continue + log.info(f"- Loading the dataloader {cur_key_ckpt_full_path}...", rank0_only=False) + _state_dict = easy_io.load( + dataloader_pkl_path, + file_format="pkl", + backend_key=self.load_s3_backend_key, + ) + dataloader_wrapper = _DataloaderWrapper(self.callbacks) + if dataloader_wrapper.has_state(): + dataloader_wrapper.load_state_dict(_state_dict) + + else: + raise ValueError(f"Invalid key: {key}. not support to resume.") + + if self.callbacks is not None: + self.callbacks.on_load_checkpoint(model, state_dict=_state_dict) + log.info(f"Loaded checkpoint from {checkpoint_path} in iteration {iteration}") + else: + log.info("Training from scratch.") + + torch.cuda.empty_cache() + + if self.callbacks is not None: + self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path) + return iteration + + # No custom save_state_dict_worker needed: the save() method below + # pre-flattens optim/scheduler into per-key top-level entries so the + # parent's save_state_dict_worker (used by both sync and async paths) + # writes each to its own subdirectory (e.g. optim_net/, optim_fake_score/). + + def save( + self, + model: ImaginaireModel, + optimizer: Any, + scheduler: Any = None, + grad_scaler: torch.amp.GradScaler | None = None, + iteration: int = 0, + ) -> None: + if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM: + self._wait_for_previous_async_checkpoint() + + if self.callbacks is not None: + self.callbacks.on_save_checkpoint_start(model, iteration) + + model_dict = model.model_dict() + checkpoint_file = f"iter_{iteration:09}" + + rng_key = f"rng_state_{dist.get_rank()}" + to_save_dict: dict[str, Any] = { + "model": ModelWrapper(model, exclude_teacher_weights=True).state_dict(), + "trainer": { + "grad_scaler": grad_scaler.state_dict(), + "iteration": iteration, + rng_key: get_rand_state_dict(), + }, + } + # Flatten optimizer and scheduler dicts into separate top-level keys + # (e.g. "optim_net", "optim_fake_score", "scheduler_net", …) so that + # the parent's save_state_dict_worker — which is used by the async + # background process — creates the correct per-key subdirectories. + for optim_key, optim in optimizer.items(): + to_save_dict[f"optim_{optim_key}"] = OptimizerWrapper(model_dict[optim_key], optim).state_dict() + for sched_key, sched in scheduler.items(): + to_save_dict[f"scheduler_{sched_key}"] = sched.state_dict() + + # Per-rank dataloader state (main's base checkpointer). No-op unless a + # callback tagged ``checkpoint_component=="dataloader"`` is registered; + # the inherited save_state_dict_worker writes the "dataloader" key as pkl. + dataloader_wrapper = _DataloaderWrapper(self.callbacks) + if dataloader_wrapper.has_state(): + to_save_dict["dataloader"] = dataloader_wrapper.state_dict() + + for key in list(to_save_dict.keys()): + output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{key}") + to_save_dict[key] = (to_save_dict[key], output_dirname) + + if self.callbacks is not None: + self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict) + + log.info(f"Saving checkpoint to {os.path.join(self.save_dirname, checkpoint_file)}") + + if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM: + # PhaseOptimizer state is structurally different before and after a + # phase's first optimizer step. Recreate the pinned staging tree for + # each distillation save so DCP does not try to copy new optimizer + # slots into a stale CPU companion structure. + self.cpu_offload_state_dict = None + self._checkpoint_async_with_pinned_memory(checkpoint_file, to_save_dict) + else: + start_time = time.monotonic() + try: + self.save_state_dict_worker(to_save_dict, checkpoint_file) + finally: + if self.callbacks is not None: + self.callbacks.on_save_checkpoint_success( + iteration=iteration, elapsed_time=time.monotonic() - start_time + ) + + if self.callbacks is not None: + self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration) diff --git a/cosmos_framework/checkpoint/dcp_distill_test.py b/cosmos_framework/checkpoint/dcp_distill_test.py new file mode 100644 index 00000000..24cb46cc --- /dev/null +++ b/cosmos_framework/checkpoint/dcp_distill_test.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""CPU tests for the distillation checkpointer ModelWrapper fake-score handling.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +import cosmos_framework.checkpoint.dcp_distill as dcp_distill_module +from cosmos_framework.checkpoint.dcp import AsyncMode +from cosmos_framework.checkpoint.dcp_distill import ( + DistributedCheckpointer, + ModelWrapper, + OptimizerWrapper, +) + + +class _ListWrappedFakeScoreModel(nn.Module): + """Mimics SF-DMD: a registered student ``net`` and an unregistered fake-score net.""" + + def __init__(self) -> None: + super().__init__() + self.net = nn.Linear(2, 2) # registered student + self._fake_score_holder = [nn.Linear(2, 2)] # held outside the registry + + @property + def net_fake_score(self) -> nn.Module: + return self._fake_score_holder[0] + + +class _RegisteredFakeScoreModel(nn.Module): + """Legacy compatibility case: ``net_fake_score`` is a registered submodule.""" + + def __init__(self) -> None: + super().__init__() + self.net = nn.Linear(2, 2) + self.net_fake_score = nn.Linear(2, 2) + + +class _StrictRejectingModel(nn.Module): + """Mimics OmniMoTModel: ``load_state_dict`` rejects strict=True. + + The fake-score net is held outside the registry and loaded explicitly, and the + model-level load (reached via ``set_model_state_dict`` during resume) must call + ``super().load_state_dict`` with ``strict=False``. + """ + + def __init__(self) -> None: + super().__init__() + self.net = nn.Linear(2, 2) + self._fake_score_holder = [nn.Linear(2, 2)] + + @property + def net_fake_score(self) -> nn.Module: + return self._fake_score_holder[0] + + def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False): # type: ignore[override] + if strict: + raise ValueError("Strict mode is not supported") + return super().load_state_dict(state_dict, strict=False, assign=assign) + + +class _SaveContractModel(_ListWrappedFakeScoreModel): + """Minimal multi-network model exposing the checkpointer's phase map.""" + + def __init__(self) -> None: + super().__init__() + self.net_teacher = nn.Linear(2, 2) + + def model_dict(self) -> dict[str, nn.Module]: + return {"net": self.net, "fake_score": self.net_fake_score} + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_checkpointer_exports_are_explicit() -> None: + assert dcp_distill_module.__all__ == ( + "DistributedCheckpointer", + "ModelWrapper", + "OptimizerWrapper", + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_optimizer_wrapper_accepts_structurally_compatible_foreign_container() -> None: + state_dict = {"foreign": True} + foreign_container = SimpleNamespace( + model=nn.Linear(2, 2), + optimizers=[MagicMock()], + step=MagicMock(), + zero_grad=MagicMock(), + state_dict=MagicMock(return_value=state_dict), + load_state_dict=MagicMock(), + ) + + wrapper = OptimizerWrapper(foreign_container.model, foreign_container) + + assert wrapper.state_dict() == state_dict + wrapper.load_state_dict(state_dict) + foreign_container.load_state_dict.assert_called_once_with(state_dict) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_model_wrapper_persists_unregistered_fake_score() -> None: + model = _ListWrappedFakeScoreModel() + + # net_fake_score is hidden from the module registry (inference loads stay student-only) + assert "net_fake_score" not in dict(model.named_children()) + + state_dict = ModelWrapper(model).state_dict() + assert any(k.startswith("net.") for k in state_dict) + assert any(k.startswith("net_fake_score.") for k in state_dict) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_model_wrapper_round_trips_unregistered_fake_score() -> None: + src = _ListWrappedFakeScoreModel() + with torch.no_grad(): + for param in src.net_fake_score.parameters(): + param.add_(1.0) + saved = {k: v.clone() for k, v in ModelWrapper(src).state_dict().items()} + + dst = _ListWrappedFakeScoreModel() + ModelWrapper(dst, strict_resume=True).load_state_dict(saved) + + for src_param, dst_param in zip(src.net_fake_score.parameters(), dst.net_fake_score.parameters()): + torch.testing.assert_close(src_param, dst_param) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_model_wrapper_load_uses_non_strict_for_strict_rejecting_model() -> None: + # Resume path: ModelWrapper.load_state_dict -> set_model_state_dict(model) -> + # model.load_state_dict, which must not pass strict=True to a model that + # rejects it (OmniMoTModel). This previously raised on resume. + src = _StrictRejectingModel() + with torch.no_grad(): + for param in src.net_fake_score.parameters(): + param.add_(0.5) + saved = {k: v.clone() for k, v in ModelWrapper(src).state_dict().items()} + + dst = _StrictRejectingModel() + ModelWrapper(dst, strict_resume=True).load_state_dict(saved) + + for src_param, dst_param in zip(src.net_fake_score.parameters(), dst.net_fake_score.parameters()): + torch.testing.assert_close(src_param, dst_param) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_model_wrapper_does_not_double_handle_registered_fake_score() -> None: + model = _RegisteredFakeScoreModel() + + # Registered fake-score already appears via get_model_state_dict; the explicit + # path must be skipped for compatibility with older dense DMD2 layouts. + assert "net_fake_score" in dict(model.named_children()) + state_dict = ModelWrapper(model).state_dict() + fake_score_keys = [k for k in state_dict if k.startswith("net_fake_score.")] + # Exactly the registry keys, with no duplicated explicit entries. + assert fake_score_keys + assert len(fake_score_keys) == len(set(fake_score_keys)) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_model_wrapper_excludes_registered_teacher_weights() -> None: + model = _SaveContractModel() + + state_dict = ModelWrapper(model, exclude_teacher_weights=True).state_dict() + + assert any(key.startswith("net.") for key in state_dict) + assert any(key.startswith("net_fake_score.") for key in state_dict) + assert not any(key.startswith("net_teacher.") for key in state_dict) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_strict_resume_rejects_missing_fake_score_tensor() -> None: + source = _ListWrappedFakeScoreModel() + state_dict = {key: value.clone() for key, value in ModelWrapper(source).state_dict().items()} + missing_key = next(key for key in state_dict if key.startswith("net_fake_score.")) + del state_dict[missing_key] + + destination = _ListWrappedFakeScoreModel() + with pytest.raises(ValueError, match="Strict resume failed"): + ModelWrapper(destination, strict_resume=True).load_state_dict(state_dict) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("async_mode", [AsyncMode.DISABLED, AsyncMode.ASYNC_WITH_PINNED_MEM]) +def test_save_emits_exact_distillation_components(async_mode: AsyncMode) -> None: + model = _SaveContractModel() + optimizers = { + "net": torch.optim.AdamW(model.net.parameters(), lr=1e-3), + "fake_score": torch.optim.AdamW(model.net_fake_score.parameters(), lr=1e-3), + } + schedulers = {"net": MagicMock(), "fake_score": MagicMock()} + grad_scaler = MagicMock(spec=torch.amp.GradScaler) + grad_scaler.state_dict.return_value = {"scale": 1.0} + + checkpointer = object.__new__(DistributedCheckpointer) + checkpointer.async_mode = async_mode + checkpointer.callbacks = None + checkpointer.save_to_object_store = False + checkpointer._local_dirname = "/tmp/distillation-checkpoints" + checkpointer.cpu_offload_state_dict = object() + checkpointer._wait_for_previous_async_checkpoint = MagicMock() + checkpointer._checkpoint_async_with_pinned_memory = MagicMock() + checkpointer.save_state_dict_worker = MagicMock() + + dataloader_wrapper = MagicMock() + dataloader_wrapper.has_state.return_value = True + dataloader_wrapper.state_dict.return_value = {"cursor": 7} + with ( + patch.object(dcp_distill_module, "_DataloaderWrapper", return_value=dataloader_wrapper), + patch.object(dcp_distill_module.dist, "get_rank", return_value=0), + patch.object(dcp_distill_module, "get_rand_state_dict", return_value={"torch_rng_state": torch.zeros(1)}), + ): + checkpointer.save(model, optimizers, schedulers, grad_scaler, iteration=7) + + if async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM: + checkpointer._wait_for_previous_async_checkpoint.assert_called_once_with() + checkpointer._checkpoint_async_with_pinned_memory.assert_called_once() + to_save_dict = checkpointer._checkpoint_async_with_pinned_memory.call_args.args[1] + assert checkpointer.cpu_offload_state_dict is None + else: + checkpointer.save_state_dict_worker.assert_called_once() + to_save_dict = checkpointer.save_state_dict_worker.call_args.args[0] + + assert set(to_save_dict) == { + "model", + "trainer", + "optim_net", + "optim_fake_score", + "scheduler_net", + "scheduler_fake_score", + "dataloader", + } + model_state = to_save_dict["model"][0] + assert any(key.startswith("net.") for key in model_state) + assert any(key.startswith("net_fake_score.") for key in model_state) + assert not any(key.startswith("net_teacher.") for key in model_state) diff --git a/cosmos_framework/configs/base/config.py b/cosmos_framework/configs/base/config.py index ec65e32e..cf672082 100644 --- a/cosmos_framework/configs/base/config.py +++ b/cosmos_framework/configs/base/config.py @@ -99,5 +99,6 @@ def make_config() -> Config: import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_nano # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_nano # noqa: F401 import cosmos_framework.configs.base.experiment.sft.vision_sft_super # noqa: F401 + import cosmos_framework.configs.base.experiment.sft.vision_sft_edge # noqa: F401 return c diff --git a/cosmos_framework/configs/base/defaults/model_config.py b/cosmos_framework/configs/base/defaults/model_config.py index 5e6fd6b4..c82fee5f 100644 --- a/cosmos_framework/configs/base/defaults/model_config.py +++ b/cosmos_framework/configs/base/defaults/model_config.py @@ -255,6 +255,12 @@ class OmniMoTModelConfig: sound_dim: int | None = None # Sound latent channel size (e.g., 64 for AVAE 48kHz) sound_latent_fps: int = 25 # Sound tokenizer's latent rate (e.g., 48kHz / 1920 hop = 25 Hz) + # When False, removes bias from vae2llm, sound2llm, and the two Linear layers inside + # time_embedder. These biases seem to inject token-constant DC offsets that dominate + # the MoE router input and create prompt invariant routing. This is observed empirically + # in the 30B-A3B checkpoint. + enable_input_bias: bool = True + log_enc_time_every_n: int = 100 # Frequency of logging encoding time to W&B # When True, ``OmniMoTModel.state_dict`` / ``load_state_dict`` skip the diff --git a/cosmos_framework/configs/base/defaults/reasoner.py b/cosmos_framework/configs/base/defaults/reasoner.py index 09e92efe..0a85364b 100644 --- a/cosmos_framework/configs/base/defaults/reasoner.py +++ b/cosmos_framework/configs/base/defaults/reasoner.py @@ -710,7 +710,9 @@ class VLMConfig: ) # Cosmos3-Edge-Reasoner at commit 4acb717. -# nemotron_siglip2 architecture: Nemotron text backbone (56-block hybrid layout, 2048 hidden) +# Edge reasoner architecture (model_type "cosmos3_edge" in the public +# nvidia/Cosmos3-Edge; historically remote-code "nemotron_siglip2"): +# Nemotron text backbone (56-block hybrid layout, 2048 hidden) # + SigLIP2 vision encoder. The text transformer is identical in shape to # Nemotron-3-Dense-VL-2B (hidden_size=2048, 56 alternating attn/MLP blocks → 28 # effective MoT layers after _transform_text_dict). Uses the same @@ -738,7 +740,7 @@ class VLMConfig: ) # Cosmos3-Edge-Reasoner at commit 9b4c028 (2026-05-29). -# Same nemotron_siglip2 architecture as 4acb717; new weights uploaded 2026-05-29. +# Same Edge reasoner (cosmos3_edge) architecture as 4acb717; new weights uploaded 2026-05-29. Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028: VLMConfig = VLMConfig( model_name="nvidia/Cosmos3-Edge-Reasoner", model_instance=L(Nemotron3DenseVLTextForCausalLM)( @@ -762,7 +764,8 @@ class VLMConfig: ) # Cosmos3-Edge-Reasoner at commit 590c1c0 (2026-06-28). -# Updated weights uploaded 2026-06-28. +# Updated weights uploaded 2026-06-28; bit-identical to the reasoner tower +# shipped inside the public nvidia/Cosmos3-Edge release. Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0: VLMConfig = VLMConfig( model_name="nvidia/Cosmos3-Edge-Reasoner", model_instance=L(Nemotron3DenseVLTextForCausalLM)( diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py index 6ce66fb9..dfca0b93 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py @@ -14,7 +14,7 @@ BASE_CHECKPOINT_PATH= \\ WAN_VAE_PATH= \\ torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \\ - --sft-toml examples/toml/sft_config/action_policy_droid_repro.toml + --sft-toml examples/toml/sft_config/action_policy_droid_nano.toml """ import copy diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py index 8cec6a0b..7451406f 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_all_nano.py @@ -6,7 +6,7 @@ Feeds ``LIBEROLeRobotDataset`` (frame-wise-relative rot6d, ``quantile_rot``, concat_view third-person + wrist) and trains the generation + action heads from the public ``nvidia/Cosmos3-Nano`` base. Trains on all 4 LIBERO suites (equal mix); ``LIBERO_ROOT`` is the -LIBERO_LeRobot_v3 parent dir. See docs/action_policy_libero_sft.md. +LIBERO_LeRobot_v3 parent dir. See docs/action_policy_libero_posttrain.md. """ import copy diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py index 594c3460..2ec89580 100644 --- a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py +++ b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py @@ -6,7 +6,7 @@ Feeds ``LIBEROLeRobotDataset`` (frame-wise-relative rot6d, ``quantile_rot``, concat_view third-person + wrist) and trains the generation + action heads from the public ``nvidia/Cosmos3-Nano`` base. Train on ``libero_10`` alone -(``LIBERO_ROOT``). See docs/action_policy_libero_sft.md. +(``LIBERO_ROOT``). See docs/action_policy_libero_posttrain.md. """ import copy diff --git a/cosmos_framework/configs/base/experiment/distillation/__init__.py b/cosmos_framework/configs/base/experiment/distillation/__init__.py new file mode 100644 index 00000000..99124078 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/distillation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Interactive training-method configuration schemas.""" + +__all__: tuple[str, ...] = () diff --git a/cosmos_framework/configs/base/experiment/distillation/dmd2_config.py b/cosmos_framework/configs/base/experiment/distillation/dmd2_config.py new file mode 100644 index 00000000..4888a0f9 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/distillation/dmd2_config.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from typing import Literal + +import attrs +from hydra.core.config_store import ConfigStore # type: ignore[import] + +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict +from cosmos_framework.configs.base.defaults.model_config import FixedStepSamplerConfig, OmniMoTModelConfig +from cosmos_framework.configs.base.defaults.reasoner import PretrainedWeightsConfig, VLMConfig + +# Cosmos3 way of registering/configuring optimizers +from cosmos_framework.utils.generator.optimizer import build_optimizer + +__all__: tuple[str, ...] = ( + "DMD2OptimizerConfig", + "DMD2RFConfig", + "register_dmd2_optimizer", +) + +IS_PREPROCESSED_KEY: str = "is_preprocessed" + +_STANDARD_KEYS_TO_SELECT: list[str] = [ + "moe_gen", + "q_norm", + "k_norm", + "time_embedder", + "vae2llm", + "llm2vae", +] + + +@attrs.define(slots=False) +class DMD2OptimizerConfig: + """Unified optimizer config for DMD2 student (net) and critic (fake_score) networks. + + Both fields are LazyDicts wrapping ``build_optimizer`` — they are instantiated at training + time via ``lazy_instantiate(config.net, model=self.net)``. Experiment overrides go under + ``optimizer=dict(net=dict(...), fake_score=dict(...))``. + """ + + net: LazyDict = L(build_optimizer)( + model=None, + optimizer_type="FusedAdam", + lr=1e-6, + weight_decay=0.01, + betas=[0.9, 0.99], + fused=True, + eps=1e-8, + keys_to_select=_STANDARD_KEYS_TO_SELECT, + lr_multipliers={}, + ) + fake_score: LazyDict = L(build_optimizer)( + model=None, + optimizer_type="FusedAdam", + lr=2e-7, + weight_decay=0.01, + betas=[0.0, 0.999], + fused=True, + eps=1e-8, + keys_to_select=_STANDARD_KEYS_TO_SELECT, + lr_multipliers={}, + ) + + +def register_dmd2_optimizer() -> None: + cs = ConfigStore.instance() + cs.store(group="optimizer", package="optimizer", name="dmd2", node=DMD2OptimizerConfig()) + + +@attrs.define(slots=False) +class DMD2RFConfig(OmniMoTModelConfig): + """ + Config for DMD2RF model. + + Inherits all base fields from ``OmniMoTModelConfig`` and adds + DMD2RF-specific knobs (teacher/fake-score, optimizers, sampler schedule, etc.). + """ + + # The student's vlm_config is inherited from OmniMoTModelConfig (model.config.vlm_config). + # Teacher and fake-score use separate VLMConfig instances because they may differ in size + # or checkpoint path from the student; all three default with pretrained_weights.enabled=False + # since DMD2 initialises them from teacher_load_from rather than the standard pretrain path. + vlm_config: VLMConfig = VLMConfig(pretrained_weights=PretrainedWeightsConfig(enabled=False)) + vlm_config_teacher: VLMConfig = VLMConfig(pretrained_weights=PretrainedWeightsConfig(enabled=False)) + vlm_config_fake_score: VLMConfig = VLMConfig(pretrained_weights=PretrainedWeightsConfig(enabled=False)) + + # ---------------- Fixed-step sampler schedule ---------------- + # Controls the discrete sigma schedule used for student inference and for sampling + # the training noise level in _sample_student_sigma. See FixedStepSamplerConfig. + fixed_step_sampler_config: FixedStepSamplerConfig = FixedStepSamplerConfig() + + # ---------------- Distillation / loss scheduling ---------------- + loss_scale_fake_score: float = 1.0 + loss_scale_sid: float = 1.0 + noise_level_parameterization: Literal["rectified_flow"] = "rectified_flow" + # "x0" uses FastGen's original clean-data VSD direction. + # "velocity" applies the raw RF velocity direction to the generated x0 tensor. + vsd_gradient_space: Literal["x0", "velocity"] = "x0" + # "mean" preserves the existing active-element mean loss. "sum" preserves + # the existing half-MSE sum. "sum_rcm" matches RCM's no-half-factor + # pseudo-target reduction and normalizer clamp. + vsd_loss_reduction: Literal["mean", "sum", "sum_rcm"] = "mean" + # "active_mean" preserves the existing fake-score FM loss: sum over generated + # elements divided by active element count. "sum_rcm" matches RCM's + # non-causal DMD critic loss by summing generated elements per instance. + fake_score_loss_reduction: Literal["active_mean", "sum_rcm"] = "active_mean" + student_update_freq: int = 5 + teacher_guidance: float = 1.0 + # Prompt used for the teacher CFG unconditional / negative branch. The empty + # default preserves the existing classifier-free guidance behavior. + teacher_negative_prompt: str = "" + # Preserve existing DMD2 clipping behavior by default; experiments can set + # this False to match RCM-style no-clipping runs. + grad_clip: bool = True + warmup_student_steps: int = 0 # Number of iterations of student-only phase before alternating with critic updates + warmup_critic_steps: int = 0 # Number of iterations of critic-only phase before alternating with student updates + + # ---------------- Student simulation mode ---------------- + # "forward": fastgen-style — single pass from a randomly sampled sigma (current default). + # "backward": rcm-style multi-step rollout from a pure-noise start (t_list[0] must be 1.0) + # over an iteration-cycled number of t_list steps, descending to sigma=0. + simulation_mode: Literal["forward", "backward"] = "forward" + # Number of trailing denoising steps in backward simulation that receive gradients. + # 1 = only the last step (low memory, matches predict2_distill). + # -1 = all steps (true BPTT, O(N) memory). + backward_grad_steps: int = 1 + + # ---------------- Model architecture / checkpointing ---------------- + # Note: In cosmos3 we use VLMConfig in DMD2RFConfig to instantiate the network. + # The teacher/fakescore/student use same instantiation function to build the net 3 times. + teacher_load_from: LazyDict | None = None + student_load_from: LazyDict | None = None + # Load teacher ckpt and copy weights into student/fake-score nets (train-time only) + load_teacher_weights: bool = True + + # ---------------- misc ---------------- + vis_debug: bool = False # Flag for visualizing intermediate results during training + + # ---------------- Enforcing value of certain fields upon creation of the config ---------------- + def __attrs_post_init__(self) -> None: + assert not (self.warmup_student_steps > 0 and self.warmup_critic_steps > 0), ( + "Only one of warmup_student_steps and warmup_critic_steps can be nonzero, " + f"got warmup_student_steps={self.warmup_student_steps}, warmup_critic_steps={self.warmup_critic_steps}" + ) + # force-disabling discriminator for Cosmos3 for now. + # Discriminator relies on intermediate features of the transformer, which is + # not yet implemented in Cosmos3. + object.__setattr__(self, "loss_scale_GAN_discriminator", 0.0) + object.__setattr__(self, "loss_scale_GAN_generator", 0.0) + object.__setattr__(self, "net_discriminator_head", None) + object.__setattr__(self, "intermediate_feature_ids", None) + object.__setattr__(self, "optimizer_discriminator_config", None) diff --git a/cosmos_framework/configs/base/experiment/distillation/dmd2_config_test.py b/cosmos_framework/configs/base/experiment/distillation/dmd2_config_test.py new file mode 100644 index 00000000..bb8f6a17 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/distillation/dmd2_config_test.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import pytest + +import cosmos_framework.configs.base.experiment.distillation.dmd2_config as dmd2_config_module +from cosmos_framework.configs.base.experiment.distillation.dmd2_config import DMD2OptimizerConfig, DMD2RFConfig + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_config_exports_are_explicit() -> None: + assert dmd2_config_module.__all__ == ( + "DMD2OptimizerConfig", + "DMD2RFConfig", + "register_dmd2_optimizer", + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_dmd2_defaults_preserve_export_and_inference_contract() -> None: + config = DMD2RFConfig() + + assert config.fixed_step_sampler_config.t_list == [0.999, 0.75, 0.5, 0.25] + assert config.fixed_step_sampler_config.sample_type == "sde" + assert config.student_update_freq == 5 + assert config.simulation_mode == "forward" + assert config.backward_grad_steps == 1 + assert config.vsd_gradient_space == "x0" + assert config.vsd_loss_reduction == "mean" + assert config.fake_score_loss_reduction == "active_mean" + assert config.load_teacher_weights is True + assert config.teacher_load_from is None + assert config.student_load_from is None + assert config.action_gen is False + assert config.sound_gen is False + assert config.vlm_config.pretrained_weights.enabled is False + assert config.vlm_config_teacher.pretrained_weights.enabled is False + assert config.vlm_config_fake_score.pretrained_weights.enabled is False + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_dmd2_optimizer_defaults_preserve_separate_student_and_critic_routes() -> None: + config = DMD2OptimizerConfig() + + assert config.net.model is None + assert config.net.optimizer_type == "FusedAdam" + assert config.net.lr == 1e-6 + assert config.net.betas == [0.9, 0.99] + assert config.fake_score.model is None + assert config.fake_score.optimizer_type == "FusedAdam" + assert config.fake_score.lr == 2e-7 + assert config.fake_score.betas == [0.0, 0.999] diff --git a/cosmos_framework/configs/base/experiment/sft/models/edge_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/edge_model_config.py new file mode 100644 index 00000000..35acedad --- /dev/null +++ b/cosmos_framework/configs/base/experiment/sft/models/edge_model_config.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Shared Edge-tier (Nemotron-2B-Dense-VL) ``model.config`` baseline for SFT experiments. + +Consumers must ``copy.deepcopy`` this constant before mutating it. Baseline +mirrors ``vision_sft_edge`` (HF-cluster deployment with empty tokenizer/vlm +paths, video-style loss scales, ``load_weights_from_pretrained=True``), except +``action_gen``: the baseline keeps the Cosmos3-Edge.yaml value (``True``) and +``vision_sft_edge`` overrides it to ``False`` (no action tokens in vision SFT). + +Derived from ``nano_model_config.NANO_MODEL_CONFIG``; every field is identical to +the Nano baseline except the Cosmos3-Edge deltas sourced verbatim from +``inference/configs/model/Cosmos3-Edge.yaml``: + * ``vlm_config`` swapped from the Qwen3-VL backbone to the Nemotron-3 Dense VL + backbone (``Nemotron3DenseVLTextForCausalLM`` / ``Nemotron3DenseVLMoTConfig``, + ``Nemotron-2B-Dense-VL.json``, ``nvidia/Cosmos3-Edge-Reasoner``, + ``build_processor_lazy(repository="nvidia/Cosmos3-Edge")`` tokenizer, + ``qk_norm_for_text=False``, ``include_visual=None``, ``layer_module=None``). + * ``resolution`` ``"720"`` -> ``"480"`` (Edge native inference res). + * ``rectified_flow_training_config``: ``loss_scale`` ``1.0`` -> ``10.0``, + ``image_loss_scale`` ``1.0`` -> ``None``, plus the Edge-only ``sound_loss_scale`` + delta (dataclass-supported keys only). +""" + +from cosmos_framework.configs.base.defaults.reasoner import create_vlm_config +from cosmos_framework.data.generator.processors import build_processor_lazy +from cosmos_framework.model.generator.mot.unified_mot import ( + Nemotron3DenseVLMoTConfig, + Nemotron3DenseVLTextForCausalLM, +) +from cosmos_framework.utils.lazy_config import LazyCall as L + +EDGE_MODEL_CONFIG = dict( + # Mirrors Cosmos3-Edge.yaml (action_gen: true), same as the nano baseline. + # Safe since the renewed (2026-07-16) checkpoint ships trained action-head + # weights (pre-renewal ones lacked them and NaN'd — see vision_sft_edge.py); + # recipes that don't train action data still override to False. + action_gen=True, + causal_training_strategy="none", + input_caption_key="ai_caption", + input_image_key="images", + input_video_key="video", + joint_attn_implementation="two_way", + latent_downsample_factor=16, + log_enc_time_every_n=100, + max_action_dim=64, + max_num_tokens_after_packing=45056, + num_embodiment_domains=32, + resolution="480", + sound_gen=False, + sound_latent_fps=25, + state_ch=48, + state_t=300, + video_temporal_causal=False, + vision_gen=True, + diffusion_expert_config=dict( + base_fps=24, + enable_fps_modulation=True, + load_weights_from_pretrained=True, + max_vae_latent_side_after_patchify=20, + patch_spatial=2, + timestep_range=1.0, + unified_3d_mrope_reset_spatial_ids=True, + unified_3d_mrope_temporal_modality_margin=15000, + ), + ema=dict( + enabled=True, + iteration_shift=0, + rate=0.1, + ), + lbl=dict( + coeff_gen=None, + coeff_und=None, + method="local", + ), + parallelism=dict( + cfg_parallel_shard_degree=1, + context_parallel_shard_degree=1, + data_parallel_shard_degree=8, + enable_inference_mode=False, + fsdp_master_dtype="float32", + ), + compile=dict( + compile_dynamic=True, + compiled_region="language", + coordinate_descent_tuning=False, + max_autotune_pointwise=False, + use_cuda_graphs=False, + enabled=True, + ), + precision="bfloat16", + activation_checkpointing=dict( + mode="full", + ), + rectified_flow_inference_config=dict( + num_train_timesteps=1000, + scheduler_type="unipc", + shift=1, + use_dynamic_shifting=False, + ), + rectified_flow_training_config=dict( + action_loss_weight=10.0, + image_loss_scale=None, # Edge delta (nano=1.0); Cosmos3-Edge.yaml image_loss_scale: null + independent_action_schedule=False, + # Edge-only keys accepted by RectifiedFlowTrainingConfig (nano omits them): + independent_sound_schedule=False, + loss_scale=10.0, # Edge delta (nano=1.0); Cosmos3-Edge.yaml loss_scale: 10.0 + normalize_loss_by_active=False, + shift={"256": 3, "480": 5, "720": 10}, + shift_action=None, + shift_sound=None, + sound_loss_scale=2.0, # Edge-only (nano leaves default None); unused while sound_gen=False + train_time_action_distribution="logitnormal", + train_time_image_distribution="logitnormal", + train_time_sound_distribution="logitnormal", + train_time_video_distribution="waver", + train_time_weight="uniform", + use_discrete_rf=False, + use_dynamic_shift=False, + # Dropped Cosmos3-Edge.yaml keys — NOT fields of RectifiedFlowTrainingConfig: + # high_sigma_ratio, high_sigma_timesteps_max, high_sigma_timesteps_min, + # use_high_sigma_strategy, use_high_sigma_strategy_action, + # use_high_sigma_strategy_sound (high-sigma strategy unsupported in the + # training-package dataclass; nano omits them and trains fine). + ), + tokenizer=dict( + bucket_name="", + chunk_duration=93, + encode_chunk_frames={"256": 68, "480": 24, "720": 12}, + encode_exact_durations=None, + keep_decoder_cache=False, + object_store_credential_path_pretrained="", + spatial_compression_factor=16, + temporal_compression_factor=4, + use_streaming_encode=False, + vae_path="pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", + ), + vlm_config=dict( + # Top-level VLMConfig fields (Cosmos3-Edge.yaml vlm_config.*): + layer_module=None, # Edge.yaml vlm_config.layer_module: null (nano="Qwen2MoTDecoderLayer") + model_name="nvidia/Cosmos3-Edge-Reasoner", + tie_word_embeddings=False, + use_system_prompt=False, + # Edge.yaml sets vlm_config.qk_norm=False; VLMConfig.qk_norm already defaults + # to False (nano omits it too), so it is intentionally left unset. + pretrained_weights=dict( + enabled=False, # SFT loads from the DCP, not the HF backbone + backbone_path=( + "s3://bucket/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/" + ), # kept for parity, unused while enabled=False + credentials_path="", + enable_gcs_patch_in_boto3=True, + ), + model_instance=L(Nemotron3DenseVLTextForCausalLM)( + config=L(create_vlm_config)( + base_config=L(Nemotron3DenseVLMoTConfig.from_json_file)( + json_file=( + "cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/configs/" + "Nemotron-2B-Dense-VL.json" + ), + ), + freeze_und=False, + layer_module="MoTDecoderLayer", + qk_norm_for_text=False, # Edge delta (nano=True); Edge.yaml create_vlm_config.qk_norm_for_text: false + use_und_k_norm_for_gen=True, # und-K norm on gen→und cross-attention; Edge.yaml parity + include_visual=None, # Edge delta (nano omits -> default); Edge.yaml create_vlm_config.include_visual: null + tie_word_embeddings=True, + ), + ), + tokenizer=L(build_processor_lazy)( + repository="nvidia/Cosmos3-Edge", + revision="main", + ), + ), +) diff --git a/cosmos_framework/configs/base/experiment/sft/vision_sft_edge.py b/cosmos_framework/configs/base/experiment/sft/vision_sft_edge.py new file mode 100644 index 00000000..c7b856f6 --- /dev/null +++ b/cosmos_framework/configs/base/experiment/sft/vision_sft_edge.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""``vision_sft_edge`` — Cosmos3-Edge vision SFT recipe. + +Nemotron-2B-Dense-VL / Cosmos3-Edge backbone, T2V/I2V/V2V SFT, PackingDataLoader + +RankPartitionedDataLoader stack. EMA enabled. + +Sibling of ``vision_sft_nano``: same Hydra defaults, same dataloader stack, same +optimizer / scheduler / trainer / checkpoint blocks. The only differences are the +model config (``EDGE_MODEL_CONFIG`` — dense Nemotron backbone, no audio) and the +recipe name. Edge has no sound tokenizer, matching nano's +``{"override /sound_tokenizer": None}`` default. + +Notes: + * ``_target_`` references for ``model``, ``optimizer``, ``scheduler``, + ``checkpoint``, ``callbacks``, ``ema``, ``tokenizer``, ``cluster``, + ``vlm_config``, and ``ckpt_type`` flow from the ``defaults:`` group + choices, matching the YAML's ``defaults:`` list (``_self_`` is placed + LAST per prerelease convention so the experiment overrides the + defaults, but no setting changes semantically). + * The YAML's giant ``model_parallel`` block, ``trainer.profiling``, + ``trainer.straggler_detection`` and ``trainer.type`` are populated by + the base ``Config`` (``cosmos_framework/configs/base/config.py``) defaults and + are therefore omitted here. + +``checkpoint.load_path`` is left as ``???`` (a Hydra MISSING marker); supply +via CLI / a downstream experiment that inherits from this one. + +Usage:: + + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. torchrun --nproc_per_node=1 \\ + --master_port=12341 -m cosmos_framework.scripts.train \\ + --config=cosmos_framework/configs/base/config.py -- \\ + experiment=vision_sft_edge \\ + checkpoint.load_path= +""" + +import copy + +from hydra.core.config_store import ConfigStore + +from cosmos_framework.configs.base.experiment.sft.models.edge_model_config import EDGE_MODEL_CONFIG +from cosmos_framework.data.generator.joint_dataloader import ( + PackingDataLoader, + RankPartitionedDataLoader, +) +from cosmos_framework.data.generator.local_datasets.sft_dataset import get_sft_dataset +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict + +cs = ConfigStore.instance() + +# Vision SFT trains no action tokens, so drop the action head (recipe semantics). +# Historical: pre-renewal (< 2026-07-16) Edge checkpoints shipped no action weights, +# so action_gen=True NaN'd (uninitialized frozen head made the graph-consistency +# dummy forward's 0.0 * preds_action NaN). +_EDGE_VISION_MODEL_CONFIG = copy.deepcopy(EDGE_MODEL_CONFIG) +_EDGE_VISION_MODEL_CONFIG["action_gen"] = False + + +vision_sft_edge = LazyDict( + dict( + defaults=[ + {"override /model": "mot_fsdp"}, + {"override /data_train": None}, + {"override /data_val": None}, + {"override /optimizer": "adamw"}, + # YAML used `scheduler: warmup_cosine_lr` but that group is only + # registered in cosmos_framework/configs/base/reasoner/defaults/optimizer.py + # (reachable from the vlm config tree). The base vfm config path + # only knows `lambdacosine`, which also sets + # lr_scheduler_type="LambdaCosine" — behaviorally identical. + {"override /scheduler": "lambdacosine"}, + {"override /checkpoint": "s3"}, + { + "override /callbacks": [ + "basic", + "optimization", + "job_monitor", + "generation", + ] + }, + {"override /ema": "power"}, + {"override /tokenizer": "wan2pt2_tokenizer"}, + {"override /sound_tokenizer": None}, + {"override /vlm_config": None}, + {"override /ckpt_type": "dcp"}, + "_self_", + ], + job=dict( + project="cosmos3", + group="sft", + name="vision_sft_edge", + wandb_mode="disabled", + ), + model=dict( + config=copy.deepcopy(_EDGE_VISION_MODEL_CONFIG), + ), + optimizer=dict( + betas=[0.9, 0.95], + eps=1.0e-06, + fused=True, + keys_to_select=[ + "moe_gen", + "time_embedder", + "vae2llm", + "llm2vae", + "k_norm_und_for_gen", # und-K norm trains with the gen pathway + ], + lr=5.0e-04, + lr_multipliers={}, + optimizer_type="AdamW", + weight_decay=0, + ), + scheduler=dict( + lr_scheduler_type="LambdaCosine", + cycle_lengths=[1000], + f_max=[1.0], + f_min=[0.0], + f_start=[0.0], + verbosity_interval=0, + warm_up_steps=[50], + ), + trainer=dict( + distributed_parallelism="fsdp", + grad_accum_iter=2, + logging_iter=1, + max_iter=500, + max_val_iter=None, + # YAML had `memory_format: preserve_format` as a string, but the + # prerelease trainer passes this verbatim to model.to(memory_format=…) + # which requires a torch.memory_format enum (not a string). + # Omit and let the framework default apply, matching how + # mixed_modality_sft_nano.py handles it. + run_validation=False, + run_validation_on_start=False, + save_zero_checkpoint=False, + seed=42, + timeout_period=999999999, + validation_iter=100, + compile_config=dict(recompile_limit=8, use_duck_shape=False), + cudnn=dict(benchmark=True, deterministic=False), + ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True), + grad_scaler_args=dict(enabled=False), + callbacks=dict( + compile_tokenizer=dict( + compile_after_iterations=3, + enabled=False, + warmup_resolutions=None, + ), + dataloader_speed=dict(every_n=100, save_s3=False, step_size=1), + device_monitor=dict( + every_n=200, + log_memory_detail=True, + save_s3=False, + step_size=1, + upload_every_n_mul=5, + ), + expert_heatmap=dict(every_n=1000), + grad_clip=dict(clip_norm=0.1, force_finite=True), + heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20), + iter_speed=dict(every_n=1, hit_thres=50, save_s3=False, save_s3_every_log_n=500), + low_precision=dict(update_iter=1), + manual_gc=dict(every_n=5, gc_level=1, warm_up=1), + norm_monitor=dict( + every_n=100, + layer_norm_only=False, + log_stat_wandb=True, + model_key=None, + save_s3=False, + step_size=1, + track_activations=True, + ), + param_count=dict(save_s3=False), + sequence_packing_padding=dict(every_n=50), + sigma_loss_analysis=dict(every_n=500, every_n_viz=500, save_s3=False), + skip_nan_step=dict(max_consecutive_nan=100), + training_stats=dict(log_freq=100), + wandb_2x=dict( + logging_iter_multipler=2, + save_logging_iter_multipler=1, + save_s3=False, + ), + wandb_val=dict(save_s3=False), + ), + ), + checkpoint=dict( + broadcast_via_filesystem=False, + dcp_async_mode_enabled=False, + enable_gcs_patch_in_boto3=True, + keys_not_to_resume=[], + keys_to_skip_loading=["net_ema."], + load_ema_to_reg=False, + load_path="???", # supply via CLI / downstream experiment + load_training_state=False, + only_load_scheduler_state=False, + save_iter=100, + strict_resume=True, + verbose=True, + hf_export=dict( + enabled=False, + export_every_n=1, + hf_repo_id=None, + upload_to_object_store=dict( + bucket="", + credentials="", + enabled=False, + ), + ), + jit=dict( + device="cuda", + dtype="bfloat16", + enabled=False, + input_shape=None, + strict=True, + ), + load_from_object_store=dict( + bucket="", + credentials="", + enabled=False, + ), + save_to_object_store=dict( + bucket="", + credentials="", + enabled=False, + ), + ), + dataloader_train=L(PackingDataLoader)( + audio_sample_rate=48000, + dataset_name="default", + max_samples_per_batch=None, + max_sequence_length=45056, + patch_spatial=2, + sound_latent_fps=0, + tokenizer_spatial_compression_factor=16, + tokenizer_temporal_compression_factor=4, + dataloader=L(RankPartitionedDataLoader)( + batch_size=1, + in_order=True, + num_workers=4, + persistent_workers=True, + pin_memory=True, + prefetch_factor=4, + sampler=None, + datasets=dict( + video=dict( + ratio=1, + dataset=L(get_sft_dataset)( + append_duration_fps_timestamps=True, + append_resolution_info=True, + # Per-caption token cap. Structured-JSON captions are long, so + # default to 2048 (measured max ~1790); tune via the TOML knob + # [dataloader_train].max_caption_tokens. See sft_dataset.py + # _MAX_CAPTION_TOKENS. + max_caption_tokens=2048, + caption_suffix="", + cfg_dropout_keep_metadata=False, + cfg_dropout_rate=0.1, + # 70% T2V, 20% I2V (first frame), 10% V2V (first 5 frames / 2 latent frames) + conditioning_config={0: 0.7, 1: 0.2, 2: 0.1}, + conditioning_fps=-1, + conditioning_fps_noise_std=0.0, + frame_selection_mode="first", + jsonl_paths=["${oc.env:DATASET_PATH}/train/video_dataset_file.jsonl"], + min_short_edge=0, + num_video_frames=-1, + resolution="256", + sample_by_window=False, + temporal_compression_factor=4, + temporal_interval_mode="max_30fps", + use_system_prompt=False, + # YAML spells this out as + # _target_: create_qwen2_tokenizer_with_download + # config_variant: gcp + # but that pins the dataset's tokenizer to the GCP + # variant, requiring credentials/gcp_checkpoint.secret. + # Use a Hydra interpolation instead so launchers + # (e.g. launch_vision_sft_edge_toml.sh) can flip + # model.config.vlm_config.tokenizer.config_variant=hf + # and have the dataset inherit the same setting. + tokenizer_config="${model.config.vlm_config.tokenizer}", + ), + ), + ), + ), + ), + dataloader_val=None, + upload_reproducible_setup=False, + ), + flags={"allow_objects": True}, +) + + +for _item in [vision_sft_edge]: + _name = [k for k, v in globals().items() if v is _item][0] + cs.store(group="experiment", package="_global_", name=_name, node=_item) diff --git a/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py b/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py index 5c119ab6..a116d65a 100644 --- a/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py +++ b/cosmos_framework/configs/base/reasoner/defaults/vlm_policy.py @@ -54,6 +54,13 @@ nemotron_nano_12b_v2_vl_bf16 = PolicyConfig(backbone=VLMConfig(model_name="nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16")) +# Cosmos3-Edge reasoner (Nemotron-2B-Dense-VL LM + SigLIP2 vision tower). +# model_name is the public (ungated) omni release nvidia/Cosmos3-Edge, +# model_type "cosmos3_edge" (native HF metadata, no remote code; classes are +# registered in-framework). Reasoner weights load directly from the snapshot — +# the training loader follows its root safetensors index; no converter step. +cosmos3_edge_reasoner = PolicyConfig(backbone=VLMConfig(model_name="nvidia/Cosmos3-Edge")) + def register_vlm_policy(): cs = ConfigStore.instance() @@ -153,3 +160,9 @@ def register_vlm_policy(): name="nemotron_nano_12b_v2_vl_bf16", node=nemotron_nano_12b_v2_vl_bf16, ) + cs.store( + group="vlm_policy", + package="model.config.policy", + name="cosmos3_edge_reasoner", + node=cosmos3_edge_reasoner, + ) diff --git a/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_edge.py b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_edge.py new file mode 100644 index 00000000..4f6e5d31 --- /dev/null +++ b/cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_edge.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""VideoPhy-2 SFT recipe on the Cosmos3-Edge reasoner. + +Mirrors ``videophy2_sft_nano`` but targets the Cosmos3-Edge reasoner +(Nemotron-2B-Dense-VL LM + SigLIP2 vision tower, ``model_type="cosmos3_edge"``) +instead of Qwen3-VL-8B. ``model_name`` is the public omni release +``nvidia/Cosmos3-Edge`` and supplies arch/config/tokenizer AND the reasoner +weights: the training loader follows the snapshot's root safetensors index, so +weights load directly from the download — no converter step. +``VLM_SAFETENSORS_PATH`` is an OPTIONAL launcher override for loading from a +local safetensors directory instead (same mechanism as nano/super). + +Deltas vs ``videophy2_sft_nano`` (everything else is identical): + * ``override /vlm_policy``: ``qwen3_vl_8b_instruct`` -> ``cosmos3_edge_reasoner``. + * ``optimizer.lr_multipliers``: nano sets ``{"model.visual": 1.0}`` to lift its + projector (which sits UNDER ``model.visual`` in Qwen) off the inherited default + ``0.1x``; edge OMITS the override. Edge's projector is a top-level + ``model.projector`` (the Edge reasoner's PatchMerger), so the default + ``{"model.visual": 0.1}`` only matches the frozen SigLIP2 tower and the + projector already trains at the uniform ``1.0x`` -- same net effect as fixed + nano (projector + LM + lm_head all at ``1e-6``), no key needed. + * ``model.config.freeze``: ``freeze_vision_encoder=True`` only supports the + known Qwen/Intern towers (vlm_model.py ``_get_vision_encoder_modules``), not + the Edge (``cosmos3_edge``) tower, so freeze the SigLIP2 tower via the regex + ``frozen_params=[r"model\\.visual\\."]`` instead — same intent (vision + frozen, projector + LM trainable). + +The dataflow helpers below are duplicated from ``videophy2_sft_nano`` on purpose +(NOT imported): the reasoner config loader reloads experiment modules with +``reload=True`` in alphabetical order, so importing nano's ``build_..`` function +here would capture a pre-reload object and fail the dataloader-worker pickle +identity check. Keeping the recipe self-contained makes every LazyCall reference +this module's own (co-reloaded) callables. + +Launch via examples/launch_sft_videophy2_edge.sh after running +prepare_videophy2_from_hf to populate $VIDEOPHYSICS_ROOT. +""" + +from __future__ import annotations + +from hydra.core.config_store import ConfigStore + +from cosmos_framework.utils.lazy_config import LazyCall as L +from cosmos_framework.utils.lazy_config import LazyDict +from cosmos_framework.data.generator.dataflow import CosmosDataLoader, IterableDistributor, PoolPackingBatcher +from cosmos_framework.data.generator.processors import build_processor +from cosmos_framework.data.reasoner.local_sft_dataset import LocalSFTDataset +from cosmos_framework.data.reasoner.data_sources_videophy2.videophy2 import DATAINFO +from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX +from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMCollator +from cosmos_framework.configs.base.reasoner.experiment.videophy2_dataflow_roles import VideoPhy2Processor + +cs = ConfigStore.instance() + + +class _UnshardedLocalSFTDataset(LocalSFTDataset): + """Yield the full shuffled manifest per iteration. + + Why: ``CosmosDataLoader``'s IterableDistributor already shards by + ``dp_rank * num_workers + worker_id``; stock ``LocalSFTDataset`` shards + again inside ``__iter__``, double-sharding to ``1 / (world*workers)^2``. + """ + + def _per_partition_indices(self, epoch: int) -> list[int]: + import random + + manifest = self._load_manifest() + indices = list(range(len(manifest))) + if self.shuffle: + rng = random.Random(self.distributor_seed + epoch) + rng.shuffle(indices) + return indices + + +def build_videophy2_local_dataset( + dataset_key: str, + split: str, +) -> _UnshardedLocalSFTDataset: + # augmentor_config=None: the Processor decodes+tokenizes inline; the + # BytesToMedia/TokenizeData augmentors aren't shipped in OSS. + source = DATAINFO[dataset_key] + if split not in source.manifest_path: + raise KeyError( + f"split={split!r} not present in DATAINFO[{dataset_key!r}].manifest_path " + f"(have: {list(source.manifest_path)})" + ) + return _UnshardedLocalSFTDataset( + manifest_path=source.manifest_path[split], + data_root=source.data_root, + media_field_name=source.media_field_name, + augmentor_config=None, + text_only=source.text_only, + shuffle=True, + distributor_seed=1993, + is_infinite_loader=True, + split=split, + dataset_name=dataset_key, + ) + + +def _dl(dataset_key, split, num_workers, persistent_workers=False, pin_memory=False, prefetch_factor=None): + return L(CosmosDataLoader)( + distributor=L(IterableDistributor)( + iterable=L(build_videophy2_local_dataset)(dataset_key=dataset_key, split=split), + ), + processor=L(VideoPhy2Processor)( + processor=L(build_processor)( + tokenizer_type="${model.config.policy.backbone.model_name}", + config_variant="hf", + ), + ignore_index=IGNORE_INDEX, + ), + batcher=L(PoolPackingBatcher)( + max_tokens="${data_setting.max_tokens}", + pool_size=16, + max_batch_size=1, + long_threshold=6400, + ), + collator=L(VLMCollator)(), + num_workers=num_workers, + persistent_workers=persistent_workers, + pin_memory=pin_memory, + prefetch_factor=prefetch_factor, + ) + + +videophy2_sft_edge = LazyDict( + dict( + defaults=[ + {"override /checkpoint": "local"}, + {"override /data_train": None}, + {"override /data_val": None}, + {"override /model": "vlm_fsdp"}, + {"override /vlm_policy": "cosmos3_edge_reasoner"}, # Edge delta (nano=qwen3_vl_8b_instruct) + {"override /callbacks": ["basic_vlm", "basic_log", "hf_export"]}, + "_self_", + ], + job=dict( + project="cosmos3_reasoner", + group="sft", + wandb_mode="disabled", + ), + trainer=dict( + callbacks=dict( + log_tensor_shape=dict(num_log=2), + ), + max_iter=50, + logging_iter=1, + run_validation=True, + validation_iter=10, + max_val_iter=50, + grad_accum_iter=8, + ), + optimizer=dict( + lr=1e-6, + fused=True, + weight_decay=0.05, + betas=[0.9, 0.999], + # Uniform 1.0x LR across projector + LM + lm_head, matching the fixed + # videophy2_sft_nano. Nano must set {"model.visual": 1.0} because its + # projector lives UNDER model.visual (Qwen model.visual.merger) and would + # otherwise inherit the reasoner default {"model.visual": 0.1}. Edge needs + # NO override: its projector is a top-level model.projector + # (the Edge reasoner's PatchMerger), not under model.visual, so the inherited + # default {"model.visual": 0.1} only matches the frozen (optimizer- + # excluded) SigLIP2 tower and never reaches the projector -- which thus + # already trains at the default 1.0x = 1e-6. So lr_multipliers is omitted. + ), + scheduler=dict( + warm_up_steps=[5], + cycle_lengths=[50], + f_min=[0.1], + ), + data_setting=dict( + # The "qwen_" names are historical knobs, not a Qwen-backbone requirement; + # kept for parity with nano. + qwen_max_video_token_length=8192, + qwen_max_image_token_length=2048, + max_tokens=16000, + max_batch_size=1, + distributor_seed=1993, + ), + model=dict( + config=dict( + # Edge delta: freeze the SigLIP2 tower via regex — + # freeze_vision_encoder=True only supports Qwen/Intern towers. + # Projector + LM stay trainable, matching nano's intent. + freeze=dict( + frozen_params=[r"model\.visual\."], + ), + # FSDP full shard, auto from WORLD_SIZE — supports the documented + # 4-GPU (NPROC_PER_NODE=4) and 8-GPU launches without a clamp warning. + parallelism=dict( + data_parallel_shard_degree=-1, + data_parallel_replicate_degree=1, + ), + policy=dict( + monkey_patch_for_text_only_data=True, + ), + ), + ), + # hf_export so eval_videophy2 can read each save as HF safetensors. + checkpoint=dict( + save_iter=100, + hf_export=dict(enabled=True), + ), + upload_reproducible_setup=False, + dataloader_train=_dl("videophy2_train", "train", 2, persistent_workers=True, pin_memory=True, prefetch_factor=2), + dataloader_val=_dl("videophy2_val", "val", 0, persistent_workers=False, pin_memory=True, prefetch_factor=None), + ), + flags={"allow_objects": True}, +) + + +for _item in [videophy2_sft_edge]: + experiment_name = [name.lower() for name, value in globals().items() if value is _item][0] + if "job" not in _item: + _item["job"] = dict(name=experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}") + else: + _item["job"]["name"] = experiment_name + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" + + cs.store(group="experiment", package="_global_", name=experiment_name, node=_item) diff --git a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py index 64e2e12b..8edf3313 100644 --- a/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py +++ b/cosmos_framework/data/generator/action/datasets/droid_lerobot_dataset.py @@ -62,6 +62,8 @@ class DROIDLeRobotDataset(BaseActionLeRobotDataset): """ """ + EMBODIMENT_TYPE: str = "droid_lerobot" + def __init__( self, root: str = "/lustre/fsw/portfolios/cosmos/projects/cosmos_base_training/cosmos3_action_datasets/droid_plus_lerobot_640x360_20260412", diff --git a/cosmos_framework/data/generator/augmentor_provider.py b/cosmos_framework/data/generator/augmentor_provider.py index fa2a1c01..214d771d 100644 --- a/cosmos_framework/data/generator/augmentor_provider.py +++ b/cosmos_framework/data/generator/augmentor_provider.py @@ -796,6 +796,7 @@ def get_video_augmentor_v3_json_caption( use_dynamic_fps: bool = False, max_stride: int = 3, min_stride: int = 1, + stride_config: dict | None = None, min_fps: float = 10.0, max_fps: float = 60.0, use_system_prompt: bool = False, @@ -879,6 +880,7 @@ def get_video_augmentor_v3_json_caption( "use_dynamic_fps": use_dynamic_fps, "max_stride": max_stride, "min_stride": min_stride, + "stride_config": stride_config, "min_fps": min_fps, "max_fps": max_fps, "seek_mode": "exact", diff --git a/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py b/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py index 694b5964..e99f421d 100644 --- a/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py +++ b/cosmos_framework/data/generator/augmentors/text_transforms_for_image.py @@ -8,7 +8,9 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.v3_text_transforms import pad_and_resize from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log - +from cosmos_framework.data.generator.augmentors.caption_embedding_keys import ( + _CAPTION_EMBEDDING_KEY_MAPPING_IMAGES, +) # For the qwen captions, we have 3 variants: short, medium, long # In addition, for synthetic data, we create prompt embeddings as well. diff --git a/cosmos_framework/data/generator/augmentors/video_parsing.py b/cosmos_framework/data/generator/augmentors/video_parsing.py index fb5808a6..4079c033 100644 --- a/cosmos_framework/data/generator/augmentors/video_parsing.py +++ b/cosmos_framework/data/generator/augmentors/video_parsing.py @@ -416,6 +416,24 @@ def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: O assert self.max_stride >= self.min_stride, ( f"max_stride ({self.max_stride}) must be >= min_stride ({self.min_stride})" ) + # Optional explicit categorical stride distribution: {stride_int: probability_float}. + # When set (e.g. {1: 0.5, 2: 0.5}), _sample_stride_with_bias samples the per-video + # stride from this map instead of the exp-decay scheme (used for FPS-mixing ablations). + # Default None => unchanged exp-decay behavior (no effect on any existing config). + _stride_config_raw = args.get("stride_config", None) + if _stride_config_raw is not None: + # Coerce to {int: float}. OmegaConf/YAML round-trip serializes int keys as strings + # (e.g. {"2": 1.0}) and may hand us a DictConfig; normalize both so downstream + # sampling always sees int strides + float probabilities. + self.stride_config = {int(k): float(v) for k, v in dict(_stride_config_raw).items()} + assert len(self.stride_config) > 0, "stride_config must be a non-empty mapping" + assert all(k >= 1 for k in self.stride_config), ( + f"stride_config keys must be strides >= 1, got {list(self.stride_config)!r}" + ) + _prob_sum = float(sum(self.stride_config.values())) + assert abs(_prob_sum - 1.0) < 1e-6, f"stride_config probabilities must sum to 1.0, got {_prob_sum}" + else: + self.stride_config = None self.min_fps = args.get("min_fps", _MIN_FPS) self.max_fps = args.get("max_fps", _MAX_FPS) if self.use_dynamic_fps: @@ -494,6 +512,12 @@ def _sample_stride_with_bias(self, max_stride: int, min_stride: int = 1) -> int: These values are chosen to approximately match our old ablations. TODO @pchattopadhy: Do ablations with this scheme """ + # FPS-mixing ablation: when an explicit categorical stride distribution is configured, + # sample the per-video stride from it directly (ignores the exp-decay scheme below). + if self.stride_config is not None: + _strides = list(self.stride_config.keys()) + _probs = list(self.stride_config.values()) + return int(np.random.choice(_strides, p=_probs)) assert max_stride >= min_stride, f"max_stride ({max_stride}) must be >= min_stride ({min_stride})" if max_stride == min_stride: return min_stride diff --git a/cosmos_framework/data/generator/augmentors/video_parsing_test.py b/cosmos_framework/data/generator/augmentors/video_parsing_test.py new file mode 100644 index 00000000..6da73081 --- /dev/null +++ b/cosmos_framework/data/generator/augmentors/video_parsing_test.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Unit tests for the FPS-mixing `stride_config` primitive in VideoParsingWithFullFrames. + +Validates the categorical per-video stride sampling used by the FPS-mixing ablation: +coercion of OmegaConf string keys, the realized per-video stride distribution, input +validation, and that omitting stride_config leaves the exp-decay path unchanged. +""" + +from collections import Counter + +import numpy as np +import pytest +from omegaconf import OmegaConf + +from cosmos_framework.data.generator.augmentors.video_parsing import VideoParsingChunkedFrames + +# Minimal args to instantiate the augmentor for stride-sampling only (no video decode). +_BASE_ARGS = dict( + max_stride=2, + min_stride=2, + use_dynamic_fps=True, + min_fps=10.0, + max_fps=60.0, + causal_vae=True, + max_num_frames=1000, +) + + +def _make(stride_config=None): + args = dict(_BASE_ARGS) + if stride_config is not None: + args["stride_config"] = stride_config + return VideoParsingChunkedFrames(input_keys=["metas", "video"], output_keys=None, args=args) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_stride_config_coercion_omegaconf_str_keys(): + # OmegaConf serializes int keys as strings; __init__ must coerce back to {int: float}. + aug = _make(OmegaConf.create({"1": 0.5, "2": 0.5})) + assert aug.stride_config == {1: 0.5, 2: 0.5} + assert all(isinstance(k, int) for k in aug.stride_config) + assert all(isinstance(v, float) for v in aug.stride_config.values()) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_stride_config_none_leaves_exp_decay_unchanged(): + aug = _make(None) + assert aug.stride_config is None + # exp-decay path still yields a valid stride in [min_stride, max_stride]. + s = aug._sample_stride_with_bias(aug.max_stride, aug.min_stride) + assert aug.min_stride <= s <= aug.max_stride + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "config,expected", + [ + ({2: 1.0}, {2: 1.0}), # baseline: always half-FPS + ({1: 0.5, 2: 0.5}, {1: 0.5, 2: 0.5}), # mix50: per-video 50/50 native+half + ({1: 0.3, 2: 0.7}, {1: 0.3, 2: 0.7}), # generic ratios + ], +) +def test_stride_config_realized_distribution(config, expected): + aug = _make(OmegaConf.create({str(k): v for k, v in config.items()})) + np.random.seed(0) # deterministic => flake-free + n = 8000 + c = Counter(aug._sample_stride_with_bias(aug.max_stride, aug.min_stride) for _ in range(n)) + realized = {k: c[k] / n for k in config} + for k, p in expected.items(): + assert abs(realized[k] - p) < 0.03, f"stride {k}: realized {realized[k]:.3f} vs expected {p}" + # only strides from the config are ever sampled + assert set(c) <= set(config) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + "bad", + [ + {}, # empty mapping + {0: 1.0}, # stride < 1 + {1: 0.5, 2: 0.4}, # probabilities do not sum to 1 + ], +) +def test_stride_config_validation_rejects_bad_input(bad): + with pytest.raises(AssertionError): + _make(OmegaConf.create({str(k): v for k, v in bad.items()})) diff --git a/cosmos_framework/data/generator/processors/__init__.py b/cosmos_framework/data/generator/processors/__init__.py index 46a3a473..2b503d01 100644 --- a/cosmos_framework/data/generator/processors/__init__.py +++ b/cosmos_framework/data/generator/processors/__init__.py @@ -10,6 +10,7 @@ from transformers import PreTrainedTokenizerFast from cosmos_framework.data.generator.processors.base import BaseVLMProcessor +from cosmos_framework.data.generator.processors.cosmos3_edge_processing import is_cosmos3_edge_native_snapshot from cosmos_framework.data.generator.processors.nemotron3densevl_processor import Nemotron3DenseVLProcessor from cosmos_framework.data.generator.processors.nemotronvl_processor import NemotronVLProcessor from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor @@ -108,8 +109,11 @@ def build_processor( # (e.g. the top level of nvidia/Cosmos3-Nano, which ships its own # preprocessor_config.json, tokenizer.json, etc). Avoids the redundant # upstream Qwen/Qwen3-VL-*-Instruct fetch. Cosmos3-Nano/Super both ship - # a Qwen3VL-compatible processor, so dispatch to Qwen3VLProcessor. + # a Qwen3VL-compatible processor, so dispatch to Qwen3VLProcessor — + # except renewed Cosmos3-Edge snapshots, which need the Nemotron bridge. if os.path.isdir(tokenizer_type): + if is_cosmos3_edge_native_snapshot(tokenizer_type): + return Nemotron3DenseVLProcessor(tokenizer_type, cache_dir=cache_dir) return Qwen3VLProcessor(tokenizer_type, cache_dir=cache_dir) if credentials is None or bucket is None: if config_variant is None: @@ -129,7 +133,9 @@ def build_processor( "NVIDIA-Nemotron-3-Dense-VL" in tokenizer_type or "Qwen3-2B-ViT" in tokenizer_type or "nvidia/Cosmos3-Reasoner-2B-Private" in tokenizer_type - or "nvidia/Cosmos3-Edge-Reasoner" in tokenizer_type + # nvidia/Cosmos3-Edge ships the same Nemotron bridge processor; the + # substring also covers nvidia/Cosmos3-Edge-Reasoner. + or "nvidia/Cosmos3-Edge" in tokenizer_type ): return Nemotron3DenseVLProcessor(tokenizer_type, credentials=credentials, bucket=bucket, cache_dir=cache_dir) elif "Qwen/Qwen3-0.6B" in tokenizer_type: diff --git a/cosmos_framework/data/generator/processors/base.py b/cosmos_framework/data/generator/processors/base.py index f944139e..9a6068ef 100644 --- a/cosmos_framework/data/generator/processors/base.py +++ b/cosmos_framework/data/generator/processors/base.py @@ -21,8 +21,12 @@ from transformers.models.auto.processing_auto import AutoProcessor -from cosmos_framework.utils import log +from cosmos_framework.data.generator.processors.cosmos3_edge_processing import ( + build_cosmos3_edge_processor, + is_cosmos3_edge_native_snapshot, +) from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.utils import log from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 @@ -118,7 +122,12 @@ def __init__( name, credentials, bucket, include_model_weights=False, cache_dir=cache_dir ) - self.processor = AutoProcessor.from_pretrained(model_name_or_path_local, trust_remote_code=True) + if is_cosmos3_edge_native_snapshot(model_name_or_path_local): + # AutoProcessor on transformers 4.x silently degrades to a bare tokenizer + # for renewed (no remote code) Cosmos3-Edge snapshots; use the native port. + self.processor = build_cosmos3_edge_processor(model_name_or_path_local) + else: + self.processor = AutoProcessor.from_pretrained(model_name_or_path_local, trust_remote_code=True) log.info("Successfully loaded processor from local cache") self.image_token_id = ( diff --git a/cosmos_framework/data/generator/processors/cosmos3_edge_processing.py b/cosmos_framework/data/generator/processors/cosmos3_edge_processing.py new file mode 100644 index 00000000..b4d5595e --- /dev/null +++ b/cosmos_framework/data/generator/processors/cosmos3_edge_processing.py @@ -0,0 +1,585 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Framework-native processor for renewed ``nvidia/Cosmos3-Edge`` snapshots. + +Port of the deleted remote-code ``processing.py`` from ``nvidia/Cosmos3-Edge`` +rev ``28a0b8e``. Renewed snapshots ship native ``cosmos3_edge`` metadata with no +remote code, and ``AutoProcessor`` on transformers 4.x silently degrades to a +bare tokenizer for them; this module rebuilds the full processor from the +snapshot files without ``trust_remote_code``. + +WARNING — pixel patch layout: patches must stay in the original convention +(raster patch-row order, ``(py, px, c)`` within-row flatten). transformers-main's +native ``Cosmos3EdgeProcessor`` emits a different layout (2x2 block-major, +channel-major within-row) that corrupts vision features with these checkpoint +weights. Do NOT "align" the patchify code with transformers main. Details: +``outputs/audit/cosmos3_edge_native_vision_layout_bug.md``. +""" + +# NOTE: no `from __future__ import annotations` here — ProcessorMixin._merge_kwargs +# introspects Qwen3VLProcessorKwargs.__annotations__ at runtime and PEP 563 +# stringified annotations (ForwardRef) break it. +import json +import math +import os +from typing import Optional, Union + +import numpy as np +import torch +from torchvision.transforms.v2 import functional as F +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_processing_utils_fast import SizeDict +from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling +from transformers.models.auto.tokenization_auto import AutoTokenizer +from transformers.models.qwen3_vl.video_processing_qwen3_vl import ( + Qwen3VLVideoProcessor, + Qwen3VLVideoProcessorInitKwargs, + get_image_size, + smart_resize, +) +from transformers.models.siglip2.image_processing_siglip2_fast import ( + Siglip2ImageProcessorFast, + convert_image_to_patches, +) +from transformers.processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from transformers.utils import TensorType, logging +from transformers.video_utils import VideoInput, group_videos_by_shape, reorder_videos + +logger = logging.get_logger(__name__) + +# Sub-processor config keys that name (remote-code or transformers-main) classes. +# Only the parameter VALUES are consumed; the framework picks the classes itself. +_CLASS_NAME_KEYS = ("processor_class", "image_processor_type", "video_processor_type", "auto_map") + + +def is_cosmos3_edge_native_snapshot(model_dir: str) -> bool: + """Return True for a renewed (native-metadata, no remote code) Cosmos3-Edge snapshot. + + Detection rule: ``config.json`` says ``model_type == "cosmos3_edge"``, or a + ``processor_config.json`` names the native ``Cosmos3EdgeProcessor``. Old + snapshots (``model_type == "nemotron_siglip2"``, remote-code ``processing.py``) + and all other models match neither and keep the ``AutoProcessor`` path. + """ + if not os.path.isdir(model_dir): + return False + config_path = os.path.join(model_dir, "config.json") + if os.path.isfile(config_path): + try: + with open(config_path) as f: + if json.load(f).get("model_type") == "cosmos3_edge": + return True + except (OSError, ValueError): + pass + processor_config_path = os.path.join(model_dir, "processor_config.json") + if os.path.isfile(processor_config_path): + try: + with open(processor_config_path) as f: + if json.load(f).get("processor_class") == "Cosmos3EdgeProcessor": + return True + except (OSError, ValueError): + pass + return False + + +def build_cosmos3_edge_processor(model_dir: str) -> "NemotronNanoV3BridgeProcessor": + """Construct the full Cosmos3-Edge processor from a snapshot dir, without remote code. + + Reproduces the old remote-code ``AutoProcessor`` object exactly (golden-pinned + by ``cosmos3_edge_processing_test.py``): the tokenizer, chat template, and + sub-processor parameter files are unchanged upstream since rev ``28a0b8e``. + """ + tokenizer = AutoTokenizer.from_pretrained(model_dir) + # The renewed tokenizer_config.json adds ``return_mm_token_type_ids: true`` + # (a transformers-main knob); left in init_kwargs it would override the + # text_kwargs default in _merge_kwargs and add an ``mm_token_type_ids`` key + # the old processor never returned. + tokenizer.init_kwargs.pop("return_mm_token_type_ids", None) + + # Wire the processor-level template explicitly; the tokenizer does not + # reliably pick up chat_template.jinja on its own. + with open(os.path.join(model_dir, "chat_template.jinja")) as f: + chat_template = f.read() + + def _load_sub_config(file_name: str) -> dict: + with open(os.path.join(model_dir, file_name)) as f: + config = json.load(f) + for key in _CLASS_NAME_KEYS: + config.pop(key, None) + return config + + image_processor = Siglip2ImageProcessorCustom(**_load_sub_config("preprocessor_config.json")) + video_processor = Qwen3VLVideoProcessorCustom(**_load_sub_config("video_preprocessor_config.json")) + return NemotronNanoV3BridgeProcessor( + image_processor=image_processor, + tokenizer=tokenizer, + video_processor=video_processor, + chat_template=chat_template, + ) + + +# Everything below is vendored from the rev 28a0b8e remote code; kept faithful +# so outputs match the HF reference bit-exactly. + + +def round_by_factor(number: int, factor: int) -> int: + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + return math.floor(number / factor) * factor + + +class NemotronImagesKwargs(Siglip2ImageProcessorFast.valid_kwargs, total=False): + # global setting for all images, can be overridden by per-image kwargs + max_pixels: Optional[int] + min_pixels: Optional[int] + + # per-image overrides, e.g. [{"min_pixels": 256*28*28, "max_pixels": 1280*28*28}, None, ...] + per_image_kwargs: Optional[list[dict]] + + +class Qwen3VLProcessorKwargs(ProcessingKwargs, total=False): + images_kwargs: NemotronImagesKwargs + _defaults = { + "text_kwargs": { + "padding": False, + "return_token_type_ids": False, + "return_mm_token_type_ids": False, + }, + "videos_kwargs": {"return_metadata": True}, + } + + +class Siglip2ImageProcessorCustom(Siglip2ImageProcessorFast): + resample = PILImageResampling.BICUBIC + valid_kwargs = NemotronImagesKwargs + + def __init__(self, **kwargs: Unpack[NemotronImagesKwargs]): + super().__init__(**kwargs) + + def _resize_image( + self, image: torch.Tensor, max_ratio=200, interpolation: Optional[PILImageResampling] = None, **kwargs + ) -> torch.Tensor: + """Resize so both dims are divisible by merge_size * patch_size and the + pixel count lands in [min_pixels, max_pixels], preserving aspect ratio.""" + image_min_pixels = kwargs.get("min_pixels", None) or self.size.get("shortest_edge", None) + image_max_pixels = kwargs.get("max_pixels", None) or self.size.get("longest_edge", None) + assert image_min_pixels is not None and image_max_pixels is not None, ( + "When do_resize is True, min_pixels and max_pixels must be provided." + ) + assert image_max_pixels >= image_min_pixels, ( + "The max_pixels of image must be greater than or equal to min_pixels." + ) + + _, height, width = image.shape + if max(height, width) / min(height, width) > max_ratio: + raise ValueError( + f"absolute aspect ratio must be smaller than {max_ratio}, got {max(height, width) / min(height, width)}" + ) + factor = self.merge_size * self.patch_size + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > image_max_pixels: + beta = math.sqrt((height * width) / image_max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < image_min_pixels: + beta = math.sqrt(image_min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + + image = self.resize(image, size=SizeDict(height=h_bar, width=w_bar), interpolation=interpolation) + return image + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + patch_size: int, + max_num_patches: int, + interpolation: Optional["F.InterpolationMode"], + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: Optional[Union[float, list[float]]], + image_std: Optional[Union[float, list[float]]], + return_tensors: Optional[Union[str, TensorType]], + **kwargs, + ) -> BatchFeature: + per_image_kwargs = kwargs.pop("per_image_kwargs", None) + pixel_values = [] + spatial_shapes = [] + for idx, image in enumerate(images): + image_kwargs = kwargs.copy() + if per_image_kwargs is not None and idx < len(per_image_kwargs) and per_image_kwargs[idx]: + image_kwargs.update(per_image_kwargs[idx]) + if do_resize: + image = self._resize_image(image, max_ratio=200, interpolation=interpolation, **image_kwargs) + + image = self.rescale_and_normalize(image, do_rescale, rescale_factor, do_normalize, image_mean, image_std) + + # (num_channels, height, width) -> (num_patches, patch_size * patch_size * num_channels) + # Raster patch-row order, (py, px, c) within-row flatten — the layout + # the checkpoint's patch_embedding weight was trained for (see module + # docstring); transformers-main's native processor differs here. + patches = convert_image_to_patches(image, patch_size) + + num_patches_height = image.shape[1] // patch_size + num_patches_width = image.shape[2] // patch_size + + spatial_shapes.append((num_patches_height, num_patches_width)) + pixel_values.append(patches) + + spatial_shapes = torch.tensor(spatial_shapes) + + batch_feature = BatchFeature( + data={ + "pixel_values": torch.cat(pixel_values, dim=0), + "spatial_shapes": spatial_shapes, + }, + tensor_type=return_tensors, + ) + return batch_feature + + +class Qwen3VLVideoProcessorCustom(Qwen3VLVideoProcessor): + def __init__(self, **kwargs: Unpack[Qwen3VLVideoProcessorInitKwargs]): + super().__init__(**kwargs) + + def _preprocess( + self, + videos: list[torch.Tensor], + do_convert_rgb: bool = True, + do_resize: bool = True, + size: Optional[SizeDict] = None, + interpolation: PILImageResampling = PILImageResampling.BICUBIC, + do_rescale: bool = True, + rescale_factor: float = 1 / 255.0, + do_normalize: bool = True, + image_mean: Optional[Union[float, list[float]]] = None, + image_std: Optional[Union[float, list[float]]] = None, + patch_size: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ): + merge_size = self.merge_size + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + + for shape, stacked_videos in grouped_videos.items(): + B, T, C, H, W = stacked_videos.shape + num_frames, height, width = T, H, W + if do_resize: + resized_height, resized_width = smart_resize( + num_frames=num_frames, + height=height, + width=width, + temporal_factor=1, + factor=patch_size * merge_size, + min_pixels=size.shortest_edge, + max_pixels=size.longest_edge, + ) + stacked_videos = stacked_videos.view(B * T, C, H, W) + stacked_videos = self.resize( + stacked_videos, + size=SizeDict(height=resized_height, width=resized_width), + interpolation=interpolation, + ) + stacked_videos = stacked_videos.view(B, T, C, resized_height, resized_width) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + # Re-group: sizes may still differ (do_resize False, or per-video resize targets) + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + processed_grids = {} + for shape, stacked_videos in grouped_videos.items(): + resized_height, resized_width = get_image_size(stacked_videos[0], channel_dim=ChannelDimension.FIRST) + stacked_videos = self.rescale_and_normalize( + stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + patches = stacked_videos + + batch_size, grid_t, channel = patches.shape[:3] + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + patches = patches.view( + batch_size, # 0 + grid_t, # 1 + channel, # 2 + grid_h, # 3 + patch_size, # 4 + grid_w, # 5 + patch_size, # 6 + ) + # -> [batch_size, grid_t, grid_h, grid_w, patch_size, patch_size, channel]: + # raster patch rows, channel-LAST within-row — the old-convention layout + # the checkpoint weights expect (see module docstring). + patches = patches.permute(0, 1, 3, 5, 4, 6, 2) + + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + patch_size * patch_size * channel, + ) + + processed_videos_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index) + processed_grids = reorder_videos(processed_grids, grouped_videos_index) + pixel_values_videos = torch.cat(processed_videos, dim=0) + video_grid_thw = torch.tensor(processed_grids) + data = { + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + } + + return BatchFeature(data=data, tensor_type=return_tensors) + + +class NemotronNanoV3BridgeProcessor(ProcessorMixin): + """Cosmos3-Edge processor: SigLIP2 image processor + Qwen3VL video processor + + tokenizer behind a Qwen3VLProcessor-style interface.""" + + attributes = ["image_processor", "tokenizer", "video_processor"] + + image_processor_class = "AutoImageProcessor" + video_processor_class = "AutoVideoProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs): + self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token + self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token + self.image_token_id = ( + tokenizer.image_token_id + if getattr(tokenizer, "image_token_id", None) + else tokenizer.convert_tokens_to_ids(self.image_token) + ) + self.video_token_id = ( + tokenizer.video_token_id + if getattr(tokenizer, "video_token_id", None) + else tokenizer.convert_tokens_to_ids(self.video_token) + ) + super().__init__(image_processor, tokenizer, video_processor, chat_template=chat_template) + self.vision_start_token = ( + "<|vision_start|>" if not hasattr(tokenizer, "vision_start_token") else tokenizer.vision_start_token + ) + self.vision_end_token = ( + "<|vision_end|>" if not hasattr(tokenizer, "vision_end_token") else tokenizer.vision_end_token + ) + self.vision_start_token_id = ( + tokenizer.vision_start_token_id + if getattr(tokenizer, "vision_start_token_id", None) + else tokenizer.convert_tokens_to_ids(self.vision_start_token) + ) + self.vision_end_token_id = ( + tokenizer.vision_end_token_id + if getattr(tokenizer, "vision_end_token_id", None) + else tokenizer.convert_tokens_to_ids(self.vision_end_token) + ) + + def apply_chat_template(self, conversation, **kwargs): + """Extract per-image ``min_pixels``/``max_pixels`` from message content items + (e.g. ``{"type": "image", "image": ..., "min_pixels": 256*28*28}``) and route + them to the image processor; items without them use the global defaults.""" + if isinstance(conversation, (list, tuple)) and conversation and isinstance(conversation[0], (list, tuple)): + conversations = conversation + else: + conversations = [conversation] + + per_image_kwargs: list = [] + for conv in conversations: + for message in conv: + content = message.get("content") + if not isinstance(content, list): + continue + for item in content: + if item.get("type") != "image": + continue + img_override = {} + if "min_pixels" in item: + img_override["min_pixels"] = item["min_pixels"] + if "max_pixels" in item: + img_override["max_pixels"] = item["max_pixels"] + per_image_kwargs.append(img_override if img_override else None) + + if any(v is not None for v in per_image_kwargs): + kwargs["per_image_kwargs"] = per_image_kwargs + + return super().apply_chat_template(conversation, **kwargs) + + def __call__( + self, + images: ImageInput = None, + text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None, + videos: VideoInput = None, + **kwargs: Unpack[Qwen3VLProcessorKwargs], + ) -> BatchFeature: + """Tokenize text and process images/videos into one [`BatchFeature`]: + ``input_ids``/``attention_mask``, plus ``pixel_values`` + ``image_grid_thw`` + and/or ``pixel_values_videos`` + ``video_grid_thw`` when vision inputs are given.""" + output_kwargs = self._merge_kwargs( + Qwen3VLProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + if images is not None: + image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"]) + pixel_values = image_inputs.pop("pixel_values") + spatial_shapes = image_inputs.pop("spatial_shapes") + final_pixel_value = pixel_values.view(-1, pixel_values.shape[-1]) + t_dim = torch.ones((spatial_shapes.shape[0], 1), dtype=spatial_shapes.dtype, device=spatial_shapes.device) + image_grid_thw = torch.cat([t_dim, spatial_shapes], dim=1) + image_inputs = { + "pixel_values": final_pixel_value, + "image_grid_thw": image_grid_thw, + } + else: + image_inputs = {} + image_grid_thw = None + + if videos is not None: + videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"]) + video_grid_thw = videos_inputs["video_grid_thw"] + if not kwargs.get("return_metadata"): + video_metadata = videos_inputs.pop("video_metadata") + else: + video_metadata = videos_inputs["video_metadata"] + else: + videos_inputs = {} + video_grid_thw = None + + if not isinstance(text, list): + text = [text] + + text = text.copy() # below lines change text in-place + if image_grid_thw is not None: + merge_length = self.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + num_image_tokens = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + + if video_grid_thw is not None: + merge_length = self.video_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while self.video_token in text[i]: + metadata = video_metadata[index] + if metadata.fps is None: + logger.warning_once( + "Qwen3VL requires frame timestamps to construct prompts, but the `fps` of the input video " + "could not be inferred. Probably `video_metadata` was missing from inputs and you passed " + "pre-sampled frames. Defaulting to `fps=24`. Please provide `video_metadata` for more " + "accurate results." + ) + metadata.fps = 24 if metadata.fps is None else metadata.fps + + curr_timestamp = self._calculate_timestamps( + metadata.frames_indices, + metadata.fps, + merge_size=1, + ) + + video_placeholder = "" + frame_seqlen = video_grid_thw[index][1:].prod() // merge_length + for frame_idx in range(video_grid_thw[index][0]): + curr_time = curr_timestamp[frame_idx] + video_placeholder += f"<{curr_time:.1f} seconds>" + video_placeholder += ( + self.vision_start_token + "<|placeholder|>" * frame_seqlen + self.vision_end_token + ) + if f"{self.vision_start_token}{self.video_token}{self.vision_end_token}" in text[i]: + text[i] = text[i].replace( + f"{self.vision_start_token}{self.video_token}{self.vision_end_token}", video_placeholder, 1 + ) + else: + # vllm may input video token directly + text[i] = text[i].replace(self.video_token, video_placeholder, 1) + index += 1 + + text[i] = text[i].replace("<|placeholder|>", self.video_token) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"]) + + if return_mm_token_type_ids: + array_ids = np.array(text_inputs["input_ids"]) + mm_token_type_ids = np.zeros_like(text_inputs["input_ids"]) + mm_token_type_ids[array_ids == self.image_token_id] = 1 + text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist() + + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors) + + def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs): + """Return `MultiModalData` with placeholder-token counts for the given + (height, width) image sizes and/or (num_frames, height, width) video sizes.""" + + vision_data = {} + if image_sizes is not None: + images_kwargs = Qwen3VLProcessorKwargs._defaults.get("images_kwargs", {}) + images_kwargs.update(kwargs) + merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size + + num_image_patches = [ + self.image_processor.get_number_of_image_patches(*image_size, images_kwargs) + for image_size in image_sizes + ] + num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches] + vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches}) + + if video_sizes is not None: + videos_kwargs = Qwen3VLProcessorKwargs._defaults.get("videos_kwargs", {}) + videos_kwargs.update(kwargs) + num_video_patches = [ + self.video_processor.get_number_of_video_patches(*video_size, videos_kwargs) + for video_size in video_sizes + ] + num_video_tokens = [(num_patches // merge_size**2) for num_patches in num_video_patches] + vision_data["num_video_tokens"] = num_video_tokens + + return MultiModalData(**vision_data) + + def post_process_image_text_to_text( + self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs + ): + """Batch-decode generated token ids to text via the tokenizer.""" + return self.tokenizer.batch_decode( + generated_outputs, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + def _calculate_timestamps(self, indices: Union[list[int], np.ndarray], video_fps: float, merge_size: int = 2): + if not isinstance(indices, list): + indices = indices.tolist() + if len(indices) % merge_size != 0: + indices.extend(indices[-1] for _ in range(merge_size - len(indices) % merge_size)) + timestamps = [idx / video_fps for idx in indices] + # frames are merged by self.merge_size, so we need to average the + # timestamps between the first/last frame within the temporal patch + timestamps = [ + (timestamps[i] + timestamps[i + merge_size - 1]) / 2 for i in range(0, len(timestamps), merge_size) + ] + return timestamps + + +__all__ = [ + "NemotronNanoV3BridgeProcessor", + "Qwen3VLVideoProcessorCustom", + "Siglip2ImageProcessorCustom", + "build_cosmos3_edge_processor", + "is_cosmos3_edge_native_snapshot", +] diff --git a/cosmos_framework/data/generator/processors/cosmos3_edge_processing_test.py b/cosmos_framework/data/generator/processors/cosmos3_edge_processing_test.py new file mode 100644 index 00000000..80cbcca7 --- /dev/null +++ b/cosmos_framework/data/generator/processors/cosmos3_edge_processing_test.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Feature-parity gate (G1) for the framework-native Cosmos3-Edge processor port. + +The embedded goldens pin exact token ids AND pixel-tensor sha256 hashes; the +fixtures must byte-match ``outputs/audit/processor_audit.py``. A pixel sha256 +mismatch with matching shapes means the patch layout regressed from the +original convention (raster patch-row order, ``(py, px, c)`` within-row) — +e.g. code "aligned" with transformers-main's native processor, which corrupts +vision features with these checkpoint weights; see +``outputs/audit/cosmos3_edge_native_vision_layout_bug.md``. +""" + +import hashlib +import json +import os + +import numpy as np +import pytest +import torch +from PIL import Image + +from cosmos_framework.data.generator.processors.cosmos3_edge_processing import ( + build_cosmos3_edge_processor, + is_cosmos3_edge_native_snapshot, +) + +# Renewed (native-metadata, no remote code) snapshot of nvidia/Cosmos3-Edge. +# Point COSMOS3_EDGE_SNAPSHOT_DIR at a local snapshot dir to run these checks; +# unset (the default) auto-skips them. +_SNAPSHOT_DIR = os.environ.get("COSMOS3_EDGE_SNAPSHOT_DIR", "") + +requires_snapshot = pytest.mark.skipif( + not os.path.isdir(_SNAPSHOT_DIR), + reason=f"renewed nvidia/Cosmos3-Edge snapshot not available at {_SNAPSHOT_DIR}", +) + +# Captured from the old remote-code processor (rev 28a0b8e, transformers 4.57.6); +# embedded so the test does not depend on the outputs/ tree. +_GOLDEN = { + "text_only": { + "num_tokens": 35, + "ids_sha256": "5916e7a832968ac9", + }, + "image": { + "num_tokens": 329, + "ids_sha256": "9cc248c1c2260d79", + "n_image_tokens": 300, + "pixel_values": {"shape": [1200, 768], "sha256": "074c173f37463c24"}, + "image_grid_thw": [[1, 30, 40]], + }, + "video": { + "num_tokens": 730, + "ids_sha256": "3c9bedc1730ecf67", + "n_video_tokens": 640, + "pixel_values_videos": {"shape": [2560, 768], "sha256": "8cf76eb9af3d0b5a"}, + "video_grid_thw": [[8, 16, 20]], + }, +} + + +def _ids_sha256(ids: list[int]) -> str: + return hashlib.sha256(json.dumps(ids).encode()).hexdigest()[:16] + + +def _tensor_sha256(t: torch.Tensor) -> str: + a = np.asarray(t.float().cpu(), dtype=np.float32) + return hashlib.sha256(a.tobytes()).hexdigest()[:16] + + +@pytest.fixture(scope="module") +def processor(): + return build_cosmos3_edge_processor(_SNAPSHOT_DIR) + + +def _apply(processor, messages, **kwargs): + return processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=False, + return_dict=True, + return_tensors="pt", + **kwargs, + ) + + +@requires_snapshot +def test_detection_rule_accepts_native_snapshot() -> None: + assert is_cosmos3_edge_native_snapshot(_SNAPSHOT_DIR) + + +def test_detection_rule_rejects_non_edge_dirs(tmp_path) -> None: + # Empty dir, non-dir, and an old-style (remote-code) config all miss the rule. + assert not is_cosmos3_edge_native_snapshot(str(tmp_path)) + assert not is_cosmos3_edge_native_snapshot(str(tmp_path / "missing")) + (tmp_path / "config.json").write_text(json.dumps({"model_type": "nemotron_siglip2"})) + assert not is_cosmos3_edge_native_snapshot(str(tmp_path)) + + +@requires_snapshot +def test_text_only_matches_golden(processor) -> None: + inputs = _apply( + processor, + [ + {"role": "user", "content": [{"type": "text", "text": "Describe the physics of a bouncing ball."}]}, + {"role": "assistant", "content": [{"type": "text", "text": "It decelerates under gravity."}]}, + ], + ) + ids = inputs["input_ids"][0].tolist() + assert len(ids) == _GOLDEN["text_only"]["num_tokens"] + assert _ids_sha256(ids) == _GOLDEN["text_only"]["ids_sha256"] + assert inputs["attention_mask"][0].tolist() == [1] * len(ids) + + +@requires_snapshot +def test_image_matches_golden(processor) -> None: + rng = np.random.RandomState(42) + image = Image.fromarray(rng.randint(0, 255, (480, 640, 3), dtype=np.uint8)) + inputs = _apply( + processor, + [ + { + "role": "user", + "content": [{"type": "image", "image": image}, {"type": "text", "text": "What is shown?"}], + }, + {"role": "assistant", "content": [{"type": "text", "text": "A test pattern."}]}, + ], + ) + golden = _GOLDEN["image"] + ids = inputs["input_ids"][0].tolist() + assert len(ids) == golden["num_tokens"] + assert _ids_sha256(ids) == golden["ids_sha256"] + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + assert ids.count(image_token_id) == golden["n_image_tokens"] + + assert list(inputs["pixel_values"].shape) == golden["pixel_values"]["shape"] + # Pins the exact float32 bytes, i.e. the OLD (py, px, c)/raster patch layout. + assert _tensor_sha256(inputs["pixel_values"]) == golden["pixel_values"]["sha256"] + assert inputs["image_grid_thw"].tolist() == golden["image_grid_thw"] + + +@requires_snapshot +def test_video_matches_golden(processor) -> None: + rng = np.random.RandomState(7) + frames = [Image.fromarray(rng.randint(0, 255, (256, 320, 3), dtype=np.uint8)) for _ in range(8)] + metadata = {"fps": 2.0, "total_num_frames": len(frames), "frames_indices": list(range(len(frames)))} + inputs = _apply( + processor, + [ + { + "role": "user", + "content": [ + {"type": "video", "video": frames}, + {"type": "text", "text": "Is this physically plausible?"}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": "Yes."}]}, + ], + videos_kwargs={"do_sample_frames": False, "video_metadata": metadata}, + ) + golden = _GOLDEN["video"] + ids = inputs["input_ids"][0].tolist() + assert len(ids) == golden["num_tokens"] + assert _ids_sha256(ids) == golden["ids_sha256"] + video_token_id = processor.tokenizer.convert_tokens_to_ids(processor.video_token) + assert ids.count(video_token_id) == golden["n_video_tokens"] + + assert list(inputs["pixel_values_videos"].shape) == golden["pixel_values_videos"]["shape"] + # Pins the exact float32 bytes, i.e. the OLD (py, px, c)/raster patch layout. + assert _tensor_sha256(inputs["pixel_values_videos"]) == golden["pixel_values_videos"]["sha256"] + assert inputs["video_grid_thw"].tolist() == golden["video_grid_thw"] + + # Per-frame "<{t:.1f} seconds>" timestamp interleaving (fps=2.0, indices 0..7). + decoded = processor.tokenizer.decode(ids) + assert "<0.0 seconds>" in decoded + assert "<0.5 seconds>" in decoded + + +@requires_snapshot +def test_wrapper_interface_surface(processor) -> None: + """The attributes the framework wrappers consume off the raw processor.""" + assert processor.image_token == "<|image_pad|>" + assert processor.video_token == "<|video_pad|>" + assert processor.image_processor.size == {"shortest_edge": 65536, "longest_edge": 16777216} + assert processor.image_processor.patch_size == 16 + assert processor.image_processor.merge_size == 2 + assert processor.video_processor.patch_size == 16 + assert processor.video_processor.temporal_patch_size == 1 + assert processor.video_processor.merge_size == 2 + assert processor.tokenizer.eos_token_id is not None + + +@requires_snapshot +def test_build_processor_routes_local_edge_dir_to_nemotron_wrapper() -> None: + from cosmos_framework.data.generator.processors import Nemotron3DenseVLProcessor, build_processor + + wrapper = build_processor(_SNAPSHOT_DIR) + assert isinstance(wrapper, Nemotron3DenseVLProcessor) + # Dataloader helper attributes resolved through the ported sub-processors. + assert wrapper.patch_size == 16 + assert wrapper.temporal_patch_size == 1 + assert wrapper.merge_size == 2 + assert wrapper.image_token_id == 19 + assert wrapper.video_token_id == 18 diff --git a/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py b/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py index 37795233..afa698b9 100644 --- a/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py +++ b/cosmos_framework/data/reasoner/processors/nemotron3densevl_processor.py @@ -10,6 +10,10 @@ from qwen_vl_utils.vision_process import smart_resize from transformers.models.auto.processing_auto import AutoProcessor +from cosmos_framework.data.generator.processors.cosmos3_edge_processing import ( + build_cosmos3_edge_processor, + is_cosmos3_edge_native_snapshot, +) from cosmos_framework.utils import log from cosmos_framework.utils.reasoner.pretrained_models_downloader import maybe_download_hf_model_from_s3 @@ -67,7 +71,12 @@ def __init__( name, credentials, bucket, include_model_weights=False ) - self.processor = AutoProcessor.from_pretrained(model_name_or_path_local, trust_remote_code=True) + if is_cosmos3_edge_native_snapshot(model_name_or_path_local): + # AutoProcessor on transformers 4.x silently degrades to a bare tokenizer + # for renewed (no remote code) Cosmos3-Edge snapshots; use the native port. + self.processor = build_cosmos3_edge_processor(model_name_or_path_local) + else: + self.processor = AutoProcessor.from_pretrained(model_name_or_path_local, trust_remote_code=True) log.info("Successfully loaded processor from local cache") if hasattr(self.processor, "image_token"): diff --git a/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py index bcb2e2c2..85e3007f 100644 --- a/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py +++ b/cosmos_framework/data/reasoner/processors/nemotronvl_processor.py @@ -278,9 +278,7 @@ def __init__( else: self.video_token_id = None self.eos_id = self.processor.tokenizer.eos_token_id - self.pad_id = self.processor.tokenizer.convert_tokens_to_ids( - "" - ) + self.pad_id = self.processor.tokenizer.convert_tokens_to_ids("") self.vision_end_id = self.processor.tokenizer.convert_tokens_to_ids("") # Helper attributes for the dataloader video decoding function diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 4ef9092c..2d634406 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -635,6 +635,35 @@ def _build_action_data(self, model_config: "OmniMoTModelConfig", sample_meta: Sa _ReasonerRepetitionPenalty = Annotated[float, pydantic.Field(gt=0)] +def _mapping_get(node: Any, key: str) -> Any: + """``.get`` that tolerates None and non-mapping nodes (dict / DictConfig / LazyDict).""" + getter = getattr(node, "get", None) + return getter(key) if callable(getter) else None + + +def _reasoner_vision_capable(model_config: "OmniMoTModelConfig") -> bool: + """Whether the checkpoint's reasoner LM can encode image/video prompts. + + Qwen-family reasoners carry the ViT inside the checkpoint, gated by + ``include_visual``; a falsy value (e.g. an old ``--no-vit`` export, or a + text-only LLM backbone) means the tower was never built, so a vision + prompt would only fail mid-batch at the ``hasattr(causal_lm, "visual")`` + gate. Edge's Nemotron reasoner sets ``include_visual=None`` by design and + loads its SigLIP2 tower lazily (``_ensure_vision_tower``), so it is always + vision-capable. Unknown families are treated as capable (never block). + """ + vlm_config = model_config.vlm_config + model_instance = getattr(vlm_config, "model_instance", None) + target = str(_mapping_get(model_instance, "_target_") or "") + model_name = getattr(vlm_config, "model_name", "") or "" + if "Nemotron3DenseVLTextForCausalLM" in target or "Cosmos3-Edge" in model_name: + return True + if "Qwen" not in target and not model_name.startswith("Qwen/"): + return True + include_visual = _mapping_get(_mapping_get(model_instance, "config"), "include_visual") + return bool(include_visual) + + class ReasonerDataArgs(ArgsBase): """Resolved reasoner (VLM) text-generation arguments. All fields are ``| None`` so non-reasoner samples (which never populate these) pass ``OmniSampleArgs`` validation; @@ -677,6 +706,14 @@ def _build_reasoner_data(self, model_config: "OmniMoTModelConfig", sample_meta: self = cast("SampleDataOverrides", self) if not self.prompt.strip(): raise ValueError("Reasoner inference requires a non-empty 'prompt'.") + if self.vision_path is not None and not _reasoner_vision_capable(model_config): + raise ValueError( + f"Reasoner sample {getattr(self, 'name', None)!r} conditions on vision input " + f"'{self.vision_path}', but the loaded checkpoint's reasoner LM has no visual tower " + "(model config include_visual is false — e.g. a --no-vit export or a text-only LLM " + "backbone). Use a text-only prompt, or re-export the checkpoint with the default " + "--vit / use a full checkpoint for image/video reasoner prompts." + ) class _TransferDataBase: @@ -1005,9 +1042,15 @@ class OmniSampleOverrides( "Qwen/Qwen3-VL-32B-Instruct": "32B", "Qwen/Qwen3-VL-30B-A3B-Instruct": "30B-A3B", "Qwen/Qwen3-VL-235B-A22B-Instruct": "235B-A22B", + "nvidia/Cosmos3-Edge-Reasoner": "2B", } _RESOLUTION_SHIFT_DEFAULTS: ClassVar[dict[(ModelSize, Resolution), float]] = { + # 2B rows mirror Cosmos3-Edge's training shift (Cosmos3-Edge.yaml + # rectified_flow_training_config.shift: 256->3, 480->5, 720->10). + ("2B", "256"): 3.0, + ("2B", "480"): 5.0, + ("2B", "720"): 10.0, ("8B", "256"): 3.0, ("8B", "480"): 5.0, ("8B", "720"): 10.0, @@ -1057,7 +1100,9 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs: ) if not shift_configured and not sample_meta.model_mode.is_reasoner: - model_size = self._VLM_MODEL_SIZE[model_config.vlm_config.model_name] + # Unregistered backbone names (custom fine-tunes) simply skip shift + # defaulting rather than KeyError-ing. + model_size = self._VLM_MODEL_SIZE.get(model_config.vlm_config.model_name) key = (model_size, self.resolution) if key in self._RESOLUTION_SHIFT_DEFAULTS: self.shift = self._RESOLUTION_SHIFT_DEFAULTS[key] @@ -1105,6 +1150,15 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs: revision="main", ), ), + "Cosmos3-Edge": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["2B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Edge.yaml"), + s3_uri="s3://bucket1/cosmos3_vfm/cosmos3_ga_midtraining/cosmos3_ga_4bm2b_v1_midtrain_0630a/checkpoints/iter_000028000/", + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Edge", + revision="main", + ), + ), "Cosmos3-Super": CheckpointConfig( model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"], config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"), @@ -1141,6 +1195,47 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs: # downloading the base Cosmos3-Super repo just for the tokenizer. vlm_processor_from_checkpoint=True, ), + "Cosmos3-Super-Text2Image-4Step": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"), + s3_uri="s3://bucket1/cosmos3_vfm/cosmos3_ga_text2image_4step/", + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Super-Text2Image-4Step", + revision="1ba94110bc118f479bbd5e461e79d685d74b2554", + ), + experiment_overrides=( + "model.config.resolution='768'", + "model.config.action_gen=false", + "model.config.sound_gen=false", + "model.config.sound_dim=null", + "model.config.sound_tokenizer=null", + "model.config.rectified_flow_training_config.shift.720=5", + "model.config.rectified_flow_training_config.shift.768=5", + "model.config.tokenizer.encode_chunk_frames.768=12", + "model.config.tokenizer.encode_exact_durations=null", + "model.config.fixed_step_sampler_config.t_list=[1.0,0.9375,0.8333333333333334,0.625]", + "model.config.fixed_step_sampler_config.sample_type=sde", + ), + ), + "Cosmos3-Super-Image2Video-4Step": CheckpointConfig( + model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"], + config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"), + s3_uri="s3://bucket1/cosmos3_vfm/cosmos3_ga_image2video_4step/", + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Super-Image2Video-4Step", + revision="f85d3335d2ad8b352462cecbd637aa980cec9688", + ), + experiment_overrides=( + "model.config.resolution='480'", + "model.config.action_gen=false", + "model.config.sound_gen=false", + "model.config.sound_dim=null", + "model.config.sound_tokenizer=null", + "model.config.diffusion_expert_config.base_fps=16", + "model.config.fixed_step_sampler_config.t_list=[1.0,0.9375,0.8333333333333334,0.625]", + "model.config.fixed_step_sampler_config.sample_type=sde", + ), + ), } DEFAULT_CHECKPOINT_NAME = "Cosmos3-Nano" DEFAULT_CHECKPOINT = _CHECKPOINTS[DEFAULT_CHECKPOINT_NAME] diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index b3f9b0e7..512d3dca 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 import json +import os import types from pathlib import Path @@ -125,12 +126,80 @@ def test_checkpoints(): for name, ckpt in OmniSetupOverrides.CHECKPOINTS.items(): assert ckpt.hf.repository.split("/")[0] == "nvidia" + # The released 4-step repositories are gated. Their pinned registry + # metadata is covered below without requiring CI to hold an HF token. + if ckpt.hf.repository.endswith("-4Step") and not os.environ.get("HF_TOKEN"): + continue + # Download a file to ensure that the repository/revision is valid. - # (The published checkpoint.json no longer encodes the source S3 URI — - # it is empty for most checkpoints and a stale local path for the rest — - # so we only validate that the repo/revision resolves and parses.) - ckpt_hf = ckpt.hf.model_copy(update=dict(include=("checkpoint.json",))) - json.loads((Path(ckpt_hf.download()) / "checkpoint.json").read_text()) + # Native checkpoints store ``config.json`` at the root; Diffusers + # checkpoints store the transformer config in ``transformer/``. + ckpt_hf = ckpt.hf.model_copy(update=dict(include=("config.json", "transformer/config.json"))) + checkpoint_dir = Path(ckpt_hf.download()) + config_paths = [ + path + for path in (checkpoint_dir / "config.json", checkpoint_dir / "transformer/config.json") + if path.is_file() + ] + assert config_paths, f"No model config found for {name}" + json.loads(config_paths[0].read_text()) + + +@pytest.mark.parametrize( + ("checkpoint_name", "repository", "revision", "expected_resolution", "expected_base_fps"), + [ + ( + "Cosmos3-Super-Text2Image-4Step", + "nvidia/Cosmos3-Super-Text2Image-4Step", + "1ba94110bc118f479bbd5e461e79d685d74b2554", + "768", + 24, + ), + ( + "Cosmos3-Super-Image2Video-4Step", + "nvidia/Cosmos3-Super-Image2Video-4Step", + "f85d3335d2ad8b352462cecbd637aa980cec9688", + "480", + 16, + ), + ], +) +def test_distilled_checkpoint_uses_published_fixed_step_schedule( + tmp_path: Path, + checkpoint_name: str, + repository: str, + revision: str, + expected_resolution: str, + expected_base_fps: int, +) -> None: + args = OmniSetupOverrides( + checkpoint_path=checkpoint_name, + output_dir=tmp_path / "outputs", + ).build_setup(world_size=4, local_world_size=4, device_memory_bytes=_H100_MEMORY_BYTES) + + assert args.checkpoint_hf is not None + assert args.checkpoint_hf.repository == repository + assert args.checkpoint_hf.revision == revision + assert args.vlm_processor_from_checkpoint is False + + model_config = args.load_model_config_dict()["config"] + assert model_config["action_gen"] is False + assert model_config["sound_gen"] is False + assert model_config["sound_dim"] is None + assert model_config["sound_tokenizer"] is None + assert model_config["resolution"] == expected_resolution + assert model_config["diffusion_expert_config"]["base_fps"] == expected_base_fps + + if checkpoint_name == "Cosmos3-Super-Text2Image-4Step": + assert model_config["rectified_flow_training_config"]["shift"]["720"] == 5 + assert model_config["rectified_flow_training_config"]["shift"]["768"] == 5 + assert model_config["tokenizer"]["encode_chunk_frames"]["768"] == 12 + assert model_config["tokenizer"]["encode_exact_durations"] is None + + fixed_step_sampler_config = model_config["fixed_step_sampler_config"] + assert fixed_step_sampler_config is not None + assert fixed_step_sampler_config["t_list"] == [1.0, 0.9375, 0.8333333333333334, 0.625] + assert fixed_step_sampler_config["sample_type"] == "sde" def test_setup_args(tmp_path: Path): diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 13ee789e..d67cbac1 100644 --- a/cosmos_framework/inference/common/args.py +++ b/cosmos_framework/inference/common/args.py @@ -382,7 +382,11 @@ def load_model_config_dict(self) -> dict: case ConfigFileType.MODULE: return unstructure_config(self.load_config().model) case ConfigFileType.YAML | ConfigFileType.JSON: - return load_model_config_from_hf_config(deserialize_config_dict(Path(self.config_file))) + config_dict = deserialize_config_dict(Path(self.config_file)) + overrides_omegaconf = OmegaConf.from_dotlist(self.experiment_overrides) + config_dict = OmegaConf.to_container(OmegaConf.merge(config_dict, overrides_omegaconf), resolve=False) + assert isinstance(config_dict, dict) + return load_model_config_from_hf_config(config_dict) case _: assert_never(self.config_file_type) @@ -424,9 +428,18 @@ class CheckpointType(StrEnum): @classmethod def from_path(cls, path: Path) -> Self: - has_hf_weights = any(path.glob("*.safetensors")) or any(path.glob("*.safetensors.index.json")) + transformer_path = path / "transformer" + has_root_hf_weights = any(path.glob("*.safetensors")) or any(path.glob("*.safetensors.index.json")) + has_diffusers_hf_weights = (path / "model_index.json").is_file() and ( + any(transformer_path.glob("*.safetensors")) + or any(transformer_path.glob("*.safetensors.index.json")) + ) + has_hf_weights = has_root_hf_weights or has_diffusers_hf_weights if has_hf_weights: - if not (path / "config.json").exists(): + has_hf_config = (path / "config.json").is_file() or ( + has_diffusers_hf_weights and (transformer_path / "config.json").is_file() + ) + if not has_hf_config: raise ValueError(f"Invalid Hugging Face checkpoint: {path}") return cls("hf") if any(path.glob("*.distcp")): @@ -465,6 +478,13 @@ class CheckpointConfig(pydantic.BaseModel): just to obtain the tokenizer. """ + experiment_overrides: tuple[str, ...] = () + """Config defaults required by this checkpoint. + + These are applied before command-line experiment overrides so callers can + still override a checkpoint default explicitly. + """ + def download(self) -> str: return self.hf.download() @@ -543,6 +563,9 @@ def _build_checkpoint(self, checkpoints: dict[str, CheckpointConfig]): self.config_file = checkpoint.config_file self.checkpoint_hf = checkpoint.hf self.vlm_processor_from_checkpoint = checkpoint.vlm_processor_from_checkpoint + for value in reversed(checkpoint.experiment_overrides): + if value not in self.experiment_overrides: + self.experiment_overrides.insert(0, value) elif self.checkpoint_path.startswith("s3://"): self.checkpoint_type = CheckpointType.DCP self.checkpoint_path = self.checkpoint_path.rstrip("/") diff --git a/cosmos_framework/inference/common/args_test.py b/cosmos_framework/inference/common/args_test.py index ae56295e..20c0d7f4 100644 --- a/cosmos_framework/inference/common/args_test.py +++ b/cosmos_framework/inference/common/args_test.py @@ -8,13 +8,23 @@ import pytest from cosmos_framework.inference.args import DEFAULT_CHECKPOINT, DEFAULT_CHECKPOINT_NAME -from cosmos_framework.inference.common.args import CheckpointConfig, CheckpointOverrides, download_file +from cosmos_framework.inference.common.args import CheckpointConfig, CheckpointOverrides, CheckpointType, download_file CHECKPOINTS: dict[str, CheckpointConfig] = { DEFAULT_CHECKPOINT_NAME: DEFAULT_CHECKPOINT, } +def test_checkpoint_type_from_diffusers_layout(tmp_path: Path) -> None: + transformer_path = tmp_path / "transformer" + transformer_path.mkdir() + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + (transformer_path / "config.json").write_text("{}", encoding="utf-8") + (transformer_path / "diffusion_pytorch_model.safetensors.index.json").write_text("{}", encoding="utf-8") + + assert CheckpointType.from_path(tmp_path) == CheckpointType.HF + + def test_download_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): # Disable the URL cache; this test asserts each download is independent. monkeypatch.delenv("COSMOS_DOWNLOAD_CACHE_DIR", raising=False) diff --git a/cosmos_framework/inference/common/checkpoints.py b/cosmos_framework/inference/common/checkpoints.py index 62913aa9..a46a1c43 100644 --- a/cosmos_framework/inference/common/checkpoints.py +++ b/cosmos_framework/inference/common/checkpoints.py @@ -207,6 +207,25 @@ def register_checkpoints(): ) ) + # The backbone URI Edge SFT recipes pin (edge_model_config), mapped to the + # public nvidia/Cosmos3-Edge repo. Deliberately NO 'include' filter: + # training's backbone seeding needs the full repo (the root safetensors + # index points into transformer/*); export_model narrows its own download + # to EDGE_VIT_BUNDLE_HF_INCLUDE at the call site. + register_checkpoint( + CheckpointConfig( + uuid=uuid4().hex, + name="Cosmos3-Edge-Reasoner-590c1c0", + s3=CheckpointDirS3( + uri="s3://bucket/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0", + ), + hf=CheckpointDirHf( + repository="nvidia/Cosmos3-Edge", + revision="main", + ), + ), + ) + register_checkpoint( CheckpointConfig( uuid=uuid4().hex, diff --git a/cosmos_framework/inference/common/config.py b/cosmos_framework/inference/common/config.py index 73d65c9b..80c3eec9 100644 --- a/cosmos_framework/inference/common/config.py +++ b/cosmos_framework/inference/common/config.py @@ -266,6 +266,16 @@ def _to_safe_string(value): config_converter = cattrs.preconf.json.make_converter() +# NoneType +def _structure_none(data: Any, _cls: Any) -> None: + if data is not None: + raise TypeError(f"Expected None, got {type(data).__name__}") + return None + + +config_converter.register_structure_hook(type(None), _structure_none) + + # type def _is_type_cls(cls: Any) -> bool: return cls in [type, abc.ABCMeta] or typing.get_origin(cls) is type diff --git a/cosmos_framework/inference/common/config_test.py b/cosmos_framework/inference/common/config_test.py index 8be427bf..73ce975e 100644 --- a/cosmos_framework/inference/common/config_test.py +++ b/cosmos_framework/inference/common/config_test.py @@ -13,13 +13,11 @@ from cosmos_framework.inference.common.args import DEFAULT_CONFIG_FILE from cosmos_framework.inference.common.config import ( _is_type_cls, - apply_config_replacements, config_converter, deserialize_config, load_config, serialize_config, structure_config, - undo_config_replacements, unstructure_config, ) from cosmos_framework.utils.flags import TRAINING @@ -113,6 +111,11 @@ def round_trip(obj): assert structure_config(config_dict, Config) == structured_config +def test_config_converter_handles_explicit_none_type() -> None: + assert config_converter.structure(None, type(None)) is None + with pytest.raises(TypeError, match="Expected None"): + config_converter.structure("not-none", type(None)) + if TRAINING: diff --git a/cosmos_framework/inference/common/distillation_export.py b/cosmos_framework/inference/common/distillation_export.py new file mode 100644 index 00000000..673d036d --- /dev/null +++ b/cosmos_framework/inference/common/distillation_export.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Portable student-only checkpoint export helpers.""" + +from collections.abc import Callable +from typing import Any + + +def build_student_checkpoint_metadata(*, use_ema_weights: bool) -> dict[str, str | bool]: + """Build portable metadata without source checkpoint or credential paths.""" + return { + "checkpoint_type": "hf", + "source_weights": "ema" if use_ema_weights else "regular", + "student_only": True, + } + + +def sanitize_student_model_config( + model_dict: dict[str, Any], + *, + base_model_target: str, + base_config_type: str, + base_config_field_names: set[str], +) -> None: + """Convert a distillation model config into a portable student-only base config.""" + config = model_dict.get("config") + if not isinstance(config, dict): + raise TypeError("Expected model config to be a dictionary.") + + model_dict["_target_"] = base_model_target + config["_type"] = base_config_type + + allowed_keys = base_config_field_names | {"_type"} + for key in tuple(config): + if key not in allowed_keys: + del config[key] + + compile_config = config.get("compile") + if isinstance(compile_config, dict): + compile_config["enabled"] = False + + +def sanitize_student_public_model_config( + model_dict: dict[str, Any], + *, + public_vlm_tokenizer_target: str | None = None, +) -> None: + """Replace internal loader settings with portable public checkpoint aliases.""" + config = model_dict.get("config") + if not isinstance(config, dict): + raise TypeError("Expected model config to be a dictionary.") + + for tokenizer_key in ("tokenizer", "sound_tokenizer"): + tokenizer_config = config.get(tokenizer_key) + if not isinstance(tokenizer_config, dict): + continue + if "bucket_name" in tokenizer_config: + tokenizer_config["bucket_name"] = "bucket" + if "object_store_credential_path_pretrained" in tokenizer_config: + tokenizer_config["object_store_credential_path_pretrained"] = "" + + vlm_config = config.get("vlm_config") + if not isinstance(vlm_config, dict): + return + + pretrained_weights = vlm_config.get("pretrained_weights") + if isinstance(pretrained_weights, dict): + pretrained_weights["enabled"] = False + pretrained_weights["backbone_path"] = "" + pretrained_weights["credentials_path"] = "" + pretrained_weights["enable_gcs_patch_in_boto3"] = False + + tokenizer_config = vlm_config.get("tokenizer") + if isinstance(tokenizer_config, dict): + if public_vlm_tokenizer_target is not None: + tokenizer_config["_target_"] = public_vlm_tokenizer_target + if "config_variant" in tokenizer_config: + tokenizer_config["config_variant"] = "hf" + + +def resolve_vision_checkpoint_path( + *, + local_path: str | None, + configured_uri: str, + download_checkpoint: Callable[[str], str], +) -> str: + """Use a local vision checkpoint when supplied, otherwise download the configured checkpoint.""" + if local_path is not None: + return local_path + return download_checkpoint(configured_uri) diff --git a/cosmos_framework/inference/common/distillation_export_test.py b/cosmos_framework/inference/common/distillation_export_test.py new file mode 100644 index 00000000..50007783 --- /dev/null +++ b/cosmos_framework/inference/common/distillation_export_test.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + + +from cosmos_framework.inference.common import distillation_export +from cosmos_framework.inference.common.distillation_export import ( + build_student_checkpoint_metadata, + sanitize_student_model_config, +) + + + +def test_sanitize_student_model_config_removes_distillation_state() -> None: + fixed_step_sampler_config = { + "sample_type": "ode", + "t_list": [1.0, 0.75, 0.5, 0.25], + } + model_dict = { + "_target_": "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", + "_recursive_": False, + "config": { + "_type": "cosmos_framework.configs.base.experiment.distillation.dmd2_config.DMD2RFConfig", + "_metadata": { + "object_type": ("cosmos_framework.configs.base.experiment.distillation.dmd2_config.DMD2RFConfig"), + }, + "ema": {"enabled": False}, + "compile": {"enabled": True, "compiled_region": "language"}, + "fixed_step_sampler_config": fixed_step_sampler_config, + "vlm_config": {"model_name": "student"}, + "vlm_config_teacher": {"model_name": "teacher"}, + "vlm_config_fake_score": {"model_name": "fake_score"}, + "load_teacher_weights": True, + "teacher_load_from": {"load_path": "internal-teacher"}, + "student_load_from": {"load_path": "internal-student"}, + "optimizer": {"net": {}, "fake_score": {}}, + }, + } + + sanitize_student_model_config( + model_dict, + base_model_target="cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", + base_config_type="cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig", + base_config_field_names={"compile", "ema", "fixed_step_sampler_config", "vlm_config"}, + ) + + assert model_dict == { + "_target_": "cosmos_framework.model.generator.omni_mot_model.OmniMoTModel", + "_recursive_": False, + "config": { + "_type": "cosmos_framework.configs.base.defaults.model_config.OmniMoTModelConfig", + "ema": {"enabled": False}, + "compile": {"enabled": False, "compiled_region": "language"}, + "fixed_step_sampler_config": fixed_step_sampler_config, + "vlm_config": {"model_name": "student"}, + }, + } + + +def test_build_student_checkpoint_metadata_omits_source_paths() -> None: + assert build_student_checkpoint_metadata(use_ema_weights=True) == { + "checkpoint_type": "hf", + "source_weights": "ema", + "student_only": True, + } + assert build_student_checkpoint_metadata(use_ema_weights=False) == { + "checkpoint_type": "hf", + "source_weights": "regular", + "student_only": True, + } + + +def test_sanitize_student_public_model_config_removes_internal_loaders() -> None: + model_dict = { + "config": { + "tokenizer": { + "bucket_name": "internal-checkpoint-bucket", + "object_store_credential_path_pretrained": "/path/to/source.secret", + "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", + }, + "sound_tokenizer": { + "bucket_name": "internal-checkpoint-bucket", + "object_store_credential_path_pretrained": "/path/to/source.secret", + "avae_path": "pretrained/tokenizers/audio/avae/avae.ckpt", + }, + "vlm_config": { + "pretrained_weights": { + "enabled": True, + "backbone_path": "s3://internal-checkpoint-bucket/reasoner", + "credentials_path": "/path/to/source.secret", + "enable_gcs_patch_in_boto3": True, + }, + "tokenizer": { + "_target_": ( + "cosmos_framework.configs.base.experiment.distillation_implementation." + "_create_oss_tokenizer_with_internal_download" + ), + "config_variant": "gcp", + "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct", + }, + }, + } + } + + distillation_export.sanitize_student_public_model_config( + model_dict, + public_vlm_tokenizer_target=( + "cosmos_framework.configs.base.defaults.reasoner.create_qwen2_tokenizer_with_download" + ), + ) + + assert model_dict == { + "config": { + "tokenizer": { + "bucket_name": "bucket", + "object_store_credential_path_pretrained": "", + "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", + }, + "sound_tokenizer": { + "bucket_name": "bucket", + "object_store_credential_path_pretrained": "", + "avae_path": "pretrained/tokenizers/audio/avae/avae.ckpt", + }, + "vlm_config": { + "pretrained_weights": { + "enabled": False, + "backbone_path": "", + "credentials_path": "", + "enable_gcs_patch_in_boto3": False, + }, + "tokenizer": { + "_target_": ( + "cosmos_framework.configs.base.defaults.reasoner.create_qwen2_tokenizer_with_download" + ), + "config_variant": "hf", + "pretrained_model_name": "Qwen/Qwen3-VL-32B-Instruct", + }, + }, + } + } + + +def test_resolve_vision_checkpoint_path_prefers_local_override() -> None: + fallback_called = False + + def download_checkpoint(_configured_uri: str) -> str: + nonlocal fallback_called + fallback_called = True + return "/downloaded/vision" + + path = distillation_export.resolve_vision_checkpoint_path( + local_path="/local/vision", + configured_uri="s3://internal/vision", + download_checkpoint=download_checkpoint, + ) + + assert path == "/local/vision" + assert fallback_called is False diff --git a/cosmos_framework/inference/common/inference.py b/cosmos_framework/inference/common/inference.py index c492756a..88c0a80c 100644 --- a/cosmos_framework/inference/common/inference.py +++ b/cosmos_framework/inference/common/inference.py @@ -5,9 +5,10 @@ import dataclasses import traceback from abc import ABC, abstractmethod -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import contextmanager, nullcontext from dataclasses import dataclass +from pathlib import Path from typing import TYPE_CHECKING, Any, ContextManager, Self, Sequence, final import torch @@ -22,6 +23,31 @@ from cosmos_framework.auxiliary.guardrail.common.core import GuardrailRunner +def _download_on_rank0(download: Callable[[], str | Path]) -> Path: + """Download once and share the resulting path across distributed ranks.""" + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return Path(download()) + + payload: list[str | None] = [None, None] + local_error: Exception | None = None + if torch.distributed.get_rank() == 0: + try: + payload[0] = str(download()) + except Exception as error: + local_error = error + payload[1] = f"{type(error).__name__}: {error}" + + torch.distributed.broadcast_object_list(payload, src=0) + path, error_message = payload + if error_message is not None: + if local_error is not None: + raise local_error + raise RuntimeError(f"Rank 0 download failed: {error_message}") + if path is None: + raise RuntimeError("Rank 0 download returned no path") + return Path(path) + + @contextlib.contextmanager def sync_distributed_errors(): """Catches local exceptions and synchronizes the error state across all distributed ranks. diff --git a/cosmos_framework/inference/common/inference_test.py b/cosmos_framework/inference/common/inference_test.py index 76a18188..93a5b2c4 100644 --- a/cosmos_framework/inference/common/inference_test.py +++ b/cosmos_framework/inference/common/inference_test.py @@ -1,14 +1,52 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 +from pathlib import Path +from unittest.mock import Mock + import numpy as np +import pytest from cosmos_framework.inference.common.args import GuardrailArgs -from cosmos_framework.inference.common.inference import GuardrailRunners -from cosmos_framework.auxiliary.guardrail.common import presets +from cosmos_framework.inference.common.inference import GuardrailRunners, _download_on_rank0 + + +def test_download_on_rank0_broadcasts_shared_path(monkeypatch: pytest.MonkeyPatch) -> None: + import torch + + download = Mock(return_value="/shared/cache/model") + broadcast = Mock() + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast) + + assert _download_on_rank0(download) == Path("/shared/cache/model") + download.assert_called_once_with() + broadcast.assert_called_once_with(["/shared/cache/model", None], src=0) + + +def test_download_on_nonzero_rank_reuses_broadcast_path(monkeypatch: pytest.MonkeyPatch) -> None: + import torch + + download = Mock() + + def broadcast(payload: list[str | None], *, src: int) -> None: + assert src == 0 + payload[:] = ["/shared/cache/model", None] + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 1) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast) + + assert _download_on_rank0(download) == Path("/shared/cache/model") + download.assert_not_called() + +def test_guardrail_runners() -> None: + from cosmos_framework.auxiliary.guardrail.common import presets -def test_guardrail_runners(): guardrail_args = GuardrailArgs(guardrails=True, offload_guardrail_models=False) runners = GuardrailRunners.create(guardrail_args) assert runners.text is not None diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index aff1a057..1d768d06 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -19,6 +19,8 @@ _TARGET_ALIASES = { "projects.cosmos3.vfm.configs.base.defaults.vlm.create_qwen2_tokenizer_with_download": "create_qwen2_tokenizer_with_download", "projects.cosmos3.vfm.configs.base.defaults.vlm.create_vlm_config": "create_vlm_config", + "projects.cosmos3.vfm.models.mot.unified_mot.Nemotron3DenseVLMoTConfig.from_json_file": "nemotron3_dense_vl_mot_config_from_json_file", + "projects.cosmos3.vfm.models.mot.unified_mot.Nemotron3DenseVLTextForCausalLM": "nemotron3_dense_vl_text_for_causal_lm", "projects.cosmos3.vfm.models.mot.unified_mot.Qwen3VLMoTConfig.from_json_file": "qwen3_vl_mot_config_from_json_file", "projects.cosmos3.vfm.models.mot.unified_mot.Qwen3VLTextForCausalLM": "qwen3_vl_text_for_causal_lm", "projects.cosmos3.vfm.models.omni_mot_model.OmniMoTModel": "omni_mot_model", @@ -32,6 +34,7 @@ "projects.cosmos3.vfm.configs.base.defaults.compile.CompileConfig": "compile_config", "projects.cosmos3.vfm.configs.base.defaults.ema.EMAConfig": "ema_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.DiffusionExpertConfig": "diffusion_expert_config", + "projects.cosmos3.vfm.configs.base.defaults.model_config.FixedStepSamplerConfig": "fixed_step_sampler_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.LBLConfig": "lbl_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.OmniMoTModelConfig": "omni_mot_model_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowInferenceConfig": "rectified_flow_inference_config", diff --git a/cosmos_framework/inference/common/public_model_config_test.py b/cosmos_framework/inference/common/public_model_config_test.py index 40e7e852..de4e018c 100644 --- a/cosmos_framework/inference/common/public_model_config_test.py +++ b/cosmos_framework/inference/common/public_model_config_test.py @@ -32,6 +32,11 @@ def test_public_model_config_round_trip_removes_internal_metadata(): "_target_": "cosmos_framework.configs.base.defaults.compile.CompileConfig", "enabled": False, }, + "fixed_step_sampler_config": { + "_type": "cosmos_framework.configs.base.defaults.model_config.FixedStepSamplerConfig", + "sample_type": "sde", + "t_list": [1.0, 0.9375, 0.8333333333333334, 0.625], + }, "tokenizer": { "_target_": "cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface", "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth", @@ -61,6 +66,7 @@ def test_public_model_config_round_trip_removes_internal_metadata(): assert public_model_config["_target"] == "omni_mot_model" assert public_model_config["config"]["_type"] == "omni_mot_model_config" assert public_model_config["config"]["compile"]["_target"] == "compile_config" + assert public_model_config["config"]["fixed_step_sampler_config"]["_type"] == "fixed_step_sampler_config" assert "projects.cosmos3" not in text assert "projects/cosmos3" not in text assert "cosmos3._src" not in text diff --git a/cosmos_framework/inference/configs/model/Cosmos3-Edge.yaml b/cosmos_framework/inference/configs/model/Cosmos3-Edge.yaml new file mode 100644 index 00000000..1cabefab --- /dev/null +++ b/cosmos_framework/inference/configs/model/Cosmos3-Edge.yaml @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +_metadata: + args: + config_file: cosmos3/_src/vfm/configs/base/config.py + config_file_type: module + experiment: cosmos3_ga_4bm2b_v1_midtrain + experiment_overrides: [] +model: + _recursive_: false + _target_: cosmos3._src.vfm.models.omni_mot_model.OmniMoTModel + config: + _type: cosmos3._src.vfm.configs.base.defaults.model_config.OmniMoTModelConfig + action_gen: true + activation_checkpointing: + _type: cosmos3._src.vfm.configs.base.defaults.activation_checkpointing.ActivationCheckpointingConfig + determinism_check: default + mode: selective + preserve_rng_state: true + save_ops_regex: + - fmha + causal_training_strategy: none + compile: + _type: cosmos3._src.vfm.configs.base.defaults.compile.CompileConfig + compile_dynamic: true + compiled_region: language + coordinate_descent_tuning: false + max_autotune_pointwise: false + use_cuda_graphs: false + enabled: true + diffusion_expert_config: + _type: cosmos3._src.vfm.configs.base.defaults.model_config.DiffusionExpertConfig + base_fps: 24 + enable_fps_modulation: true + load_weights_from_pretrained: false + max_vae_latent_side_after_patchify: 20 + patch_spatial: 2 + timestep_range: 1.0 + unified_3d_mrope_reset_spatial_ids: true + unified_3d_mrope_temporal_modality_margin: 15000 + ema: + _type: cosmos3._src.vfm.configs.base.defaults.ema.EMAConfig + enabled: true + iteration_shift: 0 + rate: 0.1 + fixed_step_sampler_config: null + input_caption_key: ai_caption + input_image_key: images + input_video_key: video + joint_attn_implementation: two_way + latent_downsample_factor: 16 + lbl: + _type: cosmos3._src.vfm.configs.base.defaults.model_config.LBLConfig + coeff_gen: null + coeff_und: null + method: local + log_enc_time_every_n: 100 + lora_alpha: 32 + lora_enabled: false + lora_rank: 16 + lora_target_modules: q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen + max_action_dim: 64 + max_num_tokens_after_packing: 74000 + natten_parameter_list: null + net: null + num_embodiment_domains: 32 + parallelism: + _type: cosmos3._src.vfm.configs.base.defaults.parallelism.ParallelismConfig + cfg_parallel_shard_degree: 1 + context_parallel_shard_degree: 1 + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: 8 + enable_inference_mode: false + fsdp_master_dtype: float32 + precision: bfloat16 + rectified_flow_inference_config: + _type: cosmos3._src.vfm.configs.base.defaults.model_config.RectifiedFlowInferenceConfig + num_train_timesteps: 1000 + scheduler_type: unipc + shift: 1 + use_dynamic_shifting: false + rectified_flow_training_config: + _type: cosmos3._src.vfm.configs.base.defaults.model_config.RectifiedFlowTrainingConfig + action_loss_weight: 10.0 + high_sigma_ratio: 0.05 + high_sigma_timesteps_max: 1000 + high_sigma_timesteps_min: 995 + image_loss_scale: null + independent_action_schedule: false + independent_sound_schedule: false + loss_scale: 10.0 + normalize_loss_by_active: false + shift: + '256': 3 + '480': 5 + '720': 10 + shift_action: null + shift_sound: null + sound_loss_scale: 2.0 + train_time_action_distribution: logitnormal + train_time_image_distribution: logitnormal + train_time_sound_distribution: logitnormal + train_time_video_distribution: waver + train_time_weight: uniform + use_discrete_rf: false + use_dynamic_shift: false + use_high_sigma_strategy: false + use_high_sigma_strategy_action: false + use_high_sigma_strategy_sound: false + resolution: '480' + sound_dim: null + sound_gen: false + sound_latent_fps: 25 + sound_tokenizer: null + state_ch: 48 + state_t: 300 + tokenizer: + _target_: cosmos3._src.vfm.tokenizers.wan2pt2_vae_4x16x16.Wan2pt2VAEInterface + bucket_name: bucket + chunk_duration: 93 + encode_bucket_multiple: null + encode_chunk_frames: + '256': 68 + '480': 24 + '720': 12 + encode_exact_durations: + - 17 + - 61 + - 73 + keep_decoder_cache: false + object_store_credential_path_pretrained: credentials/gcp_training.secret + spatial_compression_factor: 16 + temporal_compression_factor: 4 + temporal_window: null + use_streaming_encode: false + vae_path: pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth + video_temporal_causal: false + vision_gen: true + vlm_config: + _type: cosmos3._src.vfm.configs.base.defaults.vlm.VLMConfig + layer_module: null + model_instance: + _target_: cosmos3._src.vfm.models.mot.unified_mot.Nemotron3DenseVLTextForCausalLM + config: + _target_: cosmos3._src.vfm.configs.base.defaults.vlm.create_vlm_config + base_config: + _target_: cosmos3._src.vfm.models.mot.unified_mot.Nemotron3DenseVLMoTConfig.from_json_file + json_file: cosmos3/_src/vfm/models/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json + include_visual: null + qk_norm_for_text: false + use_und_k_norm_for_gen: true + model_name: nvidia/Cosmos3-Edge-Reasoner + pretrained_weights: + _type: cosmos3._src.vfm.configs.base.defaults.vlm.PretrainedWeightsConfig + backbone_path: s3://bucket/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/ + checkpoint_format: null + credentials_path: credentials/gcp_checkpoint.secret + enable_gcs_patch_in_boto3: true + enabled: false + qk_norm: false + tie_word_embeddings: false + tokenizer: + _target_: cosmos3._src.vfm.processors.build_processor_lazy + repository: nvidia/Cosmos3-Edge + revision: main + use_system_prompt: false diff --git a/cosmos_framework/inference/defaults/reasoner/sample_args.json b/cosmos_framework/inference/defaults/reasoner/sample_args.json index e7a25adf..60724530 100644 --- a/cosmos_framework/inference/defaults/reasoner/sample_args.json +++ b/cosmos_framework/inference/defaults/reasoner/sample_args.json @@ -1,6 +1,6 @@ { "model_mode": "reasoner", - "max_new_tokens": 64, + "max_new_tokens": 1024, "do_sample": false, "temperature": 1.0, "top_k": null, diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index 30bd04b5..e5e8706e 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -4,7 +4,7 @@ import hashlib import json import pickle -from collections.abc import Callable, Generator, Iterable +from collections.abc import Callable, Generator, Iterable, MutableMapping from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, Sequence, TypeVar, cast, override @@ -40,7 +40,7 @@ SampleOutputs, SetupArgs, ) -from cosmos_framework.inference.common.inference import Inference, sync_distributed_errors +from cosmos_framework.inference.common.inference import Inference, _download_on_rank0, sync_distributed_errors from cosmos_framework.inference.common.init import get_rank, get_world_size from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel from cosmos_framework.inference.vision import ( @@ -57,6 +57,7 @@ from cosmos_framework.model.generator.upsampler.prompts import is_upsampled_prompt from cosmos_framework.tools.visualize.video import save_img_or_video from cosmos_framework.utils import log +from cosmos_framework.utils.checkpoint_db import CheckpointDirHf if TYPE_CHECKING: from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig @@ -66,6 +67,110 @@ _BatchItem = TypeVar("_BatchItem") +_PROCESSOR_HF_INCLUDE = ( + "*.json", + "*.jinja", + "merges.txt", + "vocab.json", +) + + +def _checkpoint_has_processor_files(checkpoint_path: Path) -> bool: + """True when the checkpoint dir bundles VLM processor/tokenizer files at its root.""" + return any((checkpoint_path / name).is_file() for name in ("tokenizer_config.json", "preprocessor_config.json")) + + +# Files ``Qwen2Tokenizer.from_pretrained`` reads from a local directory: its +# ``vocab_files_names`` (vocab.json + merges.txt, both opened unconditionally in +# ``__init__``) plus tokenizer_config.json, which carries the added special +# tokens (<|vision_start|> etc.) the SFT data pipeline depends on. +_QWEN2_TOKENIZER_FILES = ("tokenizer_config.json", "vocab.json", "merges.txt") + +# ``tokenizer_type`` substrings that ``build_processor`` maps to Qwen3VLProcessor +# in hub mode — the same processor its local-directory mode defaults to, so a +# rewrite to a bundled dir preserves the dispatch for these nodes. +_QWEN3VL_TOKENIZER_TYPES = ("Qwen/Qwen3-VL", "Siglip2-Qwen3-1.7B") + +# Repos served by the native Cosmos3-Edge (Nemotron bridge) processor; covers +# nvidia/Cosmos3-Edge and nvidia/Cosmos3-Edge-Reasoner (mirrors the substring +# dispatch in ``build_processor``). +_EDGE_REPO_SUBSTRING = "nvidia/Cosmos3-Edge" + + +def _tokenizer_node_target(tokenizer_cfg: MutableMapping[str, Any]) -> str: + """Last component of the node's ``_target_`` (e.g. ``build_processor_lazy``).""" + return str(tokenizer_cfg.get("_target_", "")).rsplit(".", 1)[-1] + + +def _dir_serves_qwen3vl_autoprocessor(path: Path) -> bool: + """True when ``path`` holds the files Qwen3VLProcessor's AutoProcessor load needs. + + That is the image/video preprocessor params, the tokenizer config, and the + tokenizer data (fast ``tokenizer.json`` or the slow vocab/merges pair). + """ + if not all((path / name).is_file() for name in ("preprocessor_config.json", "tokenizer_config.json")): + return False + return (path / "tokenizer.json").is_file() or all((path / name).is_file() for name in ("vocab.json", "merges.txt")) + + +def _bundled_processor_serves_node(tokenizer_cfg: MutableMapping[str, Any], checkpoint_path: Path) -> bool: + """True when checkpoint-bundled processor files can serve this ``vlm_config.tokenizer`` node. + + Gates the local-first override (``_point_tokenizer_node_at_dir``) on the + bundle being dispatchable for the node's shape: Edge-family nodes need the + native-snapshot markers, Qwen3-VL-family nodes the AutoProcessor file set, + and Qwen2 nodes take a local dir only with ``config_variant="hf"``. Unknown + shapes return False (keep the configured hub behavior). + """ + if not isinstance(tokenizer_cfg, MutableMapping): + return False + target = _tokenizer_node_target(tokenizer_cfg) + if target == "build_processor_lazy": + from cosmos_framework.data.generator.processors.cosmos3_edge_processing import is_cosmos3_edge_native_snapshot + + repository = tokenizer_cfg.get("repository") + name = repository or tokenizer_cfg.get("tokenizer_type") + if not isinstance(name, str) or not name: + return False + if _EDGE_REPO_SUBSTRING in name: + return is_cosmos3_edge_native_snapshot(str(checkpoint_path)) + if repository or any(qwen_type in name for qwen_type in _QWEN3VL_TOKENIZER_TYPES): + # Local-artifact (repository) nodes and Qwen3-VL tokenizer types both + # dispatch to Qwen3VLProcessor in dir mode; an edge-native bundle + # would instead route to the Nemotron bridge, so reject it. + return _dir_serves_qwen3vl_autoprocessor(checkpoint_path) and not is_cosmos3_edge_native_snapshot( + str(checkpoint_path) + ) + return False + if target == "create_qwen2_tokenizer_with_download": + if tokenizer_cfg.get("config_variant") != "hf" or not tokenizer_cfg.get("pretrained_model_name"): + return False + return all((checkpoint_path / name).is_file() for name in _QWEN2_TOKENIZER_FILES) + return False + + +def _point_tokenizer_node_at_dir(tokenizer_cfg: MutableMapping[str, Any], processor_path: Path | str) -> bool: + """Rewrite ``vlm_config.tokenizer`` in place to source its processor from a local dir. + + Dispatches on the node's shape so the mutated kwargs match the target's + signature (``build_processor_lazy`` takes ``tokenizer_type``; the Qwen2 node + takes ``pretrained_model_name``). Returns False — node left completely + untouched — for any other shape, so callers keep the configured hub behavior. + """ + if not isinstance(tokenizer_cfg, MutableMapping): + return False + target = _tokenizer_node_target(tokenizer_cfg) + if target == "build_processor_lazy" and (tokenizer_cfg.get("repository") or tokenizer_cfg.get("tokenizer_type")): + tokenizer_cfg.pop("repository", None) + tokenizer_cfg.pop("revision", None) + tokenizer_cfg.pop("subdir", None) + tokenizer_cfg["tokenizer_type"] = str(processor_path) + return True + if target == "create_qwen2_tokenizer_with_download" and tokenizer_cfg.get("config_variant") == "hf": + tokenizer_cfg["pretrained_model_name"] = str(processor_path) + return True + return False + def _iter_packed_batches( items: Iterable[_BatchItem], @@ -198,7 +303,6 @@ def _compute_num_tokens_for_sample(sample_args: OmniSampleArgs, model: OmniMoTMo latent_t = 1 + (T - 1) // vae_temporal_downsample num_vision_tokens = latent_h * latent_w * latent_t - # small compared to vision tokens, so we can ignore them for now. return num_vision_tokens @@ -564,7 +668,6 @@ def get_sample_data( device=device, ) - if sample_args.model_mode == ModelMode.IMAGE2IMAGE: return _get_image_edit_sample_data(sample_args, model, device=device) @@ -1102,22 +1205,61 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: ) model = cast("OmniMoTModel", model) Cosmos3OmniModel.after_load_model(model) + # Thread a local checkpoint dir to the reasoner LM (consumed by + # Edge's lazy ``_ensure_vision_tower``) so a checkpoint-bundled + # ``vision_encoder/`` is preferred over a hub download. s3 DCP + # loads have no local dir to offer — skip them. + if "://" not in str(setup_args.checkpoint_path) and Path(setup_args.checkpoint_path).is_dir(): + language_model = getattr(getattr(model, "net", None), "language_model", None) + if language_model is not None: + language_model._local_checkpoint_dir = str(setup_args.checkpoint_path) save_config(config, setup_args.output_dir) else: - checkpoint_path = setup_args.download_checkpoint() + checkpoint_path = _download_on_rank0(setup_args.download_checkpoint) if setup_args.config_file_type == ConfigFileType.MODULE: config = None else: model_dict = setup_args.load_model_config_dict() + tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"] + has_bundled_processor = _checkpoint_has_processor_files(checkpoint_path) if setup_args.vlm_processor_from_checkpoint: # Source the VLM processor from the loaded checkpoint's own # bundled files instead of the repository hardcoded in the # model config. Drops the redundant base-model download. - tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"] - tokenizer_cfg.pop("repository", None) - tokenizer_cfg.pop("revision", None) - tokenizer_cfg.pop("subdir", None) - tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) + processor_path = checkpoint_path + elif has_bundled_processor and _bundled_processor_serves_node(tokenizer_cfg, checkpoint_path): + # Local-first: exported checkpoints bundle processor/tokenizer + # files at their root (export_model); prefer them over the + # repository hardcoded in the model config so exported dirs + # stay offline-capable. Gated on the bundle actually being + # dispatchable for this tokenizer node. + log.info(f"VLM processor: using checkpoint-bundled processor files at {checkpoint_path}") + processor_path = checkpoint_path + else: + if has_bundled_processor: + log.info( + f"VLM processor: checkpoint-bundled files at {checkpoint_path} cannot serve this " + "tokenizer config; falling back to the configured source" + ) + if repository := tokenizer_cfg.get("repository"): + revision = tokenizer_cfg.get("revision") + if revision is None: + raise ValueError("VLM processor 'revision' is required when 'repository' is set") + processor_path = _download_on_rank0( + CheckpointDirHf( + repository=repository, + revision=revision, + subdirectory=tokenizer_cfg.get("subdir", ""), + include=_PROCESSOR_HF_INCLUDE, + ).download + ) + else: + processor_path = None + if processor_path is not None and not _point_tokenizer_node_at_dir(tokenizer_cfg, processor_path): + log.info( + "VLM processor: unrecognized 'vlm_config.tokenizer' node shape; leaving the configured " + "tokenizer source untouched" + ) # AVAE source: the configured ``avae_path`` when set, else the loaded # checkpoint's bundled ``sound_tokenizer/``. The inference-only # ``from_checkpoint`` key (default False) forces bundled; pop it so it @@ -1438,7 +1580,6 @@ def decode_vision(vision_latent: torch.Tensor) -> torch.Tensor: seed = [sa.seed if sa.seed is not None else _fallback_seed(cast(OmniSampleArgs, sa)) for sa in sample_args_list] outputs: dict[str, Any] | None = None - if outputs is None: assert all(sa.num_outputs == 1 for sa in sample_args_list), "num_outputs must be 1" n_sample = sum(cast(OmniSampleArgs, sa).num_outputs for sa in sample_args_list) @@ -1517,9 +1658,7 @@ def decode_vision(vision_latent: torch.Tensor) -> torch.Tensor: "(one must divide the other). Non-nesting CP/CFGP overlays " "with divergent per-sample num_steps are unsupported." ) - _steps_t = torch.tensor( - [local_num_steps], device=self.model.tensor_kwargs["device"], dtype=torch.int32 - ) + _steps_t = torch.tensor([local_num_steps], device=self.model.tensor_kwargs["device"], dtype=torch.int32) torch.distributed.all_reduce( _steps_t, op=torch.distributed.ReduceOp.MAX, group=parallel_dims.dp_shard_mesh.get_group() ) @@ -1734,9 +1873,7 @@ def _generate_reasoner_batch( n_img = sum(img is not None for img in raw_images) n_vid = sum(v is not None for v in (raw_videos or [])) if n_img and n_vid: - raise ValueError( - "Reasoner batch mixes image- and video-conditioned samples. Split into separate batches." - ) + raise ValueError("Reasoner batch mixes image- and video-conditioned samples. Split into separate batches.") if 0 < n_img < len(raw_images): raise ValueError( "Reasoner batch mixes image-conditioned and text-only samples " diff --git a/cosmos_framework/inference/inference_test.py b/cosmos_framework/inference/inference_test.py index da10591e..cd4d02ff 100644 --- a/cosmos_framework/inference/inference_test.py +++ b/cosmos_framework/inference/inference_test.py @@ -246,7 +246,7 @@ def test_reasoner_defaults_json_round_trip() -> None: defaults = _load_modality_defaults("reasoner") assert defaults["model_mode"] == "reasoner" - assert defaults["max_new_tokens"] == 64 + assert defaults["max_new_tokens"] == 1024 on_disk = _json.loads((PACKAGE_DIR / "defaults/reasoner/sample_args.json").read_text()) assert defaults == on_disk diff --git a/cosmos_framework/inference/local_first_test.py b/cosmos_framework/inference/local_first_test.py new file mode 100644 index 00000000..ef707ced --- /dev/null +++ b/cosmos_framework/inference/local_first_test.py @@ -0,0 +1,381 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Local-first Edge vision tower / processor plumbing and curated vision errors.""" + +import copy +import json +import types +from pathlib import Path + +import pytest + +from cosmos_framework.inference.args import ( + ModelMode, + OmniSampleOverrides, + _reasoner_vision_capable, +) +from cosmos_framework.inference.inference import ( + _bundled_processor_serves_node, + _checkpoint_has_processor_files, + _point_tokenizer_node_at_dir, +) +from cosmos_framework.inference.model import _raise_on_missing_vision_keys +from cosmos_framework.model.generator.mot.unified_mot import _load_bundled_vision_tower + +_VISION_SPEC = { + "vision_config": {"hidden_size": 16}, + "projector_config": { + "spatial_merge_size": 2, + "input_hidden_size": 16, + "merger_intermediate_size": 32, + "out_hidden_size": 8, + }, +} +_TOKEN_IDS = {"image_token_id": 1, "video_token_id": 2, "vision_start_token_id": 3} + + +def _write_bundle(root: Path, *, standalone: dict | None, top_level: dict | None) -> None: + ve_dir = root / "vision_encoder" + ve_dir.mkdir(parents=True) + (ve_dir / "model.safetensors").write_bytes(b"stub") + if standalone is not None: + (ve_dir / "config.json").write_text(json.dumps(standalone), encoding="utf-8") + if top_level is not None: + (root / "config.json").write_text(json.dumps(top_level), encoding="utf-8") + + +def test_load_bundled_vision_tower_absent(tmp_path: Path): + assert _load_bundled_vision_tower(None) is None + assert _load_bundled_vision_tower(str(tmp_path)) is None # no vision_encoder/ + + +def test_load_bundled_vision_tower_exported_bundle(tmp_path: Path): + # export_model --vit layout: self-describing standalone config (spec + token ids). + _write_bundle(tmp_path, standalone={**_VISION_SPEC, **_TOKEN_IDS}, top_level={"model": {}}) + + bundled = _load_bundled_vision_tower(str(tmp_path)) + assert bundled is not None + weights_file, ve_cfg, top_cfg = bundled + assert weights_file == str(tmp_path / "vision_encoder" / "model.safetensors") + assert ve_cfg["projector_config"]["out_hidden_size"] == 8 + # Self-describing: token ids come from the bundle, never the root config.json. + assert top_cfg is ve_cfg + assert top_cfg["image_token_id"] == 1 + + +def test_load_bundled_vision_tower_hub_snapshot(tmp_path: Path): + # Raw hub snapshot layout: no standalone file; spec + ids folded into the top-level config. + _write_bundle(tmp_path, standalone=None, top_level={**_VISION_SPEC, **_TOKEN_IDS}) + + bundled = _load_bundled_vision_tower(str(tmp_path)) + assert bundled is not None + _, ve_cfg, top_cfg = bundled + assert ve_cfg["vision_config"] == _VISION_SPEC["vision_config"] + assert top_cfg["video_token_id"] == 2 + + +def test_load_bundled_vision_tower_standalone_without_token_ids(tmp_path: Path): + # Older standalone layout: spec-only vision_encoder/config.json; ids live top-level. + _write_bundle(tmp_path, standalone=_VISION_SPEC, top_level={**_VISION_SPEC, **_TOKEN_IDS}) + + bundled = _load_bundled_vision_tower(str(tmp_path)) + assert bundled is not None + _, ve_cfg, top_cfg = bundled + assert "image_token_id" not in ve_cfg + assert top_cfg["vision_start_token_id"] == 3 + + +def _vlm_model_config(*, target: str = "", model_name: str = "", include_visual=None, model_instance="auto"): + if model_instance == "auto": + model_instance = {"_target_": target, "config": {"include_visual": include_visual}} + return types.SimpleNamespace(vlm_config=types.SimpleNamespace(model_name=model_name, model_instance=model_instance)) + + +def test_reasoner_vision_capable_edge_lazy_tower(): + # Edge: include_visual=None BY DESIGN; SigLIP2 tower loads lazily -> capable. + config = _vlm_model_config( + target="cosmos3._src.vfm.models.mot.unified_mot.Nemotron3DenseVLTextForCausalLM", + model_name="nvidia/Cosmos3-Edge-Reasoner", + include_visual=None, + ) + assert _reasoner_vision_capable(config) + + +def test_reasoner_vision_capable_qwen_with_visual(): + config = _vlm_model_config( + target="cosmos3._src.vfm.models.mot.unified_mot.Qwen3VLTextForCausalLM", + model_name="nvidia/Cosmos3-Nano-Reasoner", + include_visual=True, + ) + assert _reasoner_vision_capable(config) + + +def test_reasoner_vision_capable_qwen_without_visual(): + config = _vlm_model_config( + target="cosmos3._src.vfm.models.mot.unified_mot.Qwen3VLTextForCausalLM", + model_name="nvidia/Cosmos3-Nano-Reasoner", + include_visual=False, + ) + assert not _reasoner_vision_capable(config) + + +def test_reasoner_vision_capable_qwen_llm_without_model_instance(): + config = _vlm_model_config(model_name="Qwen/Qwen3-0.6B", model_instance=None) + assert not _reasoner_vision_capable(config) + + +def test_reasoner_vision_capable_unknown_family_never_blocks(): + config = _vlm_model_config(target="some.other.Backbone", model_name="acme/Custom-1B") + assert _reasoner_vision_capable(config) + + +def test_build_reasoner_data_fails_fast_for_visionless_qwen(tmp_path: Path): + img = tmp_path / "cat.jpg" + img.write_bytes(b"\xff\xd8\xff\xe0") # minimal non-empty file; not actually decoded here + overrides = OmniSampleOverrides(name="cat", prompt="describe", vision_path=str(img)) + sample_meta = types.SimpleNamespace(model_mode=ModelMode.REASONER) + + visionless = _vlm_model_config( + target="cosmos3._src.vfm.models.mot.unified_mot.Qwen3VLTextForCausalLM", + model_name="nvidia/Cosmos3-Nano-Reasoner", + include_visual=False, + ) + with pytest.raises(ValueError, match="visual tower"): + overrides._build_reasoner_data(model_config=visionless, sample_meta=sample_meta) + + # Text-only reasoner prompts remain valid on vision-less checkpoints. + text_only = OmniSampleOverrides(name="txt", prompt="hello") + text_only._build_reasoner_data(model_config=visionless, sample_meta=sample_meta) + + # Edge reasoner vision samples must not be blocked (lazy tower). + edge = _vlm_model_config( + target="cosmos3._src.vfm.models.mot.unified_mot.Nemotron3DenseVLTextForCausalLM", + model_name="nvidia/Cosmos3-Edge-Reasoner", + include_visual=None, + ) + overrides._build_reasoner_data(model_config=edge, sample_meta=sample_meta) + + +def _write_root_index(root: Path, keys: list[str]) -> None: + (root / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": {key: "model-00001-of-00001.safetensors" for key in keys}}), + encoding="utf-8", + ) + + +def test_raise_on_missing_vision_keys_no_vit_export(tmp_path: Path): + _write_root_index(tmp_path, ["model.net.vae2llm.weight"]) + state_dict = { + "model.net.vae2llm.weight": None, + "model.net.language_model.visual.blocks.0.attn.qkv.weight": None, + } + with pytest.raises(ValueError, match="--no-vit"): + _raise_on_missing_vision_keys(tmp_path, state_dict) + + +def test_raise_on_missing_vision_keys_tolerates_complete_and_visionless(tmp_path: Path): + visual_key = "model.net.language_model.visual.blocks.0.attn.qkv.weight" + + # Checkpoint provides vision keys -> fine. + _write_root_index(tmp_path, ["model.net.vae2llm.weight", visual_key]) + _raise_on_missing_vision_keys(tmp_path, {visual_key: None}) + + # Model built without a visual tower (include_visual=false) -> fine. + _raise_on_missing_vision_keys(tmp_path, {"model.net.vae2llm.weight": None}) + + # No root index (e.g. DCP dir) -> defer to the loader. + _raise_on_missing_vision_keys(tmp_path / "missing", {visual_key: None}) + + +def test_checkpoint_has_processor_files(tmp_path: Path): + assert not _checkpoint_has_processor_files(tmp_path) + (tmp_path / "tokenizer_config.json").write_text("{}", encoding="utf-8") + assert _checkpoint_has_processor_files(tmp_path) + + other = tmp_path / "other" + other.mkdir() + (other / "preprocessor_config.json").write_text("{}", encoding="utf-8") + assert _checkpoint_has_processor_files(other) + + +def test_vlm_model_size_covers_edge(): + assert OmniSampleOverrides._VLM_MODEL_SIZE["nvidia/Cosmos3-Edge-Reasoner"] == "2B" + # Edge training shift (Cosmos3-Edge.yaml): 256->3, 480->5, 720->10. + assert OmniSampleOverrides._RESOLUTION_SHIFT_DEFAULTS[("2B", "480")] == 5.0 + + +# --------------------------------------------------------------------------- +# Checkpoint-local processor override: node-shape dispatch + dispatchability gate +# --------------------------------------------------------------------------- + +_LAZY_TARGET = "cosmos_framework.data.generator.processors.build_processor_lazy" +_LEGACY_TARGET = "cosmos_framework.configs.base.defaults.reasoner.create_qwen2_tokenizer_with_download" + + +def _edge_lazy_node() -> dict: + # Edge SFT export shape (edge_model_config.py / Cosmos3-Edge.yaml). + return {"_target_": _LAZY_TARGET, "repository": "nvidia/Cosmos3-Edge", "revision": "main"} + + +def _legacy_qwen2_node(config_variant: str = "hf") -> dict: + # Nano/Super SFT export shape (nano_model_config.py / super_model_config.py). + return { + "_target_": _LEGACY_TARGET, + "pretrained_model_name": "Qwen/Qwen3-VL-8B-Instruct", + "config_variant": config_variant, + } + + +def _write_json(root: Path, name: str, payload: dict | None = None) -> None: + (root / name).write_text(json.dumps(payload or {}), encoding="utf-8") + + +def _write_edge_native_bundle(root: Path) -> None: + # Renewed (5bb63b97-style) layout as bundled by export_model: the export + # root config.json says cosmos3_omni, so processor_config.json alone must + # carry the native-snapshot marker. + _write_json(root, "config.json", {"model_type": "cosmos3_omni"}) + _write_json(root, "processor_config.json", {"processor_class": "Cosmos3EdgeProcessor"}) + _write_json(root, "tokenizer_config.json") + _write_json(root, "preprocessor_config.json") + + +def _write_pre_renewal_edge_bundle(root: Path) -> None: + # 28a0b8e-style layout: no processor_config.json; preprocessor_config.json + # auto_maps to a processing.py that is never bundled -> AutoProcessor crash. + _write_json(root, "config.json", {"model_type": "cosmos3_omni"}) + _write_json(root, "tokenizer_config.json") + _write_json(root, "preprocessor_config.json", {"auto_map": {"AutoProcessor": "processing.NemotronNanoV3Processor"}}) + + +def _write_qwen_autoprocessor_bundle(root: Path) -> None: + _write_json(root, "preprocessor_config.json") + _write_json(root, "tokenizer_config.json") + _write_json(root, "tokenizer.json") + + +def _write_qwen2_slow_tokenizer_bundle(root: Path) -> None: + # Minimal but real file set for Qwen2Tokenizer.from_pretrained (vocab_files_names + # + tokenizer_config.json). + _write_json(root, "vocab.json", {"a": 0, "b": 1, "ab": 2}) + (root / "merges.txt").write_text("#version: 0.2\na b\n", encoding="utf-8") + _write_json(root, "tokenizer_config.json", {"model_max_length": 131072}) + + +def test_edge_lazy_node_rewritten_for_native_bundle(tmp_path: Path): + _write_edge_native_bundle(tmp_path) + node = _edge_lazy_node() + + assert _checkpoint_has_processor_files(tmp_path) + assert _bundled_processor_serves_node(node, tmp_path) + assert _point_tokenizer_node_at_dir(node, tmp_path) + assert node == {"_target_": _LAZY_TARGET, "tokenizer_type": str(tmp_path)} + + +def test_edge_lazy_node_gate_rejects_pre_renewal_bundle(tmp_path: Path): + _write_pre_renewal_edge_bundle(tmp_path) + node = _edge_lazy_node() + before = copy.deepcopy(node) + + # Passes the coarse file check but is NOT dispatchable for the Edge node + # (build_processor's dir mode needs the native-snapshot markers). + assert _checkpoint_has_processor_files(tmp_path) + assert not _bundled_processor_serves_node(node, tmp_path) + assert node == before # the gate never mutates + + +def test_legacy_qwen2_node_rewritten_to_local_dir(tmp_path: Path): + _write_qwen2_slow_tokenizer_bundle(tmp_path) + node = _legacy_qwen2_node() + + assert _bundled_processor_serves_node(node, tmp_path) + assert _point_tokenizer_node_at_dir(node, tmp_path) + # Only pretrained_model_name is rewritten; no tokenizer_type is injected and + # config_variant="hf" is kept so download_tokenizer_files passes the local + # dir through verbatim. + assert node == { + "_target_": _LEGACY_TARGET, + "pretrained_model_name": str(tmp_path), + "config_variant": "hf", + } + + # The rewritten kwargs must match create_qwen2_tokenizer_with_download's + # exact (pretrained_model_name, config_variant) signature and load the + # bundled files from the local dir (the original defect was a TypeError: + # unexpected keyword argument 'tokenizer_type'). + from cosmos_framework.configs.base.defaults.reasoner import create_qwen2_tokenizer_with_download + from cosmos_framework.data.generator.processors import LLMTokenizerProcessor + + kwargs = {key: value for key, value in node.items() if key != "_target_"} + processor = create_qwen2_tokenizer_with_download(**kwargs) + assert isinstance(processor, LLMTokenizerProcessor) + assert processor.tokenizer.get_vocab()["ab"] == 2 + + +def test_legacy_qwen2_node_non_hf_variant_untouched(tmp_path: Path): + _write_qwen2_slow_tokenizer_bundle(tmp_path) + node = _legacy_qwen2_node(config_variant="gcp") + before = copy.deepcopy(node) + + # Non-"hf" variants route pretrained_model_name into the object-store + # download branch, which cannot take a local dir -> old (hub) behavior. + assert not _bundled_processor_serves_node(node, tmp_path) + assert not _point_tokenizer_node_at_dir(node, tmp_path) + assert node == before + + +def test_legacy_qwen2_gate_requires_slow_tokenizer_files(tmp_path: Path): + # tokenizer_config.json alone passes the coarse file check but + # Qwen2Tokenizer.from_pretrained also needs vocab.json + merges.txt. + _write_json(tmp_path, "tokenizer_config.json") + assert _checkpoint_has_processor_files(tmp_path) + assert not _bundled_processor_serves_node(_legacy_qwen2_node(), tmp_path) + + +def test_llm_backbone_lazy_node_untouched(tmp_path: Path): + # Qwen3-0.6B LLM backbone: dir-mode build_processor would mis-dispatch to + # Qwen3VLProcessor instead of LLMTokenizerProcessor -> never use the bundle. + _write_qwen_autoprocessor_bundle(tmp_path) + node = {"_target_": _LAZY_TARGET, "tokenizer_type": "Qwen/Qwen3-0.6B", "config_variant": "hf"} + assert not _bundled_processor_serves_node(node, tmp_path) + + +def test_qwen3vl_tokenizer_type_lazy_node_gate(tmp_path: Path): + node = {"_target_": _LAZY_TARGET, "tokenizer_type": "Qwen/Qwen3-VL-8B-Instruct", "config_variant": "hf"} + + # Missing preprocessor_config.json -> AutoProcessor cannot build the image + # processor -> gate fails. + _write_json(tmp_path, "tokenizer_config.json") + _write_json(tmp_path, "tokenizer.json") + assert not _bundled_processor_serves_node(node, tmp_path) + + _write_json(tmp_path, "preprocessor_config.json") + assert _bundled_processor_serves_node(node, tmp_path) + assert _point_tokenizer_node_at_dir(node, tmp_path) + assert node == {"_target_": _LAZY_TARGET, "config_variant": "hf", "tokenizer_type": str(tmp_path)} + + +def test_repository_lazy_node_qwen_family_gate(tmp_path: Path): + node = {"_target_": _LAZY_TARGET, "repository": "nvidia/Cosmos3-Nano", "revision": "main"} + _write_qwen_autoprocessor_bundle(tmp_path) + assert _bundled_processor_serves_node(node, tmp_path) + + # An edge-native bundle would re-route dir-mode dispatch to the Nemotron + # bridge -- wrong for a Qwen-family node. + _write_json(tmp_path, "processor_config.json", {"processor_class": "Cosmos3EdgeProcessor"}) + assert not _bundled_processor_serves_node(node, tmp_path) + + +def test_unknown_tokenizer_node_untouched(tmp_path: Path): + _write_edge_native_bundle(tmp_path) + node = {"_target_": "acme.processors.build_custom_processor", "repository": "acme/Custom", "revision": "main"} + before = copy.deepcopy(node) + + assert not _bundled_processor_serves_node(node, tmp_path) + assert not _point_tokenizer_node_at_dir(node, tmp_path) + assert node == before + + # Fail-safe on malformed nodes too. + assert not _bundled_processor_serves_node(None, tmp_path) + assert not _point_tokenizer_node_at_dir(None, tmp_path) diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index 042bab5b..5b7c74ce 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -192,9 +192,29 @@ def _read_safetensors_index(index_path: Path) -> dict[str, str]: def _diffusers_weight_map(checkpoint_path: Path) -> dict[str, str]: index_path = checkpoint_path / _DIFFUSERS_ROOT_INDEX - if not index_path.exists(): - raise FileNotFoundError(f"Diffusers safetensors index not found: {index_path}") - return _read_safetensors_index(index_path) + weight_map = _read_safetensors_index(index_path) if index_path.exists() else {} + # The generator-only K norm (`k_norm_und_for_gen`) belongs to the transformer + # component and must be sourced from its own index under the diffusers name + # (`layers.N...`). The converter is supposed to exclude it from the root index + # (see _convert_model_to_diffusers._remap_language_model_state_dict), but + # checkpoints exported before that fix (e.g. nvidia/Cosmos3-Edge-Policy-DROID) + # leak it into the root index under the raw model-internal name + # (`layers.layers.N...`), which points at tensors absent from the shards and + # breaks loading. Drop any root-index k-norm entries and let the transformer + # index below supply the correctly named ones. + weight_map = {k: v for k, v in weight_map.items() if ".k_norm_und_for_gen." not in k} + # The root index may be a reasoner manifest (e.g. Cosmos3-Edge) that lists only the + # shared understanding-pathway weights and the vision tower. The transformer + # component's own index holds the complete DiT state — including the + # generation-pathway tensors (`*_moe_gen`, `add_*`) absent from such a root index. + # Merge it (root entries win) so the full generator can be loaded. + transformer_index_path = checkpoint_path / "transformer" / "diffusion_pytorch_model.safetensors.index.json" + if transformer_index_path.exists(): + for key, rel_path in _read_safetensors_index(transformer_index_path).items(): + weight_map.setdefault(key, f"transformer/{rel_path}") + if not weight_map: + raise FileNotFoundError(f"Diffusers safetensors index not found at {index_path} or {transformer_index_path}") + return weight_map def _diffusers_files_to_keys(weight_map: dict[str, str]) -> dict[str, list[str]]: @@ -208,13 +228,44 @@ def _diffusers_files_to_keys(weight_map: dict[str, str]) -> dict[str, list[str]] def _is_diffusers_checkpoint(checkpoint_path: Path) -> bool: index_path = checkpoint_path / _DIFFUSERS_ROOT_INDEX + transformer_index_path = checkpoint_path / "transformer" / "diffusion_pytorch_model.safetensors.index.json" + if (checkpoint_path / _DIFFUSERS_MODEL_INDEX).exists() and (index_path.exists() or transformer_index_path.exists()): + return True if not index_path.exists(): return False - if (checkpoint_path / _DIFFUSERS_MODEL_INDEX).exists(): - return True return any(_is_diffusers_model_weight_path(path) for path in _read_safetensors_index(index_path).values()) +_LM_VISUAL_KEY_MARKER = "language_model.visual." + + +def _raise_on_missing_vision_keys(checkpoint_path: Path, state_dict: dict[str, Any]) -> None: + """Curated error for root-safetensors checkpoints that lack the reasoner vision tower. + + Mirrors ``_DiffusersLoadPlanner``'s tolerant/curated missing-keys handling for + the plain ``HuggingFaceStorageReader`` path: a model built with a visual tower + (``include_visual=True``) cannot load from a checkpoint exported with + ``--no-vit``; without this pre-check, ``dcp.load`` dies mid-load with a raw + ``Missing key in checkpoint state_dict: model.net.language_model.visual...``. + """ + wanted = sorted(key for key in state_dict if _LM_VISUAL_KEY_MARKER in key) + if not wanted: + return + # Root index shares its filename with the diffusers layout ("model.safetensors.index.json"). + index_path = checkpoint_path / _DIFFUSERS_ROOT_INDEX + if not index_path.exists(): + return + if any(_LM_VISUAL_KEY_MARKER in key for key in _read_safetensors_index(index_path)): + return + raise ValueError( + f"Checkpoint at {checkpoint_path} provides none of the {len(wanted)} reasoner vision-tower " + f"tensor(s) the model config requires (first: {wanted[0]!r}). It was likely exported with " + "--no-vit while its config still enables the visual tower. Fix: re-export with the default " + "--vit (or use a full checkpoint); for generation-only use, re-export with --no-vit — new " + "--no-vit exports write include_visual=false so the model loads without the tower." + ) + + def _normalize_diffusers_target_key(name: str) -> str: return name.removeprefix("model.net.").replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "") @@ -355,10 +406,17 @@ def _normalize_target_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]: def _build_remapped_state_dict(self, target_state_dict: dict[str, Any]) -> tuple[dict[str, Any], set[str]]: remapped_state_dict: dict[str, Any] = {} loaded_keys: set[str] = set() + # When the model is built without a visual tower (e.g. Cosmos3-Edge t2i with + # include_visual disabled), its state dict has no `language_model.visual.*` + # targets, so the checkpoint's vision_encoder weights have nowhere to go — skip + # them rather than erroring on unmapped projector/visual keys. + has_visual_target = any(key.startswith("language_model.visual.") for key in target_state_dict) for diff_key, rel_path in sorted(self.weight_map.items()): net_key = _diffusers_to_net_key(diff_key, rel_path) if net_key is None: if _is_diffusers_model_weight_path(rel_path): + if rel_path.startswith("vision_encoder/") and not has_visual_target: + continue raise KeyError(f"Diffusers model key {diff_key!r} from {rel_path!r} has no Cosmos3 mapping.") continue target_tensor = target_state_dict.get(net_key) @@ -473,6 +531,13 @@ def from_pretrained_dcp( config.compile = attrs.asdict(compile_config) config.quantization = attrs.asdict(quantization_config) model = cls(config) + # Thread the local checkpoint dir to the reasoner LM (consumed by Edge's + # lazy ``_ensure_vision_tower``): checkpoints that bundle a + # ``vision_encoder/`` next to the weights then load it without any hub + # access. Harmless when the dir has no bundle (hub fallback applies). + language_model = getattr(getattr(model.model, "net", None), "language_model", None) + if language_model is not None: + language_model._local_checkpoint_dir = str(checkpoint_path) checkpoint_type = CheckpointType.from_path(checkpoint_path) match checkpoint_type: case CheckpointType.DCP: @@ -481,13 +546,21 @@ def from_pretrained_dcp( case CheckpointType.HF: if _is_diffusers_checkpoint(checkpoint_path): state_dict = get_model_state_dict(model.model.net) + # Single-rank load must skip DCP's collective path: it gathers + # (pickles) the load plan across ranks, and _DiffusersLoadPlanner + # carries code objects that cannot be pickled ("cannot pickle + # code objects"). no_dist loads locally without collectives. + _dist = torch.distributed + no_dist = not (_dist.is_available() and _dist.is_initialized() and _dist.get_world_size() > 1) dcp.load( state_dict=state_dict, storage_reader=_DiffusersHuggingFaceStorageReader(checkpoint_path), planner=_DiffusersLoadPlanner(checkpoint_path), + no_dist=no_dist, ) return model state_dict = get_model_state_dict(model) + _raise_on_missing_vision_keys(checkpoint_path, state_dict) storage_reader = HuggingFaceStorageReader(str(checkpoint_path)) case _: assert_never(checkpoint_type) diff --git a/cosmos_framework/inference/model_test.py b/cosmos_framework/inference/model_test.py index 9021097b..ca2c8c9c 100644 --- a/cosmos_framework/inference/model_test.py +++ b/cosmos_framework/inference/model_test.py @@ -50,6 +50,30 @@ def test_checkpoint_type_from_path_hf_index(tmp_path: Path): assert CheckpointType.from_path(tmp_path) == CheckpointType.HF +def test_diffusers_weight_map_transformer_only(tmp_path: Path) -> None: + transformer_path = tmp_path / "transformer" + transformer_path.mkdir() + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + (transformer_path / "config.json").write_text("{}", encoding="utf-8") + (transformer_path / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": { + "proj_in.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + }, + } + ), + encoding="utf-8", + ) + + assert CheckpointType.from_path(tmp_path) == CheckpointType.HF + assert _is_diffusers_checkpoint(tmp_path) + assert _diffusers_weight_map(tmp_path) == { + "proj_in.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + } + + def test_normalize_diffusers_target_key(): assert ( _normalize_diffusers_target_key( diff --git a/cosmos_framework/model/_base.py b/cosmos_framework/model/_base.py index edcc0dba..ed9e74dc 100644 --- a/cosmos_framework/model/_base.py +++ b/cosmos_framework/model/_base.py @@ -94,6 +94,18 @@ def on_train_start(self, memory_format: torch.memory_format = torch.preserve_for """ pass + def on_before_optimizer_step( + self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int + ) -> None: + """Hook before optimizer.step() is called. + + Args: + optimizer (torch.optim.Optimizer): The model optimizer. + scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler. + iteration (int): Current iteration number. + """ + pass + def on_before_zero_grad( self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int ) -> None: diff --git a/cosmos_framework/model/generator/distillation/__init__.py b/cosmos_framework/model/generator/distillation/__init__.py new file mode 100644 index 00000000..51f5095f --- /dev/null +++ b/cosmos_framework/model/generator/distillation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Distribution-matching training methods.""" + +__all__: tuple[str, ...] = () diff --git a/cosmos_framework/model/generator/distillation/common_loss.py b/cosmos_framework/model/generator/distillation/common_loss.py new file mode 100644 index 00000000..399a06ef --- /dev/null +++ b/cosmos_framework/model/generator/distillation/common_loss.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Loss functions for distribution matching distillation. + +Adapted from fastgen/methods/common_loss.py for the Cosmos3 codebase. +Reference: https://github.com/NVlabs/FastGen/blob/main/fastgen/methods/common_loss.py + +All loss functions accept per-sample lists of tensors (variable shapes across +samples) to stay consistent with the MoT pipeline convention. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +__all__: tuple[str, ...] = ( + "VSDLossReduction", + "variational_score_distillation_loss", + "variational_score_distillation_loss_from_gradient", +) + +VSDLossReduction = str +_VSD_LOSS_REDUCTIONS: set[str] = {"mean", "sum", "sum_rcm"} + + +def variational_score_distillation_loss( + gen_data: list[torch.Tensor], + teacher_x0: list[torch.Tensor], + fake_score_x0: list[torch.Tensor], + loss_mask: list[torch.Tensor] | None = None, + additional_scale: torch.Tensor | None = None, + reduction: VSDLossReduction = "mean", +) -> torch.Tensor: + """Compute the variational score distillation (VSD) loss. + + The VSD gradient is ``(fake_score_x0 - teacher_x0) / weight``, where the + weight normalises the per-sample gradient magnitude. With ``reduction="mean"``, + the gradient is divided by the active element count. With ``reduction="sum"``, + the legacy half-MSE sum keeps the gradient independent of the active element + count. With ``reduction="sum_rcm"``, the pseudo-target loss follows RCM's + no-half-factor sum and normalizer clamp. Teacher and fake-score are treated + as constants in all modes. + + All inputs are per-sample lists to support variable shapes across samples. + + Args: + gen_data: Student-generated data per sample (carries gradient). + teacher_x0: x0 predictions from the teacher per sample (detached). + fake_score_x0: x0 predictions from the fake-score per sample (detached). + loss_mask: Optional per-sample masks, each ``(T_i, 1, 1)`` where + ``1`` = generated frame (include in loss/weight), ``0`` = conditioned + frame (exclude). When provided, both the weight and MSE loss are + computed only over generated frames to avoid inflation from + near-zero conditioned differences. + additional_scale: Optional per-sample scaling tensor of shape ``(B,)``. + reduction: ``"mean"`` preserves the existing active-element mean loss. + ``"sum"`` preserves the existing half-MSE per-sample sum. + ``"sum_rcm"`` uses the RCM-style no-half-factor per-sample sum. + + Returns: + Scalar VSD loss averaged over samples. + """ + with torch.no_grad(): + vsd_grad = [fake_score_x0[i] - teacher_x0[i] for i in range(len(gen_data))] # list of [C,T_i,H_i,W_i] + return variational_score_distillation_loss_from_gradient( + gen_data=gen_data, + vsd_grad=vsd_grad, + weight_reference=teacher_x0, + loss_mask=loss_mask, + additional_scale=additional_scale, + reduction=reduction, + ) + + +def variational_score_distillation_loss_from_gradient( + gen_data: list[torch.Tensor], + vsd_grad: list[torch.Tensor], + weight_reference: list[torch.Tensor], + loss_mask: list[torch.Tensor] | None = None, + additional_scale: torch.Tensor | None = None, + reduction: VSDLossReduction = "mean", +) -> torch.Tensor: + """Compute VSD pseudo-target loss from a caller-provided raw gradient. + + ``vsd_grad`` is the unnormalised gradient direction. This helper applies the + same per-sample adaptive normalisation as :func:`variational_score_distillation_loss`, + using ``weight_reference`` as the teacher-distance reference, and constructs a + detached pseudo-target. With ``reduction="sum"``, ``d loss / d gen_data`` + is ``vsd_grad * weight``; with ``reduction="sum_rcm"``, it is + ``2 * vsd_grad * weight`` to match RCM's no-half-factor pseudo-target loss. + With ``reduction="mean"``, it is additionally divided by the active element + count. + + Args: + gen_data: Student-generated data per sample (carries gradient). + vsd_grad: Raw VSD gradient per sample, before adaptive normalisation. + weight_reference: Reference tensors used to compute the adaptive weight. + loss_mask: Optional per-sample masks, each ``(T_i, 1, 1)`` where + ``1`` = generated frame (include in loss/weight), ``0`` = conditioned + frame (exclude). + additional_scale: Optional per-sample scaling tensor of shape ``(B,)``. + reduction: ``"mean"`` divides each per-sample pseudo-target loss by + active element count. ``"sum"`` preserves the existing half-MSE sum. + ``"sum_rcm"`` matches the RCM DMD pseudo-target reduction. + + Returns: + Scalar VSD loss averaged over samples. + """ + per_sample_losses = [] + if reduction not in _VSD_LOSS_REDUCTIONS: + raise ValueError(f"Unknown VSD loss reduction: {reduction}") + for i in range(len(gen_data)): + g_i = gen_data[i] # [C,T_i,H_i,W_i] + grad_i = vsd_grad[i] # [C,T_i,H_i,W_i] + ref_i = weight_reference[i] # [C,T_i,H_i,W_i] + + with torch.no_grad(): + # Weight calculation in fp32 for numerical stability + g_fp32 = g_i.float() # [C,T_i,H_i,W_i] + ref_fp32 = ref_i.float() # [C,T_i,H_i,W_i] + + # Mean absolute difference over generated frames only (avoid inflation + # from conditioned frames where gen ≈ teacher ≈ clean data) + if loss_mask is not None: + mask_i = loss_mask[i].expand_as(g_fp32) # [C,T_i,H_i,W_i] + diff_abs_mean = ((g_fp32 - ref_fp32).abs() * mask_i).sum() / mask_i.sum().clamp(min=1) # scalar + else: + diff_abs_mean = (g_fp32 - ref_fp32).abs().mean() # scalar + if reduction == "sum_rcm": + w = 1.0 / diff_abs_mean.clamp(min=1e-5) # scalar + else: + w = 1.0 / (diff_abs_mean + 1e-6) # scalar + + if additional_scale is not None: + w = w * additional_scale[i].float() # scalar + + w = w.to(dtype=g_i.dtype) # scalar + + weighted_grad = grad_i * w # [C,T_i,H_i,W_i] + if reduction == "sum_rcm": + has_nonfinite_weighted_grad = ~torch.isfinite(weighted_grad).all() # scalar + weighted_grad = torch.where( + has_nonfinite_weighted_grad, + torch.zeros_like(weighted_grad), + weighted_grad, + ) # [C,T_i,H_i,W_i] + pseudo_target = g_i - weighted_grad # [C,T_i,H_i,W_i] + + if loss_mask is not None: + mask_i = loss_mask[i].expand_as(g_i) # [C,T_i,H_i,W_i] + loss_sq = (g_i - pseudo_target) ** 2 * mask_i # [C,T_i,H_i,W_i] + loss_sum = loss_sq.sum() # scalar + if reduction == "mean": + loss_i = 0.5 * loss_sum / mask_i.sum().clamp(min=1) # scalar + elif reduction == "sum_rcm": + loss_i = loss_sum # scalar + else: + loss_i = 0.5 * loss_sum # scalar + else: + if reduction == "mean": + loss_i = 0.5 * F.mse_loss(g_i, pseudo_target, reduction="mean") # scalar + elif reduction == "sum_rcm": + loss_sq = (g_i - pseudo_target) ** 2 # [C,T_i,H_i,W_i] + loss_sum = loss_sq.sum() # scalar + loss_i = loss_sum # scalar + else: + loss_i = 0.5 * F.mse_loss(g_i, pseudo_target, reduction="sum") # scalar + per_sample_losses.append(loss_i) + + return torch.stack(per_sample_losses).mean() # scalar diff --git a/cosmos_framework/model/generator/distillation/common_loss_test.py b/cosmos_framework/model/generator/distillation/common_loss_test.py new file mode 100644 index 00000000..360edea6 --- /dev/null +++ b/cosmos_framework/model/generator/distillation/common_loss_test.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from __future__ import annotations + +import pytest +import torch + +from cosmos_framework.model.generator.distillation import common_loss as common_loss_module +from cosmos_framework.model.generator.distillation.common_loss import ( + variational_score_distillation_loss, + variational_score_distillation_loss_from_gradient, +) + +# --------------------------------------------------------------------------- +# Tests for variational_score_distillation_loss +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_vsd_loss_is_zero_when_fake_equals_teacher(): + """When fake_score_x0 == teacher_x0 the VSD gradient is zero, so loss == 0.""" + rng = torch.Generator() + rng.manual_seed(0) + gen = [torch.randn(2, 4, 2, 2, generator=rng)] + teacher = [torch.randn(2, 4, 2, 2, generator=rng)] + fake = [teacher[0].clone()] # identical to teacher + + loss = variational_score_distillation_loss(gen, teacher, fake) + assert loss.item() == pytest.approx(0.0, abs=1e-6) + + +@pytest.mark.L0 +def test_vsd_loss_positive_when_fake_differs_from_teacher(): + """When fake_score_x0 != teacher_x0 the VSD gradient is non-zero, loss > 0.""" + rng = torch.Generator() + rng.manual_seed(1) + gen = [torch.randn(2, 4, 2, 2, generator=rng)] + teacher = [torch.randn(2, 4, 2, 2, generator=rng)] + fake = [torch.randn(2, 4, 2, 2, generator=rng)] # different from teacher + + loss = variational_score_distillation_loss(gen, teacher, fake) + assert loss.item() > 0.0 + + +@pytest.mark.L0 +def test_vsd_additional_scale_applied(): + """additional_scale per sample should scale the VSD gradient accordingly.""" + rng = torch.Generator() + rng.manual_seed(2) + gen = [torch.randn(2, 4, 2, 2, generator=rng).requires_grad_(True)] + teacher = [torch.randn(2, 4, 2, 2, generator=rng)] + fake = [torch.randn(2, 4, 2, 2, generator=rng)] + + # With scale=1.0 (no scaling) + loss_no_scale = variational_score_distillation_loss(gen, teacher, fake) + + # With scale=2.0 — pseudo_target shifts proportionally → loss changes + additional_scale = torch.tensor([2.0]) + loss_scaled = variational_score_distillation_loss(gen, teacher, fake, additional_scale=additional_scale) + + # Scaled loss should differ from unscaled loss (scale affects pseudo_target) + assert loss_scaled.item() != pytest.approx(loss_no_scale.item(), abs=1e-6) + + +@pytest.mark.L0 +def test_vsd_multi_sample_averaging(): + """Multi-sample batch: loss is the mean over per-sample losses.""" + rng = torch.Generator() + rng.manual_seed(3) + + # Build two samples with different shapes + gen0 = torch.randn(2, 2, 2, 2, generator=rng) + teacher0 = torch.randn(2, 2, 2, 2, generator=rng) + fake0 = torch.randn(2, 2, 2, 2, generator=rng) + + gen1 = torch.randn(2, 4, 2, 2, generator=rng) + teacher1 = torch.randn(2, 4, 2, 2, generator=rng) + fake1 = torch.randn(2, 4, 2, 2, generator=rng) + + # Compute joint and individual losses + loss_joint = variational_score_distillation_loss([gen0, gen1], [teacher0, teacher1], [fake0, fake1]) + loss0 = variational_score_distillation_loss([gen0], [teacher0], [fake0]) + loss1 = variational_score_distillation_loss([gen1], [teacher1], [fake1]) + + expected = (loss0 + loss1) / 2.0 + assert loss_joint.item() == pytest.approx(expected.item(), rel=1e-5) + + +@pytest.mark.L0 +def test_vsd_loss_from_gradient_uses_caller_gradient(): + """The pseudo-target loss should expose the caller-provided gradient on gen_data.""" + gen = [torch.ones(1, 1, 1, 1, requires_grad=True)] # [C,T,H,W] + raw_grad = [torch.full((1, 1, 1, 1), 2.0)] # [C,T,H,W] + weight_reference = [torch.zeros(1, 1, 1, 1)] # [C,T,H,W] + + loss = variational_score_distillation_loss_from_gradient(gen, raw_grad, weight_reference) + loss.backward() + + expected_grad = torch.full((1, 1, 1, 1), 2.0 / (1.0 + 1e-6)) # [C,T,H,W] + assert torch.allclose(gen[0].grad, expected_grad) + + +@pytest.mark.L0 +def test_vsd_sum_reduction_avoids_active_element_averaging() -> None: + """Sum reduction should avoid active-element averaging of the DMD gradient.""" + raw_grad = [torch.full((1, 1, 1, 2), 2.0)] # [C,T,H,W] + weight_reference = [torch.zeros(1, 1, 1, 2)] # [C,T,H,W] + + gen_mean = [torch.ones(1, 1, 1, 2, requires_grad=True)] # [C,T,H,W] + loss_mean = variational_score_distillation_loss_from_gradient( + gen_mean, raw_grad, weight_reference, reduction="mean" + ) + loss_mean.backward() + + gen_sum = [torch.ones(1, 1, 1, 2, requires_grad=True)] # [C,T,H,W] + loss_sum = variational_score_distillation_loss_from_gradient(gen_sum, raw_grad, weight_reference, reduction="sum") + loss_sum.backward() + + expected_sum_grad = torch.full((1, 1, 1, 2), 2.0 / (1.0 + 1e-6)) # [C,T,H,W] + expected_mean_grad = expected_sum_grad / gen_mean[0].numel() # [C,T,H,W] + assert torch.allclose(gen_mean[0].grad, expected_mean_grad) + assert torch.allclose(gen_sum[0].grad, expected_sum_grad) + + +@pytest.mark.L0 +def test_vsd_sum_rcm_reduction_matches_rcm_no_half_factor() -> None: + """sum_rcm should use RCM's no-half-factor pseudo-target reduction.""" + raw_grad = [torch.full((1, 1, 1, 2), 2.0)] # [C,T,H,W] + weight_reference = [torch.zeros(1, 1, 1, 2)] # [C,T,H,W] + + gen_sum_rcm = [torch.ones(1, 1, 1, 2, requires_grad=True)] # [C,T,H,W] + loss_sum_rcm = variational_score_distillation_loss_from_gradient( + gen_sum_rcm, raw_grad, weight_reference, reduction="sum_rcm" + ) + loss_sum_rcm.backward() + + expected_rcm_grad = torch.full((1, 1, 1, 2), 4.0) # [C,T,H,W] + assert torch.allclose(gen_sum_rcm[0].grad, expected_rcm_grad) + + +@pytest.mark.L0 +@pytest.mark.parametrize("use_loss_mask", [False, True]) +def test_vsd_sum_rcm_zeroes_nan_sample_loss_and_grad(use_loss_mask: bool) -> None: + """sum_rcm should zero loss and grad when a sample's pseudo-target gradient is NaN.""" + gen_sum_rcm = [torch.ones(1, 1, 1, 1, requires_grad=True)] # [C,T,H,W] + raw_grad = [torch.full((1, 1, 1, 1), float("nan"))] # [C,T,H,W] + weight_reference = [torch.zeros(1, 1, 1, 1)] # [C,T,H,W] + loss_mask = [torch.ones(1, 1, 1)] if use_loss_mask else None # [T,H,W] + + loss_sum_rcm = variational_score_distillation_loss_from_gradient( + gen_sum_rcm, + raw_grad, + weight_reference, + loss_mask=loss_mask, + reduction="sum_rcm", + ) + loss_sum_rcm.backward() + + assert loss_sum_rcm.item() == pytest.approx(0.0) + assert torch.allclose(gen_sum_rcm[0].grad, torch.zeros_like(gen_sum_rcm[0])) # [C,T,H,W] + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_loss_exports_are_explicit() -> None: + assert common_loss_module.__all__ == ( + "VSDLossReduction", + "variational_score_distillation_loss", + "variational_score_distillation_loss_from_gradient", + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_vsd_loss_preserves_public_dtype_contract(dtype: torch.dtype) -> None: + gen_data = [torch.ones(1, 2, 1, 2, dtype=dtype, requires_grad=True)] # [C,T,H,W] + vsd_grad = [torch.full((1, 2, 1, 2), 2.0, dtype=dtype)] # [C,T,H,W] + weight_reference = [torch.zeros(1, 2, 1, 2, dtype=dtype)] # [C,T,H,W] + + loss = variational_score_distillation_loss_from_gradient(gen_data, vsd_grad, weight_reference) + loss.backward() + + assert loss.dtype is dtype + assert gen_data[0].grad is not None + assert gen_data[0].grad.dtype is dtype + assert torch.isfinite(loss) + assert torch.isfinite(gen_data[0].grad).all() + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_vsd_list_inputs_honor_mask_and_mean_sum_reductions() -> None: + vsd_grad = [torch.full((1, 2, 1, 2), 2.0), torch.ones(1, 1, 1, 1)] # list of [C,T_i,H_i,W_i] + weight_reference = [torch.zeros(1, 2, 1, 2), torch.zeros(1, 1, 1, 1)] # list of [C,T_i,H_i,W_i] + loss_mask = [torch.tensor([[[1.0]], [[0.0]]]), torch.ones(1, 1, 1)] # list of [T_i,H_i,W_i] + + gen_mean = [ + torch.ones(1, 2, 1, 2, requires_grad=True), + torch.ones(1, 1, 1, 1, requires_grad=True), + ] # list of [C,T_i,H_i,W_i] + mean_loss = variational_score_distillation_loss_from_gradient( + gen_mean, + vsd_grad, + weight_reference, + loss_mask=loss_mask, + reduction="mean", + ) + mean_loss.backward() + + gen_sum = [ + torch.ones(1, 2, 1, 2, requires_grad=True), + torch.ones(1, 1, 1, 1, requires_grad=True), + ] # list of [C,T_i,H_i,W_i] + sum_loss = variational_score_distillation_loss_from_gradient( + gen_sum, + vsd_grad, + weight_reference, + loss_mask=loss_mask, + reduction="sum", + ) + sum_loss.backward() + + assert gen_mean[0].grad is not None + assert gen_sum[0].grad is not None + assert torch.count_nonzero(gen_mean[0].grad[:, 1:]) == 0 + assert torch.count_nonzero(gen_sum[0].grad[:, 1:]) == 0 + assert torch.allclose(gen_sum[0].grad[:, :1], 2.0 * gen_mean[0].grad[:, :1]) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_vsd_gradient_space_matches_fake_minus_teacher_direction() -> None: + gen_data = [torch.ones(1, 1, 1, 2, requires_grad=True)] # [C,T,H,W] + teacher_x0 = [torch.zeros(1, 1, 1, 2)] # [C,T,H,W] + fake_score_x0 = [torch.full((1, 1, 1, 2), 2.0)] # [C,T,H,W] + + loss = variational_score_distillation_loss( + gen_data, + teacher_x0, + fake_score_x0, + reduction="sum", + ) + loss.backward() + + expected_grad = torch.full((1, 1, 1, 2), 2.0 / (1.0 + 1e-6)) # [C,T,H,W] + assert torch.allclose(gen_data[0].grad, expected_grad) diff --git a/cosmos_framework/model/generator/distillation/dmd2_rf.py b/cosmos_framework/model/generator/distillation/dmd2_rf.py new file mode 100644 index 00000000..3edc472a --- /dev/null +++ b/cosmos_framework/model/generator/distillation/dmd2_rf.py @@ -0,0 +1,1800 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""DMD2 model using Rectified Flow parameterization for Cosmos3. + +Supports two student simulation modes (``simulation_mode`` config field): + +- ``"forward"`` (default): fastgen-style DMD2. The student generates x0 via a single forward pass + from a randomly sampled noise level, and the VSD / fake-score losses are computed on re-noised + student outputs. + +- ``"backward"``: rcm-style multi-step rollout. The student always denoises from a pure-noise start + (``t_list[0]`` must be 1.0) toward 0, over an iteration-cycled number of ``t_list`` steps (schedule + prefix). No clean data is blended into the start state. Gradients flow through the last + ``backward_grad_steps`` student forward passes (1 = last step only, −1 = full BPTT through all steps). + +Reference: https://arxiv.org/abs/2405.14867 +""" + +from __future__ import annotations + +import collections +import time +from typing import Any, Mapping, Optional, cast + +import torch +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner +from torch.distributed.checkpoint.filesystem import FileSystemReader # noqa: F401 - used after OSS release transform +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict, set_model_state_dict +from torch.nn.modules.module import _IncompatibleKeys +from torch.nn.utils.clip_grad import clip_grad_norm_ + +from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader # noqa: F401 - used after OSS release transform +from cosmos_framework.utils.lazy_config import LazyDict +from cosmos_framework.utils.lazy_config import instantiate as lazy_instantiate +from cosmos_framework.utils import log, misc +from cosmos_framework.utils.misc import get_local_tensor_if_DTensor +from cosmos_framework.model.generator.diffusion.samplers.fixed_step import FixedStepSampler +from cosmos_framework.model.generator.mot.context_parallel_utils import context_parallel_broadcast_tensor_list +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import tokenize_caption +from cosmos_framework.model.generator.utils.data_and_condition import ( + GenerationDataClean, + GenerationDataNoised, + build_dense_sound_schedule, +) +from cosmos_framework.data.generator.sequence_packing import ( + PackedSequence, + SequencePlan, + build_sequence_plans_from_data_batch, +) +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution +from cosmos_framework.configs.base.experiment.distillation.dmd2_config import DMD2OptimizerConfig, DMD2RFConfig +from cosmos_framework.model.generator.distillation.common_loss import ( + variational_score_distillation_loss, + variational_score_distillation_loss_from_gradient, +) +from cosmos_framework.model.generator.distillation.optimizer import ( + OptimizerModelView, + PhaseOptimizer, + PhaseScheduler, + iter_torch_optimizers, +) + +__all__: tuple[str, ...] = ("DMD2RFModel",) + + +class DMD2RFModel(OmniMoTModel): + """ + DMD2 distillation model using Rectified Flow parameterization. + + https://arxiv.org/abs/2405.14867 + """ + + # ------------------------ Initialization & configuration ------------------------ + + def __init__(self, config: DMD2RFConfig) -> None: + """ + Args: + config (DMD2RFConfig): The configuration for the DMD model + """ + # Keep training-only networks out of the nn.Module registry. This matches + # self-forcing DMD and keeps inference state dicts student-only. + self._net_teacher_holder: list[torch.nn.Module] = [] + self._net_fake_score_holder: list[torch.nn.Module] = [] + config.vlm_config.pretrained_weights.enabled = False + super().__init__(config) + self.config: DMD2RFConfig = config + assert config.noise_level_parameterization.lower() == "rectified_flow", ( + "DMD2RF requires rectified_flow parameterization." + ) + + @property + def net_teacher(self) -> torch.nn.Module: + """Frozen teacher kept out of the module registry / inference state dict.""" + if not self._net_teacher_holder: + raise AttributeError("DMD2RF teacher has not been initialized.") + return self._net_teacher_holder[0] + + @property + def net_fake_score(self) -> torch.nn.Module: + """Fake-score critic kept out of the registry; persisted by the distill checkpointer.""" + if not self._net_fake_score_holder: + raise AttributeError("DMD2RF fake-score network has not been initialized.") + return self._net_fake_score_holder[0] + + def _set_up_fixed_step_sampler(self) -> None: + """Build the fixed-step sampler used during distilled student inference.""" + sampler_cfg = self.config.fixed_step_sampler_config + if sampler_cfg is None: + self.fixed_step_sampler = None + return + self.fixed_step_sampler = FixedStepSampler( + t_list=list(sampler_cfg.t_list), + sample_type=sampler_cfg.sample_type, + num_train_timesteps=float(self.config.rectified_flow_inference_config.num_train_timesteps), + ) + + def _needs_fake_score(self) -> bool: + """Whether this distillation variant builds and trains a fake-score net.""" + return True + + @misc.timer("DMD2RFModel: set_up_model") + def set_up_model(self) -> None: + t0 = time.time() + # Build the student net (self.net) and shared components (tokenizer, etc.). + super().set_up_model() + + is_inference_mode = bool(getattr(self.config.parallelism, "enable_inference_mode", False)) + if is_inference_mode: + log.info("[DMD2RF] set_up_model: inference mode; skipping teacher and fake-score nets.") + self.net_discriminator_head = None + self.denoiser_nets = {"student": self.net} + self._set_up_fixed_step_sampler() + torch.cuda.empty_cache() + log.info(f"[DMD2RF] set_up_model: inference setup done in {time.time() - t0:.1f}s") + return + + load_teacher_weights = self.config.load_teacher_weights + + if load_teacher_weights: + assert self.config.teacher_load_from.load_path, ( + "A pretrained teacher model checkpoint is required for distillation" + ) + + needs_fake_score = self._needs_fake_score() + + # Build teacher and optional fake-score nets by temporarily swapping self.vlm_config, + # since build_net() reads architecture from that attribute. + saved_vlm_config = self.vlm_config + + self.config.vlm_config_teacher.pretrained_weights.enabled = False + self.vlm_config = self.config.vlm_config_teacher + self._net_teacher_holder = [self.build_net(self.precision, lora_enabled=False)] + + if needs_fake_score: + self.config.vlm_config_fake_score.pretrained_weights.enabled = False + self.vlm_config = self.config.vlm_config_fake_score + self._net_fake_score_holder = [self.build_net(self.precision)] + else: + self._net_fake_score_holder = [] + + self.vlm_config = saved_vlm_config # restore + if needs_fake_score: + assert self.net_fake_score is not None, "DMD2 requires a fake_score network." + + log.info("==========Instantiating networks...==========") + log.info(f"net: {self.net}") + log.info(f"net_teacher: {self.net_teacher}") + if needs_fake_score: + log.info(f"net_fake_score: {self.net_fake_score}") + + # Load the teacher checkpoint, then copy its weights into the student and + # optional fake-score nets as the warm-start for distillation. + if load_teacher_weights: + self._load_checkpoint_to_net( + self.net_teacher, + self.config.teacher_load_from.load_path, + credential_path=self.config.teacher_load_from.credentials, + ) + + self._copy_teacher_weights(target_net=self.net, target_name="student") + if needs_fake_score: + self._copy_teacher_weights(target_net=self.net_fake_score, target_name="fake score") + + if self.config.ema.enabled: + self.net_ema_worker.copy_to(src_model=self.net_teacher, tgt_model=self.net_ema) + + # Optionally override the student with a dedicated student checkpoint + # (e.g. resuming a distillation run from a mid-training snapshot). + if self.config.student_load_from and self.config.student_load_from.load_path: + self._load_checkpoint_to_net( + self.net, + self.config.student_load_from.load_path, + credential_path=self.config.student_load_from.credentials, + ) + if self.config.ema.enabled: + self.net_ema_worker.copy_to(src_model=self.net, tgt_model=self.net_ema) + + # Discriminator support not implemented yet for cosmos3. + self.net_discriminator_head = None + + # Freeze the teacher — it is used only for inference during training. + # Set eval() so dropout / norm running stats stay in inference mode even + # when the trainer recursively calls model.train() at the start of each step. + log.info("[DMD2RF] set_up_model: freezing teacher net ...") + self.net_teacher.eval().requires_grad_(False) + + # Register all denoiser nets in a dict for easy lookup by name in _pack_and_denoise. + self.denoiser_nets = { + "teacher": self.net_teacher, + "student": self.net, + } + if needs_fake_score: + self.denoiser_nets["fake_score"] = self.net_fake_score + + self._set_up_fixed_step_sampler() + + torch.cuda.empty_cache() + log.info(f"[DMD2RF] set_up_model: all done in {time.time() - t0:.1f}s") + + def _copy_teacher_weights(self, target_net: torch.nn.Module, target_name: str) -> None: + """Copy teacher weights into a target network, with logging on key mismatches. + + Args: + target_net: The target network to copy weights to. + target_name: The name of the target network. + """ + log.info(f"==========Copying teacher weights to {target_name} net==========") + to_load = {k: v for k, v in self.net_teacher.state_dict().items() if not k.endswith("_extra_state")} + key_match_status = target_net.load_state_dict(to_load, strict=False) + missing_all = [k for k in key_match_status.missing_keys if not k.endswith("_extra_state")] + missing = [k for k in missing_all if ".lora_" not in k] + missing_lora = [k for k in missing_all if ".lora_" in k] + unexpected = [k for k in key_match_status.unexpected_keys if not k.endswith("_extra_state")] + if missing or unexpected: + log.warning(f"==========teacher -> {target_name}: Missing: {missing[:10]}, Unexpected: {unexpected}") + elif missing_lora: + log.info( + f"==========teacher -> {target_name}: Base keys matched successfully; " + f"left {len(missing_lora)} LoRA adapter keys initialized." + ) + if not missing_all and not unexpected: + log.info(f"==========teacher -> {target_name}: All keys matched successfully.") + + def _load_checkpoint_to_net( + self, + net: torch.nn.Module, + ckpt_path: str, + prefix: str = "net_ema", + credential_path: str | None = None, + ) -> None: + """Load a DCP checkpoint into a single network.""" + + # Open the checkpoint storage and snapshot the model's current state dict. + # Infer the key prefix from the checkpoint path: ".dcp/model" layouts use "net", + # everything else (the default) uses "net_ema". + storage_reader = ( + S3StorageReader(credential_path=credential_path or "", path=ckpt_path) + if ckpt_path.startswith("s3://") + else FileSystemReader(ckpt_path) + ) + if ckpt_path.endswith(".dcp/model"): + prefix = "net" + _state_dict = get_model_state_dict(net) + + # Compare checkpoint keys against the (prefixed) model keys to surface + # missing / unexpected keys before attempting to load. + metadata = storage_reader.read_metadata() + checkpoint_keys = metadata.state_dict_metadata.keys() + + model_keys = set(_state_dict.keys()) + prefixed_model_keys = {f"{prefix}.{k}" for k in model_keys} + + missing_keys = prefixed_model_keys - checkpoint_keys + if missing_keys: + log.warning(f"Missing keys in checkpoint: {missing_keys}") + + # Filter out the complementary prefix ("net." when loading "net_ema." and vice versa) + # and _extra_state keys, which are TE metadata not present in all checkpoints. + unexpected_keys = checkpoint_keys - prefixed_model_keys + assert prefix in ["net", "net_ema"], "prefix must be either net or net_ema" + if prefix == "net_ema": + unexpected_keys = [k for k in unexpected_keys if "net." not in k] + else: + unexpected_keys = [k for k in unexpected_keys if "net_ema." not in k] + log.warning("Ignoring _extra_state keys..") + unexpected_keys = [k for k in unexpected_keys if "_extra_state" not in k] + if unexpected_keys: + log.warning(f"Unexpected keys in checkpoint: {unexpected_keys}") + + if not missing_keys and not unexpected_keys: + log.info("All keys matched successfully.") + + # DCP requires keys to be stored under their prefixed names. Build a + # prefixed view of the state dict, load into it, then strip the prefix back. + _new_state_dict = collections.OrderedDict() + for k in _state_dict.keys(): + _new_state_dict[f"{prefix}.{k}"] = _state_dict[k] + dcp.load(_new_state_dict, storage_reader=storage_reader, planner=DefaultLoadPlanner(allow_partial_load=True)) + for k in _state_dict.keys(): + _state_dict[k] = _new_state_dict[f"{prefix}.{k}"] + + # Apply the loaded weights to the network (non-strict to tolerate minor mismatches). + log.info(set_model_state_dict(net, _state_dict, options=StateDictOptions(strict=False))) + del _state_dict, _new_state_dict + + # ------------------------ Optimizers & schedulers ------------------------ + + def init_optimizer_scheduler( + self, optimizer_config: DMD2OptimizerConfig, scheduler_config: LazyDict + ) -> tuple[PhaseOptimizer, PhaseScheduler]: + """Create optimizers/schedulers for student (net) and critic (fake_score).""" + opt_net = lazy_instantiate(optimizer_config.net, model=OptimizerModelView(self.net)) + sched_net = lazy_instantiate(scheduler_config, optimizer=opt_net) + opt_fake = lazy_instantiate(optimizer_config.fake_score, model=OptimizerModelView(self.net_fake_score)) + sched_fake = lazy_instantiate(scheduler_config, optimizer=opt_fake) + return ( + PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}), + PhaseScheduler({"net": sched_net, "fake_score": sched_fake}), + ) + + def get_phase(self, iteration: int) -> str: + """Return the current training phase: ``"student"`` or ``"critic"``.""" + assert not (self.config.warmup_student_steps and self.config.warmup_critic_steps) + if iteration < self.config.warmup_critic_steps: + return "critic" + elif iteration < self.config.warmup_student_steps: + return "student" + else: + return "student" if iteration % self.config.student_update_freq == 0 else "critic" + + def get_optimizer_key(self, iteration: int) -> str: + return "net" if self.get_phase(iteration) == "student" else "fake_score" + + def get_student_iteration(self, iteration: int) -> int: + """Effective student iteration index used for EMA scheduling.""" + if iteration < self.config.warmup_student_steps: + return iteration + else: + steps_after_warmup = (iteration - self.config.warmup_student_steps) // self.config.student_update_freq + return self.config.warmup_student_steps + steps_after_warmup + + def get_critic_iteration(self, iteration: int) -> int: + """Effective critic (fake-score) update index (mirrors rcm ``get_effective_iteration_fake``).""" + return iteration - self.get_student_iteration(iteration) - 1 + + # ------------------------ training hooks ------------------------ + + def on_before_zero_grad(self, optimizer: PhaseOptimizer, scheduler: PhaseScheduler, iteration: int) -> None: + del scheduler + + if self.get_phase(iteration) == "critic": + # Critic phase: for BF16 optimizers that maintain FP32 master weights, copy the master params + # back into the model params before zeroing gradients. This keeps the low-precision model weights + # in sync with the authoritative FP32 copy. + for net, opt_key in ((self.net_fake_score, "fake_score"), (self.net_discriminator_head, "discriminator")): + opt = None if net is None else optimizer.get(opt_key) + if opt is None: + continue + params, master_params = [], [] + for inner_opt in iter_torch_optimizers(opt): + param_groups_master = getattr(inner_opt, "param_groups_master", None) + if not getattr(inner_opt, "master_weights", False) or param_groups_master is None: + continue + for group, group_master in zip(inner_opt.param_groups, param_groups_master): + for p, p_master in zip(group["params"], group_master["params"]): + params.append(get_local_tensor_if_DTensor(p.data)) + master_params.append(get_local_tensor_if_DTensor(p_master.data)) + if params: + torch._foreach_copy_(params, master_params) + elif self.get_phase(iteration) == "student": + # Student phase: update the EMA model using the current student weights. + if self.config.ema.enabled: + ema_beta = self.ema_beta(self.get_student_iteration(iteration)) + self.net_ema_worker.update_average(self.net, self.net_ema, beta=ema_beta) + + def zero_grad_for_phase(self, iteration: int) -> None: + if self.get_phase(iteration) == "student": + self.net.zero_grad(set_to_none=True) + else: + self.net_fake_score.zero_grad(set_to_none=True) + if self.net_discriminator_head is not None: + self.net_discriminator_head.zero_grad(set_to_none=True) + + # ------------------------ Helper methods and utils for callbacks ------------------------ + + def clip_grad_norm_( + self, + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: Optional[bool] = None, + ) -> torch.Tensor: + if not self.config.grad_clip: + max_norm = 1e12 + # Clip each network separately so their individual grad norms are bounded. + # Return the student norm as the canonical value logged by the trainer. + clip_grad_norm_( + self.net_fake_score.parameters(), + max_norm, + norm_type=norm_type, + error_if_nonfinite=error_if_nonfinite, + foreach=foreach, + ) + if self.net_discriminator_head is not None: + clip_grad_norm_( + self.net_discriminator_head.parameters(), + max_norm, + norm_type=norm_type, + error_if_nonfinite=error_if_nonfinite, + foreach=foreach, + ) + return clip_grad_norm_( + self.net.parameters(), + max_norm, + norm_type=norm_type, + error_if_nonfinite=error_if_nonfinite, + foreach=foreach, + ) + + @staticmethod + def _action_sample_indices(sequence_plans: list[SequencePlan]) -> list[int]: + """Original batch positions of samples that carry action data. + + ``packed_sequence.action.*`` and ``gen_data_*.x0_tokens_action`` are dense over + action-bearing samples. Mapping a dense action entry back to the original batch + position is required to align action-side sigmas with the full-batch sigma tensor. + """ + return [i for i, plan in enumerate(sequence_plans) if plan.has_action] + + def _action_sigmas_from_full( + self, + sigmas_full: torch.Tensor, + sequence_plans: list[SequencePlan], + ) -> torch.Tensor | None: + """Reindex full-batch sigmas to dense action positions. + + Returns ``None`` when action_gen is off or the batch contains no action samples, + so callers can pass the result straight to ``_add_noise_to_input(..., sigmas_action=...)`` + and fall back to legacy vision-σ behaviour automatically. + """ + if not self.config.action_gen: + return None + indices = self._action_sample_indices(sequence_plans) + if not indices: + return None + idx = torch.tensor(indices, dtype=torch.long, device=sigmas_full.device) + return sigmas_full[idx] # [n_action, *sigmas_full.shape[1:]] + + def _sound_sigmas_from_full( + self, + sigmas_full: torch.Tensor, + sequence_plans: list[SequencePlan], + x0_tokens_sound: list[torch.Tensor] | None, + ) -> torch.Tensor | None: + """Reindex full-batch sigmas to dense sound positions. + + Sound tokens are dense over samples with ``has_sound=True``. Passing a full-batch + sigma tensor directly to ``_add_noise_to_input`` misaligns mixed sound/no-sound + batches; this mirrors the base Omni helper path. + """ + if not self.config.sound_gen: + return None + _, sigmas_sound = build_dense_sound_schedule( + sequence_plans, x0_tokens_sound, sigmas_full, sigmas_full + ) # None or [n_sound,*sigmas_full.shape[1:]] + return sigmas_sound + + @staticmethod + def _has_any_noisy_action(packed_action: object) -> bool: + """Whether any action-bearing sample has at least one noisy (non-conditioning) frame. + + DMD2 action loss paths multiply by ``(1 - condition_mask)`` so all-conditioning + samples already contribute zero gradient, but the forward pass is still launched + and the adaptive VSD weight ``1/(|g - t| + ε)`` can spike to 1/ε on a fully-masked + batch. Skipping the loss explicitly avoids both the wasted compute and the numerical + edge case. + """ + if packed_action is None: + return False + condition_masks = getattr(packed_action, "condition_mask", None) + if condition_masks is None: + return False + for cm in condition_masks: + if (1.0 - cm).any(): + return True + return False + + @staticmethod + def _flow_matching_sum_rcm_loss( + pred: list[torch.Tensor], + target: list[torch.Tensor], + condition_mask: list[torch.Tensor], + has_valid_tokens: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute RCM-style fake-score FM loss: per-instance active-element sum.""" + if not has_valid_tokens: + dummy_loss = 0.0 * sum(p.sum() for p in pred) # scalar + return dummy_loss, dummy_loss.unsqueeze(0) # scalar, [1] + + per_instance_losses = [] + for i in range(len(pred)): + sqerr_i = (pred[i] - target[i]) ** 2 # [sample_shape] + noisy_mask_i = 1.0 - condition_mask[i] # [T_i,...] + loss_i = (sqerr_i * noisy_mask_i).sum() # scalar + per_instance_losses.append(loss_i) + + per_instance_loss = torch.stack(per_instance_losses) # [B] + return per_instance_loss.mean(), per_instance_loss # scalar, [B] + + def _get_vision_noise_sampling_metadata( + self, + data_batch: dict[str, Any], + gen_data_clean: GenerationDataClean, + ) -> tuple[list[str], list[int]]: + """Return per-sample resolution and token count for RF noise sampling.""" + data_resolutions: list[str] = [] + if "image_size" in data_batch: + for i in range(gen_data_clean.batch_size): + img_size = data_batch["image_size"][i] # [1,4] or [4] + if img_size.dim() == 2: + img_size = img_size[0] # [4] + target_h = int(img_size[0].item()) + target_w = int(img_size[1].item()) + data_resolutions.append(get_vision_data_resolution((target_h, target_w))) + else: + data_resolutions = [self.config.resolution] * gen_data_clean.batch_size + + num_tokens_per_sample = [x.shape[2] * x.shape[3] * x.shape[4] for x in gen_data_clean.x0_tokens_vision] + return data_resolutions, num_tokens_per_sample + + # ------------------------ Checkpointing helpers ------------------------ + + def model_dict(self) -> dict[str, torch.nn.Module]: + model_dict: dict[str, torch.nn.Module] = { + "net": self.net, + "fake_score": self.net_fake_score, + } + if self.net_discriminator_head: + model_dict["discriminator"] = self.net_discriminator_head + return model_dict + + def load_state_dict( + self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False + ) -> _IncompatibleKeys: + """Load weights for student/EMA (via base class) and for fake_score/discriminator heads. + + ``net_fake_score`` is held outside the module registry, so route its keys + to the fake-score net explicitly and keep them out of the base model load. + Strictness is enforced by the distillation checkpointer wrapper using the + returned missing/unexpected keys. + """ + strict_requested = strict + main_state_dict = collections.OrderedDict( + (key, value) + for key, value in state_dict.items() + if not key.startswith(("net_fake_score.", "net_discriminator_head.")) + ) + base_results: _IncompatibleKeys = super().load_state_dict(main_state_dict, strict=False, assign=assign) + + # Partition the flat state dict by prefix into per-head sub-dicts, + # stripping the prefix so each head's load_state_dict sees bare keys. + fake_score_state_dict = collections.OrderedDict() + discriminator_state_dict = collections.OrderedDict() + + for k, v in state_dict.items(): + if k.startswith("net_fake_score."): + fake_score_state_dict[k.replace("net_fake_score.", "")] = v + elif k.startswith("net_discriminator_head."): + discriminator_state_dict[k.replace("net_discriminator_head.", "")] = v + + # Accumulate missing/unexpected keys from all heads so the caller gets a + # unified report matching the contract of nn.Module.load_state_dict. + missing_keys: list[str] = list(base_results.missing_keys) + unexpected_keys: list[str] = list(base_results.unexpected_keys) + + fake_score_holder = getattr(self, "_net_fake_score_holder", []) + is_inference_mode = bool(getattr(self.config.parallelism, "enable_inference_mode", False)) + if fake_score_state_dict and fake_score_holder: + fake_score_results: _IncompatibleKeys = self.net_fake_score.load_state_dict( + fake_score_state_dict, strict=False, assign=assign + ) + missing_keys += [f"net_fake_score.{key}" for key in fake_score_results.missing_keys] + unexpected_keys += [f"net_fake_score.{key}" for key in fake_score_results.unexpected_keys] + elif fake_score_state_dict and is_inference_mode: + log.info("[DMD2RF] load_state_dict: ignoring net_fake_score.* keys in inference mode.") + elif fake_score_state_dict: + unexpected_keys += [f"net_fake_score.{key}" for key in fake_score_state_dict] + + if self.net_discriminator_head and discriminator_state_dict: + discriminator_results: _IncompatibleKeys = self.net_discriminator_head.load_state_dict( + discriminator_state_dict, strict=False, assign=assign + ) + missing_keys += [f"net_discriminator_head.{key}" for key in discriminator_results.missing_keys] + unexpected_keys += [f"net_discriminator_head.{key}" for key in discriminator_results.unexpected_keys] + + if strict_requested and (missing_keys or unexpected_keys): + log.warning( + "DMD2RFModel.load_state_dict received strict=True, but this loader returns incompatible keys " + "for the distillation checkpointer to validate instead of raising directly. " + f"Missing: {missing_keys[:20]}, Unexpected: {unexpected_keys[:20]}" + ) + + return _IncompatibleKeys(missing_keys=missing_keys, unexpected_keys=unexpected_keys) + + # ------------------------ Helpers for DMD2 training ------------------------ + + @staticmethod + def _velocity_to_x0( + x_t: list[torch.Tensor], + v_pred: list[torch.Tensor], + sigma: list[torch.Tensor], + ) -> list[torch.Tensor]: + """Convert velocity predictions to x0 (clean-data) predictions per sample. + + Under rectified-flow interpolation ``x_t = (1 - sigma) * x0 + sigma * eps`` and ``v = eps - x0``, + we have ``x0 = x_t - sigma * v``. + + All inputs are lists of per-sample tensors (variable shapes across samples). + + Args: + x_t: Noisy inputs, each ``(C, T_i, H_i, W_i)``. + v_pred: Predicted velocities, same shapes as *x_t*. + sigma: Per-sample noise levels, each ``(T_i, 1, 1)`` with zeros for conditioned frames. + Broadcasts with token shapes. + + Returns: + Per-sample x0 predictions with the same shapes as *x_t*. + """ + return [x_t[i] - sigma[i] * v_pred[i] for i in range(len(x_t))] + + def _sample_student_sigma( + self, + batch_size: int, + ) -> torch.Tensor: + """Sample a noise level for the student from the discrete sigma schedule. + + Mirrors ``sample_from_t_list`` in fastgen: randomly picks one of the predefined sigma values + that define the student's multi-step schedule. + + Returns: + Tensor of shape ``(batch_size, 1)`` with sigma values in ``[0, 1]``. + """ + t_list = self.config.fixed_step_sampler_config.t_list + t_tensor = torch.tensor(t_list, **self.tensor_kwargs_fp32) # [len(t_list)] + ids = torch.randint(0, len(t_list), (batch_size,), device=t_tensor.device) # [B] + return t_tensor[ids].unsqueeze(1) # [B,1] + + def _pack_and_denoise( + self, + gen_data_clean: GenerationDataClean, + gen_data_noised: GenerationDataNoised, + timesteps: torch.Tensor, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + net_type: str, + ) -> dict[str, torch.Tensor | object]: + """Pack a sequence and run a denoiser network. + + Args: + gen_data_clean: Original clean data (provides metadata: ``raw_state_*``, ``fps_*``, + ``is_image_batch``, ``num_vision_items_per_sample``). + gen_data_noised: Pre-computed noised data for all modalities (provides + ``xt_tokens_*`` for packing and token injection). + timesteps: RF timesteps ``(B, 1)`` in ``[0, num_train_timesteps]`` range. + input_text_indexes: Tokenized captions per sample. + sequence_plans: Conditioning structure per sample. + net_type: One of ``"student"``, ``"teacher"``, ``"fake_score"``. + + Returns: + Network output dict (e.g. ``{"preds_vision": [...]}``) . + """ + # Build a GenerationDataClean proxy using CPU versions of the noised tokens + # so that _pack_input_sequence can determine shapes and conditioning structure. + # The actual token values are overwritten below with the CUDA tensors. + gen_data_for_packing = GenerationDataClean( + batch_size=gen_data_clean.batch_size, + is_image_batch=gen_data_clean.is_image_batch, + raw_state_vision=gen_data_clean.raw_state_vision, + x0_tokens_vision=[vt.cpu() for vt in gen_data_noised.xt_tokens_vision], + fps_vision=gen_data_clean.fps_vision, + num_vision_items_per_sample=gen_data_clean.num_vision_items_per_sample, + raw_state_sound=gen_data_clean.raw_state_sound, + x0_tokens_sound=( + [st.cpu() for st in gen_data_noised.xt_tokens_sound] + if gen_data_noised.xt_tokens_sound is not None + else None + ), + fps_sound=gen_data_clean.fps_sound, + raw_state_action=gen_data_clean.raw_state_action, + x0_tokens_action=( + [at.cpu() for at in gen_data_noised.xt_tokens_action] + if gen_data_noised.xt_tokens_action is not None + else None + ), + fps_action=gen_data_clean.fps_action, + action_domain_id=gen_data_clean.action_domain_id, + ) + + packed_sequence = self._pack_input_sequence( + sequence_plans, input_text_indexes, gen_data_for_packing, timesteps.cpu() + ) + + # Overwrite packed tokens with the actual noised values (already on CUDA) + assert packed_sequence.vision is not None, "Packed vision data is required" + packed_sequence.vision.tokens = gen_data_noised.xt_tokens_vision + if packed_sequence.action is not None and gen_data_noised.xt_tokens_action is not None: + packed_sequence.action.tokens = gen_data_noised.xt_tokens_action + if packed_sequence.sound is not None and gen_data_noised.xt_tokens_sound is not None: + packed_sequence.sound.tokens = gen_data_noised.xt_tokens_sound + + packed_sequence.to_cuda() + + return self.denoise( + net=self.denoiser_nets[net_type], + data_batch_packed=packed_sequence, + ) + + def _gen_data_from_student( + self, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + iteration: int, + ) -> tuple[GenerationDataClean, PackedSequence, torch.Tensor, dict]: + """Generate x0 from the student network. + + Dispatches to :meth:`_forward_simulation` or :meth:`_backward_simulation` based on + ``config.simulation_mode``. ``iteration`` drives the rcm-style rollout-length cycle in + backward simulation (ignored by forward simulation). + + Returns: + Tuple of: + - ``gen_data_clean_student``: :class:`GenerationDataClean` with student-generated x0 + for all modalities (``x0_tokens_*``). + - ``packed_student``: :class:`PackedSequence` from the student forward pass; provides + per-modality ``condition_mask`` used for re-noising and loss computation. + - ``sigma_student``: The student sigma ``(B, 1)`` used. + - ``out_student``: Raw network output dict; kept for FSDP dummy terms when a + modality has no data in the current batch. + """ + if self.config.simulation_mode == "forward": + return self._forward_simulation(input_text_indexes, sequence_plans, gen_data_clean) + elif self.config.simulation_mode == "backward": + return self._backward_simulation(input_text_indexes, sequence_plans, gen_data_clean, iteration) + else: + raise NotImplementedError + + def _forward_simulation( + self, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + ) -> tuple[GenerationDataClean, PackedSequence, torch.Tensor, dict]: + """Generate x0 via a single student forward pass (fastgen-style DMD2). + + Samples one sigma from the fixed-step schedule, adds noise to clean data, runs the + student once, and converts the velocity prediction to x0. No gradient flows through + any re-noising step — all gradient signal is carried by this single forward pass. + + For single-step distillation the student sees near-pure noise (``sigma = t_list[0]``). + For multi-step distillation it sees real data noised to a randomly sampled sigma from ``t_list``. + """ + B = gen_data_clean.batch_size + + # Sample student noise level + sigma_student = self._sample_student_sigma(B) # [B,1], continuous in [0, 1] + num_train_timesteps = self.config.rectified_flow_inference_config.num_train_timesteps + timesteps_student = sigma_student * num_train_timesteps # [B,1] discrete in [0, num_train_timesteps] + + # Pack sequence and add noise (following the standard MoT training pipeline) + packed_sequence = self._pack_input_sequence( + sequence_plans, input_text_indexes, gen_data_clean, timesteps_student.cpu() + ) + sigmas_action_dense = self._action_sigmas_from_full(sigma_student, sequence_plans) + sigmas_sound_dense = self._sound_sigmas_from_full( + sigma_student, sequence_plans, cast(list[torch.Tensor] | None, gen_data_clean.x0_tokens_sound) + ) # None or [n_sound,1] + gen_data_noised = self._add_noise_to_input( + gen_data_clean, + packed_sequence, + sigma_student, + sigmas_action=sigmas_action_dense, + sigmas_sound=sigmas_sound_dense, + ) + self._replace_clean_with_noised(packed_sequence, gen_data_noised) + packed_sequence.to_cuda() + + # Student forward pass + assert packed_sequence.vision is not None + out_student = self.denoise( + net=self.denoiser_nets["student"], + data_batch_packed=packed_sequence, + ) + + # Vision x0: zero out conditioned frames so their x0 ≈ xt (clean) + condition_mask_vision = packed_sequence.vision.condition_mask # list of [T_i,1,1] + assert isinstance(condition_mask_vision, list) + sigma_vision = [ + sigma_student[i].view(1, 1, 1) * (1.0 - condition_mask_vision[i]) for i in range(B) + ] # list of [T_i,1,1] + xt_vision = [xt.to(**self.tensor_kwargs) for xt in gen_data_noised.xt_tokens_vision] # list of [C,T_i,H_i,W_i] + gen_x0_vision = self._velocity_to_x0( + xt_vision, out_student["preds_vision"], sigma_vision + ) # list of [C,T_i,H_i,W_i] + + # Action x0 (when action_gen=True and batch contains action data) + gen_x0_action = None + + # Sound x0 (when sound_gen=True and batch contains sound data) + gen_x0_sound = None + + gen_data_clean_student = GenerationDataClean( + batch_size=gen_data_clean.batch_size, + is_image_batch=gen_data_clean.is_image_batch, + raw_state_vision=gen_data_clean.raw_state_vision, + x0_tokens_vision=gen_x0_vision, + fps_vision=gen_data_clean.fps_vision, + num_vision_items_per_sample=gen_data_clean.num_vision_items_per_sample, + raw_state_sound=gen_data_clean.raw_state_sound, + x0_tokens_sound=gen_x0_sound, + fps_sound=gen_data_clean.fps_sound, + raw_state_action=gen_data_clean.raw_state_action, + x0_tokens_action=gen_x0_action, + fps_action=gen_data_clean.fps_action, + action_domain_id=gen_data_clean.action_domain_id, + ) + + return gen_data_clean_student, packed_sequence, sigma_student, out_student + + def _ode_step( + self, + xt: list[torch.Tensor], # list of [C,...], current noised tokens + v_pred: list[torch.Tensor], # list of [C,...], velocity predictions + noisy_mask: list[torch.Tensor], # list of [...], broadcastable to xt; 0 for conditioned frames + delta_sigma: float, + ) -> list[torch.Tensor]: + """Euler ODE step: xt_next = xt + delta_sigma * v * noisy_mask.""" + result = [] + for xt_i, v_i, mask_i in zip(xt, v_pred, noisy_mask): + xt_next_i = xt_i + delta_sigma * v_i * mask_i + result.append(xt_next_i.to(**self.tensor_kwargs)) + return result + + def _sde_step( + self, + xt: list[torch.Tensor], # list of [C,...], current noised tokens + x0_pred: list[torch.Tensor], # list of [C,...], predicted clean tokens + noisy_mask: list[torch.Tensor], # list of [...], 1 for generated frames, 0 for conditioned + cond_mask: list[torch.Tensor], # list of [...], 1 for conditioned frames, 0 for generated + sigma_next: float, + ) -> list[torch.Tensor]: + """SDE re-noise step: re-noise x0 to sigma_next for generated frames; keep conditioned frames.""" + result = [] + noises = [torch.randn_like(x0_i, **self.tensor_kwargs_fp32) for x0_i in x0_pred] # list of [C,...] + context_parallel_broadcast_tensor_list(noises, getattr(self, "parallel_dims", None)) + for xt_i, x0_i, noisy_i, cond_i, noise_i in zip(xt, x0_pred, noisy_mask, cond_mask, noises): + xt_next_i = ( + noisy_i * ((1.0 - sigma_next) * x0_i + sigma_next * noise_i) + cond_i.to(**self.tensor_kwargs) * xt_i + ) # [C,...] + xt_next_i = xt_next_i.to(**self.tensor_kwargs) # [C,...] + result.append(xt_next_i) + return result + + def _backward_n_steps(self, max_n_steps: int, iteration: int) -> int: + """rcm-style rollout length: cycle deterministically by per-net update index modulo max. + + Matches rcm's ``effective_iteration % max_simulation_steps + 1`` (not iid sampling), so + every rollout length in ``[1, max_n_steps]`` is visited in a fixed cycle. Cycling on the + per-net update counter (rather than the global iteration) keeps the cycle complete even + when ``student_update_freq`` shares a common factor with ``max_n_steps``. Deterministic in + ``iteration``, hence identical across ranks without a broadcast. + """ + phase = self.get_phase(iteration) + update_idx = ( + self.get_student_iteration(iteration) if phase == "student" else self.get_critic_iteration(iteration) + ) + # get_critic_iteration is -1 at iteration 0 during critic-only warmup (rcm-faithful but a + # negative-modulo quirk); clamp so the cycle starts at n_steps=1 rather than wrapping to max. + update_idx = max(update_idx, 0) + return update_idx % max_n_steps + 1 + + def _backward_simulation( + self, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + iteration: int, + ) -> tuple[GenerationDataClean, PackedSequence, torch.Tensor, dict]: + """Generate x0 via an rcm-style multi-step denoising rollout from pure noise. + + The rollout always starts from pure noise (``t_list[0]`` must be 1.0, so the RF seed + ``(1 - sigma) * x0 + sigma * eps`` reduces to ``eps`` with zero clean-data leakage) and a + iteration-cycled number of schedule steps (schedule prefix). Each step applies the ODE + Euler step or SDE re-noising to progress toward sigma=0. Conditioned frames are held fixed + throughout via ``noisy_mask``. + + Gradient flows through the last ``config.backward_grad_steps`` student forward passes + (1 = last step only, matching predict2_distill; -1 = full BPTT through all steps). + + Returns the same tuple as :meth:`_forward_simulation`. + """ + backward_grad_steps = self.config.backward_grad_steps + assert backward_grad_steps == -1 or backward_grad_steps >= 1, ( + f"backward_grad_steps must be -1 (all steps) or >= 1, got {backward_grad_steps}" + ) + + B = gen_data_clean.batch_size + num_train_timesteps = self.config.rectified_flow_inference_config.num_train_timesteps + sample_type = self.config.fixed_step_sampler_config.sample_type + t_list = list(self.config.fixed_step_sampler_config.t_list) + assert len(t_list) > 0, "fixed_step_sampler_config.t_list must not be empty" + full_t_list = t_list if t_list[-1] == 0.0 else t_list + [0.0] + assert len(full_t_list) > 1, "fixed_step_sampler_config.t_list must contain a nonzero sigma" + + # Pure-noise start contract (rcm / self-forcing style): the schedule must begin at + # sigma=1.0 so the rollout seed is pure noise. Under RF interpolation the seed is + # (1 - sigma) * x0 + sigma * eps; at sigma=1.0 the clean term vanishes (seed == eps), + # so no clean-data signal leaks into the student generation. t_list[0] < 1.0 would + # blend x0 into the start state. + assert abs(full_t_list[0] - 1.0) < 1e-6, ( + f"backward simulation requires t_list[0] == 1.0 for a pure-noise start, got {full_t_list[0]}" + ) + + # rcm-style iteration-cycled trajectory length: keep the pure-noise start (schedule prefix) + # and cycle how many denoising steps to run (by per-net update index) before descending to + # sigma=0. Fewer steps = a shorter rollout from the SAME pure-noise start, never a + # half-noised real-data start. + nonzero_levels = full_t_list[:-1] # descending sigmas; nonzero_levels[0] == 1.0 + n_steps = self._backward_n_steps(len(nonzero_levels), iteration) + full_t_list = nonzero_levels[:n_steps] + [0.0] + log.info(f"[DMD2RF] backward_simulation: n_steps={n_steps}, t_list={full_t_list}") + grad_steps = n_steps if backward_grad_steps == -1 else backward_grad_steps + + # Pack once at sigma_max to obtain condition masks (fixed throughout the rollout) + sigma_max = torch.full((B, 1), full_t_list[0], **self.tensor_kwargs_fp32) # [B,1] + timesteps_max = sigma_max * num_train_timesteps # [B,1] + packed_sequence = self._pack_input_sequence( + sequence_plans, input_text_indexes, gen_data_clean, timesteps_max.cpu() + ) + # Initialize noised data: generation frames ≈ pure noise, conditioned frames = x0_clean + sigmas_action_dense_init = self._action_sigmas_from_full(sigma_max, sequence_plans) + sigmas_sound_dense_init = self._sound_sigmas_from_full( + sigma_max, sequence_plans, cast(list[torch.Tensor] | None, gen_data_clean.x0_tokens_sound) + ) # None or [n_sound,1] + gen_data_noised = self._add_noise_to_input( + gen_data_clean, + packed_sequence, + sigma_max, + sigmas_action=sigmas_action_dense_init, + sigmas_sound=sigmas_sound_dense_init, + ) + + assert packed_sequence.vision is not None + condition_mask_vision = packed_sequence.vision.condition_mask # list of [T_i,1,1] + assert isinstance(condition_mask_vision, list) + noisy_mask_vision = [1.0 - cm for cm in condition_mask_vision] # list of [T_i,1,1] + + # Condition masks for action / sound — pre-computed once, fixed throughout the rollout + condition_mask_action = ( + packed_sequence.action.condition_mask if packed_sequence.action is not None else None + ) # list of [T_i,1] or None + condition_mask_sound = ( + packed_sequence.sound.condition_mask if packed_sequence.sound is not None else None + ) # list of [T_i,1] or None + noisy_mask_action = ( + [1.0 - cm for cm in condition_mask_action] if condition_mask_action is not None else None + ) # list of [T_i,1] or None + cond_mask_sound = ( + [cm.T for cm in condition_mask_sound] if condition_mask_sound is not None else None + ) # list of [1,T_i] or None (transposed for [C_sound,T_sound] broadcasting) + noisy_mask_sound = ( + [1.0 - cm for cm in cond_mask_sound] if cond_mask_sound is not None else None + ) # list of [1,T_i] or None + has_action = ( + self.config.action_gen + and packed_sequence.action is not None + and gen_data_noised.xt_tokens_action is not None + and condition_mask_action is not None + ) + has_sound = ( + self.config.sound_gen + and packed_sequence.sound is not None + and gen_data_noised.xt_tokens_sound is not None + and condition_mask_sound is not None + and cond_mask_sound is not None + and noisy_mask_sound is not None + ) + + out_student: dict = {} + x0_pred_vision: list[torch.Tensor] = [] + x0_pred_action: Optional[list[torch.Tensor]] = None + x0_pred_sound: Optional[list[torch.Tensor]] = None + + for count, (sigma_cur_val, sigma_next_val) in enumerate(zip(full_t_list[:-1], full_t_list[1:])): + sigma_cur = torch.full((B, 1), sigma_cur_val, **self.tensor_kwargs_fp32) # [B,1] + timesteps_cur = sigma_cur * num_train_timesteps # [B,1] + + # Gradient enabled for the last `grad_steps` steps; all prior steps run under no_grad + enable_grad = count >= n_steps - grad_steps + with torch.set_grad_enabled(enable_grad): + out_student = self._pack_and_denoise( + gen_data_clean, + gen_data_noised, + timesteps_cur, + input_text_indexes, + sequence_plans, + net_type="student", + ) + + # --- Vision x0 --- + xt_vision = [ + xt.to(**self.tensor_kwargs) for xt in gen_data_noised.xt_tokens_vision + ] # list of [C,T_i,H_i,W_i] + sigma_vision = [sigma_cur[i].view(1, 1, 1) * noisy_mask_vision[i] for i in range(B)] # list of [T_i,1,1] + x0_pred_vision = self._velocity_to_x0( + xt_vision, out_student["preds_vision"], sigma_vision + ) # list of [C,T_i,H_i,W_i] + + # --- Action x0 --- + x0_pred_action = None + xt_action: Optional[list[torch.Tensor]] = None + + # --- Sound x0 --- + x0_pred_sound = None + xt_sound: Optional[list[torch.Tensor]] = None + + if sigma_next_val > 0.0: + # Compute xt_next for the next denoising step. + # Conditioned frames (noisy_mask=0) are kept unchanged via masking. + if sample_type == "ode": + # Euler step: xt_next = xt + (sigma_next - sigma_cur) * v * noisy_mask + delta_sigma = sigma_next_val - sigma_cur_val # negative scalar + xt_next_vision = self._ode_step( + xt_vision, out_student["preds_vision"], noisy_mask_vision, delta_sigma + ) # list of [C,T_i,H_i,W_i] + xt_next_action = None + xt_next_sound = None + else: + # SDE step: re-noise x0 to sigma_next for generation frames; + # keep conditioned frames at their current xt value (= x0_clean). + xt_next_vision = self._sde_step( + xt_vision, x0_pred_vision, noisy_mask_vision, condition_mask_vision, sigma_next_val + ) # list of [C,T_i,H_i,W_i] + xt_next_action = None + xt_next_sound = None + + # Build updated GenerationDataNoised for the next step. + # Only xt_tokens_* fields are consumed by _pack_and_denoise; sigmas_vision is + # set for correctness but is not read by the packing logic. + sigma_next_tensor = torch.full((B, 1), sigma_next_val, **self.tensor_kwargs_fp32) # [B,1] + sigmas_vision_next = [ + sigma_next_tensor[i].view(1, 1, 1) * noisy_mask_vision[i] for i in range(B) + ] # list of [T_i,1,1] + sigmas_sound_next = None + if has_sound: + assert condition_mask_sound is not None + assert gen_data_noised.xt_tokens_sound is not None + sigmas_sound_dense_next = self._sound_sigmas_from_full( + sigma_next_tensor, + sequence_plans, + cast(list[torch.Tensor] | None, gen_data_clean.x0_tokens_sound), + ) # [n_sound,1] + assert sigmas_sound_dense_next is not None + sigmas_sound_next = [ + sigmas_sound_dense_next[i].view(-1, 1)[: gen_data_noised.xt_tokens_sound[i].shape[1]].T + * (1.0 - condition_mask_sound[i].T) + for i in range(len(gen_data_noised.xt_tokens_sound)) + ] # list of [1,T_i] + gen_data_noised = GenerationDataNoised( + batch_size=B, + epsilon_vision=gen_data_noised.epsilon_vision, + xt_tokens_vision=xt_next_vision, + vt_target_vision=gen_data_noised.vt_target_vision, + sigmas_vision=sigmas_vision_next, + epsilon_action=gen_data_noised.epsilon_action, + xt_tokens_action=xt_next_action, + vt_target_action=gen_data_noised.vt_target_action, + sigmas_action=gen_data_noised.sigmas_action, + epsilon_sound=gen_data_noised.epsilon_sound, + xt_tokens_sound=xt_next_sound, + vt_target_sound=gen_data_noised.vt_target_sound, + sigmas_sound=sigmas_sound_next, + ) + + gen_data_clean_student = GenerationDataClean( + batch_size=gen_data_clean.batch_size, + is_image_batch=gen_data_clean.is_image_batch, + raw_state_vision=gen_data_clean.raw_state_vision, + x0_tokens_vision=x0_pred_vision, + fps_vision=gen_data_clean.fps_vision, + num_vision_items_per_sample=gen_data_clean.num_vision_items_per_sample, + raw_state_sound=gen_data_clean.raw_state_sound, + x0_tokens_sound=x0_pred_sound, + fps_sound=gen_data_clean.fps_sound, + raw_state_action=gen_data_clean.raw_state_action, + x0_tokens_action=x0_pred_action, + fps_action=gen_data_clean.fps_action, + action_domain_id=gen_data_clean.action_domain_id, + ) + + return gen_data_clean_student, packed_sequence, sigma_max, out_student + + # ------------------------ Training step ------------------------ + + @staticmethod + def _validate_vision_only_sequence_plans(sequence_plans: list[SequencePlan]) -> None: + """Reject action or sound samples from the vision-only public DMD2 path.""" + if any(plan.has_action or plan.has_sound for plan in sequence_plans): + raise ValueError("DMD2-RF OSS supports only vision T2I/I2V batches") + + def _setup_grad_requirements(self, iteration: int) -> None: + # NOTE: requires_grad_(True) must be applied to the entire active network, not per-parameter, + # because torch.compile + activation checkpointing requires all Q/K/V tensors in a layer to + # share the same requires_grad state. This means grad is computed for ALL params (including + # frozen backbone), making the effective clip threshold looser — but optimizer.step() only + # updates the selected params, so weight correctness is unaffected. + if self.get_phase(iteration) == "student": + self.net.train().requires_grad_(True) + self.net_fake_score.eval().requires_grad_(False) + else: + self.net.eval().requires_grad_(False) + self.net_fake_score.train().requires_grad_(True) + # Teacher must remain frozen + in eval() across all phase switches; the trainer + # may have recursively re-enabled train() on the parent module. + self.net_teacher.eval().requires_grad_(False) + + def training_step(self, data_batch: dict[str, Any], iteration: int) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Perform a single training step for DMD2 distillation. + + Delegates to :meth:`training_step_generator` (student update) or :meth:`training_step_critic` + (fake-score update) depending on the current iteration phase. + """ + # Update stats on how many videos the model has seen + if self.parallel_dims is None or self.parallel_dims.cp_rank == 0: + self._update_train_stats(data_batch) + + # Freeze / unfreeze networks according to current phase + self._setup_grad_requirements(iteration) + + # Load, apply dropout, and tokenize input captions + input_text_indexes = self._load_and_tokenize_text_data(data_batch, iteration) + + # Build sequence plans if not present. SequencePlan has the conditioning information. + sequence_plans = build_sequence_plans_from_data_batch( + data_batch=data_batch, + input_video_key=self.input_video_key, + input_image_key=self.input_image_key, + ) + + self._validate_vision_only_sequence_plans(sequence_plans) + + # Get data from raw data batch and tokenize into corresponding tokens for *generation* task + gen_data_clean = self.get_data_and_condition(data_batch) + data_resolutions, num_vision_tokens_per_sample = self._get_vision_noise_sampling_metadata( + data_batch, gen_data_clean + ) + + # The noise addition and denoising happens inside respective phases + if self.get_phase(iteration) == "student": + output_batch, loss = self.training_step_generator( + input_text_indexes, + sequence_plans, + gen_data_clean, + data_resolutions, + num_vision_tokens_per_sample, + iteration, + ) + else: + output_batch, loss = self.training_step_critic( + input_text_indexes, + sequence_plans, + gen_data_clean, + data_resolutions, + num_vision_tokens_per_sample, + iteration, + ) + + return output_batch, loss + + # -------------------- Generator (student) update -------------------- + + def training_step_generator( + self, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + data_resolutions: list[str], + num_vision_tokens_per_sample: list[int], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Student update step following the fastgen DMD2 algorithm. + + 1. Student generates x0 (forward simulation: single pass; backward simulation: multi-step rollout). + 2. Student x0 is re-noised at a random sigma for the critic networks. + 3. Teacher and fake-score predict x0 on the re-noised data. + 4. VSD loss drives the student. + + All vision data is handled as per-sample lists to support variable shapes, consistent with the MoT pipeline. + + """ + B = gen_data_clean.batch_size + + # 1. Student generates x0 (gradient flows through here) + gen_data_student, packed_student, _, out_student = self._gen_data_from_student( + input_text_indexes, sequence_plans, gen_data_clean, iteration + ) + + # 2. Sample critic noise level and re-noise student output for all modalities + num_vision_latent_frames = [x.shape[2] for x in gen_data_student.x0_tokens_vision] + timesteps_critic, sigmas_critic = self._get_train_noise_level_vision( + batch_size=B, + is_image_batch=gen_data_student.is_image_batch, + num_vision_latent_frames=num_vision_latent_frames, + resolutions=data_resolutions, + num_tokens=num_vision_tokens_per_sample, + ) # [B,1], [B,1] + + # Broadcast timesteps/sigmas across CP group to ensure consistency + if self.parallel_dims is not None and self.parallel_dims.cp_enabled: + src_rank = 0 + cp_group = self.parallel_dims.cp_mesh.get_group() + global_src_rank = torch.distributed.get_global_rank(cp_group, src_rank) + timesteps_critic = timesteps_critic.contiguous() + sigmas_critic = sigmas_critic.contiguous() + torch.distributed.broadcast(timesteps_critic, src=global_src_rank, group=cp_group) + torch.distributed.broadcast(sigmas_critic, src=global_src_rank, group=cp_group) + + sigmas_action_critic = self._action_sigmas_from_full(sigmas_critic, sequence_plans) + sigmas_sound_critic = self._sound_sigmas_from_full( + sigmas_critic, sequence_plans, cast(list[torch.Tensor] | None, gen_data_student.x0_tokens_sound) + ) # None or [n_sound,1] + gen_data_noised_critic = self._add_noise_to_input( + gen_data_student, + packed_student, + sigmas_critic, + sigmas_action=sigmas_action_critic, + sigmas_sound=sigmas_sound_critic, + ) + + # 3. Fake-score prediction (no grad -- no GAN in cosmos3) + with torch.no_grad(): + out_fake = self._pack_and_denoise( + gen_data_student, + gen_data_noised_critic, + timesteps_critic, + input_text_indexes, + sequence_plans, + net_type="fake_score", + ) + fake_v = cast(list[torch.Tensor], out_fake["preds_vision"]) # list of [C,T_i,H_i,W_i] + fake_x0 = self._velocity_to_x0( + gen_data_noised_critic.xt_tokens_vision, + fake_v, + gen_data_noised_critic.sigmas_vision, # type: ignore[arg-type] + ) + + # 4. Teacher prediction (no grad) + with torch.no_grad(): + out_teacher = self._pack_and_denoise( + gen_data_student, + gen_data_noised_critic, + timesteps_critic, + input_text_indexes, + sequence_plans, + net_type="teacher", + ) + teacher_v = cast(list[torch.Tensor], out_teacher["preds_vision"]) # list of [C,T_i,H_i,W_i] + teacher_v_guided = teacher_v # list of [C,T_i,H_i,W_i] + teacher_x0 = self._velocity_to_x0( + gen_data_noised_critic.xt_tokens_vision, + teacher_v, + gen_data_noised_critic.sigmas_vision, # type: ignore[arg-type] + ) + log.info(f"student phase, after pack_and_denoise teacher, teacher_x0: {teacher_x0[0].shape}") + # Snapshot the pre-CFG conditional x0. teacher_x0 below is overwritten with the + # CFG-extrapolated target; keeping the conditional view lets the diagnostics separate + # critic weight-drift (fake vs teacher-conditional) from the CFG term (guided vs conditional). + teacher_x0_cond = teacher_x0 # list of [C,T_i,H_i,W_i] + # Optional classifier-free guidance for teacher + if self.config.teacher_guidance > 1.0: + # Tokenize the uncond prompt with the SAME modality formatting as the batch + # (image vs video). Hard-coding is_video=False corrupted the uncond prediction + # on video batches, biasing the CFG-combined teacher target. + uncond_is_video = not gen_data_student.is_image_batch + uncond_text = [ + tokenize_caption( + self.config.teacher_negative_prompt, + self.vlm_tokenizer, + is_video=uncond_is_video, + use_system_prompt=self.vlm_config.use_system_prompt, + ) + for _ in range(B) + ] + out_teacher_uncond = self._pack_and_denoise( + gen_data_student, + gen_data_noised_critic, + timesteps_critic, + uncond_text, + sequence_plans, + net_type="teacher", + ) + teacher_v_uncond = cast( + list[torch.Tensor], out_teacher_uncond["preds_vision"] + ) # list of [C,T_i,H_i,W_i] + teacher_x0_uncond = self._velocity_to_x0( + gen_data_noised_critic.xt_tokens_vision, + teacher_v_uncond, + gen_data_noised_critic.sigmas_vision, # type: ignore[arg-type] + ) + teacher_v_guided = [ + teacher_v[i] + (self.config.teacher_guidance - 1.0) * (teacher_v[i] - teacher_v_uncond[i]) + for i in range(B) + ] # list of [C,T_i,H_i,W_i] + teacher_x0 = [ + teacher_x0[i] + (self.config.teacher_guidance - 1.0) * (teacher_x0[i] - teacher_x0_uncond[i]) + for i in range(B) + ] # list of [C,T_i,H_i,W_i] + + # 5. VSD loss — vision (gradient flows into gen_data_student.x0_tokens_vision -> student) + assert gen_data_student.x0_tokens_vision is not None + assert packed_student.vision is not None + noisy_mask_vision = [ + (1.0 - cm).to(gen_data_student.x0_tokens_vision[0].dtype) + for cm in packed_student.vision.condition_mask # type: ignore[union-attr] + ] # list of [T_i,1,1] + if self.config.vsd_gradient_space == "velocity": + vsd_grad_vision = [teacher_v_guided[i] - fake_v[i] for i in range(B)] # list of [C,T_i,H_i,W_i] + vsd_loss = variational_score_distillation_loss_from_gradient( + gen_data_student.x0_tokens_vision, + vsd_grad_vision, + weight_reference=teacher_x0, + loss_mask=noisy_mask_vision, + reduction=self.config.vsd_loss_reduction, + ) + elif self.config.vsd_gradient_space == "x0": + vsd_loss = variational_score_distillation_loss( + gen_data_student.x0_tokens_vision, + teacher_x0, + fake_x0, + loss_mask=noisy_mask_vision, + reduction=self.config.vsd_loss_reduction, + ) + else: + raise ValueError(f"Unknown vsd_gradient_space={self.config.vsd_gradient_space}") + total_loss = self.config.loss_scale_sid * vsd_loss + + + # VSD loss — sound (when configured) + if self.config.sound_gen: + if gen_data_student.x0_tokens_sound is None: + # FSDP dummy: keep student/teacher/fake_score sound params in backward graph + total_loss = ( + total_loss + + 0.0 * sum(p.sum() for p in out_student["preds_sound"]) # type: ignore[union-attr] + + 0.0 * sum(p.sum() for p in out_fake["preds_sound"]) # type: ignore[union-attr] + + 0.0 * sum(p.sum() for p in out_teacher["preds_sound"]) # type: ignore[union-attr] + ) + + log.info( + f"Iteration {iteration}, student phase, after vsd_loss, " + f"vsd loss: {vsd_loss.item():.6f}, total loss: {total_loss.item():.6f}" + ) + + # Diagnostic: track the VSD gradient signal components + with torch.no_grad(): + fake_teacher_diff = torch.stack( + [(fake_x0[i] - teacher_x0[i]).abs().mean() for i in range(B)] + ).mean() # scalar + gen_x0_vision = gen_data_student.x0_tokens_vision + gen_teacher_diff = torch.stack( + [(gen_x0_vision[i].detach() - teacher_x0[i]).abs().mean() for i in range(B)] + ).mean() # scalar + fake_student_diff = torch.stack( + [(fake_x0[i] - gen_x0_vision[i].detach()).abs().mean() for i in range(B)] + ).mean() # scalar + # Clean critic-vs-teacher-weight drift: compares fake (conditional, no CFG) against the + # teacher's conditional x0, so it is not confounded by the CFG term the way fake_teacher_diff is. + fake_teacher_cond_diff = torch.stack( + [(fake_x0[i] - teacher_x0_cond[i]).abs().mean() for i in range(B)] + ).mean() # scalar + # Magnitude of the CFG extrapolation in x0 space, |teacher_guided - teacher_conditional|. + teacher_cfg_term = torch.stack( + [(teacher_x0[i] - teacher_x0_cond[i]).abs().mean() for i in range(B)] + ).mean() # scalar + diff_ratio_den = gen_teacher_diff.clamp(min=1e-6) # scalar + fake_student_to_gen_teacher_ratio = fake_student_diff / diff_ratio_den # scalar + fake_teacher_to_gen_teacher_ratio = fake_teacher_diff / diff_ratio_den # scalar + dmd_weighted_diag_keys = ( + "dmd_gen_teacher_diff_masked", + "dmd_gen_teacher_diff_masked_image", + "dmd_gen_teacher_diff_masked_video", + "dmd_gen_teacher_diff_masked_sigma_000_025", + "dmd_gen_teacher_diff_masked_sigma_025_050", + "dmd_gen_teacher_diff_masked_sigma_050_075", + "dmd_gen_teacher_diff_masked_sigma_075_100", + "dmd_fake_student_diff_masked", + "dmd_fake_student_diff_masked_image", + "dmd_fake_student_diff_masked_video", + "dmd_fake_student_diff_masked_sigma_000_025", + "dmd_fake_student_diff_masked_sigma_025_050", + "dmd_fake_student_diff_masked_sigma_050_075", + "dmd_fake_student_diff_masked_sigma_075_100", + "dmd_fake_teacher_diff_masked_sigma_000_025", + "dmd_fake_teacher_diff_masked_sigma_025_050", + "dmd_fake_teacher_diff_masked_sigma_050_075", + "dmd_fake_teacher_diff_masked_sigma_075_100", + "dmd_fake_teacher_cond_diff_masked_sigma_000_025", + "dmd_fake_teacher_cond_diff_masked_sigma_025_050", + "dmd_fake_teacher_cond_diff_masked_sigma_050_075", + "dmd_fake_teacher_cond_diff_masked_sigma_075_100", + "dmd_teacher_cfg_term_masked_sigma_000_025", + "dmd_teacher_cfg_term_masked_sigma_025_050", + "dmd_teacher_cfg_term_masked_sigma_050_075", + "dmd_teacher_cfg_term_masked_sigma_075_100", + "dmd_fake_student_to_gen_teacher_ratio_masked", + "dmd_fake_student_to_gen_teacher_ratio_masked_image", + "dmd_fake_student_to_gen_teacher_ratio_masked_video", + "dmd_fake_student_to_gen_teacher_ratio_masked_sigma_000_025", + "dmd_fake_student_to_gen_teacher_ratio_masked_sigma_025_050", + "dmd_fake_student_to_gen_teacher_ratio_masked_sigma_050_075", + "dmd_fake_student_to_gen_teacher_ratio_masked_sigma_075_100", + "dmd_fake_teacher_to_gen_teacher_ratio_masked", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_image", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_video", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_000_025", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_025_050", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_050_075", + "dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_075_100", + "dmd_vsd_teacher_pull_cosine_masked", + "dmd_vsd_teacher_pull_cosine_masked_image", + "dmd_vsd_teacher_pull_cosine_masked_video", + "dmd_vsd_teacher_pull_cosine_masked_sigma_000_025", + "dmd_vsd_teacher_pull_cosine_masked_sigma_025_050", + "dmd_vsd_teacher_pull_cosine_masked_sigma_050_075", + "dmd_vsd_teacher_pull_cosine_masked_sigma_075_100", + ) + dmd_weighted_diag = { + f"dmd_weighted_{key}_{suffix}": torch.zeros((), **self.tensor_kwargs_fp32) # scalar + for key in dmd_weighted_diag_keys + for suffix in ("sum", "count") + } + assert gen_data_noised_critic.sigmas_vision is not None + # VSD gradient vector norm: ||(f - t) * w|| per sample, averaged + vsd_grad_norms = [] + for i in range(B): + g_fp32 = gen_x0_vision[i].detach().float() # [C,T,H,W] + t_fp32 = teacher_x0[i].float() # [C,T,H,W] + f_fp32 = fake_x0[i].float() # [C,T,H,W] + diff_abs_i = (g_fp32 - t_fp32).abs() # [C,T,H,W] + fake_teacher_diff_abs_i = (f_fp32 - t_fp32).abs() # [C,T,H,W] + fake_student_diff_abs_i = (f_fp32 - g_fp32).abs() # [C,T,H,W] + tc_fp32 = teacher_x0_cond[i].float() # [C,T,H,W] + fake_teacher_cond_diff_abs_i = (f_fp32 - tc_fp32).abs() # [C,T,H,W] + teacher_cfg_term_abs_i = (t_fp32 - tc_fp32).abs() # [C,T,H,W] + mask_i = noisy_mask_vision[i].float().expand_as(diff_abs_i) # [C,T,H,W] + mask_count_i = mask_i.sum() # scalar + sample_count_i = (mask_count_i > 0).float() # scalar + masked_diff_i = (diff_abs_i * mask_i).sum() / mask_count_i.clamp(min=1.0) # scalar + weighted_diff_i = masked_diff_i * sample_count_i # scalar + masked_fake_teacher_diff_i = (fake_teacher_diff_abs_i * mask_i).sum() / mask_count_i.clamp( + min=1.0 + ) # scalar + masked_fake_student_diff_i = (fake_student_diff_abs_i * mask_i).sum() / mask_count_i.clamp( + min=1.0 + ) # scalar + weighted_fake_student_diff_i = masked_fake_student_diff_i * sample_count_i # scalar + masked_fake_teacher_cond_diff_i = (fake_teacher_cond_diff_abs_i * mask_i).sum() / mask_count_i.clamp( + min=1.0 + ) # scalar + masked_teacher_cfg_term_i = (teacher_cfg_term_abs_i * mask_i).sum() / mask_count_i.clamp( + min=1.0 + ) # scalar + masked_diff_ratio_den_i = masked_diff_i.clamp(min=1e-6) # scalar + masked_fake_student_to_gen_teacher_ratio_i = ( + masked_fake_student_diff_i / masked_diff_ratio_den_i + ) # scalar + masked_fake_teacher_to_gen_teacher_ratio_i = ( + masked_fake_teacher_diff_i / masked_diff_ratio_den_i + ) # scalar + weighted_fake_student_to_gen_teacher_ratio_i = ( + masked_fake_student_to_gen_teacher_ratio_i * sample_count_i + ) # scalar + weighted_fake_teacher_to_gen_teacher_ratio_i = ( + masked_fake_teacher_to_gen_teacher_ratio_i * sample_count_i + ) # scalar + teacher_pull_i = (t_fp32 - g_fp32) * mask_i # [C,T,H,W] + sigma_mask_i = noisy_mask_vision[i].float() # [T,1,1] + sigma_count_i = sigma_mask_i.sum() # scalar + sigma_vision_i = gen_data_noised_critic.sigmas_vision[i].detach().float() # [T,1,1] + sigma_i = (sigma_vision_i * sigma_mask_i).sum() / sigma_count_i.clamp(min=1.0) # scalar + if self.config.vsd_gradient_space == "velocity": + teacher_v_fp32 = teacher_v_guided[i].float() # [C,T,H,W] + fake_v_fp32 = fake_v[i].float() # [C,T,H,W] + # Convert the velocity-space update to x0-space so this cosine compares like vectors: + # teacher_x0 - fake_x0 = sigma * (fake_v - teacher_v). + vsd_update_i = sigma_vision_i * (fake_v_fp32 - teacher_v_fp32) * mask_i # [C,T,H,W] + raw_vsd_grad_i = teacher_v_guided[i] - fake_v[i] # [C,T,H,W] + else: + vsd_update_i = (t_fp32 - f_fp32) * mask_i # [C,T,H,W] + raw_vsd_grad_i = fake_x0[i] - teacher_x0[i] # [C,T,H,W] + cosine_den_i = teacher_pull_i.norm() * vsd_update_i.norm() # scalar + cosine_valid_i = sample_count_i * (cosine_den_i > 1e-12).float() # scalar + cosine_i = (teacher_pull_i * vsd_update_i).sum() / cosine_den_i.clamp(min=1e-12) # scalar + weighted_cosine_i = cosine_i * cosine_valid_i # scalar + + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_sum"] += weighted_diff_i + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_sum"] += weighted_fake_student_diff_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_sum"] += ( + weighted_fake_student_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_sum"] += ( + weighted_fake_teacher_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_sum"] += weighted_cosine_i + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_count"] += cosine_valid_i + if gen_x0_vision[i].shape[2] == 1: + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_image_sum"] += weighted_diff_i + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_image_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_image_sum"] += ( + weighted_fake_student_diff_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_image_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_image_sum"] += ( + weighted_fake_student_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_image_count"] += ( + sample_count_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_image_sum"] += ( + weighted_fake_teacher_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_image_count"] += ( + sample_count_i + ) + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_image_sum"] += weighted_cosine_i + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_image_count"] += cosine_valid_i + else: + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_video_sum"] += weighted_diff_i + dmd_weighted_diag["dmd_weighted_dmd_gen_teacher_diff_masked_video_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_video_sum"] += ( + weighted_fake_student_diff_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_student_diff_masked_video_count"] += sample_count_i + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_video_sum"] += ( + weighted_fake_student_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_student_to_gen_teacher_ratio_masked_video_count"] += ( + sample_count_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_video_sum"] += ( + weighted_fake_teacher_to_gen_teacher_ratio_i + ) + dmd_weighted_diag["dmd_weighted_dmd_fake_teacher_to_gen_teacher_ratio_masked_video_count"] += ( + sample_count_i + ) + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_video_sum"] += weighted_cosine_i + dmd_weighted_diag["dmd_weighted_dmd_vsd_teacher_pull_cosine_masked_video_count"] += cosine_valid_i + + sigma_bin_count_000_025 = sample_count_i * (sigma_i < 0.25).float() # scalar + sigma_bin_count_025_050 = sample_count_i * ((sigma_i >= 0.25) & (sigma_i < 0.50)).float() # scalar + sigma_bin_count_050_075 = sample_count_i * ((sigma_i >= 0.50) & (sigma_i < 0.75)).float() # scalar + sigma_bin_count_075_100 = sample_count_i * (sigma_i >= 0.75).float() # scalar + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_gen_teacher_diff_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_gen_teacher_diff_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_gen_teacher_diff_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_gen_teacher_diff_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += masked_diff_i * sigma_bin_count_i + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_fake_student_diff_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_fake_student_diff_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_fake_student_diff_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_fake_student_diff_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_fake_student_diff_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_fake_teacher_diff_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_fake_teacher_diff_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_fake_teacher_diff_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_fake_teacher_diff_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_fake_teacher_diff_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_fake_teacher_cond_diff_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_fake_teacher_cond_diff_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_fake_teacher_cond_diff_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_fake_teacher_cond_diff_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_fake_teacher_cond_diff_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_teacher_cfg_term_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_teacher_cfg_term_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_teacher_cfg_term_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_teacher_cfg_term_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_teacher_cfg_term_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_fake_student_to_gen_teacher_ratio_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_fake_student_to_gen_teacher_ratio_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_fake_student_to_gen_teacher_ratio_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_fake_student_to_gen_teacher_ratio_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_fake_student_to_gen_teacher_ratio_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_fake_teacher_to_gen_teacher_ratio_masked_sigma_075_100", sigma_bin_count_075_100), + ): + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += ( + masked_fake_teacher_to_gen_teacher_ratio_i * sigma_bin_count_i + ) + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += sigma_bin_count_i + for sigma_bin_key, sigma_bin_count_i in ( + ("dmd_vsd_teacher_pull_cosine_masked_sigma_000_025", sigma_bin_count_000_025), + ("dmd_vsd_teacher_pull_cosine_masked_sigma_025_050", sigma_bin_count_025_050), + ("dmd_vsd_teacher_pull_cosine_masked_sigma_050_075", sigma_bin_count_050_075), + ("dmd_vsd_teacher_pull_cosine_masked_sigma_075_100", sigma_bin_count_075_100), + ): + cosine_sigma_bin_count_i = cosine_valid_i * sigma_bin_count_i # scalar + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_sum"] += cosine_i * cosine_sigma_bin_count_i + dmd_weighted_diag[f"dmd_weighted_{sigma_bin_key}_count"] += cosine_sigma_bin_count_i + + # DMD normalization uses the x0-space teacher pull in both modes. In velocity mode + # raw_vsd_grad_i remains velocity-space, so this norm is useful as a trend metric + # within that mode but should not be compared directly against x0-mode norms. + w_i = 1.0 / ((g_fp32 - t_fp32).abs().mean() + 1e-6) # scalar + vsd_grad_i = raw_vsd_grad_i * w_i # [C,T,H,W] + vsd_grad_norms.append(vsd_grad_i.norm()) # scalar + vsd_grad_norm = torch.stack(vsd_grad_norms).mean() # scalar + + output_batch: dict[str, torch.Tensor] = { + "vsd_loss": vsd_loss.detach(), + "total_generator_loss": total_loss.detach(), + "sigma": sigmas_critic.detach(), + "flow_matching_loss_vision_per_instance": vsd_loss.detach().expand(B), # [B] + # Aliases for the WandB callback. + "dmd_loss_generator": total_loss.detach(), + "dmd_loss": vsd_loss.detach(), + # VSD diagnostic metrics + "dmd_fake_teacher_diff": fake_teacher_diff.detach(), + "dmd_fake_teacher_cond_diff": fake_teacher_cond_diff.detach(), + "dmd_teacher_cfg_term": teacher_cfg_term.detach(), + "dmd_gen_teacher_diff": gen_teacher_diff.detach(), + "dmd_fake_student_diff": fake_student_diff.detach(), + "dmd_fake_student_to_gen_teacher_ratio": fake_student_to_gen_teacher_ratio.detach(), + "dmd_fake_teacher_to_gen_teacher_ratio": fake_teacher_to_gen_teacher_ratio.detach(), + "dmd_vsd_grad_norm": vsd_grad_norm.detach(), + **{key: value.detach() for key, value in dmd_weighted_diag.items()}, + } + return output_batch, total_loss + + # -------------------- Critic (fake-score) update -------------------- + + def training_step_critic( + self, + input_text_indexes: list[list[int]], + sequence_plans: list[SequencePlan], + gen_data_clean: GenerationDataClean, + data_resolutions: list[str], + num_vision_tokens_per_sample: list[int], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Fake-score update step following the fastgen DMD2 algorithm. + + 1. Student generates x0 (no grad -- stop gradient). + 2. Student x0 is re-noised at a random sigma. + 3. Fake-score learns to denoise the student's distribution via flow-matching loss. + + All vision data is handled as per-sample lists to support variable shapes, consistent with the MoT pipeline. + + """ + B = gen_data_clean.batch_size + + # 1. Student generates x0 (no gradient for critic update) + with torch.no_grad(): + gen_data_student, packed_student, _, _ = self._gen_data_from_student( + input_text_indexes, sequence_plans, gen_data_clean, iteration + ) + + # 2. Sample critic noise level and re-noise student output for all modalities + num_vision_latent_frames = [x.shape[2] for x in gen_data_student.x0_tokens_vision] + timesteps_critic, sigmas_critic = self._get_train_noise_level_vision( + batch_size=B, + is_image_batch=gen_data_student.is_image_batch, + num_vision_latent_frames=num_vision_latent_frames, + resolutions=data_resolutions, + num_tokens=num_vision_tokens_per_sample, + ) # [B,1], [B,1] + + # Broadcast timesteps/sigmas across CP group to ensure consistency + if self.parallel_dims is not None and self.parallel_dims.cp_enabled: + src_rank = 0 + cp_group = self.parallel_dims.cp_mesh.get_group() + global_src_rank = torch.distributed.get_global_rank(cp_group, src_rank) + timesteps_critic = timesteps_critic.contiguous() + sigmas_critic = sigmas_critic.contiguous() + torch.distributed.broadcast(timesteps_critic, src=global_src_rank, group=cp_group) + torch.distributed.broadcast(sigmas_critic, src=global_src_rank, group=cp_group) + + with torch.no_grad(): + sigmas_action_critic = self._action_sigmas_from_full(sigmas_critic, sequence_plans) + sigmas_sound_critic = self._sound_sigmas_from_full( + sigmas_critic, sequence_plans, cast(list[torch.Tensor] | None, gen_data_student.x0_tokens_sound) + ) # None or [n_sound,1] + gen_data_noised_critic = self._add_noise_to_input( + gen_data_student, + packed_student, + sigmas_critic, + sigmas_action=sigmas_action_critic, + sigmas_sound=sigmas_sound_critic, + ) + + # 3. Fake-score forward (with gradient) + out_fake = self._pack_and_denoise( + gen_data_student, + gen_data_noised_critic, + timesteps_critic, + input_text_indexes, + sequence_plans, + net_type="fake_score", + ) + + # 4. Flow-matching loss — vision. active_mean preserves sum/active_count + # over generated frames; sum_rcm uses RCM's per-instance active-element sum. + rectified_flow = self.rectified_flow_image if gen_data_student.is_image_batch else self.rectified_flow_video + fake_score_loss_reduction = self.config.fake_score_loss_reduction + if fake_score_loss_reduction == "active_mean": + # Use loss_per_instance (unweighted) to match flow_matching_denoising_loss semantics (no time weighting). + _, loss_per_instance = self._compute_flow_matching_loss( + pred=out_fake["preds_vision"], # type: ignore[arg-type] + target=gen_data_noised_critic.vt_target_vision, # type: ignore[arg-type] + condition_mask=packed_student.vision.condition_mask, # type: ignore[union-attr] + timesteps=timesteps_critic, + has_valid_tokens=True, + rectified_flow=rectified_flow, + normalize_by_active=True, + ) + elif fake_score_loss_reduction == "sum_rcm": + _, loss_per_instance = self._flow_matching_sum_rcm_loss( + pred=out_fake["preds_vision"], # type: ignore[arg-type] + target=gen_data_noised_critic.vt_target_vision, # type: ignore[arg-type] + condition_mask=packed_student.vision.condition_mask, # type: ignore[union-attr] + has_valid_tokens=True, + ) + else: + raise ValueError(f"Unknown fake-score loss reduction: {fake_score_loss_reduction}") + loss_fake_score = loss_per_instance.mean() * self.config.loss_scale_fake_score + total_loss = loss_fake_score + + + # Flow-matching loss — sound (when configured) + if self.config.sound_gen: + if not ( + gen_data_student.x0_tokens_sound is not None + and gen_data_noised_critic.xt_tokens_sound is not None + and packed_student.sound is not None + ): + # FSDP dummy: keep fake-score sound params in backward graph + total_loss = total_loss + 0.0 * sum(p.sum() for p in out_fake["preds_sound"]) # type: ignore[union-attr] + + log.info( + f"Iteration {iteration}, critic phase, after flow_matching_loss, " + f"fake_score_loss: {loss_fake_score.item():.6f}, total_loss: {total_loss.item():.6f}" + ) + + output_batch: dict[str, torch.Tensor] = { + "fake_score_loss": loss_fake_score.detach(), + "total_critic_loss": total_loss.detach(), + "sigma": sigmas_critic.detach(), + "flow_matching_loss_vision_per_instance": loss_per_instance.detach(), # [B] + # Aliases for the WandB callback (wandb_log_rcm.py). + "dmd_loss_critic": total_loss.detach(), + "dmd_loss": loss_fake_score.detach(), + } + return output_batch, total_loss diff --git a/cosmos_framework/model/generator/distillation/dmd2_rf_test.py b/cosmos_framework/model/generator/distillation/dmd2_rf_test.py new file mode 100644 index 00000000..bdd86c83 --- /dev/null +++ b/cosmos_framework/model/generator/distillation/dmd2_rf_test.py @@ -0,0 +1,1820 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Unit tests for DMD2RFModel.""" + +import inspect +from types import ( + SimpleNamespace, +) +from typing import Optional +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from cosmos_framework.utils.flags import Device +from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.distillation.common_loss import variational_score_distillation_loss +from cosmos_framework.model.generator.distillation import dmd2_rf as dmd2_rf_module +from cosmos_framework.model.generator.distillation.dmd2_rf import DMD2RFModel + + + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_vision_only_sequence_plan_guard_accepts_vision_batch() -> None: + sequence_plans = [SimpleNamespace(has_action=False, has_sound=False)] + + DMD2RFModel._validate_vision_only_sequence_plans(sequence_plans) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize(("has_action", "has_sound"), [(True, False), (False, True)]) +def test_vision_only_sequence_plan_guard_rejects_nonvision_batch(has_action: bool, has_sound: bool) -> None: + sequence_plans = [SimpleNamespace(has_action=has_action, has_sound=has_sound)] + + with pytest.raises(ValueError, match="only vision T2I/I2V batches"): + DMD2RFModel._validate_vision_only_sequence_plans(sequence_plans) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_training_only_network_holders_are_not_module_state() -> None: + model = object.__new__(DMD2RFModel) + torch.nn.Module.__init__(model) + model.net = torch.nn.Linear(2, 2) + model._net_teacher_holder = [torch.nn.Linear(2, 2)] + model._net_fake_score_holder = [torch.nn.Linear(2, 2)] + + state_keys = set(torch.nn.Module.state_dict(model)) + + assert state_keys == {"net.weight", "net.bias"} + assert model.net_teacher is model._net_teacher_holder[0] + assert model.net_fake_score is model._net_fake_score_holder[0] + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize( + ("iteration", "expected_phase", "expected_optimizer_key"), + [ + (0, "student", "net"), + (1, "critic", "fake_score"), + (4, "critic", "fake_score"), + (5, "student", "net"), + ], +) +def test_public_phase_contract_routes_student_and_critic_optimizers( + iteration: int, + expected_phase: str, + expected_optimizer_key: str, +) -> None: + model = object.__new__(DMD2RFModel) + model.config = SimpleNamespace( + warmup_student_steps=0, + warmup_critic_steps=0, + student_update_freq=5, + ) + + assert model.get_phase(iteration) == expected_phase + assert model.get_optimizer_key(iteration) == expected_optimizer_key + + +@pytest.mark.L0 +def test_is_subclass_of_omni_mot_model(): + assert issubclass(DMD2RFModel, OmniMoTModel) + + +@pytest.mark.L0 +def test_does_not_inherit_removed_legacy_wrapper(): + # The legacy wrapper is not part of the public DMD2RFModel hierarchy. + mro_names = [cls.__name__ for cls in DMD2RFModel.__mro__] + assert "Cosmos3InteractiveModel" not in mro_names + + +@pytest.mark.L0 +def test_clip_grad_norm_defined_in_dmd2(): + # clip_grad_norm_ must be defined directly on DMD2RFModel, not inherited from OmniMoTModel. + assert "clip_grad_norm_" in DMD2RFModel.__dict__ + + +@pytest.mark.L0 +def test_on_before_zero_grad_defined_in_dmd2(): + # on_before_zero_grad must be defined directly on DMD2RFModel, not inherited from OmniMoTModel. + assert "on_before_zero_grad" in DMD2RFModel.__dict__ + + +@pytest.mark.L0 +def test_clip_grad_norm_skips_empty_params(): + """Empty param list (no grads) should return 0.0 tensor, not raise an error.""" + model = MagicMock(spec=DMD2RFModel) + model.config = MagicMock() + model.config.grad_clip = True + model.net = MagicMock() + model.net_fake_score = MagicMock() + model.net_fake_score.parameters.return_value = iter([]) + model.net_discriminator_head = None + + # A param with no grad attached + param = torch.nn.Parameter(torch.randn(4)) + param.grad = None + model.net.parameters.return_value = iter([param]) + + result = DMD2RFModel.clip_grad_norm_(model, max_norm=1.0) + + assert isinstance(result, torch.Tensor) + assert result.item() == 0.0 + + +@pytest.mark.L0 +def test_clip_grad_norm_cleans_nan_grads(): + """NaN gradients should be replaced with 0 before clipping.""" + p = torch.nn.Parameter(torch.randn(4)) + p.grad = torch.full((4,), float("nan")) + torch.nan_to_num(p.grad, nan=0.0, posinf=0.0, neginf=0.0, out=p.grad) + assert not torch.isnan(p.grad).any() + + +@pytest.mark.L0 +def test_on_before_zero_grad_student_phase_calls_ema_update(): + """In student phase, net_ema_worker.update_average should be called.""" + model = MagicMock() + model.get_phase.return_value = "student" + model.config.ema.enabled = True + model.get_student_iteration.return_value = 100 + model.ema_beta.return_value = 0.999 + + optimizer = MagicMock() + scheduler = MagicMock() + + DMD2RFModel.on_before_zero_grad(model, optimizer, scheduler, iteration=5) + + model.net_ema_worker.update_average.assert_called_once_with(model.net, model.net_ema, beta=0.999) + + +@pytest.mark.L0 +def test_on_before_zero_grad_critic_phase_skips_ema(): + """In critic phase (not student), EMA update should NOT be called.""" + model = MagicMock() + model.get_phase.return_value = "critic" + model.net_fake_score = None + model.net_discriminator_head = None + + optimizer = MagicMock() + optimizer.get.return_value = None + scheduler = MagicMock() + + DMD2RFModel.on_before_zero_grad(model, optimizer, scheduler, iteration=5) + + model.net_ema_worker.update_average.assert_not_called() + + +@pytest.mark.L0 +def test_model_dict_includes_fake_score(): + model = MagicMock(spec=DMD2RFModel) + model.net = MagicMock() + model.net_fake_score = MagicMock() + model.net_discriminator_head = None + + result = DMD2RFModel.model_dict(model) + + assert "net" in result + assert result["net"] is model.net + assert "fake_score" in result + assert result["fake_score"] is model.net_fake_score + assert "discriminator" not in result + + +@pytest.mark.L0 +def test_dmd2_rf_needs_fake_score_by_default(): + model = MagicMock(spec=DMD2RFModel) + + assert DMD2RFModel._needs_fake_score(model) is True + + +@pytest.mark.L0 +def test_model_dict_includes_discriminator_when_set(): + model = MagicMock(spec=DMD2RFModel) + model.net = MagicMock() + model.net_fake_score = MagicMock() + model.net_discriminator_head = MagicMock() # truthy + + result = DMD2RFModel.model_dict(model) + + assert "discriminator" in result + assert result["discriminator"] is model.net_discriminator_head + + +@pytest.mark.L0 +def test_fixed_step_sampler_set_in_set_up_model(): + """Verify set_up_model creates self.fixed_step_sampler through the setup helper.""" + from cosmos_framework.model.generator.diffusion.samplers.fixed_step import FixedStepSampler # noqa: F401 + + setup_src = inspect.getsource(DMD2RFModel.set_up_model) + fixed_step_setup_src = inspect.getsource(DMD2RFModel._set_up_fixed_step_sampler) + + assert "_set_up_fixed_step_sampler()" in setup_src + assert "FixedStepSampler" in fixed_step_setup_src + assert "fixed_step_sampler" in fixed_step_setup_src + + +@pytest.mark.L0 +def test_training_only_nets_stay_out_of_inference_state_dict(): + """Teacher/fake-score should stay unregistered and be skipped for inference setup.""" + init_src = inspect.getsource(DMD2RFModel.__init__) + setup_src = inspect.getsource(DMD2RFModel.set_up_model) + + assert "_net_teacher_holder" in init_src + assert "_net_fake_score_holder" in init_src + assert "if is_inference_mode:" in setup_src + assert 'self.denoiser_nets = {"student": self.net}' in setup_src + assert "needs_fake_score = self._needs_fake_score()" in setup_src + assert "_net_teacher_holder = [self.build_net(self.precision, lora_enabled=False)]" in setup_src + assert "_net_fake_score_holder = [self.build_net(self.precision)]" in setup_src + assert 'self.denoiser_nets["fake_score"] = self.net_fake_score' in setup_src + assert "_build_net_with_lora" not in DMD2RFModel.__dict__ + assert "state_dict" not in DMD2RFModel.__dict__ + + +@pytest.mark.L0 +@pytest.mark.parametrize("teacher_negative_prompt", ["", "low quality video"]) +def test_teacher_cfg_negative_prompt_is_tokenized_for_uncond_teacher_pass(teacher_negative_prompt: str) -> None: + model = MagicMock() + model.config = SimpleNamespace( + action_gen=False, + loss_scale_sid=1.0, + sound_gen=False, + teacher_guidance=3.0, + teacher_negative_prompt=teacher_negative_prompt, + vsd_gradient_space="x0", + vsd_loss_reduction="mean", + ) + model.parallel_dims = None + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.vlm_config = SimpleNamespace(use_system_prompt=False) + model.vlm_tokenizer = object() + + gen_data_student = SimpleNamespace( + batch_size=1, + is_image_batch=False, + x0_tokens_vision=[torch.full((1, 1, 1, 1), 0.2)], + x0_tokens_action=None, + x0_tokens_sound=None, + ) + packed_student = SimpleNamespace( + action=None, + sound=None, + vision=SimpleNamespace(condition_mask=[torch.zeros(1, 1, 1)]), + ) + gen_data_noised = SimpleNamespace( + sigmas_vision=[torch.full((1, 1, 1), 0.5)], + xt_tokens_vision=[torch.full((1, 1, 1, 1), 0.4)], + ) + model._gen_data_from_student.return_value = ( + gen_data_student, + packed_student, + torch.zeros(1, 1), + {"preds_vision": [torch.zeros(1, 1, 1, 1)]}, + ) + model._get_train_noise_level_vision.return_value = (torch.full((1, 1), 0.5), torch.full((1, 1), 0.5)) + model._action_sigmas_from_full.return_value = None + model._sound_sigmas_from_full.return_value = None + model._add_noise_to_input.return_value = gen_data_noised + model._pack_and_denoise.side_effect = [ + {"preds_vision": [torch.full((1, 1, 1, 1), 0.1)]}, + {"preds_vision": [torch.full((1, 1, 1, 1), 0.2)]}, + {"preds_vision": [torch.full((1, 1, 1, 1), 0.05)]}, + ] + model._velocity_to_x0.side_effect = [ + [torch.full((1, 1, 1, 1), 0.15)], + [torch.full((1, 1, 1, 1), 0.3)], + [torch.full((1, 1, 1, 1), 0.1)], + ] + + tokenized_negative_prompt = [len(teacher_negative_prompt), 1, 0] + + def fake_tokenize_caption( + prompt: str, + tokenizer: object, + *, + is_video: bool, + use_system_prompt: bool, + ) -> list[int]: + assert tokenizer is model.vlm_tokenizer + assert is_video is True + assert use_system_prompt is False + return [len(prompt), int(is_video), int(use_system_prompt)] + + with patch.object(dmd2_rf_module, "tokenize_caption", side_effect=fake_tokenize_caption) as tokenize_mock: + output_batch, loss = DMD2RFModel.training_step_generator( + model, + input_text_indexes=[[101]], + sequence_plans=[MagicMock()], + gen_data_clean=SimpleNamespace(batch_size=1), + data_resolutions=["480"], + num_vision_tokens_per_sample=[1], + iteration=7, + ) + + assert torch.isfinite(loss) + assert torch.isfinite(output_batch["total_generator_loss"]) + tokenize_mock.assert_called_once() + assert tokenize_mock.call_args.args[0] == teacher_negative_prompt + + pack_calls = model._pack_and_denoise.call_args_list + assert [call.kwargs["net_type"] for call in pack_calls] == ["fake_score", "teacher", "teacher"] + assert pack_calls[0].args[3] == [[101]] + assert pack_calls[1].args[3] == [[101]] + assert pack_calls[2].args[3] == [tokenized_negative_prompt] + + +@pytest.mark.L0 +def test_build_net_accepts_optional_lora_override() -> None: + signature = inspect.signature(OmniMoTModel.build_net) + + lora_enabled = signature.parameters["lora_enabled"] + assert lora_enabled.kind is inspect.Parameter.KEYWORD_ONLY + assert lora_enabled.default is None + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("config_lora_enabled", "lora_enabled_override", "expected_lora_enabled"), + [ + (True, None, True), + (True, False, False), + (False, True, True), + ], +) +def test_build_net_lora_override_controls_injection_and_initialization( + config_lora_enabled: bool, + lora_enabled_override: bool | None, + expected_lora_enabled: bool, +) -> None: + model = MagicMock(spec=OmniMoTModel) + model.config = MagicMock() + model.config.lora_enabled = config_lora_enabled + model.config.lora_rank = 8 + model.config.lora_alpha = 16 + model.config.lora_target_modules = "q_proj_moe_gen" + model.vlm_config = MagicMock() + model.tokenizer_vision_gen = MagicMock() + model.parallel_dims = MagicMock() + + language_model = MagicMock() + net = MagicMock() + net.to.return_value = net + model.add_lora.return_value = net + + module = "cosmos_framework.model.generator.omni_mot_model" + with ( + patch(f"{module}.lazy_instantiate", return_value=language_model), + patch(f"{module}.Cosmos3VFMNetworkConfig"), + patch(f"{module}.Cosmos3VFMNetwork", return_value=net), + patch(f"{module}.parallelize_vfm_network", return_value=net), + patch(f"{module}.DEVICE", Device.CUDA), + ): + if lora_enabled_override is None: + result = OmniMoTModel.build_net(model, torch.float32) + else: + result = OmniMoTModel.build_net(model, torch.float32, lora_enabled=lora_enabled_override) + + assert result is net + if expected_lora_enabled: + model.add_lora.assert_called_once_with( + net, + lora_rank=8, + lora_alpha=16, + lora_target_modules="q_proj_moe_gen", + ) + model._init_lora_weights_post_materialization.assert_called_once_with(net) + else: + model.add_lora.assert_not_called() + model._init_lora_weights_post_materialization.assert_not_called() + + +@pytest.mark.L0 +def test_copy_teacher_weights_allows_missing_lora_adapter_keys(): + class NetWithOptionalLoRA(torch.nn.Module): + def __init__(self, with_lora: bool): + super().__init__() + self.proj = torch.nn.Linear(2, 2, bias=False) + if with_lora: + self.proj.lora_A = torch.nn.Linear(2, 1, bias=False) + self.proj.lora_B = torch.nn.Linear(1, 2, bias=False) + + teacher = NetWithOptionalLoRA(with_lora=False) + target = NetWithOptionalLoRA(with_lora=True) + with torch.no_grad(): + teacher.proj.weight.fill_(3.0) + target.proj.lora_A.weight.fill_(1.0) + target.proj.lora_B.weight.fill_(2.0) + + model = MagicMock(spec=DMD2RFModel) + type(model).net_teacher = property(lambda self: teacher) + + DMD2RFModel._copy_teacher_weights(model, target_net=target, target_name="fake score") + + assert torch.equal(target.proj.weight, teacher.proj.weight) + assert torch.equal(target.proj.lora_A.weight, torch.ones_like(target.proj.lora_A.weight)) + assert torch.equal(target.proj.lora_B.weight, torch.full_like(target.proj.lora_B.weight, 2.0)) + + +# --------------------------------------------------------------------------- +# _pack_and_denoise proxy forwarding +# --------------------------------------------------------------------------- + + +def _make_gen_data( + *, + with_sound: bool = False, + num_vision_items: Optional[list[int]] = None, +) -> MagicMock: + """Build a minimal GenerationDataClean mock for proxy-forwarding tests.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean + + gd = MagicMock(spec=GenerationDataClean) + gd.batch_size = 2 + gd.is_image_batch = False + gd.raw_state_vision = [MagicMock(), MagicMock()] + gd.x0_tokens_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd.fps_vision = [24.0, 24.0] + gd.num_vision_items_per_sample = num_vision_items or [1, 1] + if with_sound: + gd.raw_state_sound = [MagicMock(), MagicMock()] + gd.x0_tokens_sound = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd.fps_sound = [24.0, 24.0] + else: + gd.raw_state_sound = None + gd.x0_tokens_sound = None + gd.fps_sound = None + gd.raw_state_action = None + gd.x0_tokens_action = None + gd.fps_action = None + gd.action_domain_id = None + return gd + + +def _make_gen_data_noised( + *, + with_sound: bool = False, + with_action: bool = False, +) -> MagicMock: + """Build a minimal GenerationDataNoised mock.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + gd = MagicMock(spec=GenerationDataNoised) + gd.xt_tokens_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd.vt_target_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd.sigmas_vision = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] + if with_sound: + gd.xt_tokens_sound = [torch.zeros(2, 1), torch.zeros(2, 1)] + gd.vt_target_sound = [torch.zeros(2, 1), torch.zeros(2, 1)] + gd.sigmas_sound = [torch.zeros(1, 1), torch.zeros(1, 1)] + else: + gd.xt_tokens_sound = None + gd.vt_target_sound = None + gd.sigmas_sound = None + if with_action: + gd.xt_tokens_action = [torch.zeros(1, 4), torch.zeros(1, 4)] + gd.vt_target_action = [torch.zeros(1, 4), torch.zeros(1, 4)] + gd.sigmas_action = [torch.zeros(1, 1), torch.zeros(1, 1)] + else: + gd.xt_tokens_action = None + gd.vt_target_action = None + gd.sigmas_action = None + return gd + + +def _make_packed_sequence(*, with_action: bool = False) -> MagicMock: + """Build a minimal PackedSequence mock.""" + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.tokens = [] + packed.vision.condition_mask = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] + if with_action: + packed.action = MagicMock() + packed.action.tokens = [] + packed.action.condition_mask = [torch.zeros(1, 1), torch.zeros(1, 1)] + else: + packed.action = None + packed.sound = None + return packed + + + + +@pytest.mark.L0 +def test_pack_and_denoise_returns_dict(): + """_pack_and_denoise must return a plain dict, not a tuple.""" + model = MagicMock() + model._pack_input_sequence.return_value = _make_packed_sequence() + model.denoise.return_value = {"preds_vision": [torch.zeros(4, 1, 2, 2)]} + + gen_data = _make_gen_data() + gen_data_noised = _make_gen_data_noised() + + result = DMD2RFModel._pack_and_denoise( + model, + gen_data_clean=gen_data, + gen_data_noised=gen_data_noised, + timesteps=torch.zeros(2, 1), + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + net_type="student", + ) + + assert isinstance(result, dict), f"Expected dict, got {type(result)}" + + + + +@pytest.mark.L0 +def test_pack_and_denoise_forwards_num_vision_items(): + """num_vision_items_per_sample must be forwarded from gen_data_clean in the proxy.""" + model = MagicMock() + model._pack_input_sequence.return_value = _make_packed_sequence() + model.denoise.return_value = {"preds_vision": []} + + gen_data = _make_gen_data(num_vision_items=[2, 2]) + gen_data_noised = _make_gen_data_noised() + + DMD2RFModel._pack_and_denoise( + model, + gen_data_clean=gen_data, + gen_data_noised=gen_data_noised, + timesteps=torch.zeros(2, 1), + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + net_type="student", + ) + + _, call_kwargs = model._pack_input_sequence.call_args + proxy = call_kwargs.get("gen_data_clean") or model._pack_input_sequence.call_args.args[2] + assert proxy.num_vision_items_per_sample == [2, 2] + + + + + + +@pytest.mark.L0 +def test_gen_data_from_student_forwards_num_vision_items(): + """num_vision_items_per_sample must be forwarded in the proxy inside _forward_simulation.""" + model = MagicMock() + model.tensor_kwargs = {"dtype": torch.float32, "device": "cpu"} + model.config.rectified_flow_inference_config.num_train_timesteps = 1000 + + model.config.action_gen = False + model.config.sound_gen = False + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model._sample_student_sigma.return_value = torch.full((2, 1), 0.5) + + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.condition_mask = [torch.ones(1, 1, 1), torch.ones(1, 1, 1)] + packed.action = None + packed.sound = None + model._pack_input_sequence.return_value = packed + + gen_data_noised = MagicMock() + gen_data_noised.xt_tokens_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gen_data_noised.xt_tokens_action = None + gen_data_noised.xt_tokens_sound = None + model._add_noise_to_input.return_value = gen_data_noised + model.denoise.return_value = {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + model._velocity_to_x0.return_value = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + + gen_data = _make_gen_data(num_vision_items=[2, 2]) + + DMD2RFModel._forward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + ) + + _, call_kwargs = model._pack_input_sequence.call_args + proxy = call_kwargs.get("gen_data_clean") or model._pack_input_sequence.call_args.args[2] + assert proxy.num_vision_items_per_sample == [2, 2] + + +# --------------------------------------------------------------------------- +# _gen_data_from_student return type and multi-modal x0 extraction +# --------------------------------------------------------------------------- + + +def _make_student_model(*, action_gen: bool = False, sound_gen: bool = False) -> MagicMock: + """Build a model mock wired up for _gen_data_from_student tests.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + model = MagicMock() + model.tensor_kwargs = {"dtype": torch.float32, "device": "cpu"} + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.rectified_flow_inference_config.num_train_timesteps = 1000 + model.config.action_gen = action_gen + model.config.sound_gen = sound_gen + model._sample_student_sigma.return_value = torch.full((2, 1), 0.5) + model._action_sample_indices.side_effect = DMD2RFModel._action_sample_indices + model._action_sigmas_from_full.side_effect = ( + lambda sigmas_full, sequence_plans: DMD2RFModel._action_sigmas_from_full(model, sigmas_full, sequence_plans) + ) + + packed = _make_packed_sequence(with_action=action_gen) + if action_gen: + packed.action.condition_mask = [torch.zeros(1, 1), torch.zeros(1, 1)] + model._pack_input_sequence.return_value = packed + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.xt_tokens_action = [torch.zeros(1, 4), torch.zeros(1, 4)] if action_gen else None + gd_noised.xt_tokens_sound = None + model._add_noise_to_input.return_value = gd_noised + + preds = {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + if action_gen: + preds["preds_action"] = [torch.zeros(1, 4), torch.zeros(1, 4)] + model.denoise.return_value = preds + model._velocity_to_x0.return_value = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + return model + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("is_image_batch", [True, False], ids=["t2i", "i2v"]) +def test_vision_only_forward_simulation_supports_t2i_and_i2v(is_image_batch: bool) -> None: + model = _make_student_model(action_gen=False, sound_gen=False) + model._sound_sigmas_from_full.return_value = None + gen_data = _make_gen_data() + gen_data.is_image_batch = is_image_batch + + gen_data_student, _, _, out_student = DMD2RFModel._forward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + ) + + assert gen_data_student.is_image_batch is is_image_batch + assert gen_data_student.x0_tokens_vision is not None + assert gen_data_student.x0_tokens_action is None + assert gen_data_student.x0_tokens_sound is None + assert set(out_student) == {"preds_vision"} + + +@pytest.mark.L0 +def test_gen_data_from_student_returns_4_tuple(): + """_forward_simulation must return (GenerationDataClean, PackedSequence, Tensor, dict).""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean + + model = _make_student_model() + gen_data = _make_gen_data() + + result = DMD2RFModel._forward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + ) + + assert isinstance(result, tuple) and len(result) == 4 + gen_data_student, _, sigma, out = result + assert isinstance(gen_data_student, GenerationDataClean) + assert isinstance(sigma, torch.Tensor) and sigma.shape == (2, 1) + assert isinstance(out, dict) + + + + +# --------------------------------------------------------------------------- +# training_step_critic: normalize_by_active and FSDP dummy +# --------------------------------------------------------------------------- + + +def _make_critic_model(*, action_gen: bool = False) -> MagicMock: + """Build a model mock wired up for training_step_critic tests.""" + model = MagicMock() + model.parallel_dims = None # disable CP broadcast path in tests + model.config.action_gen = action_gen + model.config.sound_gen = False + model.config.loss_scale_fake_score = 1.0 + model.config.fake_score_loss_reduction = "active_mean" + + gen_data_student = _make_gen_data() + gen_data_student.x0_tokens_action = None # absent by default + packed_student = _make_packed_sequence() + model._gen_data_from_student.return_value = (gen_data_student, packed_student, torch.zeros(2, 1), {}) + + model._get_train_noise_level_vision.return_value = (torch.zeros(2, 1), torch.zeros(2, 1)) + model._add_noise_to_input.return_value = _make_gen_data_noised() + + out_fake = {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + if action_gen: + out_fake["preds_action"] = [torch.zeros(1, 4), torch.zeros(1, 4)] + model._pack_and_denoise.return_value = out_fake + + model._compute_flow_matching_loss.return_value = (torch.tensor(0.5), torch.zeros(2)) + return model + + +@pytest.mark.L0 +def test_training_step_critic_calls_compute_flow_matching_loss_with_normalize_by_active(): + """_compute_flow_matching_loss must be called with normalize_by_active=True in critic.""" + model = _make_critic_model() + + DMD2RFModel.training_step_critic( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=MagicMock(), + data_resolutions=["480", "480"], + num_vision_tokens_per_sample=[4, 4], + iteration=0, + ) + + noise_call_kwargs = model._get_train_noise_level_vision.call_args.kwargs + assert noise_call_kwargs["resolutions"] == ["480", "480"] + assert noise_call_kwargs["num_tokens"] == [4, 4] + + model._compute_flow_matching_loss.assert_called_once() + call_kwargs = model._compute_flow_matching_loss.call_args.kwargs + assert call_kwargs.get("normalize_by_active") is True + + + + +@pytest.mark.L0 +def test_flow_matching_sum_rcm_loss_sums_active_elements_per_instance() -> None: + """sum_rcm should sum generated elements per instance, then average the batch.""" + pred = [torch.ones(1, 2, 1, 1), torch.full((1, 2, 1, 1), 2.0)] # list of [C,T,H,W] + target = [torch.zeros_like(pred[0]), torch.ones_like(pred[1])] # list of [C,T,H,W] + condition_mask = [ + torch.tensor([[[0.0]], [[1.0]]]), + torch.tensor([[[0.0]], [[0.0]]]), + ] # list of [T,1,1] + + loss, per_instance = DMD2RFModel._flow_matching_sum_rcm_loss( + pred=pred, + target=target, + condition_mask=condition_mask, + has_valid_tokens=True, + ) + + expected_per_instance = torch.tensor([1.0, 2.0]) # [B] + assert torch.allclose(per_instance, expected_per_instance) + assert loss.item() == pytest.approx(1.5) + + +# --------------------------------------------------------------------------- +# _sample_student_sigma +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_sample_student_sigma_shape(): + """Return shape must be (B, 1) regardless of batch size or t_list length.""" + model = MagicMock() + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.fixed_step_sampler_config.t_list = [1.0, 0.9, 0.75, 0.5] + + for B in [1, 4]: + sigma = DMD2RFModel._sample_student_sigma(model, B) + assert sigma.shape == (B, 1), f"Expected ({B}, 1), got {sigma.shape}" + + +@pytest.mark.L0 +def test_sample_student_sigma_single_step_always_same_value(): + """Single-entry t_list: every sample in the batch gets that exact sigma.""" + model = MagicMock() + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.fixed_step_sampler_config.t_list = [1.0] + + sigma = DMD2RFModel._sample_student_sigma(model, batch_size=4) + assert sigma.shape == (4, 1) + assert (sigma == 1.0).all() + + +@pytest.mark.L0 +def test_sample_student_sigma_multi_step_values_in_t_list(): + """Multi-entry t_list: all returned sigmas must be one of the listed values.""" + t_list = [1.0, 0.9, 0.75, 0.5] + model = MagicMock() + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.fixed_step_sampler_config.t_list = t_list + + sigma = DMD2RFModel._sample_student_sigma(model, batch_size=100) + for s in sigma.squeeze(1).tolist(): + assert any(abs(s - t) < 1e-6 for t in t_list), f"sigma {s} not in t_list {t_list}" + + +@pytest.mark.L0 +def test_sample_student_sigma_multi_step_covers_all_values(): + """Multi-entry t_list with large batch: all t_list values should appear at least once.""" + t_list = [1.0, 0.9, 0.75, 0.5] + model = MagicMock() + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.fixed_step_sampler_config.t_list = t_list + + # With 1000 samples, the probability of missing any value is negligible + sigma = DMD2RFModel._sample_student_sigma(model, batch_size=1000) + unique_vals = set(round(s, 6) for s in sigma.squeeze(1).tolist()) + for t in t_list: + assert any(abs(t - v) < 1e-5 for v in unique_vals), f"{t} never sampled" + + +@pytest.mark.L0 +def test_sample_student_sigma_returns_float32(): + """Sigma tensor should be float32 (consistent with tensor_kwargs_fp32).""" + model = MagicMock() + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.config.fixed_step_sampler_config.t_list = [1.0] + + sigma = DMD2RFModel._sample_student_sigma(model, batch_size=2) + assert sigma.dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# variational_score_distillation_loss with loss_mask +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_vsd_loss_no_mask(): + """Without mask the loss matches the original formula (non-regression).""" + torch.manual_seed(0) + B, C, T, H, W = 2, 4, 3, 2, 2 + gen_data = [torch.randn(C, T, H, W, requires_grad=True) for _ in range(B)] + teacher_x0 = [torch.randn(C, T, H, W) for _ in range(B)] + fake_x0 = [torch.randn(C, T, H, W) for _ in range(B)] + + loss = variational_score_distillation_loss(gen_data, teacher_x0, fake_x0) + assert loss.shape == () + assert loss.item() >= 0.0 + # Gradient must flow back to gen_data + loss.backward() + assert gen_data[0].grad is not None + + +@pytest.mark.L0 +def test_vsd_loss_with_mask_zeros_conditioned(): + """Conditioned frames (mask=0) should contribute zero loss. + + When mask=0 for all frames, pseudo_target == gen_data for the masked + frames only if the loss is correctly zeroed. We verify by checking that + a mask of all-zeros yields zero loss (no generated frames to fit). + """ + torch.manual_seed(1) + B, C, T, H, W = 1, 4, 4, 2, 2 + gen_data = [torch.randn(C, T, H, W, requires_grad=True)] + teacher_x0 = [torch.randn(C, T, H, W)] + fake_x0 = [torch.randn(C, T, H, W)] + + # All frames conditioned (mask = 0 everywhere) + loss_mask = [torch.zeros(T, 1, 1)] + loss = variational_score_distillation_loss(gen_data, teacher_x0, fake_x0, loss_mask=loss_mask) + assert loss.item() == pytest.approx(0.0, abs=1e-6), f"Expected zero loss when mask is all-zeros, got {loss.item()}" + + +@pytest.mark.L0 +def test_vsd_loss_weight_unaffected_by_conditioned_frames(): + """With loss_mask, weight is computed only over generated frames. + + When conditioned frames have gen ≈ teacher (near-zero diff) and generated + frames have a non-trivial diff, the unmasked version produces an inflated + weight (small diff_abs_mean → large w). The masked version should yield + a weight based solely on the generated-frame differences, which is larger + (non-trivial diff → smaller w). + + We verify that masked and unmasked losses differ when conditioned frames + have small diff and generated frames have large diff. + """ + torch.manual_seed(2) + B, C, T, H, W = 1, 4, 4, 2, 2 + + gen_data_val = torch.zeros(C, T, H, W) + teacher_val = torch.zeros(C, T, H, W) + fake_val = torch.randn(C, T, H, W) + + # Frames 0-1: conditioned (small gen-teacher diff by construction) + # Frames 2-3: generated (large gen-teacher diff) + gen_data_val[:, 2:, :, :] = 10.0 + teacher_val[:, 2:, :, :] = -10.0 # large diff on generated frames + + gen = [gen_data_val.clone().requires_grad_(True)] + tea = [teacher_val] + fake = [fake_val] + + mask_generated = torch.zeros(T, 1, 1) + mask_generated[2:] = 1.0 # only frames 2-3 are generated + + loss_masked = variational_score_distillation_loss(gen, tea, fake, loss_mask=[mask_generated]) + # Now all frames (including near-zero diff frames 0-1) inflate the denom + gen2 = [gen_data_val.clone().requires_grad_(True)] + loss_unmasked = variational_score_distillation_loss(gen2, tea, fake) + + # Masked and unmasked losses should differ when conditioned frames have small diff + assert loss_masked.item() != pytest.approx(loss_unmasked.item(), rel=0.01), ( + "Masked and unmasked losses should differ when conditioned frames have small diff" + ) + + +# --------------------------------------------------------------------------- +# simulation_mode config fields +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_simulation_mode_default_is_forward(): + """DMD2RFConfig.simulation_mode defaults to 'forward'.""" + from cosmos_framework.configs.base.experiment.distillation.dmd2_config import DMD2RFConfig + + assert DMD2RFConfig().simulation_mode == "forward" + + +@pytest.mark.L0 +def test_simulation_mode_accepts_backward(): + """DMD2RFConfig accepts simulation_mode='backward' without error.""" + import attrs + + from cosmos_framework.configs.base.experiment.distillation.dmd2_config import DMD2RFConfig + + cfg = attrs.evolve(DMD2RFConfig(), simulation_mode="backward") + assert cfg.simulation_mode == "backward" + + +@pytest.mark.L0 +def test_backward_grad_steps_default(): + """DMD2RFConfig.backward_grad_steps defaults to 1.""" + from cosmos_framework.configs.base.experiment.distillation.dmd2_config import DMD2RFConfig + + assert DMD2RFConfig().backward_grad_steps == 1 + + +@pytest.mark.L0 +def test_forward_simulation_defined(): + """_forward_simulation must be defined directly on DMD2RFModel.""" + assert "_forward_simulation" in DMD2RFModel.__dict__ + + +@pytest.mark.L0 +def test_backward_simulation_defined(): + """_backward_simulation must be defined directly on DMD2RFModel.""" + assert "_backward_simulation" in DMD2RFModel.__dict__ + + +# --------------------------------------------------------------------------- +# _gen_data_from_student dispatch +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_gen_data_from_student_dispatches_forward(): + """_gen_data_from_student must call _forward_simulation when simulation_mode='forward'.""" + model = MagicMock() + model.config.simulation_mode = "forward" + model._forward_simulation.return_value = ("gd", "packed", torch.zeros(2, 1), {}) + + DMD2RFModel._gen_data_from_student( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=MagicMock(), + iteration=0, + ) + + model._forward_simulation.assert_called_once() + model._backward_simulation.assert_not_called() + + +@pytest.mark.L0 +def test_gen_data_from_student_dispatches_backward(): + """_gen_data_from_student must call _backward_simulation when simulation_mode='backward'.""" + model = MagicMock() + model.config.simulation_mode = "backward" + model._backward_simulation.return_value = ("gd", "packed", torch.zeros(2, 1), {}) + + DMD2RFModel._gen_data_from_student( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=MagicMock(), + iteration=0, + ) + + model._backward_simulation.assert_called_once() + model._forward_simulation.assert_not_called() + + +# --------------------------------------------------------------------------- +# _backward_simulation helper tests +# --------------------------------------------------------------------------- + + +def _make_backward_model( + *, + t_list: Optional[list[float]] = None, + sample_type: str = "ode", + backward_grad_steps: int = 1, + sound_gen: bool = False, +) -> tuple[MagicMock, int]: + """Build a minimal mock for _backward_simulation tests (B=2, vision-only). + + ``_backward_n_steps`` is mocked to return the full schedule length by default, so the + rollout is deterministic; tests that exercise rcm-style variable length override its + ``side_effect``. + """ + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + if t_list is None: + t_list = [1.0, 0.5] # 2-step schedule -> full_t_list = [1.0, 0.5, 0.0] + + B = 2 + model = MagicMock() + model.tensor_kwargs = {"dtype": torch.float32, "device": "cpu"} + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.parallel_dims = None + model.config.rectified_flow_inference_config.num_train_timesteps = 1000 + model.config.action_gen = False + model.config.sound_gen = sound_gen + model.config.fixed_step_sampler_config.t_list = t_list + model.config.fixed_step_sampler_config.sample_type = sample_type + model.config.backward_grad_steps = backward_grad_steps + # Default: run the full schedule (n_steps == number of nonzero sigma levels). rcm-style + # variable-length tests override this side_effect. + model._backward_n_steps.side_effect = lambda max_n_steps, iteration: max_n_steps + + # All-zero condition mask = generation frames only + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.condition_mask = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] # [T=1,1,1] each + packed.action = None + packed.sound = None + model._pack_input_sequence.return_value = packed + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [torch.ones(4, 1, 2, 2), torch.ones(4, 1, 2, 2)] # [C,T,H,W] + gd_noised.xt_tokens_action = None + gd_noised.xt_tokens_sound = None + gd_noised.epsilon_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.vt_target_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.sigmas_vision = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] + gd_noised.epsilon_action = None + gd_noised.vt_target_action = None + gd_noised.sigmas_action = None + gd_noised.epsilon_sound = None + gd_noised.vt_target_sound = None + gd_noised.sigmas_sound = None + model._add_noise_to_input.return_value = gd_noised + + # Return zero velocity predictions by default + model._pack_and_denoise.return_value = {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + import types + + # Use real implementations so ODE/SDE steps and x0 = xt - sigma * v are exercised + # _velocity_to_x0 doesn't use self; bind the others so self.tensor_kwargs resolves + model._velocity_to_x0 = DMD2RFModel._velocity_to_x0 + model._ode_step = types.MethodType(DMD2RFModel._ode_step, model) + model._sde_step = types.MethodType(DMD2RFModel._sde_step, model) + + return model, B + + +@pytest.mark.L0 +def test_backward_simulation_vision_output_shape(): + """_backward_simulation returns a 4-tuple with x0_tokens_vision of the correct shape.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean + + model, B = _make_backward_model() + gen_data = _make_gen_data() + + result = DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert isinstance(result, tuple) and len(result) == 4 + gen_data_student, packed, sigma, out = result + assert isinstance(gen_data_student, GenerationDataClean) + assert gen_data_student.x0_tokens_vision is not None + assert len(gen_data_student.x0_tokens_vision) == B + assert gen_data_student.x0_tokens_vision[0].shape == (4, 1, 2, 2) # [C,T,H,W] + assert sigma.shape == (B, 1) # [B,1] + + + + +@pytest.mark.L0 +def test_backward_simulation_single_step_x0(): + """With a 1-step t_list, x0 must equal xt - sigma * v_pred (RF formula).""" + model, B = _make_backward_model(t_list=[1.0]) # full_t_list = [1.0, 0.0], n_steps=1 + + xt = torch.ones(4, 1, 2, 2) * 2.0 # [C,T,H,W] + v_pred = torch.ones(4, 1, 2, 2) * 3.0 # [C,T,H,W] + + # xt_tokens_vision is the initial noised (pure-noise) data at sigma=1.0 + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [xt, xt] + gd_noised.xt_tokens_action = None + gd_noised.xt_tokens_sound = None + gd_noised.epsilon_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.vt_target_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.sigmas_vision = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] + gd_noised.epsilon_action = None + gd_noised.vt_target_action = None + gd_noised.sigmas_action = None + gd_noised.epsilon_sound = None + gd_noised.vt_target_sound = None + gd_noised.sigmas_sound = None + model._add_noise_to_input.return_value = gd_noised + model._pack_and_denoise.return_value = {"preds_vision": [v_pred, v_pred]} + + gen_data = _make_gen_data() + result = DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + gen_data_student = result[0] + # x0 = xt - sigma * v_pred; sigma_eff = 1.0 * noisy_mask = 1.0 (generation frames) + expected_x0 = xt - 1.0 * v_pred # [C,T,H,W] + assert torch.allclose(gen_data_student.x0_tokens_vision[0], expected_x0, atol=1e-5) + + +@pytest.mark.L0 +def test_backward_simulation_last_step_only_grad(): + """With backward_grad_steps=1, only the last _pack_and_denoise call has grad enabled.""" + model, _ = _make_backward_model(backward_grad_steps=1) # 2-step t_list + + grad_states = [] + + def capture_grad(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + model._pack_and_denoise.side_effect = capture_grad + + gen_data = _make_gen_data() + with torch.enable_grad(): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert len(grad_states) == 2 + assert grad_states[0] is False, "Step 0 (early) should run under no_grad" + assert grad_states[1] is True, "Step 1 (last) should run with grad enabled" + + +@pytest.mark.L0 +def test_backward_simulation_no_grad_early_steps(): + """All but the last step must be wrapped in torch.no_grad.""" + model, _ = _make_backward_model(backward_grad_steps=1, t_list=[1.0, 0.9, 0.75]) + + grad_states = [] + + def capture_grad(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + model._pack_and_denoise.side_effect = capture_grad + + gen_data = _make_gen_data() + with torch.enable_grad(): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + # n_steps=3; backward_grad_steps=1; steps 0,1 → no_grad; step 2 → grad + assert len(grad_states) == 3 + assert all(not g for g in grad_states[:2]), "Steps 0-1 should be under no_grad" + assert grad_states[2] is True, "Step 2 (last) should have grad enabled" + + +@pytest.mark.L0 +def test_backward_simulation_bptt_all_steps(): + """With backward_grad_steps=-1, every _pack_and_denoise call has grad enabled.""" + model, _ = _make_backward_model(backward_grad_steps=-1) + + grad_states = [] + + def capture_grad(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + model._pack_and_denoise.side_effect = capture_grad + + gen_data = _make_gen_data() + with torch.enable_grad(): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert all(grad_states), "All steps should have grad enabled when backward_grad_steps=-1" + + +@pytest.mark.L0 +def test_backward_simulation_prefix_schedule_from_pure_noise(): + """An n_steps of 2 runs the schedule PREFIX from sigma=1.0 down to 0.""" + model, B = _make_backward_model(t_list=[1.0, 0.9, 0.75], backward_grad_steps=1) + # full_t_list = [1.0, 0.9, 0.75, 0.0]; n_steps=2 -> prefix schedule [1.0, 0.9, 0.0]. + model._backward_n_steps.side_effect = lambda max_n_steps, iteration: 2 + + timesteps_seen = [] + + def capture_timesteps(*args, **kwargs): + timesteps_seen.append(args[2].clone()) # [B,1] + return {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + model._pack_and_denoise.side_effect = capture_timesteps + + gen_data = _make_gen_data() + result = DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + # n_steps=2 -> denoise at sigma=1.0 (t=1000) then sigma=0.9 (t=900), descending to 0. + assert len(timesteps_seen) == 2 + assert torch.allclose(timesteps_seen[0], torch.full((B, 1), 1000.0)) # [B,1] + assert torch.allclose(timesteps_seen[1], torch.full((B, 1), 900.0)) # [B,1] + # Returned sigma_student is always the pure-noise start sigma (1.0), regardless of n_steps. + assert torch.allclose(result[2], torch.full((B, 1), 1.0)) # [B,1] + + +@pytest.mark.L0 +def test_backward_simulation_seeds_from_pure_noise(): + """The rollout seed must be built at sigma=1.0 (pure noise, no clean-data blending).""" + model, _ = _make_backward_model(t_list=[1.0, 0.9, 0.75], backward_grad_steps=1) + + gen_data = _make_gen_data() + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + # _add_noise_to_input(gen_data_clean, packed_sequence, sigma_max, ...) seeds the rollout. + sigma_max = model._add_noise_to_input.call_args.args[2] # [B,1] + assert torch.allclose(sigma_max, torch.ones_like(sigma_max)) + + +@pytest.mark.L0 +def test_backward_simulation_requires_pure_noise_start(): + """A schedule whose first sigma is < 1.0 must raise (would blend clean data into the seed).""" + model, _ = _make_backward_model(t_list=[0.9, 0.5], backward_grad_steps=1) + gen_data = _make_gen_data() + + with pytest.raises(AssertionError): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + +@pytest.mark.L0 +def test_backward_simulation_invalid_grad_steps(): + """backward_grad_steps=0 must raise AssertionError.""" + model, _ = _make_backward_model(backward_grad_steps=0) + gen_data = _make_gen_data() + + with pytest.raises(AssertionError): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + +@pytest.mark.L0 +def test_backward_n_steps_cycles_by_per_net_update_index(): + """rcm-style cycling: n_steps deterministically cycles 1..max by per-net update index.""" + import types + + model = MagicMock() + model.config.warmup_student_steps = 0 + model.config.warmup_critic_steps = 0 + model.config.student_update_freq = 5 + model.get_phase = types.MethodType(DMD2RFModel.get_phase, model) + model.get_student_iteration = types.MethodType(DMD2RFModel.get_student_iteration, model) + model.get_critic_iteration = types.MethodType(DMD2RFModel.get_critic_iteration, model) + model._backward_n_steps = types.MethodType(DMD2RFModel._backward_n_steps, model) + + # Student steps fall on iteration % 5 == 0 -> student_idx 0,1,2,3,4 -> n_steps cycles 1,2,3,4,1. + student_iters = [0, 5, 10, 15, 20] + assert [model._backward_n_steps(4, it) for it in student_iters] == [1, 2, 3, 4, 1] + # Critic steps (iteration % 5 != 0) cycle on the critic update counter, also spanning 1..4. + critic_iters = [1, 2, 3, 4, 6] + assert all(1 <= model._backward_n_steps(4, it) <= 4 for it in critic_iters) + + # Critic-only warmup: iteration 0 is a critic step with update index -1; the clamp keeps the + # cycle starting at n_steps=1 (not wrapping to max via negative modulo). + model.config.warmup_critic_steps = 100 + assert model.get_phase(0) == "critic" + assert model._backward_n_steps(4, 0) == 1 + + +@pytest.mark.L0 +def test_backward_simulation_grad_steps_exceeds_n_steps(): + """backward_grad_steps > n_steps must not raise and enables grad on every step.""" + model, _ = _make_backward_model(backward_grad_steps=10, t_list=[1.0, 0.5]) + + grad_states = [] + + def capture_grad(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return {"preds_vision": [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)]} + + model._pack_and_denoise.side_effect = capture_grad + + gen_data = _make_gen_data() + with torch.enable_grad(): + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert all(grad_states), "All steps should have grad when backward_grad_steps > n_steps" + + +@pytest.mark.L0 +def test_backward_simulation_ode_step_keeps_conditioned_frames(): + """ODE Euler step must leave conditioned frames unchanged between denoising steps.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + model, B = _make_backward_model(sample_type="ode", t_list=[1.0, 0.5]) + + # Frame 0 is conditioned (mask=1), frame 1 is generated (mask=0) + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.condition_mask = [ + torch.tensor([[[1.0]], [[0.0]]]), # [T=2,1,1]: frame0 conditioned, frame1 generated + torch.tensor([[[1.0]], [[0.0]]]), + ] + packed.action = None + packed.sound = None + model._pack_input_sequence.return_value = packed + + # Initial xt: frame 0 = value 5.0, frame 1 = value 1.0 + xt_init = torch.zeros(4, 2, 2, 2) # [C,T,H,W] + xt_init[:, 0, :, :] = 5.0 # conditioned frame + xt_init[:, 1, :, :] = 1.0 # generated frame + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [xt_init.clone(), xt_init.clone()] + gd_noised.xt_tokens_action = None + gd_noised.xt_tokens_sound = None + gd_noised.epsilon_vision = [torch.zeros(4, 2, 2, 2), torch.zeros(4, 2, 2, 2)] + gd_noised.vt_target_vision = [torch.zeros(4, 2, 2, 2), torch.zeros(4, 2, 2, 2)] + gd_noised.sigmas_vision = [torch.zeros(2, 1, 1), torch.zeros(2, 1, 1)] + gd_noised.epsilon_action = None + gd_noised.vt_target_action = None + gd_noised.sigmas_action = None + gd_noised.epsilon_sound = None + gd_noised.vt_target_sound = None + gd_noised.sigmas_sound = None + model._add_noise_to_input.return_value = gd_noised + + # Non-zero v_pred so ODE step actually moves generated frames + v_pred = torch.ones(4, 2, 2, 2) * 3.0 # [C,T,H,W] + + xt_received_at_step1 = [] + + def capture_step(*args, **kwargs): + if model._pack_and_denoise.call_count == 2: + # Second call: record what xt was passed in gen_data_noised + gd_n = args[1] + xt_received_at_step1.append(gd_n.xt_tokens_vision[0].clone()) + return {"preds_vision": [v_pred.clone(), v_pred.clone()]} + + model._pack_and_denoise.side_effect = capture_step + + gen_data = _make_gen_data() + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert len(xt_received_at_step1) == 1 + xt_step1 = xt_received_at_step1[0] # [C,T,H,W] + # Conditioned frame (frame 0) must be unchanged from initial value 5.0 + assert torch.allclose(xt_step1[:, 0, :, :], torch.full((4, 2, 2), 5.0), atol=1e-5), ( + "Conditioned frame should remain 5.0 after ODE step" + ) + # Generated frame (frame 1) must have changed (delta = -0.499, v=3 → moved by -1.497) + assert not torch.allclose(xt_step1[:, 1, :, :], torch.full((4, 2, 2), 1.0), atol=1e-3), ( + "Generated frame should have changed after ODE step" + ) + + +@pytest.mark.L0 +def test_backward_simulation_sde_step_keeps_conditioned_frames(): + """SDE re-noising step must leave conditioned frames unchanged between denoising steps.""" + from unittest.mock import patch + + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + model, B = _make_backward_model(sample_type="sde", t_list=[1.0, 0.5]) + + # Frame 0 conditioned, frame 1 generated + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.condition_mask = [ + torch.tensor([[[1.0]], [[0.0]]]), + torch.tensor([[[1.0]], [[0.0]]]), + ] + packed.action = None + packed.sound = None + model._pack_input_sequence.return_value = packed + + xt_init = torch.zeros(4, 2, 2, 2) + xt_init[:, 0, :, :] = 5.0 + xt_init[:, 1, :, :] = 1.0 + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [xt_init.clone(), xt_init.clone()] + gd_noised.xt_tokens_action = None + gd_noised.xt_tokens_sound = None + gd_noised.epsilon_vision = [torch.zeros(4, 2, 2, 2), torch.zeros(4, 2, 2, 2)] + gd_noised.vt_target_vision = [torch.zeros(4, 2, 2, 2), torch.zeros(4, 2, 2, 2)] + gd_noised.sigmas_vision = [torch.zeros(2, 1, 1), torch.zeros(2, 1, 1)] + gd_noised.epsilon_action = None + gd_noised.vt_target_action = None + gd_noised.sigmas_action = None + gd_noised.epsilon_sound = None + gd_noised.vt_target_sound = None + gd_noised.sigmas_sound = None + model._add_noise_to_input.return_value = gd_noised + + v_pred = torch.ones(4, 2, 2, 2) * 3.0 + + xt_received_at_step1 = [] + + def capture_step(*args, **kwargs): + if model._pack_and_denoise.call_count == 2: + gd_n = args[1] + xt_received_at_step1.append(gd_n.xt_tokens_vision[0].clone()) + return {"preds_vision": [v_pred.clone(), v_pred.clone()]} + + model._pack_and_denoise.side_effect = capture_step + + # Use deterministic randn (zeros) so the SDE fresh noise term is zero + with patch("torch.randn_like", return_value=torch.zeros(4, 2, 2, 2)): + gen_data = _make_gen_data() + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert len(xt_received_at_step1) == 1 + xt_step1 = xt_received_at_step1[0] # [C,T,H,W] + # Conditioned frame (frame 0) must remain 5.0: SDE uses cond_mask * xt for those frames + assert torch.allclose(xt_step1[:, 0, :, :], torch.full((4, 2, 2), 5.0), atol=1e-5), ( + "Conditioned frame should remain 5.0 after SDE step" + ) + + +# --------------------------------------------------------------------------- +# _backward_simulation dtype preservation (regression for bf16/fp32 mixing bug) +# --------------------------------------------------------------------------- + + +def _make_backward_model_bf16( + *, + sample_type: str = "ode", +) -> tuple[MagicMock, int]: + """Like _make_backward_model but with tensor_kwargs=bfloat16 to match real training.""" + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataNoised + + B = 2 + model = MagicMock() + model.tensor_kwargs = {"dtype": torch.bfloat16, "device": "cpu"} + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.parallel_dims = None + model.config.rectified_flow_inference_config.num_train_timesteps = 1000 + model.config.action_gen = False + model.config.sound_gen = False + model.config.fixed_step_sampler_config.t_list = [1.0, 0.5] + model.config.fixed_step_sampler_config.sample_type = sample_type + model.config.backward_grad_steps = 1 + model._backward_n_steps.side_effect = lambda max_n_steps, iteration: max_n_steps + + # Condition mask is float32 (typical from packing pipeline) — this is the source of the bug + packed = MagicMock() + packed.vision = MagicMock() + packed.vision.condition_mask = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] # float32, [T=1,1,1] + packed.action = None + packed.sound = None + model._pack_input_sequence.return_value = packed + + gd_noised = MagicMock(spec=GenerationDataNoised) + gd_noised.xt_tokens_vision = [ + torch.ones(4, 1, 2, 2, dtype=torch.bfloat16), + torch.ones(4, 1, 2, 2, dtype=torch.bfloat16), + ] + gd_noised.xt_tokens_action = None + gd_noised.xt_tokens_sound = None + gd_noised.epsilon_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.vt_target_vision = [torch.zeros(4, 1, 2, 2), torch.zeros(4, 1, 2, 2)] + gd_noised.sigmas_vision = [torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)] + gd_noised.epsilon_action = None + gd_noised.vt_target_action = None + gd_noised.sigmas_action = None + gd_noised.epsilon_sound = None + gd_noised.vt_target_sound = None + gd_noised.sigmas_sound = None + model._add_noise_to_input.return_value = gd_noised + + model._pack_and_denoise.return_value = { + "preds_vision": [ + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + ] + } + import types + + model._velocity_to_x0 = DMD2RFModel._velocity_to_x0 + model._ode_step = types.MethodType(DMD2RFModel._ode_step, model) + model._sde_step = types.MethodType(DMD2RFModel._sde_step, model) + + return model, B + + +@pytest.mark.L0 +def test_backward_simulation_ode_xt_dtype_matches_model(): + """ODE: xt_tokens passed to _pack_and_denoise on step 2 must match tensor_kwargs dtype. + + Regression test for the bf16/fp32 mixing bug where fp32 condition masks upcast xt_next + to fp32, causing a dtype mismatch in the model's linear layers. + """ + model, _ = _make_backward_model_bf16(sample_type="ode") + + xt_dtypes_step2 = [] + + def capture_step(*args, **kwargs): + if model._pack_and_denoise.call_count == 2: + gd_n = args[1] + xt_dtypes_step2.append(gd_n.xt_tokens_vision[0].dtype) + return { + "preds_vision": [ + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + ] + } + + model._pack_and_denoise.side_effect = capture_step + + gen_data = _make_gen_data() + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert len(xt_dtypes_step2) == 1 + assert xt_dtypes_step2[0] == torch.bfloat16, ( + f"xt_tokens_vision dtype should be bfloat16 after ODE step, got {xt_dtypes_step2[0]}" + ) + + +@pytest.mark.L0 +def test_backward_simulation_sde_xt_dtype_matches_model(): + """SDE: xt_tokens passed to _pack_and_denoise on step 2 must match tensor_kwargs dtype. + + Same regression test as ODE variant, for the SDE re-noising path. + """ + from unittest.mock import patch + + model, _ = _make_backward_model_bf16(sample_type="sde") + + xt_dtypes_step2 = [] + + def capture_step(*args, **kwargs): + if model._pack_and_denoise.call_count == 2: + gd_n = args[1] + xt_dtypes_step2.append(gd_n.xt_tokens_vision[0].dtype) + return { + "preds_vision": [ + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + torch.zeros(4, 1, 2, 2, dtype=torch.bfloat16), + ] + } + + model._pack_and_denoise.side_effect = capture_step + + with patch("torch.randn_like", return_value=torch.zeros(4, 1, 2, 2)): + gen_data = _make_gen_data() + DMD2RFModel._backward_simulation( + model, + input_text_indexes=[[], []], + sequence_plans=[MagicMock(), MagicMock()], + gen_data_clean=gen_data, + iteration=0, + ) + + assert len(xt_dtypes_step2) == 1 + assert xt_dtypes_step2[0] == torch.bfloat16, ( + f"xt_tokens_vision dtype should be bfloat16 after SDE step, got {xt_dtypes_step2[0]}" + ) + + +# --------------------------------------------------------------------------- +# _ode_step and _sde_step unit tests +# --------------------------------------------------------------------------- + + +def _make_step_model(*, dtype: torch.dtype = torch.float32) -> MagicMock: + model = MagicMock() + model.tensor_kwargs = {"dtype": dtype, "device": "cpu"} + model.tensor_kwargs_fp32 = {"dtype": torch.float32, "device": "cpu"} + model.parallel_dims = None + return model + + +@pytest.mark.L0 +def test_ode_step_basic(): + """_ode_step computes xt + delta_sigma * v * noisy_mask element-wise.""" + model = _make_step_model() + xt = [torch.tensor([2.0, 4.0])] + v_pred = [torch.tensor([1.0, 1.0])] + noisy_mask = [torch.tensor([1.0, 0.0])] # second element conditioned + delta_sigma = -0.5 + + result = DMD2RFModel._ode_step(model, xt, v_pred, noisy_mask, delta_sigma) + + expected = torch.tensor([2.0 + (-0.5) * 1.0 * 1.0, 4.0 + (-0.5) * 1.0 * 0.0]) + assert torch.allclose(result[0], expected, atol=1e-6) + + +@pytest.mark.L0 +def test_ode_step_conditioned_frame_unchanged(): + """_ode_step leaves frames where noisy_mask=0 unchanged.""" + model = _make_step_model() + xt = [torch.ones(4, 2, 2, 2) * 5.0] # [C,T,H,W] + v_pred = [torch.ones(4, 2, 2, 2) * 3.0] + noisy_mask = [torch.tensor([[[0.0]], [[1.0]]])] # frame0 conditioned, frame1 generated + delta_sigma = -0.5 + + result = DMD2RFModel._ode_step(model, xt, v_pred, noisy_mask, delta_sigma) + + # Conditioned frame (frame 0): unchanged from 5.0 + assert torch.allclose(result[0][:, 0, :, :], torch.full((4, 2, 2), 5.0), atol=1e-6) + # Generated frame (frame 1): changed + assert not torch.allclose(result[0][:, 1, :, :], torch.full((4, 2, 2), 5.0), atol=1e-3) + + +@pytest.mark.L0 +def test_ode_step_output_dtype_matches_tensor_kwargs(): + """_ode_step output dtype must match tensor_kwargs, not the input dtype.""" + model = _make_step_model(dtype=torch.bfloat16) + xt = [torch.ones(4, 1, 2, 2)] # float32 input + v_pred = [torch.ones(4, 1, 2, 2)] + noisy_mask = [torch.zeros(1, 1, 1)] + delta_sigma = -0.5 + + result = DMD2RFModel._ode_step(model, xt, v_pred, noisy_mask, delta_sigma) + assert result[0].dtype == torch.bfloat16 + + +@pytest.mark.L0 +def test_sde_step_basic_zero_noise(): + """_sde_step with zero fresh noise: xt_next = noisy_mask*(1-σ)*x0 + cond_mask*xt.""" + from unittest.mock import patch + + model = _make_step_model() + xt = [torch.tensor([5.0, 1.0])] # [conditioned, generated] + x0_pred = [torch.tensor([0.0, 2.0])] + noisy_mask = [torch.tensor([0.0, 1.0])] # frame0 conditioned, frame1 generated + cond_mask = [torch.tensor([1.0, 0.0])] + sigma_next = 0.5 + + with patch("torch.randn_like", return_value=torch.zeros(2)): + result = DMD2RFModel._sde_step(model, xt, x0_pred, noisy_mask, cond_mask, sigma_next) + + # frame0 (conditioned): cond_mask=1 → keeps xt=5.0; noisy_mask=0 → no re-noise + # frame1 (generated): noisy_mask=1 → (1-0.5)*2.0 + 0.5*0.0 = 1.0 + assert torch.allclose(result[0], torch.tensor([5.0, 1.0]), atol=1e-6) + + +@pytest.mark.L0 +def test_sde_step_conditioned_frame_unchanged(): + """_sde_step leaves frames where cond_mask=1 at their original xt value.""" + from unittest.mock import patch + + model = _make_step_model() + B = 2 + xt = [torch.ones(4, 2, 2, 2) * 5.0, torch.ones(4, 2, 2, 2) * 5.0] # [C,T,H,W] + x0_pred = [torch.zeros(4, 2, 2, 2), torch.zeros(4, 2, 2, 2)] + noisy_mask = [torch.tensor([[[0.0]], [[1.0]]]), torch.tensor([[[0.0]], [[1.0]]])] + cond_mask = [torch.tensor([[[1.0]], [[0.0]]]), torch.tensor([[[1.0]], [[0.0]]])] + sigma_next = 0.5 + + with patch("torch.randn_like", return_value=torch.zeros(4, 2, 2, 2)): + result = DMD2RFModel._sde_step(model, xt, x0_pred, noisy_mask, cond_mask, sigma_next) + + for i in range(B): + assert torch.allclose(result[i][:, 0, :, :], torch.full((4, 2, 2), 5.0), atol=1e-6), ( + f"Conditioned frame {i} should remain 5.0" + ) + + +@pytest.mark.L0 +def test_sde_step_output_dtype_matches_tensor_kwargs(): + """_sde_step output dtype must match tensor_kwargs (regression: fp32 masks + bf16 xt).""" + from unittest.mock import patch + + model = _make_step_model(dtype=torch.bfloat16) + xt = [torch.ones(4, 1, 2, 2, dtype=torch.bfloat16)] + x0_pred = [torch.ones(4, 1, 2, 2, dtype=torch.bfloat16)] + noisy_mask = [torch.zeros(1, 1, 1)] # float32 mask (typical from packing pipeline) + cond_mask = [torch.ones(1, 1, 1)] # float32 mask + sigma_next = 0.5 + + with patch("torch.randn_like", return_value=torch.zeros(4, 1, 2, 2)): + result = DMD2RFModel._sde_step(model, xt, x0_pred, noisy_mask, cond_mask, sigma_next) + + assert result[0].dtype == torch.bfloat16 + + +@pytest.mark.L0 +def test_sde_step_broadcasts_and_uses_fresh_noise(monkeypatch: pytest.MonkeyPatch) -> None: + """_sde_step must CP-broadcast rank-local fresh SDE noise before composing xt_next.""" + + model = _make_step_model() + model.parallel_dims = MagicMock() + xt = [torch.zeros(2)] # [C] + x0_pred = [torch.zeros(2)] # [C] + noisy_mask = [torch.ones(2)] # [C] + cond_mask = [torch.zeros(2)] # [C] + sigma_next = 1.0 + calls: list[tuple[list[torch.Tensor], object | None]] = [] + + def _spy_broadcast(tensors: list[torch.Tensor] | None, parallel_dims: object | None) -> None: + assert tensors is not None + calls.append(([tensor.clone() for tensor in tensors], parallel_dims)) # list of [C] + tensors[0].fill_(7.0) # [C] + + monkeypatch.setattr(dmd2_rf_module, "context_parallel_broadcast_tensor_list", _spy_broadcast) + + with patch("torch.randn_like", return_value=torch.ones(2)): + result = DMD2RFModel._sde_step(model, xt, x0_pred, noisy_mask, cond_mask, sigma_next) + + assert len(calls) == 1 + assert calls[0][1] is model.parallel_dims + assert torch.allclose(calls[0][0][0], torch.ones(2), atol=1e-6) + assert torch.allclose(result[0], torch.full((2,), 7.0), atol=1e-6) + + +@pytest.mark.L0 +def test_ode_step_defined(): + """_ode_step must be defined directly on DMD2RFModel.""" + assert "_ode_step" in DMD2RFModel.__dict__ + + +@pytest.mark.L0 +def test_sde_step_defined(): + """_sde_step must be defined directly on DMD2RFModel.""" + assert "_sde_step" in DMD2RFModel.__dict__ diff --git a/cosmos_framework/model/generator/distillation/optimizer.py b/cosmos_framework/model/generator/distillation/optimizer.py new file mode 100644 index 00000000..7ddf6845 --- /dev/null +++ b/cosmos_framework/model/generator/distillation/optimizer.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from __future__ import annotations + +from collections.abc import ItemsView, Iterator +from typing import Any, Protocol, TypeGuard, cast + +import torch + +from cosmos_framework.utils.generator.optimizer import LRSchedulersContainer, OptimizersContainer + +__all__: tuple[str, ...] = ( + "OptimizerContainerLike", + "OptimizerModelView", + "PhaseOptimizer", + "PhaseScheduler", + "is_optimizer_container", + "iter_torch_optimizers", +) + + +class OptimizerContainerLike(Protocol): + """Structural interface shared by i4 and release-mapped optimizer containers.""" + + model: torch.nn.Module + optimizers: list[torch.optim.Optimizer] + + def step(self) -> None: ... + + def zero_grad(self, set_to_none: bool = True) -> None: ... + + def state_dict(self) -> dict[str, Any]: ... + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: ... + + +_OPTIMIZER_CONTAINER_METHODS = ("step", "zero_grad", "state_dict", "load_state_dict") + + +def is_optimizer_container(optimizer: object) -> TypeGuard[OptimizerContainerLike]: + """Recognize native or release-mapped optimizer containers by their full interface.""" + return isinstance(optimizer, OptimizersContainer) or ( + isinstance(getattr(optimizer, "model", None), torch.nn.Module) + and isinstance(getattr(optimizer, "optimizers", None), list) + and all(callable(getattr(optimizer, method, None)) for method in _OPTIMIZER_CONTAINER_METHODS) + ) + + +class OptimizerModelView(torch.nn.Module): + """Expose a standalone denoiser as a VFM optimizer-compatible ``.net`` model.""" + + def __init__(self, net: torch.nn.Module) -> None: + super().__init__() + self.net: torch.nn.Module = net + + +def iter_torch_optimizers(optimizer: object) -> Iterator[torch.optim.Optimizer]: + """Yield raw torch optimizers from either a raw optimizer or a VFM container.""" + if is_optimizer_container(optimizer): + yield from optimizer.optimizers + else: + yield cast(torch.optim.Optimizer, optimizer) + + +class PhaseOptimizer: + """Optimizer container for alternating-phase distillation training. + + Holds a dict of optimizers and exposes step/zero_grad with an explicit key + argument. Routing logic (which key is active at a given iteration) lives in + the trainer, not here. + """ + + def __init__(self, optimizer_dict: dict[str, torch.optim.Optimizer | OptimizerContainerLike]) -> None: + self._optimizers: dict[str, torch.optim.Optimizer | OptimizerContainerLike] = optimizer_dict + + def step(self, key: str, grad_scaler: torch.amp.GradScaler) -> None: + for optimizer in iter_torch_optimizers(self._optimizers[key]): + grad_scaler.step(optimizer) + grad_scaler.update() + + def zero_grad(self, key: str) -> None: + self._optimizers[key].zero_grad(set_to_none=True) + + def parameters_for_key(self, key: str) -> list[torch.nn.Parameter]: + return [ + param + for optimizer in iter_torch_optimizers(self._optimizers[key]) + for group in optimizer.param_groups + for param in group["params"] + ] + + def get(self, key: str) -> torch.optim.Optimizer | OptimizerContainerLike | None: + return self._optimizers.get(key) + + def items(self) -> ItemsView[str, torch.optim.Optimizer | OptimizerContainerLike]: + return self._optimizers.items() + + +class PhaseScheduler: + """Scheduler container for alternating-phase distillation training. + + Mirrors PhaseOptimizer: holds a dict of LR schedulers and exposes + step with an explicit key argument. + """ + + def __init__( + self, + scheduler_dict: dict[str, torch.optim.lr_scheduler.LRScheduler | LRSchedulersContainer], + ) -> None: + self._schedulers: dict[str, torch.optim.lr_scheduler.LRScheduler | LRSchedulersContainer] = scheduler_dict + + def step(self, key: str) -> None: + self._schedulers[key].step() + + def get(self, key: str) -> torch.optim.lr_scheduler.LRScheduler | LRSchedulersContainer | None: + return self._schedulers.get(key) + + def items(self) -> ItemsView[str, torch.optim.lr_scheduler.LRScheduler | LRSchedulersContainer]: + return self._schedulers.items() diff --git a/cosmos_framework/model/generator/distillation/optimizer_test.py b/cosmos_framework/model/generator/distillation/optimizer_test.py new file mode 100644 index 00000000..3c396bd6 --- /dev/null +++ b/cosmos_framework/model/generator/distillation/optimizer_test.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from unittest.mock import MagicMock, call + +import pytest +import torch + +from cosmos_framework.utils.generator.optimizer import OptimizersContainer +from cosmos_framework.model.generator.distillation import optimizer as optimizer_module +from cosmos_framework.model.generator.distillation.optimizer import OptimizerModelView, PhaseOptimizer, PhaseScheduler + + +def _make_optimizer() -> MagicMock: + opt = MagicMock() + opt.param_groups = [] + return opt + + +def _make_scheduler() -> MagicMock: + return MagicMock() + + +def _make_grad_scaler() -> MagicMock: + return MagicMock() + + +def _make_optimizer_container(*optimizers: MagicMock) -> OptimizersContainer: + container = object.__new__(OptimizersContainer) + container.optimizers = list(optimizers) + container.zero_grad = MagicMock() + return container + + +class _ForeignOptimizerContainer: + def __init__(self, *optimizers: MagicMock) -> None: + self.model: torch.nn.Module = torch.nn.Linear(1, 1) + self.optimizers: list[MagicMock] = list(optimizers) + + def step(self) -> None: + for optimizer in self.optimizers: + optimizer.step() + + def zero_grad(self, set_to_none: bool = True) -> None: + for optimizer in self.optimizers: + optimizer.zero_grad(set_to_none=set_to_none) + + def state_dict(self) -> dict[str, object]: + return {"foreign": True} + + def load_state_dict(self, state_dict: dict[str, object]) -> None: + del state_dict + + +@pytest.mark.L0 +def test_optimizer_model_view_exposes_net_submodule() -> None: + net = torch.nn.Linear(2, 3) + view = OptimizerModelView(net) + assert view.net is net + assert dict(view.net.named_parameters()) == dict(net.named_parameters()) + + +# ------------------------------------------------------------------------- +# PhaseOptimizer.step — only the given key's optimizer is stepped +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_step_net_key() -> None: + opt_net, opt_fake = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}) + scaler = _make_grad_scaler() + po.step("net", scaler) + + scaler.step.assert_called_once_with(opt_net) + scaler.update.assert_called_once() + opt_fake.step.assert_not_called() + + +@pytest.mark.L0 +def test_step_fake_score_key() -> None: + opt_net, opt_fake = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}) + scaler = _make_grad_scaler() + po.step("fake_score", scaler) + + scaler.step.assert_called_once_with(opt_fake) + scaler.update.assert_called_once() + opt_net.step.assert_not_called() + + +@pytest.mark.L0 +def test_step_calls_grad_scaler_update_once() -> None: + opt_net = _make_optimizer() + po = PhaseOptimizer({"net": opt_net}) + scaler = _make_grad_scaler() + po.step("net", scaler) + scaler.update.assert_called_once() + + +@pytest.mark.L0 +def test_step_optimizer_container_steps_inner_optimizers() -> None: + opt_a, opt_b = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": _make_optimizer_container(opt_a, opt_b)}) + scaler = _make_grad_scaler() + po.step("net", scaler) + assert scaler.step.call_args_list == [call(opt_a), call(opt_b)] + scaler.update.assert_called_once() + + +# ------------------------------------------------------------------------- +# PhaseOptimizer.zero_grad — only the given key's optimizer is zeroed +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_zero_grad_net_key() -> None: + opt_net, opt_fake = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}) + po.zero_grad("net") + + opt_net.zero_grad.assert_called_once_with(set_to_none=True) + opt_fake.zero_grad.assert_not_called() + + +@pytest.mark.L0 +def test_zero_grad_fake_score_key() -> None: + opt_net, opt_fake = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}) + po.zero_grad("fake_score") + + opt_fake.zero_grad.assert_called_once_with(set_to_none=True) + opt_net.zero_grad.assert_not_called() + + +@pytest.mark.L0 +def test_zero_grad_optimizer_container_uses_container_method() -> None: + container = _make_optimizer_container(_make_optimizer()) + po = PhaseOptimizer({"net": container}) + po.zero_grad("net") + container.zero_grad.assert_called_once_with(set_to_none=True) + + +@pytest.mark.L0 +def test_parameters_for_key_flattens_optimizer_container_param_groups() -> None: + param_a = MagicMock() + param_b = MagicMock() + opt_a, opt_b = _make_optimizer(), _make_optimizer() + opt_a.param_groups = [{"params": [param_a]}] + opt_b.param_groups = [{"params": [param_b]}] + po = PhaseOptimizer({"net": _make_optimizer_container(opt_a, opt_b)}) + assert po.parameters_for_key("net") == [param_a, param_b] + + +@pytest.mark.L0 +def test_parameters_for_key_accepts_structurally_compatible_foreign_container() -> None: + param_a = MagicMock() + param_b = MagicMock() + opt_a, opt_b = _make_optimizer(), _make_optimizer() + opt_a.param_groups = [{"params": [param_a]}] + opt_b.param_groups = [{"params": [param_b]}] + po = PhaseOptimizer({"net": _ForeignOptimizerContainer(opt_a, opt_b)}) + + assert po.parameters_for_key("net") == [param_a, param_b] + + +# ------------------------------------------------------------------------- +# PhaseOptimizer.get +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_get_existing_key() -> None: + opt_net = _make_optimizer() + po = PhaseOptimizer({"net": opt_net}) + assert po.get("net") is opt_net + + +@pytest.mark.L0 +def test_get_missing_key_returns_none() -> None: + po = PhaseOptimizer({"net": _make_optimizer()}) + assert po.get("discriminator") is None + + +# ------------------------------------------------------------------------- +# PhaseOptimizer.items — returns all entries +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_optimizer_items_returns_all_entries() -> None: + opt_net, opt_fake = _make_optimizer(), _make_optimizer() + po = PhaseOptimizer({"net": opt_net, "fake_score": opt_fake}) + items = dict(po.items()) + assert items["net"] is opt_net + assert items["fake_score"] is opt_fake + + +# ------------------------------------------------------------------------- +# PhaseScheduler.step — only the given key's scheduler is stepped +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_scheduler_step_net_key() -> None: + sched_net, sched_fake = _make_scheduler(), _make_scheduler() + ps = PhaseScheduler({"net": sched_net, "fake_score": sched_fake}) + ps.step("net") + + sched_net.step.assert_called_once() + sched_fake.step.assert_not_called() + + +@pytest.mark.L0 +def test_scheduler_step_fake_score_key() -> None: + sched_net, sched_fake = _make_scheduler(), _make_scheduler() + ps = PhaseScheduler({"net": sched_net, "fake_score": sched_fake}) + ps.step("fake_score") + + sched_fake.step.assert_called_once() + sched_net.step.assert_not_called() + + +# ------------------------------------------------------------------------- +# PhaseScheduler.get +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_scheduler_get_existing_key() -> None: + sched_net = _make_scheduler() + ps = PhaseScheduler({"net": sched_net}) + assert ps.get("net") is sched_net + + +@pytest.mark.L0 +def test_scheduler_get_missing_key_returns_none() -> None: + ps = PhaseScheduler({"net": _make_scheduler()}) + assert ps.get("discriminator") is None + + +# ------------------------------------------------------------------------- +# PhaseScheduler.items — returns all entries +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_scheduler_items_returns_all_entries() -> None: + sched_net, sched_fake = _make_scheduler(), _make_scheduler() + ps = PhaseScheduler({"net": sched_net, "fake_score": sched_fake}) + items = dict(ps.items()) + assert items["net"] is sched_net + assert items["fake_score"] is sched_fake + + +# ------------------------------------------------------------------------- +# Single-key construction (base class scenario: net only) +# ------------------------------------------------------------------------- +@pytest.mark.L0 +def test_single_key_optimizer() -> None: + opt_net = _make_optimizer() + po = PhaseOptimizer({"net": opt_net}) + scaler = _make_grad_scaler() + po.step("net", scaler) + po.zero_grad("net") + scaler.step.assert_called_once_with(opt_net) + opt_net.zero_grad.assert_called_once_with(set_to_none=True) + + +@pytest.mark.L0 +def test_single_key_scheduler() -> None: + sched_net = _make_scheduler() + ps = PhaseScheduler({"net": sched_net}) + ps.step("net") + sched_net.step.assert_called_once() + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_optimizer_exports_are_explicit() -> None: + assert optimizer_module.__all__ == ( + "OptimizerContainerLike", + "OptimizerModelView", + "PhaseOptimizer", + "PhaseScheduler", + "is_optimizer_container", + "iter_torch_optimizers", + ) + + +@pytest.mark.L0 +@pytest.mark.CPU +@pytest.mark.parametrize("phase_key", ["net", "fake_score"]) +def test_phase_optimizer_and_scheduler_route_the_same_key(phase_key: str) -> None: + optimizers = {"net": _make_optimizer(), "fake_score": _make_optimizer()} + schedulers = {"net": _make_scheduler(), "fake_score": _make_scheduler()} + phase_optimizer = PhaseOptimizer(optimizers) + phase_scheduler = PhaseScheduler(schedulers) + grad_scaler = _make_grad_scaler() + + phase_optimizer.step(phase_key, grad_scaler) + phase_scheduler.step(phase_key) + + grad_scaler.step.assert_called_once_with(optimizers[phase_key]) + schedulers[phase_key].step.assert_called_once_with() diff --git a/cosmos_framework/model/generator/hf_model.py b/cosmos_framework/model/generator/hf_model.py index 6482479b..65180cf2 100644 --- a/cosmos_framework/model/generator/hf_model.py +++ b/cosmos_framework/model/generator/hf_model.py @@ -26,8 +26,9 @@ from accelerate import init_on_device from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText -from cosmos_framework.utils import log +import cosmos_framework.model.generator.reasoner.cosmos3_edge # noqa: F401 registers cosmos3_edge with transformers Auto classes from cosmos_framework.model.generator.utils.safetensors_loader import load_language_model, load_vlm_model +from cosmos_framework.utils import log from cosmos_framework.utils.generator.parallelism import ParallelDims diff --git a/cosmos_framework/model/generator/mot/attention.py b/cosmos_framework/model/generator/mot/attention.py index 7eb7fa6c..45123e37 100644 --- a/cosmos_framework/model/generator/mot/attention.py +++ b/cosmos_framework/model/generator/mot/attention.py @@ -68,6 +68,30 @@ def __init__( AttentionMaskType = SplitInfo +_SPLIT_INFO_ATTRIBUTES = ( + "max_causal_len", + "max_full_len", + "max_sample_len", + "split_lens", + "attn_modes", + "sample_lens", + "is_three_way", + "vision_token_shapes", + "action_token_shapes", + "num_action_tokens_per_supertoken", + "null_action_supertokens", + "control_stream_token_ranges", + "noisy_token_range", + "control_weights", +) + + +def _is_split_info_compatible(attention_mask: object) -> bool: + return isinstance(attention_mask, SplitInfo) or all( + hasattr(attention_mask, attribute) for attribute in _SPLIT_INFO_ATTRIBUTES + ) + + _dotproduct_attention_cache = {} @@ -481,14 +505,16 @@ def dispatch_attention( packed_key_states_normalized: SequencePack | None = None, ) -> tuple[SequencePack, KVToStore | None]: assert memory_value is None, "Base dispatch_attention does not handle MemoryValue" - if isinstance(attention_mask, SplitInfo) and attention_mask.control_stream_token_ranges is not None: + if not _is_split_info_compatible(attention_mask): + raise TypeError(f"Unsupported attention metadata: {type(attention_mask)}") + if attention_mask.control_stream_token_ranges is not None: output = multi_control_two_way_attention( packed_query_states, packed_key_states, packed_value_states, attention_mask, ) - elif isinstance(attention_mask, SplitInfo) and attention_mask.is_three_way: + elif attention_mask.is_three_way: output = three_way_attention( packed_query_states, packed_key_states, @@ -497,15 +523,13 @@ def dispatch_attention( attention_meta=attention_mask, packed_key_states_normalized=packed_key_states_normalized, ) - elif isinstance(attention_mask, SplitInfo): + else: output = two_way_attention( packed_query_states, packed_key_states, packed_value_states, packed_key_states_normalized=packed_key_states_normalized, ) - else: - raise TypeError(f"Unsupported attention metadata: {type(attention_mask)}") return output, None diff --git a/cosmos_framework/model/generator/mot/attention_test.py b/cosmos_framework/model/generator/mot/attention_test.py index 77c32198..27da6dc8 100644 --- a/cosmos_framework/model/generator/mot/attention_test.py +++ b/cosmos_framework/model/generator/mot/attention_test.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -import random +import random # noqa: I001 - release import rewriting changes the package sort order. +from types import SimpleNamespace from typing import cast import pytest @@ -25,6 +26,27 @@ SEQS_PER_BATCH = 4 +def _foreign_split_info(**overrides: object) -> SimpleNamespace: + fields: dict[str, object] = { + "max_causal_len": 1, + "max_full_len": 1, + "max_sample_len": 2, + "split_lens": [1, 1], + "attn_modes": ["causal", "full"], + "sample_lens": [2], + "is_three_way": False, + "vision_token_shapes": None, + "action_token_shapes": None, + "num_action_tokens_per_supertoken": 0, + "null_action_supertokens": False, + "control_stream_token_ranges": None, + "noisy_token_range": None, + "control_weights": None, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + def unwrap(fn): import torch.utils._pytree as pytree @@ -357,6 +379,41 @@ def test_build_packed_sequence_rejects_flex(): ) +@pytest.mark.L0 +def test_dispatch_attention_accepts_structurally_compatible_split_info(monkeypatch: pytest.MonkeyPatch) -> None: + expected_output = object() + + def fake_two_way_attention(*args: object, **kwargs: object) -> object: + return expected_output + + monkeypatch.setattr(attention, "two_way_attention", fake_two_way_attention) + foreign_split_info = _foreign_split_info() + + output, kv_to_store = attention.dispatch_attention( + object(), + object(), + object(), + foreign_split_info, + ) + + assert output is expected_output + assert kv_to_store is None + + +@pytest.mark.L0 +def test_dispatch_attention_rejects_incomplete_split_info() -> None: + foreign_split_info = _foreign_split_info() + del foreign_split_info.control_weights + + with pytest.raises(TypeError, match="Unsupported attention metadata"): + attention.dispatch_attention( + object(), + object(), + object(), + foreign_split_info, + ) + + @pytest.mark.L0 def test_decoder_layer_optimized_path_empty_und_tensor_shape(): """Empty und tensors in the optimized AR path must be 2D, not 1D. diff --git a/cosmos_framework/model/generator/mot/context_parallel_test.py b/cosmos_framework/model/generator/mot/context_parallel_test.py index 7f3dce0a..4f4888bb 100644 --- a/cosmos_framework/model/generator/mot/context_parallel_test.py +++ b/cosmos_framework/model/generator/mot/context_parallel_test.py @@ -19,8 +19,7 @@ context_parallel_attention, get_context_parallel_sharded_sequence, ) -from cosmos_framework.model.generator.mot.parallelize_unified_mot import ARReplicatedIODispatch -from cosmos_framework.model.generator.mot.unified_mot import _apply_head_sharded_o_proj + from cosmos_framework.model.generator.utils.data_and_condition import GenerationDataClean from cosmos_framework.data.generator.sequence_packing import ( PackedSequence, diff --git a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py index 39245765..cdb2cbcf 100644 --- a/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py +++ b/cosmos_framework/model/generator/mot/cosmos3_vfm_network.py @@ -53,6 +53,7 @@ def __init__( sound_dim: int | None = None, temporal_compression_factor_sound=1, sound_latent_fps: int = 25, + enable_input_bias: bool = True, **kwargs, ): self.vision_gen = vision_gen @@ -76,6 +77,7 @@ def __init__( self.temporal_compression_factor_vision = temporal_compression_factor_vision self.natten_parameter_list = natten_parameter_list self.video_temporal_causal = video_temporal_causal + self.enable_input_bias = enable_input_bias # action related parameters self.action_gen = action_gen # whether to generate action tokens @@ -144,8 +146,9 @@ def __init__(self, language_model, config: Cosmos3VFMNetworkConfig): self.latent_channel = config.latent_channel_size self.patch_latent_dim = self.latent_patch_size**2 * self.latent_channel - self.time_embedder = TimestepEmbedder(self.hidden_size) - self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) + _input_bias = config.enable_input_bias + self.time_embedder = TimestepEmbedder(self.hidden_size, bias=_input_bias) + self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size, bias=_input_bias) self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) if config.action_gen: @@ -158,7 +161,7 @@ def __init__(self, language_model, config: Cosmos3VFMNetworkConfig): if config.sound_gen: self.sound_dim = config.sound_dim - self.sound2llm = nn.Linear(config.sound_dim, self.hidden_size) + self.sound2llm = nn.Linear(config.sound_dim, self.hidden_size, bias=config.enable_input_bias) self.llm2sound = nn.Linear(self.hidden_size, config.sound_dim) self.sound_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) @@ -172,7 +175,8 @@ def init_weights(self, buffer_device: torch.device | None): if self.config.vision_gen: std = 1.0 / math.sqrt(self.patch_latent_dim) torch.nn.init.trunc_normal_(self.vae2llm.weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.vae2llm.bias) + if self.config.enable_input_bias: + torch.nn.init.zeros_(self.vae2llm.bias) std = 1.0 / math.sqrt(self.hidden_size) torch.nn.init.trunc_normal_(self.llm2vae.weight, std=std, a=-3 * std, b=3 * std) @@ -197,7 +201,8 @@ def init_weights(self, buffer_device: torch.device | None): # sound2llm: input_size=sound_dim, output_size=hidden_size std = 1.0 / math.sqrt(self.sound_dim) torch.nn.init.trunc_normal_(self.sound2llm.weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.sound2llm.bias) + if self.config.enable_input_bias: + torch.nn.init.zeros_(self.sound2llm.bias) # llm2sound: input_size=hidden_size, output_size=sound_dim std = 1.0 / math.sqrt(self.hidden_size) diff --git a/cosmos_framework/model/generator/mot/modeling_utils.py b/cosmos_framework/model/generator/mot/modeling_utils.py index 037cd863..68e51c7c 100644 --- a/cosmos_framework/model/generator/mot/modeling_utils.py +++ b/cosmos_framework/model/generator/mot/modeling_utils.py @@ -30,24 +30,27 @@ class TimestepEmbedder(nn.Module): Embeds scalar timesteps into vector representations. """ - def __init__(self, hidden_size, frequency_embedding_size=256): + def __init__(self, hidden_size, frequency_embedding_size=256, bias: bool = True): super().__init__() self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.Linear(frequency_embedding_size, hidden_size, bias=bias), nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True), + nn.Linear(hidden_size, hidden_size, bias=bias), ) self.frequency_embedding_size = frequency_embedding_size self.hidden_size = hidden_size + self._has_bias = bias def _init_weights(self): std = 1.0 / math.sqrt(self.frequency_embedding_size) torch.nn.init.trunc_normal_(self.mlp[0].weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.mlp[0].bias) + if self._has_bias: + torch.nn.init.zeros_(self.mlp[0].bias) std = 1.0 / math.sqrt(self.hidden_size) torch.nn.init.trunc_normal_(self.mlp[2].weight, std=std, a=-3 * std, b=3 * std) - torch.nn.init.zeros_(self.mlp[2].bias) + if self._has_bias: + torch.nn.init.zeros_(self.mlp[2].bias) @staticmethod def timestep_embedding(t, dim, max_period=10000): diff --git a/cosmos_framework/model/generator/mot/unified_mot.py b/cosmos_framework/model/generator/mot/unified_mot.py index 2bdebfdc..425e84d7 100644 --- a/cosmos_framework/model/generator/mot/unified_mot.py +++ b/cosmos_framework/model/generator/mot/unified_mot.py @@ -5,6 +5,7 @@ import time from collections.abc import Mapping from dataclasses import dataclass +from pathlib import Path from typing import Any import torch @@ -50,6 +51,9 @@ from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( prepare_multimodal_reasoner_inputs, ) +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.reasoner_multimodal_utils import ( + prepare_multimodal_reasoner_inputs as _nemotron_prepare_multimodal_reasoner_inputs, +) # Qwen3-VL-MoE imports from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import ( @@ -58,6 +62,8 @@ Qwen3VLMoeVisionConfig, ) from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import ( + AuxLossFreeLoadBalancingConfig, + CosineRouterConfig, LBLMetadata, Qwen3VLMoePreTrainedModel, Qwen3VLMoeTextMLP, @@ -231,8 +237,10 @@ def __init__( qk_norm_for_diffusion: bool = True, include_visual: bool = False, gen_noisy_gating: bool = False, + gen_cosine_router_config: CosineRouterConfig | None = None, + gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, text_config_overrides: Mapping[str, Any] | None = None, - ): + ) -> None: # Defensive copy so downstream materialization can't mutate the # caller's input. self.config_dict = dict(config_dict) @@ -242,6 +250,16 @@ def __init__( # Noisy top-k gating on the generation-tower MoE blocks (Shazeer 2017). # Gen-tower only; the understanding tower never receives this flag. self.gen_noisy_gating = gen_noisy_gating + # Cosine router on the generation-tower MoE blocks: token-mean-centered, + # L2-normalized input × L2-normalized gate rows, scaled by a learnable + # temperature. Gen-tower only; the understanding tower keeps the standard + # dot-product router. + self.gen_cosine_router_config: CosineRouterConfig = gen_cosine_router_config or CosineRouterConfig() + # Aux-loss-free load balancing on the gen-tower MoE blocks. + # Gen-tower only; the understanding tower uses the disabled default. + self.gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig = ( + gen_aux_loss_free_load_balancing_config or AuxLossFreeLoadBalancingConfig() + ) # Plain attribute (not a property) so the ``create_vlm_config`` # post-construction ``setattr`` flow can replace the whole # mapping in one shot; default to ``{}`` so the merge in @@ -801,9 +819,11 @@ def _impl_init( layer_types: LayerTypes, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, - gen_noisy_gating: bool = False, use_und_k_norm_for_gen: bool = False, -): + gen_noisy_gating: bool = False, + gen_cosine_router_config: CosineRouterConfig | None = None, + gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, +) -> None: """Shared ``__init__`` body for the three MoT text-model variants. Used by ``Qwen3VLTextModel``, ``Qwen3VLMoeTextModel``, and @@ -824,8 +844,10 @@ def _impl_init( layer_idx=layer_idx, qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, - gen_noisy_gating=gen_noisy_gating, use_und_k_norm_for_gen=use_und_k_norm_for_gen, + gen_noisy_gating=gen_noisy_gating, + gen_cosine_router_config=gen_cosine_router_config, + gen_aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, ) ) @@ -1002,9 +1024,11 @@ def __init__( layer_types: LayerTypes, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, - gen_noisy_gating: bool = False, use_und_k_norm_for_gen: bool = False, - ): + gen_noisy_gating: bool = False, + gen_cosine_router_config: CosineRouterConfig | None = None, + gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + ) -> None: super().__init__() self.hidden_size = config.hidden_size self.self_attn = PackedAttentionMoT( @@ -1022,8 +1046,14 @@ def __init__( and (config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0) ): self.mlp = Qwen3VLMoeTextSparseMoeBlock(config) - # Noisy gating is gen-tower only. - self.mlp_moe_gen = Qwen3VLMoeTextSparseMoeBlock(config, noisy_gating=gen_noisy_gating) + # Noisy gating, the cosine router, and aux-loss-free load balancing + # are gen-tower only. + self.mlp_moe_gen = Qwen3VLMoeTextSparseMoeBlock( + config, + noisy_gating=gen_noisy_gating, + cosine_router_config=gen_cosine_router_config, + aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, + ) else: self.mlp = layer_types.mlp(config) self.mlp_moe_gen = layer_types.mlp(config) @@ -1274,9 +1304,11 @@ def __init__( *, qk_norm_for_text: bool, qk_norm_for_diffusion: bool, - gen_noisy_gating: bool = False, use_und_k_norm_for_gen: bool, - ): + gen_noisy_gating: bool = False, + gen_cosine_router_config: CosineRouterConfig | None = None, + gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + ) -> None: super().__init__(config) _impl_init( self, @@ -1284,8 +1316,10 @@ def __init__( layer_types=LayerTypes("qwen3_vl_moe"), qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, - gen_noisy_gating=gen_noisy_gating, use_und_k_norm_for_gen=use_und_k_norm_for_gen, + gen_noisy_gating=gen_noisy_gating, + gen_cosine_router_config=gen_cosine_router_config, + gen_aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, ) def init_taylorseer(self, cache_dic=None, current=None): @@ -1634,9 +1668,15 @@ def _impl_generate_reasoner_text( presence_penalty: float = 0.0, seed: int | None = None, return_only_new_tokens: bool = False, + prepare_multimodal_fn: Any = prepare_multimodal_reasoner_inputs, ) -> torch.Tensor: """Run a reasoner-tower autoregressive decode loop with a per-layer KV cache. + ``prepare_multimodal_fn`` builds the image/video-conditioned prefill inputs + (embeds + mrope positions). Defaults to the Qwen3-VL implementation; the + Nemotron 3 Dense VL reasoner injects its own SigLIP2 variant so the same + decode loop serves both families. + The loop has two phases: 1. Prefill — process the prompt in one forward pass. For I2V, this first reuses Qwen3-VL visual helpers to scatter image embeddings. @@ -1792,7 +1832,7 @@ def _impl_generate_reasoner_text( deepstack_visual_embeds, position_ids, mrope_position_deltas, - ) = prepare_multimodal_reasoner_inputs( + ) = prepare_multimodal_fn( causal_lm, input_ids=input_ids, pixel_values=pixel_values, @@ -2138,8 +2178,10 @@ def __init__(self, config: Qwen3VLMoeMoTConfig): text_config, qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, - gen_noisy_gating=config.gen_noisy_gating, use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), + gen_noisy_gating=config.gen_noisy_gating, + gen_cosine_router_config=getattr(config, "gen_cosine_router_config", None), + gen_aux_loss_free_load_balancing_config=config.gen_aux_loss_free_load_balancing_config, ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) @@ -2171,6 +2213,14 @@ def init_moe(self) -> None: # Noisy-gating projection is gen-tower only (the und tower has no # gate_noise counterpart), so keep its zero-init rather than copy. pass + elif "log_temperature" in original_name: + # Cosine-router temperature is gen-tower only (the und tower has + # no log_temperature counterpart), so keep its init value. + pass + elif "router_bias" in original_name: + # Learned cosine-router de-sink bias is gen-tower only (the und tower + # has no router_bias counterpart), so keep its zero-init. + pass else: raise ValueError(f"Could not find {original_name} in state_dict for initialization of {name}") @@ -2263,6 +2313,41 @@ def generate_reasoner_text( ) +def _load_bundled_vision_tower(local_dir: str | None) -> tuple[str, dict, dict] | None: + """Locate a checkpoint-local Edge vision tower bundle. + + Returns ``(weights_file, ve_cfg, top_cfg)`` when ``local_dir`` contains + ``vision_encoder/model.safetensors``, else ``None``. Config resolution + mirrors the hub branch's standalone-first convention: + + * ``vision_encoder/config.json`` when present — ``export_model --vit`` + writes a self-describing one (``vision_config`` + ``projector_config`` + + the multimodal token ids), so ``top_cfg`` is the same dict. + * else the dir's top-level ``config.json`` (a raw hub snapshot, which + folds the spec and token ids into the top-level config). + * a standalone file without token ids (older layout) still sources the + ids from the top-level config. + """ + if not local_dir: + return None + weights = Path(local_dir) / "vision_encoder" / "model.safetensors" + if not weights.is_file(): + return None + standalone = weights.parent / "config.json" + if standalone.is_file(): + with open(standalone) as f: + ve_cfg = json.load(f) + else: + with open(Path(local_dir) / "config.json") as f: + ve_cfg = json.load(f) + if "image_token_id" in ve_cfg: + top_cfg = ve_cfg + else: + with open(Path(local_dir) / "config.json") as f: + top_cfg = json.load(f) + return str(weights), ve_cfg, top_cfg + + class Nemotron3DenseVLTextForCausalLM(Nemotron3DenseVLPreTrainedModel): """Causal LM head on top of the Nemotron 3 Dense VL MoT text model.""" @@ -2337,6 +2422,11 @@ def generate_reasoner_text( input_ids: torch.Tensor, max_new_tokens: int, *, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, do_sample: bool = False, @@ -2348,4 +2438,160 @@ def generate_reasoner_text( seed: int | None = None, return_only_new_tokens: bool = False, ) -> torch.Tensor: - raise NotImplementedError("This method is not implemented for Nemotron 3 Dense VL.") + """Autoregressively generate text tokens using only the reasoner tower. + + Mirrors :meth:`Qwen3VLTextForCausalLM.generate_reasoner_text` for the + **text-only** path: the reasoner-tower decode loop is shared — + ``Nemotron3DenseVLTextModel.reasoner_forward`` delegates to the same + ``_impl_reasoner_forward`` the Qwen text models use — so a text prompt + simply flows through :func:`_impl_generate_reasoner_text`. + + The keyword arguments mirror the Qwen signature because + ``Cosmos3VFMNetwork.generate_reasoner_text`` forwards the multimodal + kwargs unconditionally; accepting them (rather than omitting them) is + what lets text-only prompts reach this method instead of raising a + ``TypeError`` on an unexpected ``pixel_values`` kwarg. + + **Image and video conditioning** are supported via a SigLIP2 vision + tower. Unlike Qwen3-VL (whose ViT ships in the DCP checkpoint), Edge's + understanding vision encoder is a self-contained SigLIP2 + projector + stack published in the HF repo under ``vision_encoder/``; + :meth:`_ensure_vision_tower` loads it lazily on the first vision prompt + (directly from HF, no DCP remap) and caches it as ``self.visual``. The + vision-conditioned prefill then runs through + :func:`_impl_generate_reasoner_text` with the Nemotron variant of + ``prepare_multimodal_reasoner_inputs`` (SigLIP2 encode → masked_scatter → + multimodal rope, no deepstack). SigLIP2 has no temporal attention, so a + video is encoded as its frames stacked along the ``video_grid_thw`` + temporal axis; the image and video pairs are mutually exclusive. + """ + # Vision-conditioned prefill uses a SigLIP2 vision tower loaded lazily from + # the HF checkpoint (see _ensure_vision_tower). Text-only prompts skip it. + use_vision = ( + pixel_values is not None + or image_grid_thw is not None + or pixel_values_videos is not None + or video_grid_thw is not None + ) + if use_vision: + self._ensure_vision_tower() + return _impl_generate_reasoner_text( + self, + input_ids=input_ids, + max_new_tokens=max_new_tokens, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + do_sample=do_sample, + temperature=temperature, + top_k=top_k, + top_p=top_p, + repetition_penalty=repetition_penalty, + presence_penalty=presence_penalty, + seed=seed, + return_only_new_tokens=return_only_new_tokens, + prepare_multimodal_fn=_nemotron_prepare_multimodal_reasoner_inputs, + ) + + def _ensure_vision_tower(self, hf_repo: str = "nvidia/Cosmos3-Edge") -> None: + """Lazily build the SigLIP2 vision tower on the first vision prompt. + + The tower is not in the DCP checkpoint; it ships under ``vision_encoder/``. + Prefer the copy bundled in the local checkpoint dir (threaded here via + ``_local_checkpoint_dir`` — loads fully offline), else download from the + hub. Cached as ``self.visual``; also plumbs the multimodal token ids onto + ``self.config`` for ``get_rope_index`` / ``get_placeholder_mask``. + """ + if getattr(self, "visual", None) is not None: + return + from types import SimpleNamespace + + from safetensors.torch import load_file + from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig + + from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.vision_siglip2 import ( + NemotronSiglip2VisionEncoder, + ) + + # Local-first: prefer a checkpoint-bundled ``vision_encoder/`` (written + # by ``export_model`` with the default ``--vit``, or present in a full + # hub snapshot dir) over a live hub download. + bundled = _load_bundled_vision_tower(getattr(self, "_local_checkpoint_dir", None)) + if bundled is not None: + weights_file, ve_cfg, top_cfg = bundled + log.info(f"Edge vision tower: loading from checkpoint-local bundle {Path(weights_file).parent}") + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + + def _hub_error(filename: str) -> RuntimeError: + return RuntimeError( + f"Failed to download '{filename}' from Hugging Face repo '{hf_repo}' " + "(required to build the Cosmos3-Edge vision tower for image/video prompts). " + "Likely causes: offline environment or cold HF cache, the repo is private or " + "gated, or HF_TOKEN is missing/expired. Re-exporting the checkpoint with the " + "default --vit bundles vision_encoder/ into it, making inference " + "self-contained (no hub access needed)." + ) + + log.info(f"Edge vision tower: no checkpoint-local bundle; downloading from HF repo '{hf_repo}'") + try: + config_file = hf_hub_download(hf_repo, "config.json") + except Exception as e: + raise _hub_error("config.json") from e + with open(config_file) as f: + top_cfg = json.load(f) + # The SigLIP2 + projector spec normally lives in a standalone + # ``vision_encoder/config.json``, but the public ``nvidia/Cosmos3-Edge`` + # repo ships only the weights under ``vision_encoder/`` and folds the + # spec into the top-level ``config.json``. Prefer the standalone file + # when present, else fall back to the top-level config. + try: + with open(hf_hub_download(hf_repo, "vision_encoder/config.json")) as f: + ve_cfg = json.load(f) + except EntryNotFoundError: + # Also covers LocalEntryNotFoundError: with a warm offline cache + # the (nonexistent) standalone file is never cached, so keep + # falling back to the top-level config instead of hard-failing. + ve_cfg = top_cfg + except Exception as e: + raise _hub_error("vision_encoder/config.json") from e + try: + weights_file = hf_hub_download(hf_repo, "vision_encoder/model.safetensors") + except Exception as e: + raise _hub_error("vision_encoder/model.safetensors") from e + vdict = dict(ve_cfg["vision_config"]) + for k in ("model_type", "transformers_version", "spatial_merge_size"): + vdict.pop(k, None) + vcfg = Siglip2VisionConfig(**vdict) + vcfg._attn_implementation = "eager" + pc = ve_cfg["projector_config"] + pcfg = SimpleNamespace( + spatial_merge_size=pc["spatial_merge_size"], + input_hidden_size=pc["input_hidden_size"], + # Older standalone configs name this ``merger_intermedia``; the + # top-level config names it ``merger_intermediate_size``. + merger_intermedia=pc.get("merger_intermedia", pc.get("merger_intermediate_size")), + out_hidden_size=pc["out_hidden_size"], + ) + ref = self.model.embed_tokens.weight + encoder = NemotronSiglip2VisionEncoder(vcfg, pcfg).to(device=ref.device, dtype=ref.dtype).eval() + state = load_file(weights_file) + # Weights are prefixed ``model.visual.*`` / ``model.projector.*``; strip + # the ``model.`` prefix to match the encoder's ``visual.*`` / ``projector.*``. + state = {k[len("model.") :]: v for k, v in state.items()} + missing, unexpected = encoder.load_state_dict(state, strict=False) + if missing or unexpected: + raise RuntimeError( + f"Vision tower weight load mismatch: missing={missing[:3]} unexpected={unexpected[:3]}" + ) + self.visual = encoder + # Plumb multimodal token ids for get_rope_index / get_placeholder_mask. + self.config.image_token_id = top_cfg["image_token_id"] + self.config.video_token_id = top_cfg["video_token_id"] + self.config.vision_start_token_id = top_cfg["vision_start_token_id"] + self.config.vision_config = SimpleNamespace(spatial_merge_size=pc["spatial_merge_size"]) diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py index 6f3a0b41..1a7b38d3 100644 --- a/cosmos_framework/model/generator/omni_mot_model.py +++ b/cosmos_framework/model/generator/omni_mot_model.py @@ -45,6 +45,14 @@ unwrap_and_densify, ) from cosmos_framework.model.generator.utils.memory import MemoryState +from cosmos_framework.model.generator.utils.moe_utils import ( + sync_expert_biases_to_ema, + sync_router_biases_to_ema, + update_expert_biases, + update_router_biases, + uses_aux_loss_free_load_balancing, + uses_ema_router_bias, +) from cosmos_framework.model.generator.utils.safetensors_loader import ( load_language_model as load_language_model_safetensors, ) @@ -178,6 +186,7 @@ def build_net(self, dtype: torch.dtype, *, lora_enabled: bool | None = None) -> 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" + language_model = lazy_instantiate(self.vlm_config.model_instance) # NOTE: We pass "RF timesteps" to the network in the same scale as the scheduler @@ -208,6 +217,7 @@ def build_net(self, dtype: torch.dtype, *, lora_enabled: bool | None = None) -> # Sound generation parameters sound_dim=self.config.sound_dim, sound_latent_fps=self.config.sound_latent_fps, + enable_input_bias=self.config.enable_input_bias, ) network_config._attn_implementation_internal = "eager" net = Cosmos3VFMNetwork( @@ -356,6 +366,8 @@ def set_up_model(self): with misc.timer("Creating PyTorch model and ema if enabled"): self.net = self.build_net(dtype=self.precision) self._param_count = count_params(self.net, verbose=False) + self._uses_aux_loss_free_load_balancing: bool = uses_aux_loss_free_load_balancing(self.net) + self._uses_ema_router_bias: bool = uses_ema_router_bias(self.net) if config.ema.enabled: self.net_ema = self.build_net(dtype=torch.float32) @@ -510,6 +522,36 @@ def _derive_include_end_of_generation_token(self) -> bool: return False # ------------------------ training hooks ------------------------ + def on_before_optimizer_step( + self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int + ) -> None: + """Run aux-loss-free load balancing + EMA router de-sink on the gen-tower MoE blocks.""" + del scheduler, optimizer + + dp_mesh = self.parallel_dims.dp_mesh if self.parallel_dims else None + + # Aux-loss-free load balancing (only when enabled). + if self._uses_aux_loss_free_load_balancing: + update_expert_biases( + net=self.net, + device_mesh=dp_mesh, + ) + # expert_bias is a buffer, but the DTensor EMA worker copies parameters + # only, so mirror it into net_ema manually (else the persisted net_ema.* + # keeps zero and drops the trained bias at checkpoint save). update_bias + # already cross-rank-reduced, so net.expert_bias matches on every rank. + if self.config.ema.enabled: + sync_expert_biases_to_ema(net=self.net, net_ema=self.net_ema) + + # EMA-tracked token-constant de-sink bias (a checkpointed buffer). Independent of + # expert-load bias correction; the update reduces the per-step token stats across the DP mesh so + # the buffer is identical on every rank. Like expert_bias it is a buffer, so it must + # be mirrored into net_ema (the model-EMA worker copies parameters only). + if self._uses_ema_router_bias: + update_router_biases(net=self.net, device_mesh=dp_mesh) + if self.config.ema.enabled: + sync_router_biases_to_ema(net=self.net, net_ema=self.net_ema) + def on_before_zero_grad( self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int ) -> None: diff --git a/cosmos_framework/model/generator/reasoner/cosmos3_edge/__init__.py b/cosmos_framework/model/generator/reasoner/cosmos3_edge/__init__.py new file mode 100644 index 00000000..1ddbcde8 --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/cosmos3_edge/__init__.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Framework-native Cosmos3-Edge VLM (renewed ``nvidia/Cosmos3-Edge``, no remote code). + +Importing this package registers ``model_type="cosmos3_edge"`` with the transformers +Auto classes, so ``AutoConfig.from_pretrained``/``AutoModelForImageTextToText`` resolve +the renewed HF snapshot without ``trust_remote_code`` (HFModel imports this package). +""" + +from transformers import AutoConfig, AutoModelForImageTextToText + +from cosmos_framework.model.generator.reasoner.cosmos3_edge.configuration_cosmos3_edge import ( + Cosmos3EdgeConfig, + Cosmos3EdgeProjectorConfig, + Cosmos3EdgeTextConfig, + Cosmos3EdgeVisionConfig, +) +from cosmos_framework.model.generator.reasoner.cosmos3_edge.modeling_cosmos3_edge import ( + Cosmos3EdgeForConditionalGeneration, + Cosmos3EdgeModel, + Cosmos3EdgePreTrainedModel, + Cosmos3EdgeTextModel, +) + +AutoConfig.register("cosmos3_edge", Cosmos3EdgeConfig, exist_ok=True) +AutoModelForImageTextToText.register(Cosmos3EdgeConfig, Cosmos3EdgeForConditionalGeneration, exist_ok=True) + +__all__ = [ + "Cosmos3EdgeConfig", + "Cosmos3EdgeForConditionalGeneration", + "Cosmos3EdgeModel", + "Cosmos3EdgePreTrainedModel", + "Cosmos3EdgeProjectorConfig", + "Cosmos3EdgeTextConfig", + "Cosmos3EdgeTextModel", + "Cosmos3EdgeVisionConfig", +] diff --git a/cosmos_framework/model/generator/reasoner/cosmos3_edge/configuration_cosmos3_edge.py b/cosmos_framework/model/generator/reasoner/cosmos3_edge/configuration_cosmos3_edge.py new file mode 100644 index 00000000..52bbb017 --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/cosmos3_edge/configuration_cosmos3_edge.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Configuration for the framework-native Cosmos3-Edge VLM (``model_type="cosmos3_edge"``). + +Parses the renewed ``nvidia/Cosmos3-Edge`` root ``config.json`` (native schema, no +remote code): ``text_config`` has 28 paired layers, each one attention block followed +by one MLP block. The ported modeling keeps the Nemotron-H 56-block layout (even +index = attention, odd = MLP) so state-dict keys stay +``model.language_model.layers.{0..55}.*``; the 56-block view is derived +(``layers_block_type``), never stored, so serialization round-trips the native +schema. ``vision_start_token_id``/``vision_end_token_id`` default to 20/21 (the old +remote-code values) when absent from ``config.json``. +""" + +from __future__ import annotations + +from transformers.configuration_utils import PretrainedConfig +from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig + + +class Cosmos3EdgeVisionConfig(Siglip2VisionConfig): + """SigLIP2 vision-tower config under the renewed repo's native model_type.""" + + model_type = "cosmos3_edge_vision" + + +class Cosmos3EdgeProjectorConfig(PretrainedConfig): + """PatchMerger projector config (native schema: ``merger_intermediate_size``).""" + + model_type = "cosmos3_edge_projector" + + def __init__( + self, + input_hidden_size: int = 1152, + merger_intermediate_size: int = 11520, + out_hidden_size: int = 2048, + spatial_merge_size: int = 2, + use_postshuffle_norm: bool = False, + **kwargs, + ) -> None: + self.input_hidden_size = input_hidden_size + self.merger_intermediate_size = merger_intermediate_size + self.out_hidden_size = out_hidden_size + self.spatial_merge_size = spatial_merge_size + self.use_postshuffle_norm = use_postshuffle_norm + super().__init__(**kwargs) + + @property + def merger_intermedia(self) -> int: + """Old remote-code field name; the vendored ``PatchMerger`` reads this.""" + return self.merger_intermediate_size + + +class Cosmos3EdgeTextConfig(PretrainedConfig): + """Text config in the renewed native schema (28 paired layers). + + Old-remote-code aliases needed by the ported modeling and the reused + ``MultiModalRotaryEmbedding`` (``rope_theta``, ``mrope_section``, + ``layers_block_type``, ``enable_mrope``) are exposed as read-only properties so + they never serialize. + """ + + model_type = "cosmos3_edge_text" + + def __init__( + self, + vocab_size: int = 131072, + hidden_size: int = 2048, + intermediate_size: int = 9216, + num_hidden_layers: int = 28, + num_attention_heads: int = 16, + num_key_value_heads: int = 8, + head_dim: int = 128, + hidden_act: str = "relu2", + attention_bias: bool = False, + mlp_bias: bool = False, + rms_norm_eps: float = 1e-5, + initializer_range: float = 0.02, + use_cache: bool = True, + max_position_embeddings: int = 131072, + attention_dropout: float = 0.0, + hidden_dropout: float = 0.0, + rope_parameters: dict | None = None, + sliding_window: int | None = None, + residual_in_fp32: bool = False, + pad_token_id: int = 0, + bos_token_id: int = 1, + eos_token_id: int = 11, + tie_word_embeddings: bool = False, + **kwargs, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.attention_bias = attention_bias + self.mlp_bias = mlp_bias + self.rms_norm_eps = rms_norm_eps + self.initializer_range = initializer_range + self.use_cache = use_cache + self.max_position_embeddings = max_position_embeddings + self.attention_dropout = attention_dropout + self.hidden_dropout = hidden_dropout + self.rope_parameters = ( + dict(rope_parameters) + if rope_parameters is not None + else { + "mrope_section": [24, 20, 20], + "rope_theta": 100000000, + "rope_type": "default", + } + ) + self.sliding_window = sliding_window + self.residual_in_fp32 = residual_in_fp32 + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def rope_theta(self) -> float: + return self.rope_parameters.get("rope_theta", 100000000) + + @property + def mrope_section(self) -> list[int]: + return self.rope_parameters.get("mrope_section", [24, 20, 20]) + + @property + def layers_block_type(self) -> list[str]: + """56-block Nemotron-H view of the 28 paired native layers ("*-" * 28).""" + return ["attention", "mlp"] * self.num_hidden_layers + + @property + def enable_mrope(self) -> bool: + """Architecture constant (old config.json pinned enable_rope/enable_mrope=true).""" + return True + + +class Cosmos3EdgeConfig(PretrainedConfig): + """Top-level Cosmos3-Edge VLM config: SigLIP2 tower + PatchMerger + Nemotron-H dense LM.""" + + model_type = "cosmos3_edge" + sub_configs = { + "vision_config": Cosmos3EdgeVisionConfig, + "projector_config": Cosmos3EdgeProjectorConfig, + "text_config": Cosmos3EdgeTextConfig, + } + + def __init__( + self, + text_config=None, + vision_config=None, + projector_config=None, + image_token_id: int = 19, + video_token_id: int = 18, + vision_start_token_id: int = 20, + vision_end_token_id: int = 21, + projector_hidden_size: int = 11520, + tie_word_embeddings: bool = False, + **kwargs, + ) -> None: + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + else: + self.vision_config = vision_config + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"]() + else: + self.text_config = text_config + + if isinstance(projector_config, dict): + self.projector_config = self.sub_configs["projector_config"](**projector_config) + elif projector_config is None: + self.projector_config = self.sub_configs["projector_config"]() + else: + self.projector_config = projector_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + self.projector_hidden_size = projector_hidden_size + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + +__all__ = [ + "Cosmos3EdgeConfig", + "Cosmos3EdgeProjectorConfig", + "Cosmos3EdgeTextConfig", + "Cosmos3EdgeVisionConfig", +] diff --git a/cosmos_framework/model/generator/reasoner/cosmos3_edge/modeling_cosmos3_edge.py b/cosmos_framework/model/generator/reasoner/cosmos3_edge/modeling_cosmos3_edge.py new file mode 100644 index 00000000..f107c9b0 --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/cosmos3_edge/modeling_cosmos3_edge.py @@ -0,0 +1,864 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Framework-native Cosmos3-Edge VLM modeling. + +Vendored/ported from the HF remote code ``modeling_nemotron_siglip2_h.py`` (old +``nvidia/Cosmos3-Edge``), byte-faithful for the dense path; Mamba/MoE and +hybrid-cache paths are dropped (the shipped config is pure dense; training runs +with ``use_cache=False``). Text generation uses ``GenerationMixin`` with +cache-free full-context recompute (see ``prepare_inputs_for_generation``). +Logits are cast to fp32 to match the reference. + +State-dict keys are contractual (safetensors loader, freeze regex +``model\\.visual\\.``, lr multipliers, HF export): ``model.visual.*``, +``model.projector.*``, ``model.language_model.{embeddings,layers.{0..55},norm_f}.*``, +``lm_head.weight``. +""" + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass +from typing import Any, Callable, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn import CrossEntropyLoss +from torch.nn.utils.rnn import pad_sequence +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.utils import ModelOutput, logging + +from cosmos_framework.model.generator.reasoner.cosmos3_edge.configuration_cosmos3_edge import ( + Cosmos3EdgeConfig, + Cosmos3EdgeTextConfig, +) +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.nemotron_3_dense_vl import ( + MultiModalRotaryEmbedding, + apply_rotary_pos_emb_partial, +) +from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.vision_siglip2 import ( + PatchMerger, + Siglip2VisionTransformer, + eager_attention_forward, + patch_merging_by_param, +) + +logger = logging.get_logger(__name__) + + +class Cosmos3EdgeRMSNorm(nn.Module): + """Port of NemotronHRMSNorm (fp32 compute, weight applied in fp32).""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight.to(torch.float32) * hidden_states).to(input_dtype) + + +class Cosmos3EdgeTextAttention(nn.Module): + """Port of NemotronHAttention (GQA + partial mRoPE), KV-cache path dropped.""" + + def __init__(self, config: Cosmos3EdgeTextConfig, layer_idx: Optional[int] = None) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.is_causal = True + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.head_dim * self.num_heads, self.hidden_size, bias=config.attention_bias) + self.scaling = self.head_dim**-0.5 + self.sliding_window = getattr(config, "sliding_window", None) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + packing_args: Optional[dict] = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_partial(query_states, key_states, cos, sin) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + if packing_args is not None: + cu_seqlens = packing_args["cu_seqlens"] + max_seqlen = packing_args["max_seqlen_in_batch"] + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + else: + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Cosmos3EdgeMLP(nn.Module): + """Port of NemotronHMLP.""" + + def __init__(self, config: Cosmos3EdgeTextConfig, layer_idx: Optional[int] = None) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor, padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + return self.down_proj(self.act_fn(self.up_proj(x))) + + +class Cosmos3EdgeBlock(nn.Module): + """Port of NemotronHBlock: shared pre-norm + attention/MLP mixer (dense pattern only).""" + + def __init__(self, config: Cosmos3EdgeTextConfig, layer_idx: int) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.residual_in_fp32 = config.residual_in_fp32 + self.norm = Cosmos3EdgeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.block_type = config.layers_block_type[layer_idx] + if self.block_type == "attention": + self.mixer = Cosmos3EdgeTextAttention(config, layer_idx=layer_idx) + elif self.block_type == "mlp": + self.mixer = Cosmos3EdgeMLP(config, layer_idx=layer_idx) + else: + raise ValueError(f"Invalid block type {self.block_type!r} at layer {layer_idx} (dense-only port)") + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + padding_mask: Optional[torch.Tensor] = None, + packing_args: Optional[dict] = None, + ) -> torch.Tensor: + # The remote code pins the block to the CUDA default stream ("avoid NaN issues + # when using multiple GPUs"); nullcontext off-CUDA so the port also runs on CPU. + if hidden_states.device.type == "cuda": + stream_ctx = torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)) + else: + stream_ctx = contextlib.nullcontext() + with stream_ctx: + residual = hidden_states + hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) + if self.residual_in_fp32: + residual = residual.to(torch.float32) + + if self.block_type == "attention": + hidden_states, _ = self.mixer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + packing_args=packing_args, + ) + else: + hidden_states = self.mixer(hidden_states, padding_mask=padding_mask) + + hidden_states = residual + hidden_states + return hidden_states + + +class Cosmos3EdgePreTrainedModel(PreTrainedModel): + config_class = Cosmos3EdgeConfig + base_model_prefix = "model" + input_modalities = ["image", "text"] + # Text block + vendored vision encoder layer: parallelize_vlm's FSDP block + # collector matches these type names. + _no_split_modules = ["Cosmos3EdgeBlock", "Siglip2EncoderLayer"] + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + + +@dataclass +class Cosmos3EdgeModelOutput(ModelOutput): + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Cosmos3EdgeCausalLMOutput(ModelOutput): + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +class Cosmos3EdgeTextModel(Cosmos3EdgePreTrainedModel): + """Port of NemotronHModel: embeddings + 56 paired blocks + norm_f, mRoPE.""" + + config_class = Cosmos3EdgeTextConfig + + def __init__(self, config: Cosmos3EdgeTextConfig) -> None: + super().__init__(config) + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + # 2 blocks per native paired layer: even index = attention, odd = MLP. + self.layers = nn.ModuleList( + [Cosmos3EdgeBlock(config, layer_idx=idx) for idx in range(2 * config.num_hidden_layers)] + ) + # Old config.json pinned enable_rope=true; rope is unconditional here. + self.rotary_emb = MultiModalRotaryEmbedding(config) + self.gradient_checkpointing = False + self.norm_f = Cosmos3EdgeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self._register_load_state_dict_pre_hook(self.load_hook) + self.post_init() + + def load_hook(self, state_dict, prefix, *args): + # Legacy-checkpoint tolerance carried over from the remote code. + for k in state_dict: + if "embedding." in k: + state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k) + break + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + padding_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, # accepted for API parity; no cache support + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + packing_args: Optional[dict] = None, + **kwargs: Any, # swallows visual_pos_masks etc., like the remote code + ) -> Union[Tuple, Cosmos3EdgeModelOutput]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + hidden_states = inputs_embeds + + if cache_position is None: + cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device) + # The hard coded `3` is for temporal, height and width (mRoPE). + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + elif position_ids.ndim == 3 and position_ids.shape[0] == 4: + position_ids = position_ids[1:] + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position) + + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for mixer_block in self.layers: + layer_mask = causal_mask if mixer_block.block_type == "attention" else None + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + hidden_states = self._gradient_checkpointing_func( + mixer_block.__call__, hidden_states, layer_mask, position_embeddings, padding_mask, packing_args + ) + else: + hidden_states = mixer_block( + hidden_states, + attention_mask=layer_mask, + position_embeddings=position_embeddings, + padding_mask=padding_mask, + packing_args=packing_args, + ) + + hidden_states = self.norm_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states] if v is not None) + + return Cosmos3EdgeModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + # Port of NemotronHModel._update_causal_mask (jamba lineage), cache-free. + def _update_causal_mask(self, attention_mask, input_tensor, cache_position): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + target_length = cache_position[-1] + 1 + causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.dim() == 2: + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0) + causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + ): + # Attend to all tokens in fully masked rows (left padding), required by + # SDPA's memory-efficient path. See pytorch/pytorch#110213. + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + +class Cosmos3EdgeModel(Cosmos3EdgePreTrainedModel): + """Port of NemotronSiglip2Model: visual + projector + language_model.""" + + _checkpoint_conversion_mapping = {} + accepts_loss_kwargs = False + # NOTE: explicit config_class, not a `config:` annotation — with + # `from __future__ import annotations` PreTrainedModel.__init_subclass__ would + # promote the (string) annotation to config_class and break Auto registration. + config_class = Cosmos3EdgeConfig + + def __init__(self, config: Cosmos3EdgeConfig) -> None: + super().__init__(config) + # Vision pos_embed interpolation reads the merge size off the vision config. + setattr(config.vision_config, "spatial_merge_size", config.projector_config.spatial_merge_size) + self.visual = Siglip2VisionTransformer._from_config(config.vision_config) + self.projector = PatchMerger(config.projector_config) + self.language_model = Cosmos3EdgeTextModel._from_config(config.text_config) + self.rope_deltas = None + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def set_decoder(self, decoder): + self.language_model = decoder + + def get_decoder(self): + return self.language_model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Timestamp-based mRoPE (Qwen3VL-style, but every video frame is its own t=1 grid).""" + # Videos are interleaved with per-frame timestamp tokens, so split grid rows per frame. + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + video_grid_thw[:, 0] = 1 + spatial_merge_size = self.config.projector_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None: + if cu_seqlens is None: + total_input_ids = input_ids + if attention_mask is None: + total_attention_mask = torch.ones_like(total_input_ids) + else: + total_attention_mask = attention_mask + total_attention_mask = total_attention_mask.to(total_input_ids.device) + else: + varlen = cu_seqlens.cpu().tolist() + total_input_ids = [] + total_attention_mask = [] + for start, end in zip(varlen[:-1], varlen[1:]): + total_input_ids.append(input_ids[:, start:end]) + total_attention_mask.append(torch.ones_like(total_input_ids[-1]).to(input_ids.device)) + position_ids = [] + image_index, video_index = 0, 0 + for input_ids, attention_mask in zip(total_input_ids, total_attention_mask): + input_ids = input_ids[attention_mask == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + # t_index is always 0 (llm_grid_t == 1; timestamps carry temporal info). + t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids.append(llm_positions.to(input_ids.device)) + mrope_position_deltas.append(llm_positions.max() + 1 - len(input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + if cu_seqlens is None: + # Batching on second dim (batch) for batch seq + position_ids = pad_sequence(position_ids, batch_first=False, padding_value=1) + else: + # Concat on last dim (seq_len) for packing seq + position_ids = torch.cat(position_ids, dim=-1) + return position_ids, mrope_position_deltas + else: + raise ValueError("input_ids is None") + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + num_image: int = 0, + ): + """Encode packed patches, 2x2-merge, project; split back into per-media features.""" + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + image_embeds, image_grid_thw = patch_merging_by_param( + image_embeds, image_grid_thw, merge_size=self.projector.spatial_merge_size + ) + image_embeds = image_embeds.view(-1, self.projector.spatial_merge_size**2, self.projector.input_hidden_size) + projected_hidden_states = self.projector(image_embeds) + split_sizes = image_grid_thw.prod(-1).tolist() + image_embeds = torch.split(projected_hidden_states, split_sizes) + image_embeddings = image_embeds[:num_image] + video_embeddings = image_embeds[num_image:] + return image_embeddings, video_embeddings + + def get_placeholder_mask( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_features: Optional[torch.FloatTensor] = None, + video_features: Optional[torch.FloatTensor] = None, + ): + if input_ids is None: + special_image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_image_mask = special_image_mask.all(-1) + special_video_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + special_video_mask = special_video_mask.all(-1) + else: + special_image_mask = input_ids == self.config.image_token_id + special_video_mask = input_ids == self.config.video_token_id + + n_image_tokens = special_image_mask.sum() + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if image_features is not None and inputs_embeds[special_image_mask].numel() != image_features.numel(): + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}" + ) + + n_video_tokens = special_video_mask.sum() + special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if video_features is not None and inputs_embeds[special_video_mask].numel() != video_features.numel(): + raise ValueError( + f"Videos features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}" + ) + + return special_image_mask, special_video_mask + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + packing_args: Optional[dict] = None, + **kwargs: Any, + ) -> Union[tuple, Cosmos3EdgeModelOutput]: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + # Remote code hardcodes padding_mask=None (config/tokenizer pad-id mismatch upstream). + padding_mask = None + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + image_mask = None + video_mask = None + skip_visual_embedding = False + if pixel_values is None and pixel_values_videos is None: + skip_visual_embedding = True + elif pixel_values is None: + final_pixel_value = pixel_values_videos + final_thw = video_grid_thw + num_image = 0 + elif pixel_values_videos is None: + final_pixel_value = pixel_values + final_thw = image_grid_thw + num_image = image_grid_thw.shape[0] + else: + final_pixel_value = torch.cat([pixel_values, pixel_values_videos], dim=0) + final_thw = torch.cat([image_grid_thw, video_grid_thw], dim=0) + num_image = image_grid_thw.shape[0] + + if not skip_visual_embedding: + image_embeds, video_embeds = self.get_image_features(final_pixel_value, final_thw, num_image) + elif skip_visual_embedding and self.training: + # Dummy forward to keep FSDP all-gather synchronised across ranks + merge_size = self.projector.spatial_merge_size + dummy_height = merge_size + dummy_width = merge_size + channels = self.visual.config.num_channels * self.visual.config.patch_size**2 + final_pixel_value = torch.zeros( + dummy_height * dummy_width, channels, device=inputs_embeds.device, dtype=self.visual.dtype + ) + final_thw = torch.tensor([[1, dummy_height, dummy_width]], device=inputs_embeds.device) + num_image = 1 + image_embeds, video_embeds = self.get_image_features(final_pixel_value, final_thw, num_image) + image_embeds = [image_embed[0:0] for image_embed in image_embeds] + else: + image_embeds = None + video_embeds = None + + if image_embeds is not None and len(image_embeds) > 0: + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if video_embeds is not None and len(video_embeds) > 0: + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if position_ids is None and self.config.text_config.enable_mrope: + attention_mask_tensor = ( + attention_mask if not isinstance(attention_mask, dict) else attention_mask["full_attention"] + ) + if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4: + attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2) + # Only apply conversion for floating point tensors (inverted masks) + if attention_mask_tensor.dtype.is_floating_point: + attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min + attention_mask_tensor = (1.0 - attention_mask_tensor).int() + + # No KV cache in this port: every forward is a prefill, so mRoPE position + # ids are always recomputed (rope_deltas kept for parity/inspection). + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + attention_mask=attention_mask_tensor, + ) + self.rope_deltas = rope_deltas + + # position_ids: [3, bsz, seq_len]; for visual tokens the dim-0 order is t, h, w. + return self.language_model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + padding_mask=padding_mask, + cache_position=cache_position, + use_cache=use_cache, + packing_args=packing_args, + **kwargs, + ) + + +class Cosmos3EdgeForConditionalGeneration(Cosmos3EdgePreTrainedModel, GenerationMixin): + """Port of NemotronSiglip2ForConditionCausalLM: training/eval forward + cache-free generate. + + ``GenerationMixin`` supplies ``generate()``; since the port has no KV-cache path + (attention takes no ``past_key_values`` and the forward output carries none), + ``prepare_inputs_for_generation`` below pins every decode step to a full-context + recompute. Correct but O(steps × context) — fine for eval, not for serving. + """ + + _tied_weights_keys = ["lm_head.weight"] + config_class = Cosmos3EdgeConfig + + def __init__(self, config: Cosmos3EdgeConfig) -> None: + super().__init__(config) + self.model = Cosmos3EdgeModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.get_input_embeddings() + + def set_input_embeddings(self, new_embeddings): + return self.model.set_input_embeddings(new_embeddings) + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def set_decoder(self, decoder): + self.model = decoder + + @property + def language_model(self): + return self.model.language_model + + @property + def multi_modal_projector(self): + return self.model.projector + + @property + def visual(self): + return self.model.visual + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + valid_input_len: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: bool = False, + **kwargs: Any, + ) -> Union[Tuple, Cosmos3EdgeCausalLMOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Sequence packing: valid_input_len drives the varlen (cu_seqlens) attention path. + packing_args = None + if valid_input_len is not None: + batch_size = valid_input_len.shape[0] + input_ids_list = [] + for i in range(batch_size): + valid_len = valid_input_len[i].item() + cur_input_ids = input_ids[i : i + 1, :valid_len].clone() + input_ids_list.append(cur_input_ids) + cu_seqlens = torch.cumsum(valid_input_len, dim=0).to(torch.int32) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + max_seqlen_in_batch = torch.max(valid_input_len).cpu().item() + + packing_args = { + "cu_seqlens": cu_seqlens, + "max_seqlen_in_batch": max_seqlen_in_batch, + } + input_ids = torch.cat(input_ids_list, dim=1) + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + cache_position=cache_position, + use_cache=use_cache, + packing_args=packing_args, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + hidden_states = outputs[0] + + logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float() + + loss = None + if labels is not None: + labels = labels.to(logits.device) + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + loss_fct = CrossEntropyLoss() + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return Cosmos3EdgeCausalLMOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[Any] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Cache-free generation: every decode step is a full-context recompute. + + The port has no KV-cache path, so the full ``input_ids`` (and vision + tensors) are re-fed each step. ``cache_position`` must be reset to + ``None``: the single-index decode value ``generate()`` maintains assumes + a cache and would zero out the mask in ``_update_causal_mask``. + """ + if inputs_embeds is not None: + raise ValueError( + "Cosmos3EdgeForConditionalGeneration.generate() only supports input_ids " + "(mRoPE position ids and vision placeholders are derived from token ids)" + ) + kwargs.pop("position_ids", None) # mRoPE position ids are recomputed in forward + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "use_cache": False, + "cache_position": None, + **{k: v for k, v in kwargs.items() if v is not None}, + } + + +__all__ = [ + "Cosmos3EdgeForConditionalGeneration", + "Cosmos3EdgeModel", + "Cosmos3EdgePreTrainedModel", + "Cosmos3EdgeTextModel", +] diff --git a/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/reasoner_multimodal_utils.py b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/reasoner_multimodal_utils.py new file mode 100644 index 00000000..a92db2ff --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/reasoner_multimodal_utils.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Image/video-conditioned prefill helpers for the Nemotron 3 Dense VL reasoner. + +Mirrors ``reasoner/qwen3_vl/utils.py::prepare_multimodal_reasoner_inputs`` for +the Edge reasoner. The multimodal rope index (``get_rope_index``) and the +placeholder-mask helper are algorithmically identical to the Qwen3-VL path — +they depend only on ``model.config`` token ids + ``image_grid_thw`` / +``video_grid_thw`` — so they are reused directly. The Nemotron-specific parts +are: + +* the vision tower is a SigLIP2 encoder (``causal_lm.visual`` is a + :class:`NemotronSiglip2VisionEncoder`) whose ``get_image_features`` returns + the projected ``[N_merged_patches, hidden]`` embeddings directly. SigLIP2 has + no temporal attention, so a video is just its frames stacked along the grid's + temporal axis (``video_grid_thw`` rows carry ``t>1``); the same + ``get_image_features`` encodes them patch-by-patch, and +* there are **no deepstack** visual embeds (returned list is empty), so the + shared ``_impl_reasoner_forward`` deepstack path is a no-op. + +Image and video are mutually exclusive: exactly one of the +(``pixel_values``, ``image_grid_thw``) / (``pixel_values_videos``, +``video_grid_thw``) pairs is consumed per call. +""" +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from cosmos_framework.model.generator.reasoner.qwen3_vl.utils import ( + get_placeholder_mask, + get_rope_index, +) + + +def prepare_multimodal_reasoner_inputs( + causal_lm: Any, + input_ids: torch.Tensor, # [B,T_prompt] + pixel_values: torch.Tensor | None = None, # [N_patches, C*patch*patch] + image_grid_thw: torch.Tensor | None = None, # [num_images,3] + pixel_values_videos: torch.Tensor | None = None, # [N_patches, C*patch*patch] + video_grid_thw: torch.Tensor | None = None, # [num_videos,3] + attention_mask: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor], torch.Tensor, torch.Tensor]: + """Build the image/video-conditioned prefill inputs for the reasoner-only AR path. + + Returns ``(inputs_embeds, visual_pos_masks, deepstack_visual_embeds, + position_ids, mrope_position_deltas)`` matching the contract consumed by + :func:`unified_mot._impl_generate_reasoner_text`. ``deepstack_visual_embeds`` + is always ``[]`` for Nemotron (no deepstack tower). + + The video recipe mirrors the image recipe but routes the pre-processed + ``pixel_values_videos`` / ``video_grid_thw`` through the shared SigLIP2 + ``get_image_features`` (temporal-agnostic), the video placeholder mask, and + the ``video_grid_thw`` rope index. + """ + is_video = pixel_values_videos is not None or video_grid_thw is not None + if is_video and (pixel_values is not None or image_grid_thw is not None): + raise ValueError( + "prepare_multimodal_reasoner_inputs conditions on one medium at a time: " + "pass the image pair OR the video pair, not both." + ) + if is_video: + if pixel_values_videos is None or video_grid_thw is None: + raise ValueError( + "prepare_multimodal_reasoner_inputs requires pixel_values_videos and video_grid_thw together." + ) + elif pixel_values is None or image_grid_thw is None: + raise ValueError("prepare_multimodal_reasoner_inputs requires pixel_values and image_grid_thw.") + if not hasattr(causal_lm, "visual") or causal_lm.visual is None: + raise ValueError("Nemotron reasoner has no vision tower (`causal_lm.visual`).") + + inputs_embeds = causal_lm.model.embed_tokens(input_ids).clone() # [B,T_prompt,hidden] + + if is_video: + pixel_values_videos = pixel_values_videos.to(device=inputs_embeds.device) + video_grid_thw = video_grid_thw.to(device=inputs_embeds.device) + # SigLIP2 has no temporal modeling: a video is its frames as extra grid + # patches, so the image encoder handles it unchanged. Returns [N_merged_patches, hidden]. + video_embeds = causal_lm.visual.get_image_features(pixel_values_videos, video_grid_thw) + video_embeds = video_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + + _image_mask, video_mask = get_placeholder_mask( + causal_lm, + input_ids, + inputs_embeds=inputs_embeds, + video_features=video_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # [B,T_prompt,hidden] + visual_pos_masks = video_mask[..., 0] # [B,T_prompt] + else: + pixel_values = pixel_values.to(device=inputs_embeds.device) + image_grid_thw = image_grid_thw.to(device=inputs_embeds.device) + + # SigLIP2 encode -> spatial-merge -> project. Returns [N_merged_patches, hidden]. + image_embeds = causal_lm.visual.get_image_features(pixel_values, image_grid_thw) + image_embeds = image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + + image_mask, _video_mask = get_placeholder_mask( + causal_lm, + input_ids, + inputs_embeds=inputs_embeds, + image_features=image_embeds, + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # [B,T_prompt,hidden] + visual_pos_masks = image_mask[..., 0] # [B,T_prompt] + + deepstack_visual_embeds: list[torch.Tensor] = [] + + position_ids, mrope_position_deltas = get_rope_index( + causal_lm, + input_ids=input_ids, + image_grid_thw=None if is_video else image_grid_thw, + video_grid_thw=video_grid_thw if is_video else None, + attention_mask=attention_mask, + ) + + return inputs_embeds, visual_pos_masks, deepstack_visual_embeds, position_ids, mrope_position_deltas diff --git a/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/vision_siglip2.py b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/vision_siglip2.py new file mode 100644 index 00000000..706987e5 --- /dev/null +++ b/cosmos_framework/model/generator/reasoner/nemotron_3_dense_vl/vision_siglip2.py @@ -0,0 +1,459 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""SigLIP2 vision encoder + PatchMerger for the Cosmos3-Edge (Nemotron 3 Dense VL) reasoner. + +Vendored from the HF `nvidia/Cosmos3-Edge` remote code +(`modeling_nemotron_siglip2_h.py`): the custom grid_thw / cu_seqlens SigLIP2 +vision transformer (naflex-style packed patches) and the Qwen3-style +`PatchMerger` projector, plus `patch_merging_by_param`. Kept byte-faithful so +outputs match the HF reference numerically. +""" +from __future__ import annotations + +import math +from typing import Callable, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.activations import ACT2FN +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs: Any, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Siglip2VisionEmbeddings(nn.Module): + def __init__(self, config: Siglip2VisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Linear( + in_features=config.num_channels * self.patch_size * self.patch_size, + out_features=self.embed_dim, + ) + + self.num_patches = config.num_patches + self.position_embedding_size = int(self.num_patches**0.5) + self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim) + + def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: + """ + Args: + pixel_values (`torch.FloatTensor`): + Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size) + """ + + # Apply patch embeddings to already patchified pixel values + patch_embeds = self.patch_embedding(pixel_values) + + return patch_embeds + +class Siglip2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + self.num_key_value_groups = 1 + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + **kwargs, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Input shape: total_pixel_value x hidden_size""" + + seq_length, embed_dim = hidden_states.shape + + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(seq_length, self.num_heads, self.head_dim).transpose(0, 1).unsqueeze(0) + keys = keys.view(seq_length, self.num_heads, self.head_dim).transpose(0, 1).unsqueeze(0) + values = values.view(seq_length, self.num_heads, self.head_dim).transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + if self.config._attn_implementation == "flash_attention_2": + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + attn_output, _ = attention_interface( + self, + queries, + keys, + values, + attention_mask=None, + is_causal=self.is_causal, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + ) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (queries, keys, values) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + is_causal=self.is_causal, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output + +class Siglip2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class Siglip2EncoderLayer(nn.Module): + def __init__(self, config: Siglip2VisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.self_attn = Siglip2Attention(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = Siglip2MLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + **kwargs, + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + cu_seqlens=cu_seqlens, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + +class Siglip2Encoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`Siglip2EncoderLayer`]. + + Args: + config: Siglip2Config + """ + + def __init__(self, config: Siglip2VisionConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + # Ignore copy + def forward( + self, + inputs_embeds: torch.Tensor, + cu_seqlens: torch.Tensor, + **kwargs, + ): + hidden_states = inputs_embeds + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_states, + cu_seqlens, + **kwargs, + ) + + return hidden_states + +class Siglip2VisionTransformer(PreTrainedModel): + config: Siglip2VisionConfig + main_input_name = "pixel_values" + base_model_prefix = "siglip_vit" + supports_gradient_checkpointing = True + + _no_split_modules = [ + "Siglip2VisionEmbeddings", + "Siglip2EncoderLayer", + "Siglip2MultiheadAttentionPoolingHead", + ] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + _can_record_outputs = { + "hidden_states": Siglip2EncoderLayer, + "attentions": Siglip2Attention, + } + def __init__(self, config: Siglip2VisionConfig): + super().__init__(config) + self.config = config + embed_dim = config.hidden_size + + self.embeddings = Siglip2VisionEmbeddings(config) + self.encoder = Siglip2Encoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.num_grid_per_side = self.embeddings.position_embedding_size + + def get_position_embedding(self, grid_thw: torch.Tensor) -> torch.Tensor: + # prepare for interpolation + positional_embedding = self.embeddings.position_embedding.weight.reshape( + self.embeddings.position_embedding_size, self.embeddings.position_embedding_size, -1 + ).permute(2, 0, 1).unsqueeze(0) + + total_tokens = int(torch.prod(grid_thw, dim=1).sum().item()) + embed_dim = self.embeddings.embed_dim + # create a resized positional embedding of size (total_tokens, embed_size) to hold positional_embedding for all visual inputs + resized_positional_embeddings = torch.empty((total_tokens, embed_dim), dtype=positional_embedding.dtype, device=grid_thw.device) + offset = 0 + for t, height, width in grid_thw: + resized_embeddings = F.interpolate( + positional_embedding, + size=(height.cpu().item(), width.cpu().item()), + mode='bilinear', + align_corners=False, + antialias=True + ) + resized_embeddings = resized_embeddings.reshape(embed_dim, -1).transpose(0, 1) + + num_spatial_tokens = height * width + total_block_tokens = t * num_spatial_tokens + + resized_positional_embeddings[offset: offset + total_block_tokens] = resized_embeddings.repeat(t, 1) + offset += total_block_tokens + assert offset == resized_positional_embeddings.shape[0] + return resized_positional_embeddings + + def forward( + self, + pixel_values: torch.FloatTensor, + grid_thw: torch.LongTensor, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ): + r""" + spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): + Tensor containing the spatial dimensions (height, width) of the input images. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + hidden_states = self.embeddings(pixel_values) + positional_embeddings = self.get_position_embedding(grid_thw) + hidden_states = hidden_states + positional_embeddings + + # View pixel_values as packed multi-visual input and build cu_seqlens (migrated from qwen3-vl). + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + # Select dtype based on the following factors: + # - FA2 requires that cu_seqlens_q must have dtype int32 + # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw + # See https://github.com/huggingface/transformers/pull/34852 for more information + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + hidden_states = self.encoder( + inputs_embeds=hidden_states, + cu_seqlens=cu_seqlens, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + last_hidden_state = self.post_layernorm(hidden_states) + return last_hidden_state + + +class PatchMerger(nn.Module): + + def __init__(self, config: "ProjectorConfig", use_postshuffle_norm=False) -> None: + super().__init__() + self.spatial_merge_size = config.spatial_merge_size + self.hidden_size = config.input_hidden_size * (self.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.input_hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, config.merger_intermedia) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(config.merger_intermedia, config.out_hidden_size) + self.input_hidden_size = config.input_hidden_size + self.out_hidden_size = config.out_hidden_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + def init_weights(self): + # init weight with he_normal + # init bias with zero + # init layernorm with standard gaussian distribution + nn.init.kaiming_uniform_(self.linear_fc1.weight, a=math.sqrt(5)) + nn.init.kaiming_uniform_(self.linear_fc2.weight, a=math.sqrt(5)) + nn.init.zeros_(self.linear_fc1.bias) + nn.init.zeros_(self.linear_fc2.bias) + self.norm.reset_parameters() + + + +def patch_merging_by_param(image_embeds, image_grid_thw, merge_size=2): + """ + image_embeds: [Total_Patches, C] -> e.g. [2008, 1152] + image_grid_thw: [Num_Media, 3] -> e.g. [[1, 26, 38], [1, 34, 30]] + merge_size: the spatial merge factor from config, e.g. 2 + """ + new_embeds_list = [] + new_grid_thw_list = [] + curr_idx = 0 + + C = image_embeds.shape[-1] + + for i in range(image_grid_thw.shape[0]): + # Current media's T, H, W (here H, W are patch counts). + t, h, w = image_grid_thw[i].tolist() + num_patches = t * h * w + + # 1. Slice out this media's features [T*H*W, C]. + media_seq = image_embeds[curr_idx : curr_idx + num_patches] + curr_idx += num_patches + + # 2. Restore the 3D structure [T, H, W, C]. + x = media_seq.view(t, h, w, C) + + # 3. Spatially merge (2x2 blocks) with einops. + # Dimension transform: + # b=t, h=(h'/ms * ms), w=(w'/ms * ms) + # -> [t, h/ms, ms, w/ms, ms, c] -> [t, h/ms, w/ms, (ms*ms*c)] + # Note: we follow the Qwen2-VL order — h1 w1 come before C. + x = rearrange( + x, + 't (h h1) (w w1) c -> t h w (h1 w1 c)', + h1=merge_size, + w1=merge_size + ) + + # 4. Flatten back to a sequence [T * (H/ms) * (W/ms), C * ms^2]. + new_embeds_list.append(x.reshape(-1, x.shape[-1])) + + # 5. Update grid info: T unchanged, H and W shrink. + new_grid_thw_list.append([t, h // merge_size, w // merge_size]) + + # Concatenate all media back together. + image_embeds_merged = torch.cat(new_embeds_list, dim=0) + image_grid_thw_merged = torch.tensor(new_grid_thw_list, device=image_grid_thw.device) + + return image_embeds_merged, image_grid_thw_merged + + +class NemotronSiglip2VisionEncoder(nn.Module): + """Vision tower: SigLIP2 transformer + PatchMerger. Mirrors HF get_image_features.""" + def __init__(self, vision_config: Siglip2VisionConfig, projector_config): + super().__init__() + self.spatial_merge_size = projector_config.spatial_merge_size + setattr(vision_config, "spatial_merge_size", self.spatial_merge_size) + self.visual = Siglip2VisionTransformer(vision_config) + self.projector = PatchMerger(projector_config) + + @property + def dtype(self): + return self.visual.post_layernorm.weight.dtype + + def get_image_features(self, pixel_values, image_grid_thw): + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + image_embeds, image_grid_thw = patch_merging_by_param( + image_embeds, image_grid_thw, merge_size=self.projector.spatial_merge_size) + image_embeds = image_embeds.view(-1, self.projector.spatial_merge_size**2, self.projector.input_hidden_size) + projected = self.projector(image_embeds) + return projected diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py index bc570db3..3831dd71 100644 --- a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/moe_test.py @@ -10,6 +10,95 @@ Qwen3VLMoeTextConfig, ) from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import create_text_experts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import ( + AuxLossFreeLoadBalancingConfig, + CosineRouter, + CosineRouterConfig, + Qwen3VLMoeTextSparseMoeBlock, +) + + +def test_router_activation_defaults_to_softmax() -> None: + router = CosineRouter(CosineRouterConfig(), hidden_size=2) + router_logits = torch.tensor([[2.0, 1.0, 0.0]]) # [1,3] + expert_bias = torch.tensor([-0.5, 0.3, 0.0]) # [3] + + router_scores = router.get_scores(router_logits) # [1,3] + biased_selection_scores = router.apply_selection_bias( + router_logits=router_logits, + router_scores=router_scores, + expert_bias=expert_bias, + ) # [1,3] + + torch.testing.assert_close(router_scores, torch.softmax(router_logits, dim=-1)) + torch.testing.assert_close(biased_selection_scores, router_logits + expert_bias.unsqueeze(0)) + + +def test_sigmoid_router_bias_is_added_in_score_space() -> None: + router = CosineRouter(CosineRouterConfig(activation="sigmoid"), hidden_size=2) + router_logits = torch.tensor([[2.0, 1.0, 0.0]]) # [1,3] + expert_bias = torch.tensor([-0.5, 0.3, 0.0]) # [3] + router_scores = router.get_scores(router_logits) # [1,3] + + biased_selection_scores = router.apply_selection_bias( + router_logits=router_logits, + router_scores=router_scores, + expert_bias=expert_bias, + ) # [1,3] + selected_experts = torch.topk(biased_selection_scores, k=1, dim=-1).indices # [1,1] + + torch.testing.assert_close(biased_selection_scores, torch.sigmoid(router_logits) + expert_bias.unsqueeze(0)) + assert selected_experts.item() == 1 + + +def test_sigmoid_router_scores_are_normalized_for_lbl_and_metrics() -> None: + router = CosineRouter(CosineRouterConfig(activation="sigmoid"), hidden_size=2) + router_logits = torch.tensor([[2.0, 1.0, 0.0], [-1.0, 0.0, 3.0]]) # [2,3] + router_scores = router.get_scores(router_logits) # [2,3] + + routing_probabilities = router.normalize_scores(router_scores) # [2,3] + + torch.testing.assert_close(routing_probabilities.sum(dim=-1), torch.ones(2)) + torch.testing.assert_close( + routing_probabilities, + torch.sigmoid(router_logits) / torch.sigmoid(router_logits).sum(dim=-1, keepdim=True), + ) + + +def test_router_rejects_unknown_activation() -> None: + try: + CosineRouter(CosineRouterConfig(activation="relu"), hidden_size=2) + except ValueError as error: + assert "Unsupported router activation" in str(error) + else: + raise AssertionError("CosineRouter accepted an unsupported activation") + + +def test_aux_loss_free_controller_uses_block_config() -> None: + model_config = Qwen3VLMoeTextConfig( + hidden_size=8, + moe_intermediate_size=4, + num_experts=4, + num_experts_per_tok=2, + hidden_act="silu", + ) + controller_config = AuxLossFreeLoadBalancingConfig( + enabled=True, + update_speed=0.25, + max_bias=None, + ) + block = Qwen3VLMoeTextSparseMoeBlock( + model_config, + aux_loss_free_load_balancing_config=controller_config, + ) + block.tokens_per_expert.copy_(torch.tensor([0.0, 1.0, 2.0, 3.0])) # [4] + + block.update_bias() + + torch.testing.assert_close(block.expert_bias, torch.tensor([0.25, 0.25, -0.25, -0.25])) # [4] + assert block.aux_loss_free_load_balancing_config is controller_config + assert "expert_bias" in block.state_dict() + assert "tokens_per_expert" not in block.state_dict() def run_moe(mod: nn.Module, hidden_states: torch.Tensor, topk_scores: torch.Tensor, expert_indices: torch.Tensor): diff --git a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py index 0add9ba7..a81d7448 100644 --- a/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py +++ b/cosmos_framework/model/generator/reasoner/qwen3_vl_moe/qwen3_vl_moe.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: OpenMDW-1.1 import functools +import math from dataclasses import dataclass from typing import Any, Callable, NamedTuple, Optional, Union @@ -72,6 +73,189 @@ class LBLMetadata(NamedTuple): mean_router_prob_per_expert: torch.Tensor +@dataclass(frozen=True) +class AuxLossFreeLoadBalancingConfig: + """Configuration for generation-tower aux-loss-free load balancing. + + Attributes: + enabled: Enable the gradient-free expert-bias controller. + update_speed: Per-step bias nudge applied via + ``sign(average_load - expert_load)``. + max_bias: Symmetric bias-magnitude limit. ``None`` disables clamping. + """ + + enabled: bool = False + update_speed: float = 0.001 + max_bias: float | None = 0.5 + + +@dataclass(frozen=True) +class CosineRouterConfig: + """Configuration for routing in a sparse MoE block. + + Attributes: + enabled: Use cosine-similarity logits instead of the standard dot-product + gate. Activation selection applies regardless of this setting. + activation: Function applied to router logits. ``"softmax"`` produces a + distribution across experts; ``"sigmoid"`` produces independent + expert affinities. + input_centering: Center cosine-router inputs using the current batch + mean (``"batch_mean"``), a running EMA mean (``"ema"``), or no + centering (``"none"``). + ema_momentum: Momentum used to update the input-centering EMA. + clamp_temperature: Prevent the learned cosine-router temperature from + exceeding its initialization value. + """ + + enabled: bool = False + activation: str = "softmax" + input_centering: str = "batch_mean" + ema_momentum: float = 0.99 + clamp_temperature: bool = True + + +class CosineRouter(nn.Module): + """Cosine-router behavior and state.""" + + def __init__(self, config: CosineRouterConfig, hidden_size: int) -> None: + super().__init__() + self.config: CosineRouterConfig = config + self.enabled: bool = config.enabled + self.activation: str = config.activation + if self.activation not in {"softmax", "sigmoid"}: + raise ValueError(f"Unsupported router activation: {self.activation!r}") + self.input_centering: str = config.input_centering + if self.input_centering not in {"batch_mean", "ema", "none"}: + raise ValueError(f"Unsupported router input centering: {self.input_centering!r}") + self.ema_momentum: float = config.ema_momentum + self.clamp_temperature: bool = config.clamp_temperature + self.hidden_size: int = hidden_size + self.initial_log_temperature: float = self.log_temperature_init(hidden_size) + + if self.enabled: + self.log_temperature = nn.Parameter( + torch.full((hidden_size,), self.initial_log_temperature, dtype=torch.float32) + ) # [D] + self._init_input_centering_buffers() + + @staticmethod + def log_temperature_init(hidden_size: int) -> float: + """Return ``log(sqrt(hidden_size))`` for pretrained-router warm starts.""" + return 0.5 * math.log(hidden_size) + + def _init_input_centering_buffers(self, buffer_device: torch.device | None = None) -> None: + """Initialize persistent and per-step EMA input-centering buffers.""" + if self.input_centering != "ema": + return + + # Persistent EMA estimate restored from checkpoints. + self.register_buffer( + "router_bias", + torch.zeros(self.hidden_size, dtype=torch.float32, device=buffer_device), # [D] + ) + # Per-step accumulators are temporary and intentionally omitted from checkpoints. + self.register_buffer( + "router_bias_sum", + torch.zeros(self.hidden_size, dtype=torch.float32, device=buffer_device), # [D] + persistent=False, + ) + self.register_buffer( + "router_bias_count", + torch.zeros(1, dtype=torch.float32, device=buffer_device), # [1] + persistent=False, + ) + + def forward(self, hidden_states: torch.Tensor, gate: nn.Linear) -> torch.Tensor: + """Compute standard or cosine router logits. [N,D] -> [N,E]""" + if not self.enabled: + return gate(hidden_states) # [N,E] + + x = hidden_states.to(torch.float32) # [N,D] + if self.input_centering == "ema": + if self.training: + with torch.no_grad(): + router_bias_delta = x.detach().sum(dim=0) # [D] + self.router_bias_sum.add_(router_bias_delta) # [D] + self.router_bias_count.add_(x.shape[0]) # [1] + x = x - self.router_bias # [N,D] + elif self.input_centering == "batch_mean" and x.shape[0] > 1: + x = x - x.mean(dim=0, keepdim=True) # [N,D] + + x = F.normalize(x, dim=-1) # [N,D] + log_temperature = self.log_temperature # [D] + if self.clamp_temperature: + log_temperature = log_temperature.clamp(max=self.initial_log_temperature) # [D] + temperature = log_temperature.exp().to(torch.float32) # [D] + gate_weight = F.normalize(gate.weight.to(torch.float32), dim=-1) # [E,D] + x = x * temperature # [N,D] + router_logits = F.linear(x, gate_weight) # [N,E] + return router_logits.to(hidden_states.dtype) # [N,E] + + def get_scores(self, router_logits: torch.Tensor) -> torch.Tensor: + """Apply the configured activation to router logits. [N,E] -> [N,E]""" + if self.activation == "sigmoid": + return torch.sigmoid(router_logits.to(torch.float32)) # [N,E] + return F.softmax(router_logits, dim=-1, dtype=torch.float32) # [N,E] + + @staticmethod + def normalize_scores(scores: torch.Tensor) -> torch.Tensor: + """Normalize independent router scores into a probability distribution. [N,E] -> [N,E]""" + score_sums = scores.sum(dim=-1, keepdim=True) # [N,1] + normalized_scores = scores / score_sums.clamp_min(torch.finfo(scores.dtype).tiny) # [N,E] + uniform_scores = torch.full_like(scores, 1.0 / scores.shape[-1]) # [N,E] + return torch.where(score_sums > 0, normalized_scores, uniform_scores) # [N,E] + + def apply_selection_bias( + self, + router_logits: torch.Tensor, + router_scores: torch.Tensor, + expert_bias: torch.Tensor, + ) -> torch.Tensor: + """Add loss-free expert bias in the activation-appropriate selection space.""" + if self.activation == "sigmoid": + return router_scores + expert_bias.unsqueeze(0) # [N,E] + return router_logits + expert_bias.unsqueeze(0) # [N,E] + + def update_ema_bias( + self, + device_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + ) -> None: + """Update the token-constant EMA bias from accumulated per-step statistics.""" + if self.input_centering != "ema": + return + + with torch.no_grad(): + total_sum = self.router_bias_sum.clone() # [D] + total_count = self.router_bias_count.clone() # [1] + if device_mesh is not None: + from torch.distributed.tensor import DTensor, Partial + + total_sum = DTensor.from_local( + total_sum, + device_mesh=device_mesh, + placements=[Partial()] * device_mesh.ndim, + ).full_tensor() # [D] + total_count = DTensor.from_local( + total_count, + device_mesh=device_mesh, + placements=[Partial()] * device_mesh.ndim, + ).full_tensor() # [1] + if total_count.item() > 0: + batch_mean = total_sum / total_count # [D] + self.router_bias.mul_(self.ema_momentum).add_(batch_mean * (1.0 - self.ema_momentum)) # [D] + self.router_bias_sum.zero_() # [D] + self.router_bias_count.zero_() # [1] + + def init_weights(self, buffer_device: torch.device | None) -> None: + """Reset router parameters and buffers.""" + if not self.enabled: + return + + with torch.no_grad(): + self.log_temperature.fill_(self.initial_log_temperature) # [D] + self._init_input_centering_buffers(buffer_device=buffer_device) + + @use_kernel_forward_from_hub("RMSNorm") class Qwen3VLMoeTextRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): @@ -94,13 +278,37 @@ def extra_repr(self): class Qwen3VLMoeTextSparseMoeBlock(nn.Module): - def __init__(self, config, noisy_gating: bool = False): + def __init__( + self, + config: Qwen3VLMoeTextConfig, + noisy_gating: bool = False, + cosine_router_config: CosineRouterConfig | None = None, + aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + ) -> None: + load_balancing_config = aux_loss_free_load_balancing_config or AuxLossFreeLoadBalancingConfig() super().__init__() + assert not (noisy_gating and load_balancing_config.enabled), ( + "Noisy gating and aux-loss-free load balancing cannot be enabled at the same time." + ) self.config = config self.hidden_size = config.hidden_size self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + # Aux-loss-free load balancing (gen tower only). A + # gradient-free per-expert bias is added to the top-k SELECTION score. It is + # updated by a sign-based controller in update_bias() driven from the + # trainer's on_before_optimizer_step hook. ``expert_bias`` is persistent + # so it is saved to / restored from checkpoints; ``tokens_per_expert`` is + # a transient per-step accumulator and is not persisted. + self.aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig = load_balancing_config + if self.aux_loss_free_load_balancing_config.enabled: + self.register_buffer("expert_bias", torch.zeros(self.num_experts)) + self.register_buffer( + "tokens_per_expert", + torch.zeros(self.num_experts), + persistent=False, + ) # Noisy top-k gating (Shazeer 2017): a second projection produces a # per-token, per-expert noise magnitude. During training the top-k # selection is made on clean_logits + N(0,1) * softplus(gate_noise(x)), @@ -110,6 +318,23 @@ def __init__(self, config, noisy_gating: bool = False): self.noisy_gating = noisy_gating if noisy_gating: self.gate_noise = nn.Linear(config.hidden_size, config.num_experts, bias=False) + + # Cosine router (gen-tower only). Replaces the raw dot-product logits + # x·W_e with a cosine similarity between the (token-mean-centered, + # L2-normalized) router input and the L2-normalized gate rows, scaled by + # a learnable temperature. This targets three pathologies measured in the + # gen tower's router: + # 1. token-constant component of the router input — removed by + # per-batch mean-centering over the token axis; + # 2. low / inconsistent logit magnitude (gate-gain erosion under weight + # decay) — removed by normalizing the input to unit norm and folding + # all scale into the learnable temperature; + # 3. gate-row-norm dominance — removed by normalizing each + # gate row, so selection depends on direction only. + self.cosine_router: CosineRouter = CosineRouter( + config=cosine_router_config or CosineRouterConfig(), + hidden_size=config.hidden_size, + ) self.experts = create_text_experts(config, implementation_type="grouped_mm") # ── Heatmap tracking ────────────────────────────────────────────────────── @@ -184,7 +409,7 @@ def _update_moe_callback_stats( self, num_tokens_per_expert: torch.Tensor, num_tokens: torch.Tensor, - routing_weights: torch.Tensor, + routing_probabilities: torch.Tensor, expert_indices: torch.Tensor, ) -> None: # ── Heatmap + stability buffers ────────────────────────────────────── @@ -198,7 +423,7 @@ def _update_moe_callback_stats( # Summed (not meaned) so the callback can normalize by any window length. # 1e-9 prevents log(0) for near-zero probabilities. token_entropy = -torch.sum( - routing_weights * torch.log(routing_weights + ENTROPY_EPSILON), dim=-1 + routing_probabilities * torch.log(routing_probabilities + ENTROPY_EPSILON), dim=-1 ) # [num_tokens] self.sum_token_entropy.add_(token_entropy.sum().to(torch.float64)) # Per-token soft effective experts = exp(H(p_t)), bounded in [1, N]. @@ -244,29 +469,59 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, LBLMetadat assert hidden_states.ndim == 2, "hidden_states must be of shape (num_tokens, hidden_size)" num_tokens = hidden_states.shape[0] - router_logits = self.gate(hidden_states) # [num_tokens,num_experts] - # Clean router distribution. Always used for monitoring (entropy/stability - # buffers) and the load-balancing-loss probability term so those stay - # comparable regardless of whether noisy gating is enabled. - routing_weights = torch.nn.functional.softmax( - router_logits, dim=-1, dtype=torch.float32 - ) # [num_tokens,num_experts] + router_logits = self.cosine_router(hidden_states, self.gate) # [num_tokens,num_experts] + # Clean router scores. For sigmoid routing these are independent expert + # affinities, while softmax routing produces a probability distribution. + routing_scores = self.cosine_router.get_scores(router_logits) # [num_tokens,num_experts] + # Monitoring and LBL require a probability distribution. Normalizing sigmoid scores + # softmax takes the existing path unchanged for backward compatibility. + if self.cosine_router.activation == "sigmoid": + routing_probabilities = self.cosine_router.normalize_scores(routing_scores) # [num_tokens,num_experts] + else: + routing_probabilities = routing_scores # [num_tokens,num_experts] # Noisy top-k gating: only the expert *selection* (and the combine - # weights over the selected experts) sees the noise. When noise is off - # or at eval time, selection_weights == routing_weights, so behavior is - # identical to plain top-k gating. + # weights over the selected experts) sees the noise. Monitoring and LBL + # always use the clean routing_probabilities above. if self.noisy_gating and self.training: noise_std = torch.nn.functional.softplus(self.gate_noise(hidden_states)) # [num_tokens,num_experts] - noisy_logits = router_logits + torch.randn_like(router_logits) * noise_std - selection_weights = torch.nn.functional.softmax(noisy_logits, dim=-1, dtype=torch.float32) + selection_logits = router_logits + torch.randn_like(router_logits) * noise_std # [num_tokens,num_experts] + selection_scores = self.cosine_router.get_scores(selection_logits) # [num_tokens,num_experts] else: - selection_weights = routing_weights - - expert_weights, expert_indices = torch.topk(selection_weights, self.top_k, dim=-1) - # expert_weights: [num_tokens,top_k], expert_indices: [num_tokens,top_k] + selection_logits = router_logits # [num_tokens,num_experts] + selection_scores = routing_scores # [num_tokens,num_experts] + + # Gate-natural top-k: the selection the gate makes on its own. It is the + # dispatch selection when bias correction is off, and is always used as + # the load-balancing-loss count (LBL must see the natural distribution it + # is trying to push toward uniform, not the bias-balanced one). + natural_weights, natural_indices = torch.topk( + selection_scores, self.top_k, dim=-1 + ) # [num_tokens,top_k], [num_tokens,top_k] + + if self.aux_loss_free_load_balancing_config.enabled: + # Aux-loss-free load balancing changes selection only. DeepSeek-style + # Sigmoid routing adds bias to independent activated scores. Softmax + # scores are compressed and coupled by sum-to-one normalization, so + # bias is applied in logit space to avoid disproportionate rank + # jumps. + biased_selection_scores = self.cosine_router.apply_selection_bias( + router_logits=selection_logits, + router_scores=selection_scores, + expert_bias=self.expert_bias.detach(), + ) # [num_tokens,num_experts] + _, expert_indices = torch.topk( + biased_selection_scores, self.top_k, dim=-1 + ) # [num_tokens,top_k], [num_tokens,top_k] + expert_weights = selection_scores.gather(1, expert_indices) # [num_tokens,top_k] + else: + expert_indices = natural_indices # [num_tokens,top_k] + expert_weights = natural_weights # [num_tokens,top_k] - expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True) # [num_tokens,top_k] + if self.cosine_router.activation == "sigmoid": + expert_weights = self.cosine_router.normalize_scores(expert_weights) # [num_tokens,top_k] + else: + expert_weights = expert_weights / expert_weights.sum(dim=-1, keepdim=True) # [num_tokens,top_k] expert_weights = expert_weights.to(hidden_states.dtype) # [num_tokens,top_k] num_tokens_per_expert = torch.histc( @@ -290,12 +545,26 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, LBLMetadat device=num_tokens_per_expert.device, ) # [1] - # Compute the average probability of routing to these experts. - # Summing over all experts should be equal to 1. - mean_router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts] + # Compute the average normalized routing probability per expert. + # Sigmoid affinities are normalized across experts for this LBL term. + mean_router_prob_per_expert = torch.mean(routing_probabilities, dim=0) # [num_experts] + + # LBL count: when bias correction is on, ``num_tokens_per_expert`` reflects + # the bias-balanced dispatch, which would artificially satisfy LBL while + # the unbiased routing mass stays concentrated. Feed LBL the gate-natural counts + # instead. With bias off this is bit-identical to ``num_tokens_per_expert``. + if self.aux_loss_free_load_balancing_config.enabled: + num_tokens_per_expert_lbl = torch.histc( + natural_indices.to(dtype=torch.int32).view(-1), + bins=self.num_experts, + min=0, + max=self.num_experts - 1, + ).to(dtype=torch.int64) # [num_experts] + else: + num_tokens_per_expert_lbl = num_tokens_per_expert lbl_metadata = LBLMetadata( - num_tokens_per_expert=num_tokens_per_expert, + num_tokens_per_expert=num_tokens_per_expert_lbl, num_tokens=num_tokens, mean_router_prob_per_expert=mean_router_prob_per_expert, ) @@ -304,9 +573,13 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, LBLMetadat self._update_moe_callback_stats( num_tokens_per_expert=num_tokens_per_expert, num_tokens=num_tokens, - routing_weights=routing_weights, + routing_probabilities=routing_probabilities, expert_indices=expert_indices, ) + # Accumulate the actual (bias-balanced) dispatch counts that drive + # the bias controller toward an even load. + if self.aux_loss_free_load_balancing_config.enabled and self.training: + self.tokens_per_expert.add_(num_tokens_per_expert.to(self.tokens_per_expert.dtype)) return routed_out, lbl_metadata @@ -359,6 +632,57 @@ def get_coactivation_counts(self, reset: bool = True) -> torch.Tensor: self.coactivation_counts.zero_() return val + def update_bias( + self, + device_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + ) -> None: + """Routing bias correction: nudge by ``speed * sign(avg - count)``. + + Every expert receives a ``±update_speed`` nudge each step from the sign + of the rank-error between its accumulated dispatch count and the + per-expert average. The bias is recentered to zero mean every step and + clamped symmetrically by ``max_bias``. Gradient-free. + """ + if not self.aux_loss_free_load_balancing_config.enabled: + return + with torch.no_grad(): + counts = self.tokens_per_expert.clone() + if device_mesh is not None: + from torch.distributed.tensor import DTensor, Partial + + # Counts are accumulated per local rank; reduce them across the + # data-parallel mesh so every rank applies an identical update + # (keeping the replicated expert_bias buffer in sync). + counts = DTensor.from_local( + counts, + device_mesh=device_mesh, + placements=[Partial()] * device_mesh.ndim, + ).full_tensor() + avg = counts.mean() + error = avg - counts + update = self.aux_loss_free_load_balancing_config.update_speed * torch.sign(error) + self.expert_bias.add_(update) + # Recenter so the running mean stays at zero (right-skewed counts + # otherwise cause monotonic drift that saturates the clamp). + self.expert_bias.sub_(self.expert_bias.mean()) + max_bias = self.aux_loss_free_load_balancing_config.max_bias + if max_bias is not None: + self.expert_bias.clamp_(min=-max_bias, max=max_bias) + self.tokens_per_expert.zero_() + + def update_router_bias( + self, + device_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, + ) -> None: + """EMA-update the token-constant buffer ``router_bias``. + + ``forward`` accumulates the per-step token sum + count locally on each rank; this + reduces them across the data-parallel mesh (so every rank computes the same global + batch mean ``m = sum_t(x) / count``), then applies the EMA + ``router_bias ← momentum·router_bias + (1-momentum)·m`` and resets the accumulators. + Called from the trainer's optimizer-step hook. Gradient-free;""" + self.cosine_router.update_ema_bias(device_mesh=device_mesh) + def init_weights(self, buffer_device: torch.device | None = None): self.register_buffer( "total_tokens_per_expert", @@ -400,6 +724,16 @@ def init_weights(self, buffer_device: torch.device | None = None): _make_coactivation_pairs(self.top_k, device=buffer_device), persistent=False, ) + if self.aux_loss_free_load_balancing_config.enabled: + self.register_buffer( + "expert_bias", + torch.zeros(self.num_experts, device=buffer_device), + ) + self.register_buffer( + "tokens_per_expert", + torch.zeros(self.num_experts, device=buffer_device), + persistent=False, + ) if hasattr(self.config, "initializer_range"): std = self.config.initializer_range @@ -413,6 +747,7 @@ def init_weights(self, buffer_device: torch.device | None = None): # Zero-init so the initial per-expert noise std is softplus(0)=ln(2) # uniformly, giving symmetric exploration before gate_noise learns. nn.init.zeros_(self.gate_noise.weight) + self.cosine_router.init_weights(buffer_device=buffer_device) def rotate_half(x): @@ -1574,9 +1909,11 @@ def load_balancing_loss_func( concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) # concatenated_gate_logits: [num_layers*B*N,num_experts] - routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) # [num_layers*B*N,num_experts] + routing_probabilities = torch.nn.functional.softmax( + concatenated_gate_logits, dim=-1 + ) # [num_layers*B*N,num_experts] - _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # [num_layers*B*N,top_k] + _, selected_experts = torch.topk(routing_probabilities, top_k, dim=-1) # [num_layers*B*N,top_k] expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) # [num_layers*B*N,top_k,num_experts] @@ -1585,7 +1922,7 @@ def load_balancing_loss_func( tokens_per_expert = torch.mean(expert_mask.float(), dim=0) # [top_k,num_experts] # Compute the average probability of routing to these experts - router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts] + router_prob_per_expert = torch.mean(routing_probabilities, dim=0) # [num_experts] else: batch_size, sequence_length = attention_mask.shape num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length) @@ -1612,7 +1949,7 @@ def load_balancing_loss_func( ) # [num_layers*B*N,num_experts] # Compute the average probability of routing to these experts - router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum( + router_prob_per_expert = torch.sum(routing_probabilities * router_per_expert_attention_mask, dim=0) / torch.sum( router_per_expert_attention_mask, dim=0 ) # [num_experts] diff --git a/cosmos_framework/model/generator/utils/moe_utils.py b/cosmos_framework/model/generator/utils/moe_utils.py new file mode 100644 index 00000000..494095fe --- /dev/null +++ b/cosmos_framework/model/generator/utils/moe_utils.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from collections.abc import Iterator + +import torch +from torch.distributed.device_mesh import DeviceMesh + +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock + + +def _iter_generation_moe_blocks( + net: torch.nn.Module, +) -> Iterator[tuple[str, Qwen3VLMoeTextSparseMoeBlock]]: + """Yield generation-tower sparse MoE blocks and their module names.""" + for name, module in net.named_modules(): + if isinstance(module, Qwen3VLMoeTextSparseMoeBlock) and "moe_gen" in name: + yield name, module + + +def update_expert_biases( + net: torch.nn.Module, + device_mesh: DeviceMesh | None = None, +) -> None: + """Update routing-load biases on every enabled generation-tower MoE block.""" + for _, module in _iter_generation_moe_blocks(net): + if module.aux_loss_free_load_balancing_config.enabled: + module.update_bias(device_mesh=device_mesh) + + +def update_router_biases(net: torch.nn.Module, device_mesh: DeviceMesh | None = None) -> None: + """Update EMA router biases on every enabled generation-tower MoE block.""" + for _, module in _iter_generation_moe_blocks(net): + if module.cosine_router.input_centering == "ema": + module.update_router_bias(device_mesh=device_mesh) + + +def uses_aux_loss_free_load_balancing(net: torch.nn.Module) -> bool: + """Return whether any generation-tower MoE block uses aux-loss-free load balancing.""" + return any(module.aux_loss_free_load_balancing_config.enabled for _, module in _iter_generation_moe_blocks(net)) + + +def uses_ema_router_bias(net: torch.nn.Module) -> bool: + """Return whether any generation-tower MoE block uses an EMA router bias.""" + return any(module.cosine_router.input_centering == "ema" for _, module in _iter_generation_moe_blocks(net)) + + +@torch.no_grad() +def sync_expert_biases_to_ema(net: torch.nn.Module, net_ema: torch.nn.Module) -> None: + """Mirror generation-tower expert-bias buffers from ``net`` into ``net_ema``.""" + ema_blocks = { + name: module + for name, module in _iter_generation_moe_blocks(net_ema) + if module.aux_loss_free_load_balancing_config.enabled + } + for name, source in _iter_generation_moe_blocks(net): + target = ema_blocks.get(name) + if source.aux_loss_free_load_balancing_config.enabled and target is not None and hasattr(target, "expert_bias"): + target.expert_bias.copy_(source.expert_bias) # [E] + + +@torch.no_grad() +def sync_router_biases_to_ema(net: torch.nn.Module, net_ema: torch.nn.Module) -> None: + """Mirror generation-tower router-bias buffers from ``net`` into ``net_ema``.""" + ema_blocks = { + name: module + for name, module in _iter_generation_moe_blocks(net_ema) + if module.cosine_router.input_centering == "ema" + } + for name, source in _iter_generation_moe_blocks(net): + target = ema_blocks.get(name) + if source.cosine_router.input_centering == "ema" and target is not None: + target.cosine_router.router_bias.copy_(source.cosine_router.router_bias) # [D] diff --git a/cosmos_framework/model/generator/utils/safetensors_loader.py b/cosmos_framework/model/generator/utils/safetensors_loader.py index 7fc12180..f2ec0979 100644 --- a/cosmos_framework/model/generator/utils/safetensors_loader.py +++ b/cosmos_framework/model/generator/utils/safetensors_loader.py @@ -31,7 +31,9 @@ language model. Auto-detects the checkpoint format (:func:`detect_vlm_checkpoint_format`). - :func:`load_vlm_model` — generic loader for HF VLM checkpoints into an - FSDP-wrapped ``HFModel``; honors a skip-pattern overlay. + FSDP-wrapped ``HFModel``; honors a skip-pattern overlay. Also handles + the ``nvidia/Cosmos3-Edge`` indexed-snapshot layout — see + :func:`_detect_indexed_snapshot`. Borrowed from cosmos_rl's ``MultiRankWeightLoader`` (renamed to ``MultiRankCheckpointLoader`` here) with modifications for loading from @@ -39,10 +41,13 @@ https://github.com/nvidia-cosmos/cosmos-rl/blob/main/cosmos_rl/utils/multi_rank_weight_loader.py """ +import hashlib +import json import os import re import time from collections.abc import Callable, Iterator +from typing import NamedTuple import torch import torch.distributed as dist @@ -50,9 +55,9 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor -from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.utils import log from cosmos_framework.utils.easy_io import easy_io +from cosmos_framework.utils.flags import INTERNAL from cosmos_framework.utils.generator.parallelism import ParallelDims # Prefixes stripped when matching checkpoint keys to model state-dict keys. @@ -137,6 +142,180 @@ def _list_safetensors_files( ) +# ── Indexed-snapshot support (nvidia/Cosmos3-Edge layout) ──────────────────── +# +# The nvidia/Cosmos3-Edge repo ships no top-level *.safetensors; a root +# ``model.safetensors.index.json`` maps the reasoner keys into subdirectory +# shard files that also carry DiT-generation tensors absent from the index. +# Only the indexed keys are read (remapped to canonical model keys); everything +# else in those shard files is left untouched. + +_ROOT_INDEX_FILENAME = "model.safetensors.index.json" + +# Known-good sha256 of ``vision_encoder/model.safetensors`` (be935d6) — see +# :func:`_verify_edge_vision_shard_hash` for why loading must be pinned to it. +_EDGE_VISION_SHARD_RELPATH = "vision_encoder/model.safetensors" +_EDGE_VISION_SHARD_SHA256 = "2180ad739ecc96b5c1e9386892d3c5c08bfa42b9cdab9aabc53b028671db89b3" + + +# Gen-pathway-only tensors in the nvidia/Cosmos3-Edge root index: the reasoner has +# no module for them, so detection skips them (unknown keys still fail loudly). +_EDGE_INDEX_GEN_ONLY_KEY_RE = re.compile(r"^layers\.\d+\.self_attn\.k_norm_und_for_gen\.weight$") + + +def convert_key_from_cosmos3_edge_index(name: str) -> str | None: + """Map a nvidia/Cosmos3-Edge root-index reasoner key to the canonical + ``Cosmos3EdgeForConditionalGeneration`` state-dict key. + + The root index uses the public Edge reasoner manifest layout (28 + ``layers.{N}`` blocks, each carrying attention + MLP), while the model + uses the 28-paired / 56-block ``model.language_model.layers.*`` layout: + block ``2N`` is attention (``mixer.{q,k,v,o}_proj``), block ``2N+1`` is + MLP (``mixer.{up,down}_proj``). + + Key mapping (index → model):: + + layers.{N}.self_attn.to_{q,k,v}.weight → model.language_model.layers.{2N}.mixer.{q,k,v}_proj.weight + layers.{N}.self_attn.to_out.weight → model.language_model.layers.{2N}.mixer.o_proj.weight + layers.{N}.input_layernorm.weight → model.language_model.layers.{2N}.norm.weight + layers.{N}.mlp.{up,down}_proj.weight → model.language_model.layers.{2N+1}.mixer.{up,down}_proj.weight + layers.{N}.post_attention_layernorm.weight → model.language_model.layers.{2N+1}.norm.weight + embed_tokens.weight → model.language_model.embeddings.weight + norm.weight → model.language_model.norm_f.weight + lm_head.weight → lm_head.weight + model.visual.* / model.projector.* → unchanged + + Returns: + The canonical model key, or None for keys outside the known reasoner + manifest (callers fail loudly on None — see :func:`_detect_indexed_snapshot`). + """ + if name.startswith(("model.visual.", "model.projector.")): + return name + flat = { + "embed_tokens.weight": "model.language_model.embeddings.weight", + "norm.weight": "model.language_model.norm_f.weight", + "lm_head.weight": "lm_head.weight", + } + if name in flat: + return flat[name] + m = re.match(r"layers\.(\d+)\.(.+)$", name) + if m is None: + return None + n, rest = int(m.group(1)), m.group(2) + # (paired-block offset, model tail): even block = attention, odd block = MLP. + table = { + "self_attn.to_q.weight": (0, "mixer.q_proj.weight"), + "self_attn.to_k.weight": (0, "mixer.k_proj.weight"), + "self_attn.to_v.weight": (0, "mixer.v_proj.weight"), + "self_attn.to_out.weight": (0, "mixer.o_proj.weight"), + "input_layernorm.weight": (0, "norm.weight"), + "mlp.up_proj.weight": (1, "mixer.up_proj.weight"), + "mlp.down_proj.weight": (1, "mixer.down_proj.weight"), + "post_attention_layernorm.weight": (1, "norm.weight"), + } + if rest not in table: + return None + offset, tail = table[rest] + return f"model.language_model.layers.{2 * n + offset}.{tail}" + + +class _IndexedSnapshot(NamedTuple): + """Parsed root-index manifest of an indexed snapshot.""" + + files: list[str] # unique subdir-relative shard paths, sorted + key_map: dict[str, str] # raw root-index key -> canonical model state-dict key + + +def _detect_indexed_snapshot(checkpoint_path: str) -> _IndexedSnapshot | None: + """Detect the indexed-snapshot layout; None keeps the existing flat path. + + Triggers only when ALL of the following hold (every other layout — + nano/super/llava flat snapshots, S3 URIs, bare HF repo ids — returns None + and flows through the unchanged top-level-glob path): + + - ``checkpoint_path`` is a local directory, + - it contains NO top-level ``*.safetensors``, + - it has a root ``model.safetensors.index.json`` whose ``weight_map`` + references shard files in subdirectories. + + Once the layout matches, any inconsistency is a hard error (unmappable + index keys, two index keys colliding onto one model key, missing shard + files) — half-loading a checkpoint must never pass silently. Sole exception: + the gen-only allowlist (:data:`_EDGE_INDEX_GEN_ONLY_KEY_RE`) is skipped. + """ + if not os.path.isdir(checkpoint_path): + return None + with os.scandir(checkpoint_path) as entries: + if any(entry.name.endswith(".safetensors") and entry.is_file() for entry in entries): + return None + index_path = os.path.join(checkpoint_path, _ROOT_INDEX_FILENAME) + if not os.path.isfile(index_path): + return None + with open(index_path) as f: + weight_map: dict[str, str] = json.load(f).get("weight_map") or {} + if not weight_map or not any("/" in rel_path for rel_path in weight_map.values()): + return None + + key_map: dict[str, str] = {} + unmapped: list[str] = [] + for raw_key in weight_map: + if _EDGE_INDEX_GEN_ONLY_KEY_RE.match(raw_key): + continue # known gen-pathway tensor with no reasoner module — skip, not an error + dest_key = convert_key_from_cosmos3_edge_index(raw_key) + if dest_key is None: + unmapped.append(raw_key) + else: + key_map[raw_key] = dest_key + if unmapped: + raise ValueError( + f"Indexed snapshot at '{checkpoint_path}': {len(unmapped)} root-index " + f"key(s) have no canonical model-key mapping. First up to 10: {sorted(unmapped)[:10]}" + ) + if len(set(key_map.values())) != len(key_map): + by_dest: dict[str, list[str]] = {} + for raw_key, dest_key in key_map.items(): + by_dest.setdefault(dest_key, []).append(raw_key) + collisions = {d: sorted(raws) for d, raws in sorted(by_dest.items()) if len(raws) > 1} + raise ValueError( + f"Indexed snapshot at '{checkpoint_path}': multiple root-index keys " + f"map to the same model key (first up to 5): {dict(list(collisions.items())[:5])}" + ) + + files = sorted(set(weight_map.values())) + missing_files = [f for f in files if not os.path.isfile(os.path.join(checkpoint_path, f))] + if missing_files: + raise FileNotFoundError( + f"Indexed snapshot at '{checkpoint_path}': shard file(s) referenced by " + f"{_ROOT_INDEX_FILENAME} not found: {missing_files}" + ) + return _IndexedSnapshot(files=files, key_map=key_map) + + +def _verify_edge_vision_shard_hash(checkpoint_path: str, files: list[str]) -> None: + """G2 guard: refuse to load a changed ``vision_encoder/model.safetensors``. + + This pipeline keeps the OLD patch-embedding weight-layout convention; an + upstream re-export in a different convention (e.g. transformers-main's + native layout) would load without error and silently corrupt vision + features — see outputs/audit/cosmos3_edge_native_vision_layout_bug.md. + Runs on every rank: a rank-0-only raise would hang the peers at the next + collective (shared-storage page cache makes the repeat reads cheap). + """ + if _EDGE_VISION_SHARD_RELPATH not in files: + return + shard_path = os.path.join(checkpoint_path, _EDGE_VISION_SHARD_RELPATH) + with open(shard_path, "rb") as f: + actual = hashlib.file_digest(f, "sha256").hexdigest() + if actual != _EDGE_VISION_SHARD_SHA256: + raise ValueError( + f"{shard_path}: sha256 mismatch (expected {_EDGE_VISION_SHARD_SHA256}, " + f"got {actual}): upstream vision weights changed — re-verify the " + "patch-embedding layout convention (see " + "outputs/audit/cosmos3_edge_native_vision_layout_bug.md) before training" + ) + log.info(f"G2 guard: {_EDGE_VISION_SHARD_RELPATH} sha256 verified ({actual[:16]}…)") + + def _get_local_rank_and_size(device_mesh: DeviceMesh) -> tuple[int, int]: """Get the local rank and size of a device mesh. @@ -327,6 +506,7 @@ def load_files_parallel( checkpoint_path: str, credential_path: str | None, loading_device: torch.device, + index: _IndexedSnapshot | None = None, ) -> tuple[ dict[str, torch.Tensor], dict[str, tuple[list, int]], @@ -342,6 +522,15 @@ def load_files_parallel( fall back to Hugging Face. credential_path: Path to the credential file for S3/GCS. loading_device: Device to load tensors on. + index: Optional indexed-snapshot manifest (from + :func:`_detect_indexed_snapshot`). When set: the shard list + comes from the root index (no top-level listing, no HF + fallback); only tensors named in ``index.key_map`` are kept + (the shard files also carry DiT tensors that must NOT be + loaded); kept tensors are renamed to canonical model keys at + read time, BEFORE the cross-rank name gather, so all ranks + collectivize on canonical names. ``None`` (the default) + preserves the existing behavior exactly. Returns: Tuple of (rank_tensors, rank_tensor_metadata, weights_of_ckpt_names): @@ -358,7 +547,10 @@ def load_files_parallel( log.info(f"Loading safetensors files from: {checkpoint_path}", rank0_only=False) log.info(f"Credential path: {credential_path}", rank0_only=False) list_error: Exception | None = None - if checkpoint_path.startswith(_HF_URI_PREFIX): + if index is not None: + # Indexed snapshot: shard list comes from the root-index manifest. + safetensors_files = list(index.files) + elif checkpoint_path.startswith(_HF_URI_PREFIX): safetensors_files: list[str] = [] else: try: @@ -421,7 +613,14 @@ def load_files_parallel( weights_data = easy_io.get(full_path, backend_args=backend_args) state_dict = load_safetensors(weights_data) for name, tensor in state_dict.items(): - # Names are stored RAW here; per-checkpoint name + if index is not None: + canonical_name = index.key_map.get(name) + if canonical_name is None: + # Not in the root index (e.g. DiT weights sharing + # the shard file) — never load it. + continue + name = canonical_name + # Otherwise names are stored RAW here; per-checkpoint name # conversion (see _make_name_converter / the # convert_weight_from_*_hf functions) is applied later # by the caller after broadcast. @@ -1048,6 +1247,13 @@ def load_vlm_model( explicit ``hf://org/model`` Hub URIs and bare ``org/model`` repo IDs fall back to Hugging Face. + Local directories in the indexed-snapshot layout (see + :func:`_detect_indexed_snapshot`) take a dedicated branch: only the + root-index keys are read from the subdir shard files, remapped to canonical + model keys (:func:`convert_key_from_cosmos3_edge_index`), and verified for + completeness in both directions; the vision-shard content hash is checked + first (:func:`_verify_edge_vision_shard_hash`). + Both ``tensor_names_to_skip`` and ``extra_skip_patterns`` are lists of regex patterns applied to the RESOLVED model key (post-name_converter). Phase-5 skips any model key matched by either list; Phase-6's @@ -1107,6 +1313,16 @@ def load_vlm_model( "handling. Use a dense VLM checkpoint (2B, 4B, 8B, 32B) until then." ) + # Indexed-snapshot detection (nvidia/Cosmos3-Edge layout). Every other + # layout returns None and takes the existing flat-glob path unchanged. + indexed = _detect_indexed_snapshot(checkpoint_path) + if indexed is not None: + log.info( + f"load_vlm_model: indexed snapshot detected at {checkpoint_path} " + f"({len(indexed.key_map)} root-index keys in {len(indexed.files)} subdir shard files)" + ) + _verify_edge_vision_shard_hash(checkpoint_path, indexed.files) + # FUTURE: to re-enable FSDP-2 CPU offload, detect CPU local_views via # ``sample.device.type == "cpu"``, force the loader to single-rank (None # instead of _get_dp_shard_mesh), and pin ``target_device`` to ``"cpu"``. @@ -1115,6 +1331,7 @@ def load_vlm_model( checkpoint_path=checkpoint_path, credential_path=credential_path if credential_path else "", loading_device="cpu", + index=indexed, ) all_tensor_names, tensor_to_rank = loader.gather_tensor_names_and_build_mapping( ckpt_names, @@ -1128,6 +1345,7 @@ def load_vlm_model( compiled_skip_patterns = [re.compile(p) for p in (skip_patterns or [])] keys_loaded: set[str] = set() skipped_model_keys: set[str] = set() + unexpected_index_keys: set[str] = set() target_device = "cuda" if torch.cuda.is_available() else "cpu" @@ -1147,14 +1365,19 @@ def load_vlm_model( rank_tensor_meta, device=target_device, ): - dest_name = name_converter(ckpt_name) + # Indexed snapshots arrive pre-remapped to canonical model keys (see + # load_files_parallel) — the suffix converter must not touch them. + dest_name = ckpt_name if indexed is not None else name_converter(ckpt_name) if any(p.fullmatch(dest_name) for p in compiled_skip_patterns): skipped_model_keys.add(dest_name) continue if dest_name not in vlm_state_dict: - continue # extra checkpoint key — ignore + if indexed is not None: + # Root-index keys must land on model keys — fail loudly below. + unexpected_index_keys.add(dest_name) + continue # extra checkpoint key — ignore (flat layouts only) target = vlm_state_dict[dest_name] is_dtensor = isinstance(target, DTensor) @@ -1173,6 +1396,29 @@ def load_vlm_model( local_view.data.copy_(shard) keys_loaded.add(dest_name) + # Indexed snapshots: every root-index key must have been consumed — loaded + # or intentionally skipped; anything else means the remap and the model + # disagree. + if indexed is not None: + if unexpected_index_keys: + raise ValueError( + f"load_vlm_model: {len(unexpected_index_keys)} root-index key(s) " + f"mapped to model keys that do not exist in the model state dict " + f"for '{checkpoint_path}'. First up to 10: {sorted(unexpected_index_keys)[:10]}" + ) + unconsumed = { + raw_key + for raw_key, dest_key in indexed.key_map.items() + if dest_key not in keys_loaded and dest_key not in skipped_model_keys + } + if unconsumed: + # weight_map named a shard file that doesn't actually contain the key. + raise ValueError( + f"load_vlm_model: {len(unconsumed)} root-index key(s) never found " + f"in their shard files for '{checkpoint_path}'. First up to 10: " + f"{sorted(unconsumed)[:10]}" + ) + # Phase 6: completeness check with tied-embedding AND skip-list tolerance. missing = set(vlm_state_dict) - keys_loaded - skipped_model_keys diff --git a/cosmos_framework/model/generator/utils/safetensors_loader_indexed_test.py b/cosmos_framework/model/generator/utils/safetensors_loader_indexed_test.py new file mode 100644 index 00000000..0cceac2a --- /dev/null +++ b/cosmos_framework/model/generator/utils/safetensors_loader_indexed_test.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Tests for the indexed-snapshot loading branch of load_vlm_model +(nvidia/Cosmos3-Edge layout: no top-level shards; a root +model.safetensors.index.json maps reasoner keys into subdir shard files). + +Tests needing the real HF snapshot or the canonical converter output skip +themselves when those files are absent, so the suite is green on CI runners +without the Edge checkpoint cached. +""" + +import json +import os +from pathlib import Path + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +import cosmos_framework.model.generator.utils.safetensors_loader as safetensors_loader +from cosmos_framework.model.generator.utils.safetensors_loader import ( + _EDGE_INDEX_GEN_ONLY_KEY_RE, + _detect_indexed_snapshot, + _verify_edge_vision_shard_hash, + convert_key_from_cosmos3_edge_index, + load_vlm_model, +) +from cosmos_framework.model.generator.utils.safetensors_loader_test import _StubConfig, _StubModel + +_REPO_ROOT = Path(__file__).resolve().parents[4] + +# Canonical converter output: its key list is the authoritative target set +# (identical to Cosmos3EdgeForConditionalGeneration.state_dict().keys()). +_CANONICAL_VLM_FILE = _REPO_ROOT / "examples" / "checkpoints" / "Cosmos3-Edge-Reasoner-VLM" / "model.safetensors" + + +def _find_local_edge_snapshot() -> str | None: + """Locate a locally cached nvidia/Cosmos3-Edge indexed snapshot, if any.""" + roots: list[str] = [] + if os.environ.get("HF_HOME"): + roots.append(os.path.join(os.environ["HF_HOME"], "hub")) + if os.environ.get("HF_HUB_CACHE"): + roots.append(os.environ["HF_HUB_CACHE"]) + roots.append(os.path.expanduser("~/.cache/huggingface/hub")) + for root in roots: + snapshots_dir = os.path.join(root, "models--nvidia--Cosmos3-Edge", "snapshots") + if not os.path.isdir(snapshots_dir): + continue + for name in sorted(os.listdir(snapshots_dir)): + snapshot = os.path.join(snapshots_dir, name) + try: + if _detect_indexed_snapshot(snapshot) is not None: + return snapshot + except (ValueError, FileNotFoundError): + continue # partially downloaded / foreign snapshot — keep looking + return None + + +def _write_indexed_snapshot(root: Path, shards: dict[str, dict[str, torch.Tensor]]) -> Path: + """Write a fake indexed snapshot: subdir shard files + root weight-map index.""" + root.mkdir(parents=True, exist_ok=True) + weight_map: dict[str, str] = {} + for rel_path, tensors in shards.items(): + shard_path = root / rel_path + shard_path.parent.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(shard_path)) + for key in tensors: + weight_map[key] = rel_path + with open(root / "model.safetensors.index.json", "w") as f: + json.dump({"metadata": {}, "weight_map": weight_map}, f) + return root + + +# --- convert_key_from_cosmos3_edge_index (pure remap) ------------------------ + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_edge_index_remap_spot_pairs(): + """Representative index → model pairs from the mapping table.""" + pairs = { + "layers.0.self_attn.to_q.weight": "model.language_model.layers.0.mixer.q_proj.weight", + "layers.3.self_attn.to_k.weight": "model.language_model.layers.6.mixer.k_proj.weight", + "layers.5.self_attn.to_v.weight": "model.language_model.layers.10.mixer.v_proj.weight", + "layers.27.self_attn.to_out.weight": "model.language_model.layers.54.mixer.o_proj.weight", + "layers.0.input_layernorm.weight": "model.language_model.layers.0.norm.weight", + "layers.4.mlp.up_proj.weight": "model.language_model.layers.9.mixer.up_proj.weight", + "layers.4.mlp.down_proj.weight": "model.language_model.layers.9.mixer.down_proj.weight", + "layers.27.post_attention_layernorm.weight": "model.language_model.layers.55.norm.weight", + "embed_tokens.weight": "model.language_model.embeddings.weight", + "norm.weight": "model.language_model.norm_f.weight", + "lm_head.weight": "lm_head.weight", + # visual + projector pass through unchanged + "model.visual.encoder.layers.0.self_attn.q_proj.weight": ( + "model.visual.encoder.layers.0.self_attn.q_proj.weight" + ), + "model.projector.linear_fc1.bias": "model.projector.linear_fc1.bias", + } + for raw, expected in pairs.items(): + assert convert_key_from_cosmos3_edge_index(raw) == expected + + # Keys outside the reasoner manifest must return None (fail loudly upstream). + for bad in ( + "layers.0.self_attn.q_norm.weight", + "layers.0.mlp.gate_proj.weight", + "action_proj_in.fc.weight", + "blocks.0.attn.to_q_moe_gen.weight", + "model.language_model.layers.0.mixer.q_proj.weight", # already-canonical keys are not index keys + ): + assert convert_key_from_cosmos3_edge_index(bad) is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_edge_index_remap_miniature_manifest_bijection(): + """Full per-layer pattern for a miniature 2-layer manifest: the remap must + be a bijection onto the expected canonical key set.""" + raw_keys = ["embed_tokens.weight", "norm.weight", "lm_head.weight"] + expected = { + "model.language_model.embeddings.weight", + "model.language_model.norm_f.weight", + "lm_head.weight", + } + for n in range(2): + raw_keys += [ + f"layers.{n}.self_attn.to_q.weight", + f"layers.{n}.self_attn.to_k.weight", + f"layers.{n}.self_attn.to_v.weight", + f"layers.{n}.self_attn.to_out.weight", + f"layers.{n}.input_layernorm.weight", + f"layers.{n}.mlp.up_proj.weight", + f"layers.{n}.mlp.down_proj.weight", + f"layers.{n}.post_attention_layernorm.weight", + ] + expected |= { + f"model.language_model.layers.{2 * n}.mixer.q_proj.weight", + f"model.language_model.layers.{2 * n}.mixer.k_proj.weight", + f"model.language_model.layers.{2 * n}.mixer.v_proj.weight", + f"model.language_model.layers.{2 * n}.mixer.o_proj.weight", + f"model.language_model.layers.{2 * n}.norm.weight", + f"model.language_model.layers.{2 * n + 1}.mixer.up_proj.weight", + f"model.language_model.layers.{2 * n + 1}.mixer.down_proj.weight", + f"model.language_model.layers.{2 * n + 1}.norm.weight", + } + raw_keys += ["model.visual.post_layernorm.weight", "model.projector.norm.weight"] + expected |= {"model.visual.post_layernorm.weight", "model.projector.norm.weight"} + + mapped = [convert_key_from_cosmos3_edge_index(k) for k in raw_keys] + assert None not in mapped + assert len(set(mapped)) == len(mapped) # injective + assert set(mapped) == expected # onto + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_edge_index_remap_full_root_index_bijection_onto_canonical(): + """Map ALL reasoner root-index keys of the real snapshot and assert an exact + bijection onto the canonical Cosmos3-Edge-Reasoner-VLM key list (the 28 + gen-only k_norm_und_for_gen keys are skipped by detection).""" + snapshot = _find_local_edge_snapshot() + if snapshot is None: + pytest.skip("no local nvidia/Cosmos3-Edge indexed snapshot in the HF cache") + if not os.path.exists(_CANONICAL_VLM_FILE): + pytest.skip(f"canonical key list not available: {_CANONICAL_VLM_FILE}") + + with open(os.path.join(snapshot, "model.safetensors.index.json")) as f: + raw_keys = sorted(json.load(f)["weight_map"]) + with safe_open(str(_CANONICAL_VLM_FILE), framework="pt") as f: + canonical_keys = set(f.keys()) + + gen_only = [k for k in raw_keys if _EDGE_INDEX_GEN_ONLY_KEY_RE.match(k)] + assert len(gen_only) in (0, 28), gen_only + reasoner_keys = [k for k in raw_keys if not _EDGE_INDEX_GEN_ONLY_KEY_RE.match(k)] + + mapped = {k: convert_key_from_cosmos3_edge_index(k) for k in reasoner_keys} + assert None not in mapped.values(), sorted(k for k, v in mapped.items() if v is None) + assert len(set(mapped.values())) == len(mapped) # injective + assert set(mapped.values()) == canonical_keys # onto the authoritative target set + assert len(mapped) == 670 + + +# --- _detect_indexed_snapshot layout gating ---------------------------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_detect_returns_none_for_flat_layout(tmp_path): + """Top-level *.safetensors present → existing flat path, even when an + index file also exists (regression guard for nano/super/llava snapshots).""" + save_file({"model.embed_tokens.weight": torch.zeros(2, 2)}, str(tmp_path / "model.safetensors")) + with open(tmp_path / "model.safetensors.index.json", "w") as f: + json.dump({"weight_map": {"model.embed_tokens.weight": "model.safetensors"}}, f) + assert _detect_indexed_snapshot(str(tmp_path)) is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_detect_returns_none_without_subdir_references(tmp_path): + """No index at all, and an index whose weight_map has no subdir references, + both fall through to the existing path.""" + assert _detect_indexed_snapshot(str(tmp_path)) is None # empty dir, no index + + with open(tmp_path / "model.safetensors.index.json", "w") as f: + json.dump({"weight_map": {"embed_tokens.weight": "model-00001-of-00002.safetensors"}}, f) + assert _detect_indexed_snapshot(str(tmp_path)) is None + + # Non-directory inputs (S3 URI / bare repo id) are never indexed snapshots. + assert _detect_indexed_snapshot("s3://bucket/model") is None + assert _detect_indexed_snapshot("nvidia/Cosmos3-Edge") is None + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_detect_raises_on_unmappable_index_key(tmp_path): + _write_indexed_snapshot( + tmp_path, + {"transformer/part-00001.safetensors": {"layers.0.mlp.gate_proj.weight": torch.zeros(2, 2)}}, + ) + with pytest.raises(ValueError, match="no canonical model-key mapping"): + _detect_indexed_snapshot(str(tmp_path)) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_detect_skips_gen_only_k_norm_keys(tmp_path): + """Gen-only k_norm_und_for_gen index keys are recognized and skipped; + genuinely unknown keys keep failing loudly.""" + _write_indexed_snapshot( + tmp_path, + { + "transformer/part-00001.safetensors": { + "layers.0.self_attn.to_q.weight": torch.zeros(2, 2), + "layers.0.self_attn.k_norm_und_for_gen.weight": torch.zeros(2), + } + }, + ) + snapshot = _detect_indexed_snapshot(str(tmp_path)) + assert snapshot is not None + assert "layers.0.self_attn.k_norm_und_for_gen.weight" not in snapshot.key_map + assert snapshot.key_map == {"layers.0.self_attn.to_q.weight": "model.language_model.layers.0.mixer.q_proj.weight"} + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_detect_raises_on_missing_shard_file(tmp_path): + _write_indexed_snapshot( + tmp_path, + {"transformer/part-00001.safetensors": {"embed_tokens.weight": torch.zeros(2, 2)}}, + ) + with open(tmp_path / "model.safetensors.index.json", "w") as f: + json.dump({"weight_map": {"embed_tokens.weight": "transformer/does-not-exist.safetensors"}}, f) + with pytest.raises(FileNotFoundError, match="shard file"): + _detect_indexed_snapshot(str(tmp_path)) + + +# --- G2 vision-shard hash guard ---------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_g2_hash_mismatch_raises(tmp_path): + """A vision_encoder/model.safetensors with any other content must be refused.""" + snapshot = _write_indexed_snapshot( + tmp_path, + {"vision_encoder/model.safetensors": {"model.visual.post_layernorm.weight": torch.zeros(4)}}, + ) + with pytest.raises(ValueError, match="upstream vision weights changed"): + _verify_edge_vision_shard_hash(str(snapshot), ["vision_encoder/model.safetensors"]) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_g2_hash_guard_skipped_without_vision_shard(tmp_path): + """Snapshots whose index does not reference the vision shard skip the guard.""" + _verify_edge_vision_shard_hash(str(tmp_path), ["transformer/part-00001.safetensors"]) # no raise + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_g2_hash_pass_with_matching_digest(tmp_path, monkeypatch): + """Positive path: with the constant patched to the actual digest, no raise.""" + import hashlib + + snapshot = _write_indexed_snapshot( + tmp_path, + {"vision_encoder/model.safetensors": {"model.visual.post_layernorm.weight": torch.zeros(4)}}, + ) + with open(snapshot / "vision_encoder" / "model.safetensors", "rb") as f: + digest = hashlib.file_digest(f, "sha256").hexdigest() + monkeypatch.setattr(safetensors_loader, "_EDGE_VISION_SHARD_SHA256", digest) + _verify_edge_vision_shard_hash(str(snapshot), ["vision_encoder/model.safetensors"]) # no raise + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_load_vlm_model_runs_g2_guard_on_indexed_snapshot(tmp_path): + """End-to-end: the indexed branch of load_vlm_model hits the G2 guard + BEFORE copying any tensor.""" + snapshot = _write_indexed_snapshot( + tmp_path, + {"vision_encoder/model.safetensors": {"model.visual.post_layernorm.weight": torch.ones(4)}}, + ) + config = _StubConfig( + text_config=_StubConfig(num_experts=None, num_local_experts=None), + tie_word_embeddings=False, + ) + model = _StubModel({"model.visual.post_layernorm.weight": (4,)}, config) + with pytest.raises(ValueError, match="upstream vision weights changed"): + load_vlm_model(model=model, checkpoint_path=str(snapshot), credential_path=None, parallel_dims=None) + # Guard fired before any tensor copy: the model param is untouched. + assert torch.equal(model._params["model.visual.post_layernorm.weight"].data, torch.zeros(4)) + + +# --- load_vlm_model indexed branch (fake snapshot, single-rank CPU) ----------- + + +def _edge_stub(param_shapes: dict[str, tuple[int, ...]]) -> _StubModel: + config = _StubConfig( + text_config=_StubConfig(num_experts=None, num_local_experts=None), + tie_word_embeddings=False, + ) + return _StubModel(param_shapes, config) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_load_vlm_model_indexed_end_to_end(tmp_path): + """Fake indexed snapshot (no vision shard → G2 skipped): every root-index + key lands on its remapped canonical key; DiT-only tensors in the shared + shard files are never loaded.""" + fills = { + "layers.0.self_attn.to_q.weight": 1.0, + "layers.0.self_attn.to_k.weight": 2.0, + "layers.0.self_attn.to_v.weight": 3.0, + "layers.0.self_attn.to_out.weight": 4.0, + "layers.0.input_layernorm.weight": 5.0, + "layers.0.mlp.up_proj.weight": 6.0, + "layers.0.mlp.down_proj.weight": 7.0, + "layers.0.post_attention_layernorm.weight": 8.0, + "embed_tokens.weight": 9.0, + "norm.weight": 10.0, + "lm_head.weight": 11.0, + "model.projector.norm.weight": 12.0, + } + + def t(raw_key: str) -> torch.Tensor: + shape = (4,) if "norm" in raw_key else (4, 4) + return torch.full(shape, fills[raw_key]) + + shard1 = {k: t(k) for k in list(fills)[:5]} + # DiT-generation tensor sharing the shard file but absent from the index. + shard1["blocks.0.attn.to_q_moe_gen.weight"] = torch.full((4, 4), -1.0) + shard2 = {k: t(k) for k in list(fills)[5:]} + snapshot = _write_indexed_snapshot( + tmp_path, + {"transformer/part-00001.safetensors": shard1, "transformer/part-00002.safetensors": shard2}, + ) + # The DiT extra must not appear in the index: rewrite weight_map without it. + with open(snapshot / "model.safetensors.index.json") as f: + weight_map = json.load(f)["weight_map"] + del weight_map["blocks.0.attn.to_q_moe_gen.weight"] + with open(snapshot / "model.safetensors.index.json", "w") as f: + json.dump({"weight_map": weight_map}, f) + + expected = {convert_key_from_cosmos3_edge_index(k): t(k) for k in fills} + model = _edge_stub({k: tuple(v.shape) for k, v in expected.items()}) + + keys_loaded = load_vlm_model( + model=model, + checkpoint_path=str(snapshot), + credential_path=None, + parallel_dims=None, + ) + + assert keys_loaded == set(expected) + for canonical_key, tensor in expected.items(): + assert torch.equal(model._params[canonical_key].data, tensor), canonical_key + # The DiT tensor never landed anywhere. + assert not any(torch.equal(p.data, torch.full((4, 4), -1.0)) for p in model._params.values()) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_load_vlm_model_indexed_unexpected_key_raises(tmp_path): + """A root-index key that remaps to a key absent from the model must fail + loudly (no silent extra-key tolerance on the indexed branch).""" + snapshot = _write_indexed_snapshot( + tmp_path, + { + "transformer/part-00001.safetensors": { + "embed_tokens.weight": torch.ones(4, 4), + # remaps to model.language_model.layers.3.mixer.up_proj.weight — not in the model + "layers.1.mlp.up_proj.weight": torch.ones(4, 4), + } + }, + ) + model = _edge_stub({"model.language_model.embeddings.weight": (4, 4)}) + with pytest.raises(ValueError, match="do not exist in the model state dict"): + load_vlm_model(model=model, checkpoint_path=str(snapshot), credential_path=None, parallel_dims=None) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_load_vlm_model_indexed_missing_model_key_raises(tmp_path): + """A model key not covered by the root index trips the Phase-6 completeness + check (nothing about the indexed branch weakens it).""" + snapshot = _write_indexed_snapshot( + tmp_path, + {"transformer/part-00001.safetensors": {"embed_tokens.weight": torch.ones(4, 4)}}, + ) + model = _edge_stub( + { + "model.language_model.embeddings.weight": (4, 4), + "model.language_model.norm_f.weight": (4,), # not in the index + } + ) + with pytest.raises(ValueError, match="required model parameter"): + load_vlm_model(model=model, checkpoint_path=str(snapshot), credential_path=None, parallel_dims=None) + + +# --- integration: real snapshot → fake canonical-shaped model ----------------- + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_load_real_edge_snapshot_into_canonical_state_dict(): + """Load the real snapshot through the indexed branch into a dict-backed + model whose keys/shapes mirror the canonical converter output; all 670 + tensors must land, three of them verified bitwise. + + Non-persistent buffers (e.g. rotary inv_freq) never appear in a state + dict, so the canonical key list already excludes them by construction. + """ + snapshot = _find_local_edge_snapshot() + if snapshot is None: + pytest.skip("no local nvidia/Cosmos3-Edge indexed snapshot in the HF cache") + if not os.path.exists(_CANONICAL_VLM_FILE): + pytest.skip(f"canonical reference not available: {_CANONICAL_VLM_FILE}") + + with safe_open(str(_CANONICAL_VLM_FILE), framework="pt") as f: + canonical_keys = sorted(f.keys()) + shapes = {k: tuple(f.get_slice(k).get_shape()) for k in canonical_keys} + + config = _StubConfig( + text_config=_StubConfig(num_experts=None, num_local_experts=None), + tie_word_embeddings=False, + ) + # bfloat16 params: halves memory AND makes the bitwise spot checks exact + # (checkpoint tensors are bf16; copy_ into bf16 preserves bits). + model = _StubModel(param_shapes={}, config=config) + model._params = { + k: torch.nn.Parameter(torch.zeros(shapes[k], dtype=torch.bfloat16), requires_grad=False) for k in canonical_keys + } + + keys_loaded = load_vlm_model( + model=model, + checkpoint_path=snapshot, + credential_path=None, + parallel_dims=None, + ) + + assert keys_loaded == set(canonical_keys) + assert len(keys_loaded) == 670 + + # Bitwise spot checks against the canonical converter output: one remapped + # attention key, one flat key, one pass-through visual key. + spot_keys = [ + "model.language_model.layers.0.mixer.q_proj.weight", + "model.language_model.embeddings.weight", + "model.visual.embeddings.patch_embedding.weight", + ] + with safe_open(str(_CANONICAL_VLM_FILE), framework="pt") as f: + for key in spot_keys: + reference = f.get_tensor(key) + loaded = model._params[key].data + assert loaded.dtype == reference.dtype, key + assert torch.equal(loaded, reference), key diff --git a/cosmos_framework/scripts/_convert_model_to_diffusers.py b/cosmos_framework/scripts/_convert_model_to_diffusers.py index ceab0ef2..4ea15f73 100644 --- a/cosmos_framework/scripts/_convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/_convert_model_to_diffusers.py @@ -4,24 +4,34 @@ """Implementation for `convert_model_to_diffusers.py`.""" import contextlib +import inspect import json import pathlib import re +import shutil +import struct +from typing import Any, Literal import pydantic import torch from accelerate import init_empty_weights from diffusers import ( AutoencoderKLWan, - Cosmos3AVAEAudioTokenizer, - Cosmos3OmniPipeline, Cosmos3OmniTransformer, + FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler, ) + +try: + from diffusers import Cosmos3OmniPipeline +except ImportError: + # Some intermediate main-branch revisions used this renamed export. + from diffusers import Cosmos3OmniDiffusersPipeline as Cosmos3OmniPipeline from transformers import AutoConfig, AutoTokenizer from cosmos_framework.inference.model import Cosmos3OmniModel from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.model.generator.tokenizers.interface import VideoTokenizerInterface from cosmos_framework.utils import log DEFAULT_SOUND_TOKENIZER_CONFIG = { @@ -116,6 +126,26 @@ VISION_ENCODER_CHECKPOINT_PREFIX = "model.visual." VISION_ENCODER_CHECKPOINT_SUBFOLDER = "vision_encoder" +COSMOS3_EDGE_REASONER = "nvidia/Cosmos3-Edge" +COSMOS3_EDGE_REASONER_REVISION = "be935d6931e4e176d7353abad41ca529d7b33b12" +COSMOS3_EDGE_VAE = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" + +# Keep the Edge metadata and vision sidecar, but do not copy the source +# Transformers shards into the Diffusers output. The source reasoner is only an +# intermediate Stage 1 export; Stage 2 owns the Diffusers transformer shards. +COSMOS3_EDGE_REASONER_METADATA_FILES = ( + "chat_template.jinja", + "config.json", + "generation_config.json", + "preprocessor_config.json", + "processor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", +) +COSMOS3_EDGE_REASONER_INDEX_FILE = "model.safetensors.index.json" + def _get_config_value(*configs, name, default=None): for config in configs: @@ -130,6 +160,46 @@ def _get_config_value(*configs, name, default=None): return default +def _resolve_fixed_step_sampler_config(model_cfg) -> dict | None: + """Extract the distilled fixed-step sampler settings from a checkpoint config. + + Distilled (few-step) Cosmos3 checkpoints carry a `fixed_step_sampler_config` + with a fixed flow-sigma schedule (`t_list`) and a `sample_type` (`ode`/`sde`). + Returns `None` for non-distilled checkpoints. + """ + fixed_step_cfg = _get_config_value(model_cfg, name="fixed_step_sampler_config", default=None) + if fixed_step_cfg is None: + return None + + t_list = _get_config_value(fixed_step_cfg, name="t_list", default=None) + if t_list is None: + raise ValueError("`fixed_step_sampler_config` is set, but `t_list` is missing.") + t_list = [float(t) for t in t_list] + if len(t_list) == 0: + raise ValueError("`fixed_step_sampler_config.t_list` must contain at least one value.") + # Training convention excludes terminal 0.0; normalize defensively if it is present. + if t_list[-1] == 0.0: + t_list = t_list[:-1] + if len(t_list) == 0: + raise ValueError("`fixed_step_sampler_config.t_list` cannot contain only terminal 0.0.") + + # Default "sde" matches i4's FixedStepSamplerConfig; "ode" is still a valid explicit value. + sample_type = str(_get_config_value(fixed_step_cfg, name="sample_type", default="sde")).lower() + if sample_type not in {"ode", "sde"}: + raise ValueError( + f"Unsupported `fixed_step_sampler_config.sample_type={sample_type!r}` (expected 'ode' or 'sde')." + ) + + rf_inference_cfg = _get_config_value(model_cfg, name="rectified_flow_inference_config", default=None) + num_train_timesteps = int(_get_config_value(rf_inference_cfg, name="num_train_timesteps", default=1000)) + + return { + "t_list": t_list, + "sample_type": sample_type, + "num_train_timesteps": num_train_timesteps, + } + + def _remap_language_model_key(key: str) -> str: """Rename a source language-model state-dict key to the diffusers `Cosmos3OmniTransformer` layout: strip the leading `model.` prefix and @@ -143,6 +213,16 @@ def _remap_language_model_key(key: str) -> str: return key +def _remap_time_embedder_key(key: str) -> str: + try: + return _TIME_EMBEDDER_KEY_REMAP[key] + except KeyError as exc: + supported_keys = sorted(_TIME_EMBEDDER_KEY_REMAP) + raise ValueError( + f"Unsupported time_embedder state-dict key {key!r}; expected one of {supported_keys}." + ) from exc + + def _remap_language_model_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Remap language-model keys and reject ambiguous source layouts.""" remapped_state_dict: dict[str, torch.Tensor] = {} @@ -370,7 +450,14 @@ def _load_sound_tokenizer_config(config_path: pathlib.Path | None) -> dict: def _build_sound_tokenizer( checkpoint_path: pathlib.Path, config_path: pathlib.Path | None, -) -> Cosmos3AVAEAudioTokenizer: +) -> Any: + try: + from diffusers import Cosmos3AVAEAudioTokenizer + except ImportError as error: + raise RuntimeError( + "This Diffusers build does not provide Cosmos3AVAEAudioTokenizer, which is required for sound exports." + ) from error + config = _load_sound_tokenizer_config(config_path) log.info(f"Loading AVAE sound tokenizer weights from {checkpoint_path} …") @@ -404,6 +491,27 @@ def _checkpoint_weight_map(checkpoint_path: pathlib.Path) -> dict[str, str]: return index.get("weight_map", {}) +def _is_edge_exported_checkpoint(checkpoint_path: pathlib.Path) -> bool: + """Return whether a Stage 1 HF export uses the Nemotron Edge backbone.""" + config_path = checkpoint_path / "config.json" + if not config_path.is_file(): + return False + config = _load_json(config_path) + model_config = config.get("model", {}).get("config", {}) + vlm_config = model_config.get("vlm_config", {}) + pretrained_weights = vlm_config.get("pretrained_weights", {}) + model_instance = vlm_config.get("model_instance", {}) + # Internal config uses `_target_` with the class name (Nemotron3DenseVL…); this + # repo's public export config uses `_target` with the aliased name + # (nemotron_3_dense_vl_…). Accept both spellings/formats. + target = str(model_instance.get("_target_") or model_instance.get("_target") or "") + return bool( + pretrained_weights.get("checkpoint_format") == "nemotron_3_dense_vl" + or "Nemotron3" in target + or "nemotron_3_dense_vl" in target.lower() + ) + + def _checkpoint_has_weight_prefix(checkpoint_path: pathlib.Path, prefix: str) -> bool: return any(key.startswith(prefix) for key in _checkpoint_weight_map(checkpoint_path)) @@ -540,31 +648,539 @@ def _load_vision_encoder( return _build_vision_encoder(state_dict, model_name_or_path, dtype) +class _MetadataVideoTokenizer(VideoTokenizerInterface): + """Expose video-tokenizer metadata without loading the source VAE.""" + + def __init__( + self, + latent_ch: int, + spatial_compression_factor: int, + temporal_compression_factor: int, + chunk_duration: int, + causal: bool, + ) -> None: + self._latent_ch = latent_ch + self._spatial_compression_factor = spatial_compression_factor + self._temporal_compression_factor = temporal_compression_factor + self._chunk_duration = chunk_duration + self._causal = causal + + def reset_dtype(self) -> None: + pass + + def encode(self, state: torch.Tensor) -> torch.Tensor: + raise RuntimeError("The metadata-only converter tokenizer cannot encode media.") + + def decode(self, latent: torch.Tensor) -> torch.Tensor: + raise RuntimeError("The metadata-only converter tokenizer cannot decode media.") + + def get_latent_num_frames(self, num_pixel_frames: int) -> int: + return 1 + (num_pixel_frames - 1) // self._temporal_compression_factor + + def get_pixel_num_frames(self, num_latent_frames: int, **kwargs: Any) -> int: + del kwargs + return (num_latent_frames - 1) * self._temporal_compression_factor + 1 + + @property + def spatial_compression_factor(self) -> int: + return self._spatial_compression_factor + + @property + def temporal_compression_factor(self) -> int: + return self._temporal_compression_factor + + @property + def spatial_resolution(self) -> int: + return 512 + + @property + def pixel_chunk_duration(self) -> int: + return self._chunk_duration + + @property + def latent_chunk_duration(self) -> int: + return self.get_latent_num_frames(self._chunk_duration) + + @property + def latent_ch(self) -> int: + return self._latent_ch + + @contextlib.contextmanager -def _skip_source_sound_tokenizer_load(): +def _skip_source_tokenizer_load(): original_set_up_tokenizers = OmniMoTModel.set_up_tokenizers - def set_up_tokenizers_without_sound(self): - if not getattr(self.config, "sound_gen", False): - return original_set_up_tokenizers(self) - - sound_gen = self.config.sound_gen - self.config.sound_gen = False + def set_up_tokenizers_without_source_checkpoints(self): + source_tokenizer = self.config.tokenizer + source_tokenizer_target = str(source_tokenizer.get("_target_", "")) + skip_video_tokenizer = "Wan2pt2VAEInterface" in source_tokenizer_target + skip_sound_tokenizer = bool(getattr(self.config, "sound_gen", False)) + original_video_tokenizer = source_tokenizer + original_sound_gen = self.config.sound_gen + if skip_video_tokenizer: + self.config.tokenizer = { + "_target_": f"{__name__}._MetadataVideoTokenizer", + "latent_ch": int(self.config.state_ch), + "spatial_compression_factor": int(source_tokenizer.get("spatial_compression_factor", 16)), + "temporal_compression_factor": int(source_tokenizer.get("temporal_compression_factor", 4)), + "chunk_duration": int(source_tokenizer.get("chunk_duration", 93)), + "causal": bool(source_tokenizer.get("causal", True)), + } + if skip_sound_tokenizer: + self.config.sound_gen = False try: return original_set_up_tokenizers(self) finally: - self.config.sound_gen = sound_gen + self.config.tokenizer = original_video_tokenizer + self.config.sound_gen = original_sound_gen - OmniMoTModel.set_up_tokenizers = set_up_tokenizers_without_sound + OmniMoTModel.set_up_tokenizers = set_up_tokenizers_without_source_checkpoints try: yield finally: OmniMoTModel.set_up_tokenizers = original_set_up_tokenizers +def _validate_edge_action_pipeline_support() -> None: + pipeline_params = inspect.signature(Cosmos3OmniPipeline.__call__).parameters + required_params = {"action", "action_latents"} + missing_params = sorted(required_params - set(pipeline_params)) + transformer_params = inspect.signature(Cosmos3OmniTransformer.__init__).parameters + required_transformer_params = {"action_gen", "action_dim", "num_embodiment_domains"} + missing_transformer_params = sorted(required_transformer_params - set(transformer_params)) + if missing_params or missing_transformer_params: + details = [] + if missing_params: + details.append(f"pipeline call parameters {missing_params}") + if missing_transformer_params: + details.append(f"transformer constructor parameters {missing_transformer_params}") + raise RuntimeError( + "The checkpoint has action generation weights, but the checked-out Diffusers main does not provide " + "the action-enabled Cosmos3 pipeline/transformer API (missing " + f"{'; '.join(details)}). Use the merged action-enabled Diffusers main before converting an Edge " + "policy checkpoint; this converter does not drop action weights." + ) + + +def _validate_edge_transformer_support() -> None: + """Fail fast when the installed Diffusers cannot build the Edge transformer. + + The Cosmos3 Edge (Nemotron dense) backbone is structurally different from the + Qwen-based one: it has no QK-norm (`qk_norm_for_text=False`) and a non-gated + ReLU² MLP (`hidden_act`). Those are only honored when the Diffusers + `Cosmos3OmniTransformer` exposes them as constructor arguments. On a build + that predates the Edge API they are silently dropped, the transformer is + built Qwen-style, and the Edge weights fail to load with an opaque + missing-key error (`norm_q` / `norm_k` / `mlp.gate_proj`). Detect that here so + the failure is actionable. + """ + params = inspect.signature(Cosmos3OmniTransformer.__init__).parameters + accepts_var_kwargs = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + missing = sorted({"hidden_act", "qk_norm_for_text"} - set(params)) + if missing and not accepts_var_kwargs: + raise RuntimeError( + "This Diffusers build cannot convert a Cosmos3 Edge checkpoint: its Cosmos3OmniTransformer does not " + f"accept the structural Edge constructor argument(s) {missing} (it would build the Qwen-style " + "transformer, and the Edge weights would fail to load with missing norm_q/norm_k/mlp.gate_proj keys). " + "Upgrade to an Edge-capable diffusers-cosmos3 build before converting an Edge checkpoint." + ) + + +def _resolve_edge_reasoner_path(args, exported_checkpoint_path: pathlib.Path | None = None) -> pathlib.Path: + if args.reasoner_path is not None: + reasoner_path = pathlib.Path(args.reasoner_path).expanduser().absolute() + elif exported_checkpoint_path is not None and (exported_checkpoint_path / "edge_reasoner").is_dir(): + reasoner_path = exported_checkpoint_path / "edge_reasoner" + else: + from huggingface_hub import snapshot_download + + print( + "Downloading the pinned Cosmos3 Edge reasoner snapshot " + f"({args.reasoner_repo_id}@{args.reasoner_revision}) …" + ) + reasoner_path = pathlib.Path( + snapshot_download( + repo_id=args.reasoner_repo_id, + revision=args.reasoner_revision, + allow_patterns=[ + *COSMOS3_EDGE_REASONER_METADATA_FILES, + COSMOS3_EDGE_REASONER_INDEX_FILE, + "*.safetensors", + "vision_encoder/*.safetensors", + ], + ) + ) + + if not reasoner_path.is_dir(): + raise FileNotFoundError(f"Cosmos3 Edge reasoner directory not found: {reasoner_path}") + + missing_files = [ + filename + for filename in COSMOS3_EDGE_REASONER_METADATA_FILES + (COSMOS3_EDGE_REASONER_INDEX_FILE,) + if filename != "processor_config.json" and not (reasoner_path / filename).is_file() + ] + if missing_files: + raise FileNotFoundError(f"Cosmos3 Edge reasoner at {reasoner_path} is missing required files: {missing_files}") + source_index = _load_json(reasoner_path / COSMOS3_EDGE_REASONER_INDEX_FILE) + source_weight_map = source_index.get("weight_map", {}) + referenced_shards = {str(shard) for shard in source_weight_map.values()} + missing_shards = [filename for filename in referenced_shards if not (reasoner_path / filename).is_file()] + if missing_shards: + raise FileNotFoundError( + f"Cosmos3 Edge reasoner at {reasoner_path} is missing weight shards: {sorted(missing_shards)}" + ) + vision_keys = {key for key in source_weight_map if key.startswith(("model.visual.", "model.projector."))} + vision_shards = {source_weight_map[key] for key in vision_keys} + missing_vision_shards = [filename for filename in vision_shards if not (reasoner_path / filename).is_file()] + if not vision_keys or missing_vision_shards: + raise FileNotFoundError( + f"Cosmos3 Edge reasoner at {reasoner_path} has invalid vision assets: " + f"vision_keys={len(vision_keys)}, missing_shards={missing_vision_shards}" + ) + return reasoner_path + + +def _load_json(path: pathlib.Path) -> dict: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +def _save_json(payload: dict, path: pathlib.Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, sort_keys=True) + f.write("\n") + + +def _native_edge_text_config(source_config: dict) -> dict: + """Express the source text layout as standard Edge decoder blocks.""" + source_rope_parameters = source_config.get("rope_parameters") or {} + source_layer_count = int(source_config.get("num_hidden_layers", 28)) + if source_config.get("model_type") == "nemotron_h": + if source_layer_count % 2: + raise ValueError(f"Legacy Nemotron Edge text layer count must be even, got {source_layer_count}.") + source_layer_count //= 2 + native_config = { + "model_type": "cosmos3_edge_text", + "num_hidden_layers": source_layer_count, + "hidden_act": source_config.get("hidden_act", source_config.get("mlp_hidden_act", "relu2")), + "rms_norm_eps": source_config.get("rms_norm_eps", source_config.get("layer_norm_epsilon", 1e-5)), + "rope_parameters": { + "rope_type": "default", + "rope_theta": source_rope_parameters.get("rope_theta", source_config.get("rope_theta", 100_000_000.0)), + "mrope_section": source_rope_parameters.get( + "mrope_section", source_config.get("mrope_section", [24, 20, 20]) + ), + }, + } + for key in ( + "attention_bias", + "attention_dropout", + "bos_token_id", + "dtype", + "eos_token_id", + "head_dim", + "hidden_size", + "initializer_range", + "intermediate_size", + "max_position_embeddings", + "mlp_bias", + "num_attention_heads", + "num_key_value_heads", + "pad_token_id", + "use_cache", + "vocab_size", + ): + if key in source_config: + native_config[key] = source_config[key] + return native_config + + +def _native_edge_projector_config(source_config: dict) -> dict: + merger_intermediate_size = source_config.get("merger_intermediate_size") + if merger_intermediate_size is None: + merger_intermediate_size = source_config.get("merger_intermedia") + if merger_intermediate_size is None: + raise ValueError( + "Cosmos3 Edge projector config must define `merger_intermediate_size` " + "(or the legacy `merger_intermedia` alias)." + ) + return { + "model_type": "cosmos3_edge_projector", + "input_hidden_size": source_config["input_hidden_size"], + "merger_intermediate_size": merger_intermediate_size, + "out_hidden_size": source_config["out_hidden_size"], + "spatial_merge_size": source_config["spatial_merge_size"], + "use_postshuffle_norm": source_config["use_postshuffle_norm"], + } + + +def _native_edge_vision_config(source_config: dict) -> dict: + native_config = {"model_type": "cosmos3_edge_vision"} + for key in ( + "attention_dropout", + "hidden_act", + "hidden_size", + "intermediate_size", + "layer_norm_eps", + "num_attention_heads", + "num_channels", + "num_hidden_layers", + "num_patches", + "patch_size", + "spatial_merge_size", + ): + if key in source_config: + native_config[key] = source_config[key] + return native_config + + +def _native_edge_config(source_config: dict) -> dict: + native_config = { + key: value + for key, value in source_config.items() + if key not in {"architectures", "auto_map", "model_type", "projector_config", "text_config", "vision_config"} + } + native_config["architectures"] = ["Cosmos3EdgeForConditionalGeneration"] + native_config["model_type"] = "cosmos3_edge" + native_config["allow_patterns_overrides"] = ["*/*.safetensors"] + native_config["text_config"] = _native_edge_text_config(source_config["text_config"]) + native_config["vision_config"] = _native_edge_vision_config(source_config["vision_config"]) + + projector_config = _native_edge_projector_config(source_config["projector_config"]) + native_config["projector_config"] = projector_config + native_config["vision_config"].setdefault("spatial_merge_size", projector_config["spatial_merge_size"]) + if projector_config["input_hidden_size"] != native_config["vision_config"]["hidden_size"]: + raise ValueError("Edge projector input size must match the vision tower hidden size.") + if projector_config["out_hidden_size"] != native_config["text_config"]["hidden_size"]: + raise ValueError("Edge projector output size must match the text tower hidden size.") + if projector_config["use_postshuffle_norm"]: + raise ValueError("The native Cosmos3 Edge architecture only supports pre-shuffle projector normalization.") + native_config["projector_hidden_size"] = projector_config["merger_intermediate_size"] + return native_config + + +def _native_edge_tokenizer_config(source_config: dict) -> dict: + native_config = dict(source_config) + native_config["return_mm_token_type_ids"] = True + return native_config + + +def _build_edge_processor_config(reasoner_path: pathlib.Path) -> dict: + """Build current Edge processor metadata from an older reasoner export.""" + image_config = _normalize_edge_preprocessor_config( + _load_json(reasoner_path / "preprocessor_config.json"), + processor_type_key="image_processor_type", + processor_type="Cosmos3EdgeImageProcessor", + ) + video_config = _normalize_edge_preprocessor_config( + _load_json(reasoner_path / "video_preprocessor_config.json"), + processor_type_key="video_processor_type", + processor_type="Cosmos3EdgeVideoProcessor", + ) + video_config["return_metadata"] = True + return {"processor_class": "Cosmos3EdgeProcessor", "image_processor": image_config, "video_processor": video_config} + + +def _normalize_edge_preprocessor_config( + config: dict[str, Any], + processor_type_key: str, + processor_type: str, +) -> dict[str, Any]: + """Map legacy reasoner preprocessing metadata to the native Edge processor.""" + normalized_config = json.loads(json.dumps(config)) + normalized_config.pop("auto_map", None) + normalized_config["processor_class"] = "Cosmos3EdgeProcessor" + normalized_config[processor_type_key] = processor_type + return normalized_config + + +def _write_edge_vision_encoder( + reasoner_path: pathlib.Path, + output_dir: pathlib.Path, + source_weight_map: dict[str, str], +) -> None: + vision_keys = {key for key in source_weight_map if key.startswith(("model.visual.", "model.projector."))} + vision_shards = {source_weight_map[key] for key in vision_keys} + if len(vision_shards) != 1: + raise RuntimeError(f"Edge visual tensors unexpectedly span shards: {sorted(vision_shards)}") + + vision_dir = output_dir / "vision_encoder" + vision_dir.mkdir(parents=True, exist_ok=True) + vision_path = vision_dir / "model.safetensors" + print(f"Copying {len(vision_keys)} Cosmos3 Edge vision/projector tensors to {vision_path} …") + source_vision_path = reasoner_path / next(iter(vision_shards)) + from safetensors import safe_open + from safetensors.torch import save_file + + vision_tensors: dict[str, torch.Tensor] = {} + with safe_open(str(source_vision_path), framework="pt", device="cpu") as source_file: + for key in source_file.keys(): + if key in vision_keys: + vision_tensors[key] = source_file.get_tensor(key) + if set(vision_tensors) != vision_keys: + missing_keys = sorted(vision_keys - set(vision_tensors)) + raise RuntimeError(f"Missing Edge vision/projector tensors while writing {vision_path}: {missing_keys}") + save_file(vision_tensors, str(vision_path), metadata={"format": "pt"}) + # This is an auxiliary shard consumed through the root safetensors index; + # it is not a standalone Diffusers component. Do not leave a stale legacy + # config beside it when converting into an existing output directory. + (vision_dir / "config.json").unlink(missing_ok=True) + + +def _write_diffusers_safetensors_index(output_dir: pathlib.Path) -> None: + """Write the existing Edge root index over shared component shards.""" + metadata = {"total_size": 0} + weight_map: dict[str, str] = {} + for subdir in ("transformer", "vision_encoder"): + component_dir = output_dir / subdir + if not component_dir.is_dir(): + continue + for safetensors_path in sorted(component_dir.glob("*.safetensors")): + with safetensors_path.open("rb") as file: + header_size = struct.unpack(" bool: + """Return whether a Diffusers key belongs in the existing Edge root index.""" + if name.startswith("model."): + return True + if name in {"embed_tokens.weight", "lm_head.weight", "norm.weight"}: + return True + if not name.startswith("layers."): + return False + generator_markers = ( + ".input_layernorm_moe_gen.", + ".post_attention_layernorm_moe_gen.", + ".mlp_moe_gen.", + ".self_attn.add_", + ".self_attn.norm_added_", + ".self_attn.to_add_out.", + ) + return not any(marker in name for marker in generator_markers) + + +def _copy_edge_reasoner_metadata(reasoner_path: pathlib.Path, output_dir: pathlib.Path) -> None: + """Copy Edge metadata and its vision sidecar into a Diffusers repository.""" + print(f"Writing the shared Cosmos3 Edge metadata into {output_dir} …") + source_index = _load_json(reasoner_path / COSMOS3_EDGE_REASONER_INDEX_FILE) + for filename in COSMOS3_EDGE_REASONER_METADATA_FILES: + if filename == "config.json": + continue + if filename == "preprocessor_config.json": + _save_json( + _normalize_edge_preprocessor_config( + _load_json(reasoner_path / filename), + processor_type_key="image_processor_type", + processor_type="Cosmos3EdgeImageProcessor", + ), + output_dir / filename, + ) + continue + if filename == "video_preprocessor_config.json": + _save_json( + _normalize_edge_preprocessor_config( + _load_json(reasoner_path / filename), + processor_type_key="video_processor_type", + processor_type="Cosmos3EdgeVideoProcessor", + ) + | {"return_metadata": True}, + output_dir / filename, + ) + continue + if filename == "processor_config.json" and not (reasoner_path / filename).is_file(): + _save_json(_build_edge_processor_config(reasoner_path), output_dir / filename) + continue + if filename == "tokenizer_config.json": + _save_json(_native_edge_tokenizer_config(_load_json(reasoner_path / filename)), output_dir / filename) + continue + shutil.copy2(reasoner_path / filename, output_dir / filename) + + for filename in ( + "configuration_nemotron_siglip2_h.py", + "modeling_cosmos3_edge_omni.py", + "modeling_nemotron_siglip2_h.py", + "processing.py", + ): + (output_dir / filename).unlink(missing_ok=True) + _save_json(_native_edge_config(_load_json(reasoner_path / "config.json")), output_dir / "config.json") + _write_edge_vision_encoder(reasoner_path, output_dir, source_index["weight_map"]) + for filename in ("00000.safetensors", "00001.safetensors", "model.safetensors"): + (output_dir / filename).unlink(missing_ok=True) + _write_diffusers_safetensors_index(output_dir) + + +def _write_edge_transformer_config(output_dir: pathlib.Path, model_cfg: Any) -> None: + """Persist Edge fields and the existing multi-axis RoPE format.""" + transformer_config_path = output_dir / "transformer" / "config.json" + transformer_config = _load_json(transformer_config_path) + tokenizer_config = _get_config_value(model_cfg, name="tokenizer", default=None) + temporal_compression_factor = _get_config_value(tokenizer_config, name="temporal_compression_factor", default=None) + if temporal_compression_factor is None: + raise ValueError("Cosmos3 Edge export config is missing tokenizer.temporal_compression_factor.") + transformer_config.update( + { + "backbone_type": "cosmos3_edge_nemotron_dense", + "temporal_compression_factor": int(temporal_compression_factor), + } + ) + rope_axes_dim = transformer_config.get("rope_axes_dim") + if rope_axes_dim is not None: + transformer_config["rope_scaling"] = {"mrope_section": rope_axes_dim} + _save_json(transformer_config, transformer_config_path) + + +def _normalize_edge_model_index(output_dir: pathlib.Path) -> None: + """Keep Edge pipeline metadata consistent with the existing Hub export.""" + model_index_path = output_dir / "model_index.json" + model_index = _load_json(model_index_path) + model_index.update( + { + "default_use_system_prompt": False, + "text_tokenizer": ["transformers", "PreTrainedTokenizerFast"], + "use_native_flow_schedule": True, + } + ) + _save_json(model_index, model_index_path) + + +def _add_edge_reasoner_to_pipeline(args) -> None: + output_dir = pathlib.Path(args.output).expanduser().absolute() + expected_paths = ("model_index.json", "scheduler", "text_tokenizer", "transformer", "vae") + missing_paths = [str(output_dir / path) for path in expected_paths if not (output_dir / path).exists()] + if missing_paths: + raise FileNotFoundError( + "Expected an existing Cosmos3 Edge Diffusers pipeline before adding its reasoner; " + f"missing paths: {missing_paths}" + ) + + reasoner_path = _resolve_edge_reasoner_path(args) + _copy_edge_reasoner_metadata(reasoner_path, output_dir) + _normalize_edge_model_index(output_dir) + print("Done.") + + class Args(pydantic.BaseModel): checkpoint_path: pathlib.Path - """Named checkpoint (e.g. 'Cosmos3-Nano') or path to DCP checkpoint dir.""" + """Path to a Stage-1 HF export or a named non-Edge HF checkpoint.""" output: str """Directory to save the converted diffusers model.""" save_pipeline: bool = False @@ -581,10 +1197,42 @@ class Args(pydantic.BaseModel): """Qwen3-VL model/config to instantiate model.visual.* weights.""" skip_vision_encoder: bool = False """Do not save vision_encoder/ when saving a full pipeline.""" + distilled_scheduler: Literal["auto", "on", "off"] = "auto" + """How to export the scheduler for distilled (few-step) checkpoints. + + * auto: export the distilled scheduler when `fixed_step_sampler_config` is present. + * on: require and always export the distilled scheduler. + * off: always export the UniPC scheduler. + """ + + include_reasoner: bool = True + """Include the pinned Edge metadata and vision sidecar in an exported Edge checkpoint.""" + + reasoner_repo_id: str = COSMOS3_EDGE_REASONER + """Hugging Face repository containing the Cosmos3 Edge reasoner checkpoint.""" + + reasoner_revision: str = COSMOS3_EDGE_REASONER_REVISION + """Pinned revision of the Cosmos3 Edge reasoner checkpoint.""" + + reasoner_path: pathlib.Path | None = None + """Optional local Cosmos3 Edge reasoner snapshot, used instead of downloading it.""" + + copy_edge_reasoner: bool = False + """Add the shared-weight Edge reasoner to an existing Diffusers pipeline at `output`.""" def convert_model_to_diffusers(args: Args) -> None: dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype] + + if args.copy_edge_reasoner: + _add_edge_reasoner_to_pipeline(args) + return + + # A raw DCP is loaded directly by ``Cosmos3OmniModel.from_pretrained_dcp`` below + # (this repo's existing DCP -> Diffusers path). Only the Edge reasoner flow + # requires a config-driven HF export first, which the wrapper enforces + # separately (its ``is_edge_model and .metadata`` guard). + sound_tokenizer_path = ( pathlib.Path(args.sound_tokenizer_path).expanduser().absolute() if args.sound_tokenizer_path else None ) @@ -601,13 +1249,26 @@ def convert_model_to_diffusers(args: Args) -> None: raise FileNotFoundError(f"Sound tokenizer config not found: {sound_tokenizer_config_path}") checkpoint_path = args.checkpoint_path + is_edge_model = _is_edge_exported_checkpoint(checkpoint_path) + if is_edge_model: + # Fail fast (before the reasoner download) if this Diffusers build lacks + # the Edge transformer API, instead of dying later on opaque missing keys. + _validate_edge_transformer_support() + edge_reasoner_path = None + if is_edge_model and args.include_reasoner: + if not args.save_pipeline: + raise ValueError( + "Including the Cosmos3 Edge reasoner requires --save-pipeline so the native Transformers " + "repository can be written alongside the Diffusers components." + ) + edge_reasoner_path = _resolve_edge_reasoner_path(args, exported_checkpoint_path=checkpoint_path) - log.info("Instantiating model and loading weights from DCP checkpoint …") + log.info("Instantiating model and loading weights from exported HF checkpoint …") log.info("Skipping source AVAE tokenizer instantiation during converter-only model load …") - with _skip_source_sound_tokenizer_load(): + with _skip_source_tokenizer_load(): _tmp = Cosmos3OmniModel.from_pretrained_dcp(checkpoint_path).model - # Extract network components and architecture config from DCP model + # Extract network components and architecture config from the Stage 1 HF export. language_model = _tmp.net.language_model vae2llm = _tmp.net.vae2llm llm2vae = _tmp.net.llm2vae @@ -646,6 +1307,8 @@ def convert_model_to_diffusers(args: Args) -> None: action_gen = bool( _get_config_value(net_cfg, model_cfg, name="action_gen", default=False) or has_action_projection_weights ) + if action_gen and args.save_pipeline: + _validate_edge_action_pipeline_support() action_dim = _get_config_value(net_cfg, model_cfg, name="action_dim", default=None) if action_dim is None and action_proj_in is not None: action_dim = getattr(action_proj_in, "input_size", None) @@ -696,57 +1359,122 @@ def convert_model_to_diffusers(args: Args) -> None: f"action projection weights: {missing_action_modules}." ) - has_vision_encoder_weights = _checkpoint_has_weight_prefix( - checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX - ) or bool(_checkpoint_vision_subfolder_files(checkpoint_path)) + # A ViT sidecar can be built from any of three sources: the loaded reasoner + # model's own `.visual` tower (only present when include_visual is truthy), + # root `model.visual.*` weights, or a vision_encoder/ subfolder. + has_source_visual = any( + module is not None and next(module.parameters(), None) is not None + for module in ( + getattr(_tmp, "visual", None), + getattr(getattr(_tmp, "net", None), "visual", None), + getattr(getattr(getattr(_tmp, "net", None), "language_model", None), "visual", None), + ) + ) + has_vision_encoder_weights = ( + has_source_visual + or _checkpoint_has_weight_prefix(checkpoint_path, VISION_ENCODER_CHECKPOINT_PREFIX) + or bool(_checkpoint_vision_subfolder_files(checkpoint_path)) + ) vision_gen = bool( _get_config_value(net_cfg, model_cfg, name="vision_gen", default=False) or has_vision_encoder_weights ) - include_vision_encoder = bool(args.save_pipeline and vision_gen and not args.skip_vision_encoder) + want_vision_encoder = bool(args.save_pipeline and vision_gen and not args.skip_vision_encoder and not is_edge_model) + # Only build the vision_encoder/ sidecar when the checkpoint actually ships + # extractable ViT weights. A generation-only checkpoint (e.g. include_visual + # unset) reports vision_gen=True but has no reasoner ViT to export — auto-skip + # it with a warning instead of failing, so --skip-vision-encoder is not required. + include_vision_encoder = want_vision_encoder and has_vision_encoder_weights vision_encoder = None if include_vision_encoder: vision_encoder = _load_vision_encoder(checkpoint_path, _tmp, args.vision_encoder_model, dtype) + elif want_vision_encoder: + log.warning( + f"No extractable vision encoder weights found (no loaded reasoner ViT, no '{VISION_ENCODER_CHECKPOINT_PREFIX}*' " + "root weights, no vision_encoder/ subfolder); skipping vision_encoder/ save. Convert a checkpoint exported " + "with include_visual=True (or one bundling a vision_encoder/) to include the reasoner vision tower." + ) elif args.save_pipeline and vision_gen and args.skip_vision_encoder: log.info("Skipping vision_encoder/ save because --skip-vision-encoder was set.") del _tmp - # Init diffusers Cosmos3OmniTransformer with full architecture config from DCP + source_vlm_config = _get_config_value(model_cfg, name="vlm_config", default=None) + source_model_instance = _get_config_value(source_vlm_config, name="model_instance", default=None) + source_vlm_instance_config = _get_config_value(source_model_instance, name="config", default=None) + hidden_act = _get_config_value(lm_cfg, name="hidden_act", default=None) + if hidden_act is None: + hidden_act = _get_config_value(lm_cfg, name="mlp_hidden_act", default="silu") + qk_norm_for_text = bool(_get_config_value(source_vlm_instance_config, name="qk_norm_for_text", default=True)) + use_und_k_norm_for_gen = bool( + _get_config_value(source_vlm_instance_config, name="use_und_k_norm_for_gen", default=False) + ) + log.info( + "Building Diffusers transformer from exported config: " + f"hidden_act={hidden_act!r}, qk_norm_for_text={qk_norm_for_text}, " + f"use_und_k_norm_for_gen={use_und_k_norm_for_gen}, action_gen={action_gen}" + ) + + # Build Diffusers Cosmos3OmniTransformer from the exported HF architecture config. + transformer_kwargs: dict[str, Any] = { + "attention_bias": lm_cfg.attention_bias, + "attention_dropout": lm_cfg.attention_dropout, + "base_fps": base_fps, + "enable_fps_modulation": enable_fps_modulation, + "head_dim": head_dim, + "hidden_size": hidden_size, + "intermediate_size": lm_cfg.intermediate_size, + "latent_channel": latent_channel, + "latent_patch_size": latent_patch_size, + "action_dim": action_dim, + "action_gen": action_gen, + "num_embodiment_domains": num_embodiment_domains, + "num_attention_heads": num_attention_heads, + "num_hidden_layers": num_hidden_layers, + "num_key_value_heads": num_key_value_heads, + "patch_latent_dim": patch_latent_dim, + "rms_norm_eps": lm_cfg.rms_norm_eps, + "rope_scaling": lm_cfg.rope_scaling, + "rope_theta": lm_cfg.rope_theta, + "sound_dim": sound_dim, + "sound_gen": sound_gen, + "sound_latent_fps": sound_latent_fps, + "timestep_scale": timestep_scale, + "unified_3d_mrope_reset_spatial_ids": unified_3d_mrope_reset_spatial_ids, + "unified_3d_mrope_temporal_modality_margin": unified_3d_mrope_temporal_modality_margin, + "vocab_size": lm_cfg.vocab_size, + } + # hidden_act / qk_norm_for_text / rope_axes_dim are only explicit constructor + # arguments on newer Diffusers Cosmos3OmniTransformer builds. Older builds + # (e.g. the upstreamed diffusers 0.39.0) fix hidden_act/qk_norm_for_text and + # derive rope_axes_dim internally from rope_scaling, and reject these as + # unexpected keyword arguments. Pass each only when the installed constructor + # accepts it so conversion works on both old and new builds; warn (not silent) + # when omitting so a non-default exported value is not quietly dropped. + optional_transformer_kwargs = { + "hidden_act": hidden_act, + "qk_norm_for_text": qk_norm_for_text, + "use_und_k_norm_for_gen": use_und_k_norm_for_gen, + "rope_axes_dim": _get_config_value(lm_cfg.rope_scaling, name="mrope_section", default=None), + } + transformer_params = inspect.signature(Cosmos3OmniTransformer.__init__).parameters + accepts_var_kwargs = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in transformer_params.values()) + for name, value in optional_transformer_kwargs.items(): + if accepts_var_kwargs or name in transformer_params: + transformer_kwargs[name] = value + else: + log.warning( + f"Installed Diffusers Cosmos3OmniTransformer does not accept {name!r}; omitting it " + "(the checked-out Diffusers derives/fixes this field internally). Upgrade Diffusers to a " + f"build that exposes it to honor the exported value ({name}={value!r})." + ) with init_empty_weights(): - transformer = Cosmos3OmniTransformer( - attention_bias=lm_cfg.attention_bias, - attention_dropout=lm_cfg.attention_dropout, - base_fps=base_fps, - enable_fps_modulation=enable_fps_modulation, - head_dim=head_dim, - hidden_size=hidden_size, - intermediate_size=lm_cfg.intermediate_size, - latent_channel=latent_channel, - latent_patch_size=latent_patch_size, - action_dim=action_dim, - action_gen=action_gen, - num_embodiment_domains=num_embodiment_domains, - num_attention_heads=num_attention_heads, - num_hidden_layers=num_hidden_layers, - num_key_value_heads=num_key_value_heads, - patch_latent_dim=patch_latent_dim, - rms_norm_eps=lm_cfg.rms_norm_eps, - rope_scaling=lm_cfg.rope_scaling, - rope_theta=lm_cfg.rope_theta, - sound_dim=sound_dim, - sound_gen=sound_gen, - sound_latent_fps=sound_latent_fps, - timestep_scale=timestep_scale, - unified_3d_mrope_reset_spatial_ids=unified_3d_mrope_reset_spatial_ids, - unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin, - vocab_size=lm_cfg.vocab_size, - ) + transformer = Cosmos3OmniTransformer(**transformer_kwargs) state_dict = _remap_language_model_state_dict(language_model.state_dict()) for k, v in vae2llm.state_dict().items(): state_dict[f"proj_in.{k}"] = v for k, v in llm2vae.state_dict().items(): state_dict[f"proj_out.{k}"] = v for k, v in time_embedder.state_dict().items(): - state_dict[f"time_embedder.{_TIME_EMBEDDER_KEY_REMAP[k]}"] = v + state_dict[f"time_embedder.{_remap_time_embedder_key(k)}"] = v if action_gen: for k, v in action_proj_in.state_dict().items(): state_dict[f"action_proj_in.{k}"] = v @@ -787,8 +1515,46 @@ def convert_model_to_diffusers(args: Args) -> None: "is required when saving a full pipeline." ) + if is_edge_model and args.save_pipeline and args.skip_vision_encoder: + log.info("Skipping standalone vision_encoder/ construction; the embedded Edge reasoner owns those weights.") + + fixed_step_sampler_cfg = _resolve_fixed_step_sampler_config(model_cfg) + if args.distilled_scheduler == "on": + if fixed_step_sampler_cfg is None: + raise ValueError( + "distilled_scheduler='on' was requested, but the checkpoint does not define " + "`fixed_step_sampler_config`." + ) + use_distilled_scheduler = True + elif args.distilled_scheduler == "off": + if fixed_step_sampler_cfg is not None: + raise ValueError( + "distilled_scheduler='off' was requested, but the checkpoint defines " + "`fixed_step_sampler_config`. Exporting a UniPC scheduler for a distilled " + "checkpoint produces a broken pipeline (it is trained for the fixed few-step " + "schedule). Use 'auto' or 'on' instead." + ) + use_distilled_scheduler = False + else: + use_distilled_scheduler = fixed_step_sampler_cfg is not None + + if use_distilled_scheduler and fixed_step_sampler_cfg is not None: + log.info( + "Detected distilled checkpoint scheduler config: " + f"sample_type={fixed_step_sampler_cfg['sample_type']}, t_list={fixed_step_sampler_cfg['t_list']}" + ) + if args.save_pipeline: - text_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-8B-Instruct") + tokenizer_source = str(edge_reasoner_path) if is_edge_model else "Qwen/Qwen3-VL-8B-Instruct" + text_tokenizer = AutoTokenizer.from_pretrained(tokenizer_source) + if is_edge_model: + for token in ("<|vision_start|>", "<|vision_end|>"): + token_id = text_tokenizer.convert_tokens_to_ids(token) + if token_id is None or token_id < 0 or token_id >= transformer.config.vocab_size: + raise ValueError( + f"Cosmos3 Edge tokenizer token {token!r} has invalid ID {token_id!r} for " + f"vocab_size={transformer.config.vocab_size}." + ) diffusers_vae = AutoencoderKLWan.from_pretrained( "Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16 @@ -798,32 +1564,68 @@ def convert_model_to_diffusers(args: Args) -> None: if include_sound_tokenizer: sound_tokenizer = _build_sound_tokenizer(sound_tokenizer_path, sound_tokenizer_config_path) - # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps. - # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281. - # EDM sigma = flow_sigma / (1 - flow_sigma), so: - # sigma_max = 0.9998 / 0.0002 = 4999 (but capped at 200 to avoid duplicate - # integer timesteps from Karras clustering near the top) - # sigma_min = 0.1281 / (1 - 0.1281) = 0.1281 / 0.8719 ≈ 0.147 - scheduler = UniPCMultistepScheduler( - use_karras_sigmas=True, - use_flow_sigmas=True, - prediction_type="flow_prediction", - sigma_max=200.0, - sigma_min=0.147, - ) + if use_distilled_scheduler: + assert fixed_step_sampler_cfg is not None + # Distilled checkpoints are trained against a fixed flow-sigma schedule + # (`fixed_step_sampler_config.t_list`). We export FlowMatchEulerDiscreteScheduler because: + # 1) it is a flow-prediction scheduler (same parameterization expected + # by distilled Cosmos3 checkpoints), + # 2) it supports explicit sigma injection through `set_timesteps(..., sigmas=...)`, + # 3) `stochastic_sampling` maps directly to the distilled sample_type: + # `sde` -> stochastic, `ode` -> deterministic. + scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=fixed_step_sampler_cfg["num_train_timesteps"], + shift=1.0, + use_dynamic_shifting=False, + use_karras_sigmas=False, + use_exponential_sigmas=False, + use_beta_sigmas=False, + invert_sigmas=False, + stochastic_sampling=(fixed_step_sampler_cfg["sample_type"] == "sde"), + ) + # Persist checkpoint-defined fixed-step settings so distilled inference can + # call scheduler.set_timesteps(..., sigmas=t_list) at runtime. + scheduler.register_to_config( + fixed_step_sampler_config={ + "t_list": fixed_step_sampler_cfg["t_list"], + "sample_type": fixed_step_sampler_cfg["sample_type"], + }, + fixed_step_requires_explicit_sigmas=True, + ) + else: + # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps. + # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281. + # EDM sigma = flow_sigma / (1 - flow_sigma), so: + # sigma_max = 0.9998 / 0.0002 = 4999 (but capped at 200 to avoid duplicate + # integer timesteps from Karras clustering near the top) + # sigma_min = 0.1281 / (1 - 0.1281) = 0.1281 / 0.8719 ≈ 0.147 + scheduler = UniPCMultistepScheduler( + use_karras_sigmas=True, + use_flow_sigmas=True, + prediction_type="flow_prediction", + sigma_max=200.0, + sigma_min=0.147, + ) # enable_safety_checker=False: constructing the checker imports # cosmos_guardrail, which the converter environment does not ship. # Consumers can opt back in at load time via - # `Cosmos3OmniPipeline.from_pretrained(..., enable_safety_checker=True)`. - pipeline = Cosmos3OmniPipeline( - transformer=transformer, - text_tokenizer=text_tokenizer, - vae=diffusers_vae, - scheduler=scheduler, - sound_tokenizer=sound_tokenizer, - enable_safety_checker=False, - ) + # Use the pipeline class exported by the checked-out Diffusers main. + # Keep optional components gated by the constructor signature because + # the main branch does not expose the legacy sound-tokenizer and safety + # checker arguments used by older Cosmos3 pipeline implementations. + pipeline_kwargs: dict[str, Any] = { + "transformer": transformer, + "text_tokenizer": text_tokenizer, + "vae": diffusers_vae, + "scheduler": scheduler, + } + pipeline_parameters = inspect.signature(Cosmos3OmniPipeline.__init__).parameters + if "sound_tokenizer" in pipeline_parameters: + pipeline_kwargs["sound_tokenizer"] = sound_tokenizer + if "enable_safety_checker" in pipeline_parameters: + pipeline_kwargs["enable_safety_checker"] = False + pipeline = Cosmos3OmniPipeline(**pipeline_kwargs) log.info(f"Saving full pipeline to {output_dir} …") pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") if vision_encoder is not None: @@ -831,6 +1633,11 @@ def convert_model_to_diffusers(args: Args) -> None: # the transformers/vLLM consumers of the exported repository. log.info(f"Saving Qwen3-VL vision encoder to {output_dir / 'vision_encoder'} …") vision_encoder.save_pretrained(str(output_dir / "vision_encoder"), safe_serialization=True) + if is_edge_model: + assert edge_reasoner_path is not None + _copy_edge_reasoner_metadata(edge_reasoner_path, output_dir) + _normalize_edge_model_index(output_dir) + _write_edge_transformer_config(output_dir, model_cfg) else: log.info(f"Saving transformer to {output_dir} …") transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") diff --git a/cosmos_framework/scripts/_export_model_helpers.py b/cosmos_framework/scripts/_export_model_helpers.py new file mode 100644 index 00000000..cf009955 --- /dev/null +++ b/cosmos_framework/scripts/_export_model_helpers.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Pure helpers for 'cosmos_framework.scripts.export_model'. + +Kept free of heavy imports (torch, init_script side effects) so they are +unit-testable without a GPU or downloads; registry-touching helpers import +lazily and accept an injectable download callable. +""" + +from __future__ import annotations + +import contextlib +import json +import re +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable + +import numpy as np + +if TYPE_CHECKING: + from cosmos_framework.utils.checkpoint_db import CheckpointDirHf + +_EDGE_MODEL_NAME = "nvidia/Cosmos3-Edge-Reasoner" + +# HF-standard processor/tokenizer files shipped at the snapshot root. The exported +# config.json is deliberately absent (it must never be overwritten by bundling). +PROCESSOR_FILE_PATTERNS = ( + "preprocessor_config.json", + "processor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.json", + "merges.txt", + "chat_template*", + "video_preprocessor*", +) + +# Keys the bundled vision_encoder/config.json must carry so inference never needs +# the hub repo's top-level config (see Nemotron3DenseVLTextForCausalLM._ensure_vision_tower). +_BUNDLE_CONFIG_DICT_KEYS = ("vision_config", "projector_config") +_BUNDLE_CONFIG_TOKEN_ID_KEYS = ("image_token_id", "video_token_id", "vision_start_token_id") + +_HF_SNAPSHOT_RE = re.compile(r"/snapshots/([0-9a-f]{40})(?:/|$)") + + +def is_edge_model(model_dict: dict[str, Any]) -> bool: + """Return True when the model config describes a Cosmos3-Edge reasoner backbone.""" + vlm_config = (model_dict.get("config") or {}).get("vlm_config") or {} + if vlm_config.get("model_name") == _EDGE_MODEL_NAME: + return True + model_instance = vlm_config.get("model_instance") or {} + if "Nemotron3DenseVLTextForCausalLM" in str(model_instance.get("_target_", "")): + return True + backbone_path = (vlm_config.get("pretrained_weights") or {}).get("backbone_path") or "" + return "Cosmos3-Edge" in backbone_path + + +def reasoner_vision_capable(model_dict: dict[str, Any]) -> bool: + """Whether the checkpoint being exported can serve reasoner image/video prompts. + + Dict-shaped mirror of 'inference.args._reasoner_vision_capable' (keep the + two in sync). Edge is always capable (lazy tower); Qwen-family (Nano/Super) + needs a truthy 'include_visual'. Unknown families are treated as capable + (never block). + """ + vlm_config = (model_dict.get("config") or {}).get("vlm_config") or {} + model_instance = vlm_config.get("model_instance") or {} + target = str(model_instance.get("_target_", "") if isinstance(model_instance, dict) else "") + model_name = vlm_config.get("model_name") or "" + if "Nemotron3DenseVLTextForCausalLM" in target or "Cosmos3-Edge" in model_name: + return True + if "Qwen" not in target and not model_name.startswith("Qwen/"): + return True + config = model_instance.get("config") if isinstance(model_instance, dict) else None + include_visual = config.get("include_visual") if isinstance(config, dict) else None + return bool(include_visual) + + +def set_include_visual(model_dict: dict[str, Any], value: bool) -> bool: + """Set 'include_visual' on the 'create_vlm_config' node; return True when applied. + + The node's kwargs are applied to the MoT config wrapper via setattr at load + time ('configs.base.defaults.reasoner.create_vlm_config'), and the wrapper's + 'include_visual' gates visual-tower construction + ('_MoTConfigBase.vision_config'), so this is what makes a '--no-vit' export + loadable without vision weights. + """ + vlm_config = (model_dict.get("config") or {}).get("vlm_config") or {} + model_instance = vlm_config.get("model_instance") + if not isinstance(model_instance, dict): + return False + config = model_instance.get("config") + if not isinstance(config, dict): + return False + config["include_visual"] = value + return True + + +def resolve_vision_bundle_dirs(vlm_checkpoint_path: Path) -> tuple[Path, Path | None]: + """Locate the 'vision_encoder/' weights dir and the snapshot root. + + Accepts either an nvidia/Cosmos3-Edge snapshot root (contains + 'vision_encoder/model.safetensors') or a 'vision_encoder/' directory passed + directly via '--vit-checkpoint-path'. The returned root is None when the + weights directory was passed directly and its parent has no 'config.json'. + """ + if (vlm_checkpoint_path / "vision_encoder" / "model.safetensors").is_file(): + return vlm_checkpoint_path / "vision_encoder", vlm_checkpoint_path + if (vlm_checkpoint_path / "model.safetensors").is_file(): + root = vlm_checkpoint_path.parent + return vlm_checkpoint_path, root if (root / "config.json").is_file() else None + raise FileNotFoundError( + f"No 'vision_encoder/model.safetensors' under '{vlm_checkpoint_path}' (expected an " + "nvidia/Cosmos3-Edge snapshot root or a vision_encoder/ directory)." + ) + + +def build_vision_encoder_bundle_config(vision_encoder_dir: Path, snapshot_root: Path | None) -> dict[str, Any]: + """Build the merged, self-describing 'vision_encoder/config.json' for the bundle. + + Configs and multimodal token ids resolve standalone-first (matching + '_ensure_vision_tower'), falling back to the snapshot's top-level + 'config.json', so inference never needs the hub repo's top-level config. + """ + standalone_file = vision_encoder_dir / "config.json" + standalone: dict[str, Any] | None = None + if standalone_file.is_file(): + standalone = json.loads(standalone_file.read_text()) + top: dict[str, Any] | None = None + if snapshot_root is not None and (snapshot_root / "config.json").is_file(): + top = json.loads((snapshot_root / "config.json").read_text()) + if standalone is None and top is None: + raise FileNotFoundError( + f"No 'config.json' found in '{vision_encoder_dir}' or its snapshot root; cannot " + "build the vision_encoder bundle config. Point --vit-checkpoint-path at a full " + "nvidia/Cosmos3-Edge snapshot directory." + ) + + merged: dict[str, Any] = {} + for key in _BUNDLE_CONFIG_DICT_KEYS: + value = (standalone or {}).get(key) or (top or {}).get(key) + if not isinstance(value, dict): + raise ValueError( + f"Missing '{key}' in both the standalone vision_encoder config and the snapshot's " + "top-level config.json; cannot build the vision_encoder bundle config." + ) + merged[key] = dict(value) + for key in _BUNDLE_CONFIG_TOKEN_ID_KEYS: + value = (standalone or {}).get(key) + if value is None: + value = (top or {}).get(key) + if not isinstance(value, int): + raise ValueError( + f"Cannot resolve '{key}' for the vision_encoder bundle config (absent from both the " + "standalone vision_encoder config and the snapshot's top-level config.json). Point " + "--vit-checkpoint-path at a full nvidia/Cosmos3-Edge snapshot directory." + ) + merged[key] = value + return merged + + +def bundle_vision_encoder(vision_encoder_dir: Path, snapshot_root: Path | None, output_dir: Path) -> None: + """Copy 'vision_encoder/model.safetensors' verbatim and write the merged config.""" + bundle_config = build_vision_encoder_bundle_config(vision_encoder_dir, snapshot_root) + output_vision_dir = output_dir / "vision_encoder" + output_vision_dir.mkdir(parents=True, exist_ok=True) + shutil.copyfile(vision_encoder_dir / "model.safetensors", output_vision_dir / "model.safetensors") + (output_vision_dir / "config.json").write_text(json.dumps(bundle_config, indent=2, sort_keys=True) + "\n") + + +def bundle_processor_files(snapshot_root: Path, output_dir: Path) -> list[str]: + """Copy HF-standard processor/tokenizer files from the snapshot root; return copied names. + + All-or-nothing: on a mid-copy failure the already-copied files are removed + (best-effort) before the exception propagates — inference prefers local + processor files over the hub, so a partial set must never be left behind. + """ + copied: list[str] = [] + try: + for pattern in PROCESSOR_FILE_PATTERNS: + for source in sorted(snapshot_root.glob(pattern)): + if not source.is_file(): + continue + assert source.name != "config.json", "processor bundling must never overwrite the exported config.json" + shutil.copyfile(source, output_dir / source.name) + copied.append(source.name) + except Exception: + for name in copied: + with contextlib.suppress(OSError): + (output_dir / name).unlink() + raise + return copied + + +# Hugging Face repo ids are 'org/name'; anything else in a tokenizer node +# (local directory, URI) has no hub snapshot to bundle from. +_HF_REPO_ID_RE = re.compile(r"[\w.\-]+/[\w.\-]+") + +# Registered tokenizer snapshots live under this sanitized S3 prefix in the +# checkpoint registry ('inference.common.checkpoints.register_checkpoints'); +# it is the prefix the training/inference download paths consult for plain +# repo ids ('configs.base.defaults.reasoner.download_tokenizer_files', +# 'maybe_download_hf_model_from_s3'). +_TOKENIZER_REGISTRY_URI_PREFIX = "s3://bucket/cosmos3/pretrained/huggingface" + +# Files fetched when downloading a processor snapshot for bundling; mirrors +# '_PROCESSOR_HF_INCLUDE' in 'cosmos_framework.inference.inference'. +PROCESSOR_HF_INCLUDE = ("*.json", "*.jinja", "merges.txt", "vocab.json") + +# Files fetched when export resolves the Edge ViT bundle from the hub: the +# 'vision_encoder/' artifact plus the HF-standard config/processor files +# (~1 GB) — never the generation shards ('transformer/', 'vae/') or demo +# assets the bundle doesn't need (tens of GB on a cold cache). Applied at the +# export call site only: the registry entry itself must stay unfiltered +# because training's backbone seeding downloads the full repo through it +# ('load_language_model' resolves the root model.safetensors.index.json, +# whose weight_map points into 'transformer/*.safetensors'). +EDGE_VIT_BUNDLE_HF_INCLUDE = ("vision_encoder/*", *PROCESSOR_HF_INCLUDE) + + +def processor_source_from_tokenizer_node(tokenizer_node: Any) -> dict[str, Any] | None: + """Extract the VLM processor origin from a 'vlm_config.tokenizer' node. + + Accepts 'repository' (+'revision'/'subdir') nodes, or a plain HF repo id in + 'tokenizer_type' / 'pretrained_model_name' (revision None, pinned later via + the checkpoint registry). Returns {'repository', 'revision', 'subdir'}, or + None when the node names no HF repository (missing node, local dir, URI). + """ + if not isinstance(tokenizer_node, dict): + return None + repository = tokenizer_node.get("repository") + if isinstance(repository, str) and _HF_REPO_ID_RE.fullmatch(repository): + return { + "repository": repository, + "revision": tokenizer_node.get("revision"), + "subdir": tokenizer_node.get("subdir") or "", + } + repo_id = tokenizer_node.get("tokenizer_type") or tokenizer_node.get("pretrained_model_name") + if isinstance(repo_id, str) and _HF_REPO_ID_RE.fullmatch(repo_id): + return {"repository": repo_id, "revision": None, "subdir": ""} + return None + + +def resolve_processor_download(source: dict[str, Any]) -> CheckpointDirHf: + """Build the 'CheckpointDirHf' download spec for a processor source. + + An explicit revision is used as-is; revision None is pinned via the + checkpoint registry, with unregistered repos falling back to 'main'. Only + HF-standard processor/config files are included in the download. + """ + from cosmos_framework.utils.checkpoint_db import CheckpointConfig, CheckpointDirHf, sanitize_uri + + repository: str = source["repository"] + revision: str | None = source.get("revision") + subdirectory: str = source.get("subdir") or "" + if revision is None: + registered = CheckpointConfig.maybe_from_uri(sanitize_uri(f"{_TOKENIZER_REGISTRY_URI_PREFIX}/{repository}")) + if registered is not None and isinstance(registered.hf, CheckpointDirHf): + repository = registered.hf.repository + revision = registered.hf.revision + subdirectory = registered.hf.subdirectory + else: + revision = "main" + return CheckpointDirHf( + repository=repository, + revision=revision, + subdirectory=subdirectory, + include=PROCESSOR_HF_INCLUDE, + ) + + +def bundle_processor_from_tokenizer_node( + tokenizer_node: Any, + output_dir: Path, + *, + download: Callable[[CheckpointDirHf], str] | None = None, +) -> dict[str, Any] | None: + """Best-effort processor bundling from the same origin inference uses. + + Resolves the 'vlm_config.tokenizer' node, copies the HF-standard processor + files into the export root, and returns the manifest 'processor_source' + entry. Any failure logs a warning and returns None — processor bundling + must never fail an export. 'download' is injectable for tests. + """ + from cosmos_framework.utils import log + + source = processor_source_from_tokenizer_node(tokenizer_node) + if source is None: + log.warning( + "Skipping processor bundling: the model config's 'vlm_config.tokenizer' node names no " + "Hugging Face repository. The exported checkpoint will fetch its processor at inference time." + ) + return None + try: + checkpoint_dir = resolve_processor_download(source) + snapshot_path = checkpoint_dir.download() if download is None else download(checkpoint_dir) + copied = bundle_processor_files(Path(snapshot_path), output_dir) + except Exception as e: + log.warning( + f"Skipping processor bundling: could not resolve the '{source['repository']}' processor " + f"snapshot ({e}). The export continues; the checkpoint will fetch its processor at inference time." + ) + return None + if not copied: + log.warning(f"Skipping processor bundling: no processor/tokenizer files found in '{snapshot_path}'.") + return None + log.info(f"Bundled processor files from '{checkpoint_dir.repository}': {', '.join(copied)}") + return build_artifact_source(repo=checkpoint_dir.repository, resolved_path=snapshot_path, bundled=True) + + +def hf_revision_from_snapshot_path(path: str | Path) -> str | None: + """Extract the commit hash from an HF cache snapshot path, else None.""" + match = _HF_SNAPSHOT_RE.search(Path(path).as_posix()) + return match.group(1) if match else None + + +def build_artifact_source(*, repo: str, resolved_path: str | Path | None, bundled: bool) -> dict[str, Any]: + """Build a provenance entry for 'export_manifest.json'.""" + return { + "repo": repo, + "revision": hf_revision_from_snapshot_path(resolved_path) if resolved_path is not None else None, + "bundled": bundled, + } + + +def sanitize_export_args(export_args: dict[str, Any]) -> dict[str, Any]: + """Redact local-path values from the manifest's 'export_args' entries. + + Mirrors '--student-only-checkpoint-metadata': '*_path' / '*_dir' keys and + any string containing a path separator become None; other values pass + through. Hub repo@revision provenance lives in '*_source' and is unaffected. + """ + return { + key: None + if key.endswith(("_path", "_dir")) or (isinstance(value, str) and ("/" in value or "\\" in value)) + else value + for key, value in export_args.items() + } + + +def clean_stale_export_artifacts(output_dir: Path) -> list[str]: + """Remove artifacts a previous export may have left in 'output_dir'; return removed names. + + Inference prefers bundled files local-first, so a re-export must not leave + a stale 'vision_encoder/' or processor set behind. Only the paths this tool + manages are removed; user files and the model shards are never touched. + """ + removed: list[str] = [] + vision_encoder_dir = output_dir / "vision_encoder" + if vision_encoder_dir.is_dir(): + shutil.rmtree(vision_encoder_dir) + removed.append("vision_encoder/") + for pattern in (*PROCESSOR_FILE_PATTERNS, "export_manifest.json"): + for path in sorted(output_dir.glob(pattern)): + if path.is_file(): + path.unlink() + removed.append(path.name) + return removed + + +# A NaN-latent decode yields a uniform frame (JPEG rounding keeps its pixel std +# well under 1), while any real generation lands orders of magnitude above. +_CONSTANT_IMAGE_STD_THRESHOLD = 1.0 + + +def constant_image_failure(pixels: np.ndarray) -> str | None: + """Heuristic '--verify' check for NaN/undecodable latents; None when the image passes. + + NaN latents decode into a valid, uniform (typically all-black) JPEG with + status 'success', invisible to existence/size checks; a (near-)zero pixel + std flags that class. Not a general image-quality gate. + """ + array = np.asarray(pixels, dtype=np.float32) + if array.size == 0: + return "decoded to an empty pixel array" + std = float(array.std()) + if std < _CONSTANT_IMAGE_STD_THRESHOLD: + return ( + f"is (nearly) constant (pixel std {std:.4f} < {_CONSTANT_IMAGE_STD_THRESHOLD}); " + "this is how NaN/undecodable latents decode (uniform/black frame)" + ) + return None + + +def build_export_manifest( + *, + vision_tower_source: dict[str, Any] | None, + processor_source: dict[str, Any] | None, + export_args: dict[str, Any], + framework_commit: str | None, +) -> dict[str, Any]: + """Build the 'export_manifest.json' provenance sidecar.""" + return { + "vision_tower_source": vision_tower_source, + "processor_source": processor_source, + "export_args": export_args, + "framework_commit": framework_commit, + } + + +def read_framework_commit(start: Path | None = None) -> str | None: + """Resolve the framework git commit by reading '.git' files (no subprocess). + + Returns None outside a git worktree (e.g. a pip-installed package). + """ + current = (start or Path(__file__)).resolve() + for parent in [current, *current.parents]: + git_dir = parent / ".git" + if git_dir.is_dir(): + return _read_git_head(git_dir) + return None + + +def _read_git_head(git_dir: Path) -> str | None: + try: + head = (git_dir / "HEAD").read_text().strip() + except OSError: + return None + if not head.startswith("ref:"): + return head or None + ref = head.removeprefix("ref:").strip() + ref_file = git_dir / ref + if ref_file.is_file(): + return ref_file.read_text().strip() or None + packed_refs = git_dir / "packed-refs" + if packed_refs.is_file(): + for line in packed_refs.read_text().splitlines(): + if line.endswith(f" {ref}"): + return line.split(" ", 1)[0] + return None diff --git a/cosmos_framework/scripts/action_policy_server_robolab.py b/cosmos_framework/scripts/action_policy_server_robolab.py index f76989ab..2e787e5f 100644 --- a/cosmos_framework/scripts/action_policy_server_robolab.py +++ b/cosmos_framework/scripts/action_policy_server_robolab.py @@ -25,6 +25,7 @@ init_script() +import json import socket import threading from dataclasses import dataclass @@ -77,6 +78,8 @@ _ROBOLAB_POLICY_HF_REPOSITORIES = { "Cosmos3-Nano-Policy-DROID": "nvidia/Cosmos3-Nano-Policy-DROID", "nvidia/Cosmos3-Nano-Policy-DROID": "nvidia/Cosmos3-Nano-Policy-DROID", + "Cosmos3-Edge-Policy-DROID": "nvidia/Cosmos3-Edge-Policy-DROID", + "nvidia/Cosmos3-Edge-Policy-DROID": "nvidia/Cosmos3-Edge-Policy-DROID", } ActionSpace = Literal["joint_pos", "midtrain"] @@ -352,6 +355,8 @@ class RobolabServerArgs(pydantic.BaseModel): """Whether the first action row contains the current state.""" history_length: int = 1 """State/history action rows to trim from the generated action output.""" + format_prompt_as_json: bool | None = None + """Serve prompts as structured JSON (matching training ``format_prompt_as_json``).""" class RobolabPolicyService: @@ -453,9 +458,17 @@ def _build_transform(self, training_config: Any | None, args: RobolabServerArgs) if dataset_config is None or dataset_entry is None: log.warning( - "[robolab-policy-server] no training action dataset config found; using default ActionTransformPipeline" + "[robolab-policy-server] no training action dataset config found; using default " + f"ActionTransformPipeline (format_prompt_as_json={bool(args.format_prompt_as_json)})" + ) + return ( + ActionTransformPipeline( + max_action_dim=max_action_dim, + cfg_dropout_rate=0.0, + format_prompt_as_json=bool(args.format_prompt_as_json), + ), + inferred, ) - return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0), inferred if action_dataset_config is not None: chunk_length = getattr(action_dataset_config, "chunk_length", None) @@ -479,6 +492,9 @@ def _build_transform(self, training_config: Any | None, args: RobolabServerArgs) dataset_entry.dataset = {"_target_": f"{__name__}._DummyDataset"} + if args.format_prompt_as_json is not None: + dataset_config.format_prompt_as_json = bool(args.format_prompt_as_json) + wrapped_dataset = instantiate(dataset_config) return wrapped_dataset.transform, inferred @@ -552,7 +568,10 @@ def _build_sample(self, obs: dict[str, Any]) -> dict[str, Any]: } if history_action is not None: sample["history_action"] = history_action - return self._transform(sample, self.cfg.resolution) + sample = self._transform(sample, self.cfg.resolution) + if isinstance(sample.get("ai_caption"), dict): + sample["ai_caption"] = json.dumps(sample["ai_caption"]) + return sample def infer(self, obs: dict[str, Any]) -> dict[str, Any]: sample = self._build_sample(obs) diff --git a/cosmos_framework/scripts/convert_model_to_dcp.py b/cosmos_framework/scripts/convert_model_to_dcp.py index 2c783972..1a71f1c1 100644 --- a/cosmos_framework/scripts/convert_model_to_dcp.py +++ b/cosmos_framework/scripts/convert_model_to_dcp.py @@ -75,7 +75,11 @@ def convert_model_to_dcp(args: Args): print("Saving model...") storage_writer = FileSystemWriter(args.output_path / "model", thread_count=thread_count) dcp.save(state_dict=state_dict, storage_writer=storage_writer, planner=CustomSavePlanner()) - shutil.copy(hf_path / "checkpoint.json", args.output_path / "checkpoint.json") + # ``checkpoint.json`` only exists for DCP-format source repos (e.g. Cosmos3-Nano); + # safetensors/diffusers-layout repos (e.g. Cosmos3-Edge) don't ship it. Copy when present. + source_checkpoint_json = hf_path / "checkpoint.json" + if source_checkpoint_json.exists(): + shutil.copy(source_checkpoint_json, args.output_path / "checkpoint.json") hf_config.save_pretrained(args.output_path / "model") print(f"Saved checkpoint to {args.output_path}") diff --git a/cosmos_framework/scripts/convert_model_to_diffusers.py b/cosmos_framework/scripts/convert_model_to_diffusers.py index 63531771..ae9e47b3 100644 --- a/cosmos_framework/scripts/convert_model_to_diffusers.py +++ b/cosmos_framework/scripts/convert_model_to_diffusers.py @@ -16,7 +16,7 @@ import shutil import struct from pathlib import Path -from typing import Annotated, Any +from typing import Annotated, Any, Literal import pydantic import tyro @@ -45,6 +45,46 @@ class Args(pydantic.BaseModel): skip_vision_encoder: bool = False """Do not save the vision encoder sidecar in the Diffusers checkpoint.""" + distilled_scheduler: Literal["auto", "on", "off"] = "auto" + """Scheduler export mode for distilled (few-step) checkpoints. + + * auto: export the distilled FlowMatchEuler scheduler when the checkpoint + defines `fixed_step_sampler_config`, else UniPC. + * on: require and always export the distilled scheduler. + * off: always export the UniPC scheduler. + """ + + repo_id: str | None = None + """HF repo id embedded in modular_model_index.json component specs. + + Defaults to the output directory name (matches diffusers' own default). Set + this to the target Hub id (e.g. 'nvidia/Cosmos3-Super-Image2Video-4Step') + when preparing a release so the modular pipeline resolves its components. + """ + + edge_include_reasoner: bool = True + """Keep the pinned Edge metadata and vision sidecar in the final repository.""" + + edge_reasoner_repo_id: str = "nvidia/Cosmos3-Edge" + """Hugging Face repository containing the Cosmos3 Edge reasoner checkpoint.""" + + edge_reasoner_revision: str = "be935d6931e4e176d7353abad41ca529d7b33b12" + """Pinned revision of the Cosmos3 Edge reasoner checkpoint.""" + + edge_reasoner_path: Path | None = None + """Optional local Cosmos3 Edge reasoner snapshot, used instead of downloading it.""" + + edge_action_chunk_size: pydantic.PositiveInt | None = None + """Optional action chunk size override for a Cosmos3 Edge policy manifest.""" + + +class EdgePolicyMetadata(pydantic.BaseModel): + """Policy metadata resolved from the Stage 1 Edge export.""" + + action_chunk_size: pydantic.PositiveInt + conditioning_fps: pydantic.PositiveFloat + domain_name: str = pydantic.Field(min_length=1) + class SafetensorsIndexMetadata(pydantic.BaseModel): total_size: int = 0 @@ -54,7 +94,7 @@ class SafetensorsIndex(pydantic.BaseModel): metadata: SafetensorsIndexMetadata = pydantic.Field(default_factory=SafetensorsIndexMetadata) weight_map: dict[str, str] = pydantic.Field(default_factory=dict) - def update(self, safetensors_path: Path, rel_path: str): + def update(self, safetensors_path: Path, rel_path: str, exclude_key_prefixes: tuple[str, ...] = ()) -> None: with safetensors_path.open("rb") as f: header_size = struct.unpack(" None: + """Write the root index over Diffusers component shards.""" + index = SafetensorsIndex() + index.update_dir(output_path / "transformer", "transformer") + vision_encoder_path = output_path / "vision_encoder/model.safetensors" + if vision_encoder_path.is_file(): + index.update(vision_encoder_path, "vision_encoder/model.safetensors") + (output_path / "model.safetensors.index.json").write_text(index.model_dump_json(indent=2)) + + def _build_public_export_model_config(model_dict: dict[str, Any]) -> dict[str, Any]: """Remove unsupported internal-only settings before building the public config.""" public_model_dict = copy.deepcopy(model_dict) @@ -91,13 +143,195 @@ def _build_public_export_model_config(model_dict: dict[str, Any]) -> dict[str, A return build_public_model_config(public_model_dict) -def convert_model_to_diffusers(args: Args) -> None: - args.output_path.mkdir(parents=True, exist_ok=True) +def _is_edge_model_config(model_dict: dict[str, Any]) -> bool: + """Return whether a model config uses the Nemotron Cosmos3 Edge backbone.""" + config = model_dict.get("config", {}) + vlm_config = config.get("vlm_config", {}) + pretrained_weights = vlm_config.get("pretrained_weights", {}) + model_instance = vlm_config.get("model_instance", {}) + return bool( + pretrained_weights.get("checkpoint_format") == "nemotron_3_dense_vl" + or "Nemotron3" in str(model_instance.get("_target_", "")) + ) + + +def _load_edge_policy_metadata(checkpoint_path: Path) -> EdgePolicyMetadata: + """Load policy metadata emitted by the config-driven Stage 1 export.""" + metadata_path = checkpoint_path / "checkpoint.json" + if not metadata_path.is_file(): + raise FileNotFoundError( + "Cosmos3 Edge conversion requires the Stage 1 export checkpoint.json; " + f"no metadata file was found at {metadata_path}." + ) + checkpoint_metadata = deserialize_config_dict(metadata_path) + policy_metadata = checkpoint_metadata.get("policy") + if not isinstance(policy_metadata, dict): + raise ValueError( + "Cosmos3 Edge checkpoint.json is missing `policy` metadata. Re-export the DCP checkpoint with " + "export_model.py so action_chunk_size, conditioning_fps, and domain_name come from the experiment config." + ) + try: + return EdgePolicyMetadata.model_validate(policy_metadata) + except pydantic.ValidationError as exc: + raise ValueError("Invalid Cosmos3 Edge policy metadata in checkpoint.json.") from exc + + +# Modular pipeline / blocks classes keyed by whether the export is distilled. +# The distilled variant samples on a fixed schedule with guidance baked into the +# weights; the base variant uses the standard Cosmos3 omni blocks. +MODULAR_PIPELINE_CLASSES = { + False: ("Cosmos3OmniModularPipeline", "Cosmos3OmniBlocks"), + True: ("Cosmos3DistilledModularPipeline", "Cosmos3DistilledBlocks"), +} + + +def _write_modular_model_index(output_path: Path, repo_id: str) -> None: + """Write modular_model_index.json next to the task-based model_index.json. + + Component classes (tokenizer, scheduler, vae, transformer) are derived from + the already-written model_index.json rather than hardcoded, so the modular + index always matches whatever the pipeline save produced (e.g. UniPC vs. + FlowMatchEuler scheduler, slow vs. fast tokenizer). The distilled vs. base + pipeline/blocks classes are selected from the exported scheduler config. + + Note: the distilled modular pipeline classes are not yet in a released + diffusers version (they ship via the modular Cosmos3 PR), so `_diffusers_version` + is copied from the pipeline save rather than asserted, and the classes are only + referenced by name here — loading the result requires a diffusers build that + provides them. + """ + model_index = json.loads((output_path / "model_index.json").read_text()) + + # Distilled checkpoints register `fixed_step_sampler_config` on the scheduler. + scheduler_config_path = output_path / "scheduler" / "scheduler_config.json" + is_distilled = False + distilled_sigmas = None + if scheduler_config_path.is_file(): + scheduler_config = json.loads(scheduler_config_path.read_text()) + fixed_step_cfg = scheduler_config.get("fixed_step_sampler_config") + is_distilled = fixed_step_cfg is not None + if fixed_step_cfg and fixed_step_cfg.get("t_list"): + distilled_sigmas = [float(s) for s in fixed_step_cfg["t_list"]] + + pipeline_class, blocks_class = MODULAR_PIPELINE_CLASSES[is_distilled] + + modular_index: dict[str, Any] = { + "_blocks_class_name": blocks_class, + "_class_name": pipeline_class, + "_diffusers_version": model_index.get("_diffusers_version"), + "is_distilled": is_distilled, + } + if is_distilled: + # diffusers distilled modular pipeline reads its fixed sampling schedule from this pipeline config rather + # than off the scheduler; `scheduler_config.json` keeps `fixed_step_sampler_config` for non-modular + # consumers (e.g. vllm-omni). + modular_index["distilled_sigmas"] = distilled_sigmas + + for name, value in model_index.items(): + # Keep only saved components: [library, class] pairs with non-null entries. + # Skips meta keys (_class_name, ...) and unset components ([null, null]). + if name.startswith("_"): + continue + if not (isinstance(value, list) and len(value) == 2): + continue + library, class_name = value + if library is None or class_name is None: + continue + modular_index[name] = [ + library, + class_name, + { + "pretrained_model_name_or_path": repo_id, + "subfolder": name, + "type_hint": [library, class_name], + "variant": None, + }, + ] + + (output_path / "modular_model_index.json").write_text(json.dumps(modular_index, indent=2) + "\n") + log.info(f"Wrote modular_model_index.json ({pipeline_class}, is_distilled={is_distilled}).") + + +def convert_model_to_diffusers(args: Args) -> None: register_checkpoints() checkpoint_config = args.checkpoint.build_checkpoint(checkpoints=OmniSetupOverrides.CHECKPOINTS) checkpoint_path = checkpoint_config.download_checkpoint() + model_dict = checkpoint_config.load_model_config_dict() + is_edge_model = _is_edge_model_config(model_dict) + + from cosmos_framework.scripts import _convert_model_to_diffusers + + if is_edge_model and (Path(checkpoint_path) / ".metadata").is_file(): + raise ValueError( + "Raw Cosmos3 Edge DCP cannot be converted directly. Run export_model.py first with the experiment " + "config, then pass the resulting HF directory to convert_model_to_diffusers.py." + ) + + if is_edge_model: + if args.config_only: + raise ValueError("Cosmos3 Edge conversion does not support --config-only.") + if args.skip_vision_encoder: + raise ValueError("Cosmos3 Edge conversion cannot skip the reasoner vision assets.") + if not args.edge_include_reasoner: + raise ValueError("Cosmos3 Edge conversion requires the pinned reasoner metadata and vision assets.") + # Fail fast with an actionable message if the installed Diffusers build lacks the + # Edge transformer API (the exported public config isn't detected as Edge by the + # low-level checkpoint sniff, so gate here where is_edge_model is already known). + _convert_model_to_diffusers._validate_edge_transformer_support() + # Action-policy Edge checkpoints carry a `policy` block (action_chunk_size / + # conditioning_fps / domain_name) that the diffusers pipeline needs. Non-action + # Edge (e.g. a video SFT) has no such metadata — convert it too, just without the + # policy block. The core reasoner+vision conversion below is not action-specific. + is_action_policy = bool(model_dict["config"]["action_gen"]) + edge_policy_metadata = _load_edge_policy_metadata(Path(checkpoint_path)) if is_action_policy else None + + args.output_path.mkdir(parents=True, exist_ok=True) + _convert_model_to_diffusers.convert_model_to_diffusers( + _convert_model_to_diffusers.Args( + checkpoint_path=Path(checkpoint_path), + output=str(args.output_path), + save_pipeline=True, + dtype="bf16", + include_reasoner=True, + reasoner_repo_id=args.edge_reasoner_repo_id, + reasoner_revision=args.edge_reasoner_revision, + reasoner_path=args.edge_reasoner_path, + ) + ) + # The published Edge repository follows the existing Diffusers layout: + # its root index points at component shards, rather than at an + # intermediate native Transformers `model.safetensors` file. + _convert_model_to_diffusers._write_diffusers_safetensors_index(args.output_path) + (args.output_path / "model.safetensors").unlink(missing_ok=True) + _write_modular_model_index(args.output_path, repo_id=args.repo_id or args.output_path.name) + + checkpoint_payload: dict[str, Any] = { + "config_file": None, + "experiment": None, + "experiment_overrides": None, + } + if edge_policy_metadata is not None: + checkpoint_payload["policy"] = { + "action_chunk_size": ( + args.edge_action_chunk_size + if args.edge_action_chunk_size is not None + else edge_policy_metadata.action_chunk_size + ), + "conditioning_fps": edge_policy_metadata.conditioning_fps, + "domain_name": edge_policy_metadata.domain_name, + } + elif args.edge_action_chunk_size is not None: + raise ValueError( + "--edge-action-chunk-size was set, but this Edge checkpoint is not an action policy " + "(config action_gen is false), so there is no policy manifest to override." + ) + serialize_config_dict(checkpoint_payload, args.output_path / "checkpoint.json") + print(f"Saved diffusers checkpoint to {args.output_path}") + return + model_dict = load_model_config_from_hf_config(deserialize_config_dict(checkpoint_path / "config.json")) + args.output_path.mkdir(parents=True, exist_ok=True) supports_action = model_dict["config"]["action_gen"] supports_sound = model_dict["config"]["sound_gen"] @@ -140,14 +374,7 @@ def convert_model_to_diffusers(args: Args) -> None: ) if not args.config_only: - from cosmos_framework.scripts._convert_model_to_diffusers import ( - Args as _Args, - ) - from cosmos_framework.scripts._convert_model_to_diffusers import ( - convert_model_to_diffusers as _convert_model_to_diffusers, - ) - - _args = _Args( + _args = _convert_model_to_diffusers.Args( checkpoint_path=str(checkpoint_path), output=str(args.output_path), save_pipeline=True, @@ -157,8 +384,9 @@ def convert_model_to_diffusers(args: Args) -> None: include_sound_tokenizer=supports_sound, vision_encoder_model=vision_encoder_model, skip_vision_encoder=args.skip_vision_encoder, + distilled_scheduler=args.distilled_scheduler, ) - _convert_model_to_diffusers(_args) + _convert_model_to_diffusers.convert_model_to_diffusers(_args) # Add vlm files vlm_repository = model_dict["config"]["vlm_config"]["model_name"] @@ -187,12 +415,11 @@ def convert_model_to_diffusers(args: Args) -> None: if not args.config_only: # Add top-level index - index = SafetensorsIndex() - index.update_dir(args.output_path / "transformer", "transformer") - vision_encoder_rel = "vision_encoder/model.safetensors" - if (args.output_path / vision_encoder_rel).is_file(): - index.update(args.output_path / vision_encoder_rel, vision_encoder_rel) - (args.output_path / "model.safetensors.index.json").write_text(index.model_dump_json(indent=2)) + _write_diffusers_weight_index(args.output_path) + + # Modular pipeline index (for diffusers ModularPipeline.from_pretrained). + # Written alongside the task-based model_index.json produced by the save. + _write_modular_model_index(args.output_path, repo_id=args.repo_id or args.output_path.name) checkpoint_metadata_path = checkpoint_path / "checkpoint.json" if not checkpoint_metadata_path.is_file(): diff --git a/cosmos_framework/scripts/export_model.py b/cosmos_framework/scripts/export_model.py index d1479784..2d0f8a31 100644 --- a/cosmos_framework/scripts/export_model.py +++ b/cosmos_framework/scripts/export_model.py @@ -12,7 +12,14 @@ } ) +import inspect import json +import os +import shlex +import shutil +import subprocess +import sys +import tempfile from pathlib import Path from typing import Annotated, Any, Callable @@ -34,37 +41,142 @@ ) from cosmos_framework.inference.common.checkpoints import register_checkpoints from cosmos_framework.inference.common.config import serialize_config_dict +from cosmos_framework.inference.common.distillation_export import ( + build_student_checkpoint_metadata, + resolve_vision_checkpoint_path, + sanitize_student_model_config, + sanitize_student_public_model_config, +) from cosmos_framework.inference.common.init import is_rank0 from cosmos_framework.inference.common.public_model_config import build_public_model_config from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel from cosmos_framework.model.generator.omni_mot_model import OmniMoTModel +from cosmos_framework.scripts._export_model_helpers import ( + EDGE_VIT_BUNDLE_HF_INCLUDE, + build_artifact_source, + build_export_manifest, + bundle_processor_files, + bundle_processor_from_tokenizer_node, + bundle_vision_encoder, + clean_stale_export_artifacts, + constant_image_failure, + is_edge_model, + read_framework_commit, + reasoner_vision_capable, + resolve_vision_bundle_dirs, + sanitize_export_args, + set_include_visual, +) from cosmos_framework.utils import log -from cosmos_framework.utils.checkpoint_db import CheckpointConfig, sanitize_uri +from cosmos_framework.utils.checkpoint_db import CheckpointConfig, CheckpointDirHf, sanitize_uri from cosmos_framework.utils.lazy_config.registry import convert_target_to_string _INTERNAL_VISUAL_PREFIX = "model.net.language_model.visual." _EXPORTED_VISUAL_PREFIX = "model.visual." +def _config_value(config: Any, key: str) -> Any: + if isinstance(config, dict): + return config.get(key) + return getattr(config, key, None) + + +def _dataset_config_value(dataset_config: Any, key: str) -> Any: + value = _config_value(dataset_config, key) + if value is not None: + return value + + target = _config_value(dataset_config, "_target_") or _config_value(dataset_config, "_target") + if target is None: + return None + try: + parameter = inspect.signature(target).parameters.get(key) + except (TypeError, ValueError): + return None + if parameter is None or parameter.default is inspect.Parameter.empty: + return None + return parameter.default + + +def _build_edge_policy_metadata(training_config: Any) -> dict[str, Any]: + """Resolve policy manifest fields from the action experiment config.""" + try: + dataset_config = training_config.dataloader_train.dataloaders.action_data.dataloader.dataset + except AttributeError as exc: + raise ValueError( + "Cosmos3 Edge export requires an action dataset config at " + "dataloader_train.dataloaders.action_data.dataloader.dataset." + ) from exc + + dataset_entries = _config_value(dataset_config, "list_of_datasets") + if not dataset_entries: + raise ValueError("Cosmos3 Edge export requires at least one action dataset entry.") + + metadata_by_dataset: list[dict[str, Any]] = [] + for entry in dataset_entries: + action_dataset_config = _config_value(entry, "dataset") + if action_dataset_config is None: + raise ValueError("Cosmos3 Edge action dataset entries must define a dataset config.") + + target = _config_value(action_dataset_config, "_target_") or _config_value(action_dataset_config, "_target") + action_chunk_size = _dataset_config_value(action_dataset_config, "chunk_length") + conditioning_fps = _dataset_config_value(action_dataset_config, "fps") + domain_name = _config_value(action_dataset_config, "embodiment_type") + if domain_name is None and target is not None: + domain_name = getattr(target, "EMBODIMENT_TYPE", None) + + if not isinstance(action_chunk_size, int) or isinstance(action_chunk_size, bool) or action_chunk_size <= 0: + raise ValueError( + "Cosmos3 Edge action dataset config must define a positive integer `chunk_length`, " + f"got {action_chunk_size!r}." + ) + if ( + not isinstance(conditioning_fps, (int, float)) + or isinstance(conditioning_fps, bool) + or conditioning_fps <= 0 + ): + raise ValueError( + f"Cosmos3 Edge action dataset config must define a positive numeric `fps`, got {conditioning_fps!r}." + ) + if not isinstance(domain_name, str) or not domain_name.strip(): + target_name = getattr(target, "__name__", repr(target)) + raise ValueError( + "Cosmos3 Edge action dataset config must define `embodiment_type` or expose " + f"`EMBODIMENT_TYPE`; dataset target is {target_name}." + ) + + metadata_by_dataset.append( + { + "action_chunk_size": action_chunk_size, + "conditioning_fps": float(conditioning_fps), + "domain_name": domain_name.strip(), + } + ) + + metadata = metadata_by_dataset[0] + for field in metadata: + if any(dataset_metadata[field] != metadata[field] for dataset_metadata in metadata_by_dataset[1:]): + raise ValueError( + "Cosmos3 Edge checkpoint.json can represent only one action policy metadata value per field; " + f"the configured datasets disagree on `{field}`." + ) + return metadata + + def _coerce_to_base_model(model_dict: dict[str, Any]) -> None: """For distillation training configs, rewrite the target to the base OmniMoTModel so the exported checkpoint only contains the student network.""" target = model_dict.get("_target_", "") - if "OmniMoTModel" in target: - return - - log.info(f"Overriding model target from {target} to OmniMoTModel for export") - model_dict["_target_"] = convert_target_to_string(OmniMoTModel) - - config = model_dict["config"] - base_field_names = {f.name for f in attrs.fields(OmniMoTModelConfig)} - extra_keys = [k for k in config if k not in base_field_names and not k.startswith("_")] - for k in extra_keys: - del config[k] + base_model_target = convert_target_to_string(OmniMoTModel) + if target != base_model_target: + log.info(f"Overriding model target from {target} to OmniMoTModel for export") - metadata = config.get("_metadata", {}) - metadata["object_type"] = convert_target_to_string(OmniMoTModelConfig) - config["_metadata"] = metadata + sanitize_student_model_config( + model_dict, + base_model_target=base_model_target, + base_config_type=convert_target_to_string(OmniMoTModelConfig), + base_config_field_names={field.name for field in attrs.fields(OmniMoTModelConfig)}, + ) class Args(ParallelismOverrides): @@ -73,11 +185,17 @@ class Args(ParallelismOverrides): """Output model directory.""" config_only: bool = False """If True, only export config.""" + student_only_checkpoint_metadata: bool = False + """If True, omit source checkpoint and credential paths from checkpoint metadata.""" vit: bool = True """If True, export ViT weights.""" + vit_checkpoint_path: ResolvedPath | None = None + """Optional local Hugging Face checkpoint directory containing ViT weights.""" + verify: bool = False + """If True, smoke-test the exported checkpoint with single-GPU inference (tiny reasoner + generation samples).""" -def _load_safetensor_weights(model_dir: Path, predicate: Callable[[str], bool]) -> dict: +def _load_safetensor_weights(model_dir: Path, predicate: Callable[[str], bool]) -> dict[str, Any]: """Load weights from a safetensors file.""" index_path = model_dir / "model.safetensors.index.json" if index_path.exists(): @@ -104,10 +222,158 @@ def _rewrite_visual_fqns_for_vfm(state_dict: dict[str, Any]) -> dict[str, Any]: return remapped_state_dict -def export_model(args: Args): +# Env vars stripped from the --verify subprocess: init_script pinned +# COSMOS_DEVICE=cpu / COSMOS_TRAINING=1 for the export process itself, and any +# torchrun rendezvous vars would make single-process inference hang. +_VERIFY_ENV_EXCLUDE = ( + "COSMOS_DEVICE", + "COSMOS_TRAINING", + "RANK", + "LOCAL_RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "GROUP_RANK", + "MASTER_ADDR", + "MASTER_PORT", +) + + +def _verify_exported_checkpoint(output_dir: Path, *, run_reasoner_check: bool) -> None: + """Smoke-test the exported checkpoint via single-GPU 'scripts.inference'. + + Runs a minimal generation sample (plus a reasoner-image sample when + 'run_reasoner_check') and asserts each succeeded with non-empty outputs. + Skips when no CUDA device is available. + """ + import torch + + if not torch.cuda.is_available(): + log.warning("Skipping --verify: no CUDA device is available on this node.") + return + + verify_dir = Path(tempfile.mkdtemp(prefix="export_model_verify_")) + inputs_dir = verify_dir / "inputs" + outputs_dir = verify_dir / "outputs" + inputs_dir.mkdir(parents=True) + + # (sample name, output files that must exist and be non-empty) + checks: list[tuple[str, list[str]]] = [] + if run_reasoner_check: + from PIL import Image + + # Local synthetic image (mirrors inputs/reasoner/reasoner_image.json + # without a network fetch). + image_path = inputs_dir / "verify_image.jpg" + Image.new("RGB", (256, 256), color=(90, 90, 90)).save(image_path) + (inputs_dir / "verify_reasoner_image.json").write_text( + json.dumps( + { + "model_mode": "reasoner", + "prompt": "Describe this image in one short sentence.", + "vision_path": str(image_path), + "seed": 0, + } + ) + ) + checks.append(("verify_reasoner_image", ["reasoner_text.txt"])) + # Smallest generation the sample schema allows: one 256px square frame. + (inputs_dir / "verify_t2i.json").write_text( + json.dumps( + { + "model_mode": "text2image", + "prompt": "A gray robotic arm on a white table.", + "resolution": "256", + "aspect_ratio": "1,1", + "num_frames": 1, + "seed": 0, + } + ) + ) + checks.append(("verify_t2i", ["vision.jpg"])) + + cmd = [ + sys.executable, + "-m", + "cosmos_framework.scripts.inference", + "--parallelism-preset=latency", + # Guardrails pull their own HF models (nvidia/Cosmos-Guardrail1) — an + # orthogonal hub dependency that would break offline/cache-only verify. + "--no-guardrails", + "--checkpoint-path", + str(output_dir), + "-o", + str(outputs_dir), + "-i", + *(str(inputs_dir / f"{name}.json") for name, _ in checks), + ] + env = {k: v for k, v in os.environ.items() if k not in _VERIFY_ENV_EXCLUDE and not k.startswith("TORCHELASTIC_")} + log.info(f"Verifying exported checkpoint: {shlex.join(cmd)}") + result = subprocess.run(cmd, env=env) + + failures: list[str] = [] + if result.returncode != 0: + failures.append(f"inference subprocess exited with code {result.returncode}") + for name, output_files in checks: + sample_dir = outputs_dir / name + sample_outputs_file = sample_dir / "sample_outputs.json" + if sample_outputs_file.is_file(): + try: + sample_outputs = json.loads(sample_outputs_file.read_text()) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + # A truncated/corrupt file is a verify failure, not a crash. + sample_outputs = None + failures.append(f"sample '{name}' wrote a corrupt sample_outputs.json ({e})") + if sample_outputs is not None: + status = sample_outputs.get("status") if isinstance(sample_outputs, dict) else None + if status != "success": + failures.append(f"sample '{name}' finished with status '{status}'") + else: + failures.append(f"sample '{name}' wrote no sample_outputs.json") + for output_file in output_files: + output_path = sample_dir / output_file + if not output_path.is_file() or output_path.stat().st_size == 0: + failures.append(f"sample '{name}' output '{output_file}' is missing or empty") + elif output_path.suffix == ".txt" and not output_path.read_text().strip(): + failures.append(f"sample '{name}' output '{output_file}' is blank") + elif output_path.suffix == ".jpg": + # Heuristic NaN check: NaN latents decode through clamp() into a + # valid, non-empty, constant (all-black/uniform) JPEG with status + # 'success' — catch that class via (near-)zero pixel variance. + # See 'constant_image_failure'; not a general quality gate. + import numpy as np + from PIL import Image + + try: + with Image.open(output_path) as image: + pixels = np.asarray(image.convert("RGB")) + except Exception as e: + failures.append(f"sample '{name}' output '{output_file}' is not a readable image ({e})") + else: + image_failure = constant_image_failure(pixels) + if image_failure is not None: + failures.append(f"sample '{name}' output '{output_file}' {image_failure}") + if failures: + raise RuntimeError( + f"--verify failed for '{output_dir}' (the export artifacts were still written): " + + "; ".join(failures) + + f". Inspect '{verify_dir}'." + ) + log.success(f"Verified exported checkpoint '{output_dir}'") + shutil.rmtree(verify_dir, ignore_errors=True) + + +def export_model(args: Args) -> None: register_checkpoints() checkpoint_args = args.checkpoint.build_checkpoint(checkpoints={}) args.output_dir.mkdir(parents=True, exist_ok=True) + if not args.config_only and is_rank0(): + # Re-export into the same -o dir: drop artifacts a prior export may have + # left (stale vision_encoder/ bundle, stale processor files, stale + # manifest) so the dir always reflects THIS export — inference prefers + # those files local-first and would silently load a stale tower. + removed = clean_stale_export_artifacts(args.output_dir) + if removed: + log.info(f"Removed stale artifacts of a previous export from '{args.output_dir}': {', '.join(removed)}") # Load config log.info("Loading config...") @@ -116,16 +382,62 @@ def export_model(args: Args): checkpoint_args.use_ema_weights = False model_dict["config"]["ema"]["enabled"] = False - # Download VLM checkpoint - if args.vit: - vlm_checkpoint_path = model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] - vlm_checkpoint_path = sanitize_uri(vlm_checkpoint_path) - checkpoint: CheckpointConfig | None = CheckpointConfig.maybe_from_uri(vlm_checkpoint_path) - if checkpoint is None: - raise ValueError(f"Invalid checkpoint path: {vlm_checkpoint_path}") - vlm_checkpoint_path = checkpoint.hf.download() - else: - vlm_checkpoint_path = None + is_edge = is_edge_model(model_dict) + # Cosmos3 Edge action policies need the diffusers converter's `policy` block + # (action_chunk_size / conditioning_fps / domain_name) in checkpoint.json. + # Only action Edge models carry an action dataloader; non-action Edge exports + # (e.g. Edge SFT video recipes) skip this and stay unaffected. + edge_policy_metadata = ( + _build_edge_policy_metadata(checkpoint_args.load_config()) + if is_edge and model_dict["config"].get("action_gen") + else None + ) + if not args.vit: + # Text/gen-only export: write include_visual=False into the exported model + # config so inference skips visual-tower construction instead of dying on + # missing 'model.net.language_model.visual.*' keys. + if set_include_visual(model_dict, False): + log.info("Exporting with include_visual=False (--no-vit)") + else: + log.warning("Could not set include_visual=False: model config has no 'create_vlm_config' node") + + # Download VLM checkpoint. Skipped under --config-only (no downloads are + # needed to export a config). + vlm_checkpoint_path: str | None = None + vit_repo: str | None = None + if args.vit and not args.config_only: + configured_vlm_checkpoint = model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] + resolved_repositories: list[str] = [] + + def download_vlm_checkpoint(configured_uri: str) -> str: + sanitized_uri = sanitize_uri(configured_uri) + checkpoint: CheckpointConfig | None = CheckpointConfig.maybe_from_uri(sanitized_uri) + if checkpoint is None: + raise ValueError( + f"VLM backbone checkpoint URI '{configured_uri}' is not in the checkpoint " + "registry, so the vision tower cannot be resolved automatically. Either " + "export without vision weights (--no-vit), or pass " + "--vit-checkpoint-path pointing at a local " + "Hugging Face snapshot that contains them." + ) + resolved_repositories.append(checkpoint.hf.repository) + if is_edge and isinstance(checkpoint.hf, CheckpointDirHf) and not checkpoint.hf.include: + # Narrow the Edge download to the ViT bundle (~1 GB) instead of + # the full repo (tens of GB). Done at this call site, not on the + # registry entry: that entry has no include filter because + # training's backbone seeding needs the full repo. Qwen-family + # exports keep the full snapshot (the ViT merge reads the root + # safetensors). + narrowed = checkpoint.hf.model_copy(update={"include": EDGE_VIT_BUNDLE_HF_INCLUDE}) + return narrowed.download() + return checkpoint.hf.download() + + vlm_checkpoint_path = resolve_vision_checkpoint_path( + local_path=str(args.vit_checkpoint_path) if args.vit_checkpoint_path is not None else None, + configured_uri=configured_vlm_checkpoint, + download_checkpoint=download_vlm_checkpoint, + ) + vit_repo = resolved_repositories[0] if resolved_repositories else str(args.vit_checkpoint_path) # Load model log.info("Loading model...") @@ -136,6 +448,8 @@ def export_model(args: Args): # Save model log.info("Saving model...") + vision_tower_source: dict[str, Any] | None = None + processor_source: dict[str, Any] | None = None if not args.config_only: # Load checkpoint if checkpoint_args.checkpoint_path.startswith("s3://"): @@ -163,14 +477,57 @@ def export_model(args: Args): if not is_rank0(): return - # Load ViT from VLM checkpoint + # Attach the vision tower if args.vit: assert vlm_checkpoint_path is not None - vit_state_dict = _load_safetensor_weights( - Path(vlm_checkpoint_path), lambda x: x.startswith("model.visual.") + assert vit_repo is not None + if is_edge: + # Edge bundle mode: the SigLIP2 tower is a self-contained + # 'vision_encoder/' artifact that inference loads lazily (see + # Nemotron3DenseVLTextForCausalLM._ensure_vision_tower); copy it + # verbatim (incl. 'model.projector.*') instead of merging into the + # root safetensors. + if any(key.startswith(_INTERNAL_VISUAL_PREFIX) for key in state_dict): + # DCP-first: never clobber a trained in-DCP tower. + log.warning( + "DCP checkpoint unexpectedly contains a vision tower " + f"('{_INTERNAL_VISUAL_PREFIX}*'); keeping the in-DCP weights and " + "skipping the vision_encoder/ bundle copy." + ) + else: + vision_encoder_dir, snapshot_root = resolve_vision_bundle_dirs(Path(vlm_checkpoint_path)) + bundle_vision_encoder(vision_encoder_dir, snapshot_root, args.output_dir) + vision_tower_source = build_artifact_source( + repo=vit_repo, resolved_path=vlm_checkpoint_path, bundled=True + ) + log.info(f"Bundled vision_encoder/ from '{vlm_checkpoint_path}'") + if snapshot_root is not None: + copied = bundle_processor_files(snapshot_root, args.output_dir) + if copied: + processor_source = build_artifact_source( + repo=vit_repo, resolved_path=vlm_checkpoint_path, bundled=True + ) + log.info(f"Bundled processor files: {', '.join(copied)}") + else: + # Qwen family: merge the ViT weights into the root safetensors. + vit_state_dict = _load_safetensor_weights( + Path(vlm_checkpoint_path), lambda x: x.startswith("model.visual.") + ) + assert vit_state_dict, "No vision weights found" + state_dict.update(_rewrite_visual_fqns_for_vfm(vit_state_dict)) + vision_tower_source = build_artifact_source( + repo=vit_repo, resolved_path=vlm_checkpoint_path, bundled=False + ) + + # Bundle processor/tokenizer files for every export ('--no-vit' too: + # generation-only inference still tokenizes prompts through the VLM + # processor). Skipped when the Edge --vit path above already bundled + # them. Best-effort: failures log a warning and leave processor_source + # None; the export itself never fails on this. + if processor_source is None: + processor_source = bundle_processor_from_tokenizer_node( + (model_dict["config"].get("vlm_config") or {}).get("tokenizer"), args.output_dir ) - assert vit_state_dict, "No vision weights found" - state_dict.update(_rewrite_visual_fqns_for_vfm(vit_state_dict)) # Save checkpoint hf_model.save_pretrained( @@ -181,16 +538,66 @@ def export_model(args: Args): # Re-write 'config.json' to apply replacements. hf_config_file = args.output_dir / "config.json" hf_config_json = json.loads(hf_config_file.read_text()) + if args.student_only_checkpoint_metadata: + sanitize_student_public_model_config(hf_config_json["model"]) hf_config_json["model_type"] = "cosmos3_omni" serialize_config_dict(hf_config_json, hf_config_file) + # Write the provenance sidecar (kept separate from 'checkpoint.json', whose + # schema is checkpoint_args.model_dump). + export_args: dict[str, Any] = { + "config_only": args.config_only, + "student_only_checkpoint_metadata": args.student_only_checkpoint_metadata, + "verify": args.verify, + "vit": args.vit, + "vit_checkpoint_path": str(args.vit_checkpoint_path) if args.vit_checkpoint_path is not None else None, + } + if args.student_only_checkpoint_metadata: + # 'checkpoint.json' is scrubbed of local paths in this mode; the + # manifest must not leak them through 'export_args' either. The + # repo@revision provenance in '*_source' stays untouched. + export_args = sanitize_export_args(export_args) + manifest = build_export_manifest( + vision_tower_source=vision_tower_source, + processor_source=processor_source, + export_args=export_args, + framework_commit=read_framework_commit(), + ) + serialize_config_dict(manifest, args.output_dir / "export_manifest.json") + # Write 'checkpoint.json' last to indicate that the model is complete. - serialize_config_dict(checkpoint_args.model_dump(mode="json"), args.output_dir / "checkpoint.json") + checkpoint_metadata = ( + build_student_checkpoint_metadata(use_ema_weights=checkpoint_args.use_ema_weights) + if args.student_only_checkpoint_metadata + else checkpoint_args.model_dump(mode="json") + ) + if edge_policy_metadata is not None: + checkpoint_metadata["policy"] = edge_policy_metadata + serialize_config_dict(checkpoint_metadata, args.output_dir / "checkpoint.json") print(f"Saved model to {args.output_dir}") + if args.verify: + if args.config_only: + log.warning("--verify skipped: --config-only exports nothing to verify") + return + # The reasoner-image check needs a checkpoint whose reasoner can encode + # vision prompts: Edge always can (lazy tower); Qwen-family (Nano/Super) + # only with include_visual truthy — the shipped SFT configs leave it + # unset, and the sample-build fail-fast correctly rejects the vision + # sample there, so gating on args.vit alone would fail healthy exports. + run_reasoner_check = args.vit and reasoner_vision_capable(model_dict) + if args.vit and not run_reasoner_check: + log.info( + "--verify: skipping the reasoner-image check — the exported reasoner LM has no visual " + "tower (model config 'include_visual' is falsy); the generation check still runs." + ) + # Runs only after every export artifact is written: verify is a check, + # not a gate, and a failure exits non-zero without touching the export. + _verify_exported_checkpoint(args.output_dir, run_reasoner_check=run_reasoner_check) + -def main(): +def main() -> None: args = tyro_cli(Args, description=__doc__, config=(tyro.conf.OmitArgPrefixes,)) export_model(args) diff --git a/cosmos_framework/scripts/export_model_test.py b/cosmos_framework/scripts/export_model_test.py new file mode 100644 index 00000000..900ac3df --- /dev/null +++ b/cosmos_framework/scripts/export_model_test.py @@ -0,0 +1,581 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Tests for the export_model pure helpers (Edge detection, reasoner-vision capability, +include_visual writing, vision_encoder bundle config merging, processor bundling and its +failure cleanup, stale-artifact cleanup, manifest building/sanitization, the --verify +constant-image heuristic) and the Edge checkpoint registry entry. GPU/download paths are +exercised by the e2e validation, not here.""" + +import json +import re +import shutil + +import numpy as np +import pytest + +from cosmos_framework.inference.common.checkpoints import register_checkpoints +from cosmos_framework.scripts import _export_model_helpers as helpers +from cosmos_framework.utils.checkpoint_db import CheckpointConfig, sanitize_uri + +_EDGE_BACKBONE_URI = "s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/" + +_VISION_CONFIG = {"hidden_size": 1152, "num_hidden_layers": 27, "patch_size": 16} +_PROJECTOR_CONFIG_TOP = { + "input_hidden_size": 1152, + "merger_intermediate_size": 11520, + "out_hidden_size": 2048, + "spatial_merge_size": 2, +} +_PROJECTOR_CONFIG_STANDALONE = { + "input_hidden_size": 1152, + "merger_intermedia": 11520, + "out_hidden_size": 2048, + "spatial_merge_size": 2, +} +_TOKEN_IDS = {"image_token_id": 19, "video_token_id": 18, "vision_start_token_id": 20} + + +def _model_dict(*, model_name="nvidia/Cosmos3-Edge-Reasoner", target="X.Nemotron3DenseVLTextForCausalLM"): + return { + "_target_": "X.OmniMoTModel", + "config": { + "vlm_config": { + "model_name": model_name, + "model_instance": {"_target_": target, "config": {"_target_": "X.create_vlm_config"}}, + "pretrained_weights": {"backbone_path": _EDGE_BACKBONE_URI}, + }, + }, + } + + +def _make_snapshot(tmp_path, *, standalone_config=None, top_config=None): + """Build a fake nvidia/Cosmos3-Edge snapshot under an HF-cache-style path.""" + snapshot = tmp_path / "snapshots" / ("0" * 40) + vision_dir = snapshot / "vision_encoder" + vision_dir.mkdir(parents=True) + (vision_dir / "model.safetensors").write_bytes(b"tower-weights") + if standalone_config is not None: + (vision_dir / "config.json").write_text(json.dumps(standalone_config)) + if top_config is not None: + (snapshot / "config.json").write_text(json.dumps(top_config)) + return snapshot + + +def test_edge_reasoner_uri_resolves_in_registry(): + register_checkpoints() + checkpoint = CheckpointConfig.maybe_from_uri(sanitize_uri(_EDGE_BACKBONE_URI)) + assert checkpoint is not None + assert checkpoint.hf.repository == "nvidia/Cosmos3-Edge" + # Deliberately unfiltered: training's backbone seeding ('load_language_model' + # -> 'download_checkpoint') downloads the full repo through this entry (the + # root weight index points into transformer/*.safetensors). Export narrows + # its ViT-bundle download at the call site instead. + assert checkpoint.hf.include == () + assert "vision_encoder/*" in helpers.EDGE_VIT_BUNDLE_HF_INCLUDE + assert set(helpers.PROCESSOR_HF_INCLUDE) <= set(helpers.EDGE_VIT_BUNDLE_HF_INCLUDE) + + +class TestIsEdgeModel: + def test_by_model_name(self): + assert helpers.is_edge_model(_model_dict()) + + def test_by_target_fallback(self): + model_dict = _model_dict(model_name="something-else") + model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] = "" + assert helpers.is_edge_model(model_dict) + + def test_by_backbone_uri_fallback(self): + model_dict = _model_dict(model_name="something-else", target="X.Qwen3VLTextForCausalLM") + assert helpers.is_edge_model(model_dict) + + def test_qwen_is_not_edge(self): + model_dict = _model_dict(model_name="nvidia/Cosmos3-Nano-Reasoner", target="X.Qwen3VLTextForCausalLM") + model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] = ( + "s3://bucket0/cosmos3/pretrained/huggingface/Cosmos-Reason/Cosmos3-Nano-Reasoner-bb9c6f5/" + ) + assert not helpers.is_edge_model(model_dict) + + +class TestReasonerVisionCapable: + """Gates the --verify reasoner-image check; mirrors 'inference.args._reasoner_vision_capable'.""" + + @staticmethod + def _qwen_dict(include_visual=None): + model_dict = _model_dict(model_name="nvidia/Cosmos3-Nano-Reasoner", target="X.Qwen3VLTextForCausalLM") + model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] = "" + if include_visual is not None: + model_dict["config"]["vlm_config"]["model_instance"]["config"]["include_visual"] = include_visual + return model_dict + + def test_edge_is_always_capable(self): + # Lazy SigLIP2 tower: capable even without in-checkpoint vision weights. + assert helpers.reasoner_vision_capable(_model_dict()) + + def test_edge_capable_even_with_include_visual_false(self): + model_dict = _model_dict() + assert helpers.set_include_visual(model_dict, False) + assert helpers.reasoner_vision_capable(model_dict) + + def test_qwen_include_visual_absent_is_not_capable(self): + # The shipped Nano/Super SFT configs leave include_visual unset. + assert not helpers.reasoner_vision_capable(self._qwen_dict()) + + def test_qwen_include_visual_false_is_not_capable(self): + assert not helpers.reasoner_vision_capable(self._qwen_dict(include_visual=False)) + + def test_qwen_include_visual_true_is_capable(self): + assert helpers.reasoner_vision_capable(self._qwen_dict(include_visual=True)) + + def test_unknown_family_is_capable(self): + # Never block a family this heuristic doesn't know. + model_dict = _model_dict(model_name="acme/some-model", target="X.SomeOtherForCausalLM") + model_dict["config"]["vlm_config"]["pretrained_weights"]["backbone_path"] = "" + assert helpers.reasoner_vision_capable(model_dict) + + +class TestSetIncludeVisual: + def test_writes_into_create_vlm_config_node(self): + model_dict = _model_dict() + assert helpers.set_include_visual(model_dict, False) + assert model_dict["config"]["vlm_config"]["model_instance"]["config"]["include_visual"] is False + + def test_no_model_instance(self): + model_dict = _model_dict() + model_dict["config"]["vlm_config"]["model_instance"] = None + assert not helpers.set_include_visual(model_dict, False) + + +class TestVisionEncoderBundleConfig: + def test_merged_top_level_only(self, tmp_path): + # Like nvidia/Cosmos3-Edge@5bb63b97: spec folded into the top-level config. + snapshot = _make_snapshot( + tmp_path, + top_config={"vision_config": _VISION_CONFIG, "projector_config": _PROJECTOR_CONFIG_TOP, **_TOKEN_IDS}, + ) + vision_dir, root = helpers.resolve_vision_bundle_dirs(snapshot) + assert root == snapshot + merged = helpers.build_vision_encoder_bundle_config(vision_dir, root) + assert merged == { + "vision_config": _VISION_CONFIG, + "projector_config": _PROJECTOR_CONFIG_TOP, + **_TOKEN_IDS, + } + + def test_standalone_first_with_token_ids_from_top(self, tmp_path): + # Like nvidia/Cosmos3-Edge@28a0b8e: standalone vision_encoder/config.json + # (no token ids) wins for the specs; ids are folded in from the top level. + snapshot = _make_snapshot( + tmp_path, + standalone_config={"vision_config": _VISION_CONFIG, "projector_config": _PROJECTOR_CONFIG_STANDALONE}, + top_config={"vision_config": {"other": 1}, "projector_config": _PROJECTOR_CONFIG_TOP, **_TOKEN_IDS}, + ) + vision_dir, root = helpers.resolve_vision_bundle_dirs(snapshot) + merged = helpers.build_vision_encoder_bundle_config(vision_dir, root) + assert merged["projector_config"] == _PROJECTOR_CONFIG_STANDALONE + assert merged["vision_config"] == _VISION_CONFIG + assert {k: merged[k] for k in _TOKEN_IDS} == _TOKEN_IDS + + def test_direct_vision_encoder_dir_uses_parent_root(self, tmp_path): + # --vit-checkpoint-path pointed straight at /vision_encoder. + snapshot = _make_snapshot( + tmp_path, + top_config={"vision_config": _VISION_CONFIG, "projector_config": _PROJECTOR_CONFIG_TOP, **_TOKEN_IDS}, + ) + vision_dir, root = helpers.resolve_vision_bundle_dirs(snapshot / "vision_encoder") + assert vision_dir == snapshot / "vision_encoder" + assert root == snapshot + merged = helpers.build_vision_encoder_bundle_config(vision_dir, root) + assert {k: merged[k] for k in _TOKEN_IDS} == _TOKEN_IDS + + def test_missing_token_ids_is_actionable(self, tmp_path): + snapshot = _make_snapshot( + tmp_path, + standalone_config={"vision_config": _VISION_CONFIG, "projector_config": _PROJECTOR_CONFIG_STANDALONE}, + ) + vision_dir, root = helpers.resolve_vision_bundle_dirs(snapshot) + with pytest.raises(ValueError, match="image_token_id.*--vit-checkpoint-path"): + helpers.build_vision_encoder_bundle_config(vision_dir, root) + + def test_no_weights_found(self, tmp_path): + with pytest.raises(FileNotFoundError, match="vision_encoder/model.safetensors"): + helpers.resolve_vision_bundle_dirs(tmp_path) + + +def test_bundle_vision_encoder_writes_verbatim_weights_and_merged_config(tmp_path): + snapshot = _make_snapshot( + tmp_path / "src", + top_config={"vision_config": _VISION_CONFIG, "projector_config": _PROJECTOR_CONFIG_TOP, **_TOKEN_IDS}, + ) + output_dir = tmp_path / "out" + output_dir.mkdir() + vision_dir, root = helpers.resolve_vision_bundle_dirs(snapshot) + helpers.bundle_vision_encoder(vision_dir, root, output_dir) + assert (output_dir / "vision_encoder" / "model.safetensors").read_bytes() == b"tower-weights" + written = json.loads((output_dir / "vision_encoder" / "config.json").read_text()) + assert set(written) == { + "vision_config", + "projector_config", + "image_token_id", + "video_token_id", + "vision_start_token_id", + } + + +def test_bundle_processor_files_skips_config_json(tmp_path): + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + for name in [ + "config.json", # must never be copied over the exported config.json + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "video_preprocessor_config.json", + "model.safetensors.index.json", # weights index: not a processor file + ]: + (snapshot / name).write_text("{}") + output_dir = tmp_path / "out" + output_dir.mkdir() + copied = helpers.bundle_processor_files(snapshot, output_dir) + assert sorted(copied) == [ + "chat_template.jinja", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", + ] + assert not (output_dir / "config.json").exists() + assert not (output_dir / "model.safetensors.index.json").exists() + + +def _flaky_copyfile(monkeypatch, *, fail_on_call, exception): + """Patch shutil.copyfile to raise on the Nth call (quota/IO mid-copy failure).""" + real_copyfile = shutil.copyfile + calls = {"count": 0} + + def flaky(src, dst, **kwargs): + calls["count"] += 1 + if calls["count"] == fail_on_call: + raise exception + return real_copyfile(src, dst, **kwargs) + + monkeypatch.setattr(shutil, "copyfile", flaky) + + +def test_bundle_processor_files_removes_partial_set_on_midcopy_failure(tmp_path, monkeypatch): + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + for name in ["preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]: + (snapshot / name).write_text("{}") + output_dir = tmp_path / "out" + output_dir.mkdir() + _flaky_copyfile(monkeypatch, fail_on_call=3, exception=OSError("disk quota exceeded")) + with pytest.raises(OSError, match="disk quota exceeded"): + helpers.bundle_processor_files(snapshot, output_dir) + # A partial set would flip inference's local-first preference onto a broken + # processor: the export root must be left with zero processor files. + assert list(output_dir.iterdir()) == [] + + +class TestProcessorSourceFromTokenizerNode: + def test_repository_mode(self): + # Edge-style local-artifact node (edge_model_config.py). + node = {"_target_": "X.build_processor_lazy", "repository": "nvidia/Cosmos3-Edge", "revision": "main"} + assert helpers.processor_source_from_tokenizer_node(node) == { + "repository": "nvidia/Cosmos3-Edge", + "revision": "main", + "subdir": "", + } + + def test_repository_mode_respects_subdir(self): + node = {"repository": "nvidia/Cosmos3-Edge", "revision": "abc", "subdir": "processor"} + source = helpers.processor_source_from_tokenizer_node(node) + assert source == {"repository": "nvidia/Cosmos3-Edge", "revision": "abc", "subdir": "processor"} + + def test_tokenizer_type_mode(self): + # Qwen-family build_processor_lazy node (configs.base.defaults.reasoner). + node = { + "_target_": "X.build_processor_lazy", + "tokenizer_type": "Qwen/Qwen3-VL-8B-Instruct", + "config_variant": "gcp", + } + assert helpers.processor_source_from_tokenizer_node(node) == { + "repository": "Qwen/Qwen3-VL-8B-Instruct", + "revision": None, + "subdir": "", + } + + def test_pretrained_model_name_mode(self): + # Legacy create_qwen2_tokenizer_with_download node. + node = {"pretrained_model_name": "Qwen/Qwen3-VL-8B-Instruct", "config_variant": "gcp"} + source = helpers.processor_source_from_tokenizer_node(node) + assert source is not None + assert source["repository"] == "Qwen/Qwen3-VL-8B-Instruct" + assert source["revision"] is None + + @pytest.mark.parametrize( + "node", + [ + None, + "not-a-dict", + {}, + {"tokenizer_type": "/local/snapshot/dir"}, # local dir, not a repo id + {"tokenizer_type": "s3://bucket/some/path"}, # URI, not a repo id + {"tokenizer_type": "single-token"}, # no org/name shape + ], + ) + def test_unusable_nodes_return_none(self, node): + assert helpers.processor_source_from_tokenizer_node(node) is None + + +class TestResolveProcessorDownload: + def test_plain_repo_id_pinned_via_registry(self): + register_checkpoints() + spec = helpers.resolve_processor_download( + {"repository": "Qwen/Qwen3-VL-8B-Instruct", "revision": None, "subdir": ""} + ) + assert spec.repository == "Qwen/Qwen3-VL-8B-Instruct" + assert re.fullmatch(r"[0-9a-f]{40}", spec.revision) # registry pin, not 'main' + assert spec.include == helpers.PROCESSOR_HF_INCLUDE + + def test_unregistered_repo_falls_back_to_main(self): + register_checkpoints() + spec = helpers.resolve_processor_download({"repository": "some-org/some-repo", "revision": None, "subdir": ""}) + assert spec.repository == "some-org/some-repo" + assert spec.revision == "main" + + def test_explicit_revision_used_as_is(self): + spec = helpers.resolve_processor_download( + {"repository": "nvidia/Cosmos3-Edge", "revision": "deadbeef", "subdir": "processor"} + ) + assert spec.repository == "nvidia/Cosmos3-Edge" + assert spec.revision == "deadbeef" + assert spec.subdirectory == "processor" + + +class TestBundleProcessorFromTokenizerNode: + _PROCESSOR_FILES = [ + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "chat_template.jinja", + "merges.txt", + "vocab.json", + ] + + def _make_processor_snapshot(self, tmp_path): + snapshot = tmp_path / "snapshots" / ("b" * 40) + snapshot.mkdir(parents=True) + for name in [*self._PROCESSOR_FILES, "config.json"]: + (snapshot / name).write_text("{}") + return snapshot + + def test_bundles_from_mocked_snapshot(self, tmp_path): + register_checkpoints() + snapshot = self._make_processor_snapshot(tmp_path) + output_dir = tmp_path / "out" + output_dir.mkdir() + (output_dir / "config.json").write_text('{"model_type": "cosmos3_omni"}') # the exported config + seen = {} + + def fake_download(spec): + seen["spec"] = spec + return str(snapshot) + + node = {"tokenizer_type": "Qwen/Qwen3-VL-8B-Instruct", "config_variant": "gcp"} + source = helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=fake_download) + assert source == {"repo": "Qwen/Qwen3-VL-8B-Instruct", "revision": "b" * 40, "bundled": True} + assert seen["spec"].repository == "Qwen/Qwen3-VL-8B-Instruct" + for name in self._PROCESSOR_FILES: + assert (output_dir / name).is_file() + # The exported config.json is never overwritten by bundling. + assert json.loads((output_dir / "config.json").read_text()) == {"model_type": "cosmos3_omni"} + + def test_bundles_edge_node_without_vit_snapshot(self, tmp_path): + # The --no-vit path: no tower snapshot was resolved, so the processor + # comes from the tokenizer node's explicit repository/revision. + snapshot = self._make_processor_snapshot(tmp_path) + output_dir = tmp_path / "out" + output_dir.mkdir() + seen = {} + + def fake_download(spec): + seen["spec"] = spec + return str(snapshot) + + node = {"repository": "nvidia/Cosmos3-Edge", "revision": "main"} + source = helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=fake_download) + assert seen["spec"].repository == "nvidia/Cosmos3-Edge" + assert seen["spec"].revision == "main" # explicit node revision, no registry pin + assert source == {"repo": "nvidia/Cosmos3-Edge", "revision": "b" * 40, "bundled": True} + for name in self._PROCESSOR_FILES: + assert (output_dir / name).is_file() + + def test_download_failure_is_best_effort(self, tmp_path): + register_checkpoints() + output_dir = tmp_path / "out" + output_dir.mkdir() + + def failing_download(spec): + raise RuntimeError("offline") + + node = {"tokenizer_type": "Qwen/Qwen3-VL-8B-Instruct", "config_variant": "gcp"} + source = helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=failing_download) + assert source is None # warning + continue; never an exception + assert list(output_dir.iterdir()) == [] + + def test_unusable_node_is_best_effort(self, tmp_path): + assert helpers.bundle_processor_from_tokenizer_node(None, tmp_path) is None + assert helpers.bundle_processor_from_tokenizer_node({"tokenizer_type": "/local/dir"}, tmp_path) is None + + def test_snapshot_without_processor_files_returns_none(self, tmp_path): + snapshot = tmp_path / "snapshots" / ("c" * 40) + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text("{}") # weights-repo config only + output_dir = tmp_path / "out" + output_dir.mkdir() + node = {"repository": "some-org/some-repo", "revision": "main"} + source = helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=lambda spec: str(snapshot)) + assert source is None + assert list(output_dir.iterdir()) == [] + + def test_midcopy_failure_is_best_effort_and_leaves_no_partial_set(self, tmp_path, monkeypatch): + snapshot = self._make_processor_snapshot(tmp_path) + output_dir = tmp_path / "out" + output_dir.mkdir() + _flaky_copyfile(monkeypatch, fail_on_call=3, exception=OSError("disk quota exceeded")) + node = {"repository": "some-org/some-repo", "revision": "main"} + source = helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=lambda spec: str(snapshot)) + assert source is None # warning + continue; never an exception + assert list(output_dir.iterdir()) == [] + + def test_keyboard_interrupt_propagates(self, tmp_path, monkeypatch): + # Only 'Exception' is caught (best-effort); BaseExceptions must escape. + snapshot = self._make_processor_snapshot(tmp_path) + output_dir = tmp_path / "out" + output_dir.mkdir() + _flaky_copyfile(monkeypatch, fail_on_call=1, exception=KeyboardInterrupt()) + node = {"repository": "some-org/some-repo", "revision": "main"} + with pytest.raises(KeyboardInterrupt): + helpers.bundle_processor_from_tokenizer_node(node, output_dir, download=lambda spec: str(snapshot)) + + +class TestCleanStaleExportArtifacts: + def test_removes_only_managed_paths(self, tmp_path): + output_dir = tmp_path / "out" + output_dir.mkdir() + # Managed leftovers of a prior export. + vision_dir = output_dir / "vision_encoder" + vision_dir.mkdir() + (vision_dir / "model.safetensors").write_bytes(b"stale-tower") + (vision_dir / "config.json").write_text("{}") + stale = ["preprocessor_config.json", "tokenizer.json", "chat_template.jinja", "export_manifest.json"] + for name in stale: + (output_dir / name).write_text("{}") + # Never touched: the exported config/metadata, model shards + # (save_pretrained manages those), and user files. + kept = ["config.json", "checkpoint.json", "model.safetensors", "model.safetensors.index.json", "notes.txt"] + for name in kept: + (output_dir / name).write_text("keep") + + removed = helpers.clean_stale_export_artifacts(output_dir) + + assert sorted(removed) == sorted(["vision_encoder/", *stale]) + assert not vision_dir.exists() + for name in stale: + assert not (output_dir / name).exists() + for name in kept: + assert (output_dir / name).read_text() == "keep" + + def test_fresh_dir_is_a_noop(self, tmp_path): + output_dir = tmp_path / "out" + output_dir.mkdir() + assert helpers.clean_stale_export_artifacts(output_dir) == [] + assert list(output_dir.iterdir()) == [] + + +class TestSanitizeExportArgs: + def test_redacts_local_paths_and_keeps_flags(self): + export_args = { + "config_only": False, + "student_only_checkpoint_metadata": True, + "verify": True, + "vit": True, + "vit_checkpoint_path": "/lustre/users/someone/hf_snapshot", + } + sanitized = helpers.sanitize_export_args(export_args) + assert sanitized == { + "config_only": False, + "student_only_checkpoint_metadata": True, + "verify": True, + "vit": True, + "vit_checkpoint_path": None, # redacted, key kept (stable schema) + } + assert json.dumps(sanitized) # JSON-serializable + + def test_none_path_and_enum_strings_pass_through(self): + assert helpers.sanitize_export_args({"vit_checkpoint_path": None}) == {"vit_checkpoint_path": None} + assert helpers.sanitize_export_args({"mode": "latency"}) == {"mode": "latency"} + # Any string carrying a path separator is treated as a local path. + assert helpers.sanitize_export_args({"note": "sub/dir"}) == {"note": None} + # '_path'/'_dir'-keyed fields are redacted even when relative. + assert helpers.sanitize_export_args({"cache_dir": "relative"}) == {"cache_dir": None} + + +class TestConstantImageFailure: + """Heuristic --verify check for the NaN-latent decode class (constant frames).""" + + def test_all_black_fails(self): + failure = helpers.constant_image_failure(np.zeros((64, 64, 3), dtype=np.uint8)) + assert failure is not None and "constant" in failure + + def test_uniform_gray_fails(self): + assert helpers.constant_image_failure(np.full((64, 64, 3), 90, dtype=np.uint8)) is not None + + def test_nearly_constant_fails(self): + # Sub-threshold ripple, e.g. JPEG rounding on a flat NaN-decode frame. + rng = np.random.default_rng(0) + pixels = np.full((64, 64, 3), 12.0) + rng.normal(0.0, 0.1, size=(64, 64, 3)) + assert helpers.constant_image_failure(pixels) is not None + + def test_empty_fails(self): + assert helpers.constant_image_failure(np.zeros((0, 0, 3), dtype=np.uint8)) is not None + + def test_noisy_image_passes(self): + rng = np.random.default_rng(0) + pixels = rng.integers(0, 256, size=(64, 64, 3)).astype(np.uint8) + assert helpers.constant_image_failure(pixels) is None + + +def test_hf_revision_from_snapshot_path(): + commit = "5bb63b97f461726dc3480a1ad872e275cd479c25" + path = f"/cache/hub/models--nvidia--Cosmos3-Edge/snapshots/{commit}" + assert helpers.hf_revision_from_snapshot_path(path) == commit + assert helpers.hf_revision_from_snapshot_path(f"{path}/vision_encoder") == commit + assert helpers.hf_revision_from_snapshot_path("/some/local/snapshot") is None + + +def test_build_export_manifest_schema(): + manifest = helpers.build_export_manifest( + vision_tower_source=helpers.build_artifact_source( + repo="nvidia/Cosmos3-Edge", + resolved_path="/cache/hub/models--nvidia--Cosmos3-Edge/snapshots/" + "a" * 40, + bundled=True, + ), + processor_source=None, + export_args={"vit": True}, + framework_commit="deadbeef", + ) + assert set(manifest) == {"vision_tower_source", "processor_source", "export_args", "framework_commit"} + assert manifest["vision_tower_source"] == {"repo": "nvidia/Cosmos3-Edge", "revision": "a" * 40, "bundled": True} + assert manifest["processor_source"] is None + assert json.dumps(manifest) # JSON-serializable + + +def test_read_framework_commit_reads_git_head(): + commit = helpers.read_framework_commit() + assert commit is not None # this test runs from a git worktree + assert re.fullmatch(r"[0-9a-f]{40}", commit) diff --git a/cosmos_framework/scripts/reasoner/eval_videophy2.py b/cosmos_framework/scripts/reasoner/eval_videophy2.py index 35bfef82..899caa45 100644 --- a/cosmos_framework/scripts/reasoner/eval_videophy2.py +++ b/cosmos_framework/scripts/reasoner/eval_videophy2.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""Inference + accuracy/Pearson metrics for a VideoPhy-2-SFT'd Qwen3-VL ckpt. +"""Inference + accuracy/Pearson metrics for a VideoPhy-2-SFT'd VLM ckpt. + +The model class is dispatched on the export's ``config.json`` ``model_type``: +``"cosmos3_edge"`` loads the framework-registered +``Cosmos3EdgeForConditionalGeneration`` + native Edge processor (no remote +code); anything else keeps the original Qwen3-VL path. Two modes share one CLI: @@ -107,6 +112,39 @@ def _barrier(): pass +# --------------------------------------------------------------------------- +# Model-type dispatch +# --------------------------------------------------------------------------- + +# Files build_cosmos3_edge_processor() reads from a snapshot dir. HFExportCallback +# bundles all of them into Edge exports (tokenizer save + source-snapshot .json +# copy); absence means an older/partial export, handled via the recipe fallback. +_EDGE_PROCESSOR_FILES = ( + "tokenizer.json", + "chat_template.jinja", + "preprocessor_config.json", + "video_preprocessor_config.json", +) +# Processor source of the videophy2_sft_edge recipe (policy.backbone.model_name). +_EDGE_PROCESSOR_FALLBACK_REPO = "nvidia/Cosmos3-Edge" + + +def _read_model_type(hf_ckpt): + """Return ``model_type`` from ``/config.json``, or ``None`` when the + file is missing or unreadable (hub ids, legacy exports).""" + try: + return json.loads((Path(hf_ckpt) / "config.json").read_text()).get("model_type") + except (OSError, ValueError): + return None + + +def _edge_processor_source(hf_ckpt): + """Return ``hf_ckpt`` when the export bundles every file the Edge processor + build needs, else ``None`` (caller falls back to the recipe's snapshot).""" + root = Path(hf_ckpt) + return str(root) if all((root / name).is_file() for name in _EDGE_PROCESSOR_FILES) else None + + # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @@ -166,11 +204,44 @@ def _run_inference(args, rank, world_size, local_rank): device = f"cuda:{local_rank}" if rank == 0: print(f"[infer] loading model from {args.hf_ckpt} on {device} ...", flush=True) - model = Qwen3VLForConditionalGeneration.from_pretrained( - args.hf_ckpt, torch_dtype=torch.bfloat16, device_map=device, - ) - model.eval() - processor = AutoProcessor.from_pretrained(args.hf_ckpt) + if _read_model_type(args.hf_ckpt) == "cosmos3_edge": + from transformers import AutoModelForImageTextToText + + # Importing the package registers model_type="cosmos3_edge" with the + # transformers Auto classes (framework-native classes, no remote code). + import cosmos_framework.model.generator.reasoner.cosmos3_edge # noqa: F401 + from cosmos_framework.data.generator.processors.cosmos3_edge_processing import build_cosmos3_edge_processor + + model = AutoModelForImageTextToText.from_pretrained( + args.hf_ckpt, torch_dtype=torch.bfloat16, device_map=device, + ) + model.eval() + # AutoProcessor silently degrades to a bare tokenizer on native Edge + # snapshots; build the framework port instead — from the export's bundled + # processor files, else from the recipe's processor source (HF-cache-first, + # mirroring how training resolves it). + processor_src = _edge_processor_source(args.hf_ckpt) + if processor_src is None: + from cosmos_framework.utils.generator.reasoner.pretrained_models_downloader import ( + maybe_download_hf_model_from_s3, + ) + + if rank == 0: + print( + f"[infer] export lacks bundled processor files; " + f"falling back to {_EDGE_PROCESSOR_FALLBACK_REPO}", + flush=True, + ) + processor_src = maybe_download_hf_model_from_s3( + _EDGE_PROCESSOR_FALLBACK_REPO, credentials="", bucket="", include_model_weights=False + ) + processor = build_cosmos3_edge_processor(processor_src) + else: + model = Qwen3VLForConditionalGeneration.from_pretrained( + args.hf_ckpt, torch_dtype=torch.bfloat16, device_map=device, + ) + model.eval() + processor = AutoProcessor.from_pretrained(args.hf_ckpt) # Left-pad so newly generated tokens land at the actual sequence end. if hasattr(processor, "tokenizer") and processor.tokenizer is not None: processor.tokenizer.padding_side = "left" diff --git a/cosmos_framework/scripts/reasoner/eval_videophy2_test.py b/cosmos_framework/scripts/reasoner/eval_videophy2_test.py new file mode 100644 index 00000000..c35fce63 --- /dev/null +++ b/cosmos_framework/scripts/reasoner/eval_videophy2_test.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""Unit tests for the model-type dispatch helpers in ``eval_videophy2``. + +Synthetic-dir tests are pure filesystem checks — no model weights, no hub +traffic. The tests against the real ``videophy2_sft_edge`` HF export are +read-only and auto-skip when the export dir is unavailable; the full +model-load smoke lives in the export-pipeline validation plan +(``SPEC_edge_export_pipeline_ux.md`` §4.5). +""" + +import json +import os + +import pytest + +from cosmos_framework.scripts.reasoner.eval_videophy2 import ( + _EDGE_PROCESSOR_FILES, + _edge_processor_source, + _read_model_type, +) + +# Existing Edge HF export (model_type="cosmos3_edge", bundled processor files). +# Point COSMOS3_EDGE_EXPORT_DIR at a local ``videophy2_sft_edge`` HF export to +# run these read-only checks; unset (the default) auto-skips them. +_EDGE_EXPORT_DIR = os.environ.get("COSMOS3_EDGE_EXPORT_DIR", "") + +requires_edge_export = pytest.mark.skipif( + not os.path.isdir(_EDGE_EXPORT_DIR), + reason=f"videophy2_sft_edge HF export not available at {_EDGE_EXPORT_DIR}", +) + + +def test_read_model_type_edge(tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"model_type": "cosmos3_edge"})) + assert _read_model_type(str(tmp_path)) == "cosmos3_edge" + + +def test_read_model_type_qwen(tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen3_vl"})) + assert _read_model_type(str(tmp_path)) == "qwen3_vl" + + +def test_read_model_type_missing_or_invalid(tmp_path): + assert _read_model_type(str(tmp_path)) is None # no config.json + assert _read_model_type("Qwen/Qwen3-VL-2B-Init") is None # hub id, not a dir + (tmp_path / "config.json").write_text("not json") + assert _read_model_type(str(tmp_path)) is None + + +def test_edge_processor_source_bundled(tmp_path): + for name in _EDGE_PROCESSOR_FILES: + (tmp_path / name).write_text("{}") + assert _edge_processor_source(str(tmp_path)) == str(tmp_path) + + +def test_edge_processor_source_incomplete_falls_back(tmp_path): + # Any one missing processor file must trigger the recipe-snapshot fallback. + for name in _EDGE_PROCESSOR_FILES[:-1]: + (tmp_path / name).write_text("{}") + assert _edge_processor_source(str(tmp_path)) is None + + +@requires_edge_export +def test_real_edge_export_dispatches_to_edge(): + assert _read_model_type(_EDGE_EXPORT_DIR) == "cosmos3_edge" + assert _edge_processor_source(_EDGE_EXPORT_DIR) == _EDGE_EXPORT_DIR + + +@requires_edge_export +def test_real_edge_export_processor_builds(): + # CPU-only, weight-free: tokenizer + sub-processor configs from the export. + from cosmos_framework.data.generator.processors.cosmos3_edge_processing import build_cosmos3_edge_processor + + processor = build_cosmos3_edge_processor(_EDGE_EXPORT_DIR) + assert processor.tokenizer is not None + assert callable(getattr(processor, "batch_decode", None)) + assert processor.image_token_id is not None and processor.video_token_id is not None diff --git a/cosmos_framework/trainer/__init__.py b/cosmos_framework/trainer/__init__.py index 5a63f229..a3e80e9f 100644 --- a/cosmos_framework/trainer/__init__.py +++ b/cosmos_framework/trainer/__init__.py @@ -386,6 +386,7 @@ def training_step( self.callbacks.on_before_optimizer_step( model, optimizer, scheduler, grad_scaler, iteration=iteration ) + model.on_before_optimizer_step(optimizer, scheduler, iteration=iteration) self._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=iteration) self.callbacks.on_before_zero_grad(model, optimizer, scheduler, iteration=iteration) model.on_before_zero_grad(optimizer, scheduler, iteration=iteration) diff --git a/cosmos_framework/trainer/distillation.py b/cosmos_framework/trainer/distillation.py new file mode 100644 index 00000000..e13ff7bd --- /dev/null +++ b/cosmos_framework/trainer/distillation.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import inspect +from collections.abc import Iterator +from contextlib import contextmanager +from typing import cast + +import torch +from loguru import logger as log + +from cosmos_framework.trainer import ImaginaireTrainer +from cosmos_framework.utils import distributed +from cosmos_framework.model.generator.distillation.optimizer import ( + PhaseOptimizer, + PhaseScheduler, + iter_torch_optimizers, +) + +__all__: tuple[str, ...] = ("DistillationTrainer",) + + +@contextmanager +def _sync_grad_for_closure(model: torch.nn.Module, enabled: bool) -> Iterator[None]: + """Toggle per-closure gradient sync for DDP and FSDP2/HSDP models.""" + fsdp_modules = [ + module for module in model.modules() if callable(getattr(module, "set_requires_gradient_sync", None)) + ] + if fsdp_modules: + for module in fsdp_modules: + module.set_requires_gradient_sync(enabled, recurse=False) + try: + yield + finally: + for module in fsdp_modules: + module.set_requires_gradient_sync(True, recurse=False) + return + + with distributed.ddp_sync_grad(model, enabled): + yield + + +class DistillationTrainer(ImaginaireTrainer): + """Trainer for distillation / interactive model training. + + Inherits ImaginaireTrainer verbatim and overrides only the optimizer hooks. + Routing is resolved through model.get_optimizer_key when available, with + model.get_phase as the compatibility fallback. + """ + + @staticmethod + def _merge_output_batches( + base: dict[str, torch.Tensor], + update: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + for key, value in update.items(): + if key in base and isinstance(base[key], torch.Tensor) and isinstance(value, torch.Tensor): + base[key] = base[key] + value.detach() + else: + base[key] = value.detach() if isinstance(value, torch.Tensor) else value + return base + + def training_step( + self, + model_ddp: torch.nn.Module | distributed.DistributedDataParallel, + optimizer: PhaseOptimizer, + scheduler: PhaseScheduler, + grad_scaler: torch.amp.GradScaler, + data: dict[str, torch.Tensor], + iteration: int = 0, + grad_accum_iter: int = 0, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, int]: + model = model_ddp.module if self.config.trainer.distributed_parallelism == "ddp" else model_ddp + closure_fn = getattr(model, "training_step_closures", None) + if not inspect.ismethod(closure_fn): + return super().training_step( + model_ddp, + optimizer, + scheduler, + grad_scaler, + data, + iteration=iteration, + grad_accum_iter=grad_accum_iter, + ) + + self.callbacks.on_before_forward(iteration=iteration) + with self.training_timer("forward"): + with self.straggler_detector.profile_section( + "fwd", self.config.trainer.straggler_detection.analyze_forward + ): + closures = list(closure_fn(data, iteration)) + if len(closures) == 0: + raise RuntimeError(f"{type(model).__name__}.training_step_closures yielded no closures.") + + output_batch: dict[str, torch.Tensor] = {} + total_loss: torch.Tensor | None = None + should_sync_grad = grad_accum_iter == self.config.trainer.grad_accum_iter - 1 + for _closure_name, closure, is_last_closure in closures: + with _sync_grad_for_closure(model_ddp, should_sync_grad and is_last_closure): + with self.training_timer("forward"): + with self.straggler_detector.profile_section( + "fwd", self.config.trainer.straggler_detection.analyze_forward + ): + closure_output, closure_loss = closure() + self.callbacks.on_after_forward(iteration=iteration) + self._merge_output_batches(output_batch, closure_output) + closure_loss_detached = closure_loss.detach() # [] + total_loss = closure_loss_detached if total_loss is None else total_loss + closure_loss_detached # [] + self.callbacks.on_before_backward(model, closure_loss, iteration=iteration) + with self.training_timer("backward"): + with self.straggler_detector.profile_section( + "bwd", self.config.trainer.straggler_detection.analyze_backward + ): + loss_scaled = grad_scaler.scale(closure_loss / self.config.trainer.grad_accum_iter) + loss_scaled.backward() + model.on_after_backward() + self.callbacks.on_after_backward(model, iteration=iteration) + + if total_loss is None: + raise RuntimeError("No closure loss was produced.") + grad_accum_iter += 1 + if grad_accum_iter == self.config.trainer.grad_accum_iter: + with self.training_timer("optimizer_step"): + with self.straggler_detector.profile_section( + "opt", self.config.trainer.straggler_detection.analyze_optimizer + ): + self.callbacks.on_before_optimizer_step( + model, optimizer, scheduler, grad_scaler, iteration=iteration + ) + self._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=iteration) + self.callbacks.on_before_zero_grad(model, optimizer, scheduler, iteration=iteration) + model.on_before_zero_grad(optimizer, scheduler, iteration=iteration) + self._zero_grad(model, optimizer, iteration) + grad_accum_iter = 0 + return output_batch, total_loss, grad_accum_iter + + def _optimizer_step( + self, + model: torch.nn.Module, + optimizer: PhaseOptimizer, + scheduler: PhaseScheduler, + grad_scaler: torch.amp.GradScaler, + iteration: int, + ) -> None: + if not getattr(self, "_eager_init_done", False): + for opt_key, opt in optimizer.items(): + self._eager_init_optimizer_state(opt, opt_key) + self._eager_init_done = True + key = self._optimizer_key(model, iteration) + optimizer.step(key, grad_scaler) + scheduler.step(key) + + def _zero_grad(self, model: torch.nn.Module, optimizer: PhaseOptimizer, iteration: int) -> None: + key = self._optimizer_key(model, iteration) + optimizer.zero_grad(key) + model.zero_grad_for_phase(iteration) + + @staticmethod + def _optimizer_key(model: torch.nn.Module, iteration: int) -> str: + get_optimizer_key = getattr(model, "get_optimizer_key", None) + if callable(get_optimizer_key): + key = get_optimizer_key(iteration) + if isinstance(key, str): + return key + + phase = model.get_phase(iteration) + return "net" if phase == "student" else "fake_score" + + @staticmethod + def _eager_init_optimizer_state(optimizer: object, key: str) -> None: + """Pre-allocate Adam-family optimizer state without changing parameters. + + Adam-family optimizers allocate per-parameter exp_avg/exp_avg_sq lazily on + the first step() call. With `warmup_critic_steps>0` the student/generator + does not step until after several checkpoints have been saved, so the + async checkpointer's pinned-memory companion state_dict is built without + student optimizer state. Once the student starts stepping, the source + state_dict structure changes and the next save raises CompanionMismatch + in `_copy_state_dict`. Pre-allocating both opts up front keeps the + state_dict structure stable from iter 0. + """ + inner_optimizers = list(iter_torch_optimizers(optimizer)) + if len(inner_optimizers) != 1 or inner_optimizers[0] is not optimizer: + for inner_idx, inner_optimizer in enumerate(inner_optimizers): + DistillationTrainer._eager_init_optimizer_state(inner_optimizer, f"{key}:{inner_idx}") + return + + torch_optimizer = cast(torch.optim.Optimizer, optimizer) + log.info(f"[eager-init] allocating optimizer state for key='{key}'") + uses_group_step = DistillationTrainer._uses_group_step(torch_optimizer) + DistillationTrainer._eager_init_master_weights(torch_optimizer) + for group in torch_optimizer.param_groups: + DistillationTrainer._eager_init_group_state(torch_optimizer, group) + for param in group["params"]: + if not param.requires_grad: + continue + state = torch_optimizer.state[param] + if "exp_avg" not in state: + state["exp_avg"] = DistillationTrainer._optimizer_zero_like(param, uses_group_step) + if "exp_avg_sq" not in state: + state["exp_avg_sq"] = DistillationTrainer._optimizer_zero_like(param, uses_group_step) + if group.get("amsgrad", False) and "max_exp_avg_sq" not in state: + state["max_exp_avg_sq"] = DistillationTrainer._optimizer_zero_like(param, uses_group_step) + if not uses_group_step and "step" not in state: + step_device = param.device if group.get("capturable", False) or group.get("fused", False) else "cpu" + state["step"] = torch.zeros((), dtype=torch.float32, device=step_device) # [] + + @staticmethod + def _uses_group_step(optimizer: torch.optim.Optimizer) -> bool: + """Return whether the optimizer keeps the step counter on each param group.""" + return hasattr(optimizer, "param_groups_master") and hasattr(optimizer, "capturable") + + @staticmethod + def _optimizer_zero_like(param: torch.nn.Parameter, force_fp32: bool) -> torch.Tensor: + """Create an optimizer moment tensor matching the target optimizer convention.""" + value = torch.zeros_like(param) # [*param.shape] + return value.float() if force_fp32 else value # [*param.shape] + + @staticmethod + def _eager_init_master_weights(optimizer: torch.optim.Optimizer) -> None: + """Create FusedAdam-style FP32 master weights when the optimizer owns them.""" + if not hasattr(optimizer, "param_groups_master") or optimizer.param_groups_master is not None: + return + master_weights = bool(getattr(optimizer, "master_weights", False)) + optimizer.param_groups_master = [] + for group in optimizer.param_groups: + optimizer.param_groups_master.append( + { + "params": [ + param.clone().detach().float() if master_weights else None # [*param.shape] + for param in group["params"] + ] + } + ) + + @staticmethod + def _eager_init_group_state(optimizer: torch.optim.Optimizer, group: dict) -> None: + """Create FusedAdam-style group state without advancing the optimizer.""" + if not DistillationTrainer._uses_group_step(optimizer) or len(group["params"]) == 0: + return + device = group["params"][0].device + if getattr(optimizer, "capturable", False): + if "step" not in group: + group["step"] = torch.zeros(1, dtype=torch.int, device=device) # [1] + elif isinstance(group["step"], torch.Tensor): + group["step"] = group["step"].to(device=device) # [1] + else: + group["step"] = torch.tensor([group["step"]], dtype=torch.int, device=device) # [1] + if not isinstance(group["lr"], torch.Tensor): + group["lr"] = torch.tensor(group["lr"], dtype=torch.float32, device=device) # [] + else: + group["lr"] = group["lr"].to(device=device) # [] + elif "step" not in group: + group["step"] = 0 diff --git a/cosmos_framework/trainer/distillation_test.py b/cosmos_framework/trainer/distillation_test.py new file mode 100644 index 00000000..a652fdab --- /dev/null +++ b/cosmos_framework/trainer/distillation_test.py @@ -0,0 +1,613 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from collections import defaultdict +from collections.abc import Callable, Iterator +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import cosmos_framework.trainer.distillation as trainer_module +from cosmos_framework.trainer import ImaginaireTrainer +from cosmos_framework.utils.generator.optimizer import OptimizersContainer +from cosmos_framework.model.generator.distillation.optimizer import PhaseOptimizer, PhaseScheduler +from cosmos_framework.trainer.distillation import DistillationTrainer + + +def _make_config(grad_accum_iter: int = 1, distributed_parallelism: str = "ddp") -> SimpleNamespace: + straggler = SimpleNamespace( + analyze_forward=False, + analyze_backward=False, + analyze_optimizer=False, + analyze_dataloading=False, + enabled=False, + report_freq=10, + profile_freq=10, + max_diff=0.2, + raise_error=False, + save_s3=False, + ) + trainer = SimpleNamespace( + grad_accum_iter=grad_accum_iter, + distributed_parallelism=distributed_parallelism, + straggler_detection=straggler, + ) + return SimpleNamespace(trainer=trainer) + + +def _make_trainer(grad_accum_iter: int = 1, distributed_parallelism: str = "ddp") -> DistillationTrainer: + trainer = object.__new__(DistillationTrainer) + trainer.config = _make_config(grad_accum_iter=grad_accum_iter, distributed_parallelism=distributed_parallelism) + trainer.callbacks = MagicMock() + trainer.training_timer = MagicMock() + trainer.training_timer.return_value = MagicMock( + __enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False) + ) + trainer.straggler_detector = MagicMock() + trainer.straggler_detector.profile_section.return_value = MagicMock( + __enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False) + ) + return trainer + + +def _make_optimizer() -> PhaseOptimizer: + opt_net = MagicMock() + return PhaseOptimizer({"net": opt_net, "fake_score": MagicMock()}) + + +def _make_scheduler() -> PhaseScheduler: + return PhaseScheduler({"net": MagicMock(), "fake_score": MagicMock()}) + + +def _make_model(student_phase: bool = True) -> MagicMock: + model = MagicMock() + model.get_phase.return_value = "student" if student_phase else "critic" + return model + + +def _make_grad_scaler() -> MagicMock: + scaler = MagicMock(spec=torch.amp.GradScaler) + scaler.scale.side_effect = lambda x: x + return scaler + + +def _make_optimizer_container(*optimizers: torch.optim.Optimizer) -> OptimizersContainer: + container = object.__new__(OptimizersContainer) + container.optimizers = list(optimizers) + return container + + +class _FakeMasterWeightOptimizer: + """Minimal FusedAdam-like object for testing eager state allocation.""" + + def __init__( + self, + params: list[torch.nn.Parameter], + lr: float = 1e-3, + capturable: bool = True, + master_weights: bool = True, + ) -> None: + self.param_groups: list[dict[str, object]] = [ + { + "params": params, + "lr": lr, + "betas": (0.9, 0.999), + "eps": 1e-8, + "weight_decay": 0.01, + } + ] + self.state: defaultdict[torch.nn.Parameter, dict[str, torch.Tensor]] = defaultdict(dict) + self.capturable: bool = capturable + self.master_weights: bool = master_weights + self.param_groups_master: list[dict[str, list[torch.Tensor | None]]] | None = None + self.step_called: bool = False + + def step(self) -> None: + self.step_called = True + raise AssertionError("eager init must not call optimizer.step()") + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_public_trainer_exports_are_explicit() -> None: + assert trainer_module.__all__ == ("DistillationTrainer",) + + +# --------------------------------------------------------------------------- +# Class-level checks +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_is_subclass_of_imaginaire_trainer() -> None: + assert issubclass(DistillationTrainer, ImaginaireTrainer) + + +@pytest.mark.L0 +def test_does_not_define_train() -> None: + assert "train" not in DistillationTrainer.__dict__ + + +@pytest.mark.L0 +def test_defines_training_step_for_closures() -> None: + assert "training_step" in DistillationTrainer.__dict__ + + +@pytest.mark.L0 +def test_does_not_define_validate() -> None: + assert "validate" not in DistillationTrainer.__dict__ + + +# --------------------------------------------------------------------------- +# _optimizer_step — student phase routes to "net" +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_optimizer_step_student_calls_optimizer_step_net() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + model = _make_model(student_phase=True) + + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=0) + + grad_scaler.step.assert_called_once_with(optimizer.get("net")) + grad_scaler.update.assert_called_once() + grad_scaler.step.assert_called_once() # only once total — not for fake_score + + +@pytest.mark.L0 +def test_optimizer_step_student_calls_scheduler_step_net() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + model = _make_model(student_phase=True) + + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=0) + + scheduler.get("net").step.assert_called_once() + scheduler.get("fake_score").step.assert_not_called() + + +# --------------------------------------------------------------------------- +# _optimizer_step — critic phase routes to "fake_score" +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_optimizer_step_critic_calls_optimizer_step_fake_score() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + model = _make_model(student_phase=False) + + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=1) + + grad_scaler.step.assert_called_once_with(optimizer.get("fake_score")) + grad_scaler.update.assert_called_once() + + +@pytest.mark.L0 +def test_optimizer_step_critic_calls_scheduler_step_fake_score() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + model = _make_model(student_phase=False) + + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=1) + + scheduler.get("fake_score").step.assert_called_once() + scheduler.get("net").step.assert_not_called() + + +# --------------------------------------------------------------------------- +# _optimizer_step — passes iteration to get_phase +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_optimizer_step_passes_iteration_to_get_phase() -> None: + trainer = _make_trainer() + model = _make_model(student_phase=True) + trainer._optimizer_step(model, _make_optimizer(), _make_scheduler(), _make_grad_scaler(), iteration=42) + model.get_phase.assert_called_once_with(42) + + +# --------------------------------------------------------------------------- +# _zero_grad — routes to correct key +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_zero_grad_student_zeroes_net() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + model = _make_model(student_phase=True) + + trainer._zero_grad(model, optimizer, iteration=5) + + optimizer.get("net").zero_grad.assert_called_once_with(set_to_none=True) + optimizer.get("fake_score").zero_grad.assert_not_called() + + +@pytest.mark.L0 +def test_zero_grad_critic_zeroes_fake_score() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + model = _make_model(student_phase=False) + + trainer._zero_grad(model, optimizer, iteration=6) + + optimizer.get("fake_score").zero_grad.assert_called_once_with(set_to_none=True) + optimizer.get("net").zero_grad.assert_not_called() + + +@pytest.mark.L0 +def test_zero_grad_passes_iteration_to_get_phase() -> None: + trainer = _make_trainer() + model = _make_model(student_phase=True) + trainer._zero_grad(model, _make_optimizer(), iteration=11) + model.get_phase.assert_called_once_with(11) + + +# --------------------------------------------------------------------------- +# Integration: hooks wired into the inherited training_step +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_eager_init_allocates_adam_state_when_empty() -> None: + p = torch.nn.Parameter(torch.randn(3, 4)) + opt = torch.optim.AdamW([p], lr=1e-3) + assert len(opt.state.get(p, {})) == 0 + DistillationTrainer._eager_init_optimizer_state(opt, "net") + state = opt.state[p] + assert "exp_avg" in state and "exp_avg_sq" in state + assert torch.all(state["exp_avg"] == 0) + assert torch.all(state["exp_avg_sq"] == 0) + + +@pytest.mark.L0 +def test_eager_init_resets_step_counter_to_zero() -> None: + p = torch.nn.Parameter(torch.randn(2)) + opt = torch.optim.AdamW([p], lr=1e-3) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + step_val = opt.state[p]["step"] + step_val = step_val.item() if isinstance(step_val, torch.Tensor) else step_val + assert step_val == 0 + + +@pytest.mark.L0 +def test_eager_init_does_not_change_param_values() -> None: + torch.manual_seed(0) + p = torch.nn.Parameter(torch.randn(5)) + original = p.data.clone() + opt = torch.optim.AdamW([p], lr=1e-3) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert torch.equal(p.data, original) + + +@pytest.mark.L0 +def test_eager_init_preserves_existing_grads() -> None: + p = torch.nn.Parameter(torch.randn(3)) + sentinel_grad = torch.full_like(p.data, 7.0) + p.grad = sentinel_grad + opt = torch.optim.AdamW([p], lr=1e-3) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert p.grad is sentinel_grad + + +@pytest.mark.L0 +def test_eager_init_skips_when_state_already_populated() -> None: + p = torch.nn.Parameter(torch.randn(2)) + opt = torch.optim.AdamW([p], lr=1e-3) + p.grad = torch.ones_like(p.data) + opt.step() # populates state with non-zero exp_avg + p.grad = None + exp_avg_before = opt.state[p]["exp_avg"].clone() + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert torch.equal(opt.state[p]["exp_avg"], exp_avg_before) + + +@pytest.mark.L0 +def test_eager_init_skips_params_without_requires_grad() -> None: + p_train = torch.nn.Parameter(torch.randn(2)) + p_frozen = torch.nn.Parameter(torch.randn(2), requires_grad=False) + opt = torch.optim.AdamW([p_train, p_frozen], lr=1e-3) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert len(opt.state.get(p_train, {})) > 0 + assert len(opt.state.get(p_frozen, {})) == 0 + + +@pytest.mark.L0 +def test_eager_init_creates_master_weights_without_step() -> None: + p = torch.nn.Parameter(torch.randn(3).to(torch.bfloat16)) + opt = _FakeMasterWeightOptimizer([p]) + original = p.detach().clone() + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert not opt.step_called + assert torch.equal(p.data, original) + assert opt.param_groups_master is not None + master_param = opt.param_groups_master[0]["params"][0] + assert master_param.dtype == torch.float32 + assert torch.equal(master_param, original.float()) + + +@pytest.mark.L0 +def test_eager_init_creates_fused_adam_group_state() -> None: + p = torch.nn.Parameter(torch.randn(3)) + opt = _FakeMasterWeightOptimizer([p]) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + group = opt.param_groups[0] + assert isinstance(group["step"], torch.Tensor) + assert group["step"].shape == (1,) + assert group["step"].item() == 0 + assert isinstance(group["lr"], torch.Tensor) + + +@pytest.mark.L0 +def test_eager_init_allocates_missing_state_without_overwriting_existing_state() -> None: + p_existing = torch.nn.Parameter(torch.randn(2)) + p_missing = torch.nn.Parameter(torch.randn(2)) + opt = _FakeMasterWeightOptimizer([p_existing, p_missing]) + opt.state[p_existing]["exp_avg"] = torch.ones_like(p_existing) + DistillationTrainer._eager_init_optimizer_state(opt, "net") + assert torch.equal(opt.state[p_existing]["exp_avg"], torch.ones_like(p_existing)) + assert "exp_avg" in opt.state[p_missing] + assert "exp_avg_sq" in opt.state[p_missing] + + +@pytest.mark.L0 +def test_eager_init_optimizer_container_initializes_inner_optimizers() -> None: + p_a = torch.nn.Parameter(torch.randn(2)) + p_b = torch.nn.Parameter(torch.randn(2)) + opt_a = torch.optim.AdamW([p_a], lr=1e-3) + opt_b = torch.optim.AdamW([p_b], lr=1e-3) + DistillationTrainer._eager_init_optimizer_state(_make_optimizer_container(opt_a, opt_b), "net") + assert "exp_avg" in opt_a.state[p_a] + assert "exp_avg" in opt_b.state[p_b] + + +@pytest.mark.L0 +def test_optimizer_step_runs_eager_init_only_once() -> None: + trainer = _make_trainer() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + model = _make_model(student_phase=True) + with patch.object(DistillationTrainer, "_eager_init_optimizer_state") as mock_init: + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=0) + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=1) + trainer._optimizer_step(model, optimizer, scheduler, grad_scaler, iteration=2) + # Called once per opt key on first invocation only (2 keys = 2 calls). + assert mock_init.call_count == 2 + keys_initialized = sorted(call.args[1] for call in mock_init.call_args_list) + assert keys_initialized == ["fake_score", "net"] + + +@pytest.mark.L0 +def test_eager_init_after_first_step_matches_fresh_first_step() -> None: + """After eager init + one AdamW step, weights should match a fresh-launch single step.""" + torch.manual_seed(0) + p_a = torch.nn.Parameter(torch.randn(4)) + p_b = torch.nn.Parameter(p_a.data.clone()) + grad_value = torch.randn_like(p_a.data) + + # Path A: eager init, then real step. + opt_a = torch.optim.AdamW([p_a], lr=1e-3, betas=(0.9, 0.99), eps=1e-8) + DistillationTrainer._eager_init_optimizer_state(opt_a, "net") + p_a.grad = grad_value.clone() + opt_a.step() + + # Path B: fresh first step. + opt_b = torch.optim.AdamW([p_b], lr=1e-3, betas=(0.9, 0.99), eps=1e-8) + p_b.grad = grad_value.clone() + opt_b.step() + + assert torch.allclose(p_a.data, p_b.data, atol=1e-6, rtol=1e-6) + + +@pytest.mark.L0 +def test_hooks_called_via_inherited_training_step() -> None: + """End-to-end: inherited training_step must invoke _optimizer_step and _zero_grad.""" + trainer = _make_trainer(grad_accum_iter=1) + # In ddp mode, training_step resolves model_ddp.module as the model + model_ddp = MagicMock() + model_ddp.module.get_phase.return_value = "student" + loss = torch.tensor(1.0, requires_grad=True) + model_ddp.training_step.return_value = ({"x": torch.tensor(0.0)}, loss) + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + data = {"x": torch.zeros(1)} + + with patch("cosmos_framework.utils.distributed.ddp_sync_grad") as mock_ctx: + mock_ctx.return_value = MagicMock( + __enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False) + ) + trainer.training_step(model_ddp, optimizer, scheduler, grad_scaler, data, iteration=5, grad_accum_iter=0) + + # student phase → "net" key + grad_scaler.step.assert_called_once_with(optimizer.get("net")) + grad_scaler.update.assert_called_once() + scheduler.get("net").step.assert_called_once() + optimizer.get("net").zero_grad.assert_called_once_with(set_to_none=True) + # grad_scaler.step must not be called directly by the trainer (PhaseOptimizer owns it) + # — already verified above by assert_called_once (only PhaseOptimizer calls it once) + + +class _ClosureModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight: torch.nn.Parameter = torch.nn.Parameter(torch.tensor(1.0)) + self.after_backward_calls: int = 0 + self.before_zero_grad_calls: int = 0 + self.zero_grad_calls: int = 0 + + def get_phase(self, iteration: int) -> str: + del iteration + return "student" + + def zero_grad_for_phase(self, iteration: int) -> None: + del iteration + self.zero_grad_calls += 1 + + def on_after_backward(self) -> None: + self.after_backward_calls += 1 + + def on_before_zero_grad(self, optimizer: object, scheduler: object, iteration: int) -> None: + del optimizer, scheduler, iteration + self.before_zero_grad_calls += 1 + + def training_step_closures( + self, data: dict[str, torch.Tensor], iteration: int + ) -> Iterator[tuple[str, Callable[[], tuple[dict[str, torch.Tensor], torch.Tensor]], bool]]: + del data, iteration + + def _closure(scale: float) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + loss = self.weight * scale # [] + return {"loss": loss.detach()}, loss + + yield "chunk_0", lambda: _closure(1.0), False + yield "chunk_1", lambda: _closure(2.0), True + + +class _FakeFSDPGradientSyncModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gradient_sync_calls: list[tuple[bool, bool]] = [] + + def set_requires_gradient_sync(self, enabled: bool, *, recurse: bool = True) -> None: + self.gradient_sync_calls.append((enabled, recurse)) + + +class _ClosureModelWithFSDP(_ClosureModel): + def __init__(self) -> None: + super().__init__() + self.fake_fsdp = _FakeFSDPGradientSyncModule() + + +@pytest.mark.L0 +def test_training_step_runs_model_closures_and_steps_once() -> None: + trainer = _make_trainer(grad_accum_iter=1, distributed_parallelism="none") + model = _ClosureModel() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + data = {"x": torch.zeros(1)} + + with patch("cosmos_framework.utils.distributed.ddp_sync_grad") as mock_ctx: + mock_ctx.return_value = MagicMock( + __enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False) + ) + output_batch, loss, grad_accum_iter = trainer.training_step( + model, + optimizer, + scheduler, + grad_scaler, + data, + iteration=5, + grad_accum_iter=0, + ) + + assert grad_accum_iter == 0 + torch.testing.assert_close(loss, torch.tensor(3.0)) + torch.testing.assert_close(output_batch["loss"], torch.tensor(3.0)) + torch.testing.assert_close(model.weight.grad, torch.tensor(3.0)) + assert model.after_backward_calls == 2 + assert model.before_zero_grad_calls == 1 + assert model.zero_grad_calls == 1 + grad_scaler.step.assert_called_once_with(optimizer.get("net")) + scheduler.get("net").step.assert_called_once() + + +@pytest.mark.L0 +def test_training_step_defers_fsdp2_gradient_sync_until_last_closure() -> None: + trainer = _make_trainer(grad_accum_iter=1, distributed_parallelism="none") + model = _ClosureModelWithFSDP() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + data = {"x": torch.zeros(1)} + + output_batch, loss, grad_accum_iter = trainer.training_step( + model, + optimizer, + scheduler, + grad_scaler, + data, + iteration=5, + grad_accum_iter=0, + ) + + assert grad_accum_iter == 0 + torch.testing.assert_close(loss, torch.tensor(3.0)) + torch.testing.assert_close(output_batch["loss"], torch.tensor(3.0)) + torch.testing.assert_close(model.weight.grad, torch.tensor(3.0)) + assert model.fake_fsdp.gradient_sync_calls == [ + (False, False), + (True, False), + (True, False), + (True, False), + ] + grad_scaler.step.assert_called_once_with(optimizer.get("net")) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_training_step_accumulates_closures_across_two_microbatches() -> None: + trainer = _make_trainer(grad_accum_iter=2, distributed_parallelism="none") + model = _ClosureModelWithFSDP() + optimizer = _make_optimizer() + scheduler = _make_scheduler() + grad_scaler = _make_grad_scaler() + data = {"x": torch.zeros(1)} + + _, first_loss, grad_accum_iter = trainer.training_step( + model, + optimizer, + scheduler, + grad_scaler, + data, + iteration=5, + grad_accum_iter=0, + ) + + assert grad_accum_iter == 1 + torch.testing.assert_close(first_loss, torch.tensor(3.0)) + torch.testing.assert_close(model.weight.grad, torch.tensor(1.5)) + grad_scaler.step.assert_not_called() + scheduler.get("net").step.assert_not_called() + + _, second_loss, grad_accum_iter = trainer.training_step( + model, + optimizer, + scheduler, + grad_scaler, + data, + iteration=5, + grad_accum_iter=grad_accum_iter, + ) + + assert grad_accum_iter == 0 + torch.testing.assert_close(second_loss, torch.tensor(3.0)) + torch.testing.assert_close(model.weight.grad, torch.tensor(3.0)) + grad_scaler.step.assert_called_once_with(optimizer.get("net")) + scheduler.get("net").step.assert_called_once() + assert model.fake_fsdp.gradient_sync_calls == [ + (False, False), + (True, False), + (False, False), + (True, False), + (False, False), + (True, False), + (True, False), + (True, False), + ] diff --git a/cosmos_framework/utils/callback.py b/cosmos_framework/utils/callback.py index b940692e..e2b96ccd 100644 --- a/cosmos_framework/utils/callback.py +++ b/cosmos_framework/utils/callback.py @@ -69,9 +69,21 @@ def __init__(self, config: Config, trainer: ImaginaireTrainer) -> None: continue log.critical(f"Instantiating callback {callback_name}: {current_callback_cfg}") _callback = instantiate(current_callback_cfg) - assert isinstance(_callback, Callback), f"{current_callback_cfg} is not a valid callback." - _callback.config = config - _callback.trainer = trainer + if not isinstance(_callback, Callback): + missing_hooks = _missing_callback_hooks(_callback) + if missing_hooks: + raise TypeError( + f"{current_callback_cfg} is not a valid callback; " + f"missing required callback hooks: {', '.join(missing_hooks)}" + ) + try: + _callback.config = config + _callback.trainer = trainer + except (AttributeError, TypeError) as error: + raise TypeError( + f"{current_callback_cfg} is not a valid callback; " + "cannot assign required callback metadata 'config' and 'trainer'" + ) from error self._callbacks.append(_callback) def __getattr__(self, method_name: str) -> Callable: @@ -294,6 +306,15 @@ def on_app_end(self) -> None: pass +def _missing_callback_hooks(callback: object) -> list[str]: + """Return required callback hooks that are absent or non-callable.""" + return sorted( + name + for name, member in vars(Callback).items() + if name.startswith("on_") and callable(member) and not callable(getattr(callback, name, None)) + ) + + class EMAModelCallback(Callback): """The callback class for tracking EMA model weights.""" @@ -600,3 +621,4 @@ def on_after_dataloading(self, iteration: int = 0) -> None: torch.cuda.nvtx.range_pop() +# End of callback definitions. diff --git a/cosmos_framework/utils/callback_test.py b/cosmos_framework/utils/callback_test.py new file mode 100644 index 00000000..80886e04 --- /dev/null +++ b/cosmos_framework/utils/callback_test.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from types import SimpleNamespace +from typing import Any + +import pytest + +from cosmos_framework.utils.callback import Callback, CallBackGroup + + +def _foreign_callback_type() -> type[Any]: + hooks = {name: member for name, member in vars(Callback).items() if name.startswith("on_") and callable(member)} + return type("ForeignCallback", (), hooks) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_callback_group_accepts_structurally_compatible_foreign_callback() -> None: + foreign_callback = _foreign_callback_type() + config = SimpleNamespace(trainer=SimpleNamespace(callbacks={"foreign": {"_target_": foreign_callback}})) + trainer = SimpleNamespace() + + group = CallBackGroup(config=config, trainer=trainer) + + assert len(group._callbacks) == 1 + assert group._callbacks[0].config is config + assert group._callbacks[0].trainer is trainer + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_callback_group_rejects_foreign_callback_missing_required_hooks() -> None: + class IncompleteCallback: + def on_train_start(self) -> None: + pass + + config = SimpleNamespace(trainer=SimpleNamespace(callbacks={"incomplete": {"_target_": IncompleteCallback}})) + + with pytest.raises(TypeError, match="missing required callback hooks"): + CallBackGroup(config=config, trainer=SimpleNamespace()) + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_callback_group_rejects_foreign_callback_without_metadata_slots() -> None: + hooks = {name: member for name, member in vars(Callback).items() if name.startswith("on_") and callable(member)} + foreign_callback = type("SlottedForeignCallback", (), {"__slots__": (), **hooks}) + callback_config = {"_target_": foreign_callback} + config = SimpleNamespace(trainer=SimpleNamespace(callbacks={"foreign": callback_config})) + + with pytest.raises(TypeError, match="required callback metadata") as error: + CallBackGroup(config=config, trainer=SimpleNamespace()) + + assert repr(callback_config) in str(error.value) diff --git a/cosmos_framework/utils/checkpoint_db.py b/cosmos_framework/utils/checkpoint_db.py index 580f92fe..53c5bf3e 100644 --- a/cosmos_framework/utils/checkpoint_db.py +++ b/cosmos_framework/utils/checkpoint_db.py @@ -455,6 +455,16 @@ def download_checkpoint_v2(checkpoint_uri: str, *, check_exists: bool = True) -> return checkpoint_uri if checkpoint_uri.startswith("hf://"): return _download_hf_checkpoint(checkpoint_uri) + # Bare relative registry path produced by Cosmos3-Edge-Policy-DROID. Its checkpoint + # conversion exports the Wan2.2 VAE tokenizer with an empty ``bucket_name`` (Nano + # exports ``bucket_name="bucket"``), so ``Wan2pt2VAEInterface`` forwards the raw relative + # ``vae_path`` ("pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth") here instead of the + # ``s3://bucket/...`` key. Re-form that key so it resolves through the registry and + # auto-downloads from HF, matching Nano's path. Local files still take precedence (below). + if "://" not in checkpoint_uri and not os.path.exists(checkpoint_uri): + registered = CheckpointConfig.maybe_from_uri(sanitize_uri(f"s3://bucket/{checkpoint_uri}")) + if registered is not None: + return registered.download() if check_exists and not os.path.exists(checkpoint_uri): raise ValueError(f"Checkpoint path {checkpoint_uri} does not exist.") return checkpoint_uri diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py index 64e0de79..40b26bd6 100644 --- a/cosmos_framework/utils/distributed.py +++ b/cosmos_framework/utils/distributed.py @@ -43,7 +43,8 @@ def init() -> int | None: os.sched_setaffinity(0, device.get_cpu_affinity()) except pynvml.NVMLError as e: log.warning(f"Failed to set device affinity: {e}") - # Set up NCCL communication. + # Set up distributed communication. CPU checkpoint conversion needs Gloo + # because NCCL cannot synchronize CPU-resident tokenizer or model tensors. os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "0" os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1" if dist.is_available(): @@ -52,9 +53,10 @@ def init() -> int | None: timeout_seconds = os.getenv("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", 1800) # Convert the timeout to an integer (if it isn't already) and then to a timedelta timeout_timedelta = timedelta(seconds=int(timeout_seconds)) - dist.init_process_group(backend="nccl", init_method="env://", timeout=timeout_timedelta) + backend = "nccl" if os.environ.get("COSMOS_DEVICE", "cuda").lower() == "cuda" else "gloo" + dist.init_process_group(backend=backend, init_method="env://", timeout=timeout_timedelta) log.critical( - f"Initialized distributed training with local rank {local_rank} with timeout {timeout_seconds}", + f"Initialized distributed training with local rank {local_rank} using {backend} with timeout {timeout_seconds}", rank0_only=False, ) # Increase the L2 fetch granularity for faster speed. diff --git a/cosmos_framework/utils/distributed_test.py b/cosmos_framework/utils/distributed_test.py new file mode 100644 index 00000000..0cadb4cc --- /dev/null +++ b/cosmos_framework/utils/distributed_test.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from cosmos_framework.utils import distributed + + +def test_init_uses_gloo_backend_for_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + init_process_group = Mock() + monkeypatch.setenv("COSMOS_DEVICE", "cpu") + monkeypatch.setenv("TORCH_NCCL_BLOCKING_WAIT", "0") + monkeypatch.setenv("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1") + monkeypatch.setattr(distributed, "INTERNAL", False) + monkeypatch.setattr(distributed.dist, "is_initialized", lambda: False) + monkeypatch.setattr(distributed.dist, "is_available", lambda: True) + monkeypatch.setattr(distributed.dist, "init_process_group", init_process_group) + monkeypatch.setattr(distributed.pynvml, "nvmlInit", lambda: None) + monkeypatch.setattr( + distributed, + "Device", + lambda _local_rank: SimpleNamespace(get_cpu_affinity=lambda: [0]), + ) + monkeypatch.setattr(distributed.os, "sched_setaffinity", lambda _pid, _affinity: None) + monkeypatch.setattr(distributed.torch.cuda, "set_device", lambda _local_rank: None) + + distributed.init() + + assert init_process_group.call_args.kwargs["backend"] == "gloo" diff --git a/cosmos_framework/utils/lazy_config/lazy.py b/cosmos_framework/utils/lazy_config/lazy.py index f1ae6a2c..90fa5f46 100644 --- a/cosmos_framework/utils/lazy_config/lazy.py +++ b/cosmos_framework/utils/lazy_config/lazy.py @@ -52,8 +52,10 @@ def sort_recursive(obj: Union[Dict[str, Any], List[Any], Any]) -> Union[OrderedD yaml.add_representer(OrderedDict, dict_representer) -OmegaConf.register_new_resolver("add", lambda *vals: sum(vals)) -OmegaConf.register_new_resolver("subtract", lambda *vals: vals[0] - sum(vals[1:])) +if not OmegaConf.has_resolver("add"): + OmegaConf.register_new_resolver("add", lambda *vals: sum(vals)) +if not OmegaConf.has_resolver("subtract"): + OmegaConf.register_new_resolver("subtract", lambda *vals: vals[0] - sum(vals[1:])) def _visit_dict_config(cfg, func): diff --git a/cosmos_framework/utils/lazy_config/lazy_test.py b/cosmos_framework/utils/lazy_config/lazy_test.py new file mode 100644 index 00000000..a5564f4c --- /dev/null +++ b/cosmos_framework/utils/lazy_config/lazy_test.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +import importlib + +import pytest + +from cosmos_framework.utils.lazy_config import lazy + + +@pytest.mark.L0 +@pytest.mark.CPU +def test_lazy_config_module_can_be_loaded_after_equivalent_resolvers_are_registered() -> None: + importlib.reload(lazy) diff --git a/docs/action_policy_droid_posttrain.md b/docs/action_policy_droid_posttrain.md index 80241a23..beb05671 100644 --- a/docs/action_policy_droid_posttrain.md +++ b/docs/action_policy_droid_posttrain.md @@ -1,8 +1,8 @@ -# Cosmos3-Nano-Policy-DROID Post-Training +# Cosmos3 DROID Action-Policy Post-Training -[Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID) is an action policy model post-trained from [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano), a 16B Mixture-of-Transformers model, on the [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset. The model predicts absolute joint-position actions conditioned on proprioceptive state and video observations at a resolution of 480p (`640×360`). This example reproduces the post-training procedure used to train the model. +This example post-trains [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano) into a DROID action policy on the [Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset, reproducing [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID). The model predicts absolute joint-position actions conditioned on proprioceptive state and video observations at a resolution of 480p (`640×360`). Experiment `action_policy_droid_nano`, launcher `launch_sft_action_policy_droid_nano.sh`. -Two external inputs are required: (1) a pre-downloaded Cosmos3-DROID dataset in LeRobotDataset v3.0 format, and (2) a DCP base checkpoint converted from Cosmos3-Nano. +Two external inputs are required: (1) a pre-downloaded Cosmos3-DROID dataset in LeRobotDataset v3.0 format, and (2) a DCP base checkpoint converted from the chosen base model (Cosmos3-Nano or Cosmos3-Edge). The recipe runs multi-node via HSDP (single node / 8 GPUs and beyond). @@ -32,14 +32,15 @@ The runnable artifacts (TOML recipe, paired launch shell) live in [`examples/`]( ## Inputs You Provide -This package ships the training stack — the registered `action_policy_droid_nano` experiment, -the DROID action dataset class with the recipe knobs (`action_space=joint_pos`, `use_state`, -`concat_view`), and the EMA warm-start in `checkpoint/dcp.py`. Two inputs are external and must -be provided per environment: +This package ships the training stack — the registered `action_policy_droid_nano` +experiment, the DROID action dataset class with the recipe knobs +(`action_space=joint_pos`, `use_state`, `concat_view`), and the EMA warm-start in +`checkpoint/dcp.py`. Two inputs are external and must be provided per environment: 1. **[Cosmos3-DROID](https://huggingface.co/datasets/nvidia/Cosmos3-DROID) dataset (in LeRobotDataset v3.0 format)** — pre-download the - dataset and point `DROID_ROOT` at the resulting `…/Cosmos3-DROID/success` directory (must - contain `meta/info.json`). + dataset into a directory named `droid_plus_lerobot_640x360_20260412` (the loader resolves the + dataset schema from the directory name) and point `DROID_ROOT` at that directory — the parent + containing `success/`, not `.../success` itself. 2. **DCP base checkpoint** — convert [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano) to DCP and point `BASE_CHECKPOINT_PATH` at it (see [Full Reproduction](#full-reproduction)). Action heads are not loaded from it (they init fresh). @@ -67,7 +68,12 @@ be provided per environment: The OSS flow mirrors the other recipes (see [docs/training.md](./training.md)): ```shell -# Step 1: prepare Cosmos3-DROID success split -> $DATASET_PATH (see "Inputs You Provide") +# Step 1: download Cosmos3-DROID into a directory named droid_plus_lerobot_640x360_20260412 +# (the loader resolves the dataset schema from the directory name), then point +# DATASET_PATH at that parent directory — the one containing success/, not .../success. +hf download nvidia/Cosmos3-DROID --repo-type dataset \ + --local-dir /path/to/droid_plus_lerobot_640x360_20260412 +export DATASET_PATH=/path/to/droid_plus_lerobot_640x360_20260412 # Step 2: convert the base checkpoint -> $BASE_CHECKPOINT_PATH python -m cosmos_framework.scripts.convert_model_to_dcp \ @@ -76,24 +82,27 @@ python -m cosmos_framework.scripts.convert_model_to_dcp \ # Step 3: download the keep_ranges_1_0_1.json window filter (drops idle/non-task frames -> trains # the curated ~74% window set, matching the released model). +export FILTER_DIR=/path/to/droid_filters hf download KarlP/droid keep_ranges_1_0_1.json --local-dir $FILTER_DIR # Step 4: launch. The TOML selects the experiment + scalars; the dataset/action # knobs come from the registered experiment. -export DATASET_PATH=/path/to/dataset/success export BASE_CHECKPOINT_PATH=/path/to/base_checkpoint export WAN_VAE_PATH=/path/to/Wan2.2_VAE.pth export NPROC_PER_NODE=8 # Enable the keep_ranges_1_0_1.json filter via EXTRA_TAIL_OVERRIDES (space-separated Hydra -# overrides; an exported string survives `bash `). +# overrides; an exported string survives `bash `). $FILTER_DIR is expanded HERE, +# at export time — it must already be set in this shell, or filter_dict_path silently +# becomes /keep_ranges_1_0_1.json and the dataset fails with FileNotFoundError. +: "${FILTER_DIR:?set FILTER_DIR (Step 3) before composing EXTRA_TAIL_OVERRIDES}" export EXTRA_TAIL_OVERRIDES=" \ dataloader_train.dataloader.datasets.droid.dataset.use_filter_dict=True \ dataloader_train.dataloader.datasets.droid.dataset.filter_dict_path=$FILTER_DIR/keep_ranges_1_0_1.json \ " -bash examples/launch_sft_action_policy_droid.sh +bash examples/launch_sft_action_policy_droid_nano.sh ``` -The recipe TOML ([`examples/toml/sft_config/action_policy_droid_repro.toml`](../examples/toml/sft_config/action_policy_droid_repro.toml)) sets the scalar +The recipe TOML ([`action_policy_droid_nano.toml`](../examples/toml/sft_config/action_policy_droid_nano.toml)) sets the scalar knobs (`max_iter`, `save_iter`, `grad_clip`, parallelism, wandb); the dataset/action knobs (`joint_pos`, `use_state`, `concat_view`, 480p, chunk 32, count-based batch) live in the registered `action_policy_droid_nano` experiment per the schema's design. For multi-node HSDP, diff --git a/docs/action_policy_droid_server.md b/docs/action_policy_droid_server.md index 60e00424..74939487 100644 --- a/docs/action_policy_droid_server.md +++ b/docs/action_policy_droid_server.md @@ -1,6 +1,11 @@ -# Cosmos3-Nano-Policy-DROID Server +# Cosmos3-Policy-DROID Server -[Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID) is served by a policy **Server** that streams actions to a **Client** driving a simulated or real robot. This example uses [`RoboLab`](https://github.com/NVlabs/RoboLab), a simulation benchmark for task-generalist policies, as the client. Start the server first, then connect the client. +Cosmos 3 offers two post-trained policy models for DROID: + +1. [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID) +2. [Cosmos3-Edge-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID) + +Each of these policy models is served by a policy **Server** that streams actions to a **Client** driving a simulated or real robot. This example uses [`RoboLab`](https://github.com/NVlabs/RoboLab), a simulation benchmark for task-generalist policies, as the client. Start the server first, then connect the client. @@ -62,10 +67,21 @@ The `--group=cu130-train` line targets CUDA 13.x drivers. On CUDA 12.x systems, Inside the container, start the policy server: -``` -python -m cosmos_framework.scripts.action_policy_server_robolab \ - --port 8000 -``` +1. For [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID), run: + + ``` + python -m cosmos_framework.scripts.action_policy_server_robolab \ + --port 8000 + ``` + +2. For [Cosmos3-Edge-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID), run: + + ``` + python -m cosmos_framework.scripts.action_policy_server_robolab \ + --checkpoint-path nvidia/Cosmos3-Edge-Policy-DROID \ + --port 8000 \ + --format-prompt-as-json True + ``` ## Simulation Client diff --git a/docs/action_policy_libero_sft.md b/docs/action_policy_libero_posttrain.md similarity index 85% rename from docs/action_policy_libero_sft.md rename to docs/action_policy_libero_posttrain.md index f7752e96..934b4b64 100644 --- a/docs/action_policy_libero_sft.md +++ b/docs/action_policy_libero_posttrain.md @@ -8,12 +8,12 @@ lr 5e-5, warmup 500, cycle 16000, gbs 2048): - **(A) libero_10-only** — trains on `libero_10` alone; peaks by ~iter 1500 (max_iter 2000). Fast. - `action_policy_libero_nano` + `action_policy_libero_repro.toml` + - `launch_sft_action_policy_libero.sh`. + `action_policy_libero_nano` + `action_policy_libero_10_nano.toml` + + `launch_sft_action_policy_libero_10_nano.sh`. - **(B) libero-all** — equal mix of all 4 LIBERO suites; needs longer training (max_iter 5000). - `action_policy_libero_all_nano` + `action_policy_libero_all_repro.toml` + - `launch_sft_action_policy_libero_all.sh`. + `action_policy_libero_all_nano` + `action_policy_libero_all_nano.toml` + + `launch_sft_action_policy_libero_all_nano.sh`. | Piece | Path | | ---------------- | ---------------------------------------------------------------------------------------------------- | @@ -21,8 +21,8 @@ lr 5e-5, warmup 500, cycle 16000, gbs 2048): | SFT wrapper | `get_action_libero_sft_dataset` in `.../datasets/action_sft_dataset.py` | | Norm stats | `.../normalizer_stats/libero_native_frame_wise_relative_rot6d.json` | | Experiment | `cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_libero_nano.py` | -| Run TOML | `examples/toml/sft_config/action_policy_libero_repro.toml` | -| Launch | `examples/launch_sft_action_policy_libero.sh` | +| Run TOML | `examples/toml/sft_config/action_policy_libero_10_nano.toml` | +| Launch | `examples/launch_sft_action_policy_libero_10_nano.sh` | | Inference server | `cosmos_framework/scripts/action_policy_server_libero.py` | | Closed-loop eval | `cosmos_framework/simulation/libero/closed_loop_eval.py` | @@ -54,21 +54,30 @@ eval server reproduces the same snap (§4). ## 2. Train +Convert the base checkpoint to DCP once (registered catalog name; downloads +`nvidia/Cosmos3-Nano` from the HF Hub — see [docs/training.md](./training.md) Step 2): + +```bash +python -m cosmos_framework.scripts.convert_model_to_dcp \ + -o examples/checkpoints/Cosmos3-Nano \ + --checkpoint-path Cosmos3-Nano +``` + Common env, then pick a preset launcher: ```bash export LD_LIBRARY_PATH='' # NGC container: avoid torch._C import error -export BASE_CHECKPOINT_PATH= +export BASE_CHECKPOINT_PATH=examples/checkpoints/Cosmos3-Nano # the DCP dir from the convert step export WAN_VAE_PATH= export IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root # Preset A — libero_10-only (LIBERO_ROOT = the libero_10 suite dir): export LIBERO_ROOT=/LIBERO_LeRobot_v3/libero_10 -bash examples/launch_sft_action_policy_libero.sh # HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node +bash examples/launch_sft_action_policy_libero_10_nano.sh # HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node # Preset B — libero-all 4-suite (LIBERO_ROOT = the LIBERO_LeRobot_v3 parent dir): export LIBERO_ROOT=/LIBERO_LeRobot_v3 -bash examples/launch_sft_action_policy_libero_all.sh # HSDP 2x8; needs ~4500 iters to converge +bash examples/launch_sft_action_policy_libero_all_nano.sh # HSDP 2x8; needs ~4500 iters to converge ``` Both recipes set lr 5e-5, warmup 500, cycle 16000, `save_iter=500`, HSDP 2x8 (global diff --git a/docs/inference.md b/docs/inference.md index cad7c44a..0f580432 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -10,6 +10,7 @@ ______________________________________________________________________ - [Quick Start](#quick-start) - [Single-GPU](#single-gpu) + - [Cosmos3-Edge](#cosmos3-edge) - [Multi-GPU](#multi-gpu) - [Cosmos3-Nano](#cosmos3-nano) - [Cosmos3-Super](#cosmos3-super) @@ -74,6 +75,30 @@ python -m cosmos_framework.scripts.inference \ **Note:** Cosmos3-Super (32B) does not fit on a single 80 GB H100 — see [Cosmos3-Super](#cosmos3-super) for the multi-GPU recipes. +#### Cosmos3-Edge + +Cosmos3-Edge is a compact 2B omni model. It fits comfortably on a single GPU and is the recommended starting point for single-GPU inference. It supports every mode **except audio** (`enable_sound`), since the checkpoint ships without a sound tokenizer — the audio-enabled examples (`inputs/omni/t2vs.json`, `inputs/omni/i2vs.json`) are therefore not supported. + +```shell +python -m cosmos_framework.scripts.inference \ + --parallelism-preset=latency \ + -i "inputs/omni/t2i.json" \ + -o outputs/omni_edge \ + --checkpoint-path Cosmos3-Edge \ + --seed=0 +``` + +To run every supported example in one batch (the `action_*.json` glob covers the action modes; the audio-enabled `t2vs.json` / `i2vs.json` are intentionally excluded): + +```shell +python -m cosmos_framework.scripts.inference \ + --parallelism-preset=latency \ + -i inputs/omni/t2i.json inputs/omni/t2v.json inputs/omni/i2v.json inputs/omni/v2v.json inputs/omni/action_*.json \ + -o outputs/omni_edge \ + --checkpoint-path Cosmos3-Edge \ + --seed=0 +``` + ### Multi-GPU Use `torchrun --nproc-per-node=N` when launching across multiple GPUs (N > 1). By default the model weights are sharded (FSDP) across all N GPUs, so any model fits. The `throughput` preset runs that single sharded replica over a batch; the `latency` preset additionally needs `--dp-shard-size=1` on multiple GPUs so the ranks are free for context parallelism (see [Parallelism Arguments](#parallelism-arguments)). @@ -130,6 +155,7 @@ The four `--{dp,cp,cfgp}-*-size` flags override the auto-selected values from `- | Model | Arguments | Modes | | ------------- | --------------------------------- | ---------------------------------------------- | | Cosmos3-Nano | `--checkpoint-path=Cosmos3-Nano` | All | +| Cosmos3-Edge | `--checkpoint-path=Cosmos3-Edge` | All except audio (`enable_sound`) | | Cosmos3-Super | `--checkpoint-path=Cosmos3-Super` | `text2image`, `text2video`, `image2video` | ## Modes @@ -222,6 +248,8 @@ See the [Modes](#modes) table above for the action mode inputs/outputs and examp Examples: [`inputs/reasoner/reasoner.json`](../inputs/reasoner/reasoner.json) (text), [`inputs/reasoner/reasoner_image.json`](../inputs/reasoner/reasoner_image.json) (image), [`inputs/reasoner/reasoner_video.json`](../inputs/reasoner/reasoner_video.json) (video). +**Cosmos3-Edge vision tower:** with a local `--checkpoint-path`, the vision tower and processor load from the checkpoint itself (`vision_encoder/` + processor files, bundled by default by [`export_model`](./training.md#vit--vision-tower-cosmos3-edge)) — such exports run fully offline (add `--no-guardrails`, since [guardrails](#guardrails) download their own models). If the checkpoint has no bundle (e.g. a `--no-vit` export), the tower is fetched from `nvidia/Cosmos3-Edge` on the Hub instead; when that fails (offline, missing `HF_TOKEN`), the error says so and suggests re-exporting with the default `--vit`. Nano/Super checkpoints without a vision tower (`include_visual=false`) reject reasoner image/video samples up front with a clear error. + ### Custom Defaults To use your own default values instead of the built-in presets, pass a JSON file via the `defaults_file` field in your sample arguments: diff --git a/docs/training.md b/docs/training.md index 056ceb0a..854a6fce 100644 --- a/docs/training.md +++ b/docs/training.md @@ -14,6 +14,8 @@ ______________________________________________________________________ - [Option B: raw `torchrun`](#option-b-raw-torchrun) - [Outputs](#outputs) - [Export checkpoint to Hugging Face safetensors](#export-checkpoint-to-hugging-face-safetensors) + - [ViT / vision tower (Cosmos3-Edge)](#vit--vision-tower-cosmos3-edge) +- [Convert to Diffusers](#convert-to-diffusers) - [Config](#config) - [Common Hydra tail overrides](#common-hydra-tail-overrides) @@ -77,6 +79,25 @@ uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \ +
Vision SFT (Cosmos3-Edge) + +Full fine-tune of the Nemotron-2B-Dense-VL backbone (Cosmos3-Edge), on the same Bridge dataset as **Vision SFT (Cosmos3-Nano)**. Step 2 must convert the Cosmos3-Edge checkpoint, not Cosmos3-Nano. + +Launch shell: `examples/launch_sft_vision_edge.sh` + +```shell +BASE_CHECKPOINT_NAME=Cosmos3-Edge + +# Defaults match the launcher (see Step 3 → Option A to override). +uvx hf@latest download --repo-type dataset nvidia/BridgeData2-Subset-Synthetic-Captions \ + --revision 40d018ac1c1a2a4b9734f17fdb21f3d933c49a01 \ + --local-dir examples/data/BridgeData2-Subset-Synthetic-Captions --quiet +uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth \ + --local-dir examples/checkpoints/wan22_vae --quiet +``` + +
+
Reasoner Alignment SFT with LLaVA-OneVision (vfm-vlm) Alignment SFT for the Reasoner variant on the [lmms-lab/LLaVA-OneVision-Data](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-Data) dataset (streamed from HF Hub). Skips Step 2: by default the backbone `Qwen/Qwen3-VL-8B-Instruct` is fetched from the HF Hub by the model downloader at startup — no DCP conversion needed and no required env vars. To instead start from a merged Cosmos3 reasoner snapshot (Cosmos3-Nano LM merged onto the Qwen3-VL visual tower), build it with `convert_model_to_vlm_safetensors` (see [Step 2](#step-2--prepare-checkpoint)) and point `VLM_SAFETENSORS_PATH` at it — same mechanism as the VideoPhy-2 recipe below. @@ -130,14 +151,42 @@ python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \
+
Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge) + +Edge-tier counterpart of the VideoPhy-2 recipe above: same 1–5 physical-plausibility scoring on [videophysics/videophy2_train](https://huggingface.co/datasets/videophysics/videophy2_train), but on the compact **Nemotron-2B-Dense-VL** reasoner (SigLIP2 vision tower, `model_type = "cosmos3_edge"` — native HF metadata, no remote code; the model classes are registered in-framework) instead of Qwen3-VL. `[job].task = "vlm"`. The vision tower is frozen; the projector + LM train at a uniform LR. `[model.backbone].model_name` is the **public, ungated** omni release `nvidia/Cosmos3-Edge`, which supplies the arch/config/tokenizer **and** the reasoner weights: the training loader follows the snapshot's root safetensors index into its weight shards, so weights load directly from the download — no conversion step and no required weights env var. Like the nano/super recipes, `VLM_SAFETENSORS_PATH` is an optional override for loading from a local safetensors directory instead. The reasoner weights shipped inside `nvidia/Cosmos3-Edge` are the canonical Edge reasoner weights — there is no separate reasoner repo to download. Only 2B, so it runs on a 4-GPU (e.g. GB200x4) or 8-GPU allocation. + +Launch shell: `examples/launch_sft_videophy2_edge.sh` + +```shell +# Step 1 (data): same as the nano recipe — materialize the public HF dataset into +# the canonical local layout (videophy2_{train,val}/{meta.json, media/, text/}). +python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ + --out_root examples/data/videophysics --split both +``` + +
+ +
Action-Policy Post-Training (DROID / LIBERO) + +Robot action-policy recipes: DROID (`joint_pos` 8-D actions + proprioceptive state) and +LIBERO (`frame_wise_relative` rot6d 10-D actions), both from Cosmos3-Nano. They follow the same Step 2/Step 3 flow below, with their own data staging +(DROID keep-ranges window filter, LIBERO suite selection), and are documented separately: + +- [DROID action policy](./action_policy_droid_posttrain.md) — + `examples/launch_sft_action_policy_droid_nano.sh` +- [LIBERO action policy](./action_policy_libero_posttrain.md) — + `examples/launch_sft_action_policy_libero_10_nano.sh` / `examples/launch_sft_action_policy_libero_all_nano.sh` + +
+ ## Step 2 — Prepare checkpoint Convert the base checkpoint to [PyTorch Distributed Checkpoint (DCP)](https://pytorch.org/docs/stable/distributed.checkpoint.html) format. `cosmos_framework.scripts.convert_model_to_dcp` ships in the unified `cosmos_framework/` package, so this step runs from the repo root (with the environment activated per [Setup](./setup.md)). -Set `BASE_CHECKPOINT_NAME` to the value from the recipe block you picked in Step 1 (`Cosmos3-Nano` or `Cosmos3-Super`): +Set `BASE_CHECKPOINT_NAME` to the value from the recipe block you picked in Step 1 (`Cosmos3-Nano`, `Cosmos3-Super`, or `Cosmos3-Edge`): ```shell -BASE_CHECKPOINT_NAME=Cosmos3-Nano # or Cosmos3-Super — match the recipe in Step 1 +BASE_CHECKPOINT_NAME=Cosmos3-Nano # or Cosmos3-Super / Cosmos3-Edge — match the recipe in Step 1 # Default output dir matches the launcher (see Step 3 → Option A to override). python -m cosmos_framework.scripts.convert_model_to_dcp \ @@ -145,7 +194,7 @@ python -m cosmos_framework.scripts.convert_model_to_dcp \ --checkpoint-path $BASE_CHECKPOINT_NAME ``` -`$BASE_CHECKPOINT_NAME` (e.g. `Cosmos3-Nano`, `Cosmos3-Super`) is a registered name in the checkpoint catalog; the converter downloads the matching repo from the Hugging Face Hub and writes the DCP into `examples/checkpoints/$BASE_CHECKPOINT_NAME`. +`$BASE_CHECKPOINT_NAME` (e.g. `Cosmos3-Nano`, `Cosmos3-Super`, `Cosmos3-Edge`) is a registered name in the checkpoint catalog; the converter downloads the matching repo from the Hugging Face Hub and writes the DCP into `examples/checkpoints/$BASE_CHECKPOINT_NAME`. **Reasoner Alignment SFT with LLaVA-OneVision (vfm-vlm):** Skip this step — the Reasoner alignment SFT loads `Qwen/Qwen3-VL-8B-Instruct` from the HF Hub at startup (no DCP conversion required). To start from a merged Cosmos3 reasoner snapshot instead, build one with `convert_model_to_vlm_safetensors` (see the VideoPhy-2 note below) and pass it via `VLM_SAFETENSORS_PATH`. @@ -169,6 +218,8 @@ python -m cosmos_framework.scripts.convert_model_to_vlm_safetensors \ -o examples/checkpoints/Cosmos3-Super-VLM ``` +**Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge):** Skip this step — the recipe loads the reasoner weights directly from the (ungated) `nvidia/Cosmos3-Edge` snapshot at startup: the training loader follows the repo's root safetensors index into its weight shards, so no conversion or extraction is required. To start from a local safetensors directory instead, point `VLM_SAFETENSORS_PATH` at it (optional, same mechanism as the nano recipe). + ## Step 3 — Run training **Weights & Biases (optional):** every recipe TOML defaults to `job.wandb_mode = "disabled"`. To log a run to W&B, flip that field to `"online"` in the TOML and export `WANDB_API_KEY` in your environment before launching. @@ -188,12 +239,16 @@ Each launcher's default paths come from the `DATASET_PATH` + `BASE_CHECKPOINT_PA | ------------------------------ | ------------------ | ---------------------------------------------------------- | ------------------------------------------------------------------ | | `launch_sft_vision_nano.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Nano` | | `launch_sft_vision_super.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Super` | +| `launch_sft_vision_edge.sh` | Generator SFT | `BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge` | `Cosmos3-Edge` | | `launch_sft_llava_ov.sh` | Reasoner SFT | (none; dataset streams from HF Hub) | (none; backbone fetched at startup, or set `VLM_SAFETENSORS_PATH`) | | `launch_sft_videophy2_nano.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env) | | `launch_sft_videophy2_super.sh`| Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; set `VLM_SAFETENSORS_PATH` env — Cosmos3-Super-VLM merge) | +| `launch_sft_videophy2_edge.sh` | Reasoner SFT | (none; set `VIDEOPHYSICS_ROOT` env) | (none; weights load directly from `nvidia/Cosmos3-Edge`) | `WAN_VAE_PATH` defaults to `examples/checkpoints/wan22_vae/Wan2.2_VAE.pth` for every non-reasoner recipe. +Cosmos3-Edge is only 2B and also fits a 4-GPU node — e.g. `NPROC_PER_NODE=4 bash examples/launch_sft_vision_edge.sh` (tested on 4×GB200). + **Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Nano):** ```shell @@ -213,6 +268,16 @@ export VLM_SAFETENSORS_PATH=$PWD/examples/checkpoints/Cosmos3-Super-VLM bash examples/launch_sft_videophy2_super.sh ``` +**Reasoner Alignment SFT with VideoPhy-2 (Cosmos3-Edge):** + +```shell +# Same as the nano recipe, but no VLM_SAFETENSORS_PATH is needed: the reasoner +# weights load directly from the (ungated) nvidia/Cosmos3-Edge snapshot. +# Only 2B — on a 4-GPU node set NPROC_PER_NODE=4. +export VIDEOPHYSICS_ROOT=$PWD/examples/data/videophysics +bash examples/launch_sft_videophy2_edge.sh +``` + #### Overriding the defaults If you'd rather put data or checkpoints on a different filesystem (e.g. a faster SSD or shared mount), download to your chosen path in Step 1 / convert the DCP to your chosen path in Step 2, then export the matching env var(s) before launching: @@ -296,6 +361,52 @@ python -m cosmos_framework.scripts.export_model \ The exported safetensors land at `$RUN_DIR/model` and can be used in [Inference](../README.md#inference) commands by passing `--checkpoint-path $RUN_DIR/model`. +### ViT / vision tower (Cosmos3-Edge) + +The plain export above is all any recipe needs — Edge included. For Cosmos3-Edge, `export_model` bundles the SigLIP2 vision tower into `$RUN_DIR/model/vision_encoder/` and the processor/tokenizer files into the export root, so the exported directory is **self-contained**: inference runs from it offline, with no Hub download (see [Inference → Reasoner](./inference.md#reasoner)). `export_manifest.json` records where the tower and processor came from. Nano/Super exports bundle their processor files the same way. + +Optional flags: + +```shell +# --no-vit: generation-only export (T2V / I2V / V2V / T2I), ~1 GB smaller. +# Writes include_visual=false so the checkpoint loads cleanly everywhere. +python -m cosmos_framework.scripts.export_model \ + --checkpoint-path $CHECKPOINT_PATH \ + --config-file $RUN_DIR/config.yaml \ + --no-vit \ + -o $RUN_DIR/model + +# --vit-checkpoint-path: use a custom tower from a local directory +# (a snapshot root or a vision_encoder/ directory). +python -m cosmos_framework.scripts.export_model \ + --checkpoint-path $CHECKPOINT_PATH \ + --config-file $RUN_DIR/config.yaml \ + --vit-checkpoint-path /path/to/snapshot \ + -o $RUN_DIR/model + +# --verify: smoke-test the export on this node's GPU (tiny reasoner query + +# tiny generation). A failure exits non-zero but keeps the export. +python -m cosmos_framework.scripts.export_model \ + --checkpoint-path $CHECKPOINT_PATH \ + --config-file $RUN_DIR/config.yaml \ + --verify \ + -o $RUN_DIR/model +``` + +## Convert to Diffusers + +Convert the Hugging Face export from the previous step into a Diffusers pipeline (`transformer/`, `vae/`, `scheduler/`, `text_tokenizer/`, `model_index.json`, …): + +```shell +python -m cosmos_framework.scripts.convert_model_to_diffusers \ + --checkpoint-path $RUN_DIR/model \ + -o $RUN_DIR/diffusers +``` + +The input is the Hugging Face safetensors directory produced by `export_model` above — not a raw DCP checkpoint, so run the export first. The result at `$RUN_DIR/diffusers` can be loaded with `diffusers` or passed to [Inference](../README.md#inference) via `--checkpoint-path $RUN_DIR/diffusers`. + +A generation-only checkpoint with no reasoner vision tower (`include_visual` unset) is converted without a `vision_encoder/` sidecar automatically; pass `--skip-vision-encoder` to force-skip it even when a tower is present. + ## Config The recipe TOML is parsed against the pydantic schema [`SFTExperimentConfig`](../cosmos_framework/configs/toml_config/sft_config.py) at load time. Every top-level key listed below maps to a sub-model in that file; unknown keys raise a `ValidationError` before training starts (`extra="forbid"` on every sub-model). Values may use OmegaConf env interpolation `${oc.env:NAME}` — the recipe TOMLs use this for `BASE_CHECKPOINT_PATH` (`[checkpoint].load_path`) and `WAN_VAE_PATH` (`[model.tokenizer].vae_path`). `DATASET_PATH` is consumed the same way but inside the experiment-SKU Python (`cosmos_framework/configs/base/experiment/sft/.py`), not in the TOML. diff --git a/examples/README.md b/examples/README.md index 08c219a8..a106712b 100644 --- a/examples/README.md +++ b/examples/README.md @@ -14,6 +14,8 @@ This directory contains: | -------------------------------------------- | ------------------------------------- | | Vision SFT (Cosmos3-Nano) | `launch_sft_vision_nano.sh` | | Vision SFT LoRA (Cosmos3-Super) | `launch_sft_vision_super.sh` | +| Vision SFT (Cosmos3-Edge) | `launch_sft_vision_edge.sh` | | Reasoner Alignment SFT | `launch_sft_llava_ov.sh` | | Reasoner Alignment SFT (Cosmos3-Nano) | `launch_sft_videophy2_nano.sh` | | Reasoner Alignment SFT (Cosmos3-Super) | `launch_sft_videophy2_super.sh` | +| Reasoner Alignment SFT (Cosmos3-Edge) | `launch_sft_videophy2_edge.sh` | diff --git a/examples/launch_sft_action_policy_droid.sh b/examples/launch_sft_action_policy_droid_nano.sh similarity index 78% rename from examples/launch_sft_action_policy_droid.sh rename to examples/launch_sft_action_policy_droid_nano.sh index b8755acd..fd06c8b1 100755 --- a/examples/launch_sft_action_policy_droid.sh +++ b/examples/launch_sft_action_policy_droid_nano.sh @@ -5,13 +5,13 @@ # ============================================================================ # Structured-TOML launch for DROID action-policy SFT on Cosmos3-Nano (8B MoT). # Drives cosmos_framework.scripts.train against -# examples/toml/sft_config/action_policy_droid_repro.toml (selects the +# examples/toml/sft_config/action_policy_droid_nano.toml (selects the # registered `action_policy_droid_nano` experiment; res480, joint_pos 8D + # use_state, trains the generation + action heads). See # docs/action_policy_droid_posttrain.md. # # Env vars (override for your filesystem): -# DATASET_PATH Cosmos3-DROID success split (…/Cosmos3-DROID/success) +# DATASET_PATH Cosmos3-DROID download dir; must be named droid_plus_lerobot_640x360_20260412 # BASE_CHECKPOINT_PATH DCP of nvidia/Cosmos3-Nano (convert_model_to_dcp; see docs) # WAN_VAE_PATH Wan2.2 VAE .pth (Wan-AI/Wan2.2-TI2V-5B) # WANDB_API_KEY for online logging (TOML wandb_mode="online") @@ -21,21 +21,21 @@ # Single-node smoke (config/data sanity, a few iters): # export EXTRA_TAIL_OVERRIDES="trainer.max_iter=10 checkpoint.save_iter=10 \ # dataloader_train.max_samples_per_batch=32" -# bash examples/launch_sft_action_policy_droid.sh +# bash examples/launch_sft_action_policy_droid_nano.sh # # Multi-node: launch on every worker; the trainer reads torchrun's # --nnodes/--node_rank. For HSDP set # model.parallelism.data_parallel_replicate_degree = (shard stays 8). # ============================================================================ -TOML_FILE="examples/toml/sft_config/action_policy_droid_repro.toml" -: "${DATASET_PATH:=examples/data/lerobot_v30/droid_lerobot/success}" +TOML_FILE="examples/toml/sft_config/action_policy_droid_nano.toml" +: "${DATASET_PATH:=examples/data/droid_plus_lerobot_640x360_20260412}" : "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" # The experiment reads ${oc.env:DROID_ROOT}; bridge the launcher's DATASET_PATH to it. export DROID_ROOT="${DROID_ROOT:-$DATASET_PATH}" -EXTRA_DATASET_CHECK='[[ -f "$DROID_ROOT/meta/info.json" ]] || { echo "ERROR: missing $DROID_ROOT/meta/info.json (prepare Cosmos3-DROID — see docs/action_policy_droid_posttrain.md)" >&2; exit 1; }' +EXTRA_DATASET_CHECK='[[ -f "$DROID_ROOT/success/meta/info.json" ]] || { echo "ERROR: missing $DROID_ROOT/success/meta/info.json (prepare Cosmos3-DROID — see docs/action_policy_droid_posttrain.md)" >&2; exit 1; }' # Extra Hydra overrides from the environment: a space-separated string word-split into # the TAIL_OVERRIDES array. An exported string survives `bash ` (a child diff --git a/examples/launch_sft_action_policy_libero.sh b/examples/launch_sft_action_policy_libero_10_nano.sh similarity index 89% rename from examples/launch_sft_action_policy_libero.sh rename to examples/launch_sft_action_policy_libero_10_nano.sh index 29a9a16a..4d3e7595 100755 --- a/examples/launch_sft_action_policy_libero.sh +++ b/examples/launch_sft_action_policy_libero_10_nano.sh @@ -4,12 +4,12 @@ # Structured-TOML launch for action_policy_libero_nano — Cosmos3-Nano LIBERO # action-policy SFT (HSDP, full SFT). Drives cosmos_framework.scripts.train -# against examples/toml/sft_config/action_policy_libero_repro.toml. +# against examples/toml/sft_config/action_policy_libero_10_nano.toml. # # Point LIBERO_ROOT at the libero_10 suite ONLY. Use the 20 FPS # nvidia/LIBERO_LeRobot_v3. The default recipe is HSDP 2x8 (global batch 2048); # set NNODES/NODE_RANK/MASTER_ADDR per node. -# See docs/action_policy_libero_sft.md. +# See docs/action_policy_libero_posttrain.md. # # Required env vars: # LIBERO_ROOT local LIBERO-10 LeRobot dataset dir, e.g. /libero_10 (no default) @@ -24,16 +24,16 @@ # export LIBERO_ROOT=/libero_10 # # Usage (HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node): -# LIBERO_ROOT=/libero_10 bash examples/launch_sft_action_policy_libero.sh +# LIBERO_ROOT=/libero_10 bash examples/launch_sft_action_policy_libero_10_nano.sh -TOML_FILE="examples/toml/sft_config/action_policy_libero_repro.toml" +TOML_FILE="examples/toml/sft_config/action_policy_libero_10_nano.toml" : "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" # LIBEROLeRobotDataset reads ${oc.env:LIBERO_ROOT} directly (a LOCAL LeRobot dir); # export it so torchrun (launched in this shell) inherits it. export LIBERO_ROOT="${LIBERO_ROOT:-}" -EXTRA_DATASET_CHECK='[[ -f "$LIBERO_ROOT/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be a local LeRobot dir containing meta/info.json (got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --include '\''libero_10/**'\'' --local-dir (then LIBERO_ROOT=/libero_10). See docs/action_policy_libero_sft.md" >&2; exit 1; }' +EXTRA_DATASET_CHECK='[[ -f "$LIBERO_ROOT/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be a local LeRobot dir containing meta/info.json (got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --include '\''libero_10/**'\'' --local-dir (then LIBERO_ROOT=/libero_10). See docs/action_policy_libero_posttrain.md" >&2; exit 1; }' # Extra Hydra overrides from the environment: a space-separated string word-split into # the TAIL_OVERRIDES array. An exported string survives `bash ` (a child diff --git a/examples/launch_sft_action_policy_libero_all.sh b/examples/launch_sft_action_policy_libero_all_nano.sh similarity index 91% rename from examples/launch_sft_action_policy_libero_all.sh rename to examples/launch_sft_action_policy_libero_all_nano.sh index 67eeb6f4..49180d36 100755 --- a/examples/launch_sft_action_policy_libero_all.sh +++ b/examples/launch_sft_action_policy_libero_all_nano.sh @@ -5,14 +5,14 @@ # Structured-TOML launch for action_policy_libero_all_nano — Cosmos3-Nano # LIBERO-all (4-suite) action-policy SFT (HSDP, full SFT). Drives # cosmos_framework.scripts.train against -# examples/toml/sft_config/action_policy_libero_all_repro.toml. +# examples/toml/sft_config/action_policy_libero_all_nano.toml. # # Trains on all 4 LIBERO suites (equal mix). Point LIBERO_ROOT at the # LIBERO_LeRobot_v3 PARENT dir (containing libero_spatial/object/goal/10), NOT a # single suite. Use the 20 FPS nvidia/LIBERO_LeRobot_v3. Default recipe is # HSDP 2x8 (global batch 2048); set NNODES/NODE_RANK/MASTER_ADDR per node. # The 4-suite mix is coverage-limited — it needs ~4500 iters to reach ~95% on -# libero_10 (max_iter defaults to 5000). See docs/action_policy_libero_sft.md. +# libero_10 (max_iter defaults to 5000). See docs/action_policy_libero_posttrain.md. # # Required env vars: # LIBERO_ROOT local LIBERO_LeRobot_v3 PARENT dir (no default) @@ -27,16 +27,16 @@ # export LIBERO_ROOT= # # Usage (HSDP 2x8; set NNODES/NODE_RANK/MASTER_ADDR per node): -# LIBERO_ROOT= bash examples/launch_sft_action_policy_libero_all.sh +# LIBERO_ROOT= bash examples/launch_sft_action_policy_libero_all_nano.sh -TOML_FILE="examples/toml/sft_config/action_policy_libero_all_repro.toml" +TOML_FILE="examples/toml/sft_config/action_policy_libero_all_nano.toml" : "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Nano}" # The libero-all experiment reads ${oc.env:LIBERO_ROOT}/ for each of the 4 # suites; export the PARENT dir so torchrun (launched in this shell) inherits it. export LIBERO_ROOT="${LIBERO_ROOT:-}" -EXTRA_DATASET_CHECK='for _s in libero_spatial libero_object libero_goal libero_10; do [[ -f "$LIBERO_ROOT/$_s/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be the LIBERO_LeRobot_v3 parent dir containing all 4 suites (missing $_s; got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --local-dir (then LIBERO_ROOT=). See docs/action_policy_libero_sft.md" >&2; exit 1; }; done' +EXTRA_DATASET_CHECK='for _s in libero_spatial libero_object libero_goal libero_10; do [[ -f "$LIBERO_ROOT/$_s/meta/info.json" ]] || { echo "ERROR: LIBERO_ROOT must be the LIBERO_LeRobot_v3 parent dir containing all 4 suites (missing $_s; got: '\''$LIBERO_ROOT'\''). Pre-sync: hf download nvidia/LIBERO_LeRobot_v3 --repo-type dataset --local-dir (then LIBERO_ROOT=). See docs/action_policy_libero_posttrain.md" >&2; exit 1; }; done' # Extra Hydra overrides from the environment: a space-separated string word-split into # the TAIL_OVERRIDES array. An exported string survives `bash ` (a child diff --git a/examples/launch_sft_videophy2_edge.sh b/examples/launch_sft_videophy2_edge.sh new file mode 100755 index 00000000..685545c0 --- /dev/null +++ b/examples/launch_sft_videophy2_edge.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for videophy2_sft_edge (VLM dialog SFT on VideoPhy-2 via +# CosmosDataLoader) targeting the Cosmos3-Edge reasoner backbone (public, +# ungated nvidia/Cosmos3-Edge, model_type cosmos3_edge — native HF metadata, +# no remote code; the classes are registered in-framework). Drives +# cosmos_framework.scripts.train against +# examples/toml/sft_config/videophy2_sft_edge.toml. +# +# [job].task = "vlm" — picks cosmos_framework/configs/base/reasoner/config.py as the base config. +# +# Reasoner weights load DIRECTLY from the nvidia/Cosmos3-Edge snapshot resolved +# via model_name: the training loader follows the repo's root safetensors index +# into its weight shards. No converter step and no required weights env var. +# +# Required env: +# VIDEOPHYSICS_ROOT dir containing videophy2_train/ and videophy2_val/ +# (each with meta.json + media/ + text/). Populate via +# `python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf`. +# +# Optional env: +# VLM_SAFETENSORS_PATH local directory of reasoner safetensors to load +# INSTEAD of the nvidia/Cosmos3-Edge snapshot (same +# optional override as the nano/super launchers). +# When set, plumbed to backbone.safetensors_path via a +# tail override; model_name still drives +# tokenizer/architecture discovery. +# HF_TOKEN NOT needed for nvidia/Cosmos3-Edge (the repo is +# ungated); set it only if other downloads in your +# environment require authentication. +# NPROC_PER_NODE torchrun GPUs per node; default 8. Set 4 on a GB200x4 +# node — Edge is only 2B and fits a 4-GPU allocation. +# EXTRA_TAIL_OVERRIDES extra Hydra-style overrides. On nodes without a +# flash-attn wheel fall back to the portable attention +# impl: +# EXTRA_TAIL_OVERRIDES='model.config.policy.attn_implementation=sdpa' +# +# Usage (8-GPU allocation, inside the training container, from the repo root): +# VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_edge.sh +# # on a 4-GPU node (e.g. GB200x4): +# NPROC_PER_NODE=4 VIDEOPHYSICS_ROOT=/path/to/videophysics bash examples/launch_sft_videophy2_edge.sh + +TOML_FILE="examples/toml/sft_config/videophy2_sft_edge.toml" + +TAIL_OVERRIDES=( + ${EXTRA_TAIL_OVERRIDES:-} +) + +# Optional: when VLM_SAFETENSORS_PATH is set, plumb it to backbone.safetensors_path +# so the framework loads reasoner weights from the local directory instead of the +# nvidia/Cosmos3-Edge snapshot (the public HF model_name still drives +# tokenizer/architecture discovery). When unset, nothing is added and weights come +# directly from the snapshot. +if [[ -n "${VLM_SAFETENSORS_PATH:-}" ]]; then + TAIL_OVERRIDES+=("model.config.policy.backbone.safetensors_path=$VLM_SAFETENSORS_PATH") +fi + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/launch_sft_vision_edge.sh b/examples/launch_sft_vision_edge.sh new file mode 100755 index 00000000..c6220348 --- /dev/null +++ b/examples/launch_sft_vision_edge.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# Structured-TOML launch for vision_sft_edge (T2V / I2V / V2V vision-only +# SFT on Nemotron-2B-Dense-VL / Cosmos3-Edge, 8-GPU FSDP). Drives cosmos_framework.scripts.train against +# examples/toml/sft_config/vision_sft_edge.toml. +# +# Optional env vars (defaults below point under examples/; override to put +# data or checkpoints on a different filesystem): +# DATASET_PATH default: examples/data/BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge +# (must contain train/video_dataset_file.jsonl) +# BASE_CHECKPOINT_PATH default: examples/checkpoints/Cosmos3-Edge +# WAN_VAE_PATH default: examples/checkpoints/wan22_vae/Wan2.2_VAE.pth +# HF_TOKEN not needed for nvidia/Cosmos3-Edge (the repo is +# ungated); set only if another download requires it +# OUTPUT_ROOT default: outputs/train +# +# Usage (8-GPU allocation, inside the training container, from the repo root): +# bash examples/launch_sft_vision_edge.sh + +TOML_FILE="examples/toml/sft_config/vision_sft_edge.toml" +: "${DATASET_PATH:=examples/data/BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge}" +: "${BASE_CHECKPOINT_PATH:=examples/checkpoints/Cosmos3-Edge}" + +EXTRA_DATASET_CHECK='[[ -f "$DATASET_PATH/train/video_dataset_file.jsonl" ]] || { echo "ERROR: missing $DATASET_PATH/train/video_dataset_file.jsonl" >&2; exit 1; }' + +source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh" diff --git a/examples/toml/sft_config/action_policy_droid_repro.toml b/examples/toml/sft_config/action_policy_droid_nano.toml similarity index 98% rename from examples/toml/sft_config/action_policy_droid_repro.toml rename to examples/toml/sft_config/action_policy_droid_nano.toml index 2326d407..ced8e35f 100644 --- a/examples/toml/sft_config/action_policy_droid_repro.toml +++ b/examples/toml/sft_config/action_policy_droid_nano.toml @@ -21,7 +21,7 @@ task = "vfm" experiment = "action_policy_droid_nano" project = "cosmos3_action" group = "action_sft" -name = "action_policy_droid_repro" +name = "action_policy_droid_nano" wandb_mode = "online" [model] diff --git a/examples/toml/sft_config/action_policy_libero_repro.toml b/examples/toml/sft_config/action_policy_libero_10_nano.toml similarity index 96% rename from examples/toml/sft_config/action_policy_libero_repro.toml rename to examples/toml/sft_config/action_policy_libero_10_nano.toml index a0c49c7a..a388fa91 100644 --- a/examples/toml/sft_config/action_policy_libero_repro.toml +++ b/examples/toml/sft_config/action_policy_libero_10_nano.toml @@ -11,7 +11,7 @@ task = "vfm" experiment = "action_policy_libero_nano" project = "cosmos3_action_libero" group = "action_sft" -name = "action_policy_libero_repro" +name = "action_policy_libero_10_nano" wandb_mode = "online" [model] diff --git a/examples/toml/sft_config/action_policy_libero_all_repro.toml b/examples/toml/sft_config/action_policy_libero_all_nano.toml similarity index 96% rename from examples/toml/sft_config/action_policy_libero_all_repro.toml rename to examples/toml/sft_config/action_policy_libero_all_nano.toml index bb4d7c6a..14a43adb 100644 --- a/examples/toml/sft_config/action_policy_libero_all_repro.toml +++ b/examples/toml/sft_config/action_policy_libero_all_nano.toml @@ -14,7 +14,7 @@ task = "vfm" experiment = "action_policy_libero_all_nano" project = "cosmos3_action_libero" group = "action_sft" -name = "action_policy_libero_all_repro" +name = "action_policy_libero_all_nano" wandb_mode = "online" [model] diff --git a/examples/toml/sft_config/videophy2_sft_edge.toml b/examples/toml/sft_config/videophy2_sft_edge.toml new file mode 100644 index 00000000..142289c1 --- /dev/null +++ b/examples/toml/sft_config/videophy2_sft_edge.toml @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# videophy2_sft_edge — VLM dialog SFT on VideoPhy-2 via CosmosDataLoader, on the +# Cosmos3-Edge reasoner backbone (public, ungated nvidia/Cosmos3-Edge, model_type +# cosmos3_edge — native HF metadata; the classes are registered in-framework). +# Base config = cosmos_framework/configs/base/reasoner/config.py +# (selected by [job].task="vlm"). +# +# Dataset prep: +# python -m cosmos_framework.scripts.reasoner.prepare_videophy2_from_hf \ +# --out_root $VIDEOPHYSICS_ROOT --split both +# +# Required env at launch: VIDEOPHYSICS_ROOT (read by the experiment Python). +# +# Example launch: +# bash examples/launch_sft_videophy2_edge.sh + +[job] +task = "vlm" +experiment = "videophy2_sft_edge" +project = "cosmos3" +group = "vlm_videophy2_sft" +name = "videophy2_sft_edge" +wandb_mode = "disabled" + +[model] +# Edge delta vs nano: the "cosmos" adapter is Qwen3-VL-only and rejects the Edge +# reasoner's (cosmos3_edge) mask. flash_attention_2 is ~14% faster than sdpa with +# matching loss curves; set "sdpa" where flash-attn is not installed (portable fallback). +attn_implementation = "flash_attention_2" +precision = "bfloat16" # was [model.parallelism].precision + +# Supplies arch/config/tokenizer AND weights: the training loader follows the +# snapshot's root safetensors index into its weight shards, so reasoner weights +# load directly from the (ungated) download — no converter step needed. +# (VLM_SAFETENSORS_PATH remains an optional local-weights override in the launcher.) +[model.backbone] +model_name = "nvidia/Cosmos3-Edge" + +[model.ema] +enabled = false +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +data_parallel_shard_degree = -1 # FSDP full shard, auto from WORLD_SIZE (4- or 8-GPU) +data_parallel_replicate_degree = 1 +context_parallel_shard_degree = 1 +cfg_parallel_shard_degree = 1 + +[model.compile] +enabled = false # was [model.parallelism].use_torch_compile +compile_dynamic = true + +[model.activation_checkpointing] +mode = "full" +save_ops_regex = ["fmha"] +preserve_rng_state = true +determinism_check = "default" + +[optimizer] +betas = [0.9, 0.95] +eps = 1.0e-8 +fused = true +lr = 1.0e-6 +weight_decay = 0.1 + +[scheduler] +cycle_lengths = [50] +f_max = [1.0] +f_min = [0.1] +f_start = [0.05] +verbosity_interval = 0 +warm_up_steps = [5] + +[trainer] +distributed_parallelism = "fsdp" +grad_accum_iter = 8 +logging_iter = 1 +max_iter = 50 + +[trainer.callbacks.compile_tokenizer] +compile_after_iterations = 3 +enabled = false + +[trainer.callbacks.grad_clip] +clip_norm = 1.0 +force_finite = false + +[checkpoint] +keys_to_skip_loading = [] +load_path = "???" +save_iter = 100 + +[dataloader_train] +# PATH_REMAPS["vlm"] routes these onto the CosmosDataLoader batcher +# (max_batch_size / max_tokens). +max_samples_per_batch = 1 +max_sequence_length = 16000 diff --git a/examples/toml/sft_config/vision_sft_edge.toml b/examples/toml/sft_config/vision_sft_edge.toml new file mode 100644 index 00000000..dac2e50e --- /dev/null +++ b/examples/toml/sft_config/vision_sft_edge.toml @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# vision_sft_edge — T2V / I2V / V2V vision-only SFT (Nemotron-2B-Dense-VL / Cosmos3-Edge) +# Consumed by cosmos_framework.configs.toml_config.sft_config.load_experiment_from_toml. +# Uses PackingDataLoader (no dataloader_train.seed slot — keep it omitted here). + +[job] +task = "vfm" +experiment = "vision_sft_edge" +project = "cosmos3" +group = "sft" +name = "vision_sft_edge" +wandb_mode = "disabled" + +[model] +max_num_tokens_after_packing = 45056 +joint_attn_implementation = "two_way" +precision = "bfloat16" # was [model.parallelism].precision + +[model.ema] +enabled = true +rate = 0.1 +iteration_shift = 0 + +[model.parallelism] +data_parallel_shard_degree = -1 # -1 = auto from WORLD_SIZE (matches legacy) +data_parallel_replicate_degree = 1 + +[model.compile] +enabled = true # was [model.parallelism].use_torch_compile +compile_dynamic = true + +[model.activation_checkpointing] +mode = "full" +save_ops_regex = ["fmha"] +preserve_rng_state = true +determinism_check = "default" + +[model.tokenizer] +vae_path = "${oc.env:WAN_VAE_PATH}" + +[optimizer] +betas = [0.9, 0.95] +eps = 1.0e-6 +fused = true +keys_to_select = [ + "moe_gen", + "time_embedder", + "vae2llm", + "llm2vae", + "k_norm_und_for_gen", # und-K norm trains with the gen pathway +] +lr = 1.0e-4 # tuned via LR sweep: full fine-tune converges faster at 1e-4 than 2e-5 (LoRA super stays 5e-4) +weight_decay = 0 # int matches legacy YAML repr +# lr_multipliers intentionally empty for vision SFT (Hydra default {} stands). + +[scheduler] +cycle_lengths = [1000] +f_max = [1.0] +f_min = [0.0] +f_start = [0.0] +verbosity_interval = 0 +warm_up_steps = [50] + +[trainer] +distributed_parallelism = "fsdp" +grad_accum_iter = 2 +logging_iter = 1 +max_iter = 500 + +[trainer.callbacks.compile_tokenizer] +compile_after_iterations = 3 +enabled = false +# warmup_resolutions omitted (None at experiment level) + +[trainer.callbacks.grad_clip] +clip_norm = 0.1 +force_finite = true + +[checkpoint] +keys_to_skip_loading = ["net_ema."] +load_path = "${oc.env:BASE_CHECKPOINT_PATH}" +save_iter = 100 + +[dataloader_train] +max_sequence_length = 45056 +# Per-caption token cap before truncation. Structured-JSON captions run longer than +# dense prose (measured max ~1790 tokens), so keep headroom; raise it for longer captions. +max_caption_tokens = 2048 +# max_samples_per_batch omitted (None — PackingDataLoader doesn't cap by sample count) +# seed omitted — PackingDataLoader has no seed ctor kwarg diff --git a/examples/toml/sft_config/vision_sft_nano.toml b/examples/toml/sft_config/vision_sft_nano.toml index dbb192dc..d6fd8b33 100644 --- a/examples/toml/sft_config/vision_sft_nano.toml +++ b/examples/toml/sft_config/vision_sft_nano.toml @@ -50,7 +50,7 @@ keys_to_select = [ "vae2llm", "llm2vae", ] -lr = 2.0e-5 +lr = 1.0e-4 # tuned via LR sweep: full fine-tune converges faster at 1e-4 than 2e-5 (LoRA super stays 5e-4) weight_decay = 0 # int matches legacy YAML repr # lr_multipliers intentionally empty for vision SFT (Hydra default {} stands). diff --git a/examples/toml/sft_config/vision_sft_nano_mapstyle_dataloader.toml b/examples/toml/sft_config/vision_sft_nano_mapstyle_dataloader.toml index 5d76be8b..df4bd688 100644 --- a/examples/toml/sft_config/vision_sft_nano_mapstyle_dataloader.toml +++ b/examples/toml/sft_config/vision_sft_nano_mapstyle_dataloader.toml @@ -58,7 +58,7 @@ keys_to_select = [ "vae2llm", "llm2vae", ] -lr = 2.0e-5 +lr = 1.0e-4 # tuned via LR sweep: full fine-tune converges faster at 1e-4 than 2e-5 (LoRA super stays 5e-4) weight_decay = 0 # int matches legacy YAML repr # lr_multipliers intentionally empty for vision SFT (Hydra default {} stands). diff --git a/tests/_stage_h100_inputs.sh b/tests/_stage_h100_inputs.sh index db9d9b09..6ccfee45 100755 --- a/tests/_stage_h100_inputs.sh +++ b/tests/_stage_h100_inputs.sh @@ -50,14 +50,19 @@ if [[ "$(python -c 'import transformers; print(transformers.__version__)')" != " fi echo ">>> $(date '+%H:%M:%S') transformers=$(python -c 'import transformers; print(transformers.__version__)')" +# retry once with --refresh: `hf@latest` doesn't refresh dependency index metadata +_hf_download() { + uvx hf@latest download "$@" --quiet \ + || uvx --refresh hf@latest download "$@" --quiet +} + # ---------------------------------------------------------------------------- # 1. Mixed-modality SFT dataset (BridgeData2-Subset-Synthetic-Captions). # ---------------------------------------------------------------------------- echo ">>> $(date '+%H:%M:%S') downloading BridgeData2-Subset-Synthetic-Captions ..." -BRIDGE_ROOT=$(uvx hf@latest download --repo-type dataset \ +BRIDGE_ROOT=$(_hf_download --repo-type dataset \ nvidia/BridgeData2-Subset-Synthetic-Captions \ - --revision 40d018ac1c1a2a4b9734f17fdb21f3d933c49a01 \ - --quiet) + --revision 40d018ac1c1a2a4b9734f17fdb21f3d933c49a01) DATASET_PATH="$BRIDGE_ROOT/sft_dataset_bridge" echo "DATASET_PATH=$DATASET_PATH" @@ -65,7 +70,7 @@ echo "DATASET_PATH=$DATASET_PATH" # 2. Wan2.2 VAE checkpoint. # ---------------------------------------------------------------------------- echo ">>> $(date '+%H:%M:%S') downloading Wan2.2_VAE.pth ..." -WAN_VAE_PATH=$(uvx hf@latest download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth --quiet) +WAN_VAE_PATH=$(_hf_download Wan-AI/Wan2.2-TI2V-5B Wan2.2_VAE.pth) echo "WAN_VAE_PATH=$WAN_VAE_PATH" # ---------------------------------------------------------------------------- @@ -77,9 +82,8 @@ echo "WAN_VAE_PATH=$WAN_VAE_PATH" # substring is present, and point `MODEL_PATH` at the symlink. # ---------------------------------------------------------------------------- echo ">>> $(date '+%H:%M:%S') downloading Qwen3-VL-8B-Instruct ..." -_HF_SNAP=$(uvx hf@latest download Qwen/Qwen3-VL-8B-Instruct \ - --revision 0c351dd01ed87e9c1b53cbc748cba10e6187ff3b \ - --quiet) +_HF_SNAP=$(_hf_download Qwen/Qwen3-VL-8B-Instruct \ + --revision 0c351dd01ed87e9c1b53cbc748cba10e6187ff3b) mkdir -p "$STAGE_DIR/Qwen" ln -sfn "$_HF_SNAP" "$STAGE_DIR/Qwen/Qwen3-VL-8B-Instruct" MODEL_PATH="$STAGE_DIR/Qwen/Qwen3-VL-8B-Instruct" diff --git a/tests/action_policy_regression_test.py b/tests/action_policy_regression_test.py index a2206512..3432fa04 100644 --- a/tests/action_policy_regression_test.py +++ b/tests/action_policy_regression_test.py @@ -19,11 +19,11 @@ Specs ----- -* ``action_policy_libero`` — ``examples/toml/sft_config/action_policy_libero_repro.toml``. +* ``action_policy_libero`` — ``examples/toml/sft_config/action_policy_libero_10_nano.toml``. Data auto-downloads: the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` (small LeRobot dir), cached across runs. Collapses the recipe's HSDP replicate 2 -> 1 to fit one 4-GPU node. -* ``action_policy_droid`` — ``examples/toml/sft_config/action_policy_droid_repro.toml``. +* ``action_policy_droid`` — ``examples/toml/sft_config/action_policy_droid_nano.toml``. The full DROID split is far too large to auto-download in CI, so this spec requires a pre-staged LOCAL copy via ``DROID_ROOT`` and SKIPS when unset. ``DROID_ROOT`` is the **versioned merged root** whose basename is a @@ -165,15 +165,21 @@ def _run(cmd: list[str], log_file: Path, extra_env: dict | None = None) -> tuple # --- input staging (shared with tests/nano_training_smoke_test.py) ----------- +def _hf_download(args: list[str], log_path: Path) -> tuple[int, str]: + """``uvx hf@latest download ``, retrying once with ``--refresh`` + (``@latest`` doesn't refresh dependency index metadata).""" + rc, out = _run(["uvx", "hf@latest", "download", *args], log_path) + if rc != 0: + rc, out = _run(["uvx", "--refresh", "hf@latest", "download", *args], log_path) + return rc, out + + def _ensure_wan_vae(log_dir: Path) -> None: """Download the Wan2.2 VAE if not already present.""" if _WAN_VAE.is_file(): return - rc, out = _run( + rc, out = _hf_download( [ - "uvx", - "hf@latest", - "download", "Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", "--local-dir", @@ -210,11 +216,8 @@ def _ensure_libero(log_dir: Path) -> None: """Download the ``libero_10`` suite of ``nvidia/LIBERO_LeRobot_v3`` if absent.""" if (_LIBERO_ROOT / "meta" / "info.json").is_file(): return - rc, out = _run( + rc, out = _hf_download( [ - "uvx", - "hf@latest", - "download", "--repo-type", "dataset", "nvidia/LIBERO_LeRobot_v3", @@ -283,12 +286,12 @@ class LaunchSpec: _SPECS: dict[str, LaunchSpec] = { "action_policy_libero": LaunchSpec( key="action_policy_libero", - sft_toml="examples/toml/sft_config/action_policy_libero_repro.toml", + sft_toml="examples/toml/sft_config/action_policy_libero_10_nano.toml", extra_hydra_args=_COMMON_OVERRIDES, ), "action_policy_droid": LaunchSpec( key="action_policy_droid", - sft_toml="examples/toml/sft_config/action_policy_droid_repro.toml", + sft_toml="examples/toml/sft_config/action_policy_droid_nano.toml", extra_hydra_args=_COMMON_OVERRIDES, requires_droid_root=True, ), diff --git a/tests/launch_regression_test.py b/tests/launch_regression_test.py index 7d1055ab..2f105c6d 100644 --- a/tests/launch_regression_test.py +++ b/tests/launch_regression_test.py @@ -109,13 +109,19 @@ def _free_port() -> int: def _hf_download(args: list[str]) -> str: - """``uvx hf download --quiet`` -> the local path it prints (from the HF cache).""" - result = subprocess.run( - ["uvx", "hf@latest", "download", *args, "--quiet"], - cwd=str(REPO_ROOT), - capture_output=True, - text=True, - ) + """``uvx hf download --quiet`` -> the local path it prints (from the HF cache). + + Retries once with ``--refresh``: ``@latest`` doesn't refresh dependency index metadata. + """ + for uvx in (["uvx"], ["uvx", "--refresh"]): + result = subprocess.run( + [*uvx, "hf@latest", "download", *args, "--quiet"], + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + ) + if result.returncode == 0: + break if result.returncode != 0: pytest.fail(f"hf download failed for {args} (exit {result.returncode}):\n{result.stdout}\n{result.stderr}") lines = [ln.strip() for ln in result.stdout.splitlines() if ln.strip()] @@ -288,6 +294,17 @@ def _build_specs(paths: dict[str, str]) -> dict[str, LaunchSpec]: key="vision_sft_nano", sft_toml="examples/toml/sft_config/vision_sft_nano.toml", extra_hydra_args=( + # Pin the learning rate to the value the goldens below were + # captured at (2e-5). The shipped recipe (vision_sft_nano.toml) + # was retuned to 1e-4, which changes the loss trajectory from + # iter 1 on (iter-0 is lr-independent). This override keeps the + # numerical-regression guard decoupled from recipe-lr tuning, so + # the committed h100 AND gb200 goldens stay valid without a GPU + # recapture. Trailing overrides win over the TOML value (see + # sft_config.load_experiment_from_toml). If you intentionally + # want CI to guard the 1e-4 trajectory instead, drop this line + # and refresh the goldens with COSMOS_REGRESSION_UPDATE_GOLDENS=1. + "optimizer.lr=2.0e-5", "model.config.parallelism.data_parallel_shard_degree=4", "model.config.compile.enabled=true", "trainer.max_iter=10", @@ -895,7 +912,10 @@ def test_launch_regression_8gpu(spec_key: str, tmp_path: Path, h100_inputs: dict # across all 10 iters. grad_norm is deterministic here (compile.enabled=false # in nano_model_config under the new release branch), so values are pinned; # flip to None if a future change re-enables compile and reintroduces - # non-determinism in the all-rank reduction. + # non-determinism in the all-rank reduction. Captured at lr=2e-5; the + # ``vision_sft_nano`` spec pins ``optimizer.lr=2.0e-5`` (see the pin + # comment in _build_specs) so this stays valid despite the shipped + # recipe now using 1e-4. "vision_sft_nano": { "loss": [0.2243, 0.2133, 0.2437, 0.2255, 0.2616, 0.2552, 0.3313, 0.2247, 0.2036, 0.2621], "grad_norm": [0.42188, 0.30469, 0.30078, 0.26953, 0.30273, 0.41406, 0.42773, 0.38477, 0.27344, 0.27344], @@ -918,7 +938,9 @@ def test_launch_regression_8gpu(spec_key: str, tmp_path: Path, h100_inputs: dict # defaults. Runs under ``--deterministic`` so loss reproduces bit-exact # across all 10 iters, but grad_norm is non-det because # ``compile.enabled=true`` makes the all-rank reduction not bit-exact - # on H100. + # on H100. Captured at lr=2e-5; the ``vision_sft_nano`` spec pins + # ``optimizer.lr=2.0e-5`` so these stay valid even though the shipped + # recipe now uses 1e-4 (see the pin comment in _build_specs). "vision_sft_nano": { "loss": [0.2242, 0.2141, 0.2429, 0.2259, 0.2608, 0.2555, 0.332, 0.2256, 0.2041, 0.2621], "grad_norm": None, diff --git a/tests/nano_training_smoke_test.py b/tests/nano_training_smoke_test.py index ab7ecf26..6bb1bed3 100644 --- a/tests/nano_training_smoke_test.py +++ b/tests/nano_training_smoke_test.py @@ -99,12 +99,21 @@ def _run(cmd: list[str], log_file: Path, extra_env: dict | None = None) -> tuple return returncode, "".join(captured) +def _hf_download(args: list[str], log_path: Path) -> tuple[int, str]: + """``uvx hf@latest download ``, retrying once with ``--refresh`` + (``@latest`` doesn't refresh dependency index metadata).""" + rc, out = _run(["uvx", "hf@latest", "download", *args], log_path) + if rc != 0: + rc, out = _run(["uvx", "--refresh", "hf@latest", "download", *args], log_path) + return rc, out + + def _ensure_inputs(log_dir: Path) -> None: """Step 1: download the dataset + Wan2.2 VAE if not already present.""" if not (_DATASET_PATH / "train" / "video_dataset_file.jsonl").is_file(): - rc, out = _run( + rc, out = _hf_download( [ - "uvx", "hf@latest", "download", "--repo-type", "dataset", + "--repo-type", "dataset", "nvidia/bridge-v2-subset-synthetic-captions", "--revision", _DATASET_REVISION, "--local-dir", str(_DATA_DIR), "--quiet", @@ -117,9 +126,9 @@ def _ensure_inputs(log_dir: Path) -> None: ) if not _WAN_VAE.is_file(): - rc, out = _run( + rc, out = _hf_download( [ - "uvx", "hf@latest", "download", "Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", + "Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", "--local-dir", str(_WAN_VAE.parent), "--quiet", ], log_dir / "download_wan_vae.log", @@ -255,6 +264,107 @@ def _assert_export_complete(model_dir: Path) -> None: assert names, f"export {model_dir}: model.safetensors holds no tensors" +def _assert_diffusers_complete(model_dir: Path, reference_dir: Path) -> None: + """Structural + index completeness of a Diffusers pipeline converted from the HF export, + and a tensor-level comparison against the published ``nvidia/Cosmos3-Nano`` diffusers. + + The ``vision_sft_nano`` export has no sound tokenizer (``sound_gen=False``) and no + standalone reasoner ViT (``include_visual`` unset), so the ``sound_tokenizer/`` and + ``vision_encoder/`` components — and the sound-only ``audio_*`` transformer tensors — + are absent; the golden comparison ignores exactly those. Every component that *is* + present is validated as thoroughly as the HF export: required files, pipeline class, + per-shard/per-tensor self-consistency of both the transformer index and the aggregated + root weight index, and (against the golden) the transformer tensor set + config + ``architectures`` / ``model_type``. + """ + assert model_dir.is_dir(), f"diffusers dir missing: {model_dir}" + required = ( + "config.json", + "model_index.json", + "modular_model_index.json", + "model.safetensors.index.json", + "scheduler/scheduler_config.json", + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors.index.json", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + ) + for rel in required: + p = model_dir / rel + assert p.is_file() and p.stat().st_size > 0, f"diffusers export missing/empty: {p}" + + model_index = json.loads((model_dir / "model_index.json").read_text()) + assert model_index.get("_class_name") == "Cosmos3OmniPipeline", ( + f"unexpected diffusers pipeline class: {model_index.get('_class_name')!r}" + ) + for component in ("scheduler", "transformer", "vae"): + assert component in model_index, f"model_index.json missing component {component!r}" + + # Root weight index: references existing shards and every declared tensor is + # actually stored across them (no omitted/extra param). + root_index = json.loads((model_dir / "model.safetensors.index.json").read_text()) + root_weight_map: dict[str, str] = root_index["weight_map"] + assert root_weight_map, "root Diffusers safetensors index has no tensors" + assert root_index.get("metadata", {}).get("total_size", 0) > 0 + root_shards = sorted(set(root_weight_map.values())) + missing_shards = [s for s in root_shards if not (model_dir / s).is_file()] + assert not missing_shards, f"root index references missing shards: {missing_shards}" + stored: set[str] = set() + for shard in root_shards: + stored |= _safetensors_tensor_names(model_dir / shard) + assert set(root_weight_map) == stored, ( + f"root index declares {len(root_weight_map)} tensors but shards hold {len(stored)} " + f"(missing from shards: {sorted(set(root_weight_map) - stored)[:10]}; " + f"not in index: {sorted(stored - set(root_weight_map))[:10]})" + ) + + # Transformer index: its weight_map is the transformer-scoped slice of the root + # index and points at exactly the transformer shards on disk. + transformer_index = json.loads( + (model_dir / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + ) + transformer_weight_map: dict[str, str] = transformer_index["weight_map"] + expected_transformer = { + name: filename.removeprefix("transformer/") + for name, filename in root_weight_map.items() + if filename.startswith("transformer/") + } + assert transformer_weight_map == expected_transformer, "transformer index inconsistent with root index" + assert set(transformer_weight_map.values()) == { + p.name for p in (model_dir / "transformer").glob("*.safetensors") + } + + # VAE weights hold tensors; no raw DCP shards leak into the Diffusers export. + assert _safetensors_tensor_names(model_dir / "vae/diffusion_pytorch_model.safetensors") + assert not list(model_dir.rglob("*.distcp")), "Diffusers export unexpectedly contains DCP shards" + + # Golden comparison against nvidia/Cosmos3-Nano: the transformer tensor set must equal + # the reference's, ignoring the sound-only ``audio_*`` tensors (this export has + # sound_gen=False) and the reference's vision_encoder/ shards (include_visual unset + # here). config architectures/model_type must match exactly. + reference_weight_map = json.loads((reference_dir / "model.safetensors.index.json").read_text())["weight_map"] + reference_transformer = { + name + for name, filename in reference_weight_map.items() + if filename.startswith("transformer/") and not name.startswith("audio_") + } + out_transformer = {name for name, filename in root_weight_map.items() if filename.startswith("transformer/")} + assert out_transformer == reference_transformer, ( + "transformer tensor set differs from nvidia/Cosmos3-Nano (ignoring sound/vision): " + f"missing={sorted(reference_transformer - out_transformer)[:8]}, " + f"extra={sorted(out_transformer - reference_transformer)[:8]}" + ) + reference_config = json.loads((reference_dir / "config.json").read_text()) + out_config = json.loads((model_dir / "config.json").read_text()) + assert out_config.get("architectures") == reference_config.get("architectures"), ( + f"config architectures differ from golden: {out_config.get('architectures')} vs " + f"{reference_config.get('architectures')}" + ) + assert out_config.get("model_type") == reference_config.get("model_type"), ( + f"config model_type differs from golden: {out_config.get('model_type')} vs {reference_config.get('model_type')}" + ) + + def _assert_valid_image(path: Path) -> None: """Assert ``path`` is a valid, non-degenerate image.""" assert path.is_file() and path.stat().st_size > 1024, f"output image missing/too small: {path}" @@ -342,6 +452,30 @@ def test_nano_sft_train_export_infer(tmp_path: Path) -> None: assert rc == 0, f"export_model failed (exit {rc}):\nLog tail:\n{out[-4000:]}" _assert_export_complete(export_dir) + # 4b. Convert the exported HF checkpoint to a Diffusers pipeline; check layout. + diffusers_dir = run_dir / "diffusers" + rc, out = _run( + [ + "python", "-m", "cosmos_framework.scripts.convert_model_to_diffusers", + "--checkpoint-path", str(export_dir), + "-o", str(diffusers_dir), + ], + tmp_path / "convert.log", + ) + assert rc == 0, f"convert_model_to_diffusers failed (exit {rc}):\nLog tail:\n{out[-4000:]}" + # Compare against the published Cosmos3-Nano diffusers (index + config only), + # ignoring the sound_tokenizer/ and vision_encoder/ components this export lacks. + golden_dir = tmp_path / "Cosmos3-Nano-ref" + rc, out = _hf_download( + [ + "nvidia/Cosmos3-Nano", "model.safetensors.index.json", "config.json", + "--local-dir", str(golden_dir), "--quiet", + ], + tmp_path / "download_golden.log", + ) + assert rc == 0, f"Cosmos3-Nano reference download failed (exit {rc}):\n{out[-2000:]}" + _assert_diffusers_complete(diffusers_dir, golden_dir) + # 5. t2i inference from the exported model; check the image is valid. infer_out = tmp_path / "exported_out" rc, out = _run(