Sync origin/main into kunlun (through #105)

cosmos_framework/callbacks/grad_clip.py

@@ -12,15 +12,17 @@
 from cosmos_framework.utils.callback import Callback
 
 
-@torch.compile
 def _fused_nan_to_num(grads: list[torch.Tensor]) -> None:
     """Replace NaN/Inf entries with 0.0 in every floating-point grad in-place.
 
-    The Python-level loop over ``grads`` is wrapped in ``@torch.compile`` so
-    Inductor can fuse the per-tensor ``nan_to_num`` ops into a single CUDA
-    kernel.  This is NOT the ``torch._foreach_*`` API; it is fusion-via-compile
-    and depends on the grad list structure (length, dtypes, shapes) staying
-    stable across iterations so Dynamo can reuse its specialized graph.
+    Runs eager, NOT ``@torch.compile``. Compiling this generates a GPU-only Triton
+    ``nan_to_num`` kernel, which crashes whenever any grad in the list is a CPU tensor (some
+    parameters carry a small CPU grad): the static CUDA launcher launches it with ``stream=0``
+    -> ``CUDA driver error: invalid argument``, and the standard launcher raises ``Pointer
+    argument cannot be accessed from Triton (cpu tensor?)``. Both were observed at 720 under
+    replay-TF + torch.compile (grad-clip runs at the optimizer step, after the large compiled
+    graphs are cached). Eager ``torch.nan_to_num`` handles CPU and CUDA grads alike; the fusion win
+    from compiling a handful of per-tensor ops once per step is negligible next to that fragility.
     """
     grads = [g for g in grads if torch.is_floating_point(g)]
     for g in grads:

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

@@ -118,8 +118,8 @@ class FixedStepSamplerConfig:
     # Convention: exclude the final 0.0 step — FixedStepSampler appends it automatically.
     # Values must be descending. Using 0.999 instead of 1.0 avoids numeric edge cases at sigma=1.
     t_list: list[float] = [0.999, 0.75, 0.5, 0.25]
-    # Integrator type: "ode" (deterministic Euler) or "sde" (stochastic re-noising at each step).
-    sample_type: str = "ode"
+    # Distilled fixed-step sampling uses stochastic re-noising at each step.
+    sample_type: str = "sde"
 
 
 # Don't have any defaults and init only in config file.

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,
+    )

cosmos_framework/configs/base/defaults/tokenizer.py

@@ -128,6 +128,7 @@
     vae_path=PRETRAINED_TOKENIZER_UNIAE_4X16X16_C48_T16TO160_MIXP_FPS_MIX_ENCODER_NONCAUSAL_DECODER_NONCAUSAL_NOGAN_S3_NEMOTRON2B_VAE_PTH,
     spatial_compression_factor=16,
     temporal_compression_factor=4,
+    pad_frames=1,
     pixel_trim=True,
     causal=False,
 )

cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py

@@ -21,15 +21,14 @@
 
 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.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG
+from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset
 from cosmos_framework.data.generator.joint_dataloader import (
     PackingDataLoader,
     RankPartitionedDataLoader,
 )
-from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_droid_sft_dataset
+from cosmos_framework.utils.lazy_config import LazyCall as L
+from cosmos_framework.utils.lazy_config import LazyDict
 
 cs = ConfigStore.instance()
 
@@ -202,7 +201,9 @@
                             use_state=True,
                             iterable_shuffle=True,  # rank x worker episode-shuffle stream
                             episode_shuffle_seed=42,
-                            use_image_augmentation=True,  # SR boost (random crop+rescale + color jitter)
+                            # SR boost: random crop+rescale + ColorJitter, applied CPU-side in the
+                            # DROIDLeRobotDataset image augmentor (matches i4's pipeline stage).
+                            use_image_augmentation=True,
                             # keep_ranges_1_0_1.json window filter (drops idle/non-task frames). Off by default;
                             # set use_filter_dict=True + filter_dict_path to enable.
                             use_filter_dict=False,
@@ -213,6 +214,9 @@
                             max_action_dim="${model.config.max_action_dim}",
                             cfg_dropout_rate=0.1,
                             tokenizer_config="${model.config.vlm_config.tokenizer}",
+                            # Match i4 GA (droid_lerobot_8b_ga / MR #9995): format the action
+                            # prompt as JSON via ActionPromptJsonFormatter instead of plain text.
+                            format_prompt_as_json=True,
                         ),
                     ),
                 ),

cosmos_framework/configs/base/reasoner/config.py

@@ -7,6 +7,7 @@
 from cosmos_framework.configs.base.defaults.checkpointer import register_checkpoint, register_ckpt_type
 from cosmos_framework.configs.base.reasoner.defaults.callbacks import register_callbacks
 from cosmos_framework.configs.base.reasoner.defaults.config import Config
+
 from cosmos_framework.configs.base.reasoner.defaults.model import register_model
 from cosmos_framework.configs.base.reasoner.defaults.optimizer import register_optimizer, register_scheduler
 from cosmos_framework.configs.base.reasoner.defaults.vlm_policy import register_vlm_policy

cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py

@@ -160,11 +160,33 @@ def _dl(dataset_key, split, num_workers, persistent_workers=False, pin_memory=Fa
             grad_accum_iter=8,
         ),
         optimizer=dict(
+            # Uniform 1e-6 across both tiers. Nano is stable at 1e-5 too, but the 32B
+            # super run spikes transiently (loss->6.6 @ iter 23) at 1e-5 before
+            # recovering to the same endpoint; 1e-6 reaches the identical final loss
+            # with no instability, so both tiers ship at 1e-6. (Super deepcopies this
+            # block, so setting it here covers both.)
             lr=1e-6,
             fused=True,
             weight_decay=0.05,
             betas=[0.9, 0.999],
-            lr_multipliers={"mm_projector": 20.0, "merger": 20.0},
+            # NOTE: the ViT is FROZEN (freeze_vision_encoder=True below), so the only
+            # trainable "model.visual.*" params are the projector (model.visual.merger
+            # + deepstack_merger_list). This key therefore sets the PROJECTOR LR only,
+            # to 1.0x optimizer.lr -- giving a uniform LR across projector + LLM +
+            # lm_head. It is NOT touching the frozen ViT backbone.
+            #
+            # Why the key is "model.visual" (broad) and not "model.visual.merger":
+            # the reasoner default optimizer (configs/base/reasoner/defaults/optimizer.py)
+            # ships lr_multipliers={"model.visual": 0.1}. Hydra merges dicts base-first
+            # (so that key stays PINNED FIRST) and _build_params_with_metadata
+            # (utils/generator/optimizer.py) is first-substring-match-wins -- so a
+            # narrower "model.visual.merger" key is shadowed by "model.visual" and never
+            # matches (verified: the merger still resolves to 0.1x). Overriding the SAME
+            # "model.visual" key to 1.0 is the only way to cancel the 0.1x default; with
+            # the ViT frozen that reaches only the merger. (To give an UNFROZEN ViT its
+            # own LR you'd likewise override this "model.visual" value -- a narrower key
+            # can't out-prioritize it.)
+            lr_multipliers={"model.visual": 1.0},
         ),
         scheduler=dict(
             warm_up_steps=[5],

cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_super.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""VideoPhy-2 SFT recipe (Cosmos3-Super tier): Qwen3-VL-32B full fine-tune.
+
+Super-tier counterpart of ``videophy2_sft_nano``. Reuses that recipe's VideoPhy-2
+``LocalSFTDataset`` + ``CosmosDataLoader`` dataflow verbatim (imported + deepcopied,
+the same variant idiom ``llava_ov_vlm.py`` uses) and changes only what the larger
+backbone needs:
+
+  * ``vlm_policy`` ``qwen3_vl_8b_instruct`` -> ``qwen3_vl_32b_instruct``
+    (``Qwen/Qwen3-VL-32B-Instruct``) — the visual tower + config the Cosmos3-Super
+    Reasoner is built on, mirroring how the nano recipe rides the 8B tower.
+  * FSDP full-shard across every rank (``data_parallel_shard_degree=-1``) so the 32B
+    weights + optimizer state fit, instead of the nano recipe's fixed dp=8. This is
+    the same super-tier sharding switch ``vision_sft_super`` makes, and lets the recipe
+    run unchanged on a 4-GPU (e.g. GB200x4) or 8-GPU allocation.
+
+Still a full fine-tune (no LoRA): the freeze config is inherited from the nano recipe
+(vision encoder frozen, LM + mm_projector trained).
+
+Launch via ``examples/launch_sft_videophy2_super.sh`` after
+``prepare_videophy2_from_hf`` populates ``$VIDEOPHYSICS_ROOT``, supplying the merged
+Cosmos3-Super Reasoner checkpoint through ``VLM_SAFETENSORS_PATH`` (see the launch shell).
+"""
+
+from __future__ import annotations
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+# Importing the nano module registers `videophy2_sft_nano` and pulls in the shared
+# VideoPhy-2 dataflow helpers; we clone its LazyDict rather than re-declaring them.
+from cosmos_framework.configs.base.reasoner.experiment.videophy2_sft_nano import videophy2_sft_nano
+
+cs = ConfigStore.instance()
+
+
+videophy2_sft_super = copy.deepcopy(videophy2_sft_nano)
+
+# Backbone: nano 8B -> super 32B. The vlm_policy override lives in the Hydra
+# `defaults` list; rewrite that one entry's value in place, leaving the rest of the
+# recipe (checkpoint backend, callbacks, dataflow) untouched.
+for _default in videophy2_sft_super["defaults"]:
+    if not isinstance(_default, str) and "override /vlm_policy" in _default:
+        _default["override /vlm_policy"] = "qwen3_vl_32b_instruct"
+
+# 32B full fine-tune: shard model + optimizer state across every rank (FSDP full
+# shard, auto-sized from WORLD_SIZE) instead of the nano recipe's fixed dp_shard=8,
+# so the recipe fits on 4- or 8-GPU nodes.
+# examples/toml/sft_config/videophy2_sft_super.toml is authoritative at launch and
+# repeats these; keeping them here lets `experiment=videophy2_sft_super` run standalone.
+videophy2_sft_super.model.config.parallelism.data_parallel_shard_degree = -1
+videophy2_sft_super.model.config.parallelism.data_parallel_replicate_degree = 1
+
+
+for _item in [videophy2_sft_super]:
+    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)

cosmos_framework/configs/toml_config/sft_config.py

@@ -83,6 +83,17 @@ class JobConfig(BaseModel):
             "'disabled' (no wandb at all)."
         ),
     )
+    upload_reproducible_setup: bool = Field(
+        default=False,
+        description=(
+            "Upload the reproducible-setup bundle and wandb save_s3 artifacts "
+            "to S3. Defaults False (OSS: no S3 access). Set True only when S3 "
+            "upload is configured. Remapped out of [job] to the top-level "
+            "config.upload_reproducible_setup, overriding the base config value "
+            "(the VLM base defaults it True). Always emitted by "
+            "load_experiment_from_toml, so omitting it forces False."
+        ),
+    )
 
 
 # ---------------------------------------------------------------- model
@@ -721,7 +732,14 @@ def load_experiment_from_toml(
 
     # Validate structure against the pydantic schema (raises ValidationError on
     # unknown keys because of ``extra="forbid"``).
-    SFTExperimentConfig.model_validate(raw)
+    cfg = SFTExperimentConfig.model_validate(raw)
+
+    # ``build_hydra_overrides`` walks the *raw* dict, so an omitted field emits
+    # no override and the base config's value wins. ``upload_reproducible_setup``
+    # must instead default to False (OSS: no S3), overriding the VLM base's True.
+    # Inject the pydantic-resolved value back into raw so the override is always
+    # emitted — the default lives in exactly one place (JobConfig.Field).
+    raw.setdefault("job", {})["upload_reproducible_setup"] = cfg.job.upload_reproducible_setup
 
     task = raw.get("job", {}).get("task", "vfm")
     try:

cosmos_framework/configs/toml_config/toml_config_helper.py

@@ -47,6 +47,9 @@
     # ``model.config.<X>`` in the Hydra tree. ``attn_implementation`` is a
     # VLM-only knob — skip it on VFM. Other sections pass through.
     "vfm": {
+        # [job].upload_reproducible_setup lives at the top-level config field,
+        # not config.job.* — hoist it out of the job section.
+        ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",),
         ("model", "attn_implementation"): None,
         ("model", "backbone"): None,                                           # VLM-only — VFM has no model.config.backbone
         # Per-caption token cap lives on the nested SFT dataset, not a top-level
@@ -63,6 +66,9 @@
     # tasks — so the catch-all ``("model",) -> ("model", "config")`` rule
     # routes them uniformly. Fields with no VLM analog map to ``None`` (skip).
     "vlm": {
+        # [job].upload_reproducible_setup lives at the top-level config field,
+        # not config.job.* — hoist it out of the job section.
+        ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",),
         # No VLM analog — skip these leaves
         ("model", "max_num_tokens_after_packing"): None,
         ("model", "joint_attn_implementation"): None,