Sync private main updates into public main

cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py

@@ -5,7 +5,9 @@
 
 Consumers must ``copy.deepcopy`` this constant before mutating it. Baseline
 mirrors ``vision_sft_nano`` (HF-cluster deployment with empty tokenizer/vlm
-paths, video-style loss scales, ``load_weights_from_pretrained=True``).
+paths, video-style loss scales, ``load_weights_from_pretrained=True``), except
+``action_gen``: the baseline keeps the released-checkpoint value (``True``) and
+``vision_sft_nano`` overrides it to ``False`` (no action tokens in vision SFT).
 """
 
 from cosmos_framework.configs.base.defaults.reasoner import (
@@ -16,6 +18,9 @@
 from cosmos_framework.utils.lazy_config import LazyCall as L
 
 NANO_MODEL_CONFIG = dict(
+    # Mirrors the released Cosmos3-Nano checkpoint, which ships real (DROID-trained)
+    # action-head weights; the action-policy recipes rely on this default. Recipes
+    # that don't train action data override to False (see vision_sft_nano.py).
     action_gen=True,
     causal_training_strategy="none",
     input_caption_key="ai_caption",

cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py

@@ -35,23 +35,31 @@
 from hydra.core.config_store import ConfigStore
 
 from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG
-from cosmos_framework.data.generator.joint_dataloader import (
-    PackingDataLoader,
-    RankPartitionedDataLoader,
-)
 from cosmos_framework.data.generator.dataflow import (
     CosmosDataLoader,
     IdentityProcessor,
     RankPartitionedDistributor,
     SequentialPackingBatcher,
     VFMListCollator,
 )
+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,
+# matching vision_sft_edge). Unlike Edge, the Nano checkpoint ships real
+# (DROID-trained) action weights; with action_gen=False they are no longer carried
+# through vision-SFT'd checkpoints or exports — start from the base checkpoint /
+# action-policy recipes if you need them.
+_NANO_VISION_MODEL_CONFIG = copy.deepcopy(NANO_MODEL_CONFIG)
+_NANO_VISION_MODEL_CONFIG["action_gen"] = False
+
 
 vision_sft_nano = LazyDict(
     dict(
@@ -89,7 +97,7 @@
             wandb_mode="disabled",
         ),
         model=dict(
-            config=copy.deepcopy(NANO_MODEL_CONFIG),
+            config=copy.deepcopy(_NANO_VISION_MODEL_CONFIG),
         ),
         optimizer=dict(
             betas=[0.9, 0.95],

cosmos_framework/configs/base/experiment/sft/vision_sft_super.py

@@ -11,7 +11,7 @@
   * LoRA-only fine-tune: ``lora_enabled=True``, ``lora_rank=16``,
     ``lora_alpha=32``, target modules
     ``q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen``.
-  * EMA disabled; ``action_gen=False``.
+  * EMA disabled.
   * Parallelism: ``data_parallel_shard_degree=4``,
     ``context_parallel_shard_degree=2``, ``compile.enabled=False``.
   * Optimizer trains only ``lora_`` keys at ``lr=5e-4``.

tests/nano_training_smoke_test.py

@@ -268,10 +268,12 @@ 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*
+    The ``vision_sft_nano`` export has no sound tokenizer (``sound_gen=False``), no
+    standalone reasoner ViT (``include_visual`` unset), and no action heads
+    (``action_gen=False`` — vision SFT trains no action tokens), so the
+    ``sound_tokenizer/`` and ``vision_encoder/`` components — and the sound-only
+    ``audio_*`` / action-only ``action_*`` 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
@@ -340,17 +342,18 @@ def _assert_diffusers_complete(model_dir: Path, reference_dir: Path) -> None:
 
     # 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.
+    # sound_gen=False), the action-only ``action_*`` tensors (action_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_")
+        if filename.startswith("transformer/") and not name.startswith(("audio_", "action_"))
     }
     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): "
+        "transformer tensor set differs from nvidia/Cosmos3-Nano (ignoring sound/action/vision): "
         f"missing={sorted(reference_transformer - out_transformer)[:8]}, "
         f"extra={sorted(out_transformer - reference_transformer)[:8]}"
     )