diff --git a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py b/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py index d7e59206..01cda15e 100644 --- a/cosmos_framework/configs/base/experiment/sft/models/nano_model_config.py +++ b/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", diff --git a/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py b/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py index 6597ffb2..807b11cc 100644 --- a/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py +++ b/cosmos_framework/configs/base/experiment/sft/vision_sft_nano.py @@ -35,10 +35,6 @@ 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, @@ -46,12 +42,24 @@ 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], diff --git a/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py b/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py index 7c18dfbf..1e5ca8d3 100644 --- a/cosmos_framework/configs/base/experiment/sft/vision_sft_super.py +++ b/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``. diff --git a/tests/nano_training_smoke_test.py b/tests/nano_training_smoke_test.py index 6bb1bed3..aea3729a 100644 --- a/tests/nano_training_smoke_test.py +++ b/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]}" )