diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py new file mode 100644 index 000000000..49a4f4424 --- /dev/null +++ b/miles/backends/fsdp_utils/configs/cosmos3.py @@ -0,0 +1,177 @@ +"""Cosmos3 training pipeline config.""" + +from __future__ import annotations + +import math + +import torch +from miles.utils.types import CondKwargs + +from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config + +# Cosmos3 reuses the Wan2.2 VAE (4x temporal compression). +_VAE_TEMPORAL_FACTOR = 4 + +# GEN-tower param name fragments; everything else (UND tower, lm_head, unused heads) stays frozen. +_GEN_PARAM_FRAGMENTS = ( + ".add_q_proj.", + ".add_k_proj.", + ".add_v_proj.", + ".to_add_out.", + ".norm_added_q.", + ".norm_added_k.", + ".mlp_moe_gen.", + ".input_layernorm_moe_gen.", + ".post_attention_layernorm_moe_gen.", + ".norm_moe_gen.", + ".proj_in.", + ".proj_out.", + ".time_embedder.", +) + + +def _is_gen_param(name: str) -> bool: + dotted = f".{name}" + return any(fragment in dotted for fragment in _GEN_PARAM_FRAGMENTS) + + +@register_train_pipeline_config("cosmos3") +class Cosmos3TrainPipelineConfig(TrainPipelineConfig): + hf_ckpt_name_patterns = ("cosmos3", "cosmos-3") + # Timesteps stay fp32 (bf16 rounds the karras grid); conds pass through, the packed forward casts its own inputs. + input_dtype_policy = {"latents": "default", "cond": None, "timestep": "fp32"} + # The packed forward is single-sample; never batch the CFG branches. + cfg_batching = False + lora_target_modules = ["add_q_proj", "add_k_proj", "add_v_proj", "to_add_out"] + + @classmethod + def validate_args(cls, args) -> None: + if list(args.update_weight_target_modules) != ["transformer"]: + raise ValueError("Cosmos3 requires --update-weight-target-module transformer.") + + def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: + if cond is None or cond.text_ids is None: + return {} + return { + "text_ids": cond.text_ids.to(device), + "text_mask": cond.text_mask.to(device), + "fps": cond.fps, + } + + def collate_cond_for_sample_batch( + self, + per_sample_cond_kwargs: list[dict], + device: torch.device, + pad_to_len: int | None = None, + ) -> dict: + # Packed single-sample forward: keep per-sample kwargs; no batched tensor to build. + return {"per_sample": per_sample_cond_kwargs} + + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict | None, + neg_cond: dict | None, + joint_cond: dict | None, + use_cfg: bool, + cfg_batching: bool, + guidance_scale: float, + true_cfg_scale: float | None, + ) -> torch.Tensor: + assert not cfg_batching, "Cosmos3 packed forward is single-sample; cfg_batching unsupported" + config = model.config + preds = [] + for i, pos in enumerate(pos_cond["per_sample"]): + latent = latents_input[i : i + 1] + timestep = float(timesteps_input[i]) + pred = self._packed_forward(model, latent, timestep, pos, config) + if use_cfg: + pred_neg = self._packed_forward(model, latent, timestep, neg_cond["per_sample"][i], config) + pred = self.cfg_combine(pred, pred_neg, guidance_scale, true_cfg_scale=true_cfg_scale) + preds.append(pred) + return torch.stack(preds, dim=0) + + def _packed_forward( + self, + model: torch.nn.Module, + latent: torch.Tensor, + timestep: float, + cond: dict, + config, + ) -> torch.Tensor: + """One packed (text, video) forward mirroring the diffusers denoising loop (T2V/T2I: all frames noisy).""" + from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import ( + get_3d_mrope_ids_text_tokens, + get_3d_mrope_ids_vae_tokens, + ) + + device = latent.device + und_len = int(cond["text_mask"].sum().item()) + input_ids = cond["text_ids"].reshape(-1)[:und_len] + + text_mrope_ids, next_offset = get_3d_mrope_ids_text_tokens( + num_tokens=und_len, + temporal_offset=0, + use_float_positions=config.enable_fps_modulation, + ) + vision_offset = next_offset + config.unified_3d_mrope_temporal_modality_margin + + p = config.latent_patch_size + _, _, latent_t, latent_h, latent_w = latent.shape + patch_h = math.ceil(latent_h / p) + patch_w = math.ceil(latent_w / p) + num_vision_tokens = latent_t * patch_h * patch_w + + vision_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=vision_offset, + reset_spatial_indices=config.unified_3d_mrope_reset_spatial_ids, + fps=cond["fps"] if config.enable_fps_modulation else None, + base_fps=float(config.base_fps), + temporal_compression_factor=_VAE_TEMPORAL_FACTOR, + ) + + sequence_length = und_len + num_vision_tokens + vision_sequence_indexes = torch.arange(und_len, sequence_length, dtype=torch.long, device=device) + preds_vision, _, _ = model( + input_ids=input_ids, + text_indexes=torch.arange(und_len, dtype=torch.long, device=device), + position_ids=torch.cat([text_mrope_ids, vision_mrope_ids], dim=1).to(device), + und_len=und_len, + sequence_length=sequence_length, + vision_tokens=[latent], + vision_token_shapes=[(latent_t, patch_h, patch_w)], + vision_sequence_indexes=vision_sequence_indexes, + vision_mse_loss_indexes=vision_sequence_indexes, + vision_timesteps=torch.full((num_vision_tokens,), timestep, device=device, dtype=torch.float32), + vision_noisy_frame_indexes=[torch.arange(latent_t, dtype=torch.long, device=device)], + ) + return preds_vision[0] + + def cfg_combine( + self, + noise_pred_pos: torch.Tensor, + noise_pred_neg: torch.Tensor, + guidance_scale: float, + true_cfg_scale: float | None = None, + ) -> torch.Tensor: + scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale + return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) + + def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: + # The UND tower sits inside the training graph, so it must be explicitly frozen. + for name, param in model.named_parameters(): + if "lora_" not in name and not _is_gen_param(name): + param.requires_grad_(False) + + # Cast the timestep sinusoid to the MLP weight dtype before linear_1, exactly as sglang-d does. + def _cast_to_weight_dtype(module, args): + dtype = module.linear_1.weight.dtype + 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) diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 787adc6e0..8e54fbd5e 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -42,7 +42,7 @@ def prepare_flow_grpo_batch( guidance_scale = args.diffusion_guidance_scale true_cfg_scale = args.diffusion_true_cfg_scale cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale - use_cfg = cfg_scale > 0 + use_cfg = cfg_scale > 1.0 # matches sglang do_cfg: guidance<=1 runs single-branch if len(ctx.models) == 1: component_name, model = next(iter(ctx.models.items())) diff --git a/miles/backends/fsdp_utils/models/diffusers/cosmos3/__init__.py b/miles/backends/fsdp_utils/models/diffusers/cosmos3/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py new file mode 100644 index 000000000..daec57f02 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py @@ -0,0 +1,4 @@ +from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan + + +FSDP_PARALLEL_PLAN = FSDPParallelPlan() diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 4003329f3..0163adf8b 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -107,6 +107,9 @@ def _parse_cond_kwargs( data.get("audio_encoder_attention_mask"), deserialize_func=deserialize_func ), pooled_projections=_parse_tensor_or_list(data.get("pooled_projections"), deserialize_func=deserialize_func), + text_ids=deserialize_func(data.get("text_ids")), + text_mask=deserialize_func(data.get("text_mask")), + fps=data.get("fps"), ) diff --git a/miles/utils/types.py b/miles/utils/types.py index a491dcff8..2e6d900f1 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -33,6 +33,10 @@ class CondKwargs: encoder_attention_mask: list[torch.Tensor] | None = None audio_encoder_attention_mask: list[torch.Tensor] | None = None pooled_projections: list[torch.Tensor] | None = None + # Cosmos3: token-level conditioning (no separate text encoder). + text_ids: torch.Tensor | None = None + text_mask: torch.Tensor | None = None + fps: float | None = None @dataclass diff --git a/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py new file mode 100644 index 000000000..da1eadeff --- /dev/null +++ b/scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py @@ -0,0 +1,160 @@ +"""Cosmos3-Nano T2I GRPO with PickScore, fully colocated on 4 GPUs. + +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 8-11. The Cosmos3 checkpoint ships a Karras flow-sigma grid whose +head steps 1-7 sit at sigma>0.96 with |dt|<0.02 and train nothing; steps 8-11 +are the true high-noise segment (sigma 0.94-0.80). 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. + +--diffusion-recompute-old-log-prob: the trainer recomputes old log-probs at +rollout ingestion so the PPO ratio is implementation-self-consistent (rollout +fa kernels vs train SDPA would otherwise leak into the ratio). With 1 step per +rollout this makes every optimizer step exactly on-policy. + +Usage: + python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.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_{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 8,9,10,11 " + "--diffusion-recompute-old-log-prob " + ) + + 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 = "--diffusion-debug-mode " + + U.execute_train( + train_args=( + f"{ckpt_args} {rollout_args} {diffusion_args} {eval_args} {grpo_args} {optimizer_args} " + f"{lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} {misc_args} " + f"{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)