Release 2026-07-11 (i4 7392c4b5)

.file_mapping.json

@@ -1,16 +1,15 @@
 {
-  "_source_commit": "c2c7152e59e4090a61822dc22a1691e7afb44702-dirty",
-  "_dest_commit": "463ac142b47c637098a7e86d0631252f82b65418",
-  "_generated_at": "2026-07-07T05:54:17Z",
+  "_source_commit": "7392c4b59ab7c1f0a94e59fbc4a75a0f3684b66f-dirty",
+  "_dest_commit": "3d9c0878fd0dde76eac98161aed0493d85a036fd",
+  "_generated_at": "2026-07-11T05:50:38Z",
   "files": {
     "imaginaire/__init__.py": "cosmos_framework/__init__.py",
     "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py",
     "imaginaire/attention/backends.py": "cosmos_framework/model/attention/backends.py",
-    "imaginaire/attention/benchmarks/benchmark_fmha_head_sharding.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha_head_sharding.py",
+    "imaginaire/attention/benchmarks/benchmark_fmha.py": "cosmos_framework/model/attention/benchmarks/benchmark_fmha.py",
     "imaginaire/attention/checks.py": "cosmos_framework/model/attention/checks.py",
     "imaginaire/attention/cudnn/__init__.py": "cosmos_framework/model/attention/cudnn/__init__.py",
     "imaginaire/attention/cudnn/checks.py": "cosmos_framework/model/attention/cudnn/checks.py",
-    "imaginaire/attention/cudnn/cudnn_forward.py": "cosmos_framework/model/attention/cudnn/cudnn_forward.py",
     "imaginaire/attention/cudnn/functions.py": "cosmos_framework/model/attention/cudnn/functions.py",
     "imaginaire/attention/cudnn/meta.py": "cosmos_framework/model/attention/cudnn/meta.py",
     "imaginaire/attention/cudnn/stubs.py": "cosmos_framework/model/attention/cudnn/stubs.py",
@@ -88,6 +87,7 @@
     "imaginaire/flops/__init__.py": "cosmos_framework/tools/flops/__init__.py",
     "imaginaire/flops/omni_mot.py": "cosmos_framework/tools/flops/omni_mot.py",
     "imaginaire/flops/qwen3_vl.py": "cosmos_framework/tools/flops/qwen3_vl.py",
+    "imaginaire/flops/test_qwen3_vl.py": "cosmos_framework/tools/flops/test_qwen3_vl.py",
     "imaginaire/flops/wan_vae.py": "cosmos_framework/tools/flops/wan_vae.py",
     "imaginaire/functional/batch_ops.py": "cosmos_framework/utils/functional/batch_ops.py",
     "imaginaire/functional/lr_scheduler.py": "cosmos_framework/utils/functional/lr_scheduler.py",
@@ -132,6 +132,7 @@
     "imaginaire/utils/easy_io/handlers/csv_handler.py": "cosmos_framework/utils/easy_io/handlers/csv_handler.py",
     "imaginaire/utils/easy_io/handlers/gzip_handler.py": "cosmos_framework/utils/easy_io/handlers/gzip_handler.py",
     "imaginaire/utils/easy_io/handlers/imageio_video_handler.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py",
+    "imaginaire/utils/easy_io/handlers/imageio_video_handler_test.py": "cosmos_framework/utils/easy_io/handlers/imageio_video_handler_test.py",
     "imaginaire/utils/easy_io/handlers/json_handler.py": "cosmos_framework/utils/easy_io/handlers/json_handler.py",
     "imaginaire/utils/easy_io/handlers/jsonl_handler.py": "cosmos_framework/utils/easy_io/handlers/jsonl_handler.py",
     "imaginaire/utils/easy_io/handlers/np_handler.py": "cosmos_framework/utils/easy_io/handlers/np_handler.py",
@@ -405,7 +406,9 @@
     "projects/cosmos3/cosmos3/tokenizers/wan2pt2_vae_4x16x16.py": "cosmos_framework/model/generator/tokenizers/wan2pt2_vae_4x16x16.py",
     "projects/cosmos3/cosmos3/upsampler/__init__.py": "cosmos_framework/model/generator/upsampler/__init__.py",
     "projects/cosmos3/cosmos3/upsampler/prompts.py": "cosmos_framework/model/generator/upsampler/prompts.py",
+    "projects/cosmos3/cosmos3/utils/aux_optimizer_utils.py": "cosmos_framework/utils/generator/aux_optimizer_utils.py",
     "projects/cosmos3/cosmos3/utils/data_utils.py": "cosmos_framework/utils/generator/data_utils.py",
+    "projects/cosmos3/cosmos3/utils/dion2_with_aux_adamw.py": "cosmos_framework/utils/generator/dion2_with_aux_adamw.py",
     "projects/cosmos3/cosmos3/utils/dtensor_helper.py": "cosmos_framework/utils/generator/dtensor_helper.py",
     "projects/cosmos3/cosmos3/utils/flash_attn.py": "cosmos_framework/utils/generator/flash_attn.py",
     "projects/cosmos3/cosmos3/utils/fused_adam.py": "cosmos_framework/utils/generator/fused_adam.py",
@@ -414,6 +417,7 @@
     "projects/cosmos3/cosmos3/utils/model_loader.py": "cosmos_framework/utils/generator/model_loader.py",
     "projects/cosmos3/cosmos3/utils/model_weights_stats.py": "cosmos_framework/utils/generator/model_weights_stats.py",
     "projects/cosmos3/cosmos3/utils/monkey_patch.py": "cosmos_framework/utils/generator/monkey_patch.py",
+    "projects/cosmos3/cosmos3/utils/muon_with_aux_adamw.py": "cosmos_framework/utils/generator/muon_with_aux_adamw.py",
     "projects/cosmos3/cosmos3/utils/optimizer.py": "cosmos_framework/utils/generator/optimizer.py",
     "projects/cosmos3/cosmos3/utils/parallelism.py": "cosmos_framework/utils/generator/parallelism.py",
     "projects/cosmos3/cosmos3/utils/rand_state.py": "cosmos_framework/utils/generator/rand_state.py",
@@ -423,9 +427,11 @@
     "projects/cosmos3/cosmos3/utils/reasoner/flop_calculator.py": "cosmos_framework/utils/generator/reasoner/flop_calculator.py",
     "projects/cosmos3/cosmos3/utils/reasoner/pretrained_models_downloader.py": "cosmos_framework/utils/generator/reasoner/pretrained_models_downloader.py",
     "projects/cosmos3/cosmos3/utils/video_preprocess.py": "cosmos_framework/utils/generator/video_preprocess.py",
+    "projects/cosmos3/tokenizer/checkpoint_io.py": "cosmos_framework/model/tokenizer/checkpoint_io.py",
     "projects/cosmos3/tokenizer/evaluation/metric.py": "cosmos_framework/model/tokenizer/evaluation/metric.py",
     "projects/cosmos3/tokenizer/evaluation/reconstruction_metrics.py": "cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py",
     "projects/cosmos3/tokenizer/models/__init__.py": "cosmos_framework/model/tokenizer/models/__init__.py",
+    "projects/cosmos3/tokenizer/models/architecture.py": "cosmos_framework/model/tokenizer/models/architecture.py",
     "projects/cosmos3/tokenizer/models/dense_backends.py": "cosmos_framework/model/tokenizer/models/dense_backends.py",
     "projects/cosmos3/tokenizer/models/dense_runtime.py": "cosmos_framework/model/tokenizer/models/dense_runtime.py",
     "projects/cosmos3/tokenizer/models/modules/__init__.py": "cosmos_framework/model/tokenizer/models/modules/__init__.py",
@@ -445,6 +451,7 @@
     "projects/cosmos3/tokenizer/models/text_decoder.py": "cosmos_framework/model/tokenizer/models/text_decoder.py",
     "projects/cosmos3/tokenizer/models/utils.py": "cosmos_framework/model/tokenizer/models/utils.py",
     "projects/cosmos3/tokenizer/utils/hf.py": "cosmos_framework/model/tokenizer/utils/hf.py",
+    "projects/cosmos3/tokenizer/utils/tensors.py": "cosmos_framework/model/tokenizer/utils/tensors.py",
     "projects/cosmos3/tokenizer/utils/vlm_prompt_format.py": "cosmos_framework/model/tokenizer/utils/vlm_prompt_format.py",
     "projects/cosmos3/utils/torchcodec_video.py": "cosmos_framework/utils/generator/torchcodec_video.py",
     "projects/cosmos3/vlm/configs/base/defaults/checkpointer.py": "cosmos_framework/utils/reasoner/configs_defaults/checkpointer.py",

cosmos_framework/callbacks/every_n_draw_sample.py

@@ -217,6 +217,91 @@ def _add_wandb_image_paths(
     info[key_prefix] = wandb.Image(image_paths, caption=caption)
 
 
+def _pixel_tensor_to_5d(t: torch.Tensor) -> torch.Tensor:
+    """Ensure a pixel tensor has shape (B, C, T, H, W) for the visualization grid.
+
+    Handles (C, H, W), (B, C, H, W), and (B, C, T, H, W) inputs.
+    """
+    if t.ndim == 3:
+        return t.unsqueeze(0).unsqueeze(2)  # (C,H,W) -> (1,C,1,H,W)
+    if t.ndim == 4:
+        return t.unsqueeze(2)  # (B,C,H,W) -> (B,C,1,H,W)
+    return t
+
+
+def _resize_5d_to_width(img5d: torch.Tensor, target_width: int) -> torch.Tensor:
+    """Resize a single-frame (1, C, 1, H, W) tensor so its width is exactly ``target_width``.
+
+    Height is scaled proportionally to preserve the overall aspect ratio. Assumes a
+    single temporal frame (T == 1).
+    """
+    h, w = img5d.shape[-2], img5d.shape[-1]
+    if w == target_width:
+        return img5d
+    new_h = max(1, round(h * target_width / w))
+    chw = img5d[0, :, 0]  # (C,H,W)
+    resized = torchvision_F.resize(chw, [new_h, target_width], antialias=True)  # (C,new_h,target_width)
+    return resized.unsqueeze(0).unsqueeze(2)  # (1,C,1,new_h,target_width)
+
+
+def _resize_pad_to_square(img5d: torch.Tensor, cell: int) -> torch.Tensor:
+    """Resize a single-frame (1, C, 1, H, W) tensor to fit inside a ``cell`` x ``cell`` square.
+
+    Aspect ratio is preserved (the image is scaled so its longer side equals ``cell``), then
+    the result is center-padded with zeros to exactly ``cell`` x ``cell``. Assumes T == 1.
+    """
+    h, w = img5d.shape[-2], img5d.shape[-1]
+    scale = cell / max(h, w)
+    new_h = max(1, round(h * scale))
+    new_w = max(1, round(w * scale))
+    chw = img5d[0, :, 0]  # (C,H,W)
+    resized = torchvision_F.resize(chw, [new_h, new_w], antialias=True)  # (C,new_h,new_w)
+    pad_h = cell - new_h
+    pad_w = cell - new_w
+    top = pad_h // 2
+    left = pad_w // 2
+    padded = torch.nn.functional.pad(resized, (left, pad_w - left, top, pad_h - top))  # (C,cell,cell)
+    return padded.unsqueeze(0).unsqueeze(2)  # (1,C,1,cell,cell)
+
+
+def _build_reference_grid(references: list[torch.Tensor], target_width: int) -> torch.Tensor:
+    """Tile reference images into a compact near-square grid sized to ``target_width``.
+
+    All references are arranged in a ``rows`` x ``cols`` grid (``cols = ceil(sqrt(n))``), each in
+    an aspect-preserved square cell, then the whole grid is resized so its width equals
+    ``target_width``. This keeps the condition montage aligned with the generated-image column
+    width (better space utilization) instead of one overly-long row.
+
+    Args:
+        references: list of single-frame pixel tensors (each (1,C,1,H,W) or (C,H,W)), in [-1, 1].
+        target_width: width of the generated/target image for this sample.
+
+    Returns:
+        A (1, C, 1, H_grid, target_width) tensor.
+    """
+    if not references:
+        raise ValueError("Expected at least one reference image to build the condition grid.")
+
+    refs = [_pixel_tensor_to_5d(r) for r in references]  # list[(1,C,1,H,W)]
+    n = len(refs)
+    cols = math.ceil(math.sqrt(n))
+    rows = math.ceil(n / cols)
+    cell = max(1, target_width // cols)
+
+    cells = [_resize_pad_to_square(r, cell) for r in refs]  # list[(1,C,1,cell,cell)]
+    blank = torch.zeros_like(cells[0])
+    while len(cells) < rows * cols:
+        cells.append(blank)
+
+    row_imgs = []
+    for r in range(rows):
+        row = torch.cat(cells[r * cols : (r + 1) * cols], dim=-1)  # (1,C,1,cell,cols*cell)
+        row_imgs.append(row)
+    grid = torch.cat(row_imgs, dim=-2)  # (1,C,1,rows*cell,cols*cell)
+
+    return _resize_5d_to_width(grid, target_width)  # (1,C,1,H_grid,target_width)
+
+
 class EveryNDrawSample(EveryN):
     """
     This callback sample condition inputs from training data, run inference and save the results to wandb and s3.
@@ -531,11 +616,20 @@ def sample(
             vis_offset = 0
             for sample_idx in range(data_clean.batch_size):
                 n_vis = num_items[sample_idx]
-                # First item(s) are condition, last item is generation target
-                # but we need to support multiple conditions per sample in the future. Current code
-                # can handle this without throwing an error.
-                condition_images.append(raw_data[vis_offset])  # source image (1, C, 1, H, W)
-                gt_target_images.append(raw_data[vis_offset + n_vis - 1])  # target image (1, C, 1, H, W)
+                # First item(s) are condition references, last item is the generation target.
+                refs = raw_data[vis_offset : vis_offset + n_vis - 1]  # all condition items
+                target = raw_data[vis_offset + n_vis - 1]  # target image (1, C, 1, H, W)
+                # Multi-reference generation (>1 single-frame image references): tile every
+                # reference into a compact grid resized to the target width, so all references
+                # are visible without blowing up the row width. For video editing/transfer
+                # (T > 1) keep the existing behavior (first item = condition) so those tasks
+                # render exactly as before and stay consistent with the t_crop frame cropping.
+                refs_are_images = all(r.shape[-3] == 1 for r in refs) and target.shape[-3] == 1
+                if refs_are_images and len(refs) > 1:
+                    condition_images.append(_build_reference_grid(refs, target.shape[-1]))
+                else:
+                    condition_images.append(raw_data[vis_offset])  # source image (1, C, 1, H, W) / video clip
+                gt_target_images.append(target)
                 vis_offset += n_vis
 
             # Use target images for max_w/max_h/t_crop (generated samples match target size)

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
 
@@ -866,16 +866,6 @@ def load(
                         raise ValueError(
                             f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}"
                         )
-                    # Warm start that skipped net_ema (e.g. loading an EMA-only HF export
-                    # with no net_ema.* keys): the EMA shadow would otherwise keep its random
-                    # build-time generation pathway (init_moe is skipped when a checkpoint is
-                    # present). Seed net_ema from the freshly loaded net so the EMA starts equal
-                    # to net ("EMA warm-starts from net") instead of from random weights.
-                    if warm_start and any("net_ema" in skip_key for skip_key in keys_to_skip_loading):
-                        ema_worker = getattr(model, "net_ema_worker", None)
-                        if ema_worker is not None and getattr(model, "net_ema", None) is not None:
-                            ema_worker.copy_to(src_model=model.net, tgt_model=model.net_ema)
-                            log.info("Warm start: re-seeded net_ema from net (net_ema was skipped on load).")
 
                 elif key == "optim":
                     log.info("- Loading the optimizer...")

cosmos_framework/configs/base/defaults/checkpointer.py

@@ -98,6 +98,7 @@
     strict_resume=True,
     enable_gcs_patch_in_boto3=True,
     dcp_async_mode_enabled=True,
+    dcp_load_dedup=True,
 )
 
 CHECKPOINT_S3_EAST2 = CheckpointConfig(

cosmos_framework/configs/base/defaults/model_config.py

@@ -5,12 +5,12 @@
 
 import attrs
 
+from cosmos_framework.utils.lazy_config import LazyDict
 from cosmos_framework.configs.base.defaults.activation_checkpointing import ActivationCheckpointingConfig
 from cosmos_framework.configs.base.defaults.compile import CompileConfig
 from cosmos_framework.configs.base.defaults.ema import EMAConfig
 from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
 from cosmos_framework.configs.base.defaults.reasoner import VLMConfig
-from cosmos_framework.utils.lazy_config import LazyDict
 
 
 @attrs.define(slots=False)

cosmos_framework/configs/base/defaults/optimizer.py

@@ -31,6 +31,57 @@
     disable_weight_decay_for_1d_params=False,
 )
 
+# Muon / Dion2 share the standard factory knobs (keys_to_select, lr_multipliers,
+# disable_weight_decay_for_1d_params) plus their own orthogonalization
+# hyperparameters. ``fused`` is required by the factory; the AdamW side is fused
+# by construction and ``capturable`` / ``master_weights`` are forced on.
+MUON_OPTIMIZER_KWARGS: dict[str, Any] = dict(
+    # Base learning rate. Muon scales matrix params by muon_lr_scale*sqrt(max(A,B));
+    # the AdamW side and the per-param-group lr_multipliers use it directly.
+    lr=1e-4,
+    weight_decay=0.1,
+    adam_betas=[0.9, 0.99],
+    eps=1e-8,
+    fused=True,
+    keys_to_select=[],
+    lr_multipliers={},
+    disable_weight_decay_for_1d_params=False,
+    # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize
+    # per expert slice. Empty = experts stay on AdamW (no behavior change).
+    # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts.
+    expert_param_keywords=[],
+    # Muon-specific.
+    muon_momentum=0.95,
+    muon_lr_scale=0.2,
+    ns_steps=5,
+    nesterov=True,
+    use_distributed=True,
+)
+
+DION2_OPTIMIZER_KWARGS: dict[str, Any] = dict(
+    lr=1e-4,
+    weight_decay=0.1,
+    adam_betas=[0.9, 0.99],
+    eps=1e-8,
+    fused=True,
+    keys_to_select=[],
+    lr_multipliers={},
+    disable_weight_decay_for_1d_params=False,
+    # Name substrings for stacked MoE expert params ([E, M, N]) to orthogonalize
+    # per expert slice. Empty = experts stay on AdamW (no behavior change).
+    # e.g. ["gate_up_proj", "down_proj"] for grouped-MM MoE experts.
+    expert_param_keywords=[],
+    # Muon/Dion2-specific.
+    muon_momentum=0.95,
+    muon_lr_scale=0.2,
+    ns_steps=5,
+    nesterov=True,
+    use_distributed=True,
+    # Dion2-specific: submatrix selection fraction and error-feedback decay.
+    fraction=1.0,
+    ef_decay=0.95,
+)
+
 LAMBDACOSINE_KWARGS: dict[str, Any] = dict(
     warm_up_steps=[2000],
     cycle_lengths=[100000],
@@ -64,6 +115,26 @@ def register_optimizers(optimizer_kwargs: dict[str, Any]) -> None:
             **optimizer_kwargs,
         ),
     )
+    cs.store(
+        group="optimizer",
+        package="optimizer",
+        name="muonwithauxadamw",
+        node=L(build_optimizer)(
+            model=PLACEHOLDER,
+            optimizer_type="MuonWithAuxAdamW",
+            **MUON_OPTIMIZER_KWARGS,
+        ),
+    )
+    cs.store(
+        group="optimizer",
+        package="optimizer",
+        name="dion2withauxadamw",
+        node=L(build_optimizer)(
+            model=PLACEHOLDER,
+            optimizer_type="Dion2WithAuxAdamW",
+            **DION2_OPTIMIZER_KWARGS,
+        ),
+    )
 
 
 def register_schedulers(lambdacosine_kwargs: dict[str, Any]) -> None:

cosmos_framework/configs/base/defaults/reasoner.py

@@ -785,6 +785,32 @@ class VLMConfig:
     ),
 )
 
+# Same as 590c1c0 but with use_und_k_norm_for_gen=True: normalises K_und before
+# it is used as a key in the gen→und cross-attention path (the qk-norm fix for the
+# generator). Adds a freshly-initialised k_norm_und_for_gen RMSNorm.
+Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm: VLMConfig = VLMConfig(
+    model_name="nvidia/Cosmos3-Edge-Reasoner",
+    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"
+            ),
+            qk_norm_for_text=False,
+            use_und_k_norm_for_gen=True,
+        ),
+    ),
+    tokenizer=L(build_processor_lazy)(
+        tokenizer_type="nvidia/Cosmos3-Edge-Reasoner",
+        config_variant="gcp",
+    ),
+    pretrained_weights=PretrainedWeightsConfig(
+        backbone_path="s3://bucket0/cosmos3/pretrained/huggingface/nvidia/Cosmos3-Edge-Reasoner-590c1c0/",
+        credentials_path="credentials/gcp_checkpoint.secret",
+        enable_gcs_patch_in_boto3=True,
+        checkpoint_format="nemotron_3_dense_vl",
+    ),
+)
+
 # Same as 9b4c028 but with use_und_k_norm_for_gen=True: normalises K_und before
 # it is used as a key in the gen→und cross-attention path.
 Cosmos3EdgeReasoner_VLM_GCP_Config_9b4c028_UndKNorm: VLMConfig = VLMConfig(
@@ -987,3 +1013,9 @@ def register_vlm():
         name="cosmos3_edge_reasoner_vlm_gcp_590c1c0",
         node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0,
     )
+    cs.store(
+        group="vlm_config",
+        package="model.config.vlm_config",
+        name="cosmos3_edge_reasoner_vlm_gcp_590c1c0_und_k_norm",
+        node=Cosmos3EdgeReasoner_VLM_GCP_Config_590c1c0_UndKNorm,
+    )