diff --git a/docs/algo/diffusion_distillation.md b/docs/algo/diffusion_distillation.md new file mode 100644 index 000000000..3f01a7571 --- /dev/null +++ b/docs/algo/diffusion_distillation.md @@ -0,0 +1,431 @@ +# Diffusion Distribution Matching: DMD2 + +Last updated: 09/11/2026. + +## Background and scope + +DMD2 trains a few-step generator using a frozen real-score teacher and a trainable +fake-score model that tracks the generator's distribution. This implementation +supports **Qwen-Image text-to-image, distribution-only DMD2** from prompts. It +includes differentiable student sampling, alternating student/fake-score updates, +LoRA EMA, complete training checkpoints, and a student inference artifact. + +This is not original DMD's paired teacher-trajectory/perceptual-regression recipe, +and it does not include DMD2's optional adversarial objective. Original DMD, GAN, +CausVid, Self-Forcing, Qwen-Image Edit and other architectures are outside the +current supported path. A successful short run establishes execution, not +paper-level reproduction, convergence or generation-quality improvement. + +### DMD2 is not the OPD trainer path + +| Contract | DMD2 distribution matching | Existing diffusion OPD | +|---|---|---| +| Trainer selector | `algorithm.trainer_type=distribution_matching` | `algorithm.trainer_type=policy_gradient` | +| Data flow | Prompt batches; samples generated inside the training engine | Online rollout trajectories replayed by frozen teachers | +| Configuration | Independent top-level `dmd` group | `distillation` group with `enabled=true` | +| Supervision | Teacher and learned fake-score x0 estimates at re-noised student samples | Teacher reverse-transition means for `distill_kl` | +| Optimization | Separate student and fake-score optimizers | Actor update with the configured OPD loss | +| Rewards / PPO tensors | Not needed | Existing policy-gradient/OPD lifecycle | + +Do not enable `distillation.enabled`, actor `use_distill_loss`, or additional KL +objectives to run DMD2. See [diffusion OPD](diffusion_opd.md) for that separate +configuration and teacher-scheduling contract. + +### What `offline` means here + +`algorithm.sample_source=offline` selects the framework's **engine-local execution +path**, not offline RL or training on a fixed dataset of pre-generated images. +No independent rollout server or reward workers are started. + +The current student generates fresh samples from prompts and random noise during +each training attempt. Sampling runs inside the FSDP training engine so the +selected student forward can retain its autograd graph. Thus sample generation +is online during training, despite the configuration name `offline`; it does not +require precomputed student images or teacher trajectories. Keep +`sample_source=offline` for this implementation. + +The supported launcher uses `verl_omni.trainer.main_diffusion`. +`main_diffusion_v1.py` integration is not implemented for this DMD2 path. + +## Objectives and gradient boundaries + +The Qwen adapter works with normalized image latents packed as `[B,N,D]`, where +`B` is the physical microbatch, `N` is the packed spatial sequence length and `D` +is the packed channel width. The adapter derives geometry from the checkpoint; +the engine does not guess an image layout. Flow corruption and x0 conversion are + +```{math} +x_\sigma = (1-\sigma)\operatorname{sg}(x_g) + \sigma\epsilon, +\qquad +\hat{x}_0 = x_\sigma - \sigma v(x_\sigma,\sigma,c), +``` + +where `sg` denotes stop-gradient and Qwen predicts velocity in the +`noise - clean_latent` convention. The noise, sigma and corrupted inputs are +shared between real/fake scoring within a student attempt. + +### Student objective + +The teacher runs positive and negative conditioning. Standard CFG is applied in +**packed velocity space**, before x0 conversion: + +```{math} +v_r = v_{\mathrm{neg}} + s(v_{\mathrm{pos}}-v_{\mathrm{neg}}). +``` + +`dmd.cfg_norm=layer_norm` then rescales each last-dimension vector by +`norm(v_pos) / max(norm(v_r), 1e-12)`; this is norm rescaling, not a learned +LayerNorm module. `none` disables it. `scalar` uses one batch-wide norm ratio, +clamped above at 1. Student and fake-score predictions remain conditional-only. + +Let `x_r` and `x_f` be the teacher and fake-score x0 estimates. For each sample: + +```{math} +Z = \max\left(\operatorname{mean}_{\mathrm{nonbatch}} + |\operatorname{sg}(x_g)-x_r|,\eta\right), +\qquad +g = \operatorname{nan\_to\_num}\left(\frac{x_f-x_r}{Z}\right), +``` + +```{math} +L_{\mathrm{student}} = \frac{1}{2}\operatorname{mean} +\left(x_g-\operatorname{sg}(x_g-g)\right)^2. +``` + +Score predictions and `g` are detached. Only the student receives gradients. +The normalizer spans all non-batch dimensions, and the current Qwen objective is +unmasked. Nonfinite gradient elements are counted **before** sanitization; monitor +this count as well as skipped optimizer updates. Corruption, x0 conversion, +normalization and losses use fp32, independently of the transformer compute dtype. + +### Fake-score objective + +A fake-score attempt generates a new student sample without gradients, re-noises +it, and trains the fake score with ordinary flow denoising MSE: + +```{math} +L_{\mathrm{fake}} = \operatorname{mean} +\left(v_f(x_\sigma,\sigma,c) + - \operatorname{sg}(\epsilon-x_g)\right)^2. +``` + +Both the corrupted model input and the target are detached from the generator. +Only the fake-score optimizer is stepped; teacher scoring is unnecessary here. +`DMDLoss`, registered as `dmd2` in `diffusion_algos.py`, dispatches these two losses +using `dmd_stage`; it does not require old log-probabilities or advantages. + +### Student sampling and score timesteps + +For `S` inference steps, start with the linear grid `[1, ..., 1/S, 0]` and apply +`f(sigma) = mu * sigma / (1 + (mu - 1) * sigma)` once. Four steps with +`dmd.rollout_timestep_shift=3` give `[1, 0.9, 0.75, 0.5, 0]`. This deliberately +differs from the stock Qwen pipeline's resolution-dependent shift. + +Each microbatch samples a uniform exit index in `[0,S)`, broadcast from rank 0 +so sharded ranks execute the same forwards and gradient exit. Earlier Euler +transitions run without gradients. Only the selected student prediction retains +a graph during a student attempt; fake-score attempts retain no student graph. +Inference instead executes the entire fixed grid, without per-step added noise. + +Score sigma is sampled separately for each sample: + +- With `score_discrete_steps=T`, require `T` to equal the scheduler's training + grid size. Draw an integer in `[0,T)`, apply `score_timestep_shift` to `t/T`, + then clamp to `[score_sigma_min, score_sigma_max]`. +- With `score_discrete_steps=0`, draw uniformly within those sigma bounds; + no timestep shift is applied. + +## Runtime and ownership + +The implementation follows the existing loss/engine/worker/trainer structure: + +```text +main_diffusion.TaskRunner + -> DistributionMatchingRayTrainer(BaseRayDiffusionTrainer) + -> DMDTrainingWorker(TrainingWorker) + -> DMDDiffusersFSDPEngine(DiffusersFSDPEngine) + -> QwenImageDMD2 + frozen conditioning provider + -> DMDLoss through the existing diffusion_loss dispatcher +``` + +| Component | Responsibility | +|---|---| +| `DistributionMatchingRayTrainer` in `trainer/diffusion/ray_diffusion_trainer.py` | Reuse offline initialization, resource pools, dataloaders and profiling; run the explicit 1:K loop; publish complete checkpoints and export | +| `DMDTrainingWorker` in `workers/dmd_worker.py` | Reuse distributed setup, dispatch and mini/microbatch handling; constrain each actor call to one optimizer attempt | +| `DMDDiffusersFSDPEngine` in `workers/engine/fsdp/dmd_impl.py` | Differentiable sampling, score calls, optimizer selection, RNG streams, numerical skips, EMA and DMD-specific checkpoint state | +| `QwenImageDMD2` in `pipelines/qwen_image_distillation/diffusers_training_adapter.py` | Stateless `(QwenImagePipeline, dmd2)` registry adapter: conditioning construction, geometry, packing, model inputs, velocity-to-x0 conversion and sigma grid | +| `QwenImageConditionProvider` | Per-run frozen text encoding and cached negative conditioning; no text-encoder gradients | +| `DMDLoss` and `trainer/diffusion/distillation/utils.py` | Registered objective dispatch and pure tensor equations | + +The adapter is a classmethod/staticmethod registry class, not an instantiated +trainer or owner of optimizer state. The engine checks for its required DMD +methods at construction. No separate generic role graph, transport interface or +`DistributionMatchingModelAdapter` mixin is required by this implementation. + +### Logical roles on one physical base + +| Logical role | Implementation | Optimized? | +|---|---|---| +| `student` | `default` LoRA adapter | Student optimizer | +| `fake_score` | `fake_score` LoRA adapter | Independent fake-score optimizer | +| `teacher_score` | `reference` adapter context, which disables adapters | No; frozen base prediction | +| `student_ema` | `student_ema` LoRA adapter | No optimizer; student EMA only | + +Fake-score and EMA adapters start as copies of the student. The engine records +nonempty, disjoint student/fake-score parameter sets and rejects overlap; +gradients on an inactive optimized role are errors. `LoRAAdapterMixin` provides +adapter selection/restoration, copying and EMA. Teacher scoring is not an OPD +teacher worker or an extra reference-KL objective. + +Shared adapters are a storage choice, not a mathematical requirement of DMD2, +but they are the **only implemented layout here**. There is no selectable +`shared_base_adapters` / `colocated_independent` role-layout configuration. +Independent full models and separately placed teachers are not enabled. + +### Cycles, skips and EMA + +For every cycle, consume one fresh student prompt batch followed by `K` fresh +fake-score prompt batches. Each actor call is constrained to one minibatch, +one epoch and at most one optimizer update. `K` therefore counts attempts, +not an assumed number of updates hidden inside an actor call. + +- `training/global_step` advances after all `1 + K` attempts complete, including + numerical skips. It counts completed cycles, not successful student updates. +- Successful optimizer counters advance independently. Schedulers advance only + with their own successful optimizer step; the fake scheduler's configured + horizon is `K` times the student's. +- Ranks agree on nonfinite loss/gradient skips before stepping. A skipped role + clears gradients and advances neither its scheduler nor EMA; it is not retried. +- EMA starts as a student copy. A successful student update applies + `EMA = decay * EMA + (1 - decay) * student` once its successful-update count + reaches `ema_start_step`. Earlier updates and skips leave EMA unchanged. +- Even all-skipped cycles consume the finite training budget. If either role + has zero successful updates at the end, training raises instead of exporting + a supposedly trained student. +- Unexpected exceptions stop the trainer. A partially applied cycle cannot be + undone by resetting counters: resume from the last complete checkpoint in a + new trainer, rather than retrying in-process. + +## Configuration and usage + +Follow the [GPU and training installation](../start/install.md), then use the +{doc}`Qwen-Image DMD2 example <../examples/qwen_image/dmd2_trainer>`. Commands below +run from the repository root. The launcher supplies the full configuration; +these selectors identify the route: + +```yaml +algorithm: + trainer_type: distribution_matching + sample_source: offline +actor_rollout_ref: + model: + algorithm: dmd2 + model_type: diffusion_dmd_model +``` + +The existing actor loss-mode interpolation resolves to `dmd2`. Student optimizer +settings remain at `actor_rollout_ref.actor.optim`; fake-score settings are at +`dmd.fake_score_optim`. Both use existing FSDP optimizer configuration and the +engine's supported constant/cosine schedulers. + +### DMD configuration reference + +All fields below belong to top-level `dmd` (`DiffusionDMDConfig`), not to OPD's +`distillation` group. + +| Field | Default | Meaning | +|---|---|---| +| `fake_update_ratio` | `2` | Positive integer fake attempts per student attempt | +| `student_micro_batch_size_per_gpu` | `1` | Student physical microbatch per DP rank | +| `fake_score_micro_batch_size_per_gpu` | `1` | Independent fake-score physical microbatch | +| `fake_score_optim` | LR `2e-5`, weight decay `0.001` | Existing FSDP optimizer config for fake score | +| `teacher_guidance_scale` | `4.0` | Positive teacher CFG scale; student/fake stay conditional-only | +| `cfg_norm` | `layer_norm` | Packed-velocity rescaling: `none`, `layer_norm`, or `scalar` | +| `negative_prompt` | `" "` | Explicit teacher negative text; empty string is valid, null is not | +| `normalization_epsilon` | `1e-6` | Positive lower bound for the per-sample normalizer | +| `rollout_timestep_shift` | `3.0` | Fixed student/inference grid shift, at least 1 | +| `score_discrete_steps` | `1000` | Scheduler-sized discrete grid; 0 selects continuous uniform sampling | +| `score_sigma_min`, `score_sigma_max` | `0.02`, `0.98` | Bounds satisfying `0 < min < max <= 1` | +| `score_timestep_shift` | `3.0` | Discrete score-sampling shift, at least 1 | +| `ema_decay` | `0.999` | EMA decay in `[0,1]` | +| `ema_start_step` | `0` | Nonnegative successful-student-update threshold | +| `export_role` | `student` | One inference artifact; `student_ema` is an explicit alternative | + +The launcher additionally sets 1024×1024, four inference steps, max sequence +length 1024, LoRA rank/alpha 32, student LR `1e-4`, and FSDP2/BF16 with native +training attention. These are **launcher overrides**, not all model-config defaults. +It uses no VAE during training; decoding is part of the separate generation tool. + +### Prompt data and launch + +Use prompt parquet with the existing `RLHFDataset` schema, for example: + +```python +{"prompt": [{"role": "user", "content": "A red apple on a wooden table"}]} +``` + +Raw inputs accept a string or one text-only user message. The condition provider +applies the checkpoint's Qwen system template and prefix removal; custom system +messages and multi-message conversations are rejected. A custom dataset can +instead supply post-collation embeddings `[B,L,D]` with optional `[B,L]` masks; +pre-tokenized inputs require masks and the correct Qwen template prefix. +Positive and negative conditioning must match the batch. No image preference +pairs, teacher trajectories, rewards or precomputed student samples are needed. + +```bash +MODEL_PATH=/path/to/Qwen-Image \ +TRAIN_FILES=/path/to/train.parquet \ +VAL_FILES=/path/to/test.parquet \ +OUTPUT_DIR=outputs/qwen_dmd2 \ +NUM_GPUS=8 TOTAL_TRAIN_STEPS=1000 \ +bash examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh +``` + +A validation parquet is still supplied for shared dataloader initialization; +there is no validation-generation replica. The launcher disables validation +before training and sets `test_freq=-1`. + +### Batching and diagnostics + +`data.train_batch_size` is the global batch **per attempt** and must divide evenly +across DP ranks. `GLOBAL_BATCH_SIZE` defaults to `NUM_GPUS` in the example. +`STUDENT_MICRO_BATCH_SIZE` and `FAKE_MICRO_BATCH_SIZE` independently default to 1. +At SP=1, eight GPUs with global batch 16 and microbatch 2 use physical batch 2 +per rank. A rank-local batch of 3 with microbatch 2 instead accumulates `2 + 1`. +Dense tails use native TensorDict splitting and sample-weighted means; no extra +training examples are padded in. Dynamic batching is not supported. + +Samples in a physical microbatch must share image geometry. Accumulation and +physical batching may sample different rollout depths; `scalar` CFG additionally +couples samples through its batch-wide norm. Neither is automatically a controlled +performance comparison at fixed effective batch. Validated distributed coverage +uses SP=1; do not infer SP>1 support from the generic configuration surface. + +Each attempt retains its own metric prefix: `student/0/...`, `fake_score/0/...`, +`fake_score/1/...`, and so on. Useful suffixes are `dmd/loss`, `dmd/normalizer`, +`dmd/gradient_norm` for the student, `fake_score/loss` for fake denoising, and +`dmd/update_applied`, `dmd/skip_nonfinite`, `dmd/nonfinite`, `dmd/rollout_exit`, +`training/samples` where applicable. Component timers include +`perf/condition_encode_s`, `perf/student_rollout_s`, `perf/teacher_score_s`, +`perf/fake_score_s`, `perf/backward_s`. Top-level +`training/student_optimizer_steps` and `training/fake_score_optimizer_steps` +record cumulative successful updates independently of `training/global_step`. + +Loss-like microbatch metrics are sample-weighted; durations and counts are summed. +DP aggregation reports means, not slowest-rank wall time or global count sums. +Peak allocated/reserved memory uses rank maxima, cumulative from worker +initialization. `perf/cycle_s` covers the attempts but excludes checkpointing; +`perf/checkpoint_s` appears only when a checkpoint is saved. Nested host timings +are not additive CUDA kernel costs. Reuse `global_profiler.steps` and +`actor_rollout_ref.actor.profiler` for traces; see [profiling](../perf/profiler.md). + +## Checkpoint, resume and inference contracts + +A training checkpoint and one inference artifact serve different purposes: + +```text +OUTPUT_DIR/ + global_step_N/ + actor/ + model_world_size__rank_.pt + optim_world_size__rank_.pt + extra_state_world_size__rank_.pt + dmd_state_rank_.pt + data.pt + trainer.pt + latest_checkpointed_iteration.txt + inference/ + adapter_model.safetensors + adapter_config.json + inference_manifest.json +``` + +The standard FSDP shards save the shared model, including all adapters, the student +optimizer, scheduler and worker RNG. `dmd_state_rank_*` adds the fake optimizer +and scheduler, successful/skipped counters, and separate initial-noise, +rollout-decision, score-sigma and score-noise generator states. `data.pt` holds the +stateful dataloader; `trainer.pt` records completed cycles, data epoch, driver RNG +and a canonical configuration fingerprint. EMA is model state, not a third +optimizer or a separately required inference export. + +Checkpoints are synchronous and atomically published on a shared local filesystem +after all-rank file and clock validation. A failed save does not replace the +latest-complete pointer. Positive `trainer.save_freq` saves on its interval and +at the final cycle; disabling saves also disables the final training checkpoint. +Positive `trainer.max_actor_ckpt_to_keep` prunes older checkpoint directories in +that output root, so use a dedicated directory for each run. + +Resume defaults to `auto`. To choose a checkpoint, append these overrides to the +same training command, normally with a fresh `OUTPUT_DIR`: + +```bash +MODEL_PATH=/path/to/Qwen-Image \ +TRAIN_FILES=/path/to/train.parquet \ +VAL_FILES=/path/to/test.parquet \ +OUTPUT_DIR=outputs/qwen_dmd2_resumed \ +NUM_GPUS=8 TOTAL_TRAIN_STEPS=1000 \ +bash examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + trainer.resume_mode=resume_path \ + trainer.resume_from_path=outputs/qwen_dmd2/global_step_500 +``` + +Keep the same model, engine, optimizer, DMD and data settings. Version, world size, +role counters, required state and configuration are checked before worker load. +A changed configured training horizon also changes optimizer compatibility. +Legacy generic-distillation checkpoints are rejected, not silently reinterpreted; +there is no automatic format or counter migration. Load training checkpoints only +from trusted sources because optimizer/RNG restoration uses Python serialization. + +Exact recovery of saved state does not guarantee bitwise-identical subsequent +updates. Real native/BF16 runs have shown differing gradients on repeated backward +with identical model, inputs and RNG. Do not equate successful resume or exact +checkpoint restoration with deterministic full-model training replay. + +### Student inference artifact + +After the finite training budget, the default export is **student**. The fake score +and teacher are never exported for generation. `dmd.export_role=student_ema` +selects EMA instead; choose this at run configuration time because DMD settings +participate in the resume fingerprint. + +Export validates finite complete adapter parameters and their own PEFT config; +custom LoRA targets must be covered by `model.fsdp_layer_prefixes`. The artifact +is atomically published and will not overwrite an incompatible existing export. +Its manifest records role, cycle/success counts, base model identity, resolved +revision when available, transformer-config hash, weight checksum, resolution, +sequence limit and fixed Euler sampling settings. A config hash is not a base +weight checksum; unversioned local bases have a null revision and must be kept +immutable by the user. + +This is a **base-dependent LoRA**, not a merged self-contained pipeline. Reload +and decode using the supplied tool, which checks the artifact and uses the +recorded conditional-only schedule: + +```bash +python examples/dmd2_trainer/qwen_image/generate.py \ + --artifact outputs/qwen_dmd2/inference \ + --prompt 'A red apple on a wooden table' \ + --seed 42 --output outputs/apple.png +``` + +Use `--base-model` if the same base checkpoint has moved. Do not substitute stock +pipeline scheduler/CFG defaults. Automatic CheckpointEngine validation-replica +synchronization, vLLM-Omni serving/request batching, standalone score transports, +full finetuning and NPU validation are not delivered by this path. FSDP1 requires +`use_orig_params=true`; FSDP2 is the example default. + +## Validation and further reading + +The validation commands in the +{doc}`Qwen-Image DMD2 example <../examples/qwen_image/dmd2_trainer>` +cover CPU/configuration checks, real two-rank tiny-Qwen FSDP1/FSDP2 updates, +unequal microbatch tails, EMA, checkpoint replay, numerical skips and adapter +reload. A separate production smoke exercises Ray routing, checkpoint publication +and export. CPU fake-engine tests alone cannot validate shared-LoRA FSDP behavior; +real checkpoint training and decoded-generation evidence must state the model, +precision, rank count and actual completed steps. Neither a decoded image nor +finite losses establishes a quality gain. + +- [DMD2 paper](https://arxiv.org/abs/2405.14867) and [reference implementation](https://github.com/tianweiy/DMD2). +- [Diffusion OPD](diffusion_opd.md): separate frozen-teacher transition supervision. +- [Direct-preference integration guide](../contributing/integrating_a_new_direct_preference_algorithm_for_diffusion_model.md): shared offline infrastructure, not the DMD2 objective or update semantics. diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index ecd54785a..04913bfee 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -4,25 +4,32 @@ Pipelines Interface Last updated: |today| (API docstrings are auto-generated). A *pipeline* in VeRL-Omni packages everything needed to plug a particular -diffusion model architecture into the training loop: +model architecture into the training loop. Two adapter families are available: -- a **training-side adapter** subclassing +- autoregressive omni models use a training-side + :class:`~verl_omni.pipelines.model_base.OmniModelBase` and an optional + rollout-side :class:`~verl_omni.pipelines.model_base.OmniRolloutPipelineBase`; +- diffusion models use a training-side adapter subclassing :class:`~verl_omni.pipelines.model_base.DiffusionModelBase` that handles scheduler setup, model-input construction, and the per-step forward / reverse-sampling logic used by RL algorithms (e.g. FlowGRPO); -- an optional **rollout-side adapter** registered via +- their optional rollout-side adapter is registered via :class:`~verl_omni.pipelines.model_base.VllmOmniPipelineBase` that hooks into vLLM-Omni's diffusion serving stack to expose log-probabilities. -Adapters are auto-selected by matching the pair -``(DiffusionModelConfig.architecture, DiffusionModelConfig.algorithm)`` against the -registered ``(architecture, algorithm)`` key. The architecture is read from the -model's ``model_index.json``; the algorithm string is taken from the model config's -``actor_rollout_ref.model.algorithm`` value. +Autoregressive training adapters are selected by ``(architecture, model_stage)``; +their rollout adapters are selected by the vLLM-Omni ``pipeline_name``. Diffusion +adapters are selected by matching +``(DiffusionModelConfig.architecture, DiffusionModelConfig.algorithm)`` against a +registered ``(architecture, algorithm)`` key. Diffusion architecture is read from +``model_index.json`` and the algorithm from +``actor_rollout_ref.model.algorithm``. .. autosummary:: :nosignatures: + verl_omni.pipelines.model_base.OmniModelBase + verl_omni.pipelines.model_base.OmniRolloutPipelineBase verl_omni.pipelines.model_base.DiffusionModelBase verl_omni.pipelines.model_base.VllmOmniPipelineBase verl_omni.pipelines.qwen_image_flow_grpo.QwenImage @@ -33,6 +40,20 @@ model's ``model_index.json``; the algorithm string is taken from the model confi Model Base ~~~~~~~~~~~~~~~~~ +.. autoclass:: verl_omni.pipelines.model_base.OmniModelBase + :members: register, get_class, get_class_by_name, + register_auto_classes, + get_strip_modules, configure_processor, configure_tokenizer, + configure_model, prepare_model_inputs + +.. autoclass:: verl_omni.pipelines.model_base.OmniRolloutPipelineBase + :members: register, get_class, + build_stage_configs, rollout_flags, weight_sync_stage_ids, policy_stage_id, + get_pipeline_id, ensure_pipeline_registered, get_engine_hf_overrides, + get_stage_engine_extras, prepare_engine_prompt, + postprocess_agent_loop_output, + combine_engine_outputs + .. autoclass:: verl_omni.pipelines.model_base.DiffusionModelBase :members: register, get_class, build_scheduler, set_timesteps, diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index 1dd472aa8..9141d08a7 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -13,10 +13,10 @@ under [`verl_omni/pipelines/`](https://github.com/verl-project/verl-omni/tree/ma Decide which **training stage** you want to train and how the model decomposes: - **Stage-split**: Multi-component omni models (thinker → talker → code2wav) - train only the text-understanding head during RL post-training. Other - components are stripped before FSDP wrapping to save memory. This is the - Qwen3-Omni pattern — the thinker is the autoregressive language model; talker - and codec are inference-only. + train one selected autoregressive stage during RL post-training. Other + components are stripped before FSDP wrapping to save memory. Qwen3-Omni + trains the thinker; adapters for other architectures may select a different + autoregressive stage. - **Encoder-frozen**: Vision/audio encoders are typically frozen during RL training (`freeze_vision_tower=True`). The training adapter's `get_strip_modules` excludes them from the trainable set if they are separate @@ -60,6 +60,13 @@ adapt each implementation to your model's architecture: `module._no_split_modules` to the correct decoder layer class for FSDP. This method runs before FSDP wrapping and LoRA injection. +- **`register_auto_classes()`** (optional): Register classes supplied by an + optional model package with the appropriate Transformers Auto APIs. The model + config resolves one `(architecture, stage)` adapter before calling this hook; + the base implementation is a no-op. Set the adapter's `auto_model_class` when + the default `AutoModelForMultimodalLM` loader does not own the architecture; + the FSDP engine still owns `from_pretrained`. + - **`prepare_model_inputs(model_inputs, micro_batch, model_config)`** (optional): Validate model-native trajectory or conditioning data retained by rollout and add it to the actor forward inputs. Per-sample rollout data starts @@ -93,10 +100,29 @@ and implement: - **`get_pipeline_id(pipeline_mode)`**: Return the vLLM-Omni pipeline `model_type` string, used when auto-generating the deploy config YAML. -Optional overrides: `ensure_pipeline_registered` (register non-standard -pipeline variants with vLLM-Omni), `get_engine_hf_overrides` (HF config -overrides like `enable_audio_output: false`), `get_stage_engine_extras` -(per-stage overrides like `model_arch`). +Optional overrides fall into four groups: + +- Pipeline setup: `ensure_pipeline_registered`, `get_engine_hf_overrides`, and + `get_stage_engine_extras`. +- Policy and resource behavior: `policy_stage_id` identifies the stage whose + sampling parameters and logprobs define the trained policy; + `weight_sync_stage_ids` identifies the stages that receive actor weights. +- Request construction: `prepare_engine_prompt`. When this hook returns a + custom prompt, the adapter must include any non-`None` + `mm_processor_kwargs`; the shared strategy adds them automatically only to + its default prompt. +- Multi-stage output assembly: override `combine_engine_outputs` to opt into + retaining outputs from every stage marked `final_output` in the pipeline + topology. Adapters that keep the default hook preserve the engine's existing + single-output behavior. A custom combiner must also handle abort outputs with + empty token IDs and, pending + [vllm-omni#6973](https://github.com/vllm-project/vllm-omni/issues/6973), an + empty output list. + +Their defaults preserve the existing single-output AR behavior. Override only +the hooks required by the model. A stage-split adapter may, for example, limit +actor weight synchronization to its trainable stage while retaining outputs +from both the policy and decoder stages. When training an omni model's autoregressive Talker stage, also override `postprocess_agent_loop_output`. Put the sampled policy sequence in @@ -234,3 +260,10 @@ model-specific — verify each against your own model's architecture. in `configure_tokenizer` and assign it to `tokenizer.chat_template`. verl's dataset loader calls `tokenizer.apply_chat_template()` and will fail without a template. + +- **Actor/rollout probability consistency**: Autoregressive codec policies may + combine several codebook embeddings before predicting the selected token. + Match actor, reference, rollout, and weight-sync dtypes, then verify selected + token log-probabilities before training. Treat numerical comparisons as + execution-consistency diagnostics, not evidence of output quality or bitwise + agreement between different precision paths. diff --git a/docs/examples/qwen_image/dmd2_trainer.md b/docs/examples/qwen_image/dmd2_trainer.md new file mode 120000 index 000000000..9b0d25ee9 --- /dev/null +++ b/docs/examples/qwen_image/dmd2_trainer.md @@ -0,0 +1 @@ +../../../examples/dmd2_trainer/qwen_image/README.md \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index bcf88ebce..25a1341b5 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,7 +1,7 @@ # Welcome to VeRL-Omni's documentation! -Last updated: 09/02/2026 +Last updated: 09/11/2026 [VeRL-Omni](https://github.com/verl-project/verl-omni) is a general RL training framework focused on multimodal generative models, built on top of [verl](https://github.com/verl-project/verl). It originated from the multi-modal generation RL effort in `verl`, and now has a dedicated home so it can evolve in a more focused way. @@ -67,6 +67,7 @@ algo/diffusionnft.md algo/grpo_guard.md algo/mixgrpo.md algo/diffusion_opd.md +algo/diffusion_distillation.md algo/performance.md ``` @@ -85,6 +86,7 @@ examples/grpoguard_trainer.md examples/gspo_trainer.md examples/mixgrpo_trainer.md examples/diffusionopd_trainer.md +examples/qwen_image/dmd2_trainer.md examples/flowgrpo_trainer_sd35_drm.md examples/bagel/flowgrpo_trainer_bagel.md examples/qwen_image_edit/flowgrpo_trainer_qwen_image_edit.md diff --git a/examples/dmd2_trainer/qwen_image/README.md b/examples/dmd2_trainer/qwen_image/README.md new file mode 100644 index 000000000..5489c6860 --- /dev/null +++ b/examples/dmd2_trainer/qwen_image/README.md @@ -0,0 +1,194 @@ +# Qwen-Image DMD2 distribution-only + +Last updated: 09/11/2026. + +See the [algorithm and runtime contract](../../../docs/algo/diffusion_distillation.md) +for the objectives, role ownership, configuration and checkpoint semantics. + +This MVP trains a conditional-only few-step **Qwen-Image T2I student** from prompts. +It is **DMD2 distribution-only**, not original DMD: there are no paired teacher +trajectories, LPIPS regression, discriminator, GAN, rewards or PPO advantages. +Qwen-Image Edit, causal video and Self-Forcing are not included. + +## Train + +Use the repository's GPU + training installation and a base `Qwen/Qwen-Image` +checkpoint. Prompt parquet follows the existing `RLHFDataset` schema: + +```python +{"prompt": [{"role": "user", "content": "A red apple on a wooden table"}]} +``` + +The frozen encoder applies the checkpoint's Qwen system template. Do not add a +custom system message: raw inputs accept a string or one text-only user message. +A custom dataset may instead return `prompt_embeds` `[B,L,D]`, optional +`prompt_embeds_mask` `[B,L]`, and matching negative embeddings for teacher scoring. +Pre-tokenized inputs must include attention masks and the Qwen template prefix. + +```bash +MODEL_PATH=/path/to/Qwen-Image \ +TRAIN_FILES=/path/to/train.parquet \ +VAL_FILES=/path/to/test.parquet \ +OUTPUT_DIR=outputs/qwen_dmd2 \ +NUM_GPUS=8 TOTAL_TRAIN_STEPS=1000 \ +bash examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh +``` + +The script selects the new route: + +```yaml +algorithm: + trainer_type: distribution_matching + sample_source: offline +actor_rollout_ref: + model: + algorithm: dmd2 + model_type: diffusion_dmd_model +``` + +Here `sample_source: offline` means **engine-local sampling**, not offline RL or +training on pre-generated images. The current student generates fresh samples +from prompts and noise during training inside FSDP, retaining the graph needed +for its objective. No independent vLLM rollout server or reward workers are +started. Keep this configuration value `offline`; it does not make the student +samples precomputed. + +`dmd` is a separate top-level configuration group. Do **not** enable the existing +OPD `distillation.enabled` or actor `use_distill_loss` flags. Existing +policy-gradient, direct-preference and OPD routing is unchanged. + +Defaults: 1024×1024, four student steps, LoRA rank/alpha 32, student LR `1e-4`, +fake-score LR `2e-5`, fake update ratio 2, teacher CFG 4 with `layer_norm`, negative +prompt `" "`, and EMA decay 0.999. Student and fake-score are conditional-only. +Teacher CFG is computed in **packed velocity space**, before conversion to x0. +Sampling uses fixed sigmas `[1, .9, .75, .5, 0]`, not Qwen's native +resolution-dependent shift. The text encoder is frozen; no training VAE is loaded. + +## Execution and accounting + +`DistributionMatchingRayTrainer` reuses `BaseRayDiffusionTrainer`'s offline +initialization, dataloaders and profiling. `DMDTrainingWorker` subclasses the +existing `TrainingWorker`; `DMDDiffusersFSDPEngine` reuses its model loading, FSDP, +LoRA and checkpoint services. The Qwen adapter remains a stateless registry class. + +Each **cycle** fetches one fresh student batch, then K fresh fake-score batches. +Each role call attempts at most one optimizer update. `training/global_step` +counts completed cycles; per-role optimizer counters count successful updates. +All ranks agree on nonfinite skips. A skipped role advances neither its scheduler +nor EMA. Even an all-skipped cycle consumes the finite budget; a run with no +successful updates for either role ends with an error rather than a success claim. +Unexpected failures stop the run: partially applied GPU updates are not rolled +back, and recovery requires the last complete checkpoint in a new trainer. + +One frozen base holds independently optimized `default` (student) and +`fake_score` adapters, a non-optimized `student_ema`, and adapter-disabled teacher +scoring. This MVP requires LoRA; independent full modules and external teachers +are deferred, not mathematical requirements of DMD2. FSDP1 requires +`use_orig_params=true`; FSDP2 is the default. The example's attention LoRA targets +are fully covered by layer-wise export; custom targets must also be covered by +`model.fsdp_layer_prefixes`, otherwise export fails closed. + +## Batching and diagnostics + +`GLOBAL_BATCH_SIZE` defaults to `NUM_GPUS`. Each role independently uses +`STUDENT_MICRO_BATCH_SIZE` / `FAKE_MICRO_BATCH_SIZE` (default 1). Accumulation uses +the existing worker and microbatch splitter. For a non-divisible tail, it uses +native TensorDict splitting, with sample-weighted losses and identical chunk +counts across DP ranks; it does not pad in extra training samples. +Physical batches must have homogeneous image geometry. For example, eight DP +ranks and physical batch 2 use global batch 16. The baseline validation uses SP=1. + +The rollout exit is synchronized across FSDP ranks. Physical batching and +accumulation can sample different rollout depths; scalar CFG additionally uses a +batch-wide norm. Neither is a controlled performance comparison without fixing +those inputs. + +Metrics are phase-qualified (`student/0/...`, `fake_score/0/...`, +`fake_score/1/...`). Within a role attempt, losses are sample-weighted and host +component durations/counts are summed over microbatches; DP aggregation reports +rank means, not slowest-rank wall time. Cycle/checkpoint timings are fresh each +cycle. Component timings overlap and are not isolated CUDA kernel costs. The +inherited MFU fallback may report zero without diffusion FLOPs metadata. Peak +allocated/reserved memory reports the maximum across ranks, cumulative since +worker initialization (not a separately reset peak for each phase). + +Use the existing `global_profiler.steps` and +`actor_rollout_ref.actor.profiler` settings for Torch traces. No new profiler or +inference server is required for training. + +## Resume and inference + +The two user-facing outputs are a complete training checkpoint and one selected +inference artifact: + +```text +OUTPUT_DIR/ + global_step_N/ + actor/ # shared model, student optimizer, fake optimizer/scheduler, + # EMA, per-rank RNG and successful/skipped counters + data.pt + trainer.pt # completed cycle, driver RNG and configuration fingerprint + latest_checkpointed_iteration.txt + inference/ + adapter_model.safetensors + adapter_config.json + inference_manifest.json +``` + +Saving is synchronous and atomic on a shared local filesystem. All ranks and +role clocks are checked before publication and restore. `max_actor_ckpt_to_keep` +controls retention of this run's complete checkpoints. Restore requires model, +optimizer and extra state. Legacy generic-distillation checkpoints are rejected; +there is no implicit format/counter migration. Mathematical/data/optimizer +configuration drift (including a changed configured training horizon) is rejected. +Existing experimental checkpoints are not converted or modified. + +The script defaults to `RESUME_MODE=auto`. To replay from a chosen checkpoint, +keep the same data/model/training configuration and use a fresh output directory: + +```bash +# Append to the training command above: +# trainer.resume_mode=resume_path \ +# trainer.resume_from_path=/path/to/global_step_N +``` + +Export defaults to **student**. EMA remains resumable smoothing state; +`dmd.export_role=student_ema` explicitly selects it instead. The inference artifact +is a base-dependent PEFT LoRA, not a merged self-contained pipeline. Its manifest +records role, step, resolved base revision, transformer-config hash, sampler, +resolution and weight checksum. Unversioned local bases record a null revision; +use an immutable matching base rather than treating the config hash as a weight +checksum. Generation +reuses Qwen input preparation and the exact fp32 training Euler grid: + +```bash +python examples/dmd2_trainer/qwen_image/generate.py \ + --artifact outputs/qwen_dmd2/inference \ + --prompt 'A red apple on a wooden table' \ + --seed 42 --output outputs/apple.png +``` + +Use `--base-model` if the matching base checkpoint moved. Stock pipeline defaults +must not silently replace the recorded fixed-shift/conditional-only schedule. +Automatic CheckpointEngine validation replicas, vLLM request batching and new +transports are deliberately deferred. A completed smoke or decoded image proves +execution, not generation-quality improvement. + +## Validation commands + +```bash +# CPU suite (includes existing OPD regressions) +python -m pytest -o 'python_files=*_on_cpu.py' --asyncio-mode=auto tests/ + +# Real two-rank FSDP1/FSDP2 tiny-model updates, EMA, resume and adapter reload +QWEN_IMAGE_MODEL_PATH=/path/to/tiny-random/Qwen-Image \ +torchrun --standalone --nproc_per_node=2 -m pytest -q tests/workers/test_dmd_fsdp.py + +# Full Ray production entrypoint: three student / six fake-score updates +MODEL_PATH=/path/to/tiny-random/Qwen-Image NUM_GPUS=2 \ +bash tests/special_e2e/run_dmd2_qwen_image.sh +``` + +The tiny model is an execution fixture, not a useful generator. Real-model smoke, +controlled prototype comparison and resume results belong in the PR's validation +evidence, with their actual scope and hardware, not inferred from CPU tests. diff --git a/examples/dmd2_trainer/qwen_image/generate.py b/examples/dmd2_trainer/qwen_image/generate.py new file mode 100644 index 000000000..43de24143 --- /dev/null +++ b/examples/dmd2_trainer/qwen_image/generate.py @@ -0,0 +1,114 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Generate an image from a DMD2 student export with the training-matched fp32 Euler path.""" + +import argparse +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu +from verl.utils.device import get_device_name + +from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import ( + QwenImageConditionProvider, + QwenImageDMD2, + build_qwen_dmd_sigmas, +) +from verl_omni.trainer.diffusion.distillation.utils import ode_euler_step +from verl_omni.utils.fs import diffusion_model_provenance, resolve_model_local_dir +from verl_omni.workers.engine.lora_adapter_mixin import load_diffusers_lora_adapter + + +@torch.inference_mode() +def generate(artifact, prompt, seed, base_model=None, device=None): + """Load a selected-adapter artifact and return one decoded RGB image.""" + from diffusers import QwenImagePipeline + + artifact = Path(artifact) + metadata = json.loads((artifact / "inference_manifest.json").read_text()) + if metadata["algorithm"] != "dmd2" or metadata["sampler"] != "ode_euler" or metadata["guidance_scale"] != 1.0: + raise ValueError("This generator requires a conditional-only Euler DMD2 export.") + with (artifact / "adapter_model.safetensors").open("rb") as file: + if hashlib.file_digest(file, "sha256").hexdigest() != metadata["weights_sha256"]: + raise ValueError("Student artifact weight checksum does not match the manifest.") + base_model = resolve_model_local_dir(base_model or metadata["base_model"]) + if "base_transformer_config_sha256" in metadata: + provenance = diffusion_model_provenance(base_model) + if provenance["base_transformer_config_sha256"] != metadata["base_transformer_config_sha256"]: + raise ValueError("Base transformer configuration does not match the student export.") + if ( + metadata["base_model_revision"] is not None + and provenance["base_model_revision"] != metadata["base_model_revision"] + ): + raise ValueError("Base checkpoint revision does not match the student export.") + device = torch.device(device or get_device_name()) + pipeline = QwenImagePipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16).to(device) + pipeline.transformer.requires_grad_(False) + pipeline.text_encoder.requires_grad_(False) + pipeline.vae.requires_grad_(False) + load_diffusers_lora_adapter(pipeline.transformer, artifact, "student") + pipeline.transformer.set_adapter("student") + pipeline.transformer.eval() + pipeline.transformer.set_attention_backend("native") + config = SimpleNamespace( + local_path=str(base_model), + path=str(base_model), + pipeline=SimpleNamespace(height=metadata["height"], width=metadata["width"], guidance_scale=None), + ) + batch = TensorDict({}, batch_size=[1]) + tu.assign_non_tensor_stack(batch, "raw_prompt", [prompt]) + provider = QwenImageConditionProvider(str(base_model), metadata["max_sequence_length"], " ") + provider.pipeline = pipeline + condition, _ = provider.encode(batch, device=device, dtype=torch.bfloat16, require_negative=False) + shape, geometry = QwenImageDMD2.latent_geometry(pipeline.transformer, config, batch) + generator = torch.Generator(device=device).manual_seed(seed) + latents = QwenImageDMD2.pack_latents(torch.randn(shape, generator=generator, device=device, dtype=torch.float32)) + sigmas = build_qwen_dmd_sigmas(metadata["num_inference_steps"], metadata["rollout_timestep_shift"], device) + for current, following in zip(sigmas[:-1], sigmas[1:], strict=True): + inputs = QwenImageDMD2.prepare_dmd_inputs(pipeline.transformer, config, latents, current, condition, geometry) + velocity = QwenImageDMD2.forward(pipeline.transformer, config, inputs) + latents = ode_euler_step(latents, velocity, current, following) + unpacked = QwenImagePipeline._unpack_latents( + latents, metadata["height"], metadata["width"], geometry["vae_scale_factor"] + ) + mean = torch.tensor(pipeline.vae.config.latents_mean, device=device).reshape(1, -1, 1, 1, 1) + std = torch.tensor(pipeline.vae.config.latents_std, device=device).reshape(1, -1, 1, 1, 1) + decoded = pipeline.vae.decode((unpacked * std + mean).to(pipeline.vae.dtype), return_dict=False)[0][:, :, 0] + return pipeline.image_processor.postprocess(decoded, output_type="pil")[0] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--artifact", required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--base-model", help="Matching base checkpoint if moved from its recorded path") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device") + args = parser.parse_args() + output = Path(args.output) + if output.exists(): + raise FileExistsError(output) + image = generate(args.artifact, args.prompt, args.seed, args.base_model, args.device) + output.parent.mkdir(parents=True, exist_ok=True) + image.save(output) + print(f"Saved {output}") + + +if __name__ == "__main__": + main() diff --git a/examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh b/examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh new file mode 100644 index 000000000..384848d3b --- /dev/null +++ b/examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen-Image} +NUM_GPUS=${NUM_GPUS:-8} +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-${NUM_GPUS}} +OUTPUT_DIR=${OUTPUT_DIR:-outputs/qwen_image_dmd2} +TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1000} + +python3 -m verl_omni.trainer.main_diffusion \ + algorithm.trainer_type=distribution_matching \ + algorithm.sample_source=offline \ + data.train_files="${TRAIN_FILES:?Set TRAIN_FILES to prompt parquet}" \ + data.val_files="${VAL_FILES:?Set VAL_FILES to prompt parquet}" \ + data.train_batch_size="${GLOBAL_BATCH_SIZE}" \ + data.dataloader_num_workers=0 \ + actor_rollout_ref.model.path="${MODEL_PATH}" \ + actor_rollout_ref.model.algorithm=dmd2 \ + actor_rollout_ref.model.model_type=diffusion_dmd_model \ + actor_rollout_ref.model.attn_backend=native \ + actor_rollout_ref.rollout.rollout_attn_backend=TORCH_SDPA \ + actor_rollout_ref.model.lora_rank=32 \ + actor_rollout_ref.model.lora_alpha=32 \ + actor_rollout_ref.model.target_modules='[to_q,to_k,to_v,to_out.0]' \ + actor_rollout_ref.model.pipeline.height=1024 \ + actor_rollout_ref.model.pipeline.width=1024 \ + actor_rollout_ref.model.pipeline.num_inference_steps=4 \ + actor_rollout_ref.model.pipeline.max_sequence_length=1024 \ + actor_rollout_ref.actor.strategy="${STRATEGY:-fsdp2}" \ + actor_rollout_ref.actor.fsdp_config.use_orig_params=true \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.optim.lr=1e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.001 \ + dmd.fake_update_ratio=2 \ + dmd.student_micro_batch_size_per_gpu="${STUDENT_MICRO_BATCH_SIZE:-1}" \ + dmd.fake_score_micro_batch_size_per_gpu="${FAKE_MICRO_BATCH_SIZE:-1}" \ + dmd.export_role=student \ + trainer.logger='[console,tensorboard]' \ + trainer.project_name=qwen-image-dmd2 \ + trainer.experiment_name="${EXPERIMENT_NAME:-distribution-only}" \ + trainer.n_gpus_per_node="${NUM_GPUS}" \ + trainer.nnodes=1 \ + trainer.val_before_train=false \ + trainer.test_freq=-1 \ + trainer.save_freq="${SAVE_FREQ:-100}" \ + trainer.default_local_dir="${OUTPUT_DIR}" \ + trainer.resume_mode="${RESUME_MODE:-auto}" \ + trainer.total_training_steps="${TOTAL_TRAIN_STEPS}" \ + ray_kwargs.ray_init.num_cpus="${RAY_NUM_CPUS:-32}" \ + "$@" diff --git a/tests/agent_loop/test_composite_agent_loop.py b/tests/agent_loop/test_composite_agent_loop.py index b7234940b..7ef37d2e4 100644 --- a/tests/agent_loop/test_composite_agent_loop.py +++ b/tests/agent_loop/test_composite_agent_loop.py @@ -78,14 +78,19 @@ def _assert_text_encoder_outputs(result: DataProto, *, batch_size: int, max_toke """Validate Qwen-Image text-encoder returns by rollout.""" llm_response_ids = result.batch["llm_response_ids"] llm_all_log_probs = result.batch.get("rollout_llm_log_probs") + llm_attention_mask = result.batch["llm_response_attention_mask"] text_encoder_responses = result.non_tensor_batch["text_encoder_responses"] # list[str] _assert_non_empty_tensor(llm_response_ids, "llm_response_ids") _assert_non_empty_tensor(llm_all_log_probs, "llm_all_log_probs") + _assert_non_empty_tensor(llm_attention_mask, "llm_response_attention_mask") assert llm_response_ids.shape == (batch_size, max_token_len) if llm_all_log_probs is not None: assert llm_all_log_probs.shape[1] <= max_token_len assert llm_all_log_probs.shape == (batch_size, llm_all_log_probs.shape[1], llm_all_log_probs.shape[-1]) + assert llm_attention_mask.shape == (batch_size, max_token_len) + assert llm_attention_mask.dtype == torch.long + assert llm_attention_mask.min() >= 0 and llm_attention_mask.max() <= 1 assert len(text_encoder_responses) == batch_size @@ -267,6 +272,7 @@ def test_single_turn(init_config, agent_reward_loop: bool): "negative_prompt_embeds_mask", "rollout_log_probs", "rollout_llm_log_probs", + "llm_response_attention_mask", ] expected_non_tensor_batch_keys = ["text_encoder_responses"] if agent_reward_loop: diff --git a/tests/agent_loop/test_composite_agent_loop_padding_on_cpu.py b/tests/agent_loop/test_composite_agent_loop_padding_on_cpu.py new file mode 100644 index 000000000..e3fcb8b88 --- /dev/null +++ b/tests/agent_loop/test_composite_agent_loop_padding_on_cpu.py @@ -0,0 +1,63 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU tests for the dual-GRPO AR generation padding.""" + +import pytest +import torch + +from verl_omni.agent_loop.composite_agent_loop import _pad_llm_generation_outputs + + +def test_pads_ragged_llm_responses_to_max_new_tokens(): + early_eos_ids = torch.tensor([[101, 2054, 8667, 102]]) # stopped at 4 tokens + full_ids = torch.tensor([[101, 2054, 8667, 2055, 6844, 2900]]) # hit max_new_tokens + early_eos_log_probs = torch.randn(1, 4, 11) + + padded_early, mask_early, padded_early_lp = _pad_llm_generation_outputs( + early_eos_ids, early_eos_log_probs, max_new_tokens=6, pad_token_id=0 + ) + padded_full, mask_full, padded_full_lp = _pad_llm_generation_outputs( + full_ids, None, max_new_tokens=6, pad_token_id=0 + ) + + assert padded_early.shape == (1, 6) + assert padded_full.shape == (1, 6) + torch.testing.assert_close(padded_early[0, :4], early_eos_ids[0]) + torch.testing.assert_close(padded_early[0, 4:], torch.zeros(2, dtype=torch.long)) + assert mask_early.tolist() == [[1, 1, 1, 1, 0, 0]] + assert mask_full.tolist() == [[1, 1, 1, 1, 1, 1]] + assert padded_early_lp.shape == (1, 6, 11) + torch.testing.assert_close(padded_early_lp[0, :4], early_eos_log_probs[0]) + torch.testing.assert_close(padded_early_lp[0, 4:], torch.zeros(2, 11)) + assert padded_full_lp is None + + # differently sized samples must now batch + batched_ids = torch.cat([padded_early, padded_full], dim=0) + assert batched_ids.shape == (2, 6) + batched_mask = torch.cat([mask_early, mask_full], dim=0) + assert batched_mask.shape == (2, 6) + + +def test_full_length_response_passes_through_unchanged(): + ids = torch.tensor([[5, 6, 7]]) + log_probs = torch.randn(1, 3, 4) + padded_ids, mask, padded_log_probs = _pad_llm_generation_outputs(ids, log_probs, 3, pad_token_id=0) + torch.testing.assert_close(padded_ids, ids) + torch.testing.assert_close(padded_log_probs, log_probs) + assert mask.tolist() == [[1, 1, 1]] + + +def test_rejects_response_longer_than_max_new_tokens(): + with pytest.raises(ValueError, match="max_new_tokens"): + _pad_llm_generation_outputs(torch.tensor([[1, 2, 3]]), None, 2, pad_token_id=0) diff --git a/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh b/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh index 31c0bad1b..32c003fdd 100644 --- a/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh +++ b/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh @@ -74,4 +74,8 @@ run_test 12 "Diffusion OPD v1 separate_async one_step_off teachers e2e" \ env CUDA_VISIBLE_DEVICES="${CUDA_DEVICE_LIST}" NUM_GPUS="${NUM_GPUS}" V1=1 V1_MODE=separate_async SMOKE=standalone SCHEDULER=one_step_off \ bash tests/special_e2e/run_diffusion_teacher_smoke.sh +run_test 13 "Qwen-Image offline DMD2 trainer e2e" \ + env CUDA_VISIBLE_DEVICES="${CUDA_DEVICE_LIST}" NUM_GPUS="${NUM_GPUS}" \ + bash tests/special_e2e/run_dmd2_qwen_image.sh "${diffusion_trainer_args[@]}" + gpu_smoke_summary diff --git a/tests/pipelines/test_qwen_dmd2_on_cpu.py b/tests/pipelines/test_qwen_dmd2_on_cpu.py new file mode 100644 index 000000000..881a48279 --- /dev/null +++ b/tests/pipelines/test_qwen_dmd2_on_cpu.py @@ -0,0 +1,113 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.pipelines.model_base import DiffusionModelBase +from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import ( + QwenImageConditionProvider, + QwenImageDMD2, +) + + +class ToyPromptTokenizer: + def __call__(self, texts, **kwargs): + width = max(len(text) for text in texts) + ids = torch.tensor([[len(text)] * len(text) + [0] * (width - len(text)) for text in texts]) + return SimpleNamespace(input_ids=ids, attention_mask=ids.ne(0).long()) + + +class ToyTextEncoder(torch.nn.Module): + dtype = torch.float32 + + def forward(self, input_ids, **kwargs): + return SimpleNamespace(hidden_states=(input_ids.float().unsqueeze(-1),)) + + +class ToyConditionPipeline: + prompt_template_encode = "x" * 34 + "{}" + prompt_template_encode_start_idx = 34 + device = torch.device("cpu") + + def __init__(self): + from diffusers import QwenImagePipeline + + self.tokenizer = ToyPromptTokenizer() + self.text_encoder = ToyTextEncoder() + self._extract_masked_hidden = QwenImagePipeline._extract_masked_hidden.__get__(self) + + +class TestQwenDMD2: + def test_registry_does_not_claim_original_dmd_or_edit_support(self): + assert DiffusionModelBase.get_class_by_name("QwenImagePipeline", "dmd2") is QwenImageDMD2 + with pytest.raises(NotImplementedError): + DiffusionModelBase.get_class_by_name("QwenImagePipeline", "dmd") + with pytest.raises(NotImplementedError): + DiffusionModelBase.get_class_by_name("QwenImageEditPlusPipeline", "dmd2") + + def test_short_prompts_keep_tokens_after_real_prefix_removal(self): + provider = QwenImageConditionProvider("unused", 64, " ") + provider.pipeline = ToyConditionPipeline() + for size in (1, 3, 2): + batch = TensorDict({"dummy_tensor": torch.zeros(size, 1)}, batch_size=[size]) + tu.assign_non_tensor_stack(batch, "raw_prompt", [[{"role": "user", "content": "cat"}]] * size) + positive, negative = provider.encode( + batch, device=torch.device("cpu"), dtype=torch.float32, require_negative=True + ) + assert positive["prompt_embeds"].shape == (size, 3, 1) + assert negative["prompt_embeds"].shape == (size, 1, 1) + assert positive["prompt_embeds"][0, 0, 0] == 37 + assert not positive["prompt_embeds"].requires_grad + + @pytest.mark.parametrize( + "row", [[], [{"role": "system", "content": "custom"}], [{"role": "assistant", "content": "cat"}]] + ) + def test_invalid_chat_is_not_generic_chat_formatted(self, row): + provider = QwenImageConditionProvider("unused", 64, " ") + with pytest.raises(ValueError, match="single user"): + provider.tokenize_rows(Mock(), [row], torch.device("cpu")) + + def test_precomputed_inputs_do_not_load_encoder_for_fake_stage(self): + provider = QwenImageConditionProvider("unused", 2, " ") + batch = TensorDict({"prompt_embeds": torch.ones(2, 3, 4, requires_grad=True)}, batch_size=[2]) + positive, negative = provider.encode( + batch, device=torch.device("cpu"), dtype=torch.float32, require_negative=False + ) + assert provider.pipeline is None and negative is None + assert positive["prompt_embeds"].shape == (2, 2, 4) + assert positive["prompt_embeds_mask"].shape == (2, 2) + assert not positive["prompt_embeds"].requires_grad + with pytest.raises(ValueError, match="negative_prompt_embeds"): + provider.encode(batch, device=torch.device("cpu"), dtype=torch.float32, require_negative=True) + + def test_geometry_uses_vae_config_and_packs_consistently(self, tmp_path): + (tmp_path / "vae").mkdir() + (tmp_path / "vae" / "config.json").write_text(json.dumps({"z_dim": 4, "temperal_downsample": [False, True]})) + model = SimpleNamespace(config=SimpleNamespace(in_channels=16)) + config = SimpleNamespace(local_path=str(tmp_path), pipeline=SimpleNamespace(height=32, width=48)) + shape, geometry = QwenImageDMD2.latent_geometry(model, config, TensorDict({}, batch_size=[2])) + assert shape == (2, 4, 1, 8, 12) + assert geometry["vae_scale_factor"] == 4 + latents = QwenImageDMD2.pack_latents(torch.ones(shape)) + assert latents.shape == (2, 24, 16) + batch = TensorDict({"height": torch.tensor([32, 64])}, batch_size=[2]) + with pytest.raises(ValueError, match="homogeneous"): + QwenImageDMD2.latent_geometry(model, config, batch) diff --git a/tests/special_e2e/create_dummy_diffusion_data.py b/tests/special_e2e/create_dummy_diffusion_data.py index 90778120b..6a71b71c8 100644 --- a/tests/special_e2e/create_dummy_diffusion_data.py +++ b/tests/special_e2e/create_dummy_diffusion_data.py @@ -41,21 +41,16 @@ ] -def build_rows(split: str, n: int, data_sources: list[str]): +def build_rows(split: str, n: int, data_sources: list[str], user_prompt_only: bool = False): rows = [] + prefix = [] if user_prompt_only else [{"role": "system", "content": SYSTEM_PROMPT}] for i in range(n): prompt_text = USER_PROMPTS[i % len(USER_PROMPTS)] rows.append( { "data_source": data_sources[i % len(data_sources)], - "prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": prompt_text}, - ], - "negative_prompt": [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": " "}, - ], + "prompt": prefix + [{"role": "user", "content": prompt_text}], + "negative_prompt": prefix + [{"role": "user", "content": " "}], "reward_model": {"style": "rule", "ground_truth": ""}, "extra_info": {"split": split, "index": i}, } @@ -77,13 +72,16 @@ def main(): default="jpeg_compressibility", help="Comma-separated data_source values assigned to rows in round-robin order", ) + parser.add_argument( + "--user_prompt_only", action="store_true", help="Leave model-specific system templating to the encoder" + ) args = parser.parse_args() data_sources = args.data_sources.split(",") os.makedirs(args.local_save_dir, exist_ok=True) - train_df = pd.DataFrame(build_rows("train", args.train_size, data_sources)) - val_df = pd.DataFrame(build_rows("test", args.val_size, data_sources)) + train_df = pd.DataFrame(build_rows("train", args.train_size, data_sources, args.user_prompt_only)) + val_df = pd.DataFrame(build_rows("test", args.val_size, data_sources, args.user_prompt_only)) train_path = os.path.join(args.local_save_dir, "train.parquet") val_path = os.path.join(args.local_save_dir, "test.parquet") diff --git a/tests/special_e2e/run_dmd2_qwen_image.sh b/tests/special_e2e/run_dmd2_qwen_image.sh new file mode 100644 index 000000000..78127d88a --- /dev/null +++ b/tests/special_e2e/run_dmd2_qwen_image.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Multi-GPU DMD2 production smoke, including atomic checkpoints and student export. +set -euo pipefail + +export NUM_GPUS=${NUM_GPUS:-2} +export TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-3} +export MODEL_PATH=${MODEL_PATH:-${HOME}/models/tiny-random/Qwen-Image} +export OUTPUT_DIR=${OUTPUT_DIR:-outputs/dmd2_smoke} +DATA_DIR=${DATA_DIR:-${OUTPUT_DIR}/data} + +python3 tests/special_e2e/create_dummy_diffusion_data.py \ + --local_save_dir "${DATA_DIR}" \ + --train_size "$((NUM_GPUS * TOTAL_TRAIN_STEPS * 3))" \ + --val_size "${NUM_GPUS}" \ + --user_prompt_only + +export TRAIN_FILES=${DATA_DIR}/train.parquet +export VAL_FILES=${DATA_DIR}/test.parquet +export SAVE_FREQ=1 +export RESUME_MODE=${RESUME_MODE:-disable} + +bash examples/dmd2_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + actor_rollout_ref.model.lora_rank=2 \ + actor_rollout_ref.model.lora_alpha=2 \ + actor_rollout_ref.model.pipeline.height=64 \ + actor_rollout_ref.model.pipeline.width=64 \ + actor_rollout_ref.model.pipeline.max_sequence_length=64 \ + data.max_prompt_length=64 \ + trainer.logger=console \ + "$@" + +test -f "${OUTPUT_DIR}/inference/adapter_model.safetensors" +test -f "${OUTPUT_DIR}/global_step_${TOTAL_TRAIN_STEPS}/trainer.pt" +echo 'DMD2 training, checkpoint and student export completed.' diff --git a/tests/special_sanity/check_device_api_usage.py b/tests/special_sanity/check_device_api_usage.py index 0f924a8fa..de04e2562 100644 --- a/tests/special_sanity/check_device_api_usage.py +++ b/tests/special_sanity/check_device_api_usage.py @@ -25,6 +25,7 @@ # directory or file path must contain keyword ".cuda" or "cuda" CUDA_KEYWORD_CHECK_WHITELIST = [ "verl_omni/workers/engine/fsdp/diffusers_impl.py", # appear in default device_name + "verl_omni/workers/engine/fsdp/dmd_impl.py", # device=[...] registry declaration "verl_omni/trainer/diffusion/ray_diffusion_trainer.py", # appear in default device_name "verl_omni/workers/engine/fsdp/omni_impl.py", # device=[...] registry declaration "verl_omni/workers/engine/veomni/diffusion_impl.py", # device=[...] registry declaration diff --git a/tests/trainer/diffusion/test_distillation_config_on_cpu.py b/tests/trainer/diffusion/test_distillation_config_on_cpu.py index a1d4788c9..71cd63900 100644 --- a/tests/trainer/diffusion/test_distillation_config_on_cpu.py +++ b/tests/trainer/diffusion/test_distillation_config_on_cpu.py @@ -13,7 +13,6 @@ # limitations under the License. """CPU tests for the diffusion on-policy distillation config.""" -import dataclasses import os import pytest @@ -23,49 +22,11 @@ from verl_omni.workers.config.diffusion import ( DiffusionDistillationConfig, DiffusionDistillationTeacherModelConfig, - DiffusionDistributionMatchingConfig, ) CONFIG_DIR = os.path.join(os.path.dirname(os.path.abspath(verl_omni.__file__)), "trainer", "config") -class TestDistributionMatchingConfig: - def test_defaults_select_dmd2_without_enabling_opd(self): - config = DiffusionDistillationConfig() - assert config.enabled is False - assert config.distribution_matching.recipe == "dmd2" - assert config.distribution_matching.profile is None - assert config.distribution_matching.fake_update_ratio is None - - @pytest.mark.parametrize( - "kwargs,error", - [ - ({"recipe": "typo"}, "Invalid recipe"), - ({"profile": "typo"}, "Invalid profile"), - ({"fake_update_ratio": 0}, "greater than 0"), - ({"fake_warmup_cycles": -1}, "non-negative"), - ({"rollout_strategy": "typo"}, "Invalid rollout_strategy"), - ({"data_mode": "typo"}, "Invalid data_mode"), - ({"export_role": "teacher_score"}, "Invalid export_role"), - ], - ) - def test_invalid_values_fail_closed(self, kwargs, error): - with pytest.raises(ValueError, match=error): - DiffusionDistributionMatchingConfig(**kwargs) - - def test_null_fields_use_recipe_specific_defaults(self): - from verl_omni.trainer.diffusion.distillation.recipes import build_plan - - config = dataclasses.asdict(DiffusionDistributionMatchingConfig(recipe="dmd")) - plan = build_plan( - "dmd", - {**config, "model_path": "/m"}, - frozenset({"distribution_matching"}), - ) - assert plan.objective["profile"] == "paper" - assert plan.update_schedule.phases[1].repeats == 1 - - class TestDiffusionDistillationConfig: def test_disabled_by_default_skips_validation(self): config = DiffusionDistillationConfig() @@ -178,25 +139,6 @@ def test_default_composition_is_disabled(self): config = omega_conf_to_dataclass(cfg.distillation) assert isinstance(config, DiffusionDistillationConfig) assert config.enabled is False - assert isinstance(config.distribution_matching, DiffusionDistributionMatchingConfig) - assert config.distribution_matching.recipe == "dmd2" - - @pytest.mark.parametrize("scheduler", ["inline", "one_step_off"]) - def test_opd_schedule_and_distribution_matching_settings_coexist(self, scheduler): - from verl.utils.config import omega_conf_to_dataclass - - cfg = self._compose( - [ - "distillation.enabled=true", - "distillation.teacher_models.teacher_model.model_path=/ckpt/teacher", - f"distillation.scheduler={scheduler}", - ] - ) - config = omega_conf_to_dataclass(cfg.distillation) - assert config.scheduler == scheduler - assert config.teacher_models["default"].model_path == "/ckpt/teacher" - assert isinstance(config.distribution_matching, DiffusionDistributionMatchingConfig) - assert config.distribution_matching.recipe == "dmd2" def test_cli_enable_with_teacher_path(self): from verl.utils.config import omega_conf_to_dataclass @@ -213,51 +155,6 @@ def test_cli_enable_with_teacher_path(self): assert config.n_gpus_per_node == 0 assert config.nnodes == 0 - def test_cli_distribution_matching_overrides_do_not_enable_opd(self): - from verl.utils.config import omega_conf_to_dataclass - - cfg = self._compose( - [ - "algorithm.trainer_type=distillation", - "distillation.distribution_matching.recipe=dmd2", - "distillation.distribution_matching.fake_update_ratio=2", - "distillation.distribution_matching.rollout_strategy=consistency_renoise", - ] - ) - config = omega_conf_to_dataclass(cfg.distillation) - assert config.enabled is False - assert config.distribution_matching.fake_update_ratio == 2 - assert config.distribution_matching.rollout_strategy == "consistency_renoise" - - def test_composed_config_builds_validated_plan(self): - from verl_omni.trainer.diffusion.distillation.recipes import build_plan_from_config - - cfg = self._compose( - [ - "algorithm.trainer_type=distillation", - "actor_rollout_ref.model.path=/m", - "distillation.distribution_matching.fake_update_ratio=2", - ] - ) - plan = build_plan_from_config(cfg, frozenset({"distribution_matching"})) - assert plan.name == "dmd2" - assert plan.role_layout.groups[0].model_ref == "/m" - assert plan.update_schedule.phases[1].repeats == 2 - - def test_null_overrides_use_each_recipe_default(self): - from verl_omni.trainer.diffusion.distillation.recipes import build_plan_from_config - - cfg = self._compose( - [ - "algorithm.trainer_type=distillation", - "actor_rollout_ref.model.path=/m", - "distillation.distribution_matching.recipe=dmd", - ] - ) - plan = build_plan_from_config(cfg, frozenset({"distribution_matching"})) - assert plan.objective["profile"] == "paper" - assert plan.update_schedule.phases[1].repeats == 1 - def test_cli_multi_teacher_entries(self): from verl.utils.config import omega_conf_to_dataclass diff --git a/tests/trainer/diffusion/test_distillation_contracts_on_cpu.py b/tests/trainer/diffusion/test_distillation_contracts_on_cpu.py deleted file mode 100644 index 8dc61d40f..000000000 --- a/tests/trainer/diffusion/test_distillation_contracts_on_cpu.py +++ /dev/null @@ -1,395 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""CPU tests for distillation contracts, role layouts, and recipes.""" - -import dataclasses -import pickle - -import pytest -import torch - -from verl_omni.trainer.diffusion.distillation.contracts import ( - ExportSpec, - LatentBundle, - RoleBinding, - RoleGroupSpec, - RoleLayoutSpec, - ScoreTransportSpec, - TrainerCounters, - UpdatePhaseSpec, - UpdateSchedule, - resolve_export_role, - validate_role_layout, -) -from verl_omni.trainer.diffusion.distillation.recipes import ( - DistillationRegistry, - build_plan, - initialization_registry, - objective_registry, - recipe_registry, - rollout_registry, -) - -ALL_CAPS = frozenset({"distribution_matching", "autoregressive", "adversarial"}) - - -class TestLatentBundle: - def test_rejects_empty_bundle(self): - with pytest.raises(ValueError, match="at least one modality"): - LatentBundle({}) - - def test_rejects_non_tensor(self): - with pytest.raises(TypeError, match="torch.Tensor"): - LatentBundle({"image": [1, 2, 3]}) - - def test_single_returns_the_only_tensor(self): - value = torch.randn(2, 3) - torch.testing.assert_close(LatentBundle({"image": value}).single, value) - - def test_single_rejects_multimodal_bundle(self): - bundle = LatentBundle({"video": torch.randn(1), "audio": torch.randn(1)}) - with pytest.raises(ValueError, match="one-modality"): - _ = bundle.single - - def test_map_applies_to_every_modality(self): - bundle = LatentBundle({"video": torch.ones(2), "audio": torch.ones(3)}) - doubled = bundle.map(lambda tensor: tensor * 2) - torch.testing.assert_close(doubled.get("video"), torch.full((2,), 2.0)) - torch.testing.assert_close(doubled.get("audio"), torch.full((3,), 2.0)) - - -class TestImmutability: - @pytest.mark.parametrize( - "instance,field_name", - [ - (RoleGroupSpec(name="base"), "name"), - (RoleBinding(role="student", group="base"), "role"), - (ExportSpec(), "role"), - (UpdatePhaseSpec(kind="student", trainable_roles=("student",)), "kind"), - ( - UpdateSchedule(phases=(UpdatePhaseSpec(kind="student", trainable_roles=("student",)),)), - "phases", - ), - ], - ) - def test_plan_pieces_are_frozen(self, instance, field_name): - with pytest.raises(dataclasses.FrozenInstanceError): - setattr(instance, field_name, None) - - def test_plan_is_deeply_immutable(self): - plan = build_plan("dmd2", {"model_path": "/m"}, ALL_CAPS) - with pytest.raises(dataclasses.FrozenInstanceError): - plan.name = "other" - with pytest.raises(TypeError): - plan.objective["name"] = "other" - - def test_frozen_plan_is_hashable_and_pickleable(self): - plan = build_plan("dmd2", {"model_path": "/m"}, ALL_CAPS) - assert isinstance(hash(plan), int) - restored = pickle.loads(pickle.dumps(plan)) - assert restored.objective == plan.objective - assert restored.rollout == plan.rollout - - -class TestRoleLayoutValidation: - def role_layout(self, **kwargs): - defaults = dict( - groups=(RoleGroupSpec(name="base"),), - bindings=( - RoleBinding(role="student", group="base", adapter="student", trainable=True, optimizer_key="student"), - RoleBinding(role="student_ema", group="base", adapter="student_ema"), - ), - ) - defaults.update(kwargs) - return RoleLayoutSpec(**defaults) - - def test_valid_layout_passes(self): - validate_role_layout(self.role_layout()) - - def test_binding_to_unknown_group_raises(self): - layout = self.role_layout( - bindings=( - RoleBinding( - role="student", group="missing", adapter="student", trainable=True, optimizer_key="student" - ), - ) - ) - with pytest.raises(ValueError, match="unknown group"): - validate_role_layout(layout) - - def test_trainable_role_without_optimizer_key_raises(self): - layout = self.role_layout( - bindings=(RoleBinding(role="student", group="base", adapter="student", trainable=True),) - ) - with pytest.raises(ValueError, match="optimizer_key"): - validate_role_layout(layout) - - def test_frozen_role_with_optimizer_key_raises(self): - layout = self.role_layout(bindings=(RoleBinding(role="teacher_score", group="base", optimizer_key="teacher"),)) - with pytest.raises(ValueError, match="Frozen role"): - validate_role_layout(layout) - - def test_shared_base_trainable_role_requires_adapter(self): - layout = self.role_layout( - bindings=(RoleBinding(role="student", group="base", trainable=True, optimizer_key="student"),) - ) - with pytest.raises(ValueError, match="must name an adapter"): - validate_role_layout(layout) - - def test_duplicate_adapter_in_shared_base_raises(self): - layout = self.role_layout( - bindings=( - RoleBinding(role="student", group="base", adapter="dup", trainable=True, optimizer_key="student"), - RoleBinding(role="fake_score", group="base", adapter="dup", trainable=True, optimizer_key="fake"), - ) - ) - with pytest.raises(ValueError, match="duplicate adapter names"): - validate_role_layout(layout) - - def test_duplicate_group_name_raises(self): - layout = self.role_layout(groups=(RoleGroupSpec(name="base"), RoleGroupSpec(name="base"))) - with pytest.raises(ValueError, match="Duplicate role-group"): - validate_role_layout(layout) - - def test_duplicate_optimizer_key_raises(self): - layout = self.role_layout( - bindings=( - RoleBinding(role="student", group="base", adapter="student", trainable=True, optimizer_key="same"), - RoleBinding(role="fake_score", group="base", adapter="fake", trainable=True, optimizer_key="same"), - ) - ) - with pytest.raises(ValueError, match="Duplicate optimizer_key"): - validate_role_layout(layout) - - @pytest.mark.parametrize( - "kwargs,error", - [ - ({"provider": "bogus"}, "Invalid score provider"), - ({"tensor_backend": "bogus"}, "Invalid score tensor backend"), - ({"provider": "colocated", "tensor_backend": "mooncake"}, "colocated"), - ({"provider": "ray", "tensor_backend": "local"}, "Ray score provider"), - ], - ) - def test_invalid_transport_raises(self, kwargs, error): - with pytest.raises(ValueError, match=error): - ScoreTransportSpec(**kwargs) - - -class TestExportRole: - @pytest.mark.parametrize("role", ["student", "student_ema"]) - def test_exportable_roles(self, role): - assert resolve_export_role(ExportSpec(role=role)) == role - - @pytest.mark.parametrize("role", ["teacher_score", "fake_score", "discriminator"]) - def test_non_exportable_roles_raise(self, role): - with pytest.raises(ValueError, match="Export role must be"): - ExportSpec(role=role) - - -class TestRegistry: - def test_all_recipe_and_strategy_names_are_registered(self): - assert set(recipe_registry.names) == {"dmd", "dmd2", "causvid", "self_forcing"} - assert set(objective_registry.names) == {"dmd", "dmd2", "ode_regression"} - assert set(initialization_registry.names) == {"base", "ode_regression"} - assert set(rollout_registry.names) == { - "backward_simulated", - "consistency_renoise", - "ode_euler", - "one_step", - "self_forced", - "teacher_forced_causal", - } - - def test_duplicate_registration_raises(self): - registry = DistillationRegistry() - - assert registry.register("thing")(int) is int - assert registry.get("thing") is int - with pytest.raises(ValueError, match="Duplicate registration"): - registry.register("thing")(str) - - def test_unknown_name_raises_with_registered_list(self): - with pytest.raises(KeyError, match="Registered"): - DistillationRegistry().get("nope") - - -class TestRecipePlans: - @pytest.mark.parametrize("name", ["dmd", "dmd2", "causvid", "self_forcing"]) - def test_every_recipe_builds_a_validated_plan(self, name): - config = {"model_path": "/m"} - if name == "dmd": - config["profile"] = "paper" - plan = build_plan(name, config, ALL_CAPS) - assert plan.name == name - validate_role_layout(plan.role_layout) - - def test_missing_capability_is_fail_closed(self): - with pytest.raises(ValueError, match="requires capabilities"): - build_plan("self_forcing", {"model_path": "/m"}, {"distribution_matching"}) - - def test_missing_model_reference_is_fail_closed(self): - with pytest.raises(ValueError, match="model_ref"): - build_plan("dmd2", {}, ALL_CAPS) - - def test_export_is_a_top_level_plan_contract(self): - plan = build_plan("dmd2", {"model_path": "/m", "export_role": "student"}, ALL_CAPS) - assert plan.export.role == "student" - assert not hasattr(plan.role_layout, "export") - - def test_missing_required_plan_role_is_rejected(self): - plan = build_plan("dmd2", {"model_path": "/m"}, ALL_CAPS) - layout = dataclasses.replace( - plan.role_layout, - bindings=tuple(binding for binding in plan.role_layout.bindings if binding.role != "student_ema"), - ) - with pytest.raises(ValueError, match="missing required role"): - dataclasses.replace(plan, role_layout=layout) - - def test_dmd2_paper_profile_adds_discriminator_to_layout_and_phase(self): - plan = build_plan("dmd2", {"profile": "paper", "model_path": "/m"}, ALL_CAPS) - roles = {binding.role for binding in plan.role_layout.bindings} - fake_phase = next(phase for phase in plan.update_schedule.phases if phase.kind == "fake_score") - assert "discriminator" in roles - assert fake_phase.trainable_roles == ("fake_score", "discriminator") - assert plan.objective["adversarial"] is True - - def test_dmd2_distribution_only_has_no_discriminator(self): - plan = build_plan("dmd2", {"profile": "distribution_only", "model_path": "/m"}, ALL_CAPS) - roles = {binding.role for binding in plan.role_layout.bindings} - fake_phase = next(phase for phase in plan.update_schedule.phases if phase.kind == "fake_score") - assert "discriminator" not in roles - assert fake_phase.trainable_roles == ("fake_score",) - - def test_causal_recipes_use_separate_causal_and_bidirectional_groups(self): - for name in ("causvid", "self_forcing"): - plan = build_plan(name, {"model_path": "/m"}, ALL_CAPS) - groups = {group.name for group in plan.role_layout.groups} - role_groups = {binding.role: binding.group for binding in plan.role_layout.bindings} - assert groups == {"causal_base", "bidirectional_base"} - assert role_groups["student"] == role_groups["student_ema"] == "causal_base" - assert role_groups["teacher_score"] == role_groups["fake_score"] == "bidirectional_base" - - @pytest.mark.parametrize( - "name,config,error", - [ - ("dmd2", {"profile": "typo", "model_path": "/m"}, "profile"), - ("dmd2", {"rollout_strategy": "typo", "model_path": "/m"}, "rollout_strategy"), - ("dmd2", {"data_mode": "regression_pairs", "model_path": "/m"}, "data_mode"), - ("dmd2", {"fake_update_ratio": 0, "model_path": "/m"}, "greater than zero"), - ("dmd2", {"fake_update_ratio": -2, "model_path": "/m"}, "greater than zero"), - ("dmd2", {"fake_update_ratio": 1.5, "model_path": "/m"}, "integer"), - ("dmd2", {"fake_update_ratio": True, "model_path": "/m"}, "integer"), - ("dmd2", {"fake_warmup_cycles": 1.5, "model_path": "/m"}, "integer"), - ], - ) - def test_invalid_recipe_values_fail_closed(self, name, config, error): - with pytest.raises(ValueError, match=error): - build_plan(name, config, ALL_CAPS) - - -class TestUpdateSchedule: - def student_phase(self): - return UpdatePhaseSpec(kind="student", trainable_roles=("student",)) - - def fake_phase(self, repeats=1): - return UpdatePhaseSpec(kind="fake_score", repeats=repeats, trainable_roles=("fake_score",)) - - def test_normal_cycle_is_student_then_fake(self): - schedule = UpdateSchedule(phases=(self.student_phase(), self.fake_phase(repeats=2))) - cycle = schedule.next_cycle(TrainerCounters()) - assert cycle.requires_student_update is True - assert cycle.is_warmup is False - assert [request.kind for request in cycle.requests] == ["student", "fake_score", "fake_score"] - - def test_warmup_transitions_to_normal_cycles(self): - schedule = UpdateSchedule( - phases=(self.student_phase(), self.fake_phase()), - warmup_phases=(self.fake_phase(repeats=2),), - warmup_cycles=2, - ) - counters = TrainerCounters(completed_cycles=0) - assert schedule.next_cycle(counters).is_warmup is True - counters.completed_cycles = 1 - assert schedule.next_cycle(counters).is_warmup is True - counters.completed_cycles = 2 - assert schedule.next_cycle(counters).is_warmup is False - - def test_empty_schedule_raises(self): - with pytest.raises(ValueError, match="normal-cycle phases"): - UpdateSchedule(phases=()) - - def test_two_student_phases_raise(self): - with pytest.raises(ValueError, match="exactly one student"): - UpdateSchedule(phases=(self.student_phase(), self.student_phase())) - - def test_repeated_student_phase_raises_before_any_execution(self): - with pytest.raises(ValueError, match="repeats=1"): - UpdateSchedule( - phases=(UpdatePhaseSpec(kind="student", repeats=2, trainable_roles=("student",)), self.fake_phase()) - ) - - @pytest.mark.parametrize("repeats", [0, -1, True, 1.5]) - def test_nonpositive_repeats_raise(self, repeats): - with pytest.raises(ValueError, match="greater than zero"): - self.fake_phase(repeats) - - @pytest.mark.parametrize( - "kwargs,error", - [ - ({"kind": "bogus", "trainable_roles": ("fake_score",)}, "Invalid phase kind"), - ({"kind": "fake_score", "batch_policy": "bogus", "trainable_roles": ("fake_score",)}, "batch_policy"), - ({"kind": "fake_score"}, "at least one trainable role"), - ({"kind": "fake_score", "trainable_roles": ("fake_score", "fake_score")}, "duplicate"), - ({"kind": "student", "trainable_roles": ("student", "fake_score")}, "exactly the 'student'"), - ({"kind": "fake_score", "trainable_roles": ("student",)}, "must not train the student"), - ({"kind": "fake_score", "trainable_roles": ("fake_score",), "update_ema": True}, "EMA"), - ], - ) - def test_invalid_phase_contracts_raise(self, kwargs, error): - with pytest.raises(ValueError, match=error): - UpdatePhaseSpec(**kwargs) - - def test_warmup_requires_fake_only_phases(self): - with pytest.raises(ValueError, match="must not contain a student"): - UpdateSchedule( - phases=(self.student_phase(), self.fake_phase()), - warmup_phases=(self.student_phase(),), - warmup_cycles=1, - ) - - def test_warmup_cycles_require_warmup_phases(self): - with pytest.raises(ValueError, match="requires at least one warmup phase"): - UpdateSchedule(phases=(self.student_phase(), self.fake_phase()), warmup_cycles=1) - - def test_warmup_phases_require_positive_cycle_count(self): - with pytest.raises(ValueError, match="require warmup_cycles"): - UpdateSchedule( - phases=(self.student_phase(), self.fake_phase()), - warmup_phases=(self.fake_phase(),), - ) - - -class TestTrainerCounters: - def test_counters_start_at_zero(self): - counters = TrainerCounters() - assert counters.global_step == 0 - assert counters.optimizer_steps == {} - assert counters.completed_cycles == 0 - - def test_record_step_accumulates_per_role(self): - counters = TrainerCounters() - counters.record_step("student") - counters.record_step("student") - counters.record_step("fake_score") - assert counters.optimizer_steps == {"student": 2, "fake_score": 1} diff --git a/tests/trainer/diffusion/test_distillation_controller_on_cpu.py b/tests/trainer/diffusion/test_distillation_controller_on_cpu.py deleted file mode 100644 index ae954d2ba..000000000 --- a/tests/trainer/diffusion/test_distillation_controller_on_cpu.py +++ /dev/null @@ -1,225 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""CPU tests for the generic distillation trainer control plane.""" - -import ast -from pathlib import Path - -import pytest - -from verl_omni.trainer.diffusion.distillation.contracts import PhaseResult, UpdatePhaseSpec, UpdateSchedule -from verl_omni.trainer.diffusion.distillation.controller import ( - DistillationTrainerController, - FakeBatchProvider, - FakeDistillationHooks, - FakePhaseExecutor, -) -from verl_omni.trainer.diffusion.distillation.recipes import build_plan - -CAPS = frozenset({"distribution_matching", "autoregressive", "adversarial"}) - - -def make_plan(fake_repeats: int = 2, name: str = "dmd2", **config): - return build_plan(name, {"fake_update_ratio": fake_repeats, "model_path": "/m", **config}, CAPS) - - -def make_controller(plan=None, executor=None, hooks=None, batches: int = 1000): - plan = plan if plan is not None else make_plan() - executor = executor if executor is not None else FakePhaseExecutor() - hooks = hooks if hooks is not None else FakeDistillationHooks() - return ( - DistillationTrainerController(plan, executor, FakeBatchProvider(num_batches=batches), hooks), - executor, - hooks, - ) - - -class TwoStepExecutor(FakePhaseExecutor): - def execute_phase(self, request, batch): - if request.kind == "student": - return PhaseResult(optimizer_steps={"student": 2}) - return super().execute_phase(request, batch) - - -class WrongRoleExecutor(FakePhaseExecutor): - def execute_phase(self, request, batch): - if request.kind == "fake_score": - return PhaseResult(optimizer_steps={"fake_score": 1, "discriminator": 1}) - return super().execute_phase(request, batch) - - -class NoStepExecutor(FakePhaseExecutor): - def execute_phase(self, request, batch): - return PhaseResult() - - -class TestPhaseExpansion: - def test_normal_cycle_is_student_then_k_fake(self): - controller, executor, _ = make_controller(make_plan(fake_repeats=3)) - controller.run_cycle() - assert [request.kind for request in executor.executed] == [ - "student", - "fake_score", - "fake_score", - "fake_score", - ] - - def test_deterministic_ordering_across_cycles(self): - controller, executor, _ = make_controller(make_plan(fake_repeats=2)) - controller.run(3) - assert [request.kind for request in executor.executed] == ["student", "fake_score", "fake_score"] * 3 - - def test_repeat_index_is_per_phase(self): - controller, executor, _ = make_controller(make_plan(fake_repeats=3)) - controller.run_cycle() - fake_repeats = [request.repeat_index for request in executor.executed if request.kind == "fake_score"] - assert fake_repeats == [0, 1, 2] - - def test_warmup_transitions_to_normal_cycle(self): - controller, executor, hooks = make_controller(make_plan(fake_repeats=2, fake_warmup_cycles=2)) - first = controller.run_cycle() - second = controller.run_cycle() - third = controller.run_cycle() - assert first.is_warmup is True - assert second.is_warmup is True - assert third.is_warmup is False - assert [request.kind for request in executor.executed] == [ - "fake_score", - "fake_score", - "fake_score", - "fake_score", - "student", - "fake_score", - "fake_score", - ] - assert [call["global_step"] for call in hooks.calls] == [1] - - -class TestCounters: - def test_global_step_and_completed_cycles_advance(self): - controller, _, _ = make_controller(make_plan(fake_repeats=2)) - controller.run(4) - assert controller.counters.global_step == 4 - assert controller.counters.completed_cycles == 4 - - def test_distribution_only_counts_only_bound_trainable_roles(self): - controller, _, _ = make_controller(make_plan(fake_repeats=3)) - controller.run(2) - assert controller.counters.optimizer_steps == {"student": 2, "fake_score": 6} - - def test_paper_profile_counts_discriminator_steps(self): - controller, _, _ = make_controller(make_plan(fake_repeats=2, profile="paper")) - controller.run_cycle() - assert controller.counters.optimizer_steps == {"student": 1, "fake_score": 2, "discriminator": 2} - - def test_phase_request_carries_current_global_step_and_roles(self): - controller, executor, _ = make_controller(make_plan(fake_repeats=1)) - controller.run(2) - student_requests = [request for request in executor.executed if request.kind == "student"] - assert [request.global_step for request in student_requests] == [0, 1] - assert all(request.trainable_roles == ("student",) for request in student_requests) - - -class TestPhaseInvariants: - def test_skipped_student_phase_rolls_back_and_marks_driver_failed(self): - controller, _, _ = make_controller(make_plan(), executor=FakePhaseExecutor(skip_student=True)) - with pytest.raises(ValueError, match="no student optimizer step"): - controller.run_cycle() - assert controller.counters.global_step == 0 - assert controller.counters.optimizer_steps == {} - with pytest.raises(RuntimeError, match="cannot be retried"): - controller.run_cycle() - - def test_failed_fake_phase_rolls_back_and_cannot_retry(self): - controller, _, _ = make_controller(make_plan(), executor=FakePhaseExecutor(fail_on="fake_score")) - with pytest.raises(RuntimeError, match="failed on phase fake_score"): - controller.run_cycle() - assert controller.counters.global_step == 0 - assert controller.counters.optimizer_steps == {} - with pytest.raises(RuntimeError, match="cannot be retried"): - controller.run_cycle() - - def test_completed_role_must_report_exactly_one_step(self): - controller, _, _ = make_controller(make_plan(), executor=TwoStepExecutor()) - with pytest.raises(ValueError, match="exactly one optimizer step"): - controller.run_cycle() - - def test_unexpected_optimizer_role_is_rejected(self): - controller, _, _ = make_controller(make_plan(), executor=WrongRoleExecutor()) - with pytest.raises(ValueError, match="must report optimizer steps for exactly"): - controller.run_cycle() - assert controller.counters.optimizer_steps == {} - - def test_zero_progress_warmup_cycle_raises(self): - plan = make_plan(fake_repeats=1, fake_warmup_cycles=1) - controller, _, _ = make_controller(plan, executor=NoStepExecutor()) - with pytest.raises(ValueError, match="must report optimizer steps"): - controller.run_cycle() - - def test_two_student_phases_are_rejected_before_execution(self): - student = UpdatePhaseSpec(kind="student", trainable_roles=("student",)) - with pytest.raises(ValueError, match="exactly one student"): - UpdateSchedule(phases=(student, student)) - - -class TestHookScheduling: - def test_hook_observes_incremented_global_step(self): - controller, _, hooks = make_controller(make_plan(fake_repeats=1)) - controller.run(3) - assert [call["global_step"] for call in hooks.calls] == [1, 2, 3] - - def test_hook_receives_executor_and_metrics(self): - controller, executor, hooks = make_controller(make_plan(fake_repeats=1)) - controller.run_cycle() - assert hooks.calls[0]["executor"] is executor - assert "student" in hooks.calls[0]["metrics"] - - def test_hook_not_called_during_warmup(self): - controller, _, hooks = make_controller(make_plan(fake_warmup_cycles=1)) - controller.run_cycle() - assert hooks.calls == [] - - -class TestControllerPurity: - def test_core_modules_have_no_direct_model_or_ray_imports(self): - import verl_omni.trainer.diffusion.distillation.contracts as contracts_mod - import verl_omni.trainer.diffusion.distillation.controller as controller_mod - import verl_omni.trainer.diffusion.distillation.utils as utils_mod - - forbidden_roots = {"diffusers", "ray", "transformers", "vllm"} - for module in (contracts_mod, controller_mod, utils_mod): - tree = ast.parse(Path(module.__file__).read_text()) - imported_roots = set() - for node in ast.walk(tree): - if isinstance(node, ast.Import): - imported_roots.update(alias.name.split(".", 1)[0] for alias in node.names) - elif isinstance(node, ast.ImportFrom) and node.module: - imported_roots.add(node.module.split(".", 1)[0]) - assert imported_roots.isdisjoint(forbidden_roots) - - def test_reset_clears_healthy_driver_state(self): - controller, _, _ = make_controller(make_plan()) - controller.run(2) - controller.reset() - assert controller.counters.global_step == 0 - assert controller.counters.completed_cycles == 0 - assert controller.counters.optimizer_steps == {} - assert controller.metrics == {} - - def test_failed_driver_cannot_be_reset(self): - controller, _, _ = make_controller(make_plan(), executor=FakePhaseExecutor(fail_on="student")) - with pytest.raises(RuntimeError): - controller.run_cycle() - with pytest.raises(RuntimeError, match="cannot be reset"): - controller.reset() diff --git a/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py b/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py deleted file mode 100644 index fb1d8f272..000000000 --- a/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py +++ /dev/null @@ -1,156 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""CPU tests for ``algorithm.trainer_type=distillation`` routing. - -The DMD-family trainer is routed by ``algorithm.trainer_type``, which is -orthogonal to the existing on-policy distillation (OPD) path. OPD keeps using -``distillation.enabled=true`` together with ``trainer_type=policy_gradient``. -""" - -import pytest -from omegaconf import OmegaConf - -from verl_omni.trainer.config.algorithm import DiffusionAlgoConfig -from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationRayTrainer -from verl_omni.trainer.main_diffusion import _get_trainer_cls - - -class FakeAlgorithmConfig: - def __init__(self, trainer_type): - self.trainer_type = trainer_type - - -class FakeTrainerConfig: - def __init__(self, trainer_type): - self.algorithm = FakeAlgorithmConfig(trainer_type) - - -class TestTrainerRouting: - def test_distillation_routes_to_distillation_trainer(self): - assert _get_trainer_cls(FakeTrainerConfig("distillation")) is DistillationRayTrainer - - def test_policy_gradient_is_unchanged(self): - from verl_omni.trainer.diffusion.ray_diffusion_trainer import PolicyGradientRayTrainer - - assert _get_trainer_cls(FakeTrainerConfig("policy_gradient")) is PolicyGradientRayTrainer - - def test_direct_preference_is_unchanged(self): - from verl_omni.trainer.diffusion.ray_diffusion_trainer import DirectPreferenceRayTrainer - - assert _get_trainer_cls(FakeTrainerConfig("direct_preference")) is DirectPreferenceRayTrainer - - def test_unknown_trainer_type_lists_distillation(self): - with pytest.raises(ValueError, match="distillation"): - _get_trainer_cls(FakeTrainerConfig("bogus")) - - -class TestAlgorithmConfig: - def test_distillation_is_a_valid_trainer_type(self): - config = DiffusionAlgoConfig(trainer_type="distillation") - assert config.trainer_type == "distillation" - - def test_existing_trainer_types_still_valid(self): - assert DiffusionAlgoConfig(trainer_type="policy_gradient").trainer_type == "policy_gradient" - assert DiffusionAlgoConfig(trainer_type="direct_preference").trainer_type == "direct_preference" - - def test_default_is_policy_gradient(self): - assert DiffusionAlgoConfig().trainer_type == "policy_gradient" - - def test_invalid_trainer_type_raises(self): - with pytest.raises(ValueError, match="Invalid trainer_type"): - DiffusionAlgoConfig(trainer_type="bogus") - - -def runtime_config(): - return OmegaConf.create( - { - "algorithm": {"trainer_type": "distillation"}, - "actor_rollout_ref": { - "actor": { - "diffusion_loss": {"loss_mode": "flow_grpo"}, - "use_distill_loss": False, - }, - "model": {"path": "/m"}, - }, - "distillation": { - "enabled": False, - "distribution_matching": { - "recipe": "dmd2", - "profile": "distribution_only", - "fake_update_ratio": 2, - "fake_warmup_cycles": 0, - "rollout_strategy": None, - "data_mode": None, - "export_role": "student_ema", - }, - }, - } - ) - - -class TestPR1DataPlaneBoundary: - def test_production_constructor_reaches_explicit_pr2_boundary(self): - config = runtime_config() - trainer = DistillationRayTrainer( - config=config, - tokenizer=object(), - processor=object(), - role_worker_mapping={}, - resource_pool_manager=object(), - ray_worker_group_cls=object, - train_dataset=object(), - val_dataset=object(), - collate_fn=object(), - train_sampler=object(), - ) - assert trainer.config is config - with pytest.raises(NotImplementedError, match="PR 2"): - trainer.init_workers() - - def test_constructor_rejects_opd_switch(self): - config = runtime_config() - config.distillation.enabled = True - with pytest.raises(ValueError, match="must keep the OPD"): - DistillationRayTrainer(config=config) - - def test_config_and_capabilities_build_a_plan(self): - config = runtime_config() - trainer = DistillationRayTrainer(config=config, capabilities=frozenset({"distribution_matching"})) - assert trainer.plan is not None - assert trainer.plan.name == "dmd2" - assert trainer.plan.role_layout.groups[0].model_ref == "/m" - - def test_fit_without_executor_reports_pr2_boundary(self): - from verl_omni.trainer.diffusion.distillation.recipes import build_plan - - plan = build_plan("dmd2", {"model_path": "/m"}, frozenset({"distribution_matching"})) - trainer = DistillationRayTrainer(plan=plan) - with pytest.raises(NotImplementedError, match="PR 2"): - trainer.fit(num_cycles=1) - - def test_controller_binds_when_collaborators_are_supplied(self): - from verl_omni.trainer.diffusion.distillation.controller import ( - FakeBatchProvider, - FakePhaseExecutor, - ) - from verl_omni.trainer.diffusion.distillation.recipes import build_plan - - plan = build_plan("dmd2", {"fake_update_ratio": 1, "model_path": "/m"}, frozenset({"distribution_matching"})) - trainer = DistillationRayTrainer( - plan=plan, - executor=FakePhaseExecutor(), - batch_provider=FakeBatchProvider(num_batches=100), - ) - trainer.fit(num_cycles=2) - assert trainer.controller.counters.global_step == 2 diff --git a/tests/trainer/diffusion/test_distillation_trainer_wiring_on_cpu.py b/tests/trainer/diffusion/test_distillation_trainer_wiring_on_cpu.py index b127f6338..2a2af9a63 100644 --- a/tests/trainer/diffusion/test_distillation_trainer_wiring_on_cpu.py +++ b/tests/trainer/diffusion/test_distillation_trainer_wiring_on_cpu.py @@ -46,32 +46,6 @@ def test_auxiliary_distill_loss_passes(self): def test_default_config_passes(self): validate_distillation_config(compose_cfg([])) - def test_distribution_matching_trainer_does_not_enable_opd(self): - validate_distillation_config(compose_cfg(["algorithm.trainer_type=distillation"])) - - def test_distribution_matching_rejects_opd_enabled_flag(self): - with pytest.raises(ValueError, match="must keep the OPD"): - validate_distillation_config( - compose_cfg( - ENABLE - + [ - "algorithm.trainer_type=distillation", - "actor_rollout_ref.actor.diffusion_loss.loss_mode=distill_kl", - ] - ) - ) - - def test_distribution_matching_rejects_actor_distillation_loss(self): - with pytest.raises(ValueError, match="must keep the OPD"): - validate_distillation_config( - compose_cfg( - [ - "algorithm.trainer_type=distillation", - "actor_rollout_ref.actor.diffusion_loss.loss_mode=distill_kl", - ] - ) - ) - def test_enabled_but_no_distill_loss_raises(self): with pytest.raises(ValueError, match="distill"): validate_distillation_config(compose_cfg(ENABLE)) diff --git a/tests/trainer/diffusion/test_distillation_utils_on_cpu.py b/tests/trainer/diffusion/test_distillation_utils_on_cpu.py index 927f14d0e..93e688c82 100644 --- a/tests/trainer/diffusion/test_distillation_utils_on_cpu.py +++ b/tests/trainer/diffusion/test_distillation_utils_on_cpu.py @@ -24,15 +24,11 @@ import torch from verl_omni.trainer.diffusion.distillation.utils import ( - consistency_renoise_step, dmd_gradient, dmd_surrogate_loss, - epsilon_to_x0, fake_score_loss, fake_score_target, - legacy_cfg, ode_euler_step, - ode_regression_loss, standard_cfg, timestep_shift, velocity_to_x0, @@ -55,24 +51,10 @@ def test_velocity_to_x0_roundtrip(self): velocity = noise - x0 torch.testing.assert_close(velocity_to_x0(x_sigma, velocity, sigma), x0, atol=1e-5, rtol=1e-5) - def test_epsilon_to_x0_roundtrip(self): - x0 = torch.randn(2, 3, 4, 4) - noise = torch.randn(2, 3, 4, 4) - sigma = torch.rand(2, 1, 1, 1) * 0.9 - x_sigma = (1 - sigma) * x0 + sigma * noise - converted = epsilon_to_x0(x_sigma, noise, sigma, lambda value: 1 - value, lambda value: value) - torch.testing.assert_close(converted, x0, atol=1e-5, rtol=1e-5) - def test_canonical_conversions_return_fp32(self): value = torch.randn(1, 2, dtype=torch.float16) sigma = torch.full((1, 1), 0.5, dtype=torch.float16) assert velocity_to_x0(value, value, sigma).dtype == torch.float32 - assert epsilon_to_x0(value, value, sigma, lambda item: 1 - item, lambda item: item).dtype == torch.float32 - - def test_epsilon_conversion_rejects_zero_signal_coefficient(self): - value = torch.randn(1, 2) - with pytest.raises(ValueError, match="a\\(sigma\\) is zero"): - epsilon_to_x0(value, value, torch.ones(1, 1), lambda sigma: 1 - sigma, lambda sigma: sigma) class TestDMDGradient: @@ -119,10 +101,18 @@ def test_nonfinite_is_replaced_and_counted(self): assert nonfinite > 0 assert torch.isfinite(g_norm).all() - def test_invalid_normalization_epsilon_raises(self): + @pytest.mark.parametrize("epsilon", [0, -1, float("nan"), float("inf")]) + def test_invalid_normalization_epsilon_raises(self, epsilon): tensor = torch.zeros(1, 2) with pytest.raises(ValueError, match="greater than zero"): - dmd_gradient(tensor, tensor, tensor, normalization_epsilon=0) + dmd_gradient(tensor, tensor, tensor, normalization_epsilon=epsilon) + + def test_score_gradient_and_normalizer_are_detached(self): + generated = torch.ones(1, 3, requires_grad=True) + fake = torch.ones(1, 3, requires_grad=True) + real = torch.zeros(1, 3, requires_grad=True) + gradient, normalizer, _ = dmd_gradient(fake, real, generated) + assert not gradient.requires_grad and not normalizer.requires_grad def test_mismatched_shapes_raise(self): with pytest.raises(ValueError, match="identical shapes"): @@ -275,38 +265,6 @@ def test_incompatible_gradient_mask_is_rejected(loss_kind): fake_score_loss(value, value, value, mask) -class TestODERegression: - def test_masked_mse_uses_only_nonzero_timestep_positions(self): - prediction = torch.tensor([[1.0, 3.0], [5.0, 7.0]], requires_grad=True) - target = torch.zeros_like(prediction) - valid_mask = torch.tensor([[True, False], [True, False]]) - loss, active = ode_regression_loss(prediction, target, valid_mask) - assert active == 2 - assert loss.item() == pytest.approx(13.0) - loss.backward() - torch.testing.assert_close(prediction.grad, torch.tensor([[1.0, 0.0], [5.0, 0.0]])) - - def test_frame_mask_broadcasts_over_latent_dimensions(self): - prediction = torch.ones(1, 2, 1, 2, 2) - target = torch.zeros_like(prediction) - loss, active = ode_regression_loss(prediction, target, torch.tensor([[True, False]])) - assert active == 4 - assert loss.item() == pytest.approx(1.0) - - def test_target_is_detached(self): - prediction = torch.ones(1, 2, requires_grad=True) - target = torch.zeros(1, 2, requires_grad=True) - loss, _ = ode_regression_loss(prediction, target) - loss.backward() - assert prediction.grad is not None - assert target.grad is None - - def test_all_masked_ode_loss_raises(self): - tensor = torch.zeros(1, 2) - with pytest.raises(ValueError, match="all-masked"): - ode_regression_loss(tensor, tensor, torch.zeros_like(tensor, dtype=torch.bool)) - - class TestRolloutTransitions: def test_ode_euler_uses_deterministic_velocity_transition(self): latents = torch.tensor([[1.0, 2.0]]) @@ -314,21 +272,6 @@ def test_ode_euler_uses_deterministic_velocity_transition(self): result = ode_euler_step(latents, velocity, torch.tensor(0.8), torch.tensor(0.3)) torch.testing.assert_close(result, torch.tensor([[0.0, 2.5]])) - def test_consistency_transition_renoises_clean_prediction(self): - x0 = torch.tensor([[2.0, 4.0]]) - noise = torch.tensor([[0.0, 2.0]]) - result = consistency_renoise_step(x0, noise, torch.tensor(0.25)) - torch.testing.assert_close(result, torch.tensor([[1.5, 3.5]])) - - def test_euler_and_consistency_transitions_are_not_interchangeable(self): - latents = torch.randn(2, 3) - velocity = torch.randn(2, 3) - x0 = torch.randn(2, 3) - noise = torch.randn(2, 3) - euler = ode_euler_step(latents, velocity, torch.tensor(0.8), torch.tensor(0.3)) - renoised = consistency_renoise_step(x0, noise, torch.tensor(0.3)) - assert not torch.allclose(euler, renoised) - class TestCFG: def test_standard_cfg_form(self): @@ -337,18 +280,6 @@ def test_standard_cfg_form(self): out = standard_cfg(cond, uncond, 3.0) assert torch.allclose(out, uncond + 3.0 * (cond - uncond), atol=1e-6) - def test_legacy_cfg_differs_from_standard(self): - """Self-Forcing's cond + s*(cond-uncond) is not the standard definition.""" - cond = torch.randn(2, 8) - uncond = torch.randn(2, 8) - assert not torch.allclose(standard_cfg(cond, uncond, 3.0), legacy_cfg(cond, uncond, 3.0)) - - def test_legacy_cfg_equals_standard_with_shifted_scale(self): - """cond + s*(cond-uncond) == uncond + (s+1)*(cond-uncond).""" - cond = torch.randn(2, 8) - uncond = torch.randn(2, 8) - assert torch.allclose(legacy_cfg(cond, uncond, 3.0), standard_cfg(cond, uncond, 4.0), atol=1e-6) - def test_cfg_scale_zero_returns_unconditional(self): cond = torch.randn(2, 8) uncond = torch.randn(2, 8) diff --git a/tests/trainer/diffusion/test_dmd_config_on_cpu.py b/tests/trainer/diffusion/test_dmd_config_on_cpu.py new file mode 100644 index 000000000..ce84715b5 --- /dev/null +++ b/tests/trainer/diffusion/test_dmd_config_on_cpu.py @@ -0,0 +1,125 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import pytest +from hydra import compose, initialize_config_dir +from hydra.errors import ConfigCompositionException +from omegaconf import OmegaConf +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config.optimizer import FSDPOptimizerConfig + +import verl_omni +from verl_omni.workers.config import DiffusionDistillationConfig, DiffusionDMDConfig + +CONFIG_DIR = str(Path(verl_omni.__file__).parent / "trainer" / "config") + + +def compose_config(overrides=()): + with initialize_config_dir(config_dir=CONFIG_DIR, version_base=None): + return compose(config_name="diffusion_trainer", overrides=list(overrides)) + + +class TestDMDConfig: + def test_defaults_are_dmd2_without_opd_activation(self): + config = DiffusionDMDConfig() + assert config.fake_update_ratio == 2 + assert config.export_role == "student" + assert isinstance(config.fake_score_optim, FSDPOptimizerConfig) + assert config.fake_score_optim.lr == pytest.approx(2e-5) + assert not DiffusionDistillationConfig().enabled + assert not hasattr(DiffusionDistillationConfig(), "distribution_matching") + + @pytest.mark.parametrize( + "field", ["fake_update_ratio", "student_micro_batch_size_per_gpu", "fake_score_micro_batch_size_per_gpu"] + ) + @pytest.mark.parametrize("value", [True, False, 0, -1, 1.5, "2"]) + def test_positive_counts_are_not_silently_coerced(self, field, value): + with pytest.raises(ValueError, match="positive integer"): + DiffusionDMDConfig(**{field: value}) + + @pytest.mark.parametrize("field", ["score_discrete_steps", "ema_start_step"]) + @pytest.mark.parametrize("value", [True, -1, 0.5]) + def test_nonnegative_counts_are_validated(self, field, value): + with pytest.raises(ValueError, match="non-negative integer"): + DiffusionDMDConfig(**{field: value}) + + @pytest.mark.parametrize( + "kwargs", + [ + {"teacher_guidance_scale": 0}, + {"teacher_guidance_scale": float("inf")}, + {"negative_prompt": None}, + {"normalization_epsilon": 0}, + {"normalization_epsilon": float("nan")}, + {"cfg_norm": "legacy"}, + {"rollout_timestep_shift": 0.5}, + {"score_timestep_shift": float("nan")}, + {"score_sigma_min": 0}, + {"score_sigma_min": 0.9, "score_sigma_max": 0.2}, + {"score_sigma_max": 1.1}, + {"ema_decay": -0.1}, + {"ema_decay": float("nan")}, + {"export_role": "fake_score"}, + ], + ) + def test_invalid_values_fail_closed(self, kwargs): + with pytest.raises(ValueError): + DiffusionDMDConfig(**kwargs) + + def test_optimizer_defaults_are_not_shared(self): + first, second = DiffusionDMDConfig(), DiffusionDMDConfig() + first.fake_score_optim.total_training_steps = 4 + assert second.fake_score_optim.total_training_steps == -1 + + def test_typed_hydra_group_matches_dataclass_defaults(self): + config = compose_config() + dmd = omega_conf_to_dataclass(config.dmd) + assert isinstance(dmd, DiffusionDMDConfig) + assert isinstance(dmd.fake_score_optim, FSDPOptimizerConfig) + assert OmegaConf.structured(dmd) == OmegaConf.structured(DiffusionDMDConfig()) + assert omega_conf_to_dataclass(config.distillation).enabled is False + + def test_dmd_loss_can_be_composed_without_model_loading(self): + config = compose_config( + [ + "actor_rollout_ref.model.algorithm=dmd2", + "dmd.fake_update_ratio=3", + "dmd.fake_score_optim.lr=1e-5", + ] + ) + loss = omega_conf_to_dataclass(config.actor_rollout_ref.actor.diffusion_loss) + assert loss.loss_mode == "dmd2" + assert omega_conf_to_dataclass(config.dmd).fake_update_ratio == 3 + + @pytest.mark.parametrize("schedule", ["inline", "one_step_off"]) + def test_opd_group_is_unchanged(self, schedule): + config = compose_config( + [ + "distillation.enabled=true", + f"distillation.scheduler={schedule}", + "distillation.teacher_models.teacher_model.model_path=/ckpt/teacher", + ] + ) + opd = omega_conf_to_dataclass(config.distillation) + assert opd.scheduler == schedule + assert opd.teacher_models["default"].model_path == "/ckpt/teacher" + assert "distribution_matching" not in OmegaConf.to_container(config.distillation) + assert isinstance(omega_conf_to_dataclass(config.dmd), DiffusionDMDConfig) + + @pytest.mark.parametrize("field", ["recipe", "profile", "adversarial", "rollout_strategy", "teacher_models"]) + def test_out_of_scope_options_are_not_accepted(self, field): + with pytest.raises(ConfigCompositionException): + compose_config([f"dmd.{field}=unsupported"]) diff --git a/tests/trainer/diffusion/test_dmd_loss_on_cpu.py b/tests/trainer/diffusion/test_dmd_loss_on_cpu.py new file mode 100644 index 000000000..0667cabb1 --- /dev/null +++ b/tests/trainer/diffusion/test_dmd_loss_on_cpu.py @@ -0,0 +1,125 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.trainer.diffusion.diffusion_algos import DiffusionLossResult, DMDLoss, get_diffusion_loss_fn +from verl_omni.workers.config import DiffusionActorConfig, DiffusionLossConfig +from verl_omni.workers.utils.losses import diffusion_loss + + +def make_actor(): + return DiffusionActorConfig(strategy="fsdp2", rollout_n=1, diffusion_loss=DiffusionLossConfig(loss_mode="dmd2")) + + +def make_batch(size, stage="student", accumulation=1, sp_size=1): + batch = TensorDict({}, batch_size=[size]) + tu.assign_non_tensor_data(batch, "dmd_stage", stage) + tu.assign_non_tensor_data(batch, "gradient_accumulation_steps", accumulation) + tu.assign_non_tensor_data(batch, "sp_size", sp_size) + return batch + + +class TestDMDLoss: + def test_registration_does_not_alias_original_dmd(self): + assert isinstance(get_diffusion_loss_fn("dmd2"), DMDLoss) + with pytest.raises(ValueError, match="Unsupported"): + get_diffusion_loss_fn("dmd") + with pytest.raises(ValueError, match="loss_mode"): + DiffusionLossConfig(loss_mode="dmd") + + def test_student_gradient_and_detached_scores(self): + student = torch.tensor([[1.0, 3.0], [4.0, 8.0]], requires_grad=True) + teacher = torch.zeros_like(student, requires_grad=True) + fake = torch.ones_like(student, requires_grad=True) + result = DMDLoss()( + config=make_actor(), + model_output={ + "generated_x0": student, + "teacher_x0": teacher, + "fake_x0": fake, + }, + data=make_batch(2), + ) + assert isinstance(result, DiffusionLossResult) + result.loss.backward() + expected = torch.tensor([[0.5, 0.5], [1 / 6, 1 / 6]]) / student.numel() + torch.testing.assert_close(student.grad, expected) + assert teacher.grad is None and fake.grad is None + assert result.metrics["dmd/nonfinite"] == 0 + + def test_fake_stage_uses_detached_flow_target(self): + student = torch.ones(2, 3, requires_grad=True) + prediction = torch.zeros(2, 3, requires_grad=True) + noise = torch.full_like(student, 2.0, requires_grad=True) + loss, _ = diffusion_loss( + make_actor(), + { + "generated_x0": student, + "noise_pred": prediction, + "noise": noise, + }, + make_batch(2, "fake_score"), + ) + loss.backward() + torch.testing.assert_close(loss, torch.tensor(1.0)) + assert student.grad is None and noise.grad is None + torch.testing.assert_close(prediction.grad, torch.full_like(prediction, -2 / prediction.numel())) + + @pytest.mark.parametrize("stage,missing", [("student", "teacher_x0"), ("fake_score", "noise_pred")]) + def test_stage_inputs_fail_closed(self, stage, missing): + output = {key: torch.ones(1, 2) for key in ("generated_x0", "teacher_x0", "fake_x0", "noise_pred", "noise")} + del output[missing] + with pytest.raises(KeyError, match=missing): + diffusion_loss(make_actor(), output, make_batch(1, stage)) + + def test_unknown_stage_is_not_student_fallback(self): + with pytest.raises(ValueError, match="dmd_stage"): + diffusion_loss(make_actor(), {}, make_batch(1, "discriminator")) + + def test_normalization_setting_reaches_registered_loss(self): + batch = make_batch(1) + tu.assign_non_tensor_data(batch, "dmd_normalization_epsilon", 0.5) + output = {"generated_x0": torch.zeros(1, 2), "teacher_x0": torch.zeros(1, 2), "fake_x0": torch.ones(1, 2)} + loss, _ = diffusion_loss(make_actor(), output, batch) + torch.testing.assert_close(loss, torch.tensor(2.0)) + + @pytest.mark.parametrize("stage", ["student", "fake_score"]) + def test_unequal_microbatches_preserve_loss_and_gradient(self, stage): + full = torch.tensor([[1.0, 2.0], [3.0, 4.0], [7.0, 8.0]], requires_grad=True) + accumulated = full.detach().clone().requires_grad_() + if stage == "student": + output = {"generated_x0": full, "fake_x0": torch.ones_like(full), "teacher_x0": torch.zeros_like(full)} + else: + output = {"noise_pred": full, "noise": torch.ones_like(full), "generated_x0": torch.zeros_like(full)} + loss, _ = diffusion_loss(make_actor(), output, make_batch(3, stage)) + loss.backward() + total = 0.0 + for start, stop in ((0, 2), (2, 3)): + micro = {key: tensor[start:stop] for key, tensor in output.items()} + micro["generated_x0" if stage == "student" else "noise_pred"] = accumulated[start:stop] + micro_loss, _ = diffusion_loss(make_actor(), micro, make_batch(stop - start, stage, 3 / (stop - start))) + total += micro_loss.detach() + micro_loss.backward() + torch.testing.assert_close(total, loss.detach()) + torch.testing.assert_close(accumulated.grad, full.grad) + + def test_sequence_parallel_factor_applied_once(self): + output = {"generated_x0": torch.ones(1, 2), "teacher_x0": torch.zeros(1, 2), "fake_x0": torch.ones(1, 2)} + single, _ = diffusion_loss(make_actor(), output, make_batch(1)) + scaled, _ = diffusion_loss(make_actor(), output, make_batch(1, accumulation=2, sp_size=4)) + torch.testing.assert_close(scaled, single * 2) diff --git a/tests/trainer/diffusion/test_dmd_trainer_on_cpu.py b/tests/trainer/diffusion/test_dmd_trainer_on_cpu.py new file mode 100644 index 000000000..6757357f0 --- /dev/null +++ b/tests/trainer/diffusion/test_dmd_trainer_on_cpu.py @@ -0,0 +1,237 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch +from hydra import compose, initialize_config_dir +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu +from verl.utils.config import omega_conf_to_dataclass + +import verl_omni +from verl_omni.trainer.diffusion.ray_diffusion_trainer import DistributionMatchingRayTrainer +from verl_omni.trainer.main_diffusion import _get_trainer_cls + +CONFIG_DIR = str(Path(verl_omni.__file__).parent / "trainer" / "config") + + +def make_config(overrides=()): + with initialize_config_dir(config_dir=CONFIG_DIR, version_base=None): + return compose( + config_name="diffusion_trainer", + overrides=[ + "algorithm.trainer_type=distribution_matching", + "algorithm.sample_source=offline", + "actor_rollout_ref.model.algorithm=dmd2", + "actor_rollout_ref.model.model_type=diffusion_dmd_model", + "actor_rollout_ref.model.lora_rank=2", + "actor_rollout_ref.actor.strategy=fsdp2", + "data.train_batch_size=8", + "trainer.total_training_steps=3", + "trainer.save_freq=-1", + "trainer.test_freq=-1", + "trainer.val_before_train=false", + "trainer.resume_mode=disable", + *overrides, + ], + ) + + +class FakeTracking: + records = [] + + def __init__(self, **kwargs): + pass + + def log(self, data, step): + self.records.append((step, data)) + + +class FakeDMDWorker: + def __init__(self, outcomes): + self.outcomes = iter(outcomes) + self.stages = [] + + def update_actor(self, data): + self.stages.append(tu.get_non_tensor_data(data, "dmd_stage", default=None)) + outcome = next(self.outcomes) + if isinstance(outcome, Exception): + raise outcome + return tu.get_tensordict( + {}, {"metrics": {"dmd/update_applied": outcome, "dmd/skip_nonfinite": 1 - outcome, "loss": 0.25}} + ) + + +def empty_batch(): + return TensorDict({}, batch_size=[8]) + + +def make_trainer(outcomes): + trainer = object.__new__(DistributionMatchingRayTrainer) + trainer.config = make_config() + trainer.dmd_config = omega_conf_to_dataclass(trainer.config.dmd) + trainer.global_steps = 0 + trainer.failed = False + trainer.optimizer_steps = {"student": 0, "fake_score": 0} + trainer.data_epoch = 0 + trainer.total_training_steps = 3 + trainer.actor_rollout_wg = FakeDMDWorker(outcomes) + trainer.next_batch = empty_batch + trainer.export_student = MagicMock() + return trainer + + +class TestDMDConfiguration: + def test_route_and_production_preflight(self): + config = make_config() + DistributionMatchingRayTrainer.validate_config(config) + assert _get_trainer_cls(config) is DistributionMatchingRayTrainer + + @pytest.mark.parametrize( + "override", + [ + "algorithm.sample_source=online", + "actor_rollout_ref.model.algorithm=dmd", + "actor_rollout_ref.actor.use_kl_loss=true", + "actor_rollout_ref.actor.use_distill_loss=true", + "distillation.enabled=true", + "actor_rollout_ref.actor.ppo_epochs=2", + "actor_rollout_ref.model.lora_rank=0", + "actor_rollout_ref.actor.strategy=fsdp", + "data.train_batch_size=7", + "actor_rollout_ref.model.model_type=diffusion_model", + "actor_rollout_ref.actor.checkpoint.load_contents=[model]", + ], + ) + def test_unsupported_modes_fail_before_workers(self, override): + with pytest.raises(ValueError): + DistributionMatchingRayTrainer.validate_config(make_config([override])) + + def test_fingerprint_is_mapping_order_independent(self): + trainer = make_trainer([1] * 9) + first = trainer.configuration_fingerprint() + from omegaconf import OmegaConf + + fields = OmegaConf.to_container(trainer.config.dmd) + trainer.config.dmd = dict(reversed(list(fields.items()))) + assert trainer.configuration_fingerprint() == first + + +class TestDMDCycles: + def test_normal_and_skipped_cycles_keep_separate_success_counts(self, monkeypatch): + monkeypatch.setattr("verl.utils.tracking.Tracking", FakeTracking) + FakeTracking.records = [] + trainer = make_trainer([1, 1, 1, 0, 1, 0, 1, 0, 1]) + trainer.fit() + assert trainer.global_steps == 3 + assert trainer.optimizer_steps == {"student": 2, "fake_score": 4} + assert trainer.actor_rollout_wg.stages == ["student", "fake_score", "fake_score"] * 3 + assert [step for step, _ in FakeTracking.records] == [1, 2, 3] + assert "fake_score/0/loss" in FakeTracking.records[0][1] + assert "fake_score/1/loss" in FakeTracking.records[0][1] + trainer.export_student.assert_called_once() + + def test_all_skipped_budget_terminates_without_claiming_training_success(self, monkeypatch): + monkeypatch.setattr("verl.utils.tracking.Tracking", FakeTracking) + trainer = make_trainer([0] * 9) + with pytest.raises(RuntimeError, match="without successful updates"): + trainer.fit() + assert trainer.global_steps == 3 + assert trainer.optimizer_steps == {"student": 0, "fake_score": 0} + trainer.export_student.assert_not_called() + + def test_partial_exception_is_not_a_numerical_retry(self, monkeypatch): + monkeypatch.setattr("verl.utils.tracking.Tracking", FakeTracking) + trainer = make_trainer([1, RuntimeError("injected rank failure"), 1]) + with pytest.raises(RuntimeError, match="injected rank failure"): + trainer.fit() + assert trainer.global_steps == 0 + assert trainer.optimizer_steps["student"] == 1 # This cannot roll back a real optimizer update. + assert trainer.actor_rollout_wg.stages == ["student", "fake_score"] + trainer.export_student.assert_not_called() + with pytest.raises(RuntimeError, match="must be reconstructed"): + trainer.fit() + assert trainer.actor_rollout_wg.stages == ["student", "fake_score"] + + def test_malformed_fractional_outcome_is_not_accepted(self): + trainer = make_trainer([0.5]) + with pytest.raises(RuntimeError, match="Malformed"): + trainer.update_stage("student", 0) + + +class TestDMDCheckpoint: + def test_counter_mismatch_rejected_before_worker_load(self, tmp_path): + trainer = make_trainer([1] * 9) + trainer.config.trainer.n_gpus_per_node = 1 + trainer.config.trainer.resume_mode = "resume_path" + trainer.config.trainer.resume_from_path = str(tmp_path) + trainer.actor_rollout_wg = MagicMock() + torch.save( + { + "version": 1, + "configuration": trainer.configuration_fingerprint(), + "global_step": 1, + "optimizer_steps": {"student": 1, "fake_score": 2}, + }, + tmp_path / "trainer.pt", + ) + (tmp_path / "data.pt").touch() + actor = tmp_path / "actor" + actor.mkdir() + for kind in ("model", "optim", "extra_state"): + (actor / f"{kind}_world_size_1_rank_0.pt").touch() + torch.save( + {"version": 1, "world_size": 1, "optimizer_steps": {"student": 0, "fake_score": 2}}, + actor / "dmd_state_rank_0.pt", + ) + with pytest.raises(ValueError, match="counters do not match"): + trainer._load_checkpoint() + trainer.actor_rollout_wg.load_checkpoint.assert_not_called() + + def test_missing_shards_are_not_published(self, tmp_path): + trainer = make_trainer([1] * 9) + trainer.config.trainer.default_local_dir = str(tmp_path) + trainer.actor_rollout_wg = MagicMock() + with pytest.raises(ValueError, match="missing model_world"): + trainer._save_checkpoint() + assert not (tmp_path / "global_step_0").exists() + assert not list(tmp_path.glob(".global_step_*")) + + def test_failed_save_preserves_latest_and_publishes_no_partial_cycle(self, tmp_path): + trainer = make_trainer([1] * 9) + trainer.config.trainer.default_local_dir = str(tmp_path) + trainer.global_steps = 2 + trainer.train_dataloader = MagicMock() + trainer.actor_rollout_wg = MagicMock() + trainer.actor_rollout_wg.save_checkpoint.side_effect = RuntimeError("failed shard") + tracker = tmp_path / "latest_checkpointed_iteration.txt" + tracker.write_text("1") + with pytest.raises(RuntimeError, match="failed shard"): + trainer._save_checkpoint() + assert tracker.read_text() == "1" + assert not (tmp_path / "global_step_2").exists() + assert not list(tmp_path.glob(".global_step_*")) + + def test_old_checkpoint_is_rejected_before_loading_workers(self, tmp_path): + trainer = make_trainer([1] * 9) + trainer.config.trainer.resume_mode = "resume_path" + trainer.config.trainer.resume_from_path = str(tmp_path) + trainer.actor_rollout_wg = MagicMock() + (tmp_path / "trainer.pt").touch() + with pytest.raises(ValueError, match="Incomplete/old"): + trainer._load_checkpoint() + trainer.actor_rollout_wg.load_checkpoint.assert_not_called() diff --git a/tests/trainer/omni/test_main_omni_on_cpu.py b/tests/trainer/omni/test_main_omni_on_cpu.py index 7fb313597..e000e3311 100644 --- a/tests/trainer/omni/test_main_omni_on_cpu.py +++ b/tests/trainer/omni/test_main_omni_on_cpu.py @@ -123,9 +123,38 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke assert model_config.tokenizer == "tokenizer" assert model_config.processor == "processor" + mock_adapter.register_auto_classes.assert_called_once_with() mock_adapter.configure_tokenizer.assert_called_once_with("local:tokenizer-path", model_config) mock_adapter.configure_processor.assert_called_once_with(str(tmp_path), model_config) + def test_config_only_load_registers_known_auto_classes_without_requiring_an_adapter(self, monkeypatch, tmp_path): + from types import SimpleNamespace + + from verl_omni.pipelines.model_base import OmniModelBase + from verl_omni.workers.config.omni import model as model_config_module + from verl_omni.workers.config.omni.model import OmniModelConfig + + mock_adapter = MagicMock() + monkeypatch.setattr(OmniModelBase, "peek_class", lambda *_args: mock_adapter) + monkeypatch.setattr(model_config_module, "resolve_model_local_dir", lambda path, use_shm=False: str(tmp_path)) + monkeypatch.setattr(model_config_module, "copy_to_local", lambda path, use_shm=False: str(tmp_path)) + + def load_config(*_args, **_kwargs): + mock_adapter.register_auto_classes.assert_called_once_with() + return SimpleNamespace(tie_word_embeddings=False, architectures=["arch"]) + + monkeypatch.setattr(model_config_module.AutoConfig, "from_pretrained", load_config) + + model_config = OmniModelConfig( + path=str(tmp_path), + architecture="RegisteredArchitecture", + model_stage="talker", + load_tokenizer=False, + ) + + assert model_config.tokenizer is None + assert model_config.processor is None + def test_ray_task_runner_delegates_omni_model_to_adapter(self, monkeypatch, tmp_path): config = OmegaConf.create( { diff --git a/tests/utils/test_diffusion_model_provenance_on_cpu.py b/tests/utils/test_diffusion_model_provenance_on_cpu.py new file mode 100644 index 000000000..d363a227d --- /dev/null +++ b/tests/utils/test_diffusion_model_provenance_on_cpu.py @@ -0,0 +1,36 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib + +import pytest + +from verl_omni.utils.fs import diffusion_model_provenance + + +class TestDiffusionModelProvenance: + @pytest.mark.parametrize("source", ["snapshot", "download", "local"]) + def test_revision_is_recorded_only_when_available(self, tmp_path, source): + revision = "a" * 40 + root = tmp_path / "snapshots" / revision if source == "snapshot" else tmp_path + (root / "transformer").mkdir(parents=True) + config = b'{"in_channels": 64}' + (root / "transformer/config.json").write_bytes(config) + if source == "download": + metadata = root / ".cache/huggingface/download/model_index.json.metadata" + metadata.parent.mkdir(parents=True) + metadata.write_text(revision + "\netag\n0\n") + value = diffusion_model_provenance(str(root)) + assert value["base_model_revision"] == (None if source == "local" else revision) + assert value["base_transformer_config_sha256"] == hashlib.sha256(config).hexdigest() diff --git a/tests/utils/test_net_utils_on_cpu.py b/tests/utils/test_net_utils_on_cpu.py new file mode 100644 index 000000000..b9b422d20 --- /dev/null +++ b/tests/utils/test_net_utils_on_cpu.py @@ -0,0 +1,61 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU tests for the non-ephemeral MASTER_PORT picker.""" + +import socket + +import pytest + +from verl_omni.utils import net_utils + + +def _bind(address: str, port: int) -> socket.socket: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind((address, port)) + return sock + + +def test_returns_bindable_port_below_ephemeral_range(): + port = net_utils.get_non_ephemeral_free_port("127.0.0.1") + lo, _ = net_utils.ephemeral_port_range() + assert 1024 <= port < lo + holder = _bind("127.0.0.1", port) # raises if the picker handed out a busy port + holder.close() + + +def test_skips_ports_already_bound(monkeypatch): + monkeypatch.setattr(net_utils, "ephemeral_port_range", lambda: (1030, 60999)) + holders = [_bind("127.0.0.1", port) for port in range(1024, 1029)] + try: + assert net_utils.get_non_ephemeral_free_port("127.0.0.1") == 1029 + finally: + for holder in holders: + holder.close() + + +def test_raises_when_all_candidates_are_bound(monkeypatch): + monkeypatch.setattr(net_utils, "ephemeral_port_range", lambda: (1027, 60999)) + holders = [_bind("127.0.0.1", port) for port in range(1024, 1027)] + try: + with pytest.raises(RuntimeError, match="No free non-ephemeral port"): + net_utils.get_non_ephemeral_free_port("127.0.0.1") + finally: + for holder in holders: + holder.close() + + +def test_raises_when_ephemeral_range_leaves_no_candidates(monkeypatch): + monkeypatch.setattr(net_utils, "ephemeral_port_range", lambda: (1024, 60999)) + with pytest.raises(RuntimeError, match="no non-privileged candidate ports"): + net_utils.get_non_ephemeral_free_port("127.0.0.1") diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_deploy_config_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_deploy_config_on_cpu.py index 787b7817d..9b5737fb2 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_deploy_config_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_deploy_config_on_cpu.py @@ -22,11 +22,11 @@ """ import types -from unittest.mock import MagicMock import yaml from verl.utils.device import get_visible_devices_keyword +from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ARStrategy @@ -39,14 +39,27 @@ def _run_write_deploy_config( fake_self = types.SimpleNamespace( config=types.SimpleNamespace(**config_kwargs), ) - adapter = MagicMock() - adapter.build_stage_configs.return_value = [types.SimpleNamespace(stage_id=0)] - adapter.get_pipeline_id.return_value = "minimax_h3" - adapter.get_stage_engine_extras.return_value = {} + + class Adapter(OmniRolloutPipelineBase): + @classmethod + def build_stage_configs(cls, pipeline_mode="thinker_only"): + return [ + types.SimpleNamespace( + stage_id=0, + final_output=True, + final_output_type="audio", + sampling_constraints={}, + ) + ] + + @classmethod + def get_pipeline_id(cls, pipeline_mode="thinker_only"): + return "minimax_h3" + monkeypatch.setenv(get_visible_devices_keyword(), "0,1,2,3") engine_kwargs: dict = {} - ARStrategy(fake_self)._write_deploy_config(engine_kwargs, "minimax_h3", adapter, "t2av") + ARStrategy(fake_self)._write_deploy_config(engine_kwargs, "minimax_h3", Adapter, "t2av") with open(engine_kwargs["deploy_config"]) as f: return yaml.safe_load(f) diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_run_server_engine_args_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_run_server_engine_args_on_cpu.py index 0e22bb497..e6166de56 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_run_server_engine_args_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_run_server_engine_args_on_cpu.py @@ -39,7 +39,7 @@ async def _run_server_capture(monkeypatch, engine_args): server._server_address = ("127.0.0.1", 0) monkeypatch.setattr(server_module.OmniEngineArgs, "from_cli_args", classmethod(lambda cls, _: engine_args)) - monkeypatch.setattr(server_module, "get_free_port", lambda *a, **k: (12345, SimpleNamespace(close=lambda: None))) + monkeypatch.setattr(server_module, "get_non_ephemeral_free_port", lambda *a, **k: 12345) # run_server sets MASTER_ADDR/MASTER_PORT; write them to a scratch copy. monkeypatch.setattr(os, "environ", dict(os.environ)) diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py index a1771f3a5..709658bf9 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py @@ -13,12 +13,16 @@ # limitations under the License. from argparse import Namespace +from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock import pytest import torch +import yaml +from verl_omni.pipelines.model_base import OmniRolloutPipelineBase +from verl_omni.pipelines.qwen3_omni.omni_rollout_adapter import Qwen3OmniRolloutAdapter from verl_omni.pipelines.rollout_media import DiffusionIOSpec, MediaSpec from verl_omni.workers.rollout.vllm_rollout import vllm_omni_ar_strategy as ar_strategy_module from verl_omni.workers.rollout.vllm_rollout import vllm_omni_async_server as server_module @@ -115,6 +119,21 @@ def test_strategies_preserve_platform_worker_extensions(): assert DiffusionStrategy(server).worker_extension_cls("npu").endswith("vLLMOmniNPUColocateWorkerExtension") +@pytest.mark.parametrize("adapter_cls", [OmniRolloutPipelineBase, Qwen3OmniRolloutAdapter]) +def test_optional_rollout_hooks_preserve_existing_ar_defaults(adapter_cls): + first, final = object(), object() + + assert adapter_cls.supports_async_chunk is True + assert adapter_cls.weight_sync_stage_ids("full") is None + assert adapter_cls.policy_stage_id("full") == 0 + assert adapter_cls.prepare_engine_prompt([], None, {}) is None + assert adapter_cls.combine_engine_outputs([final], {}) == (final, {}) + with pytest.raises(NotImplementedError, match="multiple final outputs"): + adapter_cls.combine_engine_outputs([first, final], {}) + with pytest.raises(RuntimeError, match="no outputs"): + adapter_cls.combine_engine_outputs([], {}) + + def test_ar_strategy_preserves_prompt_and_sampling_preprocessing(): processor = SimpleNamespace(dedup_pad_tokens=lambda token_ids: token_ids[:2]) server = SimpleNamespace( @@ -195,6 +214,191 @@ def test_ar_strategy_preserves_engine_kwarg_normalization(monkeypatch): } +@pytest.mark.parametrize( + ("timeout_kwargs", "expected"), + [ + ({"stage-init-timeout": 45}, {"stage-init-timeout": 45, "init-timeout": 600}), + ( + {"stage_init_timeout": 45, "init-timeout": 90}, + {"stage-init-timeout": 45, "init-timeout": 90}, + ), + ], +) +def test_ar_strategy_preserves_hyphenated_timeout_kwargs(monkeypatch, timeout_kwargs, expected): + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: None) + strategy = ARStrategy(SimpleNamespace()) + engine_kwargs = {"pipeline_name": "missing", **timeout_kwargs} + + strategy.preprocess_engine_kwargs(engine_kwargs) + + assert engine_kwargs == expected + + +def test_ar_strategy_honors_adapter_chunking_and_weight_sync_contracts(monkeypatch): + class Adapter: + supports_async_chunk = False + + @staticmethod + def rollout_flags(pipeline_mode): + return {0: {"mode": pipeline_mode}} + + @staticmethod + def weight_sync_stage_ids(pipeline_mode): + assert pipeline_mode == "full" + return [1] + + @staticmethod + def get_engine_hf_overrides(pipeline_mode): + assert pipeline_mode == "full" + return {} + + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: Adapter) + strategy = ARStrategy(SimpleNamespace(_rollout_flags={})) + + def write_deploy_config(*args): + strategy._weight_sync_stage_ids = Adapter.weight_sync_stage_ids("full") + + monkeypatch.setattr(strategy, "_write_deploy_config", write_deploy_config) + + with pytest.raises(ValueError, match="requires async_chunk=false"): + strategy.preprocess_engine_kwargs({"pipeline_name": "adapter", "pipeline_mode": "full"}) + + engine_kwargs = {"pipeline_name": "adapter", "pipeline_mode": "full", "async_chunk": False} + strategy.preprocess_engine_kwargs(engine_kwargs) + + assert strategy._rollout_adapter is Adapter + assert strategy._weight_sync_stage_ids == [1] + assert strategy.server._rollout_flags == {0: {"mode": "full"}} + assert engine_kwargs == {"async-chunk": False} + + +def test_ar_strategy_resolves_nonzero_policy_and_weight_sync_stages(monkeypatch): + stages = [ + SimpleNamespace(stage_id=0, final_output=False, final_output_type=None, sampling_constraints={}), + SimpleNamespace(stage_id=1, final_output=True, final_output_type="latent", sampling_constraints={}), + ] + + class Adapter(OmniRolloutPipelineBase): + @classmethod + def build_stage_configs(cls, pipeline_mode="thinker_only"): + return stages + + @classmethod + def get_pipeline_id(cls, pipeline_mode="thinker_only"): + return "test_pipeline" + + @classmethod + def policy_stage_id(cls, pipeline_mode="thinker_only"): + return 1 + + @classmethod + def weight_sync_stage_ids(cls, pipeline_mode="thinker_only"): + return [1] + + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: Adapter) + monkeypatch.setattr(ar_strategy_module, "get_visible_devices_keyword", lambda: "CUDA_VISIBLE_DEVICES") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + server = SimpleNamespace( + config=SimpleNamespace( + tensor_model_parallel_size=1, + text_encoder_tp_size=1, + max_model_len=64, + max_num_batched_tokens=64, + ), + _rollout_flags={}, + ) + strategy = ARStrategy(server) + engine_kwargs = {"pipeline_name": "adapter"} + + strategy.preprocess_engine_kwargs(engine_kwargs) + + assert strategy._policy_stage_index == 1 + assert strategy._policy_sampling_constraints == {} + assert strategy._weight_sync_stage_ids == [1] + server._temp_deploy_ctx.cleanup() + + +def test_ar_strategy_moves_capacity_overrides_to_each_stage(monkeypatch): + stages = [ + SimpleNamespace(stage_id=0, final_output=False, final_output_type=None, sampling_constraints={}), + SimpleNamespace(stage_id=1, final_output=True, final_output_type="latent", sampling_constraints={}), + ] + + class Adapter(OmniRolloutPipelineBase): + @classmethod + def build_stage_configs(cls, pipeline_mode="thinker_only"): + return stages + + @classmethod + def get_pipeline_id(cls, pipeline_mode="thinker_only"): + return "test_pipeline" + + @classmethod + def get_stage_engine_extras(cls, stage_id, pipeline_mode="thinker_only"): + return {"max_model_len": 65536} if stage_id == 1 else {} + + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: Adapter) + monkeypatch.setattr(ar_strategy_module, "get_visible_devices_keyword", lambda: "CUDA_VISIBLE_DEVICES") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + server = SimpleNamespace( + config=SimpleNamespace( + tensor_model_parallel_size=1, + text_encoder_tp_size=1, + max_model_len=4096, + max_num_batched_tokens=8192, + ), + _rollout_flags={}, + ) + strategy = ARStrategy(server) + engine_kwargs = {"pipeline_name": "adapter"} + + strategy.preprocess_engine_kwargs(engine_kwargs) + + deploy = yaml.safe_load(Path(engine_kwargs["deploy-config"]).read_text(encoding="utf-8")) + assert engine_kwargs["max_model_len"] is None + assert engine_kwargs["max_num_batched_tokens"] is None + assert deploy["stages"][0]["engine_extras"] == { + "max_model_len": 4096, + "max_num_batched_tokens": 8192, + } + assert deploy["stages"][1]["engine_extras"] == { + "max_model_len": 65536, + "max_num_batched_tokens": 8192, + } + server._temp_deploy_ctx.cleanup() + + +@pytest.mark.parametrize( + ("policy_stage_id", "weight_sync_stage_ids", "message"), + [ + (2, [1], "unknown stage 2"), + (0, [2], "unknown stages"), + ], +) +def test_ar_strategy_rejects_unknown_adapter_stages(monkeypatch, policy_stage_id, weight_sync_stage_ids, message): + class Adapter(OmniRolloutPipelineBase): + @classmethod + def build_stage_configs(cls, pipeline_mode="thinker_only"): + return [ + SimpleNamespace(stage_id=0, final_output=False, final_output_type=None, sampling_constraints={}), + SimpleNamespace(stage_id=1, final_output=True, final_output_type="latent", sampling_constraints={}), + ] + + @classmethod + def policy_stage_id(cls, pipeline_mode="thinker_only"): + return policy_stage_id + + @classmethod + def weight_sync_stage_ids(cls, pipeline_mode="thinker_only"): + return weight_sync_stage_ids + + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: Adapter) + strategy = ARStrategy(SimpleNamespace(_rollout_flags={})) + + with pytest.raises(ValueError, match=message): + strategy.preprocess_engine_kwargs({"pipeline_name": "adapter"}) + + def test_ar_strategy_preserves_engine_argument_normalization(): server = SimpleNamespace(config=SimpleNamespace(logprobs_mode="raw_logprobs")) strategy = ARStrategy(server) @@ -218,6 +422,249 @@ def test_ar_strategy_preserves_engine_argument_normalization(): } +def test_ar_strategy_prepares_sampling_params_for_nonzero_policy_stage(): + class Adapter: + @staticmethod + def prepare_engine_prompt(**kwargs): + return { + "prompt_token_ids": [1, 1, 1, 1], + "additional_information": {"text": ["hello"]}, + } + + server = SimpleNamespace( + model_config=SimpleNamespace(), + global_steps=0, + config=SimpleNamespace( + max_model_len=64, + prompt_length=16, + response_length=8, + repetition_penalty=1.0, + ), + engine=SimpleNamespace( + default_sampling_params_list=[SimpleNamespace(stage="thinker"), ar_strategy_module.SamplingParams()] + ), + ) + strategy = ARStrategy(server) + strategy._rollout_adapter = Adapter + strategy._rollout_output_modalities = ["latent", "audio"] + strategy._policy_stage_index = 1 + strategy._policy_sampling_constraints = {} + + prompt, params = strategy.preprocess_input( + [5, 6], + {"temperature": 0.8, "logprobs": True}, + {}, + None, + None, + ) + + assert prompt["additional_information"]["max_new_tokens"] == [8] + assert len(params) == 2 + assert params[0].stage == "thinker" + assert params[1].max_tokens == 8 + assert params[1].temperature == pytest.approx(0.8) + assert params[1].logprobs == 0 + + _, next_params = strategy.preprocess_input( + [5, 6], + {"temperature": 0.2, "logprobs": True}, + {}, + None, + None, + ) + assert params[0] is next_params[0] + assert params[1] is not next_params[1] + assert params[1].temperature == pytest.approx(0.8) + assert next_params[1].temperature == pytest.approx(0.2) + + completion = SimpleNamespace( + token_ids=[7], + logprobs=[{7: SimpleNamespace(logprob=-0.25)}], + finish_reason="stop", + num_preempted=0, + ) + final_res = SimpleNamespace( + request_id="request-0", + request_output=SimpleNamespace(outputs=[completion]), + ) + strategy._rollout_fields_by_request_id["request-0"] = {} + assert strategy.process_output(final_res, params, {}).log_probs == [-0.25] + assert strategy._rollout_fields_by_request_id == {} + + +@pytest.mark.parametrize( + ("adapter_prompt", "message"), + [ + ({"additional_information": {"text": ["hello"]}}, "must contain prompt_token_ids"), + ([1, 2], "must return a dict or None"), + ], +) +def test_ar_strategy_rejects_invalid_adapter_prompt(adapter_prompt, message): + class Adapter: + @staticmethod + def prepare_engine_prompt(**kwargs): + return adapter_prompt + + server = SimpleNamespace( + model_config=SimpleNamespace(), + config=SimpleNamespace(max_model_len=64, prompt_length=16, response_length=8), + ) + strategy = ARStrategy(server) + strategy._rollout_adapter = Adapter + + with pytest.raises((RuntimeError, TypeError), match=message): + strategy.preprocess_input([5, 6], {}, {}, None, None) + + +@pytest.mark.asyncio +async def test_ar_strategy_retains_requested_stage_outputs_and_targets_weight_sync(): + completion = SimpleNamespace(token_ids=[7], logprobs=None, finish_reason="stop", num_preempted=0) + policy = SimpleNamespace(request_id="request-0", outputs=[completion]) + + class Engine: + def __init__(self): + self.generate_kwargs = None + self.rpc_kwargs = None + + async def generate(self, **kwargs): + self.generate_kwargs = kwargs + yield policy + + async def collective_rpc(self, **kwargs): + self.rpc_kwargs = kwargs + return "rpc-result" + + class Adapter: + @staticmethod + def combine_engine_outputs(outputs, prompt): + assert outputs == [policy] + return policy, {"audio_sample_rate": 24_000} + + server = object.__new__(server_module.vLLMOmniHttpServer) + server.engine = Engine() + server.global_steps = 3 + strategy = ARStrategy(server) + strategy._rollout_output_modalities = ["latent", "audio"] + strategy._rollout_adapter = Adapter + strategy._weight_sync_stage_ids = [0] + server._generate_strategy = strategy + + result = await strategy.run_generation( + {"prompt_token_ids": [1]}, ar_strategy_module.SamplingParams(), "request-0", None, 0 + ) + rpc_result = await server.collective_rpc("update_weights_from_ipc", kwargs={"base_sync_done": True}) + + assert result is policy + assert not hasattr(result, "_verl_omni_rollout_fields") + assert strategy._rollout_fields_by_request_id == {"request-0": {"audio_sample_rate": 24_000}} + assert server.engine.generate_kwargs["output_modalities"] == ["latent", "audio"] + assert server.engine.rpc_kwargs["stage_ids"] == [0] + assert rpc_result is None + + output = strategy.process_output(result, ar_strategy_module.SamplingParams(), {}) + assert output.extra_fields == {"global_steps": 3, "audio_sample_rate": 24_000} + assert strategy._rollout_fields_by_request_id == {} + + +def test_ar_strategy_preserves_qwen3_omni_thinker_only_contract(): + server = SimpleNamespace( + config=SimpleNamespace( + max_model_len=8, + prompt_length=4, + response_length=4, + repetition_penalty=1.0, + logprobs_mode="processed_logprobs", + ), + model_config=SimpleNamespace(processor=None), + global_steps=12, + ) + strategy = ARStrategy(server) + strategy._rollout_adapter = Qwen3OmniRolloutAdapter + strategy._rollout_output_modalities = None + + engine_args = {"model_stage": "thinker"} + strategy.prepare_engine_args(engine_args, Namespace(stage_init_timeout=None, init_timeout=None)) + assert engine_args["model_stage"] == "thinker" + + prompt, params = strategy.preprocess_input( + prompt_ids=[1, 2], + sampling_params={"max_new_tokens": 2, "logprobs": True}, + multi_modal_data={}, + lora_request=None, + negative_prompt_ids=None, + ) + assert prompt == {"prompt_token_ids": [1, 2]} + assert isinstance(params, ar_strategy_module.SamplingParams) + + completion = SimpleNamespace( + token_ids=[7], + logprobs=[{7: SimpleNamespace(logprob=-0.25)}], + finish_reason="stop", + num_preempted=0, + ) + output = strategy.process_output( + SimpleNamespace(request_output=SimpleNamespace(outputs=[completion])), + params=params, + sampling_params={}, + ) + assert output.extra_fields == {"global_steps": 12} + + +def test_ar_strategy_writes_qwen3_omni_thinker_only_deploy_config(monkeypatch): + monkeypatch.setattr(ar_strategy_module, "get_visible_devices_keyword", lambda: "CUDA_VISIBLE_DEVICES") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + server = SimpleNamespace( + config=SimpleNamespace( + tensor_model_parallel_size=1, + text_encoder_tp_size=1, + max_model_len=8, + max_num_batched_tokens=8, + ), + _rollout_flags={}, + ) + strategy = ARStrategy(server) + engine_kwargs = { + "pipeline_name": "qwen3_omni_moe", + "pipeline_mode": "thinker_only", + } + + strategy.preprocess_engine_kwargs(engine_kwargs) + + deploy_path = engine_kwargs["deploy-config"] + deploy = yaml.safe_load(Path(deploy_path).read_text(encoding="utf-8")) + assert deploy["pipeline"] == Qwen3OmniRolloutAdapter.get_pipeline_id("thinker_only") + assert [stage["stage_id"] for stage in deploy["stages"]] == [0] + assert strategy._rollout_output_modalities is None + server._temp_deploy_ctx.cleanup() + + +def test_ar_strategy_preserves_qwen3_omni_full_without_combiner(monkeypatch): + monkeypatch.setattr(ar_strategy_module, "get_visible_devices_keyword", lambda: "CUDA_VISIBLE_DEVICES") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + server = SimpleNamespace( + config=SimpleNamespace( + tensor_model_parallel_size=1, + text_encoder_tp_size=1, + max_model_len=8, + max_num_batched_tokens=8, + ), + _rollout_flags={}, + ) + strategy = ARStrategy(server) + + engine_kwargs = { + "pipeline_name": "qwen3_omni_moe", + "pipeline_mode": "full", + } + strategy.preprocess_engine_kwargs(engine_kwargs) + + deploy_path = engine_kwargs["deploy-config"] + deploy = yaml.safe_load(Path(deploy_path).read_text(encoding="utf-8")) + assert [stage["stage_id"] for stage in deploy["stages"]] == [0, 1, 2] + assert strategy._rollout_output_modalities is None + server._temp_deploy_ctx.cleanup() + + def test_diffusion_strategy_preserves_engine_argument_preparation(monkeypatch): imported = [] monkeypatch.setattr(diffusion_strategy_module, "import_external_libs", imported.append) diff --git a/tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py b/tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py index e77169511..603d15ad7 100644 --- a/tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py +++ b/tests/workers/test_diffusers_fsdp_merged_lora_on_cpu.py @@ -32,6 +32,10 @@ def __init__(self): # Carrying a ``peft_config`` is all ``get_per_tensor_param`` needs to # take the LoRA branch. self.peft_config = {"default": SimpleNamespace(to_dict=lambda: {"r": 8})} + self.active_adapter = "default" + + def set_adapter(self, name): + self.active_adapter = name def _make_engine(module, lora_config: dict) -> PPODiffusersFSDPEngine: diff --git a/tests/workers/test_diffusion_distillation_lora_on_cpu.py b/tests/workers/test_diffusion_distillation_lora_on_cpu.py new file mode 100644 index 000000000..38fcf3ef6 --- /dev/null +++ b/tests/workers/test_diffusion_distillation_lora_on_cpu.py @@ -0,0 +1,132 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU regressions for role-aware LoRA switching and export.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from verl_omni.workers.engine.fsdp.diffusers_impl import DiffusersFSDPEngine +from verl_omni.workers.engine.lora_adapter_mixin import LoRAAdapterMixin + + +class PeftConfig: + def __init__(self, name): + self.name = name + + def to_dict(self): + return {"name": self.name} + + +class PeftModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(1.0)) + self.peft_config = { + "default": PeftConfig("default"), + "student": PeftConfig("student"), + "student_ema": PeftConfig("student_ema"), + } + self.active_adapter = "student" + self.adapters_enabled = True + + @property + def active_adapters(self): + return [self.active_adapter] + + def set_adapter(self, name): + self.active_adapter = name[0] if isinstance(name, list) else name + + def disable_adapters(self): + self.adapters_enabled = False + + def enable_adapters(self): + self.adapters_enabled = True + + +class MixinHarness(LoRAAdapterMixin): + def __init__(self): + self.module = PeftModule() + self._is_offload_param = False + + +class TestAdapterContext: + def test_nested_named_adapter_context_restores_exact_selection(self): + harness = MixinHarness() + with harness.use_adapter("student_ema"): + assert harness.module.active_adapter == "student_ema" + with harness.use_adapter("default"): + assert harness.module.active_adapter == "default" + assert harness.module.active_adapter == "student_ema" + assert harness.module.active_adapter == "student" + + def test_exception_restores_previous_adapter(self): + harness = MixinHarness() + with pytest.raises(RuntimeError, match="boom"): + with harness.use_adapter("student_ema"): + raise RuntimeError("boom") + assert harness.module.active_adapter == "student" + + def test_reference_context_reenables_prior_named_adapter(self): + harness = MixinHarness() + with harness.use_adapter("reference"): + assert not harness.module.adapters_enabled + assert harness.module.active_adapter == "student" + assert harness.module.adapters_enabled + assert harness.module.active_adapter == "student" + + +class TestAdapterAwareExport: + @pytest.fixture(autouse=True) + def mock_gpu_memory_logging(self, monkeypatch): + monkeypatch.setattr("verl_omni.workers.engine.fsdp.diffusers_impl.log_gpu_memory_usage", Mock()) + + @pytest.mark.parametrize("adapter_name", [None, "default", "student", "student_ema"]) + def test_selected_adapter_exports_its_own_peft_config(self, monkeypatch, adapter_name): + harness = MixinHarness() + harness._uses_fsdp2_cpu_offload_policy = True + harness.model_config = SimpleNamespace(lora={"merge": False}, fsdp_layer_prefixes=["transformer_blocks."]) + collect_lora_params = Mock(return_value={"adapter.weight": torch.tensor([2.0])}) + + monkeypatch.setattr( + "verl_omni.workers.engine.fsdp.diffusers_impl.collect_lora_params", + collect_lora_params, + ) + monkeypatch.setattr( + "verl_omni.workers.engine.fsdp.diffusers_impl.convert_weight_keys", + lambda params, module: params, + ) + params, peft_config = DiffusersFSDPEngine.get_per_tensor_param( + harness, + base_sync_done=True, + adapter_name=adapter_name, + ) + assert dict(params) == {"transformer.adapter.weight": torch.tensor([2.0])} + assert peft_config == {"name": adapter_name or "default"} + collect_lora_params.assert_called_once() + assert collect_lora_params.call_args.kwargs["adapter_name"] == (adapter_name or "default") + assert harness.module.active_adapter == "student" + + def test_unknown_adapter_fails_before_export(self): + harness = MixinHarness() + harness._uses_fsdp2_cpu_offload_policy = True + harness.model_config = SimpleNamespace(lora={"merge": False}, fsdp_layer_prefixes=[]) + with pytest.raises(ValueError, match="unknown LoRA adapter"): + DiffusersFSDPEngine.get_per_tensor_param( + harness, + base_sync_done=True, + adapter_name="missing", + ) diff --git a/tests/workers/test_dmd_engine_on_cpu.py b/tests/workers/test_dmd_engine_on_cpu.py new file mode 100644 index 000000000..058e7c17c --- /dev/null +++ b/tests/workers/test_dmd_engine_on_cpu.py @@ -0,0 +1,189 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from peft import LoraConfig +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu +from verl.utils.metric import Metric + +from verl_omni.workers.config import DiffusionDMDConfig +from verl_omni.workers.dmd_worker import DMDTrainingWorker +from verl_omni.workers.engine.fsdp import dmd_impl +from verl_omni.workers.engine.fsdp.dmd_impl import DMDDiffusersFSDPEngine +from verl_omni.workers.engine_workers import TrainingWorker + + +class TinyAdapters(torch.nn.Module): + def __init__(self): + super().__init__() + self.adapters = torch.nn.ParameterDict( + {key: torch.nn.Parameter(torch.ones(2)) for key in ("default", "fake_score", "student_ema")} + ) + self.peft_config = {key: LoraConfig(r=2) for key in self.adapters} + self.set_adapter("default") + + def set_adapter(self, name): + self.active_adapter = name + for key, value in self.adapters.items(): + value.requires_grad_(key == name) + + +def constant_schedule(step): + return 1.0 + + +def cpu_device(): + return "cpu" + + +def no_collective(tensor, **kwargs): + return None + + +def force_peer_nonfinite(tensor, **kwargs): + tensor.zero_() + + +def microbatch_metrics(batch, loss_function, forward_only): + mean = batch["value"].mean() + return mean.clone().requires_grad_(), { + "dmd/loss": Metric("mean", mean), + "dmd/active_elements": float(len(batch)), + "perf/forward_s": float(len(batch)), + } + + +def engine_shell(): + engine = object.__new__(DMDDiffusersFSDPEngine) + engine.module = TinyAdapters() + engine.active_stage = "student" + engine.dmd_config = DiffusionDMDConfig(ema_decay=0.5) + engine.role_parameters = { + "student": (engine.module.adapters["default"],), + "fake_score": (engine.module.adapters["fake_score"],), + } + engine.optimizers = {key: torch.optim.SGD(values, lr=0.1) for key, values in engine.role_parameters.items()} + engine.schedulers = { + key: torch.optim.lr_scheduler.LambdaLR(opt, constant_schedule) for key, opt in engine.optimizers.items() + } + engine.optimizer_configs = {key: SimpleNamespace(clip_grad=1.0) for key in engine.optimizers} + engine.optimizer_steps = {"student": 0, "fake_score": 0} + engine.skipped_steps = {"student": 0, "fake_score": 0} + engine.forward_finite = True + engine.last_step_succeeded = False + engine._is_offload_param = False + engine.select_stage("student") + return engine + + +class TestDMDOptimizer: + @pytest.mark.parametrize("stage", ["student", "fake_score"]) + def test_nonfinite_gradient_does_not_step_scheduler_or_ema(self, monkeypatch, stage): + monkeypatch.setattr(dmd_impl, "get_device_id", cpu_device) + monkeypatch.setattr(torch.distributed, "all_reduce", no_collective) + engine = engine_shell() + engine.select_stage(stage) + before = deepcopy(engine.module.state_dict()) + schedule = engine.lr_scheduler.last_epoch + engine.role_parameters[stage][0].grad = torch.full((2,), float("nan")) + engine.optimizer_step() + assert not engine.last_step_succeeded + assert engine.optimizer_steps[stage] == 0 and engine.skipped_steps[stage] == 1 + assert engine.lr_scheduler.last_epoch == schedule + for name, value in engine.module.state_dict().items(): + torch.testing.assert_close(value, before[name]) + + def test_peer_skip_is_agreed_before_local_step(self, monkeypatch): + monkeypatch.setattr(dmd_impl, "get_device_id", cpu_device) + monkeypatch.setattr(torch.distributed, "all_reduce", force_peer_nonfinite) + engine = engine_shell() + engine.role_parameters["student"][0].grad = torch.ones(2) + engine.optimizer_step() + assert not engine.last_step_succeeded + assert engine.optimizer_steps["student"] == 0 + torch.testing.assert_close(engine.module.adapters["default"], torch.ones(2)) + + def test_success_steps_only_owner_and_ema(self, monkeypatch): + monkeypatch.setattr(dmd_impl, "get_device_id", cpu_device) + monkeypatch.setattr(torch.distributed, "all_reduce", no_collective) + engine = engine_shell() + engine.role_parameters["student"][0].grad = torch.ones(2) + engine.optimizer_step() + assert engine.last_step_succeeded + assert engine.optimizer_steps == {"student": 1, "fake_score": 0} + assert engine.schedulers["student"].last_epoch == 1 + assert engine.schedulers["fake_score"].last_epoch == 0 + torch.testing.assert_close(engine.module.adapters["fake_score"], torch.ones(2)) + torch.testing.assert_close( + engine.module.adapters["student_ema"], (torch.ones(2) + engine.module.adapters["default"]) / 2 + ) + engine.lr_scheduler_step() + assert engine.schedulers["student"].last_epoch == 1 + + def test_inactive_gradients_fail_instead_of_cross_role_clipping(self): + engine = engine_shell() + engine.role_parameters["fake_score"][0].grad = torch.ones(2) + with pytest.raises(RuntimeError, match="Gradient leaked"): + engine.optimizer_step() + + +class TestDMDAccumulation: + def test_metric_objects_and_unequal_microbatch_means(self, monkeypatch): + monkeypatch.setattr(dmd_impl, "get_device_id", cpu_device) + monkeypatch.setattr( + dmd_impl, + "get_torch_device", + MagicMock( + return_value=MagicMock( + max_memory_allocated=MagicMock(return_value=1024**3), + max_memory_reserved=MagicMock(return_value=2 * 1024**3), + ) + ), + ) + engine = engine_shell() + engine.ulysses_sequence_parallel_size = 1 + engine.get_data_parallel_group = MagicMock(return_value=None) + engine.forward_step = microbatch_metrics + data = TensorDict({"value": torch.tensor([1.0, 2.0, 3.0])}, batch_size=[3]) + tu.assign_non_tensor(data, micro_batch_size_per_gpu=2) + output = engine.forward_backward_batch(data, None) + assert output["metrics"]["dmd/loss"].aggregate() == pytest.approx(2.0) + assert output["metrics"]["dmd/active_elements"].aggregate() == 3 + assert output["metrics"]["perf/forward_s"].aggregate() == 3 + assert output["metrics"]["perf/max_memory_allocated_gib"].aggregate() == 1 + + +class TestDMDWorker: + def test_reuses_one_minibatch_and_selects_before_context(self, monkeypatch): + result = tu.get_tensordict({}, {"metrics": {"loss": [1.0]}}) + train = MagicMock(return_value=result) + monkeypatch.setattr(TrainingWorker, "train_mini_batch", train) + worker = object.__new__(DMDTrainingWorker) + worker.engine = MagicMock(active_stage="student", last_step_succeeded=False) + worker.dmd_config = DiffusionDMDConfig(fake_score_micro_batch_size_per_gpu=2) + data = TensorDict({}, batch_size=[4]) + tu.assign_non_tensor(data, dmd_stage="fake_score") + output = worker.update_actor(data) + train.assert_called_once() + assert tu.get_non_tensor_data(data, "num_mini_batch", None) == 1 + assert tu.get_non_tensor_data(data, "epochs", None) == 1 + assert tu.get_non_tensor_data(data, "micro_batch_size_per_gpu", None) == 2 + assert tu.get(output, "metrics")["dmd/update_applied"] == 0 + assert [call.args[0] for call in worker.engine.select_stage.call_args_list] == ["fake_score", "student"] diff --git a/tests/workers/test_dmd_fsdp.py b/tests/workers/test_dmd_fsdp.py new file mode 100644 index 000000000..8c8e11ee4 --- /dev/null +++ b/tests/workers/test_dmd_fsdp.py @@ -0,0 +1,177 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Run with torchrun to validate real multi-rank DMD2 engine updates and resume.""" + +import gc +import os +import shutil +import tempfile +from datetime import timedelta +from functools import partial +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from tensordict import TensorDict +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.workers.config import FSDPEngineConfig, FSDPOptimizerConfig + +from verl_omni.workers.config import ( + DiffusionActorConfig, + DiffusionDMDConfig, + DiffusionLossConfig, + DiffusionModelConfig, + DiffusionPipelineConfig, +) +from verl_omni.workers.engine.fsdp.dmd_impl import DMDDiffusersFSDPEngine +from verl_omni.workers.engine.lora_adapter_mixin import load_diffusers_lora_adapter +from verl_omni.workers.utils.losses import diffusion_loss + + +@pytest.fixture(scope="module") +def process_group(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required.") + if dist.is_initialized(): + pytest.skip("This fixture owns its process group.") + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + world = int(os.environ.get("WORLD_SIZE", "1")) + with tempfile.TemporaryDirectory(prefix="dmd2_pg_") as directory: + dist.init_process_group( + "nccl", + init_method="env://" if world > 1 else f"file://{directory}/rdzv", + rank=int(os.environ.get("RANK", "0")), + world_size=world, + timeout=timedelta(seconds=180), + ) + try: + yield + finally: + dist.destroy_process_group() + + +def build_engine(strategy, model_path): + model = DiffusionModelConfig( + path=model_path, + algorithm="dmd2", + model_type="diffusion_dmd_model", + load_tokenizer=False, + lora_rank=2, + lora_alpha=2, + target_modules=["to_q", "to_k", "to_v", "to_out.0"], + attn_backend="native", + enable_gradient_checkpointing=True, + pipeline=DiffusionPipelineConfig(height=64, width=64, num_inference_steps=4, max_sequence_length=64), + ) + config = FSDPEngineConfig(strategy=strategy, use_orig_params=True, model_dtype="bfloat16", seed=7) + optimizer = FSDPOptimizerConfig(lr=1e-4, total_training_steps=3) + engine = DMDDiffusersFSDPEngine( + model, config, optimizer, CheckpointConfig(), dmd_config=DiffusionDMDConfig(ema_decay=0.5) + ) + engine.initialize() + return engine + + +def adapter_values(engine, role): + tensors, _ = engine.get_per_tensor_param(base_sync_done=True, adapter_name=engine.adapter_names[role]) + return {key: value.detach().cpu().clone() for key, value in tensors} + + +def train_stage(engine, stage, batch): + engine.select_stage(stage) + tu.assign_non_tensor(batch, dmd_stage=stage, micro_batch_size_per_gpu=2) + actor = DiffusionActorConfig( + strategy=engine.engine_config.strategy, rollout_n=1, diffusion_loss=DiffusionLossConfig(loss_mode="dmd2") + ) + with engine.train_mode(): + output = engine.train_batch(batch, partial(diffusion_loss, config=actor)) + assert engine.last_step_succeeded + exits = torch.tensor(output["metrics"]["dmd/rollout_exit"].aggregate(), device="cuda") + gathered = [torch.zeros_like(exits) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, exits) + assert all(torch.equal(value, exits) for value in gathered) + return output + + +@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2"]) +def test_qwen_dmd2_engine_update_resume_export(strategy, process_group): + model_path = os.environ.get("QWEN_IMAGE_MODEL_PATH", os.path.expanduser("~/models/tiny-random/Qwen-Image")) + if not Path(model_path, "model_index.json").is_file(): + pytest.skip(f"Tiny Qwen checkpoint not found: {model_path}") + paths = [tempfile.mkdtemp(prefix="dmd2_fsdp_") if dist.get_rank() == 0 else None] + dist.broadcast_object_list(paths, src=0) + directory = Path(paths[0]) + engine = build_engine(strategy, model_path) + batch = TensorDict({"dummy_tensor": torch.zeros(3, 1)}, batch_size=[3]) + tu.assign_non_tensor_stack( + batch, + "raw_prompt", + [ + [{"role": "user", "content": "cat" if dist.get_rank() % 2 else "a red apple on a table"}], + [{"role": "user", "content": "a blue bird"}], + [{"role": "user", "content": "a green triangle"}], + ], + ) + try: + initial = {role: adapter_values(engine, role) for role in ("student", "fake_score")} + for _ in range(3): + for stage in ("student", "fake_score", "fake_score"): + train_stage(engine, stage, batch) + assert engine.optimizer_steps == {"student": 3, "fake_score": 6} + saved = {role: adapter_values(engine, role) for role in ("student", "fake_score", "student_ema")} + for role, values in initial.items(): + assert any(torch.count_nonzero(saved[role][key] - value) > 0 for key, value in values.items()) + engine.save_checkpoint(str(directory / "actor"), global_step=3) + replay = train_stage(engine, "student", batch) + expected = adapter_values(engine, "student") + restored = engine.load_checkpoint(str(directory / "actor"), del_local_after_load=False) + assert restored == {"student": 3, "fake_score": 6} + for role, parameters in saved.items(): + for key, value in adapter_values(engine, role).items(): + torch.testing.assert_close(value, parameters[key], rtol=0, atol=0) + repeated = train_stage(engine, "student", batch) + assert repeated["loss"] == pytest.approx(replay["loss"], rel=1e-6, abs=1e-8) + for key, value in adapter_values(engine, "student").items(): + torch.testing.assert_close(value, expected[key], rtol=0, atol=0) + engine.export_student(str(directory / "inference"), role="student") + assert (directory / "inference" / "adapter_model.safetensors").is_file() + from diffusers import QwenImageTransformer2DModel + from peft import get_peft_model_state_dict + from safetensors.torch import load_file + + reloaded = QwenImageTransformer2DModel.from_pretrained( + model_path, subfolder="transformer", torch_dtype=torch.bfloat16 + ) + load_diffusers_lora_adapter(reloaded, directory / "inference", "reloaded") + state = get_peft_model_state_dict(reloaded, adapter_name="reloaded") + exported = load_file(directory / "inference" / "adapter_model.safetensors") + assert state.keys() == exported.keys() + for key, value in state.items(): + torch.testing.assert_close(value, exported[key], rtol=0, atol=0, check_dtype=False) + before = dict(engine.optimizer_steps) + scheduler_step = engine.lr_scheduler.last_epoch + engine.forward_finite = dist.get_rank() != 0 + engine.optimizer_step() + assert not engine.last_step_succeeded + assert engine.optimizer_steps == before and engine.lr_scheduler.last_epoch == scheduler_step + finally: + dist.barrier() + del engine + gc.collect() + torch.cuda.empty_cache() + if dist.get_rank() == 0: + shutil.rmtree(directory) + dist.barrier() diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 3cd9137d1..897d9a423 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -58,6 +58,7 @@ def _make_mock_model_config(**overrides): cfg.model_stage = "thinker" cfg.local_path = "/fake/model/path" cfg.trust_remote_code = False + cfg.external_lib = None cfg.use_liger = False cfg.use_fused_kernels = False cfg.enable_gradient_checkpointing = False @@ -83,7 +84,6 @@ def _get(key, default=None): # Isolated module loader # --------------------------------------------------------------------------- - _omni_impl_cache = None @@ -328,6 +328,26 @@ def prepare_model_inputs(cls, model_inputs, replay_batch, model_config): assert output_args == {"base": True} +def test_weight_sync_casts_floating_dtensor_to_bfloat16(): + omni_impl = _get_omni_impl_module() + tensor = torch.tensor([1.25], dtype=torch.float32) + + synced = omni_impl.OmniFSDPEngine._cast_dtensor_weight_for_sync(tensor) + + assert synced.dtype is torch.bfloat16 + assert synced.item() == pytest.approx(1.25) + + +def test_weight_sync_keeps_integer_dtensor_buffers(): + omni_impl = _get_omni_impl_module() + tensor = torch.tensor([1, 2], dtype=torch.int64) + + synced = omni_impl.OmniFSDPEngine._cast_dtensor_weight_for_sync(tensor) + + assert synced is tensor + assert synced.dtype is torch.int64 + + # --------------------------------------------------------------------------- # ``collect_lora_params`` import source # --------------------------------------------------------------------------- @@ -359,8 +379,8 @@ def test_collect_lora_params_import_not_from_verl(): # --------------------------------------------------------------------------- -def test_build_module_uses_auto_model_for_multimodal_lm(): - """``_build_module`` uses ``AutoModelForMultimodalLM``, not ``AutoModelForCausalLM``.""" +def test_build_module_uses_adapter_selected_auto_model_class(): + """Adapters select non-default auto model classes without relying on stage names.""" omni_impl = _get_omni_impl_module() assert omni_impl.AutoModelForMultimodalLM is not None @@ -373,19 +393,27 @@ def test_build_module_uses_auto_model_for_multimodal_lm(): assert "AutoModelForMultimodalLM" in import_names, ( f"AutoModelForMultimodalLM not imported from transformers; imports: {import_names}" ) + assert "AutoModelForTextToWaveform" not in import_names assert "AutoModelForCausalLM" not in import_names, "AutoModelForCausalLM should NOT be imported from transformers" -@pytest.mark.parametrize("architecture", ["Qwen3OmniMoeForConditionalGeneration"]) -def test_build_module_calls_adapter_configure_model(architecture): +@pytest.mark.parametrize( + ("architecture", "model_stage"), + [ + ("Qwen3OmniMoeForConditionalGeneration", "thinker"), + ("FutureOmniForConditionalGeneration", "talker"), + ], +) +def test_build_module_calls_adapter_configure_model(architecture, model_stage): """Mock ``from_pretrained``; verify ``adapter_cls.configure_model(module, cfg)``.""" omni_impl = _get_omni_impl_module() - model_config = _make_mock_model_config(architecture=architecture) + model_config = _make_mock_model_config(architecture=architecture, model_stage=model_stage) fake_module = MagicMock(spec=torch.nn.Module) fake_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls = MagicMock() + fake_adapter_cls.auto_model_class = None fake_configured_module = MagicMock(spec=torch.nn.Module) fake_configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls.configure_model.return_value = fake_configured_module @@ -399,13 +427,15 @@ def test_build_module_calls_adapter_configure_model(architecture): patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=fake_adapter_cls) as mock_get_cls, patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), - patch("verl.utils.torch_dtypes.PrecisionType"), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, ): + precision_type.to_dtype.side_effect = lambda value: value engine = object.__new__(omni_impl.OmniFSDPEngine) engine.model_config = model_config engine.engine_config = MagicMock() engine.engine_config.model_dtype = None engine.engine_config.forward_only = False + engine.engine_config.strategy = "fsdp2" engine.device_mesh = None result = engine._build_module() @@ -423,9 +453,82 @@ def test_build_module_calls_adapter_configure_model(architecture): ) fake_adapter_cls.configure_model.assert_called_once_with(fake_module, model_config) + assert engine.model_adapter_cls is fake_adapter_cls assert result is fake_configured_module +def test_build_module_calls_adapter_selected_auto_model_loader(): + omni_impl = _get_omni_impl_module() + model_config = _make_mock_model_config( + architecture="FutureOmniForConditionalGeneration", + model_stage="talker", + ) + loaded_module = MagicMock(spec=torch.nn.Module) + loaded_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] + configured_module = MagicMock(spec=torch.nn.Module) + configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] + auto_model_cls = MagicMock() + auto_model_cls.from_pretrained.return_value = loaded_module + adapter_cls = MagicMock() + adapter_cls.auto_model_class = auto_model_cls + adapter_cls.configure_model.return_value = configured_module + model_base_mod = sys.modules["verl_omni.pipelines.model_base"] + + with ( + patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), + patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), + patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, + ): + precision_type.to_dtype.side_effect = lambda value: value + engine = object.__new__(omni_impl.OmniFSDPEngine) + engine.model_config = model_config + engine.engine_config = MagicMock(model_dtype=None, forward_only=False) + engine.engine_config.strategy = "fsdp2" + engine.device_mesh = None + + result = engine._build_module() + + auto_model_cls.from_pretrained.assert_called_once_with( + pretrained_model_name_or_path=model_config.local_path, + torch_dtype=torch.float32, + config=model_config.hf_config, + trust_remote_code=model_config.trust_remote_code, + ) + adapter_cls.configure_model.assert_called_once_with(loaded_module, model_config) + assert result is configured_module + + +def test_build_module_rejects_mixed_frozen_parameters_without_fsdp1_orig_params(): + omni_impl = _get_omni_impl_module() + model_config = _make_mock_model_config() + loaded_module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + configured_module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + configured_module[0].requires_grad_(False) + adapter_cls = MagicMock() + adapter_cls.auto_model_class = None + adapter_cls.configure_model.return_value = configured_module + model_base_mod = sys.modules["verl_omni.pipelines.model_base"] + + with ( + patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), + patch.object(omni_impl.AutoModelForMultimodalLM, "from_pretrained", return_value=loaded_module), + patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), + patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, + ): + precision_type.to_dtype.side_effect = lambda value: value + engine = object.__new__(omni_impl.OmniFSDPEngine) + engine.model_config = model_config + engine.engine_config = MagicMock(model_dtype=None, forward_only=False) + engine.engine_config.strategy = "fsdp" + engine.engine_config.use_orig_params = False + engine.device_mesh = None + + with pytest.raises(ValueError, match="use_orig_params=true"): + engine._build_module() + + @pytest.mark.parametrize("option", ["use_liger", "use_fused_kernels"]) def test_build_module_rejects_unsupported_optimizations_before_model_load(option): omni_impl = _get_omni_impl_module() diff --git a/verl_omni/agent_loop/composite_agent_loop.py b/verl_omni/agent_loop/composite_agent_loop.py index 55306546f..51da08693 100644 --- a/verl_omni/agent_loop/composite_agent_loop.py +++ b/verl_omni/agent_loop/composite_agent_loop.py @@ -48,6 +48,29 @@ def _config_to_sampling_dict(config: Optional[BaseConfig]) -> dict: return {k: v for k, v in config.items() if not k.startswith("_")} +def _pad_llm_generation_outputs( + response_ids: torch.Tensor, + log_probs: torch.Tensor | None, + max_new_tokens: int, + pad_token_id: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Right-pad per-request AR generation tensors to ``max_new_tokens``. + + Per-request ``generate`` stops at EOS, so samples reach ``_postprocess`` + with different lengths while it batches them with ``torch.cat``. + """ + gen_len = int(response_ids.shape[-1]) + if gen_len > max_new_tokens: + raise ValueError(f"llm_response_ids length {gen_len} exceeds rollout.max_new_tokens={max_new_tokens}") + attention_mask = torch.zeros(*response_ids.shape[:-1], max_new_tokens, dtype=torch.long, device=response_ids.device) + attention_mask[..., :gen_len] = 1 + padded_ids = F.pad(response_ids, (0, max_new_tokens - gen_len), value=pad_token_id) + padded_log_probs = None + if log_probs is not None: + padded_log_probs = F.pad(log_probs, (0, 0, 0, max_new_tokens - log_probs.shape[-2]), value=0.0) + return padded_ids, attention_mask, padded_log_probs + + class CompositeAgentLoopOutput(DiffusionAgentLoopOutput): """Agent loop output. Supplement additional fields for AR part.""" @@ -188,6 +211,18 @@ async def _agent_loop_postprocess( """Perform post-processing operations on the output of each individual agent loop.""" output = CompositeAgentLoopOutput(**dict(output)) + llm_response_ids = output.extra_fields.get("llm_response_ids") + llm_log_probs = output.extra_fields.get("llm_all_log_probs") + llm_response_mask: torch.Tensor | None = None + if isinstance(llm_response_ids, torch.Tensor): + pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0 + llm_response_ids, llm_response_mask, llm_log_probs = _pad_llm_generation_outputs( + llm_response_ids, + llm_log_probs if isinstance(llm_log_probs, torch.Tensor) else None, + self.rollout_config.max_new_tokens, + pad_token_id, + ) + # Pad extra tensor outputs from vllm-omni (e.g. prompt embeddings). extra_fields = {} for k, v in output.extra_fields.items(): @@ -198,9 +233,15 @@ async def _agent_loop_postprocess( elif k in ["prompt_embeds_mask", "negative_prompt_embeds_mask"]: pad_tuple = (0, self.max_prompt_embed_length - v.shape[0]) v = F.pad(v, pad_tuple, value=0) + elif k == "llm_response_ids" and llm_response_ids is not None: + v = llm_response_ids + elif k == "llm_all_log_probs" and llm_log_probs is not None: + v = llm_log_probs extra_fields[k] = v.unsqueeze(0) else: extra_fields[k] = v + if llm_response_mask is not None: + extra_fields["llm_response_attention_mask"] = llm_response_mask.unsqueeze(0) extra_fields["raw_prompt"] = kwargs["raw_prompt"] @@ -226,8 +267,8 @@ async def _agent_loop_postprocess( if output.response_logprobs is not None: response_logprobs = output.response_logprobs.unsqueeze(0) llm_response_logprobs = None - if output.extra_fields.get("llm_all_log_probs", None) is not None: - llm_response_logprobs = output.extra_fields["llm_all_log_probs"].unsqueeze(0) + if llm_log_probs is not None: + llm_response_logprobs = llm_log_probs.unsqueeze(0) prompt_ids = prompt_output["input_ids"] extra_fields["attention_mask"] = prompt_output["attention_mask"] diff --git a/verl_omni/pipelines/__init__.py b/verl_omni/pipelines/__init__.py index a1acf099d..c9794e198 100644 --- a/verl_omni/pipelines/__init__.py +++ b/verl_omni/pipelines/__init__.py @@ -21,6 +21,7 @@ minimax_h3_flow_grpo, qwen3_omni, qwen_image_diffusion_nft, + qwen_image_distillation, qwen_image_dpo, qwen_image_dual_grpo, qwen_image_edit_flow_grpo, @@ -38,6 +39,7 @@ from .minimax_h3_flow_grpo import * # noqa: F401, F403 from .qwen3_omni import * # noqa: F401, F403 from .qwen_image_diffusion_nft import * # noqa: F401, F403 +from .qwen_image_distillation import * # noqa: F401, F403 from .qwen_image_dpo import * # noqa: F401, F403 from .qwen_image_dual_grpo import * # noqa: F401, F403 from .qwen_image_edit_flow_grpo import * # noqa: F401, F403 @@ -50,6 +52,7 @@ __all__ = list(qwen3_omni.__all__) __all__ += list(qwen_image_flow_grpo.__all__) __all__ += list(qwen_image_diffusion_nft.__all__) +__all__ += list(qwen_image_distillation.__all__) __all__ += list(qwen_image_mix_grpo.__all__) __all__ += list(bagel_flow_grpo.__all__) __all__ += list(ltx2_flow_grpo.__all__) diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index bd5671d98..cda74d895 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -479,6 +479,7 @@ class Qwen3OmniThinkerAdapter(OmniModelBase): """ _registry: dict[tuple[str, str], type["OmniModelBase"]] = {} + auto_model_class: Any = None @classmethod def register(cls, architecture: str, stage: str = "thinker"): @@ -511,6 +512,11 @@ def get_class(cls, model_config) -> type["OmniModelBase"]: getattr(model_config, "external_lib", None), ) + @classmethod + def peek_class(cls, architecture: str, stage: str) -> Optional[type["OmniModelBase"]]: + """Return the registered adapter for ``(architecture, stage)`` or ``None``.""" + return cls._registry.get((architecture, stage)) + @classmethod def get_class_by_name( cls, @@ -546,6 +552,11 @@ def get_class_by_name( f"Set ``external_lib`` to load your training adapter." ) from None + @classmethod + def register_auto_classes(cls) -> None: + """Register optional model-package classes with Transformers auto APIs.""" + return + @classmethod @abstractmethod def get_strip_modules(cls, model_config) -> list[str]: @@ -615,7 +626,6 @@ def configure_model(cls, module, model_config): Default implementation strips the submodules returned by ``get_strip_modules``. Override to also: - - Register the model class with ``AutoModelForCausalLM``. - Redirect ``forward()`` and embedding accessors to the trainable sub-component. - Force ``tie_word_embeddings=False`` for FSDP compatibility. @@ -670,6 +680,7 @@ class Qwen3OmniRolloutAdapter(OmniRolloutPipelineBase): """ _registry: dict[str, type["OmniRolloutPipelineBase"]] = {} + supports_async_chunk = True @classmethod def register(cls, model_type: str): @@ -746,6 +757,16 @@ def rollout_flags(cls, pipeline_mode="thinker_only") -> dict[int, dict]: """ return {} + @classmethod + def weight_sync_stage_ids(cls, pipeline_mode="thinker_only") -> list[int] | None: + """Return stages that receive actor weights, or all stages by default.""" + return None + + @classmethod + def policy_stage_id(cls, pipeline_mode="thinker_only") -> int: + """Return the stage whose sampling parameters and logprobs define the policy.""" + return 0 + @classmethod def get_pipeline_id(cls, pipeline_mode: str = "thinker_only") -> str: """Return the vLLM-Omni pipeline model_type for *pipeline_mode*. @@ -800,3 +821,30 @@ def get_stage_engine_extras(cls, stage_id: int, pipeline_mode: str = "thinker_on dict: Extra key-value pairs merged into the stage's engine args. """ return {} + + @classmethod + def prepare_engine_prompt( + cls, + prompt_ids: list[int], + model_config, + multi_modal_data: dict, + mm_processor_kwargs: Optional[dict] = None, + ) -> dict | None: + """Build an architecture-specific rollout prompt when required.""" + return None + + @classmethod + def combine_engine_outputs(cls, outputs: list, prompt: dict) -> tuple[Any, dict[str, Any]]: + """Select the policy output and collect architecture-specific fields. + + Overriding this hook opts an adapter into retaining and assembling + multiple final stage outputs. The default preserves single-output AR + behavior. + """ + if not outputs: + raise RuntimeError("The omni rollout engine returned no outputs.") + if len(outputs) != 1: + raise NotImplementedError( + "An omni rollout adapter with multiple final outputs must implement combine_engine_outputs()." + ) + return outputs[0], {} diff --git a/verl_omni/pipelines/qwen_image_distillation/__init__.py b/verl_omni/pipelines/qwen_image_distillation/__init__.py new file mode 100644 index 000000000..22a29fa0d --- /dev/null +++ b/verl_omni/pipelines/qwen_image_distillation/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .diffusers_training_adapter import QwenImageDMD2 + +__all__ = ["QwenImageDMD2"] diff --git a/verl_omni/pipelines/qwen_image_distillation/diffusers_training_adapter.py b/verl_omni/pipelines/qwen_image_distillation/diffusers_training_adapter.py new file mode 100644 index 000000000..87a73ceb3 --- /dev/null +++ b/verl_omni/pipelines/qwen_image_distillation/diffusers_training_adapter.py @@ -0,0 +1,338 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen-Image T2I adapter for offline DMD2 distribution matching.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Optional + +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.pipelines.model_base import DiffusionModelBase +from verl_omni.pipelines.qwen_image_flow_grpo.common import QwenImageTokenIdPromptMixin +from verl_omni.pipelines.qwen_image_flow_grpo.diffusers_training_adapter import QwenImage + +__all__ = ["QwenImageDMD2"] + + +def build_qwen_dmd_sigmas(num_inference_steps, shift, device): + """Build the fixed, once-shifted Euler grid shared by training and inference.""" + from verl_omni.trainer.diffusion.distillation.utils import timestep_shift + + if isinstance(num_inference_steps, bool) or not isinstance(num_inference_steps, int) or num_inference_steps <= 0: + raise ValueError("num_inference_steps must be a positive integer.") + return timestep_shift(torch.linspace(1, 0, num_inference_steps + 1, device=device), 1, shift) + + +@DiffusionModelBase.register("QwenImagePipeline", algorithm="dmd2") +class QwenImageDMD2(QwenImage): + """Base Qwen T2I flow adapter; optimization stays in the DMD engine.""" + + @classmethod + def sampling_sigmas(cls, model_config, dmd_config, device): + """Use Qwen DMD2's fixed shift rather than the native resolution-dependent shift.""" + return build_qwen_dmd_sigmas( + model_config.pipeline.num_inference_steps, dmd_config.rollout_timestep_shift, device + ) + + @classmethod + def configure_train_mode(cls, module): + """Keep sampling and checkpoint recomputation in evaluation mode with autograd enabled.""" + module.eval() + + @classmethod + def build_conditioning_provider(cls, model_config, dmd_config): + """Build the frozen/local-or-cached prompt provider once per engine.""" + return QwenImageConditionProvider( + model_config.local_path or model_config.path, + model_config.pipeline.max_sequence_length, + dmd_config.negative_prompt, + ) + + @staticmethod + def batch_dimension(batch, key, default): + """Require same-resolution integer geometry within a physical batch.""" + value = tu.get(batch, key, default) + values = torch.as_tensor(value).reshape(-1) + if values.numel() == 0 or not torch.all(values == values[0]) or not torch.all(values == values.long()): + raise ValueError(f"Qwen DMD2 requires homogeneous integer {key} values.") + return int(values[0]) + + @classmethod + def latent_geometry(cls, module, model_config, batch): + """Read VAE geometry and return normalized latent shape plus forward metadata.""" + if len(batch.batch_size) != 1 or batch.batch_size[0] <= 0: + raise ValueError("Qwen DMD2 requires a nonempty leading batch dimension.") + with open(Path(model_config.local_path or model_config.path) / "vae" / "config.json") as file: + vae_config = json.load(file) + scale = 2 ** len(vae_config["temperal_downsample"]) + channels = vae_config["z_dim"] + model = getattr(module, "_fsdp_wrapped_module", module) + if model.config.in_channels != channels * 4: + raise ValueError("Qwen transformer packing and VAE channel counts do not match.") + height = cls.batch_dimension(batch, "height", model_config.pipeline.height) + width = cls.batch_dimension(batch, "width", model_config.pipeline.width) + if height <= 0 or width <= 0 or height % (scale * 2) or width % (scale * 2): + raise ValueError(f"Qwen image dimensions must be positive multiples of {scale * 2}.") + shape = (batch.batch_size[0], channels, 1, height // scale, width // scale) + return shape, {"height": height, "width": width, "vae_scale_factor": scale} + + @staticmethod + def pack_latents(latents): + """Pack declared normalized image latents using the native Qwen helper.""" + from diffusers import QwenImagePipeline + + batch, channels, _, height, width = latents.shape + return QwenImagePipeline._pack_latents(latents, batch, channels, height, width) + + @classmethod + def prepare_dmd_inputs(cls, module, model_config, latents, sigma, condition, geometry): + """Reuse Qwen's input builder without its policy-gradient CFG or SDE step.""" + metadata = TensorDict({}, batch_size=[latents.shape[0]]) + tu.assign_non_tensor(metadata, **geometry) + sigma = sigma.reshape(-1).expand(latents.shape[0]) + inputs, _ = super().prepare_model_inputs( + module, + model_config, + latents.unsqueeze(1), + (sigma * 1000).unsqueeze(1), + condition["prompt_embeds"], + condition["prompt_embeds_mask"], + None, + None, + metadata, + 0, + ) + dtype = getattr(module, "dtype", latents.dtype) + inputs["hidden_states"] = inputs["hidden_states"].to(dtype=dtype) + inputs["encoder_hidden_states"] = inputs["encoder_hidden_states"].to(dtype=dtype) + return inputs + + @staticmethod + def prediction_to_x0(noisy, prediction, sigma): + """Translate packed Qwen flow velocity to canonical fp32 clean latents.""" + from verl_omni.trainer.diffusion.distillation.utils import velocity_to_x0 + + return velocity_to_x0(noisy, prediction, sigma) + + +class QwenImageConditionProvider: + """Encode frozen local or precomputed Qwen prompt conditioning.""" + + def __init__( + self, + model_path: str, + max_sequence_length: int, + negative_prompt: str, + ) -> None: + self.model_path = model_path + self.max_sequence_length = max_sequence_length + self.negative_prompt = negative_prompt + self.pipeline = None + self.negative_condition: Optional[dict[str, torch.Tensor]] = None + + @staticmethod + def make_condition(prompt_embeds: torch.Tensor, prompt_mask: Optional[torch.Tensor]) -> dict[str, torch.Tensor]: + """Build a detached [B, L, D] condition with a matching [B, L] mask.""" + if prompt_embeds.ndim != 3 or any(size == 0 for size in prompt_embeds.shape): + raise ValueError(f"Qwen prompt embeddings must have shape [B, L, D], got {tuple(prompt_embeds.shape)}.") + prompt_embeds = prompt_embeds.detach() + if prompt_mask is None: + prompt_mask = torch.ones(prompt_embeds.shape[:2], device=prompt_embeds.device, dtype=torch.long) + elif prompt_mask.shape != prompt_embeds.shape[:2]: + raise ValueError( + f"Qwen prompt mask shape {tuple(prompt_mask.shape)} does not match {tuple(prompt_embeds.shape[:2])}." + ) + return {"prompt_embeds": prompt_embeds, "prompt_embeds_mask": prompt_mask.detach()} + + @staticmethod + def require_tensor(batch: TensorDict, key: str) -> torch.Tensor: + """Require a tensor-valued precomputed conditioning field.""" + value = tu.get(batch, key) + if not isinstance(value, torch.Tensor): + raise ValueError(f"Precomputed Qwen conditioning requires tensor batch field {key!r}.") + return value + + def encode_precomputed( + self, + batch: TensorDict, + *, + require_negative: bool, + ) -> tuple[dict[str, torch.Tensor], Optional[dict[str, torch.Tensor]]]: + """Validate and truncate cached positive and negative conditioning.""" + prompt_embeds = self.require_tensor(batch, "prompt_embeds")[:, : self.max_sequence_length] + prompt_mask = tu.get(batch, "prompt_embeds_mask") + if isinstance(prompt_mask, torch.Tensor): + prompt_mask = prompt_mask[:, : self.max_sequence_length] + negative_embeds = ( + self.require_tensor(batch, "negative_prompt_embeds")[:, : self.max_sequence_length] + if require_negative + else None + ) + negative_mask = tu.get(batch, "negative_prompt_embeds_mask") if require_negative else None + if isinstance(negative_mask, torch.Tensor): + negative_mask = negative_mask[:, : self.max_sequence_length] + if prompt_mask is not None and not isinstance(prompt_mask, torch.Tensor): + raise TypeError("prompt_embeds_mask must be a tensor when supplied.") + if negative_mask is not None and not isinstance(negative_mask, torch.Tensor): + raise TypeError("negative_prompt_embeds_mask must be a tensor when supplied.") + positive = self.make_condition(prompt_embeds, prompt_mask) + negative = self.make_condition(negative_embeds, negative_mask) if negative_embeds is not None else None + if positive["prompt_embeds"].shape[0] != batch.batch_size[0]: + raise ValueError("Precomputed Qwen conditioning batch size does not match the input batch.") + if negative is not None and negative["prompt_embeds"].shape[0] != batch.batch_size[0]: + raise ValueError("Precomputed Qwen negative conditioning batch size does not match the input batch.") + return positive, negative + + def ensure_pipeline(self, device: torch.device, dtype: torch.dtype): + """Load the frozen checkpoint text encoder once on the execution device.""" + if self.pipeline is not None: + return self.pipeline + + from diffusers import QwenImagePipeline + + pipeline = QwenImagePipeline.from_pretrained( + self.model_path, + transformer=None, + vae=None, + torch_dtype=dtype, + local_files_only=os.path.isdir(self.model_path), + ).to(device) + pipeline.text_encoder.requires_grad_(False) + pipeline.text_encoder.eval() + self.pipeline = pipeline + return pipeline + + @staticmethod + def prompt_rows(value: Any, batch_size: int, key: str) -> list[Any]: + """Unwrap one text or chat-message row per batch sample.""" + if hasattr(value, "tolist") and not isinstance(value, torch.Tensor): + value = value.tolist() + if batch_size == 1 and ( + isinstance(value, str) or (isinstance(value, list) and value and isinstance(value[0], dict)) + ): + return [value] + if not isinstance(value, list) or len(value) != batch_size: + raise ValueError(f"{key} must contain exactly {batch_size} prompt row(s).") + return value + + def tokenize_rows(self, pipeline, rows: list[Any], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the fixed Qwen template before tokenization and prefix removal.""" + rendered = [] + for row in rows: + if isinstance(row, list): + if len(row) != 1 or not isinstance(row[0], dict) or row[0].get("role") != "user": + raise ValueError( + "Qwen DMD raw prompts require a single user message; use precomputed conditioning otherwise." + ) + row = row[0].get("content") + if not isinstance(row, str): + raise TypeError("Qwen DMD prompts must be strings or single text-only user messages.") + # The encoder removes this template's fixed prefix, not a generic chat prefix. + rendered.append(pipeline.prompt_template_encode.format(row)) + tokens = pipeline.tokenizer( + rendered, + max_length=self.max_sequence_length + pipeline.prompt_template_encode_start_idx, + padding=True, + truncation=True, + return_tensors="pt", + ) + return tokens.input_ids.to(device), tokens.attention_mask.to(device) + + def encode_ids( + self, + pipeline, + prompt_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Reuse Qwen token-ID encoding under no-grad and truncate the result.""" + with torch.no_grad(): + prompt_embeds, prompt_mask = QwenImageTokenIdPromptMixin._get_qwen_prompt_embeds( + pipeline, prompt_ids, attention_mask=attention_mask + ) + prompt_embeds = prompt_embeds[:, : self.max_sequence_length] + if prompt_mask is not None: + prompt_mask = prompt_mask[:, : self.max_sequence_length] + if prompt_embeds.shape[1] == 0: + raise ValueError("Qwen prompt encoding produced no tokens after removing the template prefix.") + return self.make_condition(prompt_embeds.detach(), prompt_mask.detach() if prompt_mask is not None else None) + + @torch.profiler.record_function("distillation/condition_encode") + def encode( + self, + batch: TensorDict, + *, + device: torch.device, + dtype: torch.dtype, + require_negative: bool, + ) -> tuple[dict[str, torch.Tensor], Optional[dict[str, torch.Tensor]]]: + """Encode the positive prompt and optional teacher negative condition.""" + if "prompt_embeds" in batch: + return self.encode_precomputed(batch, require_negative=require_negative) + + pipeline = self.ensure_pipeline(device, dtype) + prompt_ids = tu.get(batch, "prompt_ids", tu.get(batch, "prompts")) + prompt_mask = tu.get( + batch, + "prompt_attention_mask", + tu.get(batch, "prompt_mask", tu.get(batch, "attention_mask")), + ) + if prompt_ids is None: + raw_prompt = tu.get(batch, "raw_prompt") + if raw_prompt is None: + raise ValueError("Qwen DMD batches require prompt_ids, prompts, prompt_embeds, or raw_prompt.") + rows = self.prompt_rows(raw_prompt, batch.batch_size[0], "raw_prompt") + prompt_ids, prompt_mask = self.tokenize_rows(pipeline, rows, device) + else: + if prompt_mask is None: + raise ValueError("Pre-tokenized Qwen prompts require prompt_attention_mask or attention_mask.") + if not isinstance(prompt_ids, torch.Tensor): + prompt_ids = torch.as_tensor(prompt_ids, device=device, dtype=torch.long) + else: + prompt_ids = prompt_ids.to(device=device, dtype=torch.long) + if prompt_mask is not None: + prompt_mask = torch.as_tensor(prompt_mask, device=device, dtype=torch.long) + positive = self.encode_ids(pipeline, prompt_ids, prompt_mask) + + if not require_negative: + return positive, None + + negative_ids = tu.get(batch, "negative_prompt_ids") + negative_mask = tu.get(batch, "negative_prompt_attention_mask", tu.get(batch, "negative_prompt_mask")) + if negative_ids is None: + raw_negative = tu.get(batch, "raw_negative_prompt", tu.get(batch, "negative_prompt")) + if raw_negative is not None: + negative_rows = self.prompt_rows(raw_negative, batch.batch_size[0], "negative_prompt") + negative_ids, negative_mask = self.tokenize_rows(pipeline, negative_rows, device) + negative = self.encode_ids(pipeline, negative_ids, negative_mask) + else: + if self.negative_condition is None: + negative_ids, negative_mask = self.tokenize_rows(pipeline, [self.negative_prompt], device) + self.negative_condition = self.encode_ids(pipeline, negative_ids, negative_mask) + negative = self.make_condition( + self.negative_condition["prompt_embeds"].expand(batch.batch_size[0], -1, -1), + self.negative_condition["prompt_embeds_mask"].expand(batch.batch_size[0], -1), + ) + else: + if negative_mask is None: + raise ValueError("Pre-tokenized negative Qwen prompts require negative_prompt_attention_mask.") + negative_ids = torch.as_tensor(negative_ids, device=device, dtype=torch.long) + negative_mask = torch.as_tensor(negative_mask, device=device, dtype=torch.long) + negative = self.encode_ids(pipeline, negative_ids, negative_mask) + return positive, negative diff --git a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml index 3cb671deb..46946fd65 100644 --- a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml @@ -498,15 +498,35 @@ distillation: world_size: 0 teacher_key: data_source scheduler: inline - distribution_matching: - _target_: verl_omni.workers.config.diffusion.DiffusionDistributionMatchingConfig - recipe: dmd2 - profile: null - fake_update_ratio: null - fake_warmup_cycles: 0 - rollout_strategy: null - data_mode: null - export_role: student_ema +dmd: + _target_: verl_omni.workers.config.DiffusionDMDConfig + fake_update_ratio: 2 + student_micro_batch_size_per_gpu: 1 + fake_score_micro_batch_size_per_gpu: 1 + fake_score_optim: + _target_: verl.workers.config.FSDPOptimizerConfig + lr: 2.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.999 + clip_grad: 1.0 + lr_scheduler_type: constant + lr_warmup_steps: -1 + lr_warmup_steps_ratio: 0.0 + total_training_steps: -1 + teacher_guidance_scale: 4.0 + cfg_norm: layer_norm + negative_prompt: ' ' + normalization_epsilon: 1.0e-06 + rollout_timestep_shift: 3.0 + score_discrete_steps: 1000 + score_sigma_min: 0.02 + score_sigma_max: 0.98 + score_timestep_shift: 3.0 + ema_decay: 0.999 + ema_start_step: 0 + export_role: student algorithm: _target_: verl_omni.trainer.config.DiffusionAlgoConfig trainer_type: policy_gradient diff --git a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml index d2c2ac6a7..935acea0b 100644 --- a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml @@ -539,15 +539,35 @@ distillation: world_size: 0 teacher_key: data_source scheduler: inline - distribution_matching: - _target_: verl_omni.workers.config.diffusion.DiffusionDistributionMatchingConfig - recipe: dmd2 - profile: null - fake_update_ratio: null - fake_warmup_cycles: 0 - rollout_strategy: null - data_mode: null - export_role: student_ema +dmd: + _target_: verl_omni.workers.config.DiffusionDMDConfig + fake_update_ratio: 2 + student_micro_batch_size_per_gpu: 1 + fake_score_micro_batch_size_per_gpu: 1 + fake_score_optim: + _target_: verl.workers.config.FSDPOptimizerConfig + lr: 2.0e-05 + weight_decay: 0.001 + betas: + - 0.9 + - 0.999 + clip_grad: 1.0 + lr_scheduler_type: constant + lr_warmup_steps: -1 + lr_warmup_steps_ratio: 0.0 + total_training_steps: -1 + teacher_guidance_scale: 4.0 + cfg_norm: layer_norm + negative_prompt: ' ' + normalization_epsilon: 1.0e-06 + rollout_timestep_shift: 3.0 + score_discrete_steps: 1000 + score_sigma_min: 0.02 + score_sigma_max: 0.98 + score_timestep_shift: 3.0 + ema_decay: 0.999 + ema_start_step: 0 + export_role: student algorithm: _target_: verl_omni.trainer.config.DiffusionAlgoConfig trainer_type: policy_gradient diff --git a/verl_omni/trainer/config/algorithm.py b/verl_omni/trainer/config/algorithm.py index e213bd410..6c1a292fa 100644 --- a/verl_omni/trainer/config/algorithm.py +++ b/verl_omni/trainer/config/algorithm.py @@ -43,7 +43,7 @@ class DiffusionAlgoConfig(BaseConfig): rollout_correction: RolloutCorrectionConfig = field(default_factory=RolloutCorrectionConfig) def __post_init__(self): - valid_trainer_types = {"policy_gradient", "direct_preference", "distillation"} + valid_trainer_types = {"policy_gradient", "direct_preference", "distribution_matching"} if self.trainer_type not in valid_trainer_types: raise ValueError(f"Invalid trainer_type: {self.trainer_type}. Must be one of {sorted(valid_trainer_types)}") valid_adv_modes = {"continuous", "positive_only", "negative_only", "one_only", "binary"} diff --git a/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml b/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml index 367ede769..c3a1f88b1 100644 --- a/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml +++ b/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml @@ -48,30 +48,3 @@ teacher_key: data_source # Teacher scoring schedule: inline scores within the training step; one_step_off overlaps the scoring of # batch k with the actor update on batch k-1 (v1 separate_async with standalone teachers only). scheduler: inline - -# DMD-family settings used only when algorithm.trainer_type=distillation. Keep enabled=false; that flag belongs to OPD. -distribution_matching: - - # Target class for this configuration - _target_: verl_omni.workers.config.diffusion.DiffusionDistributionMatchingConfig - - # Registered recipe name - recipe: dmd2 - - # Optional recipe profile; null selects the recipe default - profile: null - - # Optional fake-score phase count; null selects the recipe default - fake_update_ratio: null - - # Number of fake/discriminator-only cycles before student updates begin - fake_warmup_cycles: 0 - - # Optional registered rollout override; null selects the recipe default - rollout_strategy: null - - # Optional data-mode override; null selects the recipe default - data_mode: null - - # Semantic role exported to inference replicas - export_role: student_ema diff --git a/verl_omni/trainer/config/diffusion/dmd/diffusion_dmd.yaml b/verl_omni/trainer/config/diffusion/dmd/diffusion_dmd.yaml new file mode 100644 index 000000000..c0ca478ef --- /dev/null +++ b/verl_omni/trainer/config/diffusion/dmd/diffusion_dmd.yaml @@ -0,0 +1,59 @@ +# DMD2 distribution-only configuration; does not activate OPD. +_target_: verl_omni.workers.config.DiffusionDMDConfig + +# Fake-score update attempts per student attempt. +fake_update_ratio: 2 + +# Student physical microbatch size per GPU. +student_micro_batch_size_per_gpu: 1 + +# Fake-score physical microbatch size per GPU. +fake_score_micro_batch_size_per_gpu: 1 + +# Independent fake-score optimizer, reusing the FSDP optimizer configuration. +fake_score_optim: + _target_: verl.workers.config.FSDPOptimizerConfig + lr: 2.0e-5 + weight_decay: 0.001 + betas: [0.9, 0.999] + clip_grad: 1.0 + lr_scheduler_type: constant + lr_warmup_steps: -1 + lr_warmup_steps_ratio: 0.0 + total_training_steps: -1 + +# Teacher CFG scale; student and fake score are conditional-only. +teacher_guidance_scale: 4.0 + +# Packed-denoiser-space teacher CFG normalization. +cfg_norm: layer_norm + +# Explicit negative teacher condition. +negative_prompt: ' ' + +# Per-sample score normalization floor. +normalization_epsilon: 1.0e-6 + +# Fixed inference sigma shift. +rollout_timestep_shift: 3.0 + +# Discrete score grid size; zero selects continuous uniform sampling. +score_discrete_steps: 1000 + +# Lower score sigma bound. +score_sigma_min: 0.02 + +# Upper score sigma bound. +score_sigma_max: 0.98 + +# Discrete score timestep shift. +score_timestep_shift: 3.0 + +# EMA after successful student updates. +ema_decay: 0.999 + +# Successful student steps before EMA starts. +ema_start_step: 0 + +# Default inference artifact. +export_role: student diff --git a/verl_omni/trainer/config/diffusion_trainer.yaml b/verl_omni/trainer/config/diffusion_trainer.yaml index a2ba44e41..74915595f 100644 --- a/verl_omni/trainer/config/diffusion_trainer.yaml +++ b/verl_omni/trainer/config/diffusion_trainer.yaml @@ -32,6 +32,9 @@ defaults: # Distillation config for diffusion on-policy distillation. - diffusion/distillation@distillation: diffusion_distillation + # DMD2 distribution matching, separate from OPD. + - diffusion/dmd@dmd: diffusion_dmd + # load the reference default config, then apply the fields in the current yaml # self config override anything above - _self_ diff --git a/verl_omni/trainer/diffusion/diffusion_algos.py b/verl_omni/trainer/diffusion/diffusion_algos.py index 298e65ae7..dc15d05da 100644 --- a/verl_omni/trainer/diffusion/diffusion_algos.py +++ b/verl_omni/trainer/diffusion/diffusion_algos.py @@ -1051,6 +1051,72 @@ def __call__( return DiffusionLossResult(loss=kl_loss, metrics=metrics) +@register_diffusion_loss("dmd2") +class DMDLoss(DiffusionLossFn): + """DMD2 student surrogate and batch dispatch for fake-score denoising. + + Inputs share an explicit normalized latent layout ``(B, ...)``. The engine + supplies detached score predictions and detached fake-stage model inputs; + this loss also enforces the target/score stop-gradient boundary. + """ + + @classmethod + def compute_loss( + cls, + *, + generated_x0: torch.Tensor, + fake_x0: torch.Tensor, + teacher_x0: torch.Tensor, + normalization_epsilon: float = 1e-6, + gradient_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Compute the fp32 distribution-matching surrogate, not paired regression.""" + from verl_omni.trainer.diffusion.distillation.utils import dmd_gradient, dmd_surrogate_loss + + gradient, normalizer, nonfinite = dmd_gradient(fake_x0, teacher_x0, generated_x0, normalization_epsilon) + loss, active = dmd_surrogate_loss(generated_x0, gradient, gradient_mask) + return loss, { + "dmd/loss": loss.detach(), + "dmd/normalizer": normalizer.mean(), + "dmd/gradient_norm": gradient.norm(), + "dmd/nonfinite": nonfinite, + "dmd/active_elements": active, + } + + def validate_inputs(self, *, loss_name: str, model_output: dict[str, Any], data: TensorDict) -> None: + """Validate the stage-specific tensors without requiring PPO inputs.""" + stage = tu.get_non_tensor_data(data, "dmd_stage", default="student") + if stage == "student": + required = ("generated_x0", "fake_x0", "teacher_x0") + elif stage == "fake_score": + required = ("generated_x0", "noise_pred", "noise") + else: + raise ValueError(f"Invalid dmd_stage {stage!r}; expected 'student' or 'fake_score'.") + missing = [key for key in required if key not in model_output] + if missing: + raise KeyError(f"Diffusion loss `{loss_name}` is missing model_output keys: {missing}") + + def __call__(self, *, config: DiffusionActorConfig, model_output: dict[str, Any], data: TensorDict): + """Adapt the selected engine computation to the existing loss dispatcher.""" + self.validate_inputs(loss_name="dmd2", model_output=model_output, data=data) + stage = tu.get_non_tensor_data(data, "dmd_stage", default="student") + if stage == "student": + loss, metrics = self.compute_loss( + generated_x0=model_output["generated_x0"], + fake_x0=model_output["fake_x0"], + teacher_x0=model_output["teacher_x0"], + normalization_epsilon=tu.get_non_tensor_data(data, "dmd_normalization_epsilon", default=1e-6), + ) + else: + from verl_omni.trainer.diffusion.distillation.utils import fake_score_loss + + loss, active = fake_score_loss( + model_output["noise_pred"], model_output["noise"], model_output["generated_x0"] + ) + metrics = {"fake_score/loss": loss.detach(), "fake_score/active_elements": active} + return DiffusionLossResult(loss=loss, metrics=metrics, add_loss_metric=True) + + @register_diffusion_loss("distill_kl") class DistillKLLoss(DiffusionLossFn): """KL divergence between student and teacher reverse-SDE means (online policy distillation).""" diff --git a/verl_omni/trainer/diffusion/diffusion_trainer_utils.py b/verl_omni/trainer/diffusion/diffusion_trainer_utils.py index b8fc13de9..c44c9de57 100644 --- a/verl_omni/trainer/diffusion/diffusion_trainer_utils.py +++ b/verl_omni/trainer/diffusion/diffusion_trainer_utils.py @@ -71,13 +71,6 @@ def validate_distillation_config(config) -> None: actor = config.actor_rollout_ref.actor distill_active = actor.diffusion_loss.get("loss_mode", "flow_grpo") == "distill_kl" or actor.use_distill_loss enabled = is_distillation_enabled(config.get("distillation")) - if config.algorithm.trainer_type == "distillation": - if enabled or distill_active: - raise ValueError( - "DMD-family training selected by algorithm.trainer_type=distillation must keep the OPD " - "distillation.enabled flag and actor distillation losses disabled." - ) - return if enabled and not distill_active: raise ValueError( "distillation.enabled=true but no distillation loss is active; set " diff --git a/verl_omni/trainer/diffusion/distillation/__init__.py b/verl_omni/trainer/diffusion/distillation/__init__.py index f37526bab..c7e456f82 100644 --- a/verl_omni/trainer/diffusion/distillation/__init__.py +++ b/verl_omni/trainer/diffusion/distillation/__init__.py @@ -11,101 +11,4 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Distribution-matching distillation runtime (DMD, DMD2, CausVid, Self-Forcing). - -PR 1 provides the architecture-neutral trainer controller, immutable execution -contracts, recipe/objective/rollout registries, and pure DMD tensor utilities. It -defines no model pipeline, Ray worker, FSDP model, or GPU runtime. -""" - -from verl_omni.trainer.diffusion.distillation import contracts, controller, ray_trainer, recipes, utils -from verl_omni.trainer.diffusion.distillation.contracts import ( - CanonicalPrediction, - ConditionBundle, - DistillationCheckpointState, - DistillationPlan, - ExportSpec, - FrozenDict, - LatentBundle, - PhaseRequest, - PhaseResult, - RoleBinding, - RoleCheckpointManifest, - RoleGroupSpec, - RoleLayoutSpec, - ScoreBatch, - ScoreTransportSpec, - StudentRollout, - TeacherScoreProvider, - TrainerCounters, - UpdateCycle, - UpdatePhaseSpec, - UpdateSchedule, - describe_role_groups, - resolve_export_role, - validate_distillation_plan, - validate_export_role, - validate_role_layout, -) -from verl_omni.trainer.diffusion.distillation.controller import ( - BatchProvider, - DistillationPhaseExecutor, - DistillationTrainerController, - DistillationTrainerHooks, - FakeBatchProvider, - FakeDistillationHooks, - FakePhaseExecutor, -) -from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationRayTrainer -from verl_omni.trainer.diffusion.distillation.recipes import build_plan, build_plan_from_config, recipe_registry - -__all__ = [ - # submodules - "contracts", - "utils", - "recipes", - "controller", - "ray_trainer", - # contracts - "FrozenDict", - "LatentBundle", - "StudentRollout", - "ScoreBatch", - "ConditionBundle", - "CanonicalPrediction", - "RoleGroupSpec", - "RoleBinding", - "RoleLayoutSpec", - "ScoreTransportSpec", - "ExportSpec", - "UpdatePhaseSpec", - "UpdateSchedule", - "UpdateCycle", - "PhaseRequest", - "PhaseResult", - "TrainerCounters", - "DistillationPlan", - # role layout validation / export / provider / checkpoint - "validate_role_layout", - "validate_export_role", - "validate_distillation_plan", - "describe_role_groups", - "resolve_export_role", - "TeacherScoreProvider", - "RoleCheckpointManifest", - "DistillationCheckpointState", - # controller / executor - "DistillationTrainerController", - "BatchProvider", - "DistillationTrainerHooks", - "DistillationPhaseExecutor", - "FakePhaseExecutor", - "FakeBatchProvider", - "FakeDistillationHooks", - # driver - "DistillationRayTrainer", - # recipes - "build_plan", - "build_plan_from_config", - "recipe_registry", -] +"""Numerical diffusion distillation utilities; runtime uses the shared trainer/engine.""" diff --git a/verl_omni/trainer/diffusion/distillation/contracts.py b/verl_omni/trainer/diffusion/distillation/contracts.py deleted file mode 100644 index 4f860171c..000000000 --- a/verl_omni/trainer/diffusion/distillation/contracts.py +++ /dev/null @@ -1,545 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Immutable execution contracts for distribution-matching distillation. - -These architecture-neutral types make up a validated -:class:`~verl_omni.trainer.diffusion.distillation.recipes.DistillationPlan`. -They carry no Ray, model-pipeline, or FSDP dependency. The generic tensor -utilities operate on unpacked tensors; :class:`LatentBundle` is only a transport -container used across role boundaries. -""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Any, Literal, Optional, Protocol, runtime_checkable - -from torch import Tensor - -__all__ = [ - "FrozenDict", - "LatentBundle", - "StudentRollout", - "ScoreBatch", - "ConditionBundle", - "CanonicalPrediction", - "RoleGroupSpec", - "RoleBinding", - "ScoreTransportSpec", - "ExportSpec", - "RoleLayoutSpec", - "DataRequirements", - "ObjectiveSpec", - "RolloutSpec", - "InitializationSpec", - "UpdatePhaseSpec", - "PhaseRequest", - "PhaseResult", - "UpdateCycle", - "UpdateSchedule", - "TrainerCounters", - "DistillationPlan", - "validate_role_layout", - "validate_export_role", - "validate_distillation_plan", - "describe_role_groups", - "resolve_export_role", - "EXPORTABLE_ROLES", - "TeacherScoreProvider", - "RoleCheckpointManifest", - "DistillationCheckpointState", -] - - -class FrozenDict(Mapping[str, Any]): - """Small recursively immutable, pickle-friendly mapping for plan specs.""" - - __slots__ = ("data",) - - def __init__(self, values: Mapping[str, Any] | None = None) -> None: - values = values or {} - self.data = MappingProxyType({key: freeze_value(value) for key, value in values.items()}) - - def __getitem__(self, key: str) -> Any: - return self.data[key] - - def __iter__(self) -> Iterator[str]: - return iter(self.data) - - def __len__(self) -> int: - return len(self.data) - - def __repr__(self) -> str: - return f"FrozenDict({self.data!r})" - - def __hash__(self) -> int: - return hash(tuple(sorted(self.data.items()))) - - def __reduce__(self): - return FrozenDict, (dict(self.data),) - - -def freeze_value(value: Any) -> Any: - """Recursively freeze plan mappings and collections.""" - if isinstance(value, FrozenDict): - return value - if isinstance(value, Mapping): - return FrozenDict(value) - if isinstance(value, list | tuple): - return tuple(freeze_value(item) for item in value) - if isinstance(value, set | frozenset): - return frozenset(freeze_value(item) for item in value) - return value - - -@dataclass -class LatentBundle: - """Transport container for architecture-native per-modality tensors.""" - - tensors: dict[str, Tensor] - - def __post_init__(self) -> None: - if not self.tensors: - raise ValueError("LatentBundle must contain at least one modality tensor.") - for key, value in self.tensors.items(): - if not isinstance(value, Tensor): - raise TypeError(f"LatentBundle[{key!r}] must be a torch.Tensor, got {type(value)}.") - - def __len__(self) -> int: - return len(self.tensors) - - def get(self, modality: str) -> Tensor: - return self.tensors[modality] - - @property - def single(self) -> Tensor: - """Return the single tensor of a one-modality bundle.""" - if len(self.tensors) != 1: - raise ValueError(f"single only valid for a one-modality bundle, got {sorted(self.tensors)}") - return next(iter(self.tensors.values())) - - def map(self, fn): - """Apply ``fn`` to every tensor and return a new bundle.""" - return LatentBundle({key: fn(value) for key, value in self.tensors.items()}) - - -@dataclass -class StudentRollout: - """Output of a student rollout pass, consumed by score-model phases.""" - - generated_x0: LatentBundle - initial_noise: LatentBundle - selected_step_indices: Tensor - denoised_sigma_from: Optional[Tensor] = None - denoised_sigma_to: Optional[Tensor] = None - gradient_mask: Optional[LatentBundle] = None - committed_context_length: Optional[Tensor] = None - - -@dataclass -class ScoreBatch: - """Generated samples and conditioning scored by real/fake score models.""" - - generated_x0: LatentBundle - generated_x0_detached: LatentBundle - noisy_latents: LatentBundle - noise: LatentBundle - sigma: Tensor - condition: ConditionBundle - negative_condition: Optional[ConditionBundle] = None - - -@dataclass -class CanonicalPrediction: - """Teacher or fake-score output converted to canonical fp32 ``x0``.""" - - x0: LatentBundle - raw: Optional[LatentBundle] = None - - -@dataclass -class ConditionBundle: - """Conditioning tensors, masks, and metadata shared by all roles.""" - - tensors: dict[str, Tensor] - masks: dict[str, Tensor] = field(default_factory=dict) - metadata: dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class RoleGroupSpec: - """A physical model group owning one wrapped model.""" - - name: str - model_ref: str = "" - storage: Literal["independent_module", "shared_base_adapters"] = "shared_base_adapters" - placement: Literal["colocated", "standalone"] = "colocated" - - def __post_init__(self) -> None: - if not self.name: - raise ValueError("RoleGroupSpec.name must not be empty.") - valid_storage = {"independent_module", "shared_base_adapters"} - if self.storage not in valid_storage: - raise ValueError(f"Invalid role-group storage {self.storage!r}; expected one of {sorted(valid_storage)}.") - valid_placement = {"colocated", "standalone"} - if self.placement not in valid_placement: - raise ValueError( - f"Invalid role-group placement {self.placement!r}; expected one of {sorted(valid_placement)}." - ) - - -@dataclass(frozen=True) -class RoleBinding: - """A logical algorithm role bound to a group and optional named adapter.""" - - role: str - group: str - adapter: Optional[str] = None - trainable: bool = False - optimizer_key: Optional[str] = None - - def __post_init__(self) -> None: - if not self.role or not self.group: - raise ValueError("RoleBinding.role and RoleBinding.group must not be empty.") - - -@dataclass(frozen=True) -class ScoreTransportSpec: - """How teacher/fake scores are transported.""" - - provider: Literal["colocated", "ray"] = "colocated" - tensor_backend: Literal["local", "ray_nixl", "mooncake"] = "local" - - def __post_init__(self) -> None: - valid_providers = {"colocated", "ray"} - if self.provider not in valid_providers: - raise ValueError(f"Invalid score provider {self.provider!r}; expected one of {sorted(valid_providers)}.") - valid_backends = {"local", "ray_nixl", "mooncake"} - if self.tensor_backend not in valid_backends: - raise ValueError( - f"Invalid score tensor backend {self.tensor_backend!r}; expected one of {sorted(valid_backends)}." - ) - if self.provider == "colocated" and self.tensor_backend != "local": - raise ValueError("A colocated score provider requires tensor_backend='local'.") - if self.provider == "ray" and self.tensor_backend == "local": - raise ValueError("A Ray score provider requires tensor_backend='ray_nixl' or 'mooncake'.") - - -@dataclass(frozen=True) -class ExportSpec: - """Which semantic role is exported and through which backend.""" - - role: Literal["student", "student_ema"] = "student_ema" - checkpoint_engine_backend: str = "naive" - - def __post_init__(self) -> None: - if self.role not in EXPORTABLE_ROLES: - raise ValueError(f"Export role must be one of {EXPORTABLE_ROLES}, got {self.role!r}.") - if not self.checkpoint_engine_backend: - raise ValueError("checkpoint_engine_backend must not be empty.") - - -@dataclass(frozen=True) -class RoleLayoutSpec: - """Validated physical groups, logical bindings, and score transport.""" - - groups: tuple[RoleGroupSpec, ...] - bindings: tuple[RoleBinding, ...] - score_transport: ScoreTransportSpec = ScoreTransportSpec() - - -DataRequirements = Mapping[str, Any] -ObjectiveSpec = Mapping[str, Any] -RolloutSpec = Mapping[str, Any] -InitializationSpec = Mapping[str, Any] - - -@dataclass(frozen=True) -class UpdatePhaseSpec: - """A static phase specification expanded by :class:`UpdateSchedule`.""" - - kind: Literal["student", "fake_score"] - repeats: int = 1 - batch_policy: Literal["fresh", "reuse_student"] = "fresh" - trainable_roles: tuple[str, ...] = () - update_ema: bool = False - - def __post_init__(self) -> None: - if self.kind not in {"student", "fake_score"}: - raise ValueError(f"Invalid phase kind {self.kind!r}; expected 'student' or 'fake_score'.") - if isinstance(self.repeats, bool) or not isinstance(self.repeats, int) or self.repeats <= 0: - raise ValueError(f"Phase repeats must be an integer greater than zero, got {self.repeats}.") - if self.batch_policy not in {"fresh", "reuse_student"}: - raise ValueError(f"Invalid batch_policy {self.batch_policy!r}; expected 'fresh' or 'reuse_student'.") - if not self.trainable_roles: - raise ValueError(f"{self.kind!r} phase must declare at least one trainable role.") - if len(set(self.trainable_roles)) != len(self.trainable_roles): - raise ValueError(f"{self.kind!r} phase contains duplicate trainable roles.") - if self.kind == "student" and self.trainable_roles != ("student",): - raise ValueError("A student phase must train exactly the 'student' role.") - if self.kind == "fake_score" and "student" in self.trainable_roles: - raise ValueError("A fake_score phase must not train the student role.") - if self.update_ema and self.kind != "student": - raise ValueError("EMA updates are only valid on student phases.") - - -@dataclass(frozen=True) -class PhaseRequest: - """A concrete phase request emitted by :meth:`UpdateSchedule.next_cycle`.""" - - kind: Literal["student", "fake_score"] - global_step: int - repeat_index: int - batch_policy: Literal["fresh", "reuse_student"] - trainable_roles: tuple[str, ...] - update_ema: bool = False - - -@dataclass -class PhaseResult: - """The phase-specific state returned from an executor to the driver.""" - - metrics: dict[str, float] = field(default_factory=dict) - optimizer_steps: dict[str, int] = field(default_factory=dict) - - -@dataclass(frozen=True) -class UpdateCycle: - """A sequence of phase requests produced for one cycle.""" - - requests: tuple[PhaseRequest, ...] - requires_student_update: bool - is_warmup: bool = False - - -@dataclass -class TrainerCounters: - """Driver counters; ``global_step`` counts completed student updates only.""" - - global_step: int = 0 - optimizer_steps: dict[str, int] = field(default_factory=dict) - completed_cycles: int = 0 - - def increment_global(self) -> None: - self.global_step += 1 - - def record_step(self, role: str) -> None: - self.optimizer_steps[role] = self.optimizer_steps.get(role, 0) + 1 - - -@dataclass(frozen=True) -class UpdateSchedule: - """Normal phases plus an optional finite fake/discriminator warmup.""" - - phases: tuple[UpdatePhaseSpec, ...] - warmup_phases: tuple[UpdatePhaseSpec, ...] = () - warmup_cycles: int = 0 - - def __post_init__(self) -> None: - if not self.phases: - raise ValueError("UpdateSchedule must contain normal-cycle phases.") - student_indices = [index for index, phase in enumerate(self.phases) if phase.kind == "student"] - if student_indices != [0] or self.phases[0].repeats != 1: - raise ValueError( - "A normal cycle must contain exactly one student phase with repeats=1 and it must be first." - ) - if isinstance(self.warmup_cycles, bool) or not isinstance(self.warmup_cycles, int) or self.warmup_cycles < 0: - raise ValueError(f"warmup_cycles must be a non-negative integer, got {self.warmup_cycles}.") - if self.warmup_cycles > 0 and not self.warmup_phases: - raise ValueError("warmup_cycles > 0 requires at least one warmup phase.") - if self.warmup_cycles == 0 and self.warmup_phases: - raise ValueError("warmup_phases require warmup_cycles > 0.") - if any(phase.kind == "student" for phase in self.warmup_phases): - raise ValueError("Warmup phases must not contain a student phase.") - - def next_cycle(self, counters: TrainerCounters) -> UpdateCycle: - """Expand either the next warmup cycle or the normal static phases.""" - is_warmup = counters.completed_cycles < self.warmup_cycles - phases = self.warmup_phases if is_warmup else self.phases - requests = tuple( - PhaseRequest( - kind=phase.kind, - global_step=counters.global_step, - repeat_index=repeat_index, - batch_policy=phase.batch_policy, - trainable_roles=phase.trainable_roles, - update_ema=phase.update_ema, - ) - for phase in phases - for repeat_index in range(phase.repeats) - ) - return UpdateCycle( - requests=requests, - requires_student_update=not is_warmup, - is_warmup=is_warmup, - ) - - -@dataclass(frozen=True) -class DistillationPlan: - """An immutable, validated plan describing one named recipe.""" - - name: str - version: int - role_layout: RoleLayoutSpec - data_requirements: DataRequirements - objective: ObjectiveSpec - rollout: RolloutSpec - initialization: InitializationSpec - update_schedule: UpdateSchedule - export: ExportSpec - required_capabilities: frozenset[str] - - def __post_init__(self) -> None: - object.__setattr__(self, "data_requirements", FrozenDict(self.data_requirements)) - object.__setattr__(self, "objective", FrozenDict(self.objective)) - object.__setattr__(self, "rollout", FrozenDict(self.rollout)) - object.__setattr__(self, "initialization", FrozenDict(self.initialization)) - object.__setattr__(self, "required_capabilities", frozenset(self.required_capabilities)) - validate_distillation_plan(self) - - -EXPORTABLE_ROLES = ("student", "student_ema") - - -def validate_role_layout(layout: RoleLayoutSpec) -> None: - """Fail closed on invalid group/binding ownership before model allocation.""" - group_names = [group.name for group in layout.groups] - if not group_names: - raise ValueError("RoleLayoutSpec must contain at least one role group.") - duplicate_groups = {name for name in group_names if group_names.count(name) > 1} - if duplicate_groups: - raise ValueError(f"Duplicate role-group names: {sorted(duplicate_groups)}.") - if not layout.bindings: - raise ValueError("RoleLayoutSpec must contain at least one role binding.") - - binding_roles: set[str] = set() - optimizer_keys: set[str] = set() - groups = {group.name: group for group in layout.groups} - for binding in layout.bindings: - if binding.group not in groups: - raise ValueError(f"RoleBinding {binding.role!r} references unknown group {binding.group!r}.") - if binding.role in binding_roles: - raise ValueError(f"Duplicate role binding for {binding.role!r}.") - binding_roles.add(binding.role) - if binding.trainable and not binding.optimizer_key: - raise ValueError(f"Trainable role {binding.role!r} must set an optimizer_key.") - if not binding.trainable and binding.optimizer_key is not None: - raise ValueError(f"Frozen role {binding.role!r} must not set an optimizer_key.") - if binding.trainable and binding.optimizer_key in optimizer_keys: - raise ValueError(f"Duplicate optimizer_key {binding.optimizer_key!r}.") - if binding.optimizer_key is not None: - optimizer_keys.add(binding.optimizer_key) - if groups[binding.group].storage == "shared_base_adapters" and binding.trainable and not binding.adapter: - raise ValueError(f"Trainable role {binding.role!r} in a shared-base group must name an adapter.") - - for group in layout.groups: - if group.storage != "shared_base_adapters": - continue - adapters = [binding.adapter for binding in layout.bindings if binding.group == group.name and binding.adapter] - duplicates = {adapter for adapter in adapters if adapters.count(adapter) > 1} - if duplicates: - raise ValueError(f"Shared-base group {group.name!r} has duplicate adapter names {sorted(duplicates)}.") - - -def validate_export_role(export: ExportSpec, binding_roles: set[str]) -> None: - """Ensure the export role is semantic, exportable, and bound.""" - if export.role not in EXPORTABLE_ROLES: - raise ValueError(f"Export role must be one of {EXPORTABLE_ROLES}, got {export.role!r}.") - if export.role not in binding_roles: - raise ValueError(f"Export role {export.role!r} is not bound. Bound roles: {sorted(binding_roles)}.") - - -def validate_distillation_plan(plan: DistillationPlan) -> None: - """Validate architecture-neutral plan invariants before execution.""" - if not plan.name: - raise ValueError("DistillationPlan.name must not be empty.") - if plan.version <= 0: - raise ValueError(f"DistillationPlan.version must be positive, got {plan.version}.") - validate_role_layout(plan.role_layout) - missing_model_refs = [group.name for group in plan.role_layout.groups if not group.model_ref] - if missing_model_refs: - raise ValueError(f"Every role group must define model_ref; missing for {sorted(missing_model_refs)}.") - - binding_roles = {binding.role for binding in plan.role_layout.bindings} - required_roles = {"student", "student_ema", "teacher_score", "fake_score"} - missing_roles = required_roles - binding_roles - if missing_roles: - raise ValueError(f"DistillationPlan is missing required role bindings: {sorted(missing_roles)}.") - validate_export_role(plan.export, binding_roles) - - trainable_roles = {binding.role for binding in plan.role_layout.bindings if binding.trainable} - all_phases = plan.update_schedule.phases + plan.update_schedule.warmup_phases - for phase in all_phases: - unknown_roles = set(phase.trainable_roles) - trainable_roles - if unknown_roles: - raise ValueError( - f"Phase {phase.kind!r} references roles that are not trainable bindings: {sorted(unknown_roles)}." - ) - if phase.update_ema and "student_ema" not in binding_roles: - raise ValueError("A phase requesting EMA requires a bound student_ema role.") - - required_spec_keys = ( - ("data_requirements", plan.data_requirements, "mode"), - ("objective", plan.objective, "name"), - ("rollout", plan.rollout, "strategy"), - ("initialization", plan.initialization, "stage"), - ) - for spec_name, spec, required_key in required_spec_keys: - if not spec.get(required_key): - raise ValueError(f"DistillationPlan.{spec_name} must define non-empty {required_key!r}.") - - -def describe_role_groups(layout: RoleLayoutSpec) -> dict[str, str]: - """Return concise role-group descriptions for logging.""" - return {group.name: f"({group.storage}, {group.placement})" for group in layout.groups} - - -def resolve_export_role(export: ExportSpec) -> str: - """Return the semantic export role after validation.""" - if export.role not in EXPORTABLE_ROLES: - raise ValueError(f"Export role must be one of {EXPORTABLE_ROLES}, got {export.role!r}.") - return export.role - - -@runtime_checkable -class TeacherScoreProvider(Protocol): - """Provides canonical ``x0`` teacher predictions for a score batch.""" - - def predict_x0(self, score_batch: ScoreBatch) -> CanonicalPrediction: - """Return the teacher's canonical fp32 ``x0`` for ``score_batch``.""" - ... - - -@dataclass -class RoleCheckpointManifest: - """Metadata describing one role's stored state.""" - - role: str - model_path: str = "" - model_revision: str = "" - config_hash: str = "" - optimizer_key: str = "" - - -@dataclass -class DistillationCheckpointState: - """Composite multi-role checkpoint state, restored atomically.""" - - global_step: int = 0 - completed_cycles: int = 0 - role_manifests: list[RoleCheckpointManifest] = field(default_factory=list) - rng: dict[str, Any] = field(default_factory=dict) diff --git a/verl_omni/trainer/diffusion/distillation/controller.py b/verl_omni/trainer/diffusion/distillation/controller.py deleted file mode 100644 index 864c615bc..000000000 --- a/verl_omni/trainer/diffusion/distillation/controller.py +++ /dev/null @@ -1,269 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Pure, deterministic distillation trainer controller. - -The controller receives an immutable :class:`DistillationPlan` and talks to a -single :class:`DistillationPhaseExecutor`. It never imports a model pipeline, -manipulates latents, selects a PEFT adapter, or computes a DMD loss. Keeping the -controller free of Ray types makes its state machine testable in a CPU process -(RFC §13.1). - -The cycle state machine follows RFC §14. Worker allocation, role binding, and -checkpoint restore are deliberately delegated to the PR 2 executor; this module -only controls validated phase execution: - -- Optional fake/discriminator warmup cycles: emit fake-only ``UpdateCycle`` - requests, advance fake/discriminator optimizer counters, never ``global_step``. -- For each ``global_step``: one student phase then ``K`` fake phases; increment - ``global_step`` after all required phases complete; checkpoint/export/validate - when due. - -Invariants enforced here (RFC §14): - -- a completed student phase reports exactly one student optimizer step; -- a skipped student phase reports no student optimizer step and cannot advance - ``global_step``; -- a partially completed cycle is fail-fast and never retried in-process; -- a zero-progress cycle raises rather than being silently skipped; -- ``after_completed_step`` runs only after ``global_step`` is incremented. -""" - -from __future__ import annotations - -from typing import Any, Optional, Protocol, runtime_checkable - -from verl_omni.trainer.diffusion.distillation.contracts import ( - DistillationPlan, - PhaseRequest, - PhaseResult, - TrainerCounters, - UpdateCycle, -) - -__all__ = [ - "DistillationTrainerController", - "BatchProvider", - "DistillationTrainerHooks", - "DistillationPhaseExecutor", - "FakePhaseExecutor", - "FakeBatchProvider", - "FakeDistillationHooks", -] - - -@runtime_checkable -class BatchProvider(Protocol): - """Supplies per-phase input batches to the executor.""" - - def next(self, phase: PhaseRequest) -> Any: ... - - -@runtime_checkable -class DistillationTrainerHooks(Protocol): - """Receives post-completed-step callbacks with counters, metrics, and executor. - - ``after_completed_step`` is the sanctioned place to schedule checkpoint, - validation, and export behavior without introducing model-specific or - Ray-specific control flow into the generic trainer. It runs only after - ``global_step`` is incremented, so it observes the new counter value. - """ - - def after_completed_step(self, counters: TrainerCounters, metrics: dict, executor: Any) -> None: ... - - -class DistillationTrainerController: - """Pure driver over a plan and an executor. No Ray, model, or FSDP types.""" - - def __init__( - self, - plan: DistillationPlan, - executor: DistillationPhaseExecutor, - batch_provider: BatchProvider, - hooks: Optional[DistillationTrainerHooks] = None, - ) -> None: - self.plan = plan - self.executor = executor - self.batch_provider = batch_provider - self.hooks = hooks - self.counters = TrainerCounters() - self.phase_metrics: dict[str, dict] = {} - self.failed = False - - def run(self, num_cycles: int) -> None: - """Drive ``num_cycles`` update cycles through the executor.""" - for _ in range(num_cycles): - self.run_cycle() - - def run_cycle(self) -> UpdateCycle: - """Run one cycle transactionally and become terminal after a failure.""" - if self.failed: - raise RuntimeError( - "This controller previously failed during a cycle and cannot be retried in-process; " - "restore the last completed-cycle checkpoint into a new driver." - ) - - before_counters = TrainerCounters( - global_step=self.counters.global_step, - optimizer_steps=dict(self.counters.optimizer_steps), - completed_cycles=self.counters.completed_cycles, - ) - before_metrics = dict(self.phase_metrics) - try: - cycle = self.plan.update_schedule.next_cycle(self.counters) - student_step_reported = self.drive_requests(cycle.requests) - - if cycle.requires_student_update: - if not student_step_reported: - raise ValueError( - "Cycle requires a student update but no student optimizer step was reported; " - "global_step must not advance." - ) - self.counters.increment_global() - - self.assert_progress(before_counters) - self.counters.completed_cycles += 1 - - if cycle.requires_student_update and self.hooks is not None: - self.hooks.after_completed_step(self.counters, self.metrics, self.executor) - return cycle - except Exception: - self.counters = before_counters - self.phase_metrics = before_metrics - self.failed = True - raise - - def drive_requests(self, requests: tuple[PhaseRequest, ...]) -> bool: - """Execute phases in order and validate each executor result fail closed.""" - student_step_reported = False - for request in requests: - batch = self.batch_provider.next(request) - result = self.executor.execute_phase(request, batch) - self.validate_result(result, request) - if request.kind == "student" and not result.optimizer_steps: - return False - self.accumulate_result(result, request) - if request.kind == "student": - student_step_reported = True - return student_step_reported - - @staticmethod - def validate_result(result: PhaseResult, request: PhaseRequest) -> None: - """Require exactly one step for every role declared by a completed phase.""" - if not isinstance(result, PhaseResult): - raise TypeError(f"execute_phase must return PhaseResult, got {type(result)}.") - if not result.optimizer_steps and request.kind == "student": - return - expected_roles = set(request.trainable_roles) - reported_roles = set(result.optimizer_steps) - if reported_roles != expected_roles: - raise ValueError( - f"Phase {request.kind!r} must report optimizer steps for exactly {sorted(expected_roles)}, " - f"got {sorted(reported_roles)}." - ) - invalid_steps = { - role: steps - for role, steps in result.optimizer_steps.items() - if isinstance(steps, bool) or not isinstance(steps, int) or steps != 1 - } - if invalid_steps: - raise ValueError(f"Each completed phase role must report exactly one optimizer step, got {invalid_steps}.") - - def accumulate_result(self, result: PhaseResult, request: PhaseRequest) -> None: - """Record validated role counters and the latest phase metrics.""" - for role, steps in result.optimizer_steps.items(): - self.counters.optimizer_steps[role] = self.counters.optimizer_steps.get(role, 0) + steps - self.phase_metrics[request.kind] = dict(result.metrics) - - def assert_progress(self, before: TrainerCounters) -> None: - """A cycle must advance global_step or at least one role optimizer counter.""" - if self.counters.global_step != before.global_step: - return - if self.counters.optimizer_steps != before.optimizer_steps: - return - raise ValueError("Zero-progress cycle: it advanced neither global_step nor any role optimizer counter.") - - @property - def metrics(self) -> dict[str, dict]: - """Metrics recorded for the most recent phase of each kind.""" - return self.phase_metrics - - def reset(self) -> None: - """Reset a healthy driver; failed drivers must be reconstructed from checkpoint.""" - if self.failed: - raise RuntimeError("A failed controller cannot be reset in-process; construct a new driver.") - self.counters = TrainerCounters() - self.phase_metrics = {} - - -@runtime_checkable -class DistillationPhaseExecutor(Protocol): - """Executes one update phase and returns a :class:`PhaseResult`.""" - - def execute_phase(self, request: PhaseRequest, batch: Any) -> PhaseResult: - """Run the phase requested by ``request`` on ``batch`` and return metrics/steps.""" - ... - - -class FakeBatchProvider: - """Minimal batch provider that yields synthetic batches on request.""" - - def __init__(self, num_batches: int = 1, batch_size: int = 1) -> None: - self.num_batches = num_batches - self.batch_size = batch_size - self.sent_batches = 0 - - def next(self, request: PhaseRequest) -> Any: - """Return a synthetic batch for the requested phase, advancing a counter.""" - if self.sent_batches >= self.num_batches: - raise StopIteration("No more batches.") - self.sent_batches += 1 - return {"phase_kind": request.kind, "global_step": request.global_step, "repeat": request.repeat_index} - - -class FakePhaseExecutor: - """Deterministic in-process fake executor used for CPU control-plane tests. - - It confirms the phase kind equals ``PhaseRequest.kind`` (fail-fast on a - misordered phase), emits a deterministic metric and optimizer-step record, and - can be configured to skip or fail a student phase for testing. - """ - - def __init__(self, skip_student: bool = False, fail_on: str | None = None) -> None: - self.skip_student = skip_student - self.fail_on = fail_on - self.executed: list[PhaseRequest] = [] - - def execute_phase(self, request: PhaseRequest, batch: Any) -> PhaseResult: - """Return synthetic phase results or a configured failure.""" - if batch.get("phase_kind") != request.kind: - raise ValueError(f"Batch phase {batch.get('phase_kind')!r} does not match request {request.kind!r}.") - self.executed.append(request) - if request.kind == self.fail_on: - raise RuntimeError(f"FakePhaseExecutor failed on phase {request.kind} (global_step={request.global_step}).") - if request.kind == "student" and self.skip_student: - return PhaseResult(metrics={"fake/student": float(request.global_step)}, optimizer_steps={}) - metrics = {f"fake/{request.kind}": float(request.global_step)} - optimizer_steps = {role: 1 for role in request.trainable_roles} - return PhaseResult(metrics=metrics, optimizer_steps=optimizer_steps) - - -class FakeDistillationHooks: - """Collects ``after_completed_step`` callbacks for deterministic validation.""" - - def __init__(self) -> None: - self.calls: list[dict[str, Any]] = [] - - def after_completed_step(self, counters, metrics, executor) -> None: - """Record the completed-step hook arguments.""" - self.calls.append({"global_step": counters.global_step, "metrics": dict(metrics), "executor": executor}) diff --git a/verl_omni/trainer/diffusion/distillation/ray_trainer.py b/verl_omni/trainer/diffusion/distillation/ray_trainer.py deleted file mode 100644 index 3f9dd9fd1..000000000 --- a/verl_omni/trainer/diffusion/distillation/ray_trainer.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Ray-entrypoint-compatible shell around the pure distillation control plane. - -PR 1 deliberately stops before model allocation. The shell accepts the same -constructor protocol and lifecycle calls as the existing diffusion trainers so -``algorithm.trainer_type=distillation`` reaches an explicit PR 2 boundary rather -than failing with a Python signature error. -""" - -from __future__ import annotations - -from typing import Any, Optional - -from verl_omni.trainer.diffusion.diffusion_trainer_utils import validate_distillation_config -from verl_omni.trainer.diffusion.distillation.contracts import DistillationPlan -from verl_omni.trainer.diffusion.distillation.controller import DistillationTrainerController -from verl_omni.trainer.diffusion.distillation.recipes import build_plan_from_config - -__all__ = ["DistillationRayTrainer"] - - -class DistillationRayTrainer: - """Production-compatible driver shell for a validated distillation plan.""" - - def __init__( - self, - config=None, - tokenizer=None, - role_worker_mapping=None, - resource_pool_manager=None, - ray_worker_group_cls=None, - processor=None, - train_dataset=None, - val_dataset=None, - collate_fn=None, - train_sampler=None, - device_name=None, - *, - plan: Optional[DistillationPlan] = None, - capabilities: Optional[frozenset[str]] = None, - executor: Optional[Any] = None, - batch_provider: Optional[Any] = None, - hooks: Optional[Any] = None, - ) -> None: - if isinstance(config, DistillationPlan) and plan is None: - plan = config - config = None - if config is not None: - validate_distillation_config(config) - if plan is not None and config is not None and capabilities is not None: - raise ValueError("Pass either an explicit plan or config+capabilities, not both.") - if plan is None and config is not None and capabilities is not None: - plan = build_plan_from_config(config, capabilities) - - self.config = config - self.tokenizer = tokenizer - self.processor = processor - self.role_worker_mapping = role_worker_mapping - self.resource_pool_manager = resource_pool_manager - self.ray_worker_group_cls = ray_worker_group_cls - self.train_dataset = train_dataset - self.val_dataset = val_dataset - self.collate_fn = collate_fn - self.train_sampler = train_sampler - self.device_name = device_name - self.plan = plan - self.capabilities = capabilities - self.executor = executor - self.batch_provider = batch_provider - self.hooks = hooks - self.controller_instance: Optional[DistillationTrainerController] = None - - def init_workers(self) -> None: - """Validate the PR 1 boundary before PR 2 supplies role-group workers.""" - if self.executor is None or self.batch_provider is None: - raise NotImplementedError( - "The multi-role distillation workers and architecture capability binding land in PR 2. " - "PR 1 accepts the production trainer interface but does not allocate model workers." - ) - if self.plan is None: - raise ValueError("A validated DistillationPlan is required when an executor is bound.") - - def build_controller(self) -> DistillationTrainerController: - """Construct the pure controller from a plan and bound collaborators.""" - self.init_workers() - assert self.plan is not None - self.controller_instance = DistillationTrainerController( - plan=self.plan, - executor=self.executor, - batch_provider=self.batch_provider, - hooks=self.hooks, - ) - return self.controller_instance - - @property - def controller(self) -> DistillationTrainerController: - """Return the lazily constructed distillation trainer controller.""" - if self.controller_instance is None: - return self.build_controller() - return self.controller_instance - - def fit(self, num_cycles: int = 0) -> None: - """Drive the injected CPU controller; production data plane arrives in PR 2.""" - self.controller.run(num_cycles) diff --git a/verl_omni/trainer/diffusion/distillation/recipes.py b/verl_omni/trainer/diffusion/distillation/recipes.py deleted file mode 100644 index 1b6c0c90b..000000000 --- a/verl_omni/trainer/diffusion/distillation/recipes.py +++ /dev/null @@ -1,533 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Named distillation recipes and their composed strategy registries.""" - -from __future__ import annotations - -import abc -from functools import partial - -from verl_omni.trainer.diffusion.distillation.contracts import ( - DistillationPlan, - ExportSpec, - RoleBinding, - RoleGroupSpec, - RoleLayoutSpec, - ScoreTransportSpec, - UpdatePhaseSpec, - UpdateSchedule, - validate_distillation_plan, -) - -__all__ = [ - "recipe_registry", - "DMDRecipe", - "DMD2Recipe", - "CausVidRecipe", - "SelfForcingRecipe", - "build_plan", - "build_plan_from_config", - "DistillationRecipeBase", - "DistillationRecipeRegistry", - "ObjectiveBase", - "ObjectiveRegistry", - "RolloutStrategyBase", - "RolloutStrategyRegistry", - "InitializationBase", - "InitializationRegistry", - "DistillationRegistry", - "objective_registry", - "rollout_registry", - "initialization_registry", - "DMDObjective", - "DMD2Objective", - "ODERegressionObjective", - "OneStepRollout", - "EulerRollout", - "ConsistencyRenoiseRollout", - "TeacherForcedCausalRollout", - "SelfForcedRollout", - "BackwardSimulatedRollout", - "BaseInitialization", - "ODERegressionInitialization", -] - - -class DistillationRecipeBase(abc.ABC): - """Base class for a named distillation recipe.""" - - @classmethod - @abc.abstractmethod - def build_plan(cls, config, capabilities) -> DistillationPlan: - """Return a validated immutable :class:`DistillationPlan`.""" - - -class DistillationRegistry: - """Minimal name-to-class registry with duplicate rejection.""" - - def __init__(self) -> None: - self.registered_classes: dict[str, type] = {} - - def register(self, name: str, subclass: type | None = None): - """Register a class directly or return its registration decorator.""" - if not name: - raise ValueError("Registry names must not be empty.") - if subclass is None: - return partial(self.register, name) - if name in self.registered_classes: - raise ValueError(f"Duplicate registration for {name!r}.") - self.registered_classes[name] = subclass - return subclass - - def get(self, name: str) -> type: - """Resolve a registered class by name.""" - try: - return self.registered_classes[name] - except KeyError: - raise KeyError( - f"No {self.kind} registered for {name!r}. Registered: {sorted(self.registered_classes)}" - ) from None - - @property - def names(self) -> tuple[str, ...]: - """Return the registered names in stable order.""" - return tuple(sorted(self.registered_classes)) - - @property - def kind(self) -> str: - """Identify this registry in resolution errors.""" - return self.__class__.__name__ - - -class DistillationRecipeRegistry(DistillationRegistry): - """Registry of named recipes that build a :class:`DistillationPlan`.""" - - @property - def kind(self) -> str: - return "distillation recipe" - - def build(self, name: str, config, capabilities) -> DistillationPlan: - """Build a plan from the recipe registered under ``name``.""" - return self.get(name).build_plan(config, capabilities) - - -class ObjectiveRegistry(DistillationRegistry): - """Registry of composed objective strategies.""" - - @property - def kind(self) -> str: - return "objective" - - -class RolloutStrategyRegistry(DistillationRegistry): - """Registry of student rollout strategies.""" - - @property - def kind(self) -> str: - return "rollout strategy" - - -class InitializationRegistry(DistillationRegistry): - """Registry of initialization strategies.""" - - @property - def kind(self) -> str: - return "initialization" - - -objective_registry = ObjectiveRegistry() -rollout_registry = RolloutStrategyRegistry() -initialization_registry = InitializationRegistry() - - -class ObjectiveBase: - """Name marker for a composed objective implemented by a phase executor.""" - - name: str = "" - - -@objective_registry.register("dmd") -class DMDObjective(ObjectiveBase): - """DMD detached normalized score-gradient objective.""" - - name = "dmd" - - -@objective_registry.register("dmd2") -class DMD2Objective(ObjectiveBase): - """DMD2 two-time-scale distribution-matching objective.""" - - name = "dmd2" - - -@objective_registry.register("ode_regression") -class ODERegressionObjective(ObjectiveBase): - """ODE regression against precomputed trajectory targets.""" - - name = "ode_regression" - - -class RolloutStrategyBase: - """Name marker for a student rollout implemented by a phase executor.""" - - name: str = "" - - -@rollout_registry.register("one_step") -class OneStepRollout(RolloutStrategyBase): - """Single-step student rollout.""" - - name = "one_step" - - -@rollout_registry.register("ode_euler") -class EulerRollout(RolloutStrategyBase): - """Deterministic Euler backward simulation.""" - - name = "ode_euler" - - -@rollout_registry.register("consistency_renoise") -class ConsistencyRenoiseRollout(RolloutStrategyBase): - """Consistency re-noising backward simulation.""" - - name = "consistency_renoise" - - -@rollout_registry.register("teacher_forced_causal") -class TeacherForcedCausalRollout(RolloutStrategyBase): - """Teacher-forced causal rollout used by CausVid.""" - - name = "teacher_forced_causal" - - -@rollout_registry.register("self_forced") -class SelfForcedRollout(RolloutStrategyBase): - """Self-forced autoregressive rollout.""" - - name = "self_forced" - - -@rollout_registry.register("backward_simulated") -class BackwardSimulatedRollout(RolloutStrategyBase): - """Inference-time backward-simulated multi-step student input.""" - - name = "backward_simulated" - - -class InitializationBase: - """Name marker for a recipe initialization strategy.""" - - name: str = "" - - -@initialization_registry.register("base") -class BaseInitialization(InitializationBase): - """Initialize trainable roles directly from their configured base.""" - - name = "base" - - -@initialization_registry.register("ode_regression") -class ODERegressionInitialization(InitializationBase): - """Initialize a causal student from a provenance-tracked ODE stage.""" - - name = "ode_regression" - - -recipe_registry = DistillationRecipeRegistry() - - -def shared_base_layout(model_ref: str, with_discriminator: bool = False) -> RoleLayoutSpec: - """Build the image-first shared-base named-adapter role layout.""" - group_name = "base" - bindings = [ - RoleBinding(role="student", group=group_name, adapter="student", trainable=True, optimizer_key="student"), - RoleBinding(role="teacher_score", group=group_name), - RoleBinding( - role="fake_score", group=group_name, adapter="fake_score", trainable=True, optimizer_key="fake_score" - ), - RoleBinding(role="student_ema", group=group_name, adapter="student_ema"), - ] - if with_discriminator: - bindings.append( - RoleBinding( - role="discriminator", - group=group_name, - adapter="discriminator", - trainable=True, - optimizer_key="discriminator", - ) - ) - return RoleLayoutSpec( - groups=(RoleGroupSpec(name=group_name, model_ref=model_ref),), - bindings=tuple(bindings), - score_transport=ScoreTransportSpec(), - ) - - -def causal_bidirectional_layout(causal_model_ref: str, bidirectional_model_ref: str) -> RoleLayoutSpec: - """Keep causal student/EMA separate from bidirectional score models.""" - return RoleLayoutSpec( - groups=( - RoleGroupSpec(name="causal_base", model_ref=causal_model_ref), - RoleGroupSpec(name="bidirectional_base", model_ref=bidirectional_model_ref), - ), - bindings=( - RoleBinding( - role="student", - group="causal_base", - adapter="student", - trainable=True, - optimizer_key="student", - ), - RoleBinding(role="student_ema", group="causal_base", adapter="student_ema"), - RoleBinding(role="teacher_score", group="bidirectional_base"), - RoleBinding( - role="fake_score", - group="bidirectional_base", - adapter="fake_score", - trainable=True, - optimizer_key="fake_score", - ), - ), - score_transport=ScoreTransportSpec(), - ) - - -def build_update_schedule( - fake_repeats: int, fake_warmup_cycles: int = 0, with_discriminator: bool = False -) -> UpdateSchedule: - """Build one student phase followed by ``fake_repeats`` fake phases.""" - fake_roles = ("fake_score", "discriminator") if with_discriminator else ("fake_score",) - fake_phase = UpdatePhaseSpec(kind="fake_score", repeats=fake_repeats, trainable_roles=fake_roles) - warmup_phases = (fake_phase,) if fake_warmup_cycles else () - return UpdateSchedule( - phases=( - UpdatePhaseSpec(kind="student", trainable_roles=("student",), update_ema=True), - fake_phase, - ), - warmup_phases=warmup_phases, - warmup_cycles=fake_warmup_cycles, - ) - - -def get_config_value(config, key: str, default=None): - """Read ``key`` from a mapping-like or attribute-style config.""" - if config is None: - return default - if hasattr(config, "get"): - return config.get(key, default) - return getattr(config, key, default) - - -def get_config_or_default(config, key: str, default): - """Use the default when a config field is absent or explicitly null.""" - value = get_config_value(config, key, default) - return default if value is None else value - - -def require_choice(field: str, value: str, valid_values: set[str]) -> str: - """Validate a named configuration choice.""" - if value not in valid_values: - raise ValueError(f"Invalid {field} {value!r}; expected one of {sorted(valid_values)}.") - return value - - -def common_plan_kwargs( - config, - *, - default_fake_repeats: int = 5, - with_discriminator: bool = False, -) -> dict: - """Build architecture-neutral scheduling and export settings.""" - fake_repeats = get_config_or_default(config, "fake_update_ratio", default_fake_repeats) - fake_warmup_cycles = get_config_or_default(config, "fake_warmup_cycles", 0) - export_role = require_choice( - "export_role", get_config_or_default(config, "export_role", "student_ema"), {"student", "student_ema"} - ) - return { - "version": 1, - "update_schedule": build_update_schedule(fake_repeats, fake_warmup_cycles, with_discriminator), - "export": ExportSpec(role=export_role, checkpoint_engine_backend="naive"), - } - - -@recipe_registry.register("dmd") -class DMDRecipe(DistillationRecipeBase): - """Original DMD with paired trajectory regression.""" - - @classmethod - def build_plan(cls, config, capabilities) -> DistillationPlan: - profile = require_choice("DMD profile", get_config_or_default(config, "profile", "paper"), {"paper"}) - data_mode = require_choice( - "DMD data_mode", get_config_or_default(config, "data_mode", "regression_pairs"), {"regression_pairs"} - ) - rollout = require_choice( - "DMD rollout_strategy", get_config_or_default(config, "rollout_strategy", "one_step"), {"one_step"} - ) - model_ref = get_config_value(config, "model_path", "") or "" - return DistillationPlan( - name="dmd", - role_layout=shared_base_layout(model_ref), - data_requirements={"mode": data_mode}, - objective={"name": "dmd", "profile": profile}, - rollout={"strategy": rollout}, - initialization={"stage": "base"}, - required_capabilities=frozenset({"distribution_matching"}), - **common_plan_kwargs(config, default_fake_repeats=1), - ) - - -@recipe_registry.register("dmd2") -class DMD2Recipe(DistillationRecipeBase): - """DMD2 distribution matching with optional diffusion-GAN profile.""" - - @classmethod - def build_plan(cls, config, capabilities) -> DistillationPlan: - profile = require_choice( - "DMD2 profile", - get_config_or_default(config, "profile", "distribution_only"), - {"distribution_only", "paper"}, - ) - adversarial = profile == "paper" - expected_data_mode = "prompt_and_real_latent" if adversarial else "prompts" - data_mode = require_choice( - "DMD2 data_mode", get_config_or_default(config, "data_mode", expected_data_mode), {expected_data_mode} - ) - rollout = require_choice( - "DMD2 rollout_strategy", - get_config_or_default(config, "rollout_strategy", "ode_euler"), - {"ode_euler", "consistency_renoise", "backward_simulated"}, - ) - model_ref = get_config_value(config, "model_path", "") or "" - return DistillationPlan( - name="dmd2", - role_layout=shared_base_layout(model_ref, with_discriminator=adversarial), - data_requirements={"mode": data_mode}, - objective={"name": "dmd2", "profile": profile, "adversarial": adversarial}, - rollout={"strategy": rollout}, - initialization={"stage": "base"}, - required_capabilities=frozenset({"distribution_matching"} | ({"adversarial"} if adversarial else set())), - **common_plan_kwargs(config, with_discriminator=adversarial), - ) - - -@recipe_registry.register("causvid") -class CausVidRecipe(DistillationRecipeBase): - """ODE-initialized causal student versus bidirectional score models.""" - - @classmethod - def build_plan(cls, config, capabilities) -> DistillationPlan: - require_choice( - "CausVid profile", get_config_or_default(config, "profile", "distribution_only"), {"distribution_only"} - ) - data_mode = require_choice( - "CausVid data_mode", - get_config_or_default(config, "data_mode", "prompt_and_real_latent"), - {"prompt_and_real_latent"}, - ) - rollout = require_choice( - "CausVid rollout_strategy", - get_config_or_default(config, "rollout_strategy", "teacher_forced_causal"), - {"teacher_forced_causal"}, - ) - model_ref = get_config_value(config, "model_path", "") or "" - causal_ref = get_config_value(config, "causal_model_path", model_ref) or model_ref - bidirectional_ref = get_config_value(config, "bidirectional_model_path", model_ref) or model_ref - return DistillationPlan( - name="causvid", - role_layout=causal_bidirectional_layout(causal_ref, bidirectional_ref), - data_requirements={"mode": data_mode}, - objective={"name": "dmd", "profile": "distribution_only"}, - rollout={"strategy": rollout}, - initialization={"stage": "ode_regression", "requires_provenance": True}, - required_capabilities=frozenset({"distribution_matching", "autoregressive"}), - **common_plan_kwargs(config), - ) - - -@recipe_registry.register("self_forcing") -class SelfForcingRecipe(DistillationRecipeBase): - """DMD objective with self-forced autoregressive rollout.""" - - @classmethod - def build_plan(cls, config, capabilities) -> DistillationPlan: - require_choice( - "Self-Forcing profile", get_config_or_default(config, "profile", "distribution_only"), {"distribution_only"} - ) - data_mode = require_choice( - "Self-Forcing data_mode", get_config_or_default(config, "data_mode", "prompts"), {"prompts"} - ) - rollout = require_choice( - "Self-Forcing rollout_strategy", - get_config_or_default(config, "rollout_strategy", "self_forced"), - {"self_forced"}, - ) - model_ref = get_config_value(config, "model_path", "") or "" - causal_ref = get_config_value(config, "causal_model_path", model_ref) or model_ref - bidirectional_ref = get_config_value(config, "bidirectional_model_path", model_ref) or model_ref - return DistillationPlan( - name="self_forcing", - role_layout=causal_bidirectional_layout(causal_ref, bidirectional_ref), - data_requirements={"mode": data_mode}, - objective={"name": "dmd", "profile": "distribution_only"}, - rollout={"strategy": rollout}, - initialization={"stage": "ode_regression", "requires_provenance": True}, - required_capabilities=frozenset({"distribution_matching", "autoregressive"}), - **common_plan_kwargs(config), - ) - - -def validate_registered_strategies(plan: DistillationPlan) -> None: - """Resolve every strategy name before execution.""" - objective_registry.get(plan.objective["name"]) - rollout_registry.get(plan.rollout["strategy"]) - initialization_registry.get(plan.initialization["stage"]) - - -def build_plan(name: str, config=None, capabilities=frozenset()) -> DistillationPlan: - """Build and fail-closed validate a named recipe plan.""" - plan = recipe_registry.build(name, config, capabilities) - validate_registered_strategies(plan) - validate_distillation_plan(plan) - missing = plan.required_capabilities - frozenset(capabilities) - if missing: - raise ValueError( - f"Recipe {name!r} requires capabilities {sorted(missing)} that the architecture adapter " - f"does not provide. Declared: {sorted(capabilities)}." - ) - return plan - - -def build_plan_from_config(config, capabilities) -> DistillationPlan: - """Build a plan from the composed trainer config and adapter capabilities.""" - distillation = get_config_value(config, "distillation") - distribution_matching = get_config_value(distillation, "distribution_matching") - if distribution_matching is None: - raise ValueError("config.distillation.distribution_matching is required for the distillation trainer.") - - actor_rollout_ref = get_config_value(config, "actor_rollout_ref") - model = get_config_value(actor_rollout_ref, "model") - model_path = get_config_value(model, "path", "") or "" - recipe_config = { - "model_path": model_path, - "fake_warmup_cycles": get_config_value(distribution_matching, "fake_warmup_cycles", 0), - "export_role": get_config_value(distribution_matching, "export_role", "student_ema"), - } - for optional_key in ("profile", "fake_update_ratio", "rollout_strategy", "data_mode"): - value = get_config_value(distribution_matching, optional_key) - if value is not None: - recipe_config[optional_key] = value - return build_plan(get_config_value(distribution_matching, "recipe", "dmd2"), recipe_config, capabilities) diff --git a/verl_omni/trainer/diffusion/distillation/utils.py b/verl_omni/trainer/diffusion/distillation/utils.py index c6443ea46..81b257fba 100644 --- a/verl_omni/trainer/diffusion/distillation/utils.py +++ b/verl_omni/trainer/diffusion/distillation/utils.py @@ -11,132 +11,64 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Pure tensor utilities used by the DMD-family distillation trainer. +"""Pure fp32 tensor utilities for Qwen-Image DMD2 distribution matching. -This module is independent of Ray and model libraries: it implements only the -detached-normalized distribution-matching gradient, the surrogate student loss, -the fake-score target, canonical x0 conversion, and the CFG forms that the -reference implementations use. The equations follow RFC §7 and §8. - -Why this is a separate module (and not part of ``recipes.py`` or -``contracts.py``): - -- ``contracts.py`` holds *data types* (immutable plan pieces, role layout, the - execution state machine). It carries no equations. -- ``recipes.py`` holds *declarations* (which objective, which rollout strategy, - which initialization, how roles map onto groups). It never computes a quantity. -- ``utils.py`` holds the only *executable equations* in the package. Every value - is a pure function of its tensors; nothing here reads config, weights, or the - prompt. Keeping these functions together means they can be unit-tested as - algebraic identities and finite-difference checks without building a plan or an - executor (see ``test_distillation_dmd_math_on_cpu.py``). - -Boundary conditions (all must match the reviewed reference implementations): - -- ``normalizer`` is formed over the **entire** ``x_g - x0_real`` tensor across all - non-batch (block, frame, channel, spatial) dimensions of one sample, ``keepdim`` - per sample. It is **not** restricted by ``gradient_mask``. -- Only the surrogate loss is masked by ``gradient_mask``. -- ``normalization_epsilon`` is applied as ``max(normalizer, normalization_epsilon)`` - before division, and a non-finite ``g / normalizer`` is replaced by - ``nan_to_num`` and counted. -- All score/objective arithmetic is fp32. +The score-difference normalizer spans every non-batch dimension per sample. +Only the surrogate reduction is masked. These functions own no model or runtime. """ from __future__ import annotations +import math from typing import Optional import torch from torch import Tensor __all__ = [ - "epsilon_to_x0", "velocity_to_x0", "dmd_gradient", "dmd_surrogate_loss", "fake_score_target", "fake_score_loss", - "ode_regression_loss", "ode_euler_step", - "consistency_renoise_step", "standard_cfg", - "legacy_cfg", "timestep_shift", ] -def epsilon_to_x0(noisy: Tensor, epsilon: Tensor, sigma: Tensor, a_fn, b_fn) -> Tensor: - """Convert an epsilon prediction to canonical ``x0`` for ``a(sigma)``/``b(sigma)``. - - For rectified flow, ``a(sigma) = 1 - sigma`` and ``b(sigma) = sigma``. For an - epsilon-prediction model, ``x_sigma = a*x0 + b*epsilon`` so - ``x0 = (x_sigma - b*epsilon) / a``. - """ - noisy = noisy.float() - epsilon = epsilon.float() - sigma = sigma.float() - a = a_fn(sigma).float() - b = b_fn(sigma).float() - if torch.any(a == 0): - raise ValueError("epsilon_to_x0 is undefined where a(sigma) is zero.") - return (noisy - b * epsilon) / a - - def velocity_to_x0(noisy: Tensor, velocity: Tensor, sigma: Tensor) -> Tensor: - """Convert a velocity prediction to canonical ``x0`` (rectified flow). - - ``x_sigma = (1 - sigma) * x0 + sigma * epsilon`` and ``v = epsilon - x0``, so - ``x0 = x_sigma - sigma * v``. - """ + """Convert flow velocity ``epsilon - x0`` to ``x0 = x_sigma - sigma * v``.""" return noisy.float() - sigma.float() * velocity.float() +@torch.no_grad() def dmd_gradient( x0_fake: Tensor, x0_real: Tensor, x_g: Tensor, normalization_epsilon: float = 1e-5, ) -> tuple[Tensor, Tensor, int]: - """Compute the detached normalized fake-minus-real score gradient. + """Return detached normalized fake-minus-real scores for tensors ``(B, ...)``. - Follows Self-Forcing ``_compute_kl_grad`` and LightX2V ``dmd_loss`` exactly: - - ``g = x0_fake - x0_real`` - ``normalizer = mean(abs(x_g - x0_real), non-batch dimensions)`` - ``g_normalized = nan_to_num(g / max(normalizer, normalization_epsilon))`` - - Returns ``(g_normalized, normalizer, nonfinite_count)``. ``normalizer`` is - formed over the entire ``x_g - x0_real`` tensor across all non-batch - dimensions of one sample, ``keepdim`` per sample, and is **not** restricted by - ``gradient_mask``. + Return the gradient, per-sample clamped normalizer and replacement count. + The normalizer is unmasked, including when the caller masks the final loss. """ - if normalization_epsilon <= 0: - raise ValueError(f"normalization_epsilon must be greater than zero, got {normalization_epsilon}.") + if not math.isfinite(normalization_epsilon) or normalization_epsilon <= 0: + raise ValueError(f"normalization_epsilon must be finite and greater than zero, got {normalization_epsilon}.") if x_g.shape != x0_fake.shape or x_g.shape != x0_real.shape: - raise ValueError( - f"x_g, x0_fake, and x0_real must have identical shapes, got " - f"{tuple(x_g.shape)}, {tuple(x0_fake.shape)}, and {tuple(x0_real.shape)}." - ) - if x_g.ndim < 2: - raise ValueError("DMD tensors must include a batch dimension and at least one non-batch dimension.") - - x_g = x_g.float() - x0_fake = x0_fake.float() - x0_real = x0_real.float() - - g = x0_fake - x0_real - # All non-batch dims (keep batch dim), per-sample keepdim. - reduction_dims = tuple(range(1, x_g.dim())) - normalizer = torch.abs(x_g - x0_real).mean(dim=reduction_dims, keepdim=True) - normalizer = torch.maximum(normalizer, torch.as_tensor(normalization_epsilon, device=normalizer.device)) - g_normalized = g / normalizer - nonfinite = (~torch.isfinite(g_normalized)).sum().item() - g_normalized = torch.nan_to_num(g_normalized) - return g_normalized, normalizer, nonfinite - - -def _expand_loss_mask(mask: Tensor, prediction: Tensor) -> Tensor: + raise ValueError("x_g, x0_fake and x0_real must have identical shapes.") + if x_g.ndim < 2 or x_g.numel() == 0: + raise ValueError("DMD tensors require a nonempty batch and at least one non-batch dimension.") + x_g, x0_fake, x0_real = x_g.float(), x0_fake.float(), x0_real.float() + normalizer = (x_g - x0_real).abs().mean(dim=tuple(range(1, x_g.ndim)), keepdim=True) + normalizer = normalizer.clamp_min(normalization_epsilon) + gradient = (x0_fake - x0_real) / normalizer + nonfinite = int((~torch.isfinite(gradient)).sum().item()) + return torch.nan_to_num(gradient), normalizer, nonfinite + + +def expand_loss_mask(mask: Tensor, prediction: Tensor) -> Tensor: """Expand a prefix or broadcastable mask to individual loss elements.""" shape = tuple(mask.shape) mask = mask.bool() @@ -153,34 +85,24 @@ def dmd_surrogate_loss( g_normalized: Tensor, gradient_mask: Optional[Tensor] = None, ) -> tuple[Tensor, int]: - """Surrogate objective ``L_DMD = 0.5 * mean((x_g - stop_gradient(x_g - g_norm))^2)``. - - Only the surrogate loss is masked by ``gradient_mask``. Prefix masks expand - over trailing dimensions. An all-masked loss is an error, not zero. - Returns ``(loss, active_elements)``. - """ + """Compute ``0.5 * mean((x_g - stop_gradient(x_g - g))**2)`` in fp32.""" + if x_g.shape != g_normalized.shape or x_g.numel() == 0: + raise ValueError("Student and normalized gradient must have identical nonempty shapes.") x_g = x_g.float() - g_normalized = g_normalized.float() - target = (x_g - g_normalized).detach() + target = (x_g - g_normalized.float()).detach() if gradient_mask is None: - loss = 0.5 * torch.mean((x_g - target) ** 2) - active = x_g.numel() - else: - mask = _expand_loss_mask(gradient_mask, x_g) - active = int(mask.sum().item()) - if active == 0: - raise ValueError("all-masked DMD loss is an error, not zero.") - diff = (x_g - target) ** 2 - loss = 0.5 * torch.sum(diff[mask]) / active - return loss, active + return 0.5 * (x_g - target).square().mean(), x_g.numel() + mask = expand_loss_mask(gradient_mask, x_g) + active = int(mask.sum().item()) + if active == 0: + raise ValueError("all-masked DMD loss is an error, not zero.") + return 0.5 * (x_g - target).square()[mask].sum() / active, active def fake_score_target(noise: Tensor, x_g: Tensor) -> Tensor: - """Rectified-flow fake-score target ``v_target = epsilon - x_g``. - - The fake score is trained to denoise the generated clean latent with - ``x_sigma = (1 - sigma) * x_g + sigma * epsilon``. - """ + """Construct the detached rectified-flow target ``epsilon - x_g``.""" + if noise.shape != x_g.shape: + raise ValueError("Fake-score noise and generated latents must have identical shapes.") return noise.detach().float() - x_g.detach().float() @@ -190,105 +112,47 @@ def fake_score_loss( x_g: Tensor, gradient_mask: Optional[Tensor] = None, ) -> tuple[Tensor, int]: - """Fake-score denoising MSE against the rectified-flow velocity target. - - ``L_fake = mean((model_output - (epsilon - x_g))^2)``. With a mask, the loss - divides by the number of active elements after applying the mask. - """ + """Compute fp32 denoising MSE; the engine must also detach its noisy inputs.""" target = fake_score_target(noise, x_g) - model_output = model_output.float() + if model_output.shape != target.shape or model_output.numel() == 0: + raise ValueError("Fake-score prediction and target must have identical nonempty shapes.") + difference = (model_output.float() - target).square() if gradient_mask is None: - loss = torch.mean((model_output - target) ** 2) - active = model_output.numel() - else: - mask = _expand_loss_mask(gradient_mask, model_output) - active = int(mask.sum().item()) - if active == 0: - raise ValueError("all-masked fake-score loss is an error, not zero.") - diff = (model_output - target) ** 2 - loss = torch.sum(diff[mask]) / active - return loss, active - - -def ode_regression_loss( - prediction: Tensor, - target: Tensor, - valid_mask: Optional[Tensor] = None, -) -> tuple[Tensor, int]: - """Compute fp32 ODE-target MSE over nonzero-timestep positions.""" - prediction = prediction.float() - target = target.detach().float() - if prediction.shape != target.shape: - raise ValueError( - f"ODE prediction and target must have identical shapes, got {tuple(prediction.shape)} and " - f"{tuple(target.shape)}." - ) - if valid_mask is None: - return torch.mean((prediction - target) ** 2), prediction.numel() - mask = _expand_loss_mask(valid_mask, prediction) + return difference.mean(), difference.numel() + mask = expand_loss_mask(gradient_mask, model_output) active = int(mask.sum().item()) if active == 0: - raise ValueError("all-masked ODE regression loss is an error, not zero.") - return torch.sum((prediction - target)[mask] ** 2) / active, active + raise ValueError("all-masked fake-score loss is an error, not zero.") + return difference[mask].sum() / active, active def ode_euler_step(latents: Tensor, velocity: Tensor, sigma_from: Tensor, sigma_to: Tensor) -> Tensor: - """Apply one deterministic Euler transition used by backward simulation.""" + """Apply one deterministic Euler transition during backward simulation.""" return latents.float() + (sigma_to.float() - sigma_from.float()) * velocity.float() -def consistency_renoise_step(x0: Tensor, noise: Tensor, sigma_to: Tensor) -> Tensor: - """Re-noise a clean prediction with fresh noise for a consistency transition.""" - return (1.0 - sigma_to.float()) * x0.float() + sigma_to.float() * noise.float() - - def standard_cfg( cond: Tensor, uncond: Tensor, guidance_scale: float, cfg_norm: Optional[str] = "none", ) -> Tensor: - """Apply ``uncond + scale * (cond - uncond)`` classifier-free guidance. - - ``cfg_norm`` follows the LightX2V reference: ``layer_norm`` rescales each - last-dimension vector, while ``scalar`` uses one norm ratio for the entire - tensor. - """ - cond = cond.float() - uncond = uncond.float() + """Apply standard CFG, with LightX2V last-vector or batch-global rescaling.""" + cond, uncond = cond.float(), uncond.float() guided = uncond + guidance_scale * (cond - uncond) if cfg_norm in (None, "none"): return guided if cfg_norm == "layer_norm": - cond_norm = torch.norm(cond, dim=-1, keepdim=True) - guided_norm = torch.norm(guided, dim=-1, keepdim=True) - return guided * (cond_norm / guided_norm.clamp_min(1e-12)) + return guided * (cond.norm(dim=-1, keepdim=True) / guided.norm(dim=-1, keepdim=True).clamp_min(1e-12)) if cfg_norm == "scalar": - ratio = torch.norm(cond) / torch.norm(guided).clamp_min(1e-12) - return guided * min(1.0, ratio.item()) + return guided * (cond.norm() / guided.norm().clamp_min(1e-12)).clamp_max(1.0) raise ValueError(f"Unknown cfg_norm {cfg_norm!r}; expected one of {{'none', 'layer_norm', 'scalar'}}.") -def legacy_cfg(cond: Tensor, uncond: Tensor, guidance_scale: float) -> Tensor: - """Self-Forcing CFG form ``cond + legacy_scale*(cond - uncond)``. - - This is *not* numerically identical to :func:`standard_cfg` with the same - scale. A parity recipe must convert explicitly rather than silently reusing the - number. - """ - return cond.float() + guidance_scale * (cond.float() - uncond.float()) - - def timestep_shift(timestep: Tensor, num_train_timesteps: int, shift: float = 1.0) -> Tensor: - """Apply the time-shift remapping exactly once. - - ``shifted = shift * (t / T) / (1 + (shift - 1) * (t / T)) * T``. The - normalization is by the model's ``num_train_timesteps``, so it diverges from a - hardcoded 1000 for any model whose ``num_train_timesteps`` is not 1000. - """ - if shift <= 1.0: - return timestep.float() + """Apply one rational shift, normalized by the model's actual training grid.""" + if not math.isfinite(shift) or shift < 1 or num_train_timesteps <= 0: + raise ValueError("The timestep shift must be finite and at least 1, with a positive training grid size.") t = timestep.float() frac = t / num_train_timesteps - shifted = shift * frac / (1.0 + (shift - 1.0) * frac) * num_train_timesteps - return shifted + return shift * frac / (1.0 + (shift - 1.0) * frac) * num_train_timesteps diff --git a/verl_omni/trainer/diffusion/ray_diffusion_trainer.py b/verl_omni/trainer/diffusion/ray_diffusion_trainer.py index b418234be..4abf075bd 100644 --- a/verl_omni/trainer/diffusion/ray_diffusion_trainer.py +++ b/verl_omni/trainer/diffusion/ray_diffusion_trainer.py @@ -16,13 +16,19 @@ This trainer supports model-agnostic model initialization with Hugging Face. """ +import hashlib import json import logging import os +import random +import shutil import subprocess +import tempfile +import time import uuid from abc import ABC, abstractmethod from collections import defaultdict +from pathlib import Path from pprint import pprint from typing import Any, Literal, Optional @@ -768,6 +774,10 @@ def init_workers(self): def _teacher_wg_name(key: str) -> str: return f"teacher_{key.replace('/', '_')}" + def actor_worker_extra_kwargs(self): + """Additional typed settings for an algorithm-specific training worker.""" + return {} + def _init_colocated_workers(self): """Create Ray pools and colocated actor/ref worker groups (online and offline).""" self.resource_pool_manager.create_resource_pool() @@ -788,6 +798,7 @@ def _init_colocated_workers(self): config=self.config.actor_rollout_ref, distillation_config=self.config.get("distillation"), role=str(actor_role), + **self.actor_worker_extra_kwargs(), ) self.resource_pool_to_cls[actor_rollout_resource_pool][str(actor_role)] = actor_rollout_cls else: @@ -1091,7 +1102,7 @@ def _update_actor(self, batch: DataProto) -> DataProto: actor_output["perf/mfu/actor"] = actor_mfu return DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output}) - def _start_profiling(self, do_profile: bool) -> None: + def _start_profiling(self, do_profile: bool, profile_step: Optional[int] = None) -> None: """Start profiling for all worker groups if profiling is enabled.""" if not do_profile: return @@ -1105,7 +1116,9 @@ def _start_profiling(self, do_profile: bool) -> None: self._controller_nsys_profile_active = True controller_profile_started = True - self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps) + self.actor_rollout_wg.start_profile( + role="e2e", profile_step=self.global_steps if profile_step is None else profile_step + ) if self.use_reference_policy and not self.ref_in_actor: self.ref_policy_wg.start_profile(profile_step=self.global_steps) if self.use_teacher_policy and Role.TeacherModel in self.role_worker_mapping: @@ -1489,6 +1502,335 @@ def fit(self): self.train_dataset.on_batch_end(batch=batch) +class DistributionMatchingRayTrainer(BaseRayDiffusionTrainer): + """Offline DMD2: one student attempt followed by K fake-score attempts.""" + + def __init__(self, config, *args, **kwargs): + self.validate_config(config) + super().__init__(config, *args, **kwargs) + self.dmd_config = omega_conf_to_dataclass(config.dmd) + self.global_steps = 0 + self.optimizer_steps = {"student": 0, "fake_score": 0} + self.data_epoch = 0 + self.batch_iterator = None + self.failed = False + self.use_reference_policy = False + + @staticmethod + def validate_config(config): + """Reject unsupported/composite execution before worker or teacher allocation.""" + if config.algorithm.sample_source != "offline": + raise ValueError("DMD2 requires algorithm.sample_source=offline.") + actor = config.actor_rollout_ref.actor + if config.actor_rollout_ref.model.algorithm != "dmd2" or actor.diffusion_loss.loss_mode != "dmd2": + raise ValueError( + "DMD2 requires model.algorithm=dmd2 and diffusion_loss.loss_mode=dmd2; original dmd is separate." + ) + if config.actor_rollout_ref.model.model_type != "diffusion_dmd_model": + raise ValueError("DMD2 requires model.model_type=diffusion_dmd_model.") + if ( + config.distillation.enabled + or actor.use_distill_loss + or actor.use_kl_loss + or config.algorithm.get("use_kl_in_reward", False) + ): + raise ValueError("DMD2 distribution-only cannot enable OPD or additional KL objectives.") + if config.algorithm.paired_preference or config.actor_rollout_ref.get("separate", False): + raise ValueError("DMD2 uses unpaired prompts and colocated training, not separate online rollout.") + if actor.strategy not in {"fsdp", "fsdp2"} or actor.ppo_epochs != 1: + raise ValueError("DMD2 requires FSDP/FSDP2 and exactly one optimizer attempt per actor call.") + engine = actor.fsdp_config + rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) or config.actor_rollout_ref.model.get( + "lora_rank", 0 + ) + if rank <= 0 or (actor.strategy == "fsdp" and not engine.use_orig_params): + raise ValueError("DMD2 shared-base LoRA requires positive rank and FSDP1 use_orig_params=true.") + if engine.get("use_dynamic_bsz", False): + raise ValueError("DMD2 currently uses static microbatches with synchronized forward counts.") + world = config.trainer.n_gpus_per_node * config.trainer.nnodes + sp = engine.ulysses_sequence_parallel_size + if world <= 0 or sp <= 0 or world % sp or config.data.train_batch_size % (world // sp): + raise ValueError("DMD2 global batch must divide evenly across data-parallel ranks.") + if config.data.get("gen_batch_size") not in (None, config.data.train_batch_size): + raise ValueError("DMD2 gen_batch_size must equal train_batch_size.") + if config.trainer.default_hdfs_dir is not None or actor.checkpoint.get("async_save", False): + raise ValueError("DMD2 uses synchronous atomic checkpoints in a shared local directory.") + for contents in (actor.checkpoint.save_contents, actor.checkpoint.load_contents): + if not {"model", "optimizer", "extra"}.issubset(contents): + raise ValueError("DMD2 checkpoints require saving and loading model, optimizer and extra state.") + if config.trainer.get("val_only", False): + raise ValueError("DMD2 validation-only serving is not implemented; use the exported student artifact.") + omega_conf_to_dataclass(config.dmd) + + def actor_worker_extra_kwargs(self): + """Pass standalone DMD settings through the shared worker constructor.""" + return {"dmd_config": self.config.dmd} + + def _validate(self): + """Offline DMD2 exports inference artifacts rather than starting an RL server.""" + return {"val/offline/skipped": 1.0} + + def next_batch(self): + """Reuse the stateful dataloader, including its sampler and resume position.""" + if self.batch_iterator is None: + self.batch_iterator = iter(self.train_dataloader) + try: + batch = next(self.batch_iterator) + except StopIteration: + self.data_epoch += 1 + self.batch_iterator = iter(self.train_dataloader) + batch = next(self.batch_iterator) + return _to_diffusion_worker_tensordict(DataProto.from_single_dict(batch)) + + def update_stage(self, stage, repeat): + """Dispatch one attempt and validate its numerical outcome, never retry it.""" + batch = self.next_batch() + tu.assign_non_tensor(batch, dmd_stage=stage) + output = self.actor_rollout_wg.update_actor(batch) + raw = tu.get(output, "metrics") + metrics = {key: float(value) for key, value in reduce_metrics(raw).items()} + applied = metrics.get("dmd/update_applied") + skipped = metrics.get("dmd/skip_nonfinite") + if applied not in (0.0, 1.0) or skipped != 1.0 - applied: + raise RuntimeError("Malformed DMD2 optimizer outcome; restore the last complete checkpoint.") + self.optimizer_steps[stage] += int(applied) + metrics["training/samples"] = self.config.data.train_batch_size + return {f"{stage}/{repeat}/{key}": value for key, value in metrics.items()} + + def configuration_fingerprint(self): + """Canonicalize the mathematical, optimizer and data settings, not output paths.""" + payload = { + "model": OmegaConf.to_container(self.config.actor_rollout_ref.model, resolve=True), + "engine": OmegaConf.to_container(self.config.actor_rollout_ref.actor.fsdp_config, resolve=True), + "optimizer": OmegaConf.to_container(self.config.actor_rollout_ref.actor.optim, resolve=True), + "dmd": OmegaConf.to_container(self.config.dmd, resolve=True), + "data": OmegaConf.to_container(self.config.data, resolve=True), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + def validate_actor_checkpoint(self, path, step, counts): + """Check all rank files and role clocks before publication or model mutation.""" + world = self.config.trainer.n_gpus_per_node * self.config.trainer.nnodes + for rank in range(world): + for filename in ( + f"model_world_size_{world}_rank_{rank}.pt", + f"optim_world_size_{world}_rank_{rank}.pt", + f"extra_state_world_size_{world}_rank_{rank}.pt", + f"dmd_state_rank_{rank}.pt", + ): + if not (path / "actor" / filename).is_file(): + raise ValueError(f"Incomplete DMD2 checkpoint: missing {filename}.") + state = torch.load(path / "actor" / f"dmd_state_rank_{rank}.pt", map_location="cpu", weights_only=False) + if state.get("version") != 1 or state.get("world_size") != world or state.get("optimizer_steps") != counts: + raise ValueError("Worker and trainer DMD2 checkpoint versions or counters do not match.") + quotas = {"student": step, "fake_score": step * self.dmd_config.fake_update_ratio} + if state.get("skipped_steps") != {role: quota - counts[role] for role, quota in quotas.items()}: + raise ValueError("DMD2 checkpoint skipped-update counters do not match completed cycles.") + if not {"fake_optimizer", "fake_scheduler", "generators"}.issubset(state): + raise ValueError("DMD2 checkpoint is missing optimizer, scheduler or sampling state.") + + def _save_checkpoint(self): + """Publish actor shards, both clocks, dataloader and RNG as one complete cycle.""" + root = Path(self.config.trainer.default_local_dir) + root.mkdir(parents=True, exist_ok=True) + target = root / f"global_step_{self.global_steps}" + if target.exists(): + raise FileExistsError(f"Refusing to overwrite checkpoint {target}.") + temporary = Path(tempfile.mkdtemp(prefix=f".global_step_{self.global_steps}_", dir=root)) + try: + self.actor_rollout_wg.save_checkpoint(str(temporary / "actor"), global_step=self.global_steps) + self.validate_actor_checkpoint(temporary, self.global_steps, self.optimizer_steps) + torch.save(self.train_dataloader.state_dict(), temporary / "data.pt") + torch.save( + { + "version": 1, + "global_step": self.global_steps, + "optimizer_steps": self.optimizer_steps, + "data_epoch": self.data_epoch, + "configuration": self.configuration_fingerprint(), + "rng": { + "torch": torch.get_rng_state(), + "numpy": np.random.get_state(), + "python": random.getstate(), + }, + }, + temporary / "trainer.pt", + ) + os.replace(temporary, target) + tracker = root / f".latest_{uuid.uuid4().hex}" + tracker.write_text(str(self.global_steps)) + os.replace(tracker, root / "latest_checkpointed_iteration.txt") + except Exception: + if temporary.exists(): + shutil.rmtree(temporary) + raise + keep = self.config.trainer.get("max_actor_ckpt_to_keep") + if keep is not None and keep > 0: + checkpoints = sorted( + ( + p + for p in root.glob("global_step_*") + if p.is_dir() + and not p.is_symlink() + and p.name.removeprefix("global_step_").isdigit() + and (p / "trainer.pt").is_file() + ), + key=lambda p: int(p.name.removeprefix("global_step_")), + ) + for old in checkpoints[:-keep]: + shutil.rmtree(old) + + def _load_checkpoint(self): + """Validate the new format before loading; old prototype checkpoints are not reinterpreted.""" + mode = self.config.trainer.resume_mode + if mode == "disable": + return + if mode == "auto": + path = find_latest_ckpt_path(self.config.trainer.default_local_dir) + if path is None: + return + elif mode == "resume_path": + path = self.config.trainer.resume_from_path + else: + raise ValueError(f"Unknown resume_mode {mode!r}.") + path = Path(path) + if not (path / "trainer.pt").is_file() or not (path / "data.pt").is_file(): + raise ValueError("Incomplete/old DMD2 checkpoint: trainer.pt and data.pt are required.") + state = torch.load(path / "trainer.pt", map_location="cpu", weights_only=False) + if state.get("version") != 1 or state.get("configuration") != self.configuration_fingerprint(): + raise ValueError( + "DMD2 checkpoint format/configuration does not match; prototype migration is not implicit." + ) + counts = state.get("optimizer_steps", {}) + step = state.get("global_step") + if isinstance(step, bool) or not isinstance(step, int) or step < 0 or set(counts) != {"student", "fake_score"}: + raise ValueError("Invalid DMD2 checkpoint counters.") + for role, maximum in (("student", step), ("fake_score", step * self.dmd_config.fake_update_ratio)): + if isinstance(counts[role], bool) or not isinstance(counts[role], int) or not 0 <= counts[role] <= maximum: + raise ValueError("DMD2 successful optimizer counts exceed completed attempt quotas.") + self.validate_actor_checkpoint(path, step, counts) + if ( + isinstance(state.get("data_epoch"), bool) + or not isinstance(state.get("data_epoch"), int) + or state["data_epoch"] < 0 + ): + raise ValueError("Invalid DMD2 checkpoint data epoch.") + if not {"torch", "numpy", "python"}.issubset(state.get("rng", {})): + raise ValueError("Missing DMD2 driver RNG state.") + data_state = torch.load(path / "data.pt", weights_only=False) + restored = self.actor_rollout_wg.load_checkpoint(str(path / "actor"), del_local_after_load=False) + if restored is not None and any(item != counts for item in restored): + raise ValueError("Worker and trainer optimizer counters do not match.") + self.train_dataloader.load_state_dict(data_state) + self.global_steps, self.optimizer_steps, self.data_epoch = step, counts, state["data_epoch"] + torch.set_rng_state(state["rng"]["torch"]) + np.random.set_state(state["rng"]["numpy"]) + random.setstate(state["rng"]["python"]) + self.batch_iterator = None + + def export_student(self): + """Write one student inference artifact, separately from the resumable checkpoint.""" + directory = Path(self.config.trainer.default_local_dir) / "inference" + metadata = { + "algorithm": "dmd2", + "role": self.dmd_config.export_role, + "global_step": self.global_steps, + "student_optimizer_steps": self.optimizer_steps["student"], + "base_model": self.config.actor_rollout_ref.model.path, + "configuration": self.configuration_fingerprint(), + "sampler": "ode_euler", + "num_inference_steps": self.config.actor_rollout_ref.model.pipeline.num_inference_steps, + "height": self.config.actor_rollout_ref.model.pipeline.height, + "width": self.config.actor_rollout_ref.model.pipeline.width, + "max_sequence_length": self.config.actor_rollout_ref.model.pipeline.max_sequence_length, + "rollout_timestep_shift": self.dmd_config.rollout_timestep_shift, + "guidance_scale": 1.0, + } + provenance = self.actor_rollout_wg.get_model_provenance() + if not provenance or any(value != provenance[0] for value in provenance): + raise ValueError("DMD2 workers disagree on base checkpoint provenance.") + metadata.update(provenance[0]) + manifest = directory / "inference_manifest.json" + if directory.exists(): + weights = directory / "adapter_model.safetensors" + if not manifest.is_file() or not weights.is_file() or not (directory / "adapter_config.json").is_file(): + raise ValueError("Incomplete inference artifact; choose a new output directory.") + with weights.open("rb") as file: + metadata["weights_sha256"] = hashlib.file_digest(file, "sha256").hexdigest() + if json.loads(manifest.read_text()) != metadata: + raise ValueError("Incompatible inference artifact; choose a new output directory.") + return + directory.parent.mkdir(parents=True, exist_ok=True) + temporary = Path(tempfile.mkdtemp(prefix=".inference_", dir=directory.parent)) + try: + artifact = temporary / "adapter" + self.actor_rollout_wg.export_student(str(artifact), role=self.dmd_config.export_role) + with (artifact / "adapter_model.safetensors").open("rb") as file: + metadata["weights_sha256"] = hashlib.file_digest(file, "sha256").hexdigest() + (artifact / "inference_manifest.json").write_text(json.dumps(metadata, indent=2, sort_keys=True)) + os.replace(artifact, directory) + finally: + shutil.rmtree(temporary) + + def fit(self): + """Use shared worker/data/profiling services with a finite explicit 1:K loop.""" + from verl.utils.tracking import Tracking + + if self.failed: + raise RuntimeError("A failed DMD2 trainer must be reconstructed from its last complete checkpoint.") + logger = Tracking( + project_name=self.config.trainer.project_name, + experiment_name=self.config.trainer.experiment_name, + default_backend=self.config.trainer.logger, + config=OmegaConf.to_container(self.config, resolve=True), + ) + self._load_checkpoint() + profile_steps = self.config.global_profiler.steps or [] + profiling = False + progress = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="DMD2 Training") + try: + while self.global_steps < self.total_training_steps: + cycle = self.global_steps + 1 + if cycle in profile_steps and not profiling: + self._start_profiling(True, profile_step=cycle) + profiling = True + start = time.perf_counter() + metrics = self.update_stage("student", 0) + for repeat in range(self.dmd_config.fake_update_ratio): + metrics.update(self.update_stage("fake_score", repeat)) + self.global_steps = cycle + metrics["perf/cycle_s"] = time.perf_counter() - start + metrics["training/global_step"] = cycle + metrics["training/data_epoch"] = self.data_epoch + for stage, count in self.optimizer_steps.items(): + metrics[f"training/{stage}_optimizer_steps"] = count + if self.config.trainer.save_freq > 0 and ( + cycle % self.config.trainer.save_freq == 0 or cycle == self.total_training_steps + ): + start = time.perf_counter() + self._save_checkpoint() + metrics["perf/checkpoint_s"] = time.perf_counter() - start + if self.config.trainer.test_freq > 0 and cycle % self.config.trainer.test_freq == 0: + metrics.update(self._validate()) + logger.log(data=metrics, step=cycle) + progress.update(1) + if profiling and ( + not self.config.global_profiler.profile_continuous_steps or cycle + 1 not in profile_steps + ): + self._stop_profiling(True) + profiling = False + if not all(self.optimizer_steps.values()): + raise RuntimeError("The finite DMD2 attempt budget ended without successful updates for both roles.") + self.export_student() + except Exception: + self.failed = True + raise + finally: + if profiling: + self._stop_profiling(True) + progress.close() + + class DirectPreferenceRayTrainer(BaseRayDiffusionTrainer): """Direct-preference diffusion trainer for DPO, DiffusionNFT, AWM, etc.""" diff --git a/verl_omni/trainer/main_diffusion.py b/verl_omni/trainer/main_diffusion.py index a0ace0966..0ce5a1a37 100644 --- a/verl_omni/trainer/main_diffusion.py +++ b/verl_omni/trainer/main_diffusion.py @@ -130,13 +130,13 @@ def _get_trainer_cls(config): return PolicyGradientRayTrainer if trainer_type == "direct_preference": return DirectPreferenceRayTrainer - if trainer_type == "distillation": - from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationRayTrainer + if trainer_type == "distribution_matching": + from verl_omni.trainer.diffusion.ray_diffusion_trainer import DistributionMatchingRayTrainer - return DistillationRayTrainer + return DistributionMatchingRayTrainer raise ValueError( f"Unsupported diffusion trainer_type {trainer_type!r}. " - f"Expected one of: 'policy_gradient', 'direct_preference', 'distillation'." + "Expected one of: 'policy_gradient', 'direct_preference', 'distribution_matching'." ) @@ -160,6 +160,15 @@ def add_actor_rollout_worker(self, config): from verl.single_controller.ray import RayWorkerGroup from verl.trainer.ppo.ray_trainer import Role + if config.algorithm.trainer_type == "distribution_matching": + from verl_omni.trainer.diffusion.ray_diffusion_trainer import DistributionMatchingRayTrainer + from verl_omni.workers.dmd_worker import DMDTrainingWorker + + DistributionMatchingRayTrainer.validate_config(config) + self.role_worker_mapping[Role.Actor] = ray.remote(DMDTrainingWorker) + self.mapping[Role.Actor] = "global_pool" + return DMDTrainingWorker, RayWorkerGroup + from verl_omni.workers.engine_workers import ActorRolloutRefWorker actor_rollout_cls = ActorRolloutRefWorker diff --git a/verl_omni/utils/fs.py b/verl_omni/utils/fs.py index 9a6f408e8..927ea10b2 100644 --- a/verl_omni/utils/fs.py +++ b/verl_omni/utils/fs.py @@ -12,11 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. +import hashlib import os +from pathlib import Path from verl.utils.fs import copy_to_local -__all__ = ["resolve_model_local_dir"] +__all__ = ["resolve_model_local_dir", "diffusion_model_provenance"] + + +def diffusion_model_provenance(local_path: str) -> dict: + """Record a resolved snapshot's revision when available and its transformer config hash.""" + root = Path(local_path) + revision = root.name if root.parent.name == "snapshots" else None + metadata = root / ".cache/huggingface/download/model_index.json.metadata" + if metadata.is_file(): + with metadata.open() as file: + revision = file.readline().strip() + if not revision or len(revision) != 40 or any(char not in "0123456789abcdef" for char in revision): + revision = None + with (root / "transformer/config.json").open("rb") as file: + config_hash = hashlib.file_digest(file, "sha256").hexdigest() + return {"base_model_revision": revision, "base_transformer_config_sha256": config_hash} def resolve_model_local_dir(path: str, use_shm: bool = False) -> str: diff --git a/verl_omni/utils/net_utils.py b/verl_omni/utils/net_utils.py new file mode 100644 index 000000000..3a814c8ed --- /dev/null +++ b/verl_omni/utils/net_utils.py @@ -0,0 +1,49 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random +import socket + + +def ephemeral_port_range() -> tuple[int, int]: + try: + with open("/proc/sys/net/ipv4/ip_local_port_range") as f: + lo, hi = f.read().split() + return int(lo), int(hi) + except (OSError, ValueError): + return 32768, 60999 # Linux default + + +def get_non_ephemeral_free_port(address: str = "127.0.0.1") -> int: + """Pick a free port below the kernel's ephemeral port range. + + The consumer binds the port only after worker spawn + model load, and ports + inside the ephemeral range can be claimed as source ports by unrelated + connections in that window (later failing to listen with ``EADDRINUSE``). + """ + lo, _ = ephemeral_port_range() + if lo <= 1024: + raise RuntimeError(f"Ephemeral port range starts at {lo}; no non-privileged candidate ports below it.") + candidates = range(1024, lo) + start = random.randrange(len(candidates)) + for offset in range(len(candidates)): + port = candidates[(start + offset) % len(candidates)] + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind((address, port)) + except OSError: + continue + return port + raise RuntimeError(f"No free non-ephemeral port on {address} below {lo}.") diff --git a/verl_omni/workers/config/diffusion/__init__.py b/verl_omni/workers/config/diffusion/__init__.py index 0cb75d919..b1fa0f545 100644 --- a/verl_omni/workers/config/diffusion/__init__.py +++ b/verl_omni/workers/config/diffusion/__init__.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import actor, distillation, model, rollout +from . import actor, distillation, dmd, model, rollout from .actor import * # noqa: F401 from .distillation import * # noqa: F401 +from .dmd import * # noqa: F401 from .model import * # noqa: F401 from .rollout import * # noqa: F401 -__all__ = actor.__all__ + distillation.__all__ + model.__all__ + rollout.__all__ +__all__ = actor.__all__ + distillation.__all__ + dmd.__all__ + model.__all__ + rollout.__all__ diff --git a/verl_omni/workers/config/diffusion/actor.py b/verl_omni/workers/config/diffusion/actor.py index 2cd1261af..04baa6fb4 100644 --- a/verl_omni/workers/config/diffusion/actor.py +++ b/verl_omni/workers/config/diffusion/actor.py @@ -55,6 +55,7 @@ def __post_init__(self): "grpo_guard", "diffusion_nft", "dpo", + "dmd2", "dance_grpo", "distill_kl", "distill_fm_mse", diff --git a/verl_omni/workers/config/diffusion/distillation.py b/verl_omni/workers/config/diffusion/distillation.py index 3e008628d..3fc87f5f9 100644 --- a/verl_omni/workers/config/diffusion/distillation.py +++ b/verl_omni/workers/config/diffusion/distillation.py @@ -17,11 +17,7 @@ from verl.base_config import BaseConfig -__all__ = [ - "DiffusionDistillationTeacherModelConfig", - "DiffusionDistributionMatchingConfig", - "DiffusionDistillationConfig", -] +__all__ = ["DiffusionDistillationTeacherModelConfig", "DiffusionDistillationConfig"] @dataclass @@ -50,68 +46,9 @@ def check_configured(self): raise ValueError("key must be specified for distillation teacher model config.") -@dataclass -class DiffusionDistributionMatchingConfig(BaseConfig): - """Architecture-neutral DMD-family recipe selection. - - This config is active only when ``algorithm.trainer_type=distillation``. - The existing parent ``enabled`` flag remains exclusively owned by on-policy - distillation and must stay false for DMD-family training. - """ - - # Registered recipe name. - recipe: str = "dmd2" - # Optional recipe profile; null selects the recipe default. - profile: Optional[str] = None - # Optional fake-score phase count; null selects the recipe default. - fake_update_ratio: Optional[int] = None - # Number of fake/discriminator-only cycles before student updates begin. - fake_warmup_cycles: int = 0 - # Optional registered rollout override; null selects the recipe default. - rollout_strategy: Optional[str] = None - # Optional data-mode override; null selects the recipe default. - data_mode: Optional[str] = None - # Semantic role exported to inference replicas. - export_role: str = "student_ema" - - def __post_init__(self): - valid_recipes = {"dmd", "dmd2", "causvid", "self_forcing"} - if self.recipe not in valid_recipes: - raise ValueError(f"Invalid recipe: {self.recipe}. Must be one of {sorted(valid_recipes)}") - valid_profiles = {"distribution_only", "paper"} - if self.profile is not None and self.profile not in valid_profiles: - raise ValueError(f"Invalid profile: {self.profile}. Must be one of {sorted(valid_profiles)}") - if self.fake_update_ratio is not None and self.fake_update_ratio <= 0: - raise ValueError(f"fake_update_ratio must be greater than 0, got {self.fake_update_ratio}") - if self.fake_warmup_cycles < 0: - raise ValueError(f"fake_warmup_cycles must be non-negative, got {self.fake_warmup_cycles}") - valid_rollout_strategies = { - "backward_simulated", - "consistency_renoise", - "ode_euler", - "one_step", - "self_forced", - "teacher_forced_causal", - } - if self.rollout_strategy is not None and self.rollout_strategy not in valid_rollout_strategies: - raise ValueError( - f"Invalid rollout_strategy: {self.rollout_strategy}. Must be one of {sorted(valid_rollout_strategies)}" - ) - valid_data_modes = {"prompts", "prompt_and_real_latent", "regression_pairs"} - if self.data_mode is not None and self.data_mode not in valid_data_modes: - raise ValueError(f"Invalid data_mode: {self.data_mode}. Must be one of {sorted(valid_data_modes)}") - valid_export_roles = {"student", "student_ema"} - if self.export_role not in valid_export_roles: - raise ValueError(f"Invalid export_role: {self.export_role}. Must be one of {sorted(valid_export_roles)}") - - @dataclass class DiffusionDistillationConfig(BaseConfig): - """Diffusion distillation settings shared by OPD and DMD-family routing. - - ``enabled`` and the teacher-pool fields remain exclusive to OPD. DMD-family - training is selected by ``algorithm.trainer_type=distillation`` and reads the - nested ``distribution_matching`` config while keeping ``enabled=false``. + """Diffusion on-policy distillation. enabled (bool): Whether on-policy distillation is enabled. @@ -149,7 +86,7 @@ class DiffusionDistillationConfig(BaseConfig): ``` """ - _mutable_fields = BaseConfig._mutable_fields | {"teacher_models", "distribution_matching"} + _mutable_fields = BaseConfig._mutable_fields | {"teacher_models"} enabled: bool = False n_gpus_per_node: int = 0 @@ -157,10 +94,6 @@ class DiffusionDistillationConfig(BaseConfig): teacher_models: dict[str, DiffusionDistillationTeacherModelConfig] = field(default_factory=dict) teacher_key: str = "data_source" scheduler: str = "inline" - # DMD-family recipe settings; selected by algorithm.trainer_type rather than enabled. - distribution_matching: DiffusionDistributionMatchingConfig = field( - default_factory=DiffusionDistributionMatchingConfig - ) def __post_init__(self): if not self.enabled: diff --git a/verl_omni/workers/config/diffusion/dmd.py b/verl_omni/workers/config/diffusion/dmd.py new file mode 100644 index 000000000..b8f045967 --- /dev/null +++ b/verl_omni/workers/config/diffusion/dmd.py @@ -0,0 +1,102 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass, field +from functools import partial + +from verl.base_config import BaseConfig +from verl.workers.config.optimizer import FSDPOptimizerConfig + +__all__ = ["DiffusionDMDConfig"] + + +@dataclass +class DiffusionDMDConfig(BaseConfig): + """DMD2 distribution-only settings, independent of on-policy distillation.""" + + # Fake-score update attempts per student attempt. + fake_update_ratio: int = 2 + # Student physical microbatch size per data-parallel rank. + student_micro_batch_size_per_gpu: int = 1 + # Fake-score physical microbatch size per data-parallel rank. + fake_score_micro_batch_size_per_gpu: int = 1 + # Fake score uses the existing optimizer type, independently of the actor. + fake_score_optim: FSDPOptimizerConfig = field( + default_factory=partial(FSDPOptimizerConfig, lr=2e-5, weight_decay=0.001) + ) + # Qwen teacher CFG; student and fake score remain conditional-only. + teacher_guidance_scale: float = 4.0 + # Teacher velocity normalization in packed denoiser space. + cfg_norm: str = "layer_norm" + # Explicit negative teacher condition; an empty string is also valid. + negative_prompt: str = " " + # Per-sample score-difference normalization floor. + normalization_epsilon: float = 1e-6 + # Fixed rational shift applied once to the inference sigma grid. + rollout_timestep_shift: float = 3.0 + # Discrete score grid size; zero selects continuous uniform sampling. + score_discrete_steps: int = 1000 + # Lower score sigma bound. + score_sigma_min: float = 0.02 + # Upper score sigma bound. + score_sigma_max: float = 0.98 + # Rational shift for discrete score sampling only. + score_timestep_shift: float = 3.0 + # EMA decay, applied only after successful student updates. + ema_decay: float = 0.999 + # Successful student update count at which EMA starts. + ema_start_step: int = 0 + # Default inference artifact; EMA is an explicit alternative. + export_role: str = "student" + + def __post_init__(self): + for name in ("fake_update_ratio", "student_micro_batch_size_per_gpu", "fake_score_micro_batch_size_per_gpu"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + for name in ("score_discrete_steps", "ema_start_step"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer, got {value!r}") + for name in ( + "teacher_guidance_scale", + "normalization_epsilon", + "rollout_timestep_shift", + "score_sigma_min", + "score_sigma_max", + "score_timestep_shift", + "ema_decay", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(value): + raise ValueError(f"{name} must be finite, got {value!r}") + if self.teacher_guidance_scale <= 0: + raise ValueError("teacher_guidance_scale must be positive") + if not isinstance(self.negative_prompt, str): + raise ValueError("negative_prompt must be an explicit string for teacher conditioning") + if self.normalization_epsilon <= 0: + raise ValueError("normalization_epsilon must be positive") + if self.rollout_timestep_shift < 1 or self.score_timestep_shift < 1: + raise ValueError("rollout_timestep_shift and score_timestep_shift must be at least 1") + if not 0 < self.score_sigma_min < self.score_sigma_max <= 1: + raise ValueError("score sigma bounds must satisfy 0 < score_sigma_min < score_sigma_max <= 1") + if not 0 <= self.ema_decay <= 1: + raise ValueError("ema_decay must be between 0 and 1") + valid_cfg_norms = {"none", "layer_norm", "scalar"} + if self.cfg_norm not in valid_cfg_norms: + raise ValueError(f"Invalid cfg_norm: {self.cfg_norm}. Must be one of {sorted(valid_cfg_norms)}") + valid_export_roles = {"student", "student_ema"} + if self.export_role not in valid_export_roles: + raise ValueError(f"Invalid export_role: {self.export_role}. Must be one of {sorted(valid_export_roles)}") diff --git a/verl_omni/workers/config/omni/model.py b/verl_omni/workers/config/omni/model.py index 3f625e7ae..0ba341646 100644 --- a/verl_omni/workers/config/omni/model.py +++ b/verl_omni/workers/config/omni/model.py @@ -168,6 +168,14 @@ def __post_init__(self): # Build hf_config so the FSDP engine can load and wrap the model. self.local_hf_config_path = copy_to_local(self.hf_config_path, use_shm=self.use_shm) attn_implementation = self.override_config.get("attn_implementation", "flash_attention_2") + from verl_omni.pipelines.model_base import OmniModelBase + + if self.load_tokenizer: + adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) + else: + adapter_cls = OmniModelBase.peek_class(self.architecture, self.model_stage) + if adapter_cls is not None: + adapter_cls.register_auto_classes() self.hf_config = AutoConfig.from_pretrained( self.local_hf_config_path, trust_remote_code=self.trust_remote_code, @@ -178,10 +186,7 @@ def __post_init__(self): self.architectures = getattr(self.hf_config, "architectures", None) if self.load_tokenizer: - from verl_omni.pipelines.model_base import OmniModelBase - self.local_tokenizer_path = copy_to_local(self.tokenizer_path, use_shm=self.use_shm) - adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) self.tokenizer = adapter_cls.configure_tokenizer(self.local_tokenizer_path, self) self.processor = adapter_cls.configure_processor(self.local_path, self) diff --git a/verl_omni/workers/dmd_worker.py b/verl_omni/workers/dmd_worker.py new file mode 100644 index 000000000..62e7aef0e --- /dev/null +++ b/verl_omni/workers/dmd_worker.py @@ -0,0 +1,111 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DMD2 specialization of the shared diffusion TrainingWorker.""" + +from functools import partial + +from verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register +from verl.utils import tensordict_utils as tu +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import TrainingWorkerConfig +from verl.workers.engine import EngineRegistry + +from verl_omni.workers.engine_workers import TrainingWorker +from verl_omni.workers.utils.losses import diffusion_loss + + +class DMDTrainingWorker(TrainingWorker): + """Pass DMD settings into the engine without duplicating worker initialization.""" + + def __init__(self, config, *, dmd_config, role="actor", distillation_config=None): + if role != "actor" or (distillation_config is not None and distillation_config.get("enabled", False)): + raise ValueError("DMD2 uses one offline actor group, not OPD teacher workers.") + self.dmd_config = omega_conf_to_dataclass(dmd_config) + self.actor_config = omega_conf_to_dataclass(config.actor) + model_config = omega_conf_to_dataclass(config.model) + profiler = self.actor_config.profiler + if profiler is not None and profiler.tool_config.get(profiler.tool) is not None: + profiler.tool_config[profiler.tool] = omega_conf_to_dataclass( + config.actor.profiler.tool_config[profiler.tool] + ) + worker_config = TrainingWorkerConfig( + model_type="diffusion_dmd_model", + model_config=model_config, + engine_config=self.actor_config.engine, + optimizer_config=self.actor_config.optim, + checkpoint_config=self.actor_config.checkpoint, + profiler_config=profiler, + ) + super().__init__(worker_config) + self.loss_fn = partial(diffusion_loss, config=self.actor_config) + + def build_engine(self): + """Use the registry with one additional typed DMD configuration.""" + from verl_omni.workers.engine.fsdp import dmd_impl # noqa: F401 + + return EngineRegistry.new( + model_type=self.config.model_type, + backend=self.engine_config.strategy, + model_config=self.model_config, + engine_config=self.engine_config, + optimizer_config=self.optimizer_config, + checkpoint_config=self.checkpoint_config, + dmd_config=self.dmd_config, + ) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_model(self): + """Reuse the standard worker reset/model initialization boundary.""" + self.reset() + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=True) + def update_actor(self, data): + """Execute exactly one optimizer attempt using existing mini/microbatch machinery.""" + stage = tu.get_non_tensor_data(data, "dmd_stage", default="student") + if stage not in {"student", "fake_score"}: + raise ValueError(f"Invalid DMD2 stage {stage!r}.") + if tu.get_non_tensor_data(data, "epochs", default=1) != 1: + raise ValueError("DMD2 update_actor performs one optimizer attempt, not multiple epochs.") + previous = self.engine.active_stage + self.engine.select_stage(stage) + tu.assign_non_tensor( + data, + global_token_num=None, + num_mini_batch=1, + mini_batch_size=None, + epochs=1, + dataloader_kwargs={"shuffle": False}, + micro_batch_size_per_gpu=getattr(self.dmd_config, f"{stage}_micro_batch_size_per_gpu"), + ) + try: + result = self.train_mini_batch(data) + if result is not None: + metrics = tu.get_non_tensor_data(result, "metrics", default=None) + metrics["dmd/update_applied"] = float(self.engine.last_step_succeeded) + metrics["dmd/skip_nonfinite"] = float(not self.engine.last_step_succeeded) + return result + finally: + self.engine.select_stage(previous) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def get_model_provenance(self): + """Read the resolved checkpoint identity used by this worker, not a moving hub alias.""" + from verl_omni.utils.fs import diffusion_model_provenance + + return diffusion_model_provenance(self.engine.model_config.local_path) + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def export_student(self, directory, role="student"): + """Export the selected student adapter without exposing score-model weights.""" + self.engine.export_student(directory, role) diff --git a/verl_omni/workers/engine/fsdp/diffusers_impl.py b/verl_omni/workers/engine/fsdp/diffusers_impl.py index 928141b8c..94b3987e6 100644 --- a/verl_omni/workers/engine/fsdp/diffusers_impl.py +++ b/verl_omni/workers/engine/fsdp/diffusers_impl.py @@ -678,13 +678,8 @@ def optimizer_zero_grad(self): """ self.optimizer.zero_grad() - def optimizer_step(self): - """ - Clip gradients, skip update if non-finite, and step optimizer. - - Returns: - grad_norm (float): Norm of gradients before clipping. - """ + def clip_grad_norm(self): + """Clip gradients using the active FSDP strategy and return the global norm.""" assert self.optimizer_config.clip_grad is not None if isinstance(self.module, FSDP): @@ -699,7 +694,11 @@ def optimizer_step(self): if isinstance(grad_norm, DTensor): grad_norm = grad_norm.full_tensor() - # if grad_norm is not finite, skip the update + return grad_norm + + def optimizer_step(self): + """Clip gradients and step the optimizer only when the norm is finite.""" + grad_norm = self.clip_grad_norm() if not torch.isfinite(grad_norm): print(f"WARN: grad_norm is not finite: {grad_norm}") self.optimizer.zero_grad() @@ -811,15 +810,21 @@ def get_per_tensor_param( peft_model = getattr(self.module, "_fsdp_wrapped_module", self.module) if hasattr(peft_model, "peft_config"): # LoRA if not merge_lora: - peft_config = peft_model.peft_config.get("default", None) - adapter_ctx = self.use_adapter(adapter_name) if adapter_name is not None else nullcontext() + resolved_adapter = adapter_name or "default" + peft_config = peft_model.peft_config.get(resolved_adapter, None) + if peft_config is None: + raise ValueError( + f"Cannot export unknown LoRA adapter {resolved_adapter!r}; " + f"available adapters: {sorted(peft_model.peft_config)}." + ) + adapter_ctx = self.use_adapter(resolved_adapter) with adapter_ctx: params = collect_lora_params( module=self.module, layered_summon=layered_summon, base_sync_done=base_sync_done, is_diffusers=True, - adapter_name=adapter_name or "default", + adapter_name=resolved_adapter, layer_prefixes=self.model_config.fsdp_layer_prefixes, ) else: # merge lora diff --git a/verl_omni/workers/engine/fsdp/dmd_impl.py b/verl_omni/workers/engine/fsdp/dmd_impl.py new file mode 100644 index 000000000..df16651da --- /dev/null +++ b/verl_omni/workers/engine/fsdp/dmd_impl.py @@ -0,0 +1,450 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DMD2 sampling and optimization on the existing Diffusers FSDP lifecycle.""" + +from __future__ import annotations + +import json +import math +import os +import time +from contextlib import nullcontext +from copy import deepcopy +from pathlib import Path + +import torch +from verl.utils import tensordict_utils as tu +from verl.utils.device import get_device_id, get_device_name, get_torch_device +from verl.utils.fsdp_utils import load_fsdp_optimizer, offload_fsdp_optimizer +from verl.utils.metric import Metric +from verl.workers.config.optimizer import build_optimizer +from verl.workers.engine.base import EngineRegistry +from verl.workers.engine.utils import prepare_micro_batches + +from verl_omni.pipelines.model_base import DiffusionModelBase +from verl_omni.trainer.diffusion.distillation.utils import ode_euler_step, standard_cfg, timestep_shift + +from .diffusers_impl import DiffusersFSDPEngine + + +@EngineRegistry.register(model_type="diffusion_dmd_model", backend=["fsdp", "fsdp2"], device=["cuda", "npu"]) +class DMDDiffusersFSDPEngine(DiffusersFSDPEngine): + """One student/fake-score optimizer pair with frozen scoring and adapter EMA. + + The validated storage path is named LoRA adapters on one frozen base. Model + conversion and conditioning are adapter-owned; no image layouts enter here. + """ + + adapter_names = { + "student": "default", + "fake_score": "fake_score", + "teacher_score": "reference", + "student_ema": "student_ema", + } + stream_offsets = {"initial_noise": 0, "rollout_decision": 1, "score_sigma": 3, "score_noise": 4} + + def __init__(self, model_config, engine_config, optimizer_config, checkpoint_config, *, dmd_config): + if model_config.lora_rank <= 0: + raise ValueError("The DMD2 MVP requires LoRA; full-module training is not enabled.") + if engine_config.strategy == "fsdp" and not engine_config.use_orig_params: + raise ValueError("DMD2 shared adapters require FSDP1 use_orig_params=true.") + if engine_config.forward_only: + raise ValueError("A DMD2 engine must own trainable student and fake-score optimizers.") + self.dmd_config = dmd_config + self.active_stage = "student" + self.optimizers = {} + self.schedulers = {} + self.role_parameters = {} + self.optimizer_steps = {"student": 0, "fake_score": 0} + self.skipped_steps = {"student": 0, "fake_score": 0} + self.generators = {} + self.pending_generator_states = {} + self.last_step_succeeded = False + self.forward_finite = True + model_config = deepcopy(model_config) + allowed = {"default", "fake_score", "student_ema", "reference"} + if not set(model_config.policy_state_adapters).issubset(allowed): + raise ValueError( + "DMD2 manages default, fake_score and student_ema adapters; other policy states are unsupported." + ) + object.__setattr__(model_config, "policy_state_adapters", ("default", "fake_score", "student_ema")) + self.model_adapter = DiffusionModelBase.get_class(model_config) + for method in ( + "build_conditioning_provider", + "latent_geometry", + "pack_latents", + "prepare_dmd_inputs", + "prediction_to_x0", + "sampling_sigmas", + ): + if not callable(getattr(self.model_adapter, method, None)): + raise TypeError(f"The selected DMD2 model adapter must implement {method}.") + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + + def initialize(self): + """Reuse model loading and checkpoint management with reproducible LoRA initialization.""" + with torch.random.fork_rng(devices=[get_device_id()], device_type=get_device_name()): + torch.manual_seed(self.engine_config.seed) + super().initialize() + self.copy_adapter("default", "fake_score") + self.copy_adapter("default", "student_ema") + self.select_stage("student") + self.condition_provider = self.model_adapter.build_conditioning_provider(self.model_config, self.dmd_config) + + def _build_optimizer(self, module): + return build_optimizer( + (parameter for parameter in module.parameters() if parameter.requires_grad), self.optimizer_config + ) + + def _build_model_optimizer(self): + super()._build_model_optimizer() + self.optimizers["student"] = self.optimizer + self.schedulers["student"] = self.lr_scheduler + self.optimizer_configs = { + "student": self.optimizer_config, + "fake_score": deepcopy(self.dmd_config.fake_score_optim), + } + self.optimizer_configs["fake_score"].total_training_steps = ( + self.optimizer_config.total_training_steps * self.dmd_config.fake_update_ratio + ) + for stage in ("student", "fake_score"): + with self.use_adapter(self.adapter_names[stage]): + self.role_parameters[stage] = tuple( + parameter for parameter in self.module.parameters() if parameter.requires_grad + ) + if not all(self.role_parameters.values()): + raise ValueError("Every DMD2 optimizer must own parameters.") + if {id(p) for p in self.role_parameters["student"]} & {id(p) for p in self.role_parameters["fake_score"]}: + raise ValueError("Student and fake-score optimizers must not share parameters.") + self.optimizers["fake_score"] = build_optimizer( + self.role_parameters["fake_score"], self.optimizer_configs["fake_score"] + ) + previous = self.optimizer_config + try: + self.optimizer_config = self.optimizer_configs["fake_score"] + self.schedulers["fake_score"] = self._build_lr_scheduler(self.optimizers["fake_score"]) + finally: + self.optimizer_config = previous + + def select_stage(self, stage): + """Select optimizer ownership before entering the ordinary engine train context.""" + if stage not in {"student", "fake_score"}: + raise ValueError(f"Invalid DMD2 update stage {stage!r}.") + self.active_stage = stage + self._set_adapter(self.adapter_names[stage]) + self.optimizer = self.optimizers[stage] + self.lr_scheduler = self.schedulers[stage] + self.optimizer_config = self.optimizer_configs[stage] + + def to(self, device, model=True, optimizer=True, grad=True): + """Reuse base model offload and move both optimizer states when requested.""" + super().to(device=device, model=model, optimizer=False, grad=grad) + if optimizer: + for item in self.optimizers.values(): + if device == "cpu": + offload_fsdp_optimizer(item) + else: + load_fsdp_optimizer(item, device) + + def generator(self, stream, device): + """Use independent, checkpointed streams with the prototype's DP seed spacing.""" + if stream not in self.generators: + generator = torch.Generator(device=device) + generator.manual_seed( + self.engine_config.seed + self.get_data_parallel_rank() * 5 + self.stream_offsets[stream] + ) + if stream in self.pending_generator_states: + generator.set_state(self.pending_generator_states.pop(stream)) + self.generators[stream] = generator + return self.generators[stream] + + def noise(self, shape, device, stream): + """Draw fp32 noise without sharing state with logging or data shuffling.""" + return torch.randn(shape, dtype=torch.float32, device=device, generator=self.generator(stream, device)) + + @staticmethod + def expand_sigma(sigma, latents): + """Broadcast explicit per-sample sigma over the architecture's latent axes.""" + sigma = sigma.reshape(-1) + if sigma.numel() not in (1, latents.shape[0]): + raise ValueError("Sigma batch size does not match the latent batch.") + return sigma.reshape(-1, *((1,) * (latents.ndim - 1))) + + def predict(self, role, latents, sigma, condition, geometry, *, grad_enabled=False): + """Run one role with adapter restoration, without nested engine train/eval contexts.""" + grad_context = nullcontext() if grad_enabled else torch.no_grad() + with self.use_adapter(self.adapter_names[role]), grad_context, torch.profiler.record_function(f"dmd/{role}"): + self.module.eval() + inputs = self.model_adapter.prepare_dmd_inputs( + self.module, self.model_config, latents, sigma, condition, geometry + ) + prediction = self.model_adapter.forward(self.module, self.model_config, inputs) + if prediction.shape != latents.shape: + raise ValueError("DMD2 model prediction must preserve the declared latent shape.") + return prediction + + def student_sample(self, noise, condition, geometry, *, grad_enabled): + """Backward-simulate to a rank-synchronized exit, retaining only its graph.""" + steps = self.model_config.pipeline.num_inference_steps + if not isinstance(steps, int) or isinstance(steps, bool) or steps <= 0: + raise ValueError("num_inference_steps must be a positive integer.") + sigmas = self.model_adapter.sampling_sigmas(self.model_config, self.dmd_config, noise.device) + exit_step = torch.randint( + steps, (1,), device=noise.device, generator=self.generator("rollout_decision", noise.device) + ) + torch.distributed.broadcast(exit_step, src=0) + exit_index = int(exit_step.item()) + sample = noise + for index in range(exit_index + 1): + prediction = self.predict( + "student", sample, sigmas[index], condition, geometry, grad_enabled=grad_enabled and index == exit_index + ) + if index == exit_index: + return self.model_adapter.prediction_to_x0( + sample, prediction, self.expand_sigma(sigmas[index], sample) + ), exit_index + sample = ode_euler_step(sample, prediction, sigmas[index], sigmas[index + 1]) + raise RuntimeError("Student sampling reached no exit prediction.") + + def score_inputs(self, generated): + """Construct detached flow corruption with discrete or continuous sigma sampling.""" + cfg = self.dmd_config + if cfg.score_discrete_steps: + total = self.scheduler.config.num_train_timesteps + if cfg.score_discrete_steps != total: + raise ValueError("score_discrete_steps must equal the model scheduler's num_train_timesteps.") + timestep = torch.randint( + total, + (generated.shape[0],), + device=generated.device, + generator=self.generator("score_sigma", generated.device), + ) + sigma = timestep_shift(timestep, total, cfg.score_timestep_shift) / total + sigma = sigma.clamp(cfg.score_sigma_min, cfg.score_sigma_max) + else: + sigma = torch.rand( + generated.shape[0], device=generated.device, generator=self.generator("score_sigma", generated.device) + ) + sigma = cfg.score_sigma_min + (cfg.score_sigma_max - cfg.score_sigma_min) * sigma + noise = self.noise(generated.shape, generated.device, "score_noise") + expanded = self.expand_sigma(sigma, generated) + return (1 - expanded) * generated.detach().float() + expanded * noise, noise, sigma + + def prepare_model_inputs(self, micro_batch, step=None): + """Delegate geometry and frozen conditioning to the registered architecture.""" + device = torch.device(get_device_id()) + dtype = getattr(self.module, "dtype", torch.bfloat16) + condition, negative = self.condition_provider.encode( + micro_batch, device=device, dtype=dtype, require_negative=self.active_stage == "student" + ) + for item in (condition, negative): + if item is None: + continue + for key, value in item.items(): + item[key] = value.to(device) + if self.use_ulysses_sp: + item["prompt_embeds"], item["prompt_embeds_mask"] = self._pad_embeds_for_sp( + item["prompt_embeds"], item["prompt_embeds_mask"], self.ulysses_sequence_parallel_size + ) + shape, geometry = self.model_adapter.latent_geometry(self.module, self.model_config, micro_batch) + return condition, negative, shape, geometry + + def prepare_model_outputs(self, output, micro_batch): + """Keep graph-bearing student outputs and detached scoring inputs explicit.""" + return output + + def forward_step(self, micro_batch, loss_function, forward_only=False, step=None): + """Compute one DMD2 microbatch through the registered diffusion loss.""" + timings = {} + start = time.perf_counter() + condition, negative, shape, geometry = self.prepare_model_inputs(micro_batch) + timings["perf/condition_encode_s"] = time.perf_counter() - start + initial = self.model_adapter.pack_latents(self.noise(shape, condition["prompt_embeds"].device, "initial_noise")) + start = time.perf_counter() + with torch.profiler.record_function("dmd/student_rollout"): + generated, exit_index = self.student_sample( + initial, condition, geometry, grad_enabled=self.active_stage == "student" and not forward_only + ) + timings["perf/student_rollout_s"] = time.perf_counter() - start + noisy, noise, sigma = self.score_inputs(generated) + start = time.perf_counter() + prediction = self.predict( + "fake_score", + noisy, + sigma, + condition, + geometry, + grad_enabled=self.active_stage == "fake_score" and not forward_only, + ) + timings["perf/fake_score_s"] = time.perf_counter() - start + if self.active_stage == "student": + start = time.perf_counter() + positive = self.predict("teacher_score", noisy, sigma, condition, geometry) + negative_prediction = self.predict("teacher_score", noisy, sigma, negative, geometry) + guided = standard_cfg( + positive, negative_prediction, self.dmd_config.teacher_guidance_scale, self.dmd_config.cfg_norm + ) + timings["perf/teacher_score_s"] = time.perf_counter() - start + expanded = self.expand_sigma(sigma, generated) + output = { + "generated_x0": generated, + "fake_x0": self.model_adapter.prediction_to_x0(noisy, prediction, expanded), + "teacher_x0": self.model_adapter.prediction_to_x0(noisy, guided, expanded), + } + else: + output = {"generated_x0": generated.detach(), "noise_pred": prediction, "noise": noise} + output = self.prepare_model_outputs(output, micro_batch) + loss, metrics = loss_function(model_output=output, data=micro_batch, dp_group=self.get_data_parallel_group()) + metrics.update({"dmd/rollout_exit": float(exit_index), **timings}) + return loss, metrics + + def forward_backward_batch(self, data, loss_function, forward_only=False): + """Use the existing DP-aware microbatch splitter with sample-weighted means.""" + stage = tu.get_non_tensor_data(data, "dmd_stage", default="student") + if stage != self.active_stage: + raise ValueError("Select the DMD2 optimizer before entering its train context.") + self.module.eval() + tu.assign_non_tensor(data, use_dynamic_bsz=False, sp_size=self.ulysses_sequence_parallel_size) + micro_size = tu.get_non_tensor_data(data, "micro_batch_size_per_gpu", default=None) + if len(data) % micro_size: + # verl's wrapper requires divisibility; native TensorDict splitting preserves the dense tail. + micro_batches = data.split(micro_size) + else: + micro_batches, _ = prepare_micro_batches( + data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True + ) + losses, metrics = [], {} + self.forward_finite = True + for micro in micro_batches: + micro = micro.to(get_device_id()) + tu.assign_non_tensor( + micro, + gradient_accumulation_steps=len(data) / len(micro), + dmd_normalization_epsilon=self.dmd_config.normalization_epsilon, + ) + loss, values = self.forward_step(micro, loss_function, forward_only) + loss_value = loss.detach().item() + self.forward_finite = self.forward_finite and math.isfinite(loss_value) + if not forward_only: + start = time.perf_counter() + loss.backward() + values["perf/backward_s"] = time.perf_counter() - start + losses.append(loss_value) + for key, value in values.items(): + summed = key.endswith("_s") or key.endswith("/active_elements") or key.endswith("/nonfinite") + weight = 1.0 if summed else len(micro) / len(data) + value = value.aggregate() if isinstance(value, Metric) else value + metrics[key] = metrics.get(key, 0.0) + float(value) * weight + metrics = Metric.from_dict(metrics, aggregation="mean") + device = get_torch_device() + metrics["perf/max_memory_allocated_gib"] = Metric("max", device.max_memory_allocated() / 1024**3) + metrics["perf/max_memory_reserved_gib"] = Metric("max", device.max_memory_reserved() / 1024**3) + return {"loss": losses, "metrics": metrics, "model_output": {}} + + def optimizer_step(self): + """Agree on numerical skips before stepping, and update only the owning scheduler/EMA.""" + for stage, parameters in self.role_parameters.items(): + if stage != self.active_stage and any(parameter.grad is not None for parameter in parameters): + raise RuntimeError(f"Gradient leaked into inactive DMD2 role {stage}.") + norm = float(self.clip_grad_norm()) + finite = torch.tensor(int(math.isfinite(norm) and self.forward_finite), device=get_device_id()) + torch.distributed.all_reduce(finite, op=torch.distributed.ReduceOp.MIN) + self.last_step_succeeded = bool(finite.item()) + if self.last_step_succeeded: + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer_steps[self.active_stage] += 1 + if self.active_stage == "student" and self.optimizer_steps["student"] >= self.dmd_config.ema_start_step: + self.ema_update_adapter("default", "student_ema", self.dmd_config.ema_decay) + else: + self.skipped_steps[self.active_stage] += 1 + self.optimizer.zero_grad() + return norm + + def lr_scheduler_step(self): + """Schedulers already advance atomically with successful optimizer updates.""" + return self.lr_scheduler.get_last_lr()[0] + + def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None, **kwargs): + """Save the physical model/student optimizer once, plus DMD-specific state.""" + if hdfs_path is not None: + raise ValueError("DMD2 atomic checkpoints currently require a local shared directory.") + previous = self.active_stage + self.select_stage("student") + try: + super().save_checkpoint(local_path, global_step=global_step) + torch.save( + { + "version": 1, + "world_size": torch.distributed.get_world_size(), + "fake_optimizer": self.optimizers["fake_score"].state_dict(), + "fake_scheduler": self.schedulers["fake_score"].state_dict(), + "optimizer_steps": self.optimizer_steps, + "skipped_steps": self.skipped_steps, + "generators": { + **self.pending_generator_states, + **{name: generator.get_state() for name, generator in self.generators.items()}, + }, + }, + Path(local_path) / f"dmd_state_rank_{self.rank}.pt", + ) + torch.distributed.barrier() + finally: + self.select_stage(previous) + + def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False, **kwargs): + """Reject old prototype/incomplete state, then restore both optimization clocks.""" + if hdfs_path is not None or del_local_after_load: + raise ValueError("DMD2 resume preserves its local atomic checkpoint.") + state = torch.load(Path(local_path) / f"dmd_state_rank_{self.rank}.pt", map_location="cpu", weights_only=False) + if state.get("version") != 1 or state.get("world_size") != torch.distributed.get_world_size(): + raise ValueError("Incompatible DMD2 checkpoint version or world size.") + if set(state.get("optimizer_steps", {})) != {"student", "fake_score"}: + raise ValueError("Missing DMD2 optimizer counters.") + self.select_stage("student") + super().load_checkpoint(local_path, del_local_after_load=False) + self.optimizers["fake_score"].load_state_dict(state["fake_optimizer"]) + self.schedulers["fake_score"].load_state_dict(state["fake_scheduler"]) + self.optimizer_steps = state["optimizer_steps"] + self.skipped_steps = state["skipped_steps"] + self.generators.clear() + self.pending_generator_states = state["generators"] + return dict(self.optimizer_steps) + + def export_student(self, directory, role="student"): + """Export only the chosen inference adapter using existing per-tensor collection.""" + if role not in {"student", "student_ema"}: + raise ValueError("Only student or student_ema can be exported.") + adapter = self.adapter_names[role] + with self.use_adapter(adapter): + expected = torch.tensor( + sum(parameter.numel() for parameter in self.module.parameters() if parameter.requires_grad), + device=get_device_id(), + ) + if self.engine_config.strategy == "fsdp": + torch.distributed.all_reduce(expected) + tensors, peft_config = self.get_per_tensor_param(base_sync_done=True, adapter_name=adapter) + state = {name.removeprefix("transformer."): value.detach().cpu().contiguous() for name, value in tensors} + if sum(value.numel() for value in state.values()) != int(expected.item()): + raise ValueError("Incomplete LoRA export: fsdp_layer_prefixes must cover every selected adapter parameter.") + if not state or any(not torch.isfinite(value).all() for value in state.values()): + raise ValueError("Inference export requires nonempty finite adapter weights.") + if self.rank == 0: + from safetensors.torch import save_file + + os.makedirs(directory, exist_ok=False) + save_file(state, str(Path(directory) / "adapter_model.safetensors")) + with open(Path(directory) / "adapter_config.json", "w") as file: + json.dump(peft_config, file, indent=2, default=list) + torch.distributed.barrier() diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index 97d62fb43..6fbdf92c5 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -43,6 +43,12 @@ class OmniFSDPEngine(FSDPEngineWithLMHead): """FSDP engine for omni models""" + @staticmethod + def _cast_dtensor_weight_for_sync(tensor: torch.Tensor) -> torch.Tensor: + if tensor.is_floating_point() and tensor.dtype != torch.bfloat16: + return tensor.to(dtype=torch.bfloat16, non_blocking=True) + return tensor + def prepare_model_inputs(self, micro_batch): """Prepare standard LM inputs, then add model-native replay fields.""" model_inputs, output_args = super().prepare_model_inputs(micro_batch) @@ -97,11 +103,10 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwa per_tensor_param = params.items() else: device = get_device_id() # used when fsdp2 set cpu_offload_policy - # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate per_tensor_param = ( ( name, - param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) + self._cast_dtensor_weight_for_sync(param.to(device, non_blocking=True).full_tensor()) if isinstance(param, DTensor) else param, ) @@ -144,7 +149,7 @@ def _merged_lora_per_tensor_param(self): for name, param in params.items(): yield ( name, - param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) + self._cast_dtensor_weight_for_sync(param.to(device, non_blocking=True).full_tensor()) if isinstance(param, DTensor) else param.detach().clone(), ) @@ -171,6 +176,12 @@ def _build_module(self): self.model_config: OmniModelConfig architecture = self.model_config.architecture + adapter_cls = OmniModelBase.get_class_by_name( + architecture, + self.model_config.model_stage, + self.model_config.get("external_lib"), + ) + self.model_adapter_cls = adapter_cls torch_dtype = self.engine_config.model_dtype @@ -192,21 +203,22 @@ def _build_module(self): with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") - module = AutoModelForMultimodalLM.from_pretrained( + auto_model_cls = getattr(adapter_cls, "auto_model_class", None) or AutoModelForMultimodalLM + module = auto_model_cls.from_pretrained( pretrained_model_name_or_path=self.model_config.local_path, torch_dtype=torch_dtype, config=self.model_config.hf_config, trust_remote_code=self.model_config.trust_remote_code, ) - - adapter_cls = OmniModelBase.get_class_by_name( - architecture, - self.model_config.model_stage, - self.model_config.get("external_lib"), - ) - self.model_adapter_cls = adapter_cls module = adapter_cls.configure_model(module, self.model_config) + if self.engine_config.strategy == "fsdp" and not self.engine_config.use_orig_params: + trainability = {parameter.requires_grad for parameter in module.parameters()} + if len(trainability) > 1: + raise ValueError( + "FSDP1 requires use_orig_params=true when a model adapter freezes only part of the model." + ) + module.to(torch_dtype) if self.model_config.enable_gradient_checkpointing: diff --git a/verl_omni/workers/engine/lora_adapter_mixin.py b/verl_omni/workers/engine/lora_adapter_mixin.py index 2590d9f42..1e63b95f3 100644 --- a/verl_omni/workers/engine/lora_adapter_mixin.py +++ b/verl_omni/workers/engine/lora_adapter_mixin.py @@ -13,8 +13,10 @@ # limitations under the License. """Reusable PEFT/LoRA adapter lifecycle helpers for training engines.""" +import json import logging from contextlib import contextmanager, nullcontext +from pathlib import Path import torch from peft import LoraConfig @@ -23,6 +25,24 @@ logger = logging.getLogger(__name__) +def load_diffusers_lora_adapter(module, path, adapter_name="default"): + """Load a native Diffusers LoRA or a PEFT export with its own scaling metadata.""" + weights_path = Path(path) / "adapter_model.safetensors" + metadata_path = Path(path) / "adapter_config.json" + if weights_path.is_file() and metadata_path.is_file(): + from safetensors.torch import load_file + + weights = { + key.removeprefix("base_model.model.").removeprefix("transformer."): value + for key, value in load_file(weights_path).items() + } + module.load_lora_adapter( + weights, adapter_name=adapter_name, prefix=None, metadata=json.loads(metadata_path.read_text()) + ) + else: + module.load_lora_adapter(path, adapter_name=adapter_name) + + class LoRAAdapterMixin: """Backend-agnostic helpers for named PEFT/LoRA policy adapters.""" @@ -38,7 +58,7 @@ def _build_lora_module(self, module): local_adapter_path = copy_to_local(lora_adapter_path, use_shm=self.model_config.use_shm) # diffusers auto-names an unnamed first adapter "default_0"; name it explicitly. - module.load_lora_adapter(local_adapter_path, adapter_name=primary_adapter) + load_diffusers_lora_adapter(module, local_adapter_path, primary_adapter) peft_config = getattr(module, "peft_config", {}).get(primary_adapter, None) for adapter_name in extra_adapters: if peft_config is not None and adapter_name not in getattr(module, "peft_config", {}): @@ -78,6 +98,25 @@ def _build_lora_module(self, module): return module + def active_adapter_selection(self): + """Read the current named adapter selection for context restoration.""" + module = getattr(self.module, "_fsdp_wrapped_module", self.module) + active = getattr(module, "active_adapters", None) + if callable(active): + active = active() + if active is None: + active = getattr(module, "active_adapter", None) + if isinstance(active, list | tuple): + return active[0] if len(active) == 1 else list(active) + return active + + def restore_adapter_selection(self, selection) -> None: + """Restore the preceding adapter rather than always selecting default.""" + if selection: + self._set_adapter(selection) + else: + self._set_adapter("default") + @contextmanager def _adapter_state_context(self): """Open writable adapter parameter access (FSDP summon when applicable).""" @@ -89,6 +128,7 @@ def _adapter_state_context(self): is_fsdp_module = fsdp_version(self.module) in (1, 2) is_offload_param = getattr(self, "_is_offload_param", False) origin_module_device = next(self.module.parameters()).device.type + previous_adapter = self.active_adapter_selection() if is_fsdp_module and (is_offload_param or origin_module_device == "cpu"): load_fsdp_model_to_gpu(self.module) @@ -98,13 +138,13 @@ def _adapter_state_context(self): try: yield finally: - self._set_adapter("default") + self.restore_adapter_selection(previous_adapter) finally: if is_offload_param: offload_fsdp_model_to_cpu(self.module) aggressive_empty_cache(force_sync=True) - def _set_adapter(self, name: str): + def _set_adapter(self, name): module = getattr(self.module, "_fsdp_wrapped_module", self.module) if not hasattr(module, "set_adapter"): raise AttributeError(f"Module does not support set_adapter({name!r})") @@ -117,6 +157,7 @@ def use_adapter(self, name: str): ``"reference"`` is a logical policy state (see ``policy_state_adapters``) that runs with all LoRA adapters disabled, not a registered PEFT adapter. """ + previous_adapter = self.active_adapter_selection() if name == "reference": with self.disable_adapter(): yield @@ -125,7 +166,7 @@ def use_adapter(self, name: str): try: yield finally: - self._set_adapter("default") + self.restore_adapter_selection(previous_adapter) def _active_adapter_trainable_params(self, adapter_name: str) -> list[torch.nn.Parameter]: peft_model = getattr(self.module, "_fsdp_wrapped_module", self.module) diff --git a/verl_omni/workers/engine_workers.py b/verl_omni/workers/engine_workers.py index ce1979343..e06f0024d 100644 --- a/verl_omni/workers/engine_workers.py +++ b/verl_omni/workers/engine_workers.py @@ -109,8 +109,6 @@ class TrainingWorker(Worker, DistProfilerExtension): def __init__(self, config: TrainingWorkerConfig): Worker.__init__(self) - from verl.workers.engine import BaseEngine, EngineRegistry - # TODO(jhz): Switch to `set_expandable_segments` when the torch_npu library # supports `torch.npu.memory._set_allocator_settings` if is_npu_available: @@ -155,14 +153,7 @@ def __init__(self, config: TrainingWorkerConfig): ) self.model_config.model_type = self.config.model_type - self.engine: BaseEngine = EngineRegistry.new( - model_type=self.config.model_type, - backend=self.engine_config.strategy, - model_config=self.model_config, - engine_config=self.engine_config, - optimizer_config=self.optimizer_config, - checkpoint_config=self.checkpoint_config, - ) + self.engine = self.build_engine() # build dispatch info self._register_dispatch_collect_info( @@ -173,7 +164,12 @@ def __init__(self, config: TrainingWorkerConfig): if getattr(self.model_config, "hf_config", None) is not None: self.flops_counter = FlopsCounter(self.model_config.hf_config) - elif self.config.model_type in ("diffusion_model", "diffusion_dpo_model", "diffusion_nft_model"): + elif self.config.model_type in ( + "diffusion_model", + "diffusion_dpo_model", + "diffusion_nft_model", + "diffusion_dmd_model", + ): self.flops_counter = DiffusionFlopsCounter( architecture=getattr(self.model_config, "architecture", None), transformer_config=getattr(self.model_config, "transformer_config", None), @@ -183,6 +179,19 @@ def __init__(self, config: TrainingWorkerConfig): self.loss_fn = None + def build_engine(self): + """Construct the engine while allowing specialized workers to pass typed settings.""" + from verl.workers.engine import EngineRegistry + + return EngineRegistry.new( + model_type=self.config.model_type, + backend=self.engine_config.strategy, + model_config=self.model_config, + engine_config=self.engine_config, + optimizer_config=self.optimizer_config, + checkpoint_config=self.checkpoint_config, + ) + @register(dispatch_mode=Dispatch.ONE_TO_ALL) def to(self, device, model=True, optimizer=True, grad=True): """Manual control of load/offload""" diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index c82630ccd..e06925223 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy import json import logging import os @@ -56,6 +57,15 @@ class ARStrategy(OmniStrategyBase): rollout_config_cls = RolloutConfig model_config_cls = OmniModelConfig + def __init__(self, server: Any) -> None: + super().__init__(server) + self._rollout_adapter: type[OmniRolloutPipelineBase] | None = None + self._rollout_output_modalities: list[str] | None = None + self._weight_sync_stage_ids: list[int] | None = None + self._rollout_fields_by_request_id: dict[str, dict[str, Any]] = {} + self._policy_stage_index = 0 + self._policy_sampling_constraints: dict[str, Any] = {} + def validate_configs(self) -> None: if self.server.config.max_model_len is None: self.server.config.max_model_len = self.server.config.prompt_length + self.server.config.response_length @@ -78,6 +88,15 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: adapter_cls = OmniRolloutPipelineBase.get_class(pipeline_name) if adapter_cls is not None: + async_chunk = engine_kwargs.get("async_chunk", engine_kwargs.get("async-chunk", True)) + if not isinstance(async_chunk, bool): + raise TypeError(f"async_chunk must be a boolean, got {type(async_chunk).__name__}.") + if async_chunk and not adapter_cls.supports_async_chunk: + raise ValueError( + f"{adapter_cls.__name__} requires async_chunk=false because chunked stage outputs " + "cannot be replayed by its actor adapter." + ) + self._rollout_adapter = adapter_cls self._write_deploy_config(engine_kwargs, pipeline_name, adapter_cls, pipeline_mode) self.server._rollout_flags = adapter_cls.rollout_flags(pipeline_mode=pipeline_mode) adapter_overrides = adapter_cls.get_engine_hf_overrides(pipeline_mode=pipeline_mode) @@ -110,13 +129,59 @@ def _write_deploy_config( """Write a deploy config YAML from the adapter's stage topology.""" adapter_cls.ensure_pipeline_registered(pipeline_mode) stages = adapter_cls.build_stage_configs(pipeline_mode=pipeline_mode) + stage_ids = [stage.stage_id for stage in stages] + policy_stage_id = adapter_cls.policy_stage_id(pipeline_mode=pipeline_mode) + if policy_stage_id not in stage_ids: + raise ValueError( + f"{adapter_cls.__name__}.policy_stage_id() returned unknown stage {policy_stage_id}; " + f"available stages are {stage_ids}." + ) + self._policy_stage_index = stage_ids.index(policy_stage_id) + self._policy_sampling_constraints = dict(stages[self._policy_stage_index].sampling_constraints) + + weight_sync_stage_ids = adapter_cls.weight_sync_stage_ids(pipeline_mode=pipeline_mode) + if weight_sync_stage_ids is not None: + unknown_stage_ids = sorted(set(weight_sync_stage_ids) - set(stage_ids)) + if unknown_stage_ids: + raise ValueError( + f"{adapter_cls.__name__}.weight_sync_stage_ids() returned unknown stages {unknown_stage_ids}; " + f"available stages are {stage_ids}." + ) + self._weight_sync_stage_ids = weight_sync_stage_ids + pipeline_id = adapter_cls.get_pipeline_id(pipeline_mode) + final_output_types = [stage.final_output_type for stage in stages if stage.final_output] + adapter_combiner = getattr(adapter_cls.combine_engine_outputs, "__func__", adapter_cls.combine_engine_outputs) + default_combiner = getattr( + OmniRolloutPipelineBase.combine_engine_outputs, + "__func__", + OmniRolloutPipelineBase.combine_engine_outputs, + ) + self._rollout_output_modalities = ( + list(dict.fromkeys(final_output_types)) + if len(final_output_types) > 1 and adapter_combiner is not default_combiner + else None + ) + stage_extras = { + stage.stage_id: dict(adapter_cls.get_stage_engine_extras(stage.stage_id, pipeline_mode=pipeline_mode)) + for stage in stages + } + capacity_fields = ("max_model_len", "max_num_batched_tokens") + if any(field in extras for extras in stage_extras.values() for field in capacity_fields): + for extras in stage_extras.values(): + for field in capacity_fields: + extras.setdefault(field, getattr(self.server.config, field)) + for field in capacity_fields: + engine_kwargs[field] = None device_control_env = get_visible_devices_keyword() visible_devices = os.environ.get(device_control_env, "") tp_size = self.server.config.tensor_model_parallel_size deploy_dict: dict[str, object] = {"pipeline": pipeline_id} + async_chunk = engine_kwargs.get("async_chunk", engine_kwargs.get("async-chunk")) + if async_chunk is not None: + deploy_dict["async_chunk"] = async_chunk if visible_devices: device_count = len([device for device in visible_devices.split(",") if device.strip()]) @@ -128,7 +193,7 @@ def _write_deploy_config( "devices": devices, "tensor_parallel_size": tp_size, "text_encoder_tp_size": getattr(self.server.config, "text_encoder_tp_size", 1), - "engine_extras": adapter_cls.get_stage_engine_extras(stage_id, pipeline_mode=pipeline_mode), + "engine_extras": stage_extras[stage_id], } for stage_id in stage_ids ] @@ -151,6 +216,10 @@ def _write_deploy_config( engine_kwargs["deploy_config"] = deploy_path def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> None: + if self._rollout_output_modalities is not None: + # The generated per-stage deploy config owns model_stage for + # multi-output pipelines. + engine_args["model_stage"] = None for timeout_key in ("stage_init_timeout", "init_timeout"): timeout_value = getattr(args, timeout_key, None) if timeout_value is not None: @@ -159,6 +228,11 @@ def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> N if isinstance(engine_args.get("compilation_config"), dict): engine_args["compilation_config"] = _drop_none_mapping_values(engine_args["compilation_config"]) + def collective_rpc_stage_ids(self, method: Any) -> list[int] | None: + if method in {"set_pending_lora_peft_config", "update_weights_from_ipc"}: + return self._weight_sync_stage_ids + return None + def preprocess_input( self, prompt_ids: list[int], @@ -170,15 +244,34 @@ def preprocess_input( mm_processor_kwargs: Optional[dict[str, Any]] = None, extra_prompt_ids: Optional[dict[str, list[int]]] = None, negative_extra_prompt_ids: Optional[dict[str, list[int]]] = None, - ) -> tuple[dict[str, Any], SamplingParams]: + ) -> tuple[dict[str, Any], SamplingParams | list[Any]]: if multi_modal_data: processor = getattr(self.server.model_config, "processor", None) if processor is not None and hasattr(processor, "dedup_pad_tokens"): prompt_ids = processor.dedup_pad_tokens(prompt_ids) - max_possible_tokens = self.server.config.max_model_len - len(prompt_ids) + + prompt = None + adapter_prepared_prompt = False + if self._rollout_adapter is not None: + prompt = self._rollout_adapter.prepare_engine_prompt( + prompt_ids=prompt_ids, + model_config=self.server.model_config, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + ) + adapter_prepared_prompt = prompt is not None + if prompt is not None: + if not isinstance(prompt, dict): + raise TypeError(f"An omni rollout adapter must return a dict or None, got {type(prompt).__name__}.") + if "prompt_token_ids" not in prompt: + raise RuntimeError("An adapter-prepared omni prompt must contain prompt_token_ids.") + effective_prompt_ids = prompt["prompt_token_ids"] + else: + effective_prompt_ids = prompt_ids + max_possible_tokens = self.server.config.max_model_len - len(effective_prompt_ids) if max_possible_tokens <= 0: raise ValueError( - f"Prompt length ({len(prompt_ids)}) meets or exceeds the model's maximum context length " + f"Prompt length ({len(effective_prompt_ids)}) meets or exceeds the model's maximum context length " f"({self.server.config.max_model_len}), leaving no space for generation." ) @@ -189,7 +282,7 @@ def preprocess_input( else: max_tokens = min( self.server.config.response_length, - self.server.config.prompt_length + self.server.config.response_length - len(prompt_ids), + self.server.config.prompt_length + self.server.config.response_length - len(effective_prompt_ids), ) max_tokens = max(0, min(max_tokens, max_possible_tokens)) @@ -201,13 +294,29 @@ def preprocess_input( else: sampling_params["logprobs"] = None sampling_params.setdefault("repetition_penalty", getattr(self.server.config, "repetition_penalty", 1.0)) - params = SamplingParams(max_tokens=max_tokens, **sampling_params) + policy_params = SamplingParams(max_tokens=max_tokens, **sampling_params) + if self._rollout_output_modalities is not None: + default_stage_sampling_params = self.server.engine.default_sampling_params_list + if len(default_stage_sampling_params) <= 1 or self._policy_stage_index >= len( + default_stage_sampling_params + ): + raise RuntimeError("A multi-output omni rollout requires per-stage sampling parameters.") + params = list(default_stage_sampling_params) + params[self._policy_stage_index] = copy.copy(params[self._policy_stage_index]) + for field in {"max_tokens", *sampling_params} - self._policy_sampling_constraints.keys(): + setattr(params[self._policy_stage_index], field, getattr(policy_params, field)) + else: + params = policy_params - prompt = {"prompt_token_ids": prompt_ids} + if prompt is None: + prompt = {"prompt_token_ids": prompt_ids} + additional_information = prompt.get("additional_information") + if isinstance(additional_information, dict): + additional_information.setdefault("max_new_tokens", [max_tokens]) if multi_modal_data: - prompt["multi_modal_data"] = multi_modal_data - if mm_processor_kwargs: - prompt["mm_processor_kwargs"] = mm_processor_kwargs + prompt.setdefault("multi_modal_data", multi_modal_data) + if mm_processor_kwargs and not adapter_prepared_prompt: + prompt.setdefault("mm_processor_kwargs", mm_processor_kwargs) return prompt, params async def run_generation( @@ -218,28 +327,55 @@ async def run_generation( lora_request: Optional[LoRARequest], priority: int, ) -> Any: - return await self._collect_last_output( - self.server.engine.generate( - prompt=prompt, - sampling_params_list=params, - request_id=request_id, - lora_request=lora_request, - priority=priority, - ) + generate_kwargs = dict( + prompt=prompt, + sampling_params_list=params, + request_id=request_id, + lora_request=lora_request, + priority=priority, ) - - def process_output(self, final_res: Any, params: SamplingParams, sampling_params: dict[str, Any]) -> TokenOutput: + if self._rollout_output_modalities is not None: + generate_kwargs["output_modalities"] = self._rollout_output_modalities + generator = self.server.engine.generate(**generate_kwargs) + if self._rollout_output_modalities is None: + return await self._collect_last_output(generator) + + outputs = [] + async for output in generator: + outputs.append(output) + if self._rollout_adapter is None: + raise RuntimeError("Retaining multiple stage outputs requires a registered rollout adapter.") + final_res, rollout_fields = self._rollout_adapter.combine_engine_outputs(outputs, prompt) + self._rollout_fields_by_request_id[request_id] = rollout_fields + return final_res + + def process_output( + self, + final_res: Any, + params: SamplingParams | list[Any], + sampling_params: dict[str, Any], + ) -> TokenOutput: if final_res is None: raise RuntimeError("AR mode: vLLM-Omni engine yielded no output for the prompt.") + rollout_fields = {} + if self._rollout_output_modalities is not None: + request_id = final_res.request_id + try: + rollout_fields = self._rollout_fields_by_request_id.pop(request_id) + except KeyError: + raise RuntimeError(f"Missing retained rollout fields for request {request_id!r}.") from None + req_output = getattr(final_res, "request_output", None) or final_res if not req_output.outputs: raise RuntimeError("AR mode expects outputs with token IDs, but got None or empty.") extra_fields = {"global_steps": self.server.global_steps} + extra_fields.update(rollout_fields) token_ids = req_output.outputs[0].token_ids log_probs = None - if params.logprobs is not None: + policy_params = params[self._policy_stage_index] if isinstance(params, list) else params + if policy_params.logprobs is not None: log_probs = [ logprobs[token_ids[index]].logprob for index, logprobs in enumerate(req_output.outputs[0].logprobs) ] diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py index 6601d9836..0a8108577 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py @@ -21,7 +21,6 @@ import ray import torch import vllm_omni.entrypoints.cli.serve -from verl.utils.net_utils import get_free_port from verl.workers.config import RolloutConfig from verl.workers.rollout.replica import RolloutMode, TokenOutput from verl.workers.rollout.utils import run_uvicorn @@ -37,6 +36,7 @@ from vllm_omni.entrypoints.openai.api_server import omni_init_app_state from vllm_omni.lora.request import LoRARequest +from verl_omni.utils.net_utils import get_non_ephemeral_free_port from verl_omni.workers.config import DiffusionModelConfig, DiffusionRolloutConfig, OmniModelConfig from verl_omni.workers.rollout.replica import DiffusionOutput from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ARStrategy @@ -157,8 +157,12 @@ async def run_server(self, args: argparse.Namespace): if engine_args.get("seed") is None: engine_args.pop("seed", None) - diffusion_master_port, diffusion_master_sock = get_free_port("127.0.0.1", with_alive_sock=True) - diffusion_master_sock.close() + # The port stays unbound until vllm-omni's rank-0 DiffusionWorker + # listens on it; see get_non_ephemeral_free_port for why it must not + # come from the ephemeral range. + # TODO (mike): drop once vllm-omni passes a FileStore-backed + # distributed_init_method to its workers instead of env:// MASTER_PORT. + diffusion_master_port = get_non_ephemeral_free_port("127.0.0.1") os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = str(diffusion_master_port) @@ -191,6 +195,22 @@ async def run_headless(self, args: argparse.Namespace): # TODO (mike): support multi node raise NotImplementedError("vLLM-Omni headless mode is not implemented yet.") + async def collective_rpc( + self, + method: Any, + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ): + """Dispatch a shared RPC to the stages selected by the active strategy.""" + await self.engine.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + stage_ids=self._generate_strategy.collective_rpc_stage_ids(method), + ) + # ----------------------------------------------------------------------- # wake_up hook: Omni does not restore KV cache on wake-up # ----------------------------------------------------------------------- diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py index 57bf23922..a5c396c9c 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py @@ -56,8 +56,8 @@ class OmniStrategyBase(ABC): * optionally override the concrete hooks (:meth:`init_config`, :meth:`init_model_config`, :meth:`validate_configs`, :meth:`post_init`, :meth:`apply_quantization`, :meth:`override_generation_config`, - :meth:`preprocess_engine_kwargs`) when the mode needs behavior beyond the - shared defaults. + :meth:`preprocess_engine_kwargs`, :meth:`collective_rpc_stage_ids`) when + the mode needs behavior beyond the shared defaults. The two concrete subclasses are :class:`~verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy.ARStrategy` @@ -155,6 +155,14 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: """ engine_kwargs.pop("output_mode", None) + def collective_rpc_stage_ids(self, method: Any) -> list[int] | None: + """Return pipeline stages targeted by a shared collective RPC. + + ``None`` preserves the engine default of broadcasting to every stage. + AR adapters can narrow actor weight synchronization to trainable stages. + """ + return None + @abstractmethod def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> None: """Mutate ``engine_args`` in place with mode-specific engine arguments.