[cosmos3] Bitwise train/rollout parity: FSDP precision spec + cosmos3_bitwise rollout patch group

miles/backends/fsdp_utils/configs/cosmos3.py

@@ -5,6 +5,7 @@
 import math
 
 import torch
+from miles.backends.sglang_diffusion_utils.monkey_patches import patch_cosmos3_bitwise
 from miles.utils.types import CondKwargs
 
 from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config
@@ -49,6 +50,9 @@ def validate_args(cls, args) -> None:
         if list(args.update_weight_target_modules) != ["transformer"]:
             raise ValueError("Cosmos3 requires --update-weight-target-module transformer.")
 
+    def configure(self, args) -> None:
+        self._bitwise_parity = "cosmos3_bitwise" in args.rollout_patch_groups
+
     def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict:
         if cond is None or cond.text_ids is None:
             return {}
@@ -176,3 +180,5 @@ def _cast_to_weight_dtype(module, args):
             return tuple(a.to(dtype) if torch.is_tensor(a) else a for a in args)
 
         model.time_embedder.register_forward_pre_hook(_cast_to_weight_dtype)
+        if self._bitwise_parity:
+            patch_cosmos3_bitwise.apply_train(model)

miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py

@@ -1,4 +1,6 @@
 from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan
 
 
-FSDP_PARALLEL_PLAN = FSDPParallelPlan()
+FSDP_PARALLEL_PLAN = FSDPParallelPlan(
+    param_dtype_patterns={"*time_embedder*": "fp32"},
+)

miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py

@@ -41,6 +41,13 @@ def apply_qwen_image_rollout_patches() -> None:
     patch_qwen_image.apply()
 
 
+@register_rollout_patch_group("cosmos3_bitwise")
+def apply_cosmos3_bitwise_patches() -> None:
+    from miles.backends.sglang_diffusion_utils.monkey_patches import patch_cosmos3_bitwise
+
+    patch_cosmos3_bitwise.apply()
+
+
 @register_rollout_patch_group("wan")
 def apply_wan_rollout_patches() -> None:
     from miles.backends.sglang_diffusion_utils.monkey_patches import patch_wan_norm_ops

miles/utils/arguments.py

@@ -1508,11 +1508,12 @@ def miles_validate_args(args):
         from miles.backends.sglang_diffusion_utils.monkey_patches import validate_rollout_patch_groups
 
         validate_rollout_patch_groups(args.rollout_patch_groups)
-        if args.use_lora and "qwen_image" in args.rollout_patch_groups:
+        bitwise_lora_groups = {"qwen_image", "cosmos3_bitwise"}.intersection(args.rollout_patch_groups)
+        if args.use_lora and bitwise_lora_groups:
             # Missing on engines whose ServerArgs predates --lora-merge-mode.
             if getattr(args, "sglang_lora_merge_mode", None) != "dynamic":
                 logger.warning(
-                    "qwen_image runs LoRA without --sglang-lora-merge-mode dynamic; the engine "
+                    f"{'/'.join(sorted(bitwise_lora_groups))} runs LoRA without --sglang-lora-merge-mode dynamic; the engine "
                     "auto-merges the adapters into the base weights, introducing precision "
                     "drift against the trainer's unmerged forward — training still works."
                 )

scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py

@@ -0,0 +1,168 @@
+"""Cosmos3-Nano T2I GRPO with PickScore, fully colocated on 4 GPUs, with
+bitwise train<->rollout parity.
+
+pretrained = nvidia/Cosmos3-Nano (16B MoT: 8B UND tower frozen, 8B GEN tower
+trained via LoRA r64), 832x480 single frame, num_steps=16, eval_steps=35,
+guidance 4.0, Flow-SDE noise_level=0.7, no KL, per-prompt mean + global std.
+
+Layout: train, rollout and PickScore reward all share the same 4 GPUs
+(--colocate --colocate-reward, one PickScore worker per rollout engine).
+
+SDE schedule: epoch_global_random_choice draws 2 steps per epoch from
+candidates 4-7 — the high-noise segment (sigma 0.94-0.80) of the FlowUniPC
+shift-3 grid the rollout inherits from serving. Step numbers are NOT
+transferable across sigma-grid families - re-derive candidates from |dt| when
+changing model/grid.
+
+Pacing: lr 1e-4 x 1 optimizer step per rollout (the whole rollout is one
+batch). CFG amplifies per-step policy displacement, so training with
+guidance > 1 needs this slower pacing than a comparable CFG-free recipe.
+
+Bitwise parity: --rollout-patch-group cosmos3_bitwise re-expresses the
+engine's kernel organization as the exact op sequence the trainer runs, and
+--sglang-lora-merge-mode dynamic keeps the adapters unmerged (merging rounds
+differently in bf16). Parity holds at tp_size==1 with --lora-ipc-weight-sync.
+
+--diffusion-recompute-old-log-prob: the trainer recomputes old log-probs at
+rollout ingestion so the PPO ratio is implementation-self-consistent; with
+parity on both sides agree bitwise, so this only guards runs without the
+patch group. With 1 step per rollout every optimizer step is exactly
+on-policy.
+
+Usage:
+    python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py
+"""
+
+from dataclasses import dataclass
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+MODEL = "nvidia/Cosmos3-Nano"
+DATASET = "rockdu/miles-diffusion-datasets"
+DATASET_SUBSET = "flowgrpo_pickscore"
+WANDB_PROJECT = "miles-diffusion-grpo"
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+    cuda_visible_devices: str = "0,1,2,3"
+    num_rollout: int = 10000
+    data_dir: str = "/root/datasets"
+    extra_args: str = ""
+
+
+def prepare(args: ScriptArgs) -> str:
+    local_dir = U.hf_download_dataset(DATASET, include=f"{DATASET_SUBSET}/**", data_dir=args.data_dir)
+    return f"{local_dir}/{DATASET_SUBSET}"
+
+
+def execute(args: ScriptArgs, data_dir: str) -> None:
+    run_name = f"diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise_{U.create_run_id()}"
+
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
+
+    rollout_args = (
+        "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "
+        f"--prompt-data {data_dir}/train.jsonl "
+        "--input-key input "
+        "--rollout-batch-size 48 "
+        "--n-samples-per-prompt 16 "
+        f"--num-rollout {args.num_rollout} "
+        "--num-steps-per-rollout 1 "
+        # The Cosmos3 transformer is a packed-sequence single-sample interface;
+        # one request cannot batch multiple outputs.
+        "--rollout-microgroup-size 1 "
+        "--micro-batch-size 1 "
+    )
+
+    diffusion_args = (
+        "--diffusion-num-steps 16 "
+        "--diffusion-output-num-frames 1 "
+        "--diffusion-guidance-scale 4.0 "
+        "--diffusion-noise-level 0.7 "
+        "--diffusion-height 480 "
+        "--diffusion-width 832 "
+        "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_random_choice "
+        "--diffusion-num-sde-steps 2 "
+        "--diffusion-sde-candidate-steps 4,5,6,7 "
+        "--diffusion-recompute-old-log-prob "
+    )
+
+    parity_args = "--rollout-patch-group cosmos3_bitwise --sglang-lora-merge-mode dynamic "
+
+    eval_args = (
+        f"--eval-prompt-data pickscore_test {data_dir}/test.jsonl "
+        "--eval-interval 30 "
+        "--diffusion-eval-num-steps 35 "
+        "--skip-eval-before-train "
+    )
+
+    grpo_args = "--advantage-estimator grpo --globalize-reward-std --diffusion-clip-range 1e-3 "
+
+    optimizer_args = "--lr 1e-4 --adam-beta2 0.999 --weight-decay 1e-4 "
+
+    # UND/GEN towers share layers and differ by parameter name (to_q vs
+    # add_q_proj, mlp vs mlp_moe_gen); LoRA targeting defaults to the GEN
+    # fragments in the cosmos3 train pipeline config.
+    lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 --lora-init-weights gaussian "
+
+    reward_args = (
+        "--rm-type pickscore "
+        "--colocate-reward "
+        "--pickscore-num-workers 4 "
+        "--pickscore-batch-size 8 "
+        "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K "
+        "--pickscore-model-path yuvalkirstain/PickScore_v1 "
+    )
+
+    wandb_args = U.get_default_wandb_args(
+        __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10
+    )
+
+    sglang_args = (
+        "--use-miles-router "
+        "--sglang-server-concurrency 8 "
+        "--update-weight-buffer-size 2147483648 "
+        "--update-weight-target-module transformer "
+    )
+
+    train_backend_args = (
+        "--train-backend fsdp --fsdp-master-dtype fp32 --fsdp-reduce-dtype fp32 --diffusion-forward-dtype bf16 "
+    )
+
+    misc_args = (
+        "--actor-num-gpus-per-node 4 "
+        "--rollout-num-gpus 4 "
+        "--rollout-num-gpus-per-engine 1 "
+        "--num-gpus-per-node 4 "
+        "--colocate "
+    )
+
+    debug_args = "--deterministic-mode --diffusion-debug-mode "
+
+    U.execute_train(
+        train_args=(
+            f"{ckpt_args} {rollout_args} {diffusion_args} {parity_args} {eval_args} {grpo_args} "
+            f"{optimizer_args} {lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} "
+            f"{misc_args} {debug_args} {args.extra_args}"
+        ),
+        num_gpus_per_node=4,
+        config=args,
+        extra_env_vars={
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False",
+            # RL rollout scores raw samples; skip the serving-side guardrail models.
+            "SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1",
+        },
+    )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs) -> None:
+    data_dir = prepare(args)
+    execute(args, data_dir)
+
+
+if __name__ == "__main__":
+    typer.run(main)

tests/ci/fixtures/e2e_standards/test_cosmos3_pickscore_grpo_t2i_4xGPU.json

@@ -0,0 +1,108 @@
+{
+  "meta": {
+    "commit": "local",
+    "source": "test_cosmos3_pickscore_grpo_t2i_4xGPU.py"
+  },
+  "metrics": {
+    "rollout/reward/raw_mean": [
+      [
+        0,
+        0.7647189497947693
+      ],
+      [
+        1,
+        0.7742270827293396
+      ]
+    ],
+    "rollout/reward/raw_median": [
+      [
+        0,
+        0.7654154300689697
+      ],
+      [
+        1,
+        0.7737646102905273
+      ]
+    ],
+    "rollout/reward/raw_num_samples": [
+      [
+        0,
+        768.0
+      ],
+      [
+        1,
+        768.0
+      ]
+    ],
+    "rollout/reward/raw_std": [
+      [
+        0,
+        0.06847453862428665
+      ],
+      [
+        1,
+        0.061292387545108795
+      ]
+    ],
+    "train/grad_norm": [
+      [
+        1.0,
+        5.782852167612873e-05
+      ],
+      [
+        2.0,
+        3.818701588897966e-05
+      ]
+    ],
+    "train/log_prob_mean_abs_diff": [
+      [
+        1.0,
+        0.0
+      ],
+      [
+        2.0,
+        0.0
+      ]
+    ],
+    "train/log_prob_new_idx_0": [
+      [
+        1.0,
+        -0.2110141608864069
+      ],
+      [
+        2.0,
+        -0.3804309892778595
+      ]
+    ],
+    "train/log_prob_old_idx_0": [
+      [
+        1.0,
+        -0.2110141608864069
+      ],
+      [
+        2.0,
+        -0.3804309892778595
+      ]
+    ],
+    "train/model_output_mean_abs_diff": [
+      [
+        1.0,
+        0.0
+      ],
+      [
+        2.0,
+        0.0
+      ]
+    ],
+    "train/model_output_rel_max": [
+      [
+        1.0,
+        0.0
+      ],
+      [
+        2.0,
+        0.0
+      ]
+    ]
+  }
+}

tests/e2e/short/test_cosmos3_pickscore_grpo_t2i_4xGPU.py

@@ -0,0 +1,40 @@
+"""E2E: Cosmos3-Nano T2I PickScore GRPO, 4-GPU fully colocated (train, rollout
+and PickScore reward share the same 4 GPUs) — runs the bitwise recipe's real
+configuration and checks its metric series against the registered standard
+(tests/ci/fixtures/e2e_standards/). Runs with --deterministic-mode, so every
+metric is compared strictly, bit for bit.
+
+Only --num-rollout is cut down, 10000 -> 2: one weight-sync round trip is
+enough to catch drift in the post-update rollout and the second optimizer step.
+--cuda-visible-devices "" unpins the recipe's default 0,1,2,3 so it inherits
+the runner's GPU set (the 5gpu runner exposes GPUs 3-7).
+
+What this test uniquely guards is the bitwise train<->rollout parity the
+cosmos3_bitwise rollout patch group + --sglang-lora-merge-mode dynamic
+establish: train/model_output_mean_abs_diff / train/model_output_rel_max
+compare the raw DiT outputs between engine and trainer — with parity both are
+exactly 0. The recipe recomputes old log-probs at ingestion, so the log_prob
+series guard the trainer-side pipeline rather than cross-side parity.
+"""
+
+from tests.ci.e2e_metrics_registry import register_e2e_ci
+
+register_e2e_ci(
+    est_time=4800,
+    suite="stage-c-5-gpu-h200",
+    script="scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu_bitwise.py",
+    args=["--num-rollout", "2", "--cuda-visible-devices", ""],
+    labels=["e2e"],
+    metrics=[
+        "rollout/reward/raw_num_samples",
+        "rollout/reward/raw_mean",
+        "rollout/reward/raw_median",
+        "rollout/reward/raw_std",
+        "train/log_prob_old_idx_0",
+        "train/log_prob_new_idx_0",
+        "train/log_prob_mean_abs_diff",
+        "train/model_output_mean_abs_diff",
+        "train/model_output_rel_max",
+        "train/grad_norm",
+    ],
+)

tests/fast/backends/sglang_diffusion_utils/test_rollout_patch_groups.py

@@ -44,6 +44,7 @@ def test_unknown_group_fails_loud(self, monkeypatch):
     def test_builtin_group_registered(self):
         # The decorator ran at import time for the in-repo groups.
         assert "qwen_image" in mp._ROLLOUT_PATCH_APPLIERS
+        assert "cosmos3_bitwise" in mp._ROLLOUT_PATCH_APPLIERS
         assert "wan" in mp._ROLLOUT_PATCH_APPLIERS
         assert "ltx" in mp._ROLLOUT_PATCH_APPLIERS