diff --git a/docs/algo/distribution_matching.md b/docs/algo/distribution_matching.md new file mode 100644 index 000000000..e6ef0de2b --- /dev/null +++ b/docs/algo/distribution_matching.md @@ -0,0 +1,145 @@ +# Distribution-Matching Distillation + +Last updated: 09/05/2026. + +verl-omni supports Qwen-Image training with DMD and the distribution-only +profile of DMD2. The implementation follows the multi-role design in +[RFC #519](https://github.com/verl-project/verl-omni/issues/519): a trainable +student, a frozen real-score teacher, a trainable fake-score model, and a +student EMA are coordinated by the dedicated `distillation` trainer. + +This path is not policy-gradient RL and does not use a vLLM-Omni rollout to +build the differentiable student graph. Sampling for training runs inside the +FSDP/FSDP2 worker with `algorithm.sample_source=offline`. + +## Objective + +For clean student output $x_g$, independent noise $\epsilon$, and noise level +$\sigma$, Qwen-Image uses the rectified-flow convention + +$$ +x_\sigma = (1-\sigma)x_g + \sigma\epsilon, +\qquad v_{target} = \epsilon - x_g, +\qquad \hat{x}_0 = x_\sigma - \sigma v. +$$ + +The fake and teacher velocity predictions are converted to canonical clean +predictions. The student receives the detached normalized score difference + +$$ +g = \frac{\hat{x}_{0,fake}-\hat{x}_{0,real}} + {\max(\operatorname{mean}(|x_g-\hat{x}_{0,real}|),\epsilon_{norm})} +$$ + +through the surrogate + +$$ +L_{DMD} = \frac{1}{2}\operatorname{MSE} + \left(x_g,\operatorname{stopgrad}(x_g-g)\right). +$$ + +The normalizer covers every non-batch latent dimension and is not masked. All +score arithmetic and loss reductions run in fp32. The fake-score update uses +MSE against the detached target $\epsilon-x_g$. + +The real teacher uses standard CFG, +`uncond + scale * (cond - uncond)`, with an explicit negative condition. The +fake score and distilled student use one conditional forward by default. + +## Qwen-Image rollout + +The Qwen adapter preserves the checkpoint's native conventions: + +- normalized 5-D VAE latents are packed into Qwen image tokens; +- the transformer receives timesteps in `[0, 1]`; +- its velocity is `noise - x0`; +- the few-step rollout uses deterministic Euler transitions; +- all ranks sharing FSDP collectives use one broadcast rollout exit step; +- only that exit step retains a student autograd graph; +- all preceding rollout steps execute under `no_grad`; +- student, teacher, and fake-score forwards remain in evaluation mode. + +The default four-step schedule applies the reference linear shift `3.0`, giving +sigmas `[1.0, 0.9, 0.75, 0.5, 0.0]`. Score timesteps are sampled discretely +from the model's 1,000 training timesteps, shifted once, and clamped to +`[0.02, 0.98]`. With `score_discrete_steps=0`, sigma is instead sampled +uniformly inside the configured bounds without timestep shifting. Sampling +from `[0, 1)` and clamping is not equivalent to that continuous distribution. + +Raw prompts use one text-only user message. Conditioning applies the fixed +Qwen pipeline template before the encoder removes its 34-token prefix; a +generic chat template can leave short prompts with no encoded tokens. + +A registered vLLM-Omni adapter uses the same sigma construction for +non-autograd inference. It requires deterministic sampling (`noise_level=0`) +and defaults to no inference CFG. A non-default training shift must also be set +as `actor_rollout_ref.rollout.algo.rollout_timestep_shift`. Request batching +reuses the vLLM-Omni scheduler and request collation with an explicit capability +flag. Use `step_execution=false`; unvalidated stepwise DMD execution is rejected. + +## DMD and DMD2 profiles + +`recipe=dmd2`, `profile=distribution_only` is the recommended first runnable +configuration. One student update is followed by `fake_update_ratio` +fake-score updates. It needs prompt conditioning only. + +`recipe=dmd` adds paired trajectory regression. Each sample must provide: + +- `reference_noise`; +- exactly one of `teacher_target_latents` or normalized `[0, 1]` + `teacher_target_pixels`; +- a non-empty `teacher_sampling_manifest`; +- prompt conditioning. + +`regression_type=decoded_lpips` decodes normalized Qwen latents through the +frozen checkpoint VAE and applies PIQ LPIPS. It is the paper-oriented mode and +requires the `distillation` dependency extra. `regression_type=latent_mse` is a +non-paper diagnostic variant. + +The DMD2 adversarial classifier profile is not part of this integration and +fails closed. It belongs to the later adversarial-runtime stage. + +## Role storage and checkpoints + +The recommended LoRA layout stores `student`, `fake_score`, and `student_ema` +as named adapters over one frozen Qwen base. `teacher_score` disables adapters. +Student and fake-score optimizers and schedulers are independent, and EMA is +updated only after a successful student optimizer step. FSDP1 requires +`use_orig_params=true`; FSDP2 is the recommended backend. + +Composite checkpoints save the physical model once together with every role's +optimizer and scheduler, EMA state, phase-runner RNG streams, control-plane +counters, dataloader state, and driver RNG. Only the semantic `student` or +`student_ema` role can be exported to inference. + +## Configuration + +The minimal routing fields are: + +```bash +algorithm.trainer_type=distillation +algorithm.sample_source=offline +actor_rollout_ref.model.algorithm=dmd2 +distillation.enabled=false +distillation.distribution_matching.recipe=dmd2 +distillation.distribution_matching.profile=distribution_only +``` + +Same-resolution physical batches and sample-weighted gradient accumulation use +the existing worker loop, with independent student/fake micro-batch sizes. Each +original-DMD sample must retain its own regression provenance. One physical batch +shares a synchronized rollout exit, so batch-size comparisons must account for +changes in sampled forward counts. + +Profiling uses the existing `DistProfiler` and `Tracking` components. Metrics keep +every repeated phase, reset at cycle boundaries, sum elapsed host intervals and +counts, and preserve peak memory across ranks. The driver reports cycle wall time +and separate student/fake sample throughput; nested host intervals must not be +summed into a GPU-time estimate. Logging steps include warmup cycles, while +`training/global_step` retains student-update semantics. See the example README +for metric definitions and controlled profiling commands. + +The runnable LoRA recipe is +`examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh`. +See its adjacent README for data fields, installation, and the complete launch +command. diff --git a/docs/examples/qwen_image/distillation_trainer_qwen_image.md b/docs/examples/qwen_image/distillation_trainer_qwen_image.md new file mode 120000 index 000000000..3df255db8 --- /dev/null +++ b/docs/examples/qwen_image/distillation_trainer_qwen_image.md @@ -0,0 +1 @@ +../../../examples/distillation_trainer/qwen_image/README.md \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index fca0839a9..c361ca529 100644 --- a/docs/index.md +++ b/docs/index.md @@ -59,6 +59,7 @@ algo/flowgrpo.md algo/flowdppo.md algo/diffusion_dpo.md algo/diffusionnft.md +algo/distribution_matching.md algo/grpo_guard.md algo/mixgrpo.md algo/diffusion_opd.md @@ -81,6 +82,7 @@ examples/diffusionopd_trainer.md examples/flowgrpo_trainer_sd35_drm.md examples/bagel/flowgrpo_trainer_bagel.md examples/qwen_image_edit/flowgrpo_trainer_qwen_image_edit.md +examples/qwen_image/distillation_trainer_qwen_image.md examples/ltx2/flowgrpo_trainer_ltx2.md examples/minimax_h3/diffusionnft_trainer_minimax_h3.md examples/boogu_image/flowgrpo_trainer_boogu_image.md diff --git a/examples/distillation_trainer/qwen_image/README.md b/examples/distillation_trainer/qwen_image/README.md new file mode 100644 index 000000000..14ae435f3 --- /dev/null +++ b/examples/distillation_trainer/qwen_image/README.md @@ -0,0 +1,123 @@ +# Qwen-Image DMD/DMD2 + +Last updated: 09/07/2026. + +This example trains a few-step Qwen-Image generator with the distribution-matching runtime introduced by RFC #519. +Sampling is **offline and differentiable inside the FSDP actor**; it does not use vLLM-Omni for the training rollout. + +## Install + +```bash +uv pip install -e ".[gpu]" --torch-backend=auto +uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" +uv pip install -e ".[train,distillation]" +``` + +The `distillation` extra installs `piq`, which is required only by the paper-oriented original-DMD `decoded_lpips` regression profile. Distribution-only DMD2 does not import it. + +## Data + +The DMD2 distribution-only recipe accepts the normal prompt parquet contract. Each row needs `prompt=[{"role": "user", "content": "..."}]`. The standard `RLHFDataset` passes this as `raw_prompt`; the Qwen adapter applies the checkpoint pipeline's fixed template before encoding and removing its 34-token prefix. Plain strings are also supported; custom system messages and multi-turn chats require precomputed conditioning. + +For PickScore / Pick-a-Pic prompts, use the existing SFW converter: + +```bash +python examples/flowgrpo_trainer/data_process/sd3_pickscore_sfw.py \ + --dataset CarperAI/pickapic_v1_no_images_training_sfw \ + --output-dir ~/data/pickscore_sfw/qwen_image +``` + +To preserve the upstream Flow-GRPO split, download its `dataset/pickscore_sfw/train.txt` and `test.txt`, then pass their directory with `--input-dir`. The converter records source and split metadata. Set `TRAIN_FILES` and `VAL_FILES` to the resulting parquet files. DMD2 uses only these prompts, **not a PickScore reward model**; prompt-only data does not satisfy original DMD's paired-regression requirements. + +Precomputed conditioning is also supported. Set: + +```bash +distillation.distribution_matching.conditioning_provider=precomputed +``` + +and provide `prompt_embeds`, `prompt_embeds_mask`, `negative_prompt_embeds`, and `negative_prompt_embeds_mask` tensors. + +Original DMD additionally requires paired `reference_noise`, either `teacher_target_latents` or normalized `[0, 1]` `teacher_target_pixels`, and a non-empty `teacher_sampling_manifest`. Use `data.custom_cls.path=pkg://verl_omni.utils.dataset.qwen_image_distillation_dataset` and `data.custom_cls.name=QwenImageDMDPairDataset` to convert inline arrays, serialized tensor bytes, or absolute `.pt` paths into fp32 tensors. Its default `decoded_lpips` regression is paper-oriented; `latent_mse` is available only as a non-paper diagnostic. + +## Run DMD2 + +```bash +MODEL_PATH=/path/to/Qwen-Image \ +TRAIN_FILES=/path/to/train.parquet \ +VAL_FILES=/path/to/val.parquet \ +NUM_GPUS=8 \ +bash examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh +``` + +The reference-aligned defaults are four student denoising steps, rollout and score-noise time shifts of `3.0`, score sigma range `[0.02, 0.98]`, teacher CFG `4.0` with per-token norm preservation, student LR `1e-4`, fake-score LR `2e-5`, and two fake-score updates after each student update. + +Physical micro-batches support multiple samples at the same resolution. The existing worker splits each rank-local batch, sample-weights gradients (including a smaller final micro-batch), and steps each role optimizer once. Student and fake-score micro-batch sizes are independent. Original DMD additionally requires one provenance manifest per sample. Mixed-resolution micro-batches fail closed. + +For eight data-parallel ranks, compare these configurations at **the same effective batch of 16**: + +```bash +# Accumulation: two physical micro-batches of one per rank. +bash examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + data.train_batch_size=16 \ + distillation.distribution_matching.student_micro_batch_size_per_gpu=1 \ + distillation.distribution_matching.fake_score_micro_batch_size_per_gpu=1 + +# Physical batching: one micro-batch of two per rank. +bash examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + data.train_batch_size=16 \ + distillation.distribution_matching.student_micro_batch_size_per_gpu=2 \ + distillation.distribution_matching.fake_score_micro_batch_size_per_gpu=2 +``` + +Pass one set of overrides to the launch script above. Rollout exit decisions are broadcast across the training group: FSDP shards must execute identical forward counts and gradient exits, even though their prompts and sample noise differ. A physical batch shares one exit; accumulation samples an exit for each micro-batch. Consequently, compare forward counts as well as wall time rather than attributing random exit-depth differences to batching. The default `layer_norm` CFG is sample-separable; optional `scalar` CFG uses a batch-wide norm, so changing physical batch size also changes that normalization. Shared-base FSDP1 additionally requires `use_orig_params=true`; the script uses FSDP2. + +Only `student` or `student_ema` is exportable. Teacher and fake-score parameters remain training-only state. The registered vLLM-Omni `dmd`/`dmd2` rollout adapter uses the same fixed-shift sigma schedule, fp32 initial noise, deterministic sampling (`noise_level=0`) and no inference CFG by default. It accepts `rollout_timestep_shift` through request `extra_args` when a non-default training shift is used. + +## Request batching + +The adapter explicitly advertises `supports_request_batch=True` and uses the existing vLLM-Omni request scheduler, request-local generators, prompt collation, denoising loop, and output splitting. Requests must share compatible sampling parameters and the DMD shift. Seeds and prompt lengths may differ; multiple images per request are supported. Use `step_execution=false`: stepwise continuous batching has not been validated for this DMD adapter and is rejected rather than silently using FlowGRPO defaults. + +Start inference with a small `actor_rollout_ref.rollout.max_num_seqs` (for example, `2`) and measure memory before increasing it; multiple images per request further increase the effective tensor batch. + +Request batching concerns non-autograd inference, **not the offline FSDP training rollout**. Export APIs do not by themselves wire validation replicas into the distillation trainer. The tiny-GPU test covers native request scheduling, variable prompt lengths, seeds, multiple images, and serial-versus-packed trajectories. It does not establish real-model inference throughput or numerical parity between the training and inference backends. + +## Profiling and metric semantics + +Reuse the existing `DistProfiler` and `Tracking` backends. Run a short job in a separate output directory after other training finishes; keep the same checkpoint, effective batch, schedule, seeds, and hardware for each comparison. Example profiling overrides to the launch script: + +```bash +bash examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + trainer.total_training_steps=6 \ + trainer.resume_mode=disable trainer.save_freq=-1 \ + trainer.logger='[console,tensorboard]' \ + global_profiler.tool=torch 'global_profiler.steps=[3,4]' \ + global_profiler.save_path=/path/to/profile/traces \ + actor_rollout_ref.actor.profiler.enable=true \ + 'actor_rollout_ref.actor.profiler.ranks=[0,1]' \ + 'actor_rollout_ref.actor.profiler.tool_config.torch.contents=[cuda]' +``` + +Both the global step selection and worker profiler enable/rank selection are required. The Torch trace includes `distillation/condition_encode`, `student_rollout`, role forwards, `backward`, role optimizers, `ema`, and original-DMD `regression` ranges. Profiled runs measure attribution, not clean throughput: compare unprofiled warm runs for speed, and report other GPU workloads. + +- `perf/cycle_s`: driver wall time for the entire completed cycle, including data fetch, all phase RPCs and any checkpoint; excludes trace export and logging. +- `perf/*_s`: component **host** durations summed across micro-batches and repeated phases, averaged across DP ranks. These ranges overlap; do not sum them to derive cycle time or GPU kernel time. +- `perf_max_rank/*_s`: sums of the corresponding per-phase slowest-rank durations, not an independently measured cycle critical path. +- `phase/[/]/...`: individual phase metrics, so no fake-score update is overwritten. Unprefixed losses are means across phase updates; element/nonfinite counts are summed across micro-batches and phases, then DP-averaged. +- `memory/max_*`: maximum across data-parallel replicas and phases, rather than average memory usage. +- `training/_samples`: actual global samples processed this cycle (reused samples count as processing again). `perf/_samples_per_s` divides that count by cycle time. Fake updates do not inflate student throughput. +- `batch/_micro_batches`: total micro-batches per rank this cycle; `training/_optimizer_steps`: cumulative successful updates. +- Checkpoint timing is emitted only on checkpoint cycles. Tracking uses `completed_cycles` as its monotonic logging step, including warmup; `training/global_step` remains the completed student-update counter. + +Choose tuning targets from the trace. Compare batching at fixed effective batch, then independently test the existing `actor_rollout_ref.model.enable_gradient_checkpointing` and `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` settings, monitoring memory and repeating correctness tests. + +On supported Hopper GPUs, the existing Diffusers/Kernels FA3 Hub backend is another option; it does not require implementing a new attention kernel in this repository. Configure both sides to satisfy the existing attention-consistency validator, even when the current run only samples offline: + +```bash +bash examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh \ + actor_rollout_ref.model.attn_backend=_flash_3_varlen_hub \ + actor_rollout_ref.rollout.rollout_attn_backend=FLASH_ATTN_3_HUB +``` + +The Hub kernel must be available locally or downloadable. `tests/pipelines/test_qwen_image_dmd_request_batch.py` checks masked/unmasked FA3 forward/backward against native attention as well as the real inference request path. `tests/workers/test_distillation_fsdp_roles.py` covers multi-rank phase execution for DMD/DMD2 at physical batch sizes one and two. Keep backend/precision tolerance checks separate from throughput benchmarks. + +Do not change the fake-update ratio or denoising schedule and label it an implementation speedup. These optimizations are opt-in; the example does not automatically disable checkpointing or resharding. diff --git a/examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh b/examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh new file mode 100755 index 000000000..417bdb82c --- /dev/null +++ b/examples/distillation_trainer/qwen_image/run_qwen_image_dmd2_lora.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# Qwen-Image four-step DMD2 LoRA training with colocated student, teacher, fake-score, and EMA roles. +set -xeuo pipefail + +WORKSPACE=${WORKSPACE:-$HOME} +MODEL_PATH=${MODEL_PATH:-Qwen/Qwen-Image} +TRAIN_FILES=${TRAIN_FILES:-${WORKSPACE}/data/ocr/qwen_image/train.parquet} +VAL_FILES=${VAL_FILES:-${WORKSPACE}/data/ocr/qwen_image/test.parquet} +NUM_GPUS=${NUM_GPUS:-8} + +python3 -m verl_omni.trainer.main_diffusion \ + data.train_files=${TRAIN_FILES} \ + data.val_files=${VAL_FILES} \ + data.train_batch_size=${NUM_GPUS} \ + data.max_prompt_length=1024 \ + algorithm.trainer_type=distillation \ + algorithm.sample_source=offline \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.model.algorithm=dmd2 \ + actor_rollout_ref.model.model_type=diffusion_distillation_model \ + 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.rollout.algo.noise_level=0.0 \ + actor_rollout_ref.rollout.algo.rollout_timestep_shift=3.0 \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.optim.lr=1e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.001 \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=false \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ + distillation.enabled=false \ + distillation.distribution_matching.recipe=dmd2 \ + distillation.distribution_matching.profile=distribution_only \ + distillation.distribution_matching.role_storage=shared_base_adapters \ + distillation.distribution_matching.conditioning_provider=local_frozen_encoder \ + distillation.distribution_matching.fake_update_ratio=2 \ + distillation.distribution_matching.student_micro_batch_size_per_gpu=1 \ + distillation.distribution_matching.fake_score_micro_batch_size_per_gpu=1 \ + distillation.distribution_matching.fake_score_optim.lr=2e-5 \ + distillation.distribution_matching.fake_score_optim.weight_decay=0.001 \ + distillation.distribution_matching.teacher_guidance_scale=4.0 \ + distillation.distribution_matching.teacher_cfg_norm=layer_norm \ + 'distillation.distribution_matching.negative_prompt=" "' \ + distillation.distribution_matching.rollout_timestep_shift=3.0 \ + distillation.distribution_matching.score_timestep_shift=3.0 \ + distillation.distribution_matching.score_sigma_min=0.02 \ + distillation.distribution_matching.score_sigma_max=0.98 \ + distillation.distribution_matching.score_discrete_steps=1000 \ + distillation.distribution_matching.ema_decay=0.999 \ + trainer.logger='["console","wandb"]' \ + trainer.project_name=qwen-image-distillation \ + trainer.experiment_name=qwen-image-dmd2-lora \ + trainer.n_gpus_per_node=${NUM_GPUS} \ + trainer.nnodes=1 \ + trainer.val_before_train=false \ + trainer.test_freq=-1 \ + trainer.save_freq=100 \ + trainer.total_training_steps=1000 \ + "$@" diff --git a/pyproject.toml b/pyproject.toml index 39af3a4f1..73056e530 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,9 @@ gpu = [ train = [ "verl @ git+https://github.com/verl-project/verl.git@fefb080262e1c015a0ea05f958822a6a512dc795", ] +distillation = [ + "piq==0.8.0", +] dev = [ "pytest", "pytest-cov", diff --git a/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh b/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh index aa732a099..eaaca8256 100644 --- a/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh +++ b/tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # ci-e2e-diffusion GPU smoke tests (4-GPU): end-to-end diffusion training paths. -# Includes FlowGRPO / online DPO / DiffusionNFT (v0), synchronous separate, -# FlowGRPO v1 separate_async, and OPD teachers on v0 and the v1 sync trainer. +# Includes FlowGRPO / online DPO / DiffusionNFT / DMD2 (v0), synchronous +# separate, FlowGRPO v1 separate_async, and OPD teachers on v0 and v1. set -euo pipefail @@ -65,4 +65,8 @@ run_test 10 "Diffusion OPD v1 sync standalone teachers e2e" \ env CUDA_VISIBLE_DEVICES="${CUDA_DEVICE_LIST}" NUM_GPUS="${NUM_GPUS}" V1=1 SMOKE=standalone \ bash tests/special_e2e/run_diffusion_teacher_smoke.sh +run_test 11 "Qwen-Image DMD2 distillation 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_image_distillation_adapter_on_cpu.py b/tests/pipelines/test_qwen_image_distillation_adapter_on_cpu.py new file mode 100644 index 000000000..8f4e82c3d --- /dev/null +++ b/tests/pipelines/test_qwen_image_distillation_adapter_on_cpu.py @@ -0,0 +1,592 @@ +# 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 contract tests for the Qwen-Image DMD/DMD2 computer.""" + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from tensordict import NonTensorData, TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.pipelines.model_base import DiffusionModelBase, DistributionMatchingModelAdapter +from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import ( + QwenImageConditionProvider, + QwenImageDistributionMatching, + QwenImageDMDComputer, + build_qwen_dmd_sigmas, +) +from verl_omni.pipelines.qwen_image_distillation.vllm_omni_rollout_adapter import QwenImageDMDPipeline +from verl_omni.pipelines.schedulers import FlowMatchSDEDiscreteScheduler +from verl_omni.trainer.diffusion.distillation.contracts import PhaseRequest +from verl_omni.trainer.diffusion.distillation.recipes import build_plan +from verl_omni.trainer.diffusion.distillation.utils import ode_euler_step + + +class ToyQwenTransformer(torch.nn.Module): + def __init__(self, scale: float, *, trainable: bool = True) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(scale), requires_grad=trainable) + self.config = SimpleNamespace(in_channels=4, guidance_embeds=False) + + def forward(self, hidden_states, encoder_hidden_states, **kwargs): + del kwargs + condition = encoder_hidden_states.float().mean(dim=(1, 2)).reshape(-1, 1, 1).to(hidden_states.dtype) + return (hidden_states * self.scale + condition,) + + +class ToyEngine: + def __init__(self, module: torch.nn.Module) -> None: + self.module = module + self.scheduler = SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])) + + def get_data_parallel_rank(self) -> int: + return 0 + + +class ToyRuntime: + def __init__(self) -> None: + self.modules = { + "student": ToyQwenTransformer(0.2), + "fake_score": ToyQwenTransformer(0.4), + "teacher_score": ToyQwenTransformer(0.7, trainable=False), + } + self.engines = {role: ToyEngine(module) for role, module in self.modules.items()} + + def engine_for_role(self, role: str): + return self.engines[role] + + def scheduler_for_role(self, role: str): + return self.engines[role].scheduler + + @contextmanager + def use_role(self, role: str, *, grad_enabled=None): + module = self.modules[role] + enabled = bool(grad_enabled and module.scale.requires_grad) + with torch.set_grad_enabled(enabled): + yield module + + +def model_config(algorithm: str = "dmd2"): + return SimpleNamespace( + architecture="QwenImagePipeline", + algorithm=algorithm, + external_lib=None, + path="/unused", + local_path="/unused", + transformer_config={"in_channels": 4}, + pipeline=SimpleNamespace( + height=16, + width=16, + num_inference_steps=2, + max_sequence_length=8, + guidance_scale=None, + ), + ) + + +def make_plan(name: str = "dmd2", **overrides): + config = { + "model_path": "/unused", + "conditioning_provider": "precomputed", + "rng_seed": 17, + **overrides, + } + return build_plan(name, config, frozenset({"distribution_matching"})) + + +def phase_batch(batch_size: int = 1, *, regression: bool = False) -> TensorDict: + data = { + "dummy_tensor": torch.zeros(batch_size, 1), + "prompt_embeds": torch.ones(batch_size, 2, 3), + "prompt_embeds_mask": torch.ones(batch_size, 2, dtype=torch.long), + "negative_prompt_embeds": torch.zeros(batch_size, 2, 3), + "negative_prompt_embeds_mask": torch.ones(batch_size, 2, dtype=torch.long), + } + batch = tu.get_tensordict(data) + if regression: + batch["reference_noise"] = torch.full((batch_size, 1, 4), 0.25) + batch["teacher_target_latents"] = torch.zeros(batch_size, 1, 4) + tu.assign_non_tensor(batch, teacher_sampling_manifest={"scheduler": "fixture"}) + return batch + + +def phase_request(kind: str) -> PhaseRequest: + role = "student" if kind == "student" else "fake_score" + return PhaseRequest( + kind=kind, + global_step=0, + repeat_index=0, + batch_policy="fresh", + trainable_roles=(role,), + update_ema=kind == "student", + ) + + +class ToyPromptTokenizer: + def __init__(self): + self.rendered = [] + + 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 ToyVAE(torch.nn.Module): + config = SimpleNamespace(z_dim=1, latents_mean=[0.0], latents_std=[1.0]) + + def decode(self, latent, return_dict=False): + return (latent.repeat(1, 3, 1, 1, 1),) + + +class ToyLPIPS(torch.nn.Module): + def forward(self, prediction, target): + return (prediction - target).square().flatten(1).mean(1) + + +def broadcast_selected_exit(value, src): + assert src == 0 + value.fill_(1) + + +def constant_noise(shape, device, runtime, stream): + return torch.full(shape, 0.3, device=device) + + +def constant_score_sigma(generated, runtime): + return torch.full((generated.shape[0],), 0.5, device=generated.device) + + +def capture_qwen_forward(self, req, *args, **kwargs): + return self.rollout_timestep_shift, kwargs + + +class TestQwenImageDistillationRegistry: + @pytest.mark.parametrize("algorithm", ["dmd", "dmd2"]) + def test_registered_for_distribution_matching_algorithms(self, algorithm): + config = model_config(algorithm) + adapter = DiffusionModelBase.get_class(config) + assert adapter is QwenImageDistributionMatching + assert issubclass(adapter, DistributionMatchingModelAdapter) + + @pytest.mark.parametrize("algorithm", ["dmd", "dmd2"]) + def test_rollout_adapter_registered_for_inference(self, algorithm): + from verl_omni.pipelines.model_base import VllmOmniPipelineBase + + assert VllmOmniPipelineBase.get_class("QwenImagePipeline", algorithm) is QwenImageDMDPipeline + + +class TestQwenImageConditionProvider: + @pytest.mark.parametrize("description", ["cat", "a red apple", ""]) + def test_single_user_chat_uses_the_same_qwen_prefix_as_plain_text(self, description): + tokenizer = Mock(return_value=SimpleNamespace(input_ids=torch.ones(1, 40), attention_mask=torch.ones(1, 40))) + pipeline = SimpleNamespace( + tokenizer=tokenizer, + prompt_template_encode="fixed-qwen-system:{}:assistant", + prompt_template_encode_start_idx=34, + ) + provider = QwenImageConditionProvider("/unused", "local_frozen_encoder", 1024, " ") + for row in (description, [{"role": "user", "content": description}]): + provider.tokenize_rows(pipeline, [row], torch.device("cpu")) + assert tokenizer.call_args.args[0] == [pipeline.prompt_template_encode.format(description)] + tokenizer.apply_chat_template.assert_not_called() + + @pytest.mark.parametrize( + "row", [[], [{"role": "system", "content": "custom"}], [{"role": "assistant", "content": "cat"}]] + ) + def test_unsupported_chat_fails_before_encoding(self, row): + pipeline = Mock() + provider = QwenImageConditionProvider("/unused", "local_frozen_encoder", 1024, " ") + with pytest.raises(ValueError, match="single user message"): + provider.tokenize_rows(pipeline, [row], torch.device("cpu")) + pipeline.tokenizer.assert_not_called() + + @pytest.mark.parametrize("chat", [False, True]) + def test_cached_negative_condition_expands_to_each_physical_batch_size(self, chat): + provider = QwenImageConditionProvider("/unused", "local_frozen_encoder", 8, " ") + provider.pipeline = ToyConditionPipeline() + provider.encode_ids = Mock(wraps=provider.encode_ids) + for batch_size in (1, 3, 2, 1): + batch = tu.get_tensordict({"dummy_tensor": torch.zeros(batch_size, 1)}) + prompt = [{"role": "user", "content": "cat"}] if chat else "cat" + tu.assign_non_tensor_stack(batch, "raw_prompt", [prompt] * batch_size) + positive, negative = provider.encode( + batch, device=torch.device("cpu"), dtype=torch.float32, require_negative=True + ) + assert positive.tensors["prompt_embeds"].shape[0] == batch_size + assert negative.tensors["prompt_embeds"].shape[0] == batch_size + assert negative.masks["prompt_embeds"].shape[0] == batch_size + assert not negative.tensors["prompt_embeds"].requires_grad + assert provider.encode_ids.call_count == 5 + + def test_raw_chat_is_rendered_once_and_negative_prompt_is_encoded(self): + provider = QwenImageConditionProvider("/unused", "local_frozen_encoder", 8, " ") + provider.pipeline = ToyConditionPipeline() + batch = tu.get_tensordict(tensor_dict={"dummy_tensor": torch.zeros(1, 1)}) + tu.assign_non_tensor_stack( + batch, + "raw_prompt", + [[{"role": "user", "content": "prompt"}]], + ) + + positive, negative = provider.encode( + batch, + device=torch.device("cpu"), + dtype=torch.float32, + require_negative=True, + ) + + assert provider.pipeline.tokenizer.rendered == [] + assert positive.tensors["prompt_embeds"].shape[0] == 1 + assert negative.tensors["prompt_embeds"].shape[0] == 1 + assert positive.tensors["prompt_embeds"].shape[1] == len("prompt") + assert positive.tensors["prompt_embeds"][0, 0, 0].item() == 34 + len("prompt") + assert negative.tensors["prompt_embeds"][0, 0, 0].item() == 35 + + +class TestQwenImageDMDComputer: + def test_rollout_exit_is_broadcast_across_sharded_and_sequence_parallel_ranks(self, monkeypatch): + runtime = ToyRuntime() + computer = QwenImageDMDComputer(model_config(), make_plan()) + broadcast = Mock(side_effect=broadcast_selected_exit) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "broadcast", broadcast) + assert computer.sample_rollout_exit(4, torch.device("cpu"), runtime) == 1 + broadcast.assert_called_once() + + def test_continuous_score_sampling_is_uniform_inside_bounds_without_shift(self, monkeypatch): + computer = QwenImageDMDComputer( + model_config(), + make_plan(score_discrete_steps=0, score_sigma_min=0.2, score_sigma_max=0.6, score_timestep_shift=8.0), + ) + monkeypatch.setattr(torch, "rand", Mock(return_value=torch.tensor([0.0, 0.25, 0.5, 0.75]))) + sigma = computer.sample_score_sigma(torch.zeros(4, 1, 4), ToyRuntime()) + torch.testing.assert_close(sigma, torch.tensor([0.2, 0.3, 0.4, 0.5])) + + def test_training_and_inference_shifts_must_match(self): + config = model_config() + config.algo = SimpleNamespace(rollout_timestep_shift=4.0) + with pytest.raises(ValueError, match="must match"): + QwenImageDMDComputer(config, make_plan(fake_update_ratio=1)) + + def test_student_phase_keeps_only_student_gradient(self, monkeypatch): + runtime = ToyRuntime() + computer = QwenImageDMDComputer(model_config(), make_plan(fake_update_ratio=1)) + monkeypatch.setattr( + computer, + "rollout_sigmas", + lambda scheduler, height, width, device: torch.tensor([1.0, 0.5, 0.0], device=device), + ) + + computation = computer.compute_phase(phase_request("student"), phase_batch(), runtime) + computation.losses["student"].backward() + + assert computation.losses["student"].requires_grad + assert runtime.modules["student"].scale.grad is not None + assert runtime.modules["fake_score"].scale.grad is None + assert runtime.modules["teacher_score"].scale.grad is None + assert 0.02 <= computation.metrics["score/sigma"] <= 1.0 + + def test_fake_phase_detaches_student_rollout(self, monkeypatch): + runtime = ToyRuntime() + computer = QwenImageDMDComputer(model_config(), make_plan(fake_update_ratio=1)) + monkeypatch.setattr( + computer, + "rollout_sigmas", + lambda scheduler, height, width, device: torch.tensor([1.0, 0.5, 0.0], device=device), + ) + + batch = phase_batch() + del batch["negative_prompt_embeds"] + del batch["negative_prompt_embeds_mask"] + computation = computer.compute_phase(phase_request("fake_score"), batch, runtime) + computation.losses["fake_score"].backward() + + assert runtime.modules["fake_score"].scale.grad is not None + assert runtime.modules["student"].scale.grad is None + assert runtime.modules["teacher_score"].scale.grad is None + + def test_original_dmd_adds_paired_latent_regression(self): + runtime = ToyRuntime() + computer = QwenImageDMDComputer( + model_config("dmd"), + make_plan("dmd", regression_type="latent_mse", regression_loss_weight=2.0), + ) + + computation = computer.compute_phase(phase_request("student"), phase_batch(regression=True), runtime) + + assert computation.losses["student"].requires_grad + assert computation.metrics["regression/loss"] >= 0 + assert computation.metrics["dmd/loss"] >= 0 + + def test_four_step_rollout_matches_reference_linear_shift(self): + config = model_config() + config.pipeline.num_inference_steps = 4 + computer = QwenImageDMDComputer(config, make_plan(fake_update_ratio=1)) + + sigmas = computer.rollout_sigmas(None, 16, 16, torch.device("cpu")) + + torch.testing.assert_close(sigmas, torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0])) + torch.testing.assert_close(sigmas, build_qwen_dmd_sigmas(4, 3.0)) + + def test_vllm_rollout_accepts_request_local_shift_and_restores_state(self, monkeypatch): + from verl_omni.pipelines.qwen_image_flow_grpo.vllm_omni_rollout_adapter import QwenImagePipelineWithLogProb + + pipeline = object.__new__(QwenImageDMDPipeline) + pipeline.rollout_timestep_shift = 3.0 + request = SimpleNamespace( + sampling_params=SimpleNamespace(extra_args={"rollout_timestep_shift": 4.0}), + ) + monkeypatch.setattr(QwenImagePipelineWithLogProb, "forward", capture_qwen_forward) + + observed_shift, captured = pipeline.forward(request) + assert observed_shift == 4.0 + assert pipeline.rollout_timestep_shift == 3.0 + assert captured["noise_level"] == 0.0 + assert captured["logprobs"] is False + assert captured["true_cfg_scale"] == 1.0 + + def test_vllm_rollout_rejects_stochastic_sampling(self): + pipeline = object.__new__(QwenImageDMDPipeline) + request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"noise_level": 0.1})) + with pytest.raises(ValueError, match="noise_level=0"): + pipeline.forward(request) + + def test_vllm_schedule_uses_identical_training_sigmas(self): + pipeline = object.__new__(QwenImageDMDPipeline) + pipeline.rollout_timestep_shift = 3.0 + pipeline._components = {} + pipeline.device = torch.device("cpu") + pipeline.scheduler = SimpleNamespace(config={"num_train_timesteps": 1000}) + + timesteps, count = pipeline.prepare_timesteps(4, None, image_seq_len=4096) + + assert count == 4 + torch.testing.assert_close(pipeline.scheduler.sigmas, build_qwen_dmd_sigmas(4, 3.0)) + torch.testing.assert_close(timesteps, torch.tensor([1000.0, 900.0, 750.0, 500.0])) + + def test_vllm_deterministic_step_matches_training_euler(self): + scheduler = FlowMatchSDEDiscreteScheduler( + num_train_timesteps=1000, + use_dynamic_shifting=True, + time_shift_type="exponential", + ) + pipeline = object.__new__(QwenImageDMDPipeline) + pipeline.rollout_timestep_shift = 3.0 + pipeline._components = {} + pipeline.device = torch.device("cpu") + pipeline.scheduler = scheduler + timesteps, _ = pipeline.prepare_timesteps(4, None, image_seq_len=4096) + scheduler.set_begin_index(0) + sample = torch.randn(1, 4, 8) + velocity = torch.randn_like(sample) + + inference = scheduler.step( + velocity, + timesteps[0], + sample, + noise_level=0.0, + return_logprobs=False, + ).prev_sample + training = ode_euler_step(sample, velocity, scheduler.sigmas[0], scheduler.sigmas[1]) + + torch.testing.assert_close(inference, training) + + def test_training_rollout_matches_vllm_deterministic_latents(self, monkeypatch): + runtime = ToyRuntime() + computer = QwenImageDMDComputer(model_config(), make_plan(fake_update_ratio=1)) + monkeypatch.setattr(computer, "sample_rollout_exit", lambda high, device, runtime: high - 1) + condition, _ = computer.condition_provider.encode( + phase_batch(), + device=torch.device("cpu"), + dtype=torch.float32, + require_negative=False, + ) + initial = torch.randn(1, 1, 4) + + training_x0, _, _, _ = computer.rollout( + runtime, + condition, + initial.clone(), + height=16, + width=16, + grad_enabled=False, + ) + scheduler = FlowMatchSDEDiscreteScheduler( + num_train_timesteps=1000, + use_dynamic_shifting=True, + time_shift_type="exponential", + ) + pipeline = object.__new__(QwenImageDMDPipeline) + pipeline.rollout_timestep_shift = 3.0 + pipeline._components = {} + pipeline.device = torch.device("cpu") + pipeline.scheduler = scheduler + timesteps, _ = pipeline.prepare_timesteps(2, None, image_seq_len=1) + scheduler.set_begin_index(0) + inference = initial.clone() + for index, timestep in enumerate(timesteps): + sigma = scheduler.sigmas[index].reshape(1) + velocity = computer.predict_velocity( + runtime, + "student", + inference, + sigma, + condition, + height=16, + width=16, + grad_enabled=False, + ) + inference = scheduler.step( + velocity, + timestep, + inference, + noise_level=0.0, + return_logprobs=False, + ).prev_sample + + torch.testing.assert_close(training_x0, inference) + + @pytest.mark.parametrize("batch_size", [1, 3]) + def test_decoded_lpips_regression_retains_student_gradient(self, monkeypatch, batch_size): + computer = QwenImageDMDComputer( + model_config("dmd"), + make_plan("dmd", regression_type="decoded_lpips"), + ) + + monkeypatch.setattr(computer, "ensure_vae_and_lpips", Mock(return_value=(ToyVAE(), ToyLPIPS()))) + prediction = torch.randn(batch_size, 1, 4, requires_grad=True) + target = torch.zeros_like(prediction) + + loss = computer.decoded_lpips_loss(prediction, target, None, height=16, width=16) + loss.backward() + + assert loss.ndim == 0 + assert prediction.grad is not None + assert torch.count_nonzero(prediction.grad) > 0 + + def test_rng_state_round_trip_replays_next_phase(self, monkeypatch): + plan = make_plan(fake_update_ratio=1) + computer = QwenImageDMDComputer(model_config(), plan) + restored = QwenImageDMDComputer(model_config(), plan) + for instance in (computer, restored): + monkeypatch.setattr( + instance, + "rollout_sigmas", + lambda scheduler, height, width, device: torch.tensor([1.0, 0.5, 0.0], device=device), + ) + runtime = ToyRuntime() + computer.compute_phase(phase_request("fake_score"), phase_batch(), runtime) + state = computer.state_dict() + expected = computer.compute_phase(phase_request("fake_score"), phase_batch(), runtime) + + restored.load_state_dict(state) + actual = restored.compute_phase(phase_request("fake_score"), phase_batch(), runtime) + + torch.testing.assert_close(actual.losses["fake_score"], expected.losses["fake_score"]) + expected_values = {key: value for key, value in expected.metrics.items() if not key.startswith("perf/")} + actual_values = {key: value for key, value in actual.metrics.items() if not key.startswith("perf/")} + assert actual_values == expected_values + + @pytest.mark.parametrize("algorithm", ["dmd", "dmd2"]) + @pytest.mark.parametrize("kind", ["student", "fake_score"]) + @pytest.mark.parametrize("batch_size", [2, 3]) + def test_physical_batch_preserves_role_gradient_ownership(self, algorithm, kind, batch_size): + runtime = ToyRuntime() + computer = QwenImageDMDComputer(model_config(algorithm), make_plan(algorithm, regression_type="latent_mse")) + batch = phase_batch(batch_size, regression=algorithm == "dmd") + if algorithm == "dmd": + tu.assign_non_tensor_stack(batch, "teacher_sampling_manifest", [{"pair": i} for i in range(batch_size)]) + result = computer.compute_phase(phase_request(kind), batch, runtime) + result.losses[kind].backward() + assert result.losses[kind].isfinite() + for role, module in runtime.modules.items(): + assert (module.scale.grad is not None) == (role == kind) + + def test_mixed_resolution_batch_fails_closed(self): + computer = QwenImageDMDComputer(model_config(), make_plan()) + batch = phase_batch(2) + batch["height"] = torch.tensor([16, 32]) + with pytest.raises(ValueError, match="homogeneous height"): + computer.compute_phase(phase_request("student"), batch, ToyRuntime()) + + @pytest.mark.parametrize("manifests", [[{"pair": 0}, {}], [{"pair": 0}], {"pair": 0}]) + def test_original_dmd_requires_provenance_for_every_batch_row(self, manifests): + computer = QwenImageDMDComputer(model_config("dmd"), make_plan("dmd", regression_type="latent_mse")) + batch = phase_batch(2, regression=True) + batch["teacher_sampling_manifest"] = NonTensorData(manifests, batch_size=batch.batch_size) + with pytest.raises(ValueError, match="provenance"): + computer.compute_phase(phase_request("student"), batch, ToyRuntime()) + + @pytest.mark.parametrize("kind", ["student", "fake_score"]) + @pytest.mark.parametrize("cfg_norm", ["none", "layer_norm"]) + def test_batched_loss_and_gradient_match_microbatch_accumulation(self, kind, cfg_norm, monkeypatch): + plan = make_plan(teacher_cfg_norm=cfg_norm) + batch = phase_batch(3) + batch["prompt_embeds"] = torch.arange(18).reshape(3, 2, 3).float() / 18 + full_runtime, accumulated_runtime = ToyRuntime(), ToyRuntime() + full_computer, accumulated_computer = [QwenImageDMDComputer(model_config(), plan) for _ in range(2)] + for computer in (full_computer, accumulated_computer): + monkeypatch.setattr(computer, "sample_noise", constant_noise) + monkeypatch.setattr(computer, "sample_score_sigma", constant_score_sigma) + monkeypatch.setattr(computer, "sample_rollout_exit", Mock(return_value=1)) + full = full_computer.compute_phase(phase_request(kind), batch, full_runtime) + full.losses[kind].backward() + loss = 0.0 + for micro_batch in batch.split(2, dim=0): + result = accumulated_computer.compute_phase(phase_request(kind), micro_batch, accumulated_runtime) + weight = micro_batch.batch_size[0] / batch.batch_size[0] + (result.losses[kind] * weight).backward() + loss += float(result.losses[kind].detach()) * weight + assert float(full.losses[kind].detach()) == pytest.approx(loss) + torch.testing.assert_close(full_runtime.modules[kind].scale.grad, accumulated_runtime.modules[kind].scale.grad) + + def test_precomputed_provider_requires_negative_condition(self): + batch = phase_batch() + del batch["negative_prompt_embeds"] + computer = QwenImageDMDComputer(model_config(), make_plan(fake_update_ratio=1)) + with pytest.raises(ValueError, match="negative_prompt_embeds"): + computer.compute_phase(phase_request("student"), batch, ToyRuntime()) + + def test_runner_rejects_dmd2_adversarial_profile(self): + plan = build_plan( + "dmd2", + {"model_path": "/unused", "conditioning_provider": "precomputed", "profile": "paper"}, + frozenset({"distribution_matching", "adversarial"}), + ) + with pytest.raises(NotImplementedError, match="adversarial profile"): + QwenImageDMDComputer(model_config(), plan) diff --git a/tests/pipelines/test_qwen_image_dmd_request_batch.py b/tests/pipelines/test_qwen_image_dmd_request_batch.py new file mode 100644 index 000000000..599ad3b4e --- /dev/null +++ b/tests/pipelines/test_qwen_image_dmd_request_batch.py @@ -0,0 +1,135 @@ +# 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. +"""GPU integration test of native vLLM-Omni request scheduling and Qwen DMD inference.""" + +import asyncio +import os +from unittest.mock import Mock + +import pytest +import torch + + +async def collect_request(engine, request): + output = None + async for outputs in engine.step_streaming(request): + assert len(outputs) == 1 + output = outputs[0] + assert output is not None and output.finished and output.trajectory_latents is not None + assert len(output.images) == request.sampling_params.num_outputs_per_prompt + for key in ("preprocess_time_ms", "diffusion_engine_exec_time_ms", "postprocess_time_ms"): + assert output.metrics[key] >= 0 + return output + + +async def compare_request_batches(engine, tokens, num_outputs): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + serial = [] + requests = [] + for index, (ids, mask) in enumerate(tokens): + for prefix in ("serial", "packed"): + request = OmniDiffusionRequest( + prompt={"prompt_token_ids": ids, "prompt_mask": mask}, + sampling_params=OmniDiffusionSamplingParams( + seed=41 + index, + num_outputs_per_prompt=num_outputs, + height=64, + width=64, + num_inference_steps=4, + true_cfg_scale=1.0, + output_type="pil", + extra_args={"noise_level": 0.0}, + ), + request_id=f"{prefix}-{index}", + ) + if prefix == "serial": + serial.append(await collect_request(engine, request)) + else: + requests.append(request) + execute = Mock(wraps=engine.execute_fn) + engine.execute_fn = execute + order = [2, 0, 1] + packed = await asyncio.gather(*(collect_request(engine, requests[index]) for index in order)) + assert any(len(call.args[0].scheduled_request_ids) > 1 for call in execute.call_args_list) + for output, index in zip(packed, order, strict=True): + torch.testing.assert_close(output.trajectory_latents, serial[index].trajectory_latents, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("num_outputs", [1, 2]) +def test_native_qwen_dmd_request_batch_matches_serial(num_outputs): + model_path = os.environ.get("QWEN_IMAGE_MODEL_PATH", os.path.expanduser("~/models/tiny-random/Qwen-Image")) + if not torch.cuda.is_available() or not os.path.isfile(os.path.join(model_path, "model_index.json")): + pytest.skip("Requires CUDA and QWEN_IMAGE_MODEL_PATH pointing to a tiny Qwen-Image checkpoint.") + from diffusers import QwenImagePipeline + from transformers import AutoTokenizer + from vllm.distributed import destroy_distributed_environment, destroy_model_parallel + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + + from verl_omni.pipelines.model_base import VllmOmniPipelineBase + from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import QwenImageConditionProvider + + tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"), local_files_only=True) + template = QwenImagePipeline(tokenizer=tokenizer, text_encoder=None, vae=None, transformer=None, scheduler=None) + provider = QwenImageConditionProvider(model_path, "local_frozen_encoder", 64, " ") + tokens = [] + for text in ("cat", "a red apple", "a house on a hill in the evening"): + ids, mask = provider.tokenize_rows(template, [text], torch.device("cpu")) + tokens.append((ids[0].tolist(), mask[0].tolist())) + config = OmniDiffusionConfig.from_kwargs( + model=model_path, + dtype=torch.float32, + num_gpus=1, + distributed_executor_backend="uni", + step_execution=False, + max_num_seqs=6, + request_batch_max_wait_ms=50, + diffusion_attention_backend="TORCH_SDPA", + custom_pipeline_args={"pipeline_class": VllmOmniPipelineBase.get_pipeline_path("QwenImagePipeline", "dmd2")}, + ) + config.enrich_config() + engine = None + try: + engine = DiffusionEngine(config) + asyncio.run(compare_request_batches(engine, tokens, num_outputs)) + finally: + if engine is not None: + engine.close() + destroy_model_parallel() + destroy_distributed_environment() + + +@pytest.mark.parametrize("lengths", [(64, 64), (37, 64)]) +def test_flash3_attention_matches_native_forward_and_backward(lengths): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("This FlashAttention-3 validation targets Hopper GPUs.") + from diffusers.models.attention_dispatch import attention_backend, dispatch_attention_fn + + generator = torch.Generator(device="cuda").manual_seed(17) + values = [torch.randn(2, 64, 4, 128, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3)] + mask = torch.arange(64, device="cuda")[None, :] < torch.tensor(lengths, device="cuda")[:, None] + mask = mask[:, None, None, :] + outputs, gradients = [], [] + for backend in ("native", "_flash_3_varlen_hub"): + inputs = [value.detach().clone().requires_grad_(True) for value in values] + with attention_backend(backend): + output = dispatch_attention_fn(*inputs, attn_mask=mask) + grads = torch.autograd.grad(output.float().square().sum(), inputs) + outputs.append(output.detach()) + gradients.append(grads) + torch.testing.assert_close(outputs[0], outputs[1], rtol=0.02, atol=0.003) + for native, flash in zip(gradients[0], gradients[1], strict=True): + torch.testing.assert_close(native, flash, rtol=0.03, atol=0.005) diff --git a/tests/pipelines/test_qwen_image_dmd_request_batch_on_cpu.py b/tests/pipelines/test_qwen_image_dmd_request_batch_on_cpu.py new file mode 100644 index 000000000..12208c6b2 --- /dev/null +++ b/tests/pipelines/test_qwen_image_dmd_request_batch_on_cpu.py @@ -0,0 +1,169 @@ +# 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 the real request collation, RNG, Euler loop and output splitting on CPU.""" + +from types import MethodType + +import pytest +import torch +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +from verl_omni.pipelines.qwen_image_distillation.vllm_omni_rollout_adapter import QwenImageDMDPipeline +from verl_omni.pipelines.schedulers import FlowMatchSDEDiscreteScheduler + + +class BatchTransformer(torch.nn.Module): + in_channels = 4 + guidance_embeds = False + + def __init__(self): + super().__init__() + self.img_in = torch.nn.Linear(4, 4) + self.batch_sizes = [] + + def forward(self, hidden_states, encoder_hidden_states, encoder_hidden_states_mask, img_shapes, **kwargs): + self.batch_sizes.append(hidden_states.shape[0]) + assert len(img_shapes) == hidden_states.shape[0] + mask = encoder_hidden_states_mask.to(encoder_hidden_states.dtype).unsqueeze(-1) + condition = (encoder_hidden_states * mask).sum(1) / mask.sum(1) + return (hidden_states * 0.2 + condition.unsqueeze(1),) + + +def encode_tokens(self, prompt_ids, attention_mask, num_images_per_prompt, **kwargs): + ids = torch.as_tensor(prompt_ids) + embeds = ids.unsqueeze(-1).float().expand(-1, -1, 4) / 10 + mask = torch.as_tensor(attention_mask) + return embeds.repeat_interleave(num_images_per_prompt, 0), mask.repeat_interleave(num_images_per_prompt, 0) + + +def make_pipeline(): + pipeline = object.__new__(QwenImageDMDPipeline) + torch.nn.Module.__init__(pipeline) + pipeline._components = {} + pipeline.device = torch.device("cpu") + pipeline.vae_scale_factor = 8 + pipeline.default_sample_size = 2 + pipeline.transformer = BatchTransformer() + pipeline.scheduler = FlowMatchSDEDiscreteScheduler() + pipeline.encode_prompt = MethodType(encode_tokens, pipeline) + return pipeline + + +def make_request(index, *, outputs=1, extra_args=None, **overrides): + tokens = [index + 1] * (index + 1) + params = dict( + seed=41 + index, + height=16, + width=16, + num_inference_steps=4, + output_type="latent", + true_cfg_scale=1.0, + num_outputs_per_prompt=outputs, + extra_args=extra_args or {}, + ) + params.update(overrides) + return OmniDiffusionRequest( + request_id=f"request-{index}", + prompt={"prompt_token_ids": tokens, "prompt_mask": [1] * len(tokens)}, + sampling_params=OmniDiffusionSamplingParams(**params), + ) + + +class TestQwenDMDRequestBatch: + @pytest.mark.parametrize("outputs", [1, 2]) + def test_packed_inference_matches_serial_and_preserves_request_order(self, outputs): + pipeline = make_pipeline() + with torch.no_grad(): + serial = [pipeline.forward(make_request(i, outputs=outputs)) for i in range(3)] + pipeline.transformer.batch_sizes.clear() + order = [2, 0, 1] + packed = pipeline.forward(DiffusionRequestBatch([make_request(i, outputs=outputs) for i in order])) + assert pipeline.supports_request_batch + assert pipeline.transformer.batch_sizes == [3 * outputs] * 4 + assert len(packed) == 3 + for actual, index in zip(packed, order, strict=True): + torch.testing.assert_close(actual.output["payload"], serial[index].output["payload"]) + actual_condition = actual.output["metadata"]["prompt_embeddings"] + serial_condition = serial[index].output["metadata"]["prompt_embeddings"] + torch.testing.assert_close( + actual_condition["prompt_embeds"][actual_condition["prompt_embeds_mask"].bool()], + serial_condition["prompt_embeds"][serial_condition["prompt_embeds_mask"].bool()], + ) + torch.testing.assert_close(actual.trajectory_latents, serial[index].trajectory_latents) + assert actual.output["payload"]["image"].shape[0] == outputs + + def test_supplied_latents_are_collated_by_the_existing_request_batch(self): + pipeline = make_pipeline() + requests = [make_request(i, latents=torch.full((1, 1, 4), float(i))) for i in range(2)] + with torch.no_grad(): + packed = pipeline.forward(DiffusionRequestBatch(requests)) + for index, result in enumerate(packed): + torch.testing.assert_close(result.trajectory_latents[:, 0], torch.full((1, 1, 4), float(index))) + + @pytest.mark.parametrize("field,value", [("height", 32), ("num_inference_steps", 2), ("true_cfg_scale", 4.0)]) + def test_scheduler_separates_incompatible_requests_and_direct_calls_fail_closed(self, field, value): + first, second = make_request(0), make_request(1, **{field: value}) + assert build_request_batch_sampling_params_key(first) != build_request_batch_sampling_params_key(second) + pipeline = make_pipeline() + with pytest.raises(ValueError, match="sampling parameters"): + pipeline.forward(DiffusionRequestBatch([first, second])) + assert pipeline.transformer.batch_sizes == [] + + @pytest.mark.parametrize("shift", [0.0, float("nan"), float("inf")]) + def test_invalid_dmd_shift_fails_before_inference(self, shift): + with pytest.raises(ValueError, match="finite and at least 1"): + make_pipeline().forward(make_request(0, extra_args={"rollout_timestep_shift": shift})) + + def test_mixed_dmd_shift_is_rejected_without_mutating_pipeline_state(self): + pipeline = make_pipeline() + requests = [make_request(0), make_request(1, extra_args={"rollout_timestep_shift": 4.0})] + with pytest.raises(ValueError, match="same rollout_timestep_shift"): + pipeline.forward(DiffusionRequestBatch(requests)) + assert pipeline.rollout_timestep_shift == 3.0 + + def test_null_extra_args_keep_defaults_on_the_real_forward_path(self): + pipeline = make_pipeline() + request = make_request(0) + request.sampling_params.extra_args = None + with torch.no_grad(): + result = pipeline.forward(request) + assert torch.isfinite(result.output["payload"]["image"]).all() + + def test_unvalidated_step_mode_is_not_silently_enabled(self): + assert not QwenImageDMDPipeline.supports_step_execution + with pytest.raises(NotImplementedError, match="step_execution=false"): + make_pipeline().prepare_encode(None) + + def test_request_shift_is_restored_after_forward_error(self): + pipeline = make_pipeline() + request = make_request(0, extra_args={"rollout_timestep_shift": 4.0}) + request.prompt["prompt_token_ids"] = ["invalid"] + with pytest.raises((TypeError, ValueError)): + pipeline.forward(request) + assert pipeline.rollout_timestep_shift == 3.0 + assert pipeline.transformer.batch_sizes == [] + + def test_empty_request_batch_fails_before_forward(self): + with pytest.raises(ValueError, match="empty"): + make_pipeline().forward(DiffusionRequestBatch([])) + + def test_initial_noise_matches_training_fp32_regardless_of_embedding_dtype(self): + pipeline = make_pipeline() + full = pipeline.prepare_latents(1, 1, 16, 16, torch.float32, "cpu", torch.Generator().manual_seed(7)) + low = pipeline.prepare_latents(1, 1, 16, 16, torch.bfloat16, "cpu", torch.Generator().manual_seed(7)) + assert low.dtype == torch.float32 + torch.testing.assert_close(full, low, rtol=0, atol=0) diff --git a/tests/special_e2e/create_dummy_diffusion_data.py b/tests/special_e2e/create_dummy_diffusion_data.py index 90778120b..db200f87a 100644 --- a/tests/special_e2e/create_dummy_diffusion_data.py +++ b/tests/special_e2e/create_dummy_diffusion_data.py @@ -41,7 +41,7 @@ ] -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 = [] for i in range(n): prompt_text = USER_PROMPTS[i % len(USER_PROMPTS)] @@ -60,6 +60,9 @@ def build_rows(split: str, n: int, data_sources: list[str]): "extra_info": {"split": split, "index": i}, } ) + if user_prompt_only: + rows[-1]["prompt"] = [{"role": "user", "content": prompt_text}] + rows[-1]["negative_prompt"] = [{"role": "user", "content": " "}] return rows @@ -77,13 +80,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="Let the model supply its fixed system template" + ) 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, user_prompt_only=args.user_prompt_only)) + val_df = pd.DataFrame(build_rows("test", args.val_size, data_sources, user_prompt_only=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 100755 index 000000000..f59087953 --- /dev/null +++ b/tests/special_e2e/run_dmd2_qwen_image.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# Qwen-Image DMD2 multi-role training smoke test. +# Requires a tiny checkpoint at ~/models/tiny-random/Qwen-Image. +set -xeuo pipefail + +NUM_GPUS=${NUM_GPUS:-4} +MODEL_PATH=${MODEL_PATH:-${HOME}/models/tiny-random/Qwen-Image} +TOKENIZER_PATH=${TOKENIZER_PATH:-${MODEL_PATH}/tokenizer} +DATA_DIR=${DATA_DIR:-${HOME}/data/dummy_diffusion} +TRAIN_FILES=${TRAIN_FILES:-${DATA_DIR}/train.parquet} +VAL_FILES=${VAL_FILES:-${DATA_DIR}/test.parquet} +TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-2} + +python3 tests/special_e2e/create_dummy_diffusion_data.py \ + --local_save_dir "${DATA_DIR}" \ + --train_size "$((NUM_GPUS * TOTAL_TRAIN_STEPS))" \ + --val_size "${NUM_GPUS}" \ + --user_prompt_only + +python3 -m verl_omni.trainer.main_diffusion \ + data.train_files=${TRAIN_FILES} \ + data.val_files=${VAL_FILES} \ + data.train_batch_size=${NUM_GPUS} \ + data.max_prompt_length=64 \ + data.dataloader_num_workers=0 \ + algorithm.trainer_type=distillation \ + algorithm.sample_source=offline \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.model.tokenizer_path=${TOKENIZER_PATH} \ + actor_rollout_ref.model.algorithm=dmd2 \ + actor_rollout_ref.model.model_type=diffusion_distillation_model \ + actor_rollout_ref.model.attn_backend=native \ + actor_rollout_ref.rollout.rollout_attn_backend=TORCH_SDPA \ + actor_rollout_ref.rollout.algo.noise_level=0.0 \ + actor_rollout_ref.rollout.algo.rollout_timestep_shift=3.0 \ + actor_rollout_ref.model.lora_rank=8 \ + actor_rollout_ref.model.lora_alpha=8 \ + actor_rollout_ref.model.target_modules=all-linear \ + actor_rollout_ref.model.pipeline.height=64 \ + actor_rollout_ref.model.pipeline.width=64 \ + actor_rollout_ref.model.pipeline.num_inference_steps=2 \ + actor_rollout_ref.model.pipeline.max_sequence_length=64 \ + actor_rollout_ref.actor.strategy=fsdp2 \ + actor_rollout_ref.actor.optim.lr=1e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.001 \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.param_offload=false \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ + distillation.enabled=false \ + distillation.distribution_matching.recipe=dmd2 \ + distillation.distribution_matching.profile=distribution_only \ + distillation.distribution_matching.role_storage=shared_base_adapters \ + distillation.distribution_matching.conditioning_provider=local_frozen_encoder \ + distillation.distribution_matching.fake_update_ratio=1 \ + distillation.distribution_matching.student_micro_batch_size_per_gpu=1 \ + distillation.distribution_matching.fake_score_micro_batch_size_per_gpu=1 \ + distillation.distribution_matching.teacher_guidance_scale=4.0 \ + distillation.distribution_matching.teacher_cfg_norm=layer_norm \ + distillation.distribution_matching.rollout_timestep_shift=3.0 \ + distillation.distribution_matching.score_timestep_shift=3.0 \ + trainer.logger=console \ + trainer.project_name=verl-test \ + trainer.experiment_name=qwen-image-dmd2-e2e \ + trainer.log_val_generations=0 \ + trainer.n_gpus_per_node=${NUM_GPUS} \ + trainer.nnodes=1 \ + trainer.val_before_train=false \ + trainer.test_freq=-1 \ + trainer.save_freq=-1 \ + trainer.resume_mode=disable \ + trainer.total_epochs=1 \ + trainer.total_training_steps=${TOTAL_TRAIN_STEPS} \ + "$@" + +echo "Qwen-Image DMD2 e2e test passed." diff --git a/tests/trainer/diffusion/test_distillation_config_on_cpu.py b/tests/trainer/diffusion/test_distillation_config_on_cpu.py index 2b9df088c..1a8e96591 100644 --- a/tests/trainer/diffusion/test_distillation_config_on_cpu.py +++ b/tests/trainer/diffusion/test_distillation_config_on_cpu.py @@ -41,6 +41,16 @@ def test_defaults_select_dmd2_without_enabling_opd(self): assert config.distribution_matching.fake_score_micro_batch_size_per_gpu == 1 assert config.distribution_matching.ema_decay == pytest.approx(0.999) assert config.distribution_matching.ema_start_step == 0 + assert config.distribution_matching.conditioning_provider == "local_frozen_encoder" + assert config.distribution_matching.teacher_guidance_scale == pytest.approx(4.0) + assert config.distribution_matching.teacher_cfg_norm == "layer_norm" + assert config.distribution_matching.negative_prompt == " " + assert config.distribution_matching.rollout_timestep_shift == pytest.approx(3.0) + assert config.distribution_matching.score_sigma_min == pytest.approx(0.02) + assert config.distribution_matching.score_sigma_max == pytest.approx(0.98) + assert config.distribution_matching.score_timestep_shift == pytest.approx(3.0) + assert config.distribution_matching.score_discrete_steps == 1000 + assert config.distribution_matching.regression_type == "decoded_lpips" assert config.distribution_matching.fake_score_optim.lr == pytest.approx(2e-5) @pytest.mark.parametrize( @@ -59,12 +69,35 @@ def test_defaults_select_dmd2_without_enabling_opd(self): ({"ema_decay": -0.1}, "ema_decay"), ({"ema_decay": 1.1}, "ema_decay"), ({"ema_start_step": -1}, "non-negative"), + ({"conditioning_provider": "remote"}, "Invalid conditioning_provider"), + ({"negative_prompt": None}, "negative_prompt"), + ({"teacher_guidance_scale": 1.0}, "teacher_guidance_scale"), + ({"teacher_cfg_norm": "batch_norm"}, "teacher_cfg_norm"), + ({"rollout_timestep_shift": 0.0}, "rollout_timestep_shift"), + ({"score_sigma_min": -0.1}, "score sigma bounds"), + ({"score_sigma_max": 1.1}, "score sigma bounds"), + ({"score_timestep_shift": 0.0}, "score_timestep_shift"), + ({"score_discrete_steps": -1}, "score_discrete_steps"), + ({"normalization_epsilon": 0.0}, "normalization_epsilon"), + ({"dmd_loss_weight": -1.0}, "dmd_loss_weight"), + ({"dmd_loss_weight": 0.0}, "DMD2 requires"), + ({"regression_type": "pixel_mse"}, "regression_type"), + ({"regression_loss_weight": -1.0}, "regression_loss_weight"), + ({"rng_seed": -1}, "rng_seed"), ], ) def test_invalid_values_fail_closed(self, kwargs, error): with pytest.raises(ValueError, match=error): DiffusionDistributionMatchingConfig(**kwargs) + def test_dmd_requires_at_least_one_objective_weight(self): + with pytest.raises(ValueError, match="at least one positive"): + DiffusionDistributionMatchingConfig( + recipe="dmd", + dmd_loss_weight=0.0, + regression_loss_weight=0.0, + ) + def test_null_fields_use_recipe_specific_defaults(self): from verl_omni.trainer.diffusion.distillation.recipes import build_plan @@ -75,6 +108,10 @@ def test_null_fields_use_recipe_specific_defaults(self): frozenset({"distribution_matching"}), ) assert plan.objective["profile"] == "paper" + assert plan.objective["teacher_guidance_scale"] == pytest.approx(4.0) + assert plan.objective["regression_type"] == "decoded_lpips" + assert plan.rollout["score_timestep_shift"] == pytest.approx(3.0) + assert plan.data_requirements["conditioning_provider"] == "local_frozen_encoder" assert plan.update_schedule.phases[1].repeats == 1 @@ -227,12 +264,22 @@ def test_composed_config_builds_validated_plan(self): "algorithm.sample_source=offline", "actor_rollout_ref.model.path=/m", "distillation.distribution_matching.fake_update_ratio=2", + "distillation.distribution_matching.conditioning_provider=precomputed", + "distillation.distribution_matching.rollout_timestep_shift=2.5", + "distillation.distribution_matching.score_timestep_shift=4.0", + "distillation.distribution_matching.teacher_guidance_scale=3.5", + "distillation.distribution_matching.regression_type=latent_mse", ] ) 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 + assert plan.data_requirements["conditioning_provider"] == "precomputed" + assert plan.rollout["rollout_timestep_shift"] == pytest.approx(2.5) + assert plan.rollout["score_timestep_shift"] == pytest.approx(4.0) + assert plan.objective["teacher_guidance_scale"] == pytest.approx(3.5) + assert plan.objective["regression_type"] == "latent_mse" def test_null_overrides_use_each_recipe_default(self): from verl_omni.trainer.diffusion.distillation.recipes import build_plan_from_config diff --git a/tests/trainer/diffusion/test_distillation_metrics_on_cpu.py b/tests/trainer/diffusion/test_distillation_metrics_on_cpu.py new file mode 100644 index 000000000..281329605 --- /dev/null +++ b/tests/trainer/diffusion/test_distillation_metrics_on_cpu.py @@ -0,0 +1,158 @@ +# 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. +"""Cycle metrics and existing Tracking/DistProfiler lifecycle regression tests.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from omegaconf import OmegaConf + +from verl_omni.trainer.diffusion.distillation.contracts import PhaseResult +from verl_omni.trainer.diffusion.distillation.controller import FakeBatchProvider, FakePhaseExecutor +from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationRayTrainer +from verl_omni.trainer.diffusion.distillation.recipes import build_plan + + +class MetricExecutor(FakePhaseExecutor): + def execute_phase(self, request, batch): + result = super().execute_phase(request, batch) + return PhaseResult( + optimizer_steps=result.optimizer_steps, + metrics={ + f"{request.kind}/loss": float(request.repeat_index + 1), + "perf/condition_encode_s": 0.25, + f"perf/{request.kind}_s": float(request.repeat_index + 2), + "memory/max_allocated_gb": float(request.repeat_index + 1), + f"training/{request.kind}_samples": 8.0, + }, + ) + + +def make_trainer(**config): + plan = build_plan("dmd2", {"model_path": "/model", "fake_update_ratio": 2, **config}, {"distribution_matching"}) + return DistillationRayTrainer(plan, executor=MetricExecutor(), batch_provider=FakeBatchProvider(num_batches=100)) + + +def production_trainer(monkeypatch): + trainer = make_trainer() + trainer._production = True + trainer.config = OmegaConf.create( + { + "trainer": {"save_freq": 2, "project_name": "test", "experiment_name": "metrics", "logger": ["console"]}, + "data": {"train_batch_size": 8}, + "global_profiler": {"steps": [2]}, + } + ) + trainer.total_training_steps = 4 + trainer._load_checkpoint = Mock(return_value=0) + trainer._save_checkpoint = Mock() + trainer.distillation_worker_group = SimpleNamespace(start_profile=Mock(), stop_profile=Mock()) + tracker = Mock() + monkeypatch.setattr("verl_omni.trainer.diffusion.distillation.ray_trainer.Tracking", Mock(return_value=tracker)) + return trainer, tracker + + +class TestDistillationMetrics: + def test_every_repeated_phase_is_recorded_and_cycle_timings_are_summed(self): + trainer = make_trainer() + trainer.fit(num_cycles=1) + assert len(trainer.controller.metrics) == 3 + metrics = trainer.flatten_metrics(trainer.controller.metrics) + assert metrics["student/loss"] == 1.0 + assert metrics["fake_score/loss"] == 1.5 + assert metrics["perf/condition_encode_s"] == 0.75 + assert metrics["perf/fake_score_s"] == 5.0 + assert metrics["memory/max_allocated_gb"] == 2.0 + assert metrics["phase/fake_score/1/perf/fake_score_s"] == 3.0 + + def test_performance_ratios_and_rates_are_not_summed_as_durations(self): + metrics = DistillationRayTrainer.flatten_metrics( + { + "student": {"perf/mfu": 0.2, "perf/rate_per_s": 10.0}, + "fake_score": {"perf/mfu": 0.4, "perf/rate_per_s": 20.0}, + } + ) + assert metrics["perf/mfu"] == pytest.approx(0.3) + assert metrics["perf/rate_per_s"] == pytest.approx(15.0) + + def test_real_console_tracking_receives_plain_numeric_scalars(self, capsys): + from verl.utils.tracking import Tracking + + trainer = make_trainer() + trainer.fit(num_cycles=1) + metrics = trainer.flatten_metrics(trainer.controller.metrics) + assert all(type(value) is float for value in metrics.values()) + tracker = Tracking(project_name="test", experiment_name="metrics", default_backend=["console"], config={}) + tracker.log(data=metrics, step=1) + output = capsys.readouterr().out + assert "np.float" not in output + assert "fake_score/loss:1.5" in output + + def test_metrics_do_not_leak_into_the_next_cycle(self): + trainer = make_trainer() + trainer.fit(num_cycles=1) + trainer.controller.metrics["system"] = {"perf/checkpoint_s": 12.0} + trainer.fit(num_cycles=1) + assert "perf/checkpoint_s" not in trainer.flatten_metrics(trainer.controller.metrics) + assert len(trainer.controller.metrics) == 3 + + def test_failed_cycle_preserves_previous_metrics_but_is_not_loggable_as_success(self): + trainer = make_trainer() + trainer.fit(num_cycles=1) + before = dict(trainer.controller.metrics) + trainer.executor.fail_on = "fake_score" + with pytest.raises(RuntimeError, match="failed on phase"): + trainer.fit(num_cycles=1) + assert trainer.controller.metrics == before + assert trainer.controller.counters.global_step == 1 + + def test_tracking_receives_cycle_latency_samples_and_nonstale_checkpoint_time(self, monkeypatch): + trainer, tracker = production_trainer(monkeypatch) + trainer.fit(num_cycles=3) + assert tracker.log.call_count == 3 + logs = [call.kwargs for call in tracker.log.call_args_list] + assert [call["step"] for call in logs] == [1, 2, 3] + assert "perf/checkpoint_s" not in logs[0]["data"] + assert "perf/checkpoint_s" in logs[1]["data"] + assert "perf/checkpoint_s" not in logs[2]["data"] + for call in logs: + assert call["data"]["perf/cycle_s"] > 0 + assert call["data"]["training/student_samples"] == 8 + assert call["data"]["training/fake_score_samples"] == 16 + assert call["data"]["perf/student_samples_per_s"] > 0 + trainer.distillation_worker_group.start_profile.assert_called_once_with(role="distillation", profile_step=2) + trainer.distillation_worker_group.stop_profile.assert_called_once() + + def test_profile_is_stopped_when_phase_execution_fails(self, monkeypatch): + trainer, tracker = production_trainer(monkeypatch) + trainer.config.global_profiler.steps = [1] + trainer.executor.fail_on = "fake_score" + with pytest.raises(RuntimeError, match="failed on phase"): + trainer.fit(num_cycles=1) + trainer.distillation_worker_group.stop_profile.assert_called_once() + tracker.log.assert_not_called() + + def test_warmup_has_no_student_throughput_or_checkpoint(self, monkeypatch): + trainer, tracker = production_trainer(monkeypatch) + trainer.plan = build_plan( + "dmd2", {"model_path": "/model", "fake_warmup_cycles": 1, "fake_update_ratio": 2}, {"distribution_matching"} + ) + trainer.fit(num_cycles=1) + warmup, student = [call.kwargs["data"] for call in tracker.log.call_args_list] + assert [call.kwargs["step"] for call in tracker.log.call_args_list] == [1, 2] + assert warmup["training/student_samples"] == 0 + assert warmup["perf/student_samples_per_s"] == 0 + assert student["training/student_samples"] == 8 + trainer._save_checkpoint.assert_not_called() diff --git a/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py b/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py index ecc780b40..15117e726 100644 --- a/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py +++ b/tests/trainer/diffusion/test_distillation_trainer_routing_on_cpu.py @@ -117,7 +117,7 @@ def runtime_config(): "diffusion_loss": {"loss_mode": "flow_grpo"}, "use_distill_loss": False, }, - "model": {"path": "/m"}, + "model": {"path": "/m", "algorithm": "dmd2"}, }, "distillation": { "enabled": False, @@ -142,12 +142,23 @@ def trainer_config(*, role_storage="shared_base_adapters", strategy="fsdp2", lor { "distillation": {"distribution_matching": {"role_storage": role_storage}}, "actor_rollout_ref": { - "model": {"lora_rank": lora_rank}, + "model": {"algorithm": "dmd2", "lora_rank": lora_rank}, "actor": {"strategy": strategy, "fsdp_config": {"use_orig_params": use_orig}}, }, } ) + def test_model_algorithm_must_match_recipe(self): + from verl_omni.trainer.diffusion.distillation.recipes import build_plan + + trainer = object.__new__(DistillationRayTrainer) + trainer.plan = build_plan("dmd2", {"model_path": "/m"}, frozenset({"distribution_matching"})) + trainer.config = self.trainer_config() + trainer.config.actor_rollout_ref.model.algorithm = "dmd" + trainer.config.distillation.distribution_matching.recipe = "dmd2" + with pytest.raises(ValueError, match="must match"): + trainer.validate_runtime_config() + def test_shared_base_requires_lora(self): from verl_omni.trainer.diffusion.distillation.recipes import build_plan diff --git a/tests/utils/dataset/test_qwen_image_distillation_dataset_on_cpu.py b/tests/utils/dataset/test_qwen_image_distillation_dataset_on_cpu.py new file mode 100644 index 000000000..4885ea1c5 --- /dev/null +++ b/tests/utils/dataset/test_qwen_image_distillation_dataset_on_cpu.py @@ -0,0 +1,100 @@ +# 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 Qwen-Image original-DMD regression-pair data.""" + +import io + +import pytest +import torch +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn + +from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationBatchProvider +from verl_omni.utils.dataset.qwen_image_distillation_dataset import QwenImageDMDPairDataset, load_float_tensor + + +class TestDMDTensorLoading: + def test_loads_nested_values_and_serialized_tensors(self): + nested = load_float_tensor([[1, 2], [3, 4]], "value") + buffer = io.BytesIO() + torch.save(torch.tensor([5.0]), buffer) + serialized = load_float_tensor(buffer.getvalue(), "value") + + assert nested.dtype == torch.float32 + torch.testing.assert_close(nested, torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + torch.testing.assert_close(serialized, torch.tensor([5.0])) + + def test_rejects_empty_tensor(self): + with pytest.raises(ValueError, match="must not be empty"): + load_float_tensor([], "value") + + +class TestQwenImageDMDPairDataset: + @staticmethod + def make_dataset(): + return object.__new__(QwenImageDMDPairDataset) + + def test_converts_regression_pair_and_preserves_manifest(self, monkeypatch): + row = { + "reference_noise": [[[1.0]]], + "teacher_target_latents": [[[2.0]]], + "teacher_sampling_manifest": {"model": "teacher"}, + "index": 7, + } + monkeypatch.setattr(RLHFDataset, "__getitem__", lambda self, item: dict(row)) + + output = self.make_dataset()[0] + + assert output["reference_noise"].dtype == torch.float32 + assert output["teacher_target_latents"].dtype == torch.float32 + assert output["pair_id"] == "7" + assert output["teacher_sampling_manifest"] == {"model": "teacher"} + + @pytest.mark.parametrize("target_key", ["teacher_target_latents", "teacher_target_pixels"]) + @pytest.mark.parametrize("null_value", [None, float("nan")]) + def test_nullable_target_survives_worker_transport(self, monkeypatch, target_key, null_value): + inactive_key = "teacher_target_pixels" if target_key == "teacher_target_latents" else "teacher_target_latents" + row = { + "reference_noise": [[0.0, 1.0]], + target_key: [[2.0, 3.0]], + inactive_key: null_value, + "teacher_sampling_manifest": {"teacher": "fixture", "steps": 20}, + } + monkeypatch.setattr(RLHFDataset, "__getitem__", lambda self, item: dict(row)) + dataset = self.make_dataset() + batch = DistillationBatchProvider([collate_fn([dataset[0], dataset[1]])]).fresh_batch() + assert tu.get(batch, inactive_key) is None + torch.testing.assert_close(tu.get(batch, target_key), torch.tensor([[[2.0, 3.0]], [[2.0, 3.0]]])) + + @pytest.mark.parametrize( + "row,error", + [ + ({"teacher_target_latents": [1], "teacher_sampling_manifest": {"x": 1}}, "reference_noise"), + ({"reference_noise": [1], "teacher_sampling_manifest": {"x": 1}}, "exactly one teacher target"), + ( + { + "reference_noise": [1], + "teacher_target_latents": [1], + "teacher_target_pixels": [1], + "teacher_sampling_manifest": {"x": 1}, + }, + "exactly one teacher target", + ), + ({"reference_noise": [1], "teacher_target_latents": [1]}, "teacher_sampling_manifest"), + ], + ) + def test_rejects_incomplete_rows(self, monkeypatch, row, error): + monkeypatch.setattr(RLHFDataset, "__getitem__", lambda self, item: dict(row)) + with pytest.raises(ValueError, match=error): + self.make_dataset()[0] diff --git a/tests/workers/config/test_diffusion_config_on_cpu.py b/tests/workers/config/test_diffusion_config_on_cpu.py index 087f4e11c..a7f82fcf8 100644 --- a/tests/workers/config/test_diffusion_config_on_cpu.py +++ b/tests/workers/config/test_diffusion_config_on_cpu.py @@ -151,6 +151,11 @@ def test_defaults(self): assert cfg.sde_window_size is None assert cfg.sde_window_range is None assert cfg.sde_contiguous is True + assert cfg.rollout_timestep_shift is None + + def test_invalid_rollout_timestep_shift_raises(self): + with pytest.raises(ValueError, match="rollout_timestep_shift"): + DiffusionRolloutAlgoConfig(rollout_timestep_shift=0.5) def test_invalid_sample_strategy_raises(self): with pytest.raises(ValueError, match="Unknown sample_strategy"): diff --git a/tests/workers/test_diffusion_distillation_runtime_on_cpu.py b/tests/workers/test_diffusion_distillation_runtime_on_cpu.py index 92d38dd33..b220b2259 100644 --- a/tests/workers/test_diffusion_distillation_runtime_on_cpu.py +++ b/tests/workers/test_diffusion_distillation_runtime_on_cpu.py @@ -80,6 +80,15 @@ def __init__(self, roles, initial=None): self.model_config = object() self.active_role = None + def train_mode(self): + return nullcontext() + + def get_data_parallel_group(self): + return None + + def get_data_parallel_size(self): + return 2 + @contextmanager def use_role(self, role): previous = self.active_role @@ -98,12 +107,6 @@ def backward_role(self, role, loss, retain_graph=False): assert self.active_role is None loss.backward(retain_graph=retain_graph) - def train_mode(self): - return nullcontext() - - def get_data_parallel_group(self): - return None - def parameters_for_role(self, role): return (self.parameters[role],) @@ -241,6 +244,7 @@ def test_element_reduction_normalizes_worker_loss_metrics_and_gradients(self, mo metrics = tu.get(result, "metrics") assert metrics["student/loss"] == pytest.approx(13.0) assert metrics["ode/loss"] == pytest.approx(13.0) + assert metrics["ode/active_elements"] == 10 assert metrics["student/grad_norm"] == pytest.approx(5.2) torch.testing.assert_close(engine.parameters["student"], torch.tensor(1.52)) assert tu.get(result, "optimizer_steps") == {"student": 1} @@ -409,6 +413,85 @@ def test_pr2_rejects_multi_optimizer_adversarial_phase(self): ) +class ToyDMComputer: + def compute_phase(self, request, batch, runtime): + role = request.trainable_roles[0] + parameter = runtime.engine_for_role(role).parameters[role] + return DistillationPhaseComputation( + losses={role: (parameter - batch["target"]).square().mean()}, + metrics={ + "perf/condition_encode_s": 1.0, + "perf/mfu": 0.25, + "dmd/active_elements": float(batch.batch_size[0]), + }, + ) + + +def simulate_dp_reduce(values, op, group): + assert group == "dp" + # Sorted names: loss, peak memory, host time; simulate a slower second rank. + peer = torch.tensor([5.0, 7.0, 9.0]) + if op == torch.distributed.ReduceOp.AVG: + values.copy_((values + peer) / 2) + else: + assert op == torch.distributed.ReduceOp.MAX + values.copy_(torch.maximum(values, peer)) + + +class TestWorkerMetrics: + @pytest.mark.parametrize("micro_batch_size", [1, 2, 3]) + def test_existing_worker_accumulation_keeps_loss_weights_but_sums_elapsed_time(self, monkeypatch, micro_batch_size): + plan = build_plan("dmd2", {"model_path": "/m"}, _CAPABILITIES) + engine = ToyRoleEngine(("student", "teacher_score", "fake_score", "student_ema"), {"student": 1.0}) + runtime = DistillationRoleRuntime( + plan, + {"base": engine}, + ema_decay=0.9, + ema_start_step=0, + micro_batch_sizes={"student": micro_batch_size, "fake_score": 1}, + ) + worker = object.__new__(DiffusionDistillationWorker) + worker.runtime = runtime + worker.dm_computer = ToyDMComputer() + device = Mock() + device.max_memory_allocated.return_value = 2 * 1024**3 + device.max_memory_reserved.return_value = 3 * 1024**3 + monkeypatch.setattr( + "verl_omni.workers.diffusion_distillation_worker.get_torch_device", Mock(return_value=device) + ) + monkeypatch.setattr("verl_omni.workers.diffusion_distillation_worker.get_device_id", Mock(return_value="cpu")) + batch = tu.get_tensordict({"target": torch.tensor([0.0, 2.0, 6.0])}) + tu.assign_non_tensor(batch, phase_request=PhaseRequest("student", 0, 0, "fresh", ("student",), True)) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: + result = worker.execute_phase(batch) + metrics = tu.get(result, "metrics") + assert tu.get(result, "optimizer_steps") == {"student": 1} + assert metrics["student/loss"] == pytest.approx(9.0) + assert metrics["perf/condition_encode_s"] == (3 + micro_batch_size - 1) // micro_batch_size + assert metrics["dmd/active_elements"] == 3 + assert metrics["perf/mfu"] == 0.25 + assert metrics["training/student_samples"] == 6 + assert metrics["memory/max_allocated_gb"] == 2 + assert engine.parameters["student"].item() == pytest.approx(4 / 3) + names = {event.key for event in profile.key_averages()} + assert {"distillation/backward", "distillation/student_optimizer", "distillation/ema"} <= names + + def test_dp_reports_mean_and_slowest_rank_time_without_averaging_peak_memory(self, monkeypatch): + plan = build_plan("dmd2", {"model_path": "/m"}, _CAPABILITIES) + engine = ToyRoleEngine(("student", "teacher_score", "fake_score", "student_ema")) + engine.get_data_parallel_group = Mock(return_value="dp") + runtime = DistillationRoleRuntime(plan, {"base": engine}, ema_decay=0.9, ema_start_step=0) + monkeypatch.setattr("verl_omni.workers.diffusion_distillation_worker.get_device_id", Mock(return_value="cpu")) + monkeypatch.setattr(torch.distributed, "all_reduce", simulate_dp_reduce) + metrics = runtime.reduce_metrics({"loss": 1.0, "memory/max_allocated_gb": 2.0, "perf/student_s": 3.0}) + assert metrics == { + "loss": 3.0, + "memory/max_allocated_gb": 7.0, + "perf/student_s": 6.0, + "perf_max_rank/student_s": 9.0, + } + + class ToyPeftModule(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/workers/test_distillation_fsdp_roles.py b/tests/workers/test_distillation_fsdp_roles.py index 7062d3cc1..0aaae0edc 100644 --- a/tests/workers/test_distillation_fsdp_roles.py +++ b/tests/workers/test_distillation_fsdp_roles.py @@ -15,16 +15,21 @@ import os import tempfile +from datetime import timedelta from types import SimpleNamespace import pytest import torch import torch.distributed as dist from peft import LoraConfig, get_peft_model +from tensordict import TensorDict from torch import nn from torch.distributed.tensor import DTensor +from verl.utils import tensordict_utils as tu -from verl_omni.trainer.diffusion.distillation.contracts import RoleBinding, RoleGroupSpec +from verl_omni.trainer.diffusion.distillation.contracts import PhaseRequest, RoleBinding, RoleGroupSpec +from verl_omni.trainer.diffusion.distillation.recipes import build_plan +from verl_omni.workers.diffusion_distillation_worker import DistillationRoleRuntime from verl_omni.workers.engine.fsdp.distillation_impl import DistillationRoleGroupEngine @@ -267,3 +272,135 @@ def test_distillation_role_switch_preserves_graph_ema_and_state(strategy): ) finally: dist.destroy_process_group() + + +def wrap_qwen_image_model(strategy, model_path): + from diffusers import QwenImageTransformer2DModel + + model = QwenImageTransformer2DModel.from_pretrained( + model_path, + subfolder="transformer", + torch_dtype=torch.float32, + ).cuda() + adapter_config = LoraConfig(r=2, lora_alpha=2, target_modules=["to_q", "to_k", "to_v", "to_out.0"]) + model.add_adapter(adapter_config, adapter_name="student") + model.add_adapter(adapter_config, adapter_name="fake_score") + model.add_adapter(adapter_config, adapter_name="student_ema") + model.set_adapter("student") + if strategy == "fsdp": + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + return FSDP(model, use_orig_params=True, device_id=torch.cuda.current_device()) + from torch.distributed.fsdp import fully_shard + + for block in model.transformer_blocks: + fully_shard(block) + fully_shard(model) + return model + + +@pytest.fixture(scope="module") +def qwen_process_group(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Qwen-Image FSDP distillation tests.") + if dist.is_initialized(): + pytest.skip("This FSDP fixture creates its own process group.") + with tempfile.TemporaryDirectory(prefix="qwen_image_dmd_fsdp_") as tmp_dir: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + dist.init_process_group( + backend="nccl", + init_method="env://" if world_size > 1 else f"file://{os.path.join(tmp_dir, 'rendezvous')}", + rank=int(os.environ.get("RANK", "0")), + world_size=world_size, + timeout=timedelta(seconds=90), + ) + try: + yield + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("strategy", ["fsdp", "fsdp2"]) +@pytest.mark.parametrize("algorithm", ["dmd", "dmd2"]) +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_qwen_image_dm_computer_on_fsdp(strategy, algorithm, batch_size, qwen_process_group): + model_path = os.environ.get("QWEN_IMAGE_MODEL_PATH", os.path.expanduser("~/models/tiny-random/Qwen-Image")) + if not os.path.isfile(os.path.join(model_path, "model_index.json")): + pytest.skip(f"Tiny Qwen-Image checkpoint not found at {model_path}.") + + from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import QwenImageDMDComputer + + world_size = dist.get_world_size() + rank = dist.get_rank() + torch.manual_seed(7) + engine = engine_shell(wrap_qwen_image_model(strategy, model_path)) + engine.scheduler = SimpleNamespace() + engine.model_config = SimpleNamespace(fsdp_layer_prefixes=["transformer_blocks."]) + engine.ulysses_device_mesh = None + engine.ulysses_sequence_parallel_size = 1 + plan = build_plan( + algorithm, + { + "model_path": model_path, + "conditioning_provider": "local_frozen_encoder", + "fake_update_ratio": 2, + "regression_type": "decoded_lpips", + "rng_seed": 3, + }, + frozenset({"distribution_matching"}), + ) + runtime = DistillationRoleRuntime(plan, {"base": engine}, ema_decay=0.9, ema_start_step=0) + model_config = SimpleNamespace( + path=model_path, + local_path=model_path, + transformer_config={"in_channels": 64}, + pipeline=SimpleNamespace( + height=64, + width=64, + num_inference_steps=4, + max_sequence_length=64, + guidance_scale=None, + ), + ) + computer = QwenImageDMDComputer(model_config, plan) + batch = TensorDict({"dummy_tensor": torch.zeros(batch_size, 1, device="cuda")}, batch_size=[batch_size]) + tu.assign_non_tensor_stack( + batch, + "raw_prompt", + [ + [{"role": "user", "content": "cat" if (rank + row) % 2 == 0 else "a red apple on a wooden table"}] + for row in range(batch_size) + ], + ) + if algorithm == "dmd": + pytest.importorskip("piq") + batch["reference_noise"] = torch.randn(batch_size, 16, 64, device="cuda") + batch["teacher_target_latents"] = torch.zeros(batch_size, 16, 64, device="cuda") + tu.assign_non_tensor_stack( + batch, "teacher_sampling_manifest", [{"scheduler": "tiny-qwen", "sample": row} for row in range(batch_size)] + ) + for kind, role in (("student", "student"), ("fake_score", "fake_score"), ("fake_score", "fake_score")) * 3: + request = PhaseRequest( + kind=kind, + global_step=0, + repeat_index=0, + batch_policy="fresh", + trainable_roles=(role,), + update_ema=kind == "student", + ) + runtime.zero_grad(request.trainable_roles) + computation = computer.compute_phase(request, batch, runtime) + exits = torch.tensor([computation.metrics["rollout/exit_index"]], device="cuda") + gathered = [torch.zeros_like(exits) for _ in range(world_size)] + dist.all_gather(gathered, exits) + assert all(value.item() == exits.item() for value in gathered) + runtime.backward_micro_batch(request, computation, weight=1.0) + optimizer_steps, _ = runtime.step_phase(request) + assert optimizer_steps == {role: 1} + + tensors, peft_config = runtime.export_tensors(base_sync_done=True) + exported = list(tensors) + assert exported + assert peft_config["r"] == 2 + assert all(name.startswith("transformer.") for name, _ in exported) diff --git a/verl_omni/pipelines/__init__.py b/verl_omni/pipelines/__init__.py index 73f555380..247b2ddda 100644 --- a/verl_omni/pipelines/__init__.py +++ b/verl_omni/pipelines/__init__.py @@ -20,6 +20,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, @@ -36,6 +37,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 @@ -48,6 +50,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/qwen_image_distillation/__init__.py b/verl_omni/pipelines/qwen_image_distillation/__init__.py new file mode 100644 index 000000000..fb8b69051 --- /dev/null +++ b/verl_omni/pipelines/qwen_image_distillation/__init__.py @@ -0,0 +1,27 @@ +# 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 ( + QwenImageDistributionMatching, + QwenImageDMDComputer, + build_qwen_dmd_sigmas, +) +from .vllm_omni_rollout_adapter import QwenImageDMDPipeline + +__all__ = [ + "QwenImageDistributionMatching", + "QwenImageDMDComputer", + "QwenImageDMDPipeline", + "build_qwen_dmd_sigmas", +] 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..8a253bf7e --- /dev/null +++ b/verl_omni/pipelines/qwen_image_distillation/diffusers_training_adapter.py @@ -0,0 +1,1033 @@ +# 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 training adapter for DMD and DMD2. + +Registers the architecture adapter and owns the differentiable phase computation +following LightX2V's public DMD equations. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Optional + +import torch +from tensordict import TensorDict +from verl.utils import tensordict_utils as tu + +from verl_omni.pipelines.model_base import DiffusionModelBase, DistributionMatchingModelAdapter +from verl_omni.pipelines.qwen_image_flow_grpo.common import ( + QWEN_IMAGE_VAE_SCALE_FACTOR, + QwenImageTokenIdPromptMixin, + build_img_shapes, +) +from verl_omni.pipelines.qwen_image_flow_grpo.diffusers_training_adapter import QwenImage +from verl_omni.trainer.diffusion.distillation.contracts import ConditionBundle, DistillationPlan, PhaseRequest +from verl_omni.trainer.diffusion.distillation.utils import ( + consistency_renoise_step, + dmd_gradient, + dmd_surrogate_loss, + fake_score_loss, + ode_euler_step, + standard_cfg, + timestep_shift, + velocity_to_x0, +) +from verl_omni.workers.config import DiffusionModelConfig + +if TYPE_CHECKING: + from verl_omni.workers.diffusion_distillation_worker import ( + DistillationPhaseComputation, + DistillationRoleRuntime, + ) + +__all__ = ["QwenImageDistributionMatching", "QwenImageDMDComputer", "build_qwen_dmd_sigmas"] + + +def build_qwen_dmd_sigmas( + num_inference_steps: int, + shift: float, + *, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Build the fixed linear-shift schedule used by Qwen-Image DMD training.""" + if num_inference_steps <= 0: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}.") + if shift < 1: + raise ValueError(f"rollout_timestep_shift must be at least 1, got {shift}.") + raw = torch.linspace( + 1.0, + 1.0 / num_inference_steps, + num_inference_steps, + device=device, + dtype=torch.float32, + ) + shifted = timestep_shift(raw * 1000.0, 1000, shift) / 1000.0 + return torch.cat((shifted, shifted.new_zeros(1))) + + +class QwenImageConditionProvider: + """Encode frozen local or precomputed Qwen prompt conditioning.""" + + def __init__( + self, + model_path: str, + provider: str, + max_sequence_length: int, + negative_prompt: str, + ) -> None: + self.model_path = model_path + self.provider = provider + self.max_sequence_length = max_sequence_length + self.negative_prompt = negative_prompt + self.pipeline = None + self._negative_condition: Optional[ConditionBundle] = None + + @staticmethod + def make_condition(prompt_embeds: torch.Tensor, prompt_mask: Optional[torch.Tensor]) -> ConditionBundle: + """Build a detached [B, L, D] condition with a matching [B, L] mask.""" + if prompt_embeds.ndim != 3: + 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 ConditionBundle( + tensors={"prompt_embeds": prompt_embeds}, + masks={"prompt_embeds": 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[ConditionBundle, Optional[ConditionBundle]]: + """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.tensors["prompt_embeds"].shape[0] != batch.batch_size[0]: + raise ValueError("Precomputed Qwen conditioning batch size does not match the phase batch.") + if negative is not None and negative.tensors["prompt_embeds"].shape[0] != batch.batch_size[0]: + raise ValueError("Precomputed Qwen negative conditioning batch size does not match the phase 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 phase 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], + ) -> ConditionBundle: + """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[ConditionBundle, Optional[ConditionBundle]]: + """Encode the positive prompt and optional teacher negative condition.""" + if self.provider == "precomputed": + return self.encode_precomputed(batch, require_negative=require_negative) + if self.provider != "local_frozen_encoder": + raise ValueError(f"Unsupported Qwen conditioning provider {self.provider!r}.") + + 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.tensors["prompt_embeds"].expand(batch.batch_size[0], -1, -1), + self._negative_condition.masks["prompt_embeds"].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 + + +class QwenImageDMDComputer: + """Differentiable Qwen-Image phase program for DMD and distribution-only DMD2.""" + + STREAM_OFFSETS = { + "initial_noise": 0, + "rollout_decision": 1, + "rollout_transition": 2, + "score_sigma": 3, + "score_noise": 4, + } + + def __init__(self, model_config: DiffusionModelConfig, plan: DistillationPlan) -> None: + if plan.name not in {"dmd", "dmd2"}: + raise ValueError(f"QwenImageDMDComputer does not implement recipe {plan.name!r}.") + if plan.objective.get("adversarial", False): + raise NotImplementedError("The Qwen DMD2 adversarial profile is not supported by this phase runner.") + self.model_config = model_config + self.plan = plan + self.height = int(model_config.pipeline.height) + self.width = int(model_config.pipeline.width) + self.num_inference_steps = int(model_config.pipeline.num_inference_steps) + self.strategy = str(plan.rollout["strategy"]) + self.rollout_timestep_shift = float(plan.rollout["rollout_timestep_shift"]) + self.score_sigma_min = float(plan.rollout["score_sigma_min"]) + self.score_sigma_max = float(plan.rollout["score_sigma_max"]) + self.score_timestep_shift = float(plan.rollout["score_timestep_shift"]) + self.score_discrete_steps = int(plan.rollout["score_discrete_steps"]) + self.rng_seed = int(plan.rollout["rng_seed"]) + self.guidance_scale = float(plan.objective["teacher_guidance_scale"]) + self.cfg_norm = str(plan.objective["teacher_cfg_norm"]) + self.normalization_epsilon = float(plan.objective["normalization_epsilon"]) + self.dmd_loss_weight = float(plan.objective["dmd_loss_weight"]) + self.regression_type = str(plan.objective["regression_type"]) + self.regression_loss_weight = float(plan.objective["regression_loss_weight"]) + rollout_config = getattr(model_config, "algo", None) + if hasattr(rollout_config, "get"): + inference_shift = rollout_config.get("rollout_timestep_shift") + else: + inference_shift = getattr(rollout_config, "rollout_timestep_shift", None) + self.inference_rollout_timestep_shift = 3.0 if inference_shift is None else float(inference_shift) + self.condition_provider = QwenImageConditionProvider( + model_config.local_path or model_config.path, + str(plan.data_requirements["conditioning_provider"]), + int(model_config.pipeline.max_sequence_length), + str(plan.data_requirements["negative_prompt"]), + ) + self._generators: dict[str, torch.Generator] = {} + self._pending_generator_states: dict[str, torch.Tensor] = {} + self._vae = None + self._lpips = None + self.validate_config() + + def validate_config(self) -> None: + """Reject unsupported Qwen sampling, geometry, and objective settings.""" + if self.height <= 0 or self.width <= 0: + raise ValueError("Qwen DMD height and width must be positive.") + divisor = QWEN_IMAGE_VAE_SCALE_FACTOR * 2 + if self.height % divisor or self.width % divisor: + raise ValueError(f"Qwen DMD height and width must be divisible by {divisor}.") + if self.num_inference_steps <= 0: + raise ValueError("Qwen DMD num_inference_steps must be positive.") + if self.rollout_timestep_shift < 1: + raise ValueError("Qwen DMD rollout_timestep_shift must be at least 1.") + if self.inference_rollout_timestep_shift != self.rollout_timestep_shift: + raise ValueError( + "Qwen DMD training and inference rollout_timestep_shift must match; " + f"got {self.rollout_timestep_shift} and {self.inference_rollout_timestep_shift}." + ) + if self.strategy not in {"one_step", "ode_euler", "consistency_renoise", "backward_simulated"}: + raise ValueError(f"Unsupported Qwen DMD rollout strategy {self.strategy!r}.") + if not 0 <= self.score_sigma_min < self.score_sigma_max <= 1: + raise ValueError("Qwen DMD score sigma bounds must satisfy 0 <= min < max <= 1.") + if self.score_timestep_shift < 1: + raise ValueError("Qwen DMD score_timestep_shift must be at least 1.") + if self.score_discrete_steps < 0: + raise ValueError("Qwen DMD score_discrete_steps must be non-negative.") + if self.guidance_scale <= 1: + raise ValueError("Qwen DMD teacher guidance_scale must be greater than 1.") + if self.cfg_norm not in {"none", "layer_norm", "scalar"}: + raise ValueError(f"Unsupported Qwen DMD CFG normalization {self.cfg_norm!r}.") + if self.normalization_epsilon <= 0: + raise ValueError("Qwen DMD normalization_epsilon must be positive.") + if self.dmd_loss_weight < 0 or self.regression_loss_weight < 0: + raise ValueError("Qwen DMD objective weights must be non-negative.") + if self.plan.name == "dmd2" and self.dmd_loss_weight == 0: + raise ValueError("Qwen DMD2 requires dmd_loss_weight > 0.") + if self.plan.name == "dmd" and self.dmd_loss_weight == 0 and self.regression_loss_weight == 0: + raise ValueError("Qwen DMD requires at least one positive objective weight.") + if self.regression_type not in {"decoded_lpips", "latent_mse"}: + raise ValueError(f"Unsupported Qwen DMD regression_type {self.regression_type!r}.") + if self.condition_provider.provider not in {"local_frozen_encoder", "precomputed"}: + raise ValueError(f"Unsupported Qwen conditioning provider {self.condition_provider.provider!r}.") + if self.plan.name == "dmd" and self.strategy != "one_step": + raise ValueError("The Qwen original-DMD profile requires rollout_strategy='one_step'.") + + @staticmethod + def module_dtype(module: torch.nn.Module) -> torch.dtype: + """Read the model parameter dtype for frozen conditioning.""" + try: + return next(module.parameters()).dtype + except StopIteration: + return torch.float32 + + @staticmethod + def module_config(module: torch.nn.Module): + """Read the underlying transformer configuration through a wrapper.""" + return getattr(getattr(module, "module", module), "config", None) + + def generator_for_stream( + self, name: str, device: torch.device, runtime: DistillationRoleRuntime + ) -> torch.Generator: + """Resolve a checkpointable RNG stream, shared by sequence-parallel peers.""" + if name in self._generators: + return self._generators[name] + if name not in self.STREAM_OFFSETS: + raise KeyError(f"Unknown Qwen DMD RNG stream {name!r}.") + engine = runtime.engine_for_role("student") + rank = int(engine.get_data_parallel_rank()) if hasattr(engine, "get_data_parallel_rank") else 0 + generator = torch.Generator(device=device) + generator.manual_seed(self.rng_seed + rank * len(self.STREAM_OFFSETS) + self.STREAM_OFFSETS[name]) + pending = self._pending_generator_states.pop(name, None) + if pending is not None: + generator.set_state(pending) + self._generators[name] = generator + return generator + + def sample_noise(self, shape: tuple[int, ...], device: torch.device, runtime: DistillationRoleRuntime, stream: str): + """Draw fp32 noise from the named independent RNG stream.""" + return torch.randn( + shape, device=device, dtype=torch.float32, generator=self.generator_for_stream(stream, device, runtime) + ) + + def sample_rollout_exit(self, high: int, device: torch.device, runtime: DistillationRoleRuntime) -> int: + """Broadcast one exit step so FSDP and sequence-parallel control flow agree.""" + value = torch.randint( + high, + (1,), + device=device, + generator=self.generator_for_stream("rollout_decision", device, runtime), + ) + if torch.distributed.is_initialized(): + # FSDP shards and SP peers must enter identical forward/backward collectives. + torch.distributed.broadcast(value, src=0) + return int(value.item()) + + @staticmethod + def batch_int(batch: TensorDict, key: str, default: int) -> int: + """Require homogeneous integer geometry metadata within a micro-batch.""" + value = tu.get(batch, key, default) + if isinstance(value, torch.Tensor): + values = value.detach().reshape(-1) + if values.numel() == 0 or not torch.all(values == values[0]): + raise ValueError(f"Qwen DMD requires one homogeneous {key} per micro-batch.") + return int(values[0].item()) + if isinstance(value, list | tuple): + values = [int(item) for item in value] + if not values or any(item != values[0] for item in values): + raise ValueError(f"Qwen DMD requires one homogeneous {key} per micro-batch.") + return values[0] + return int(value) + + def latent_geometry(self, batch: TensorDict, module: torch.nn.Module) -> tuple[int, int, int, tuple[int, ...]]: + """Resolve the declared Qwen packed-token and VAE latent dimensions.""" + if len(batch.batch_size) != 1 or batch.batch_size[0] <= 0: + raise ValueError("Qwen DMD requires a nonempty leading batch dimension.") + height = self.batch_int(batch, "height", self.height) + width = self.batch_int(batch, "width", self.width) + divisor = QWEN_IMAGE_VAE_SCALE_FACTOR * 2 + if height <= 0 or width <= 0 or height % divisor or width % divisor: + raise ValueError( + f"Qwen DMD height and width must be positive multiples of {divisor}, got {height}x{width}." + ) + module_config = self.module_config(module) + in_channels = getattr(module_config, "in_channels", None) + if in_channels is None: + in_channels = (self.model_config.transformer_config or {}).get("in_channels") + if not isinstance(in_channels, int) or in_channels <= 0 or in_channels % 4: + raise ValueError(f"Qwen transformer in_channels must be a positive multiple of four, got {in_channels!r}.") + latent_channels = in_channels // 4 + latent_height = height // QWEN_IMAGE_VAE_SCALE_FACTOR + latent_width = width // QWEN_IMAGE_VAE_SCALE_FACTOR + return height, width, in_channels, (batch.batch_size[0], latent_channels, 1, latent_height, latent_width) + + @staticmethod + def pack_latents(latents: torch.Tensor) -> torch.Tensor: + """Pack normalized [B, C, 1, H, W] latents with the Diffusers Qwen helper.""" + from diffusers import QwenImagePipeline + + batch, channels, _, height, width = latents.shape + return QwenImagePipeline._pack_latents(latents, batch, channels, height, width) + + @staticmethod + def expand_sigma(sigma: torch.Tensor, tensor: torch.Tensor) -> torch.Tensor: + """Broadcast a scalar or per-sample sigma over non-batch latent dimensions.""" + if sigma.ndim == 0: + sigma = sigma.reshape(1) + if sigma.shape[0] == 1 and tensor.shape[0] != 1: + sigma = sigma.expand(tensor.shape[0]) + if sigma.shape[0] != tensor.shape[0]: + raise ValueError(f"Sigma batch {sigma.shape[0]} does not match latent batch {tensor.shape[0]}.") + return sigma.reshape(sigma.shape[0], *((1,) * (tensor.ndim - 1))) + + def predict_velocity( + self, + runtime: DistillationRoleRuntime, + role: str, + latents: torch.Tensor, + sigma: torch.Tensor, + condition: ConditionBundle, + *, + height: int, + width: int, + grad_enabled: bool, + ) -> torch.Tensor: + """Predict Qwen flow velocity with explicit role and gradient ownership.""" + with ( + torch.profiler.record_function(f"distillation/{role}_forward"), + runtime.use_role(role, grad_enabled=grad_enabled) as module, + ): + module.eval() + timestep = sigma.reshape(-1) + if timestep.shape[0] == 1 and latents.shape[0] != 1: + timestep = timestep.expand(latents.shape[0]) + module_config = self.module_config(module) + guidance = None + if getattr(module_config, "guidance_embeds", False): + configured = self.model_config.pipeline.guidance_scale + if configured is None: + raise ValueError("Qwen guidance-embedded transformers require model.pipeline.guidance_scale.") + guidance = torch.full( + (latents.shape[0],), float(configured), device=latents.device, dtype=torch.float32 + ) + output = module( + hidden_states=latents, + timestep=timestep, + guidance=guidance, + encoder_hidden_states_mask=condition.masks.get("prompt_embeds"), + encoder_hidden_states=condition.tensors["prompt_embeds"], + img_shapes=build_img_shapes(height, width, latents.shape[0], QWEN_IMAGE_VAE_SCALE_FACTOR), + return_dict=False, + )[0] + if output.shape != latents.shape: + raise ValueError( + f"Qwen DMD transformer output shape {tuple(output.shape)} does not match latent shape " + f"{tuple(latents.shape)}." + ) + return output + + def rollout_sigmas(self, scheduler, height: int, width: int, device: torch.device) -> torch.Tensor: + """Return the inference-matched fixed-shift student sigma schedule.""" + del scheduler, height, width + steps = 1 if self.strategy == "one_step" else self.num_inference_steps + return build_qwen_dmd_sigmas(steps, self.rollout_timestep_shift, device=device) + + @torch.profiler.record_function("distillation/student_rollout") + def rollout( + self, + runtime: DistillationRoleRuntime, + condition: ConditionBundle, + initial_noise: torch.Tensor, + *, + height: int, + width: int, + grad_enabled: bool, + ) -> tuple[torch.Tensor, int, torch.Tensor, torch.Tensor]: + """Run to the shared exit step, retaining only its student graph.""" + scheduler = runtime.scheduler_for_role("student") + sigmas = self.rollout_sigmas(scheduler, height, width, initial_noise.device) + exit_index = ( + 0 + if self.strategy == "one_step" + else self.sample_rollout_exit(sigmas.numel() - 1, initial_noise.device, runtime) + ) + sample = initial_noise + x0 = None + for index in range(exit_index + 1): + sigma = sigmas[index].reshape(1) + use_grad = grad_enabled and index == exit_index + velocity = self.predict_velocity( + runtime, + "student", + sample, + sigma, + condition, + height=height, + width=width, + grad_enabled=use_grad, + ) + expanded_sigma = self.expand_sigma(sigma, sample) + x0 = velocity_to_x0(sample, velocity, expanded_sigma) + if index == exit_index: + break + sigma_next = sigmas[index + 1].reshape(1) + expanded_next = self.expand_sigma(sigma_next, sample) + if self.strategy == "consistency_renoise": + transition_noise = self.sample_noise(tuple(sample.shape), sample.device, runtime, "rollout_transition") + sample = consistency_renoise_step(x0, transition_noise, expanded_next) + else: + sample = ode_euler_step(sample, velocity, expanded_sigma, expanded_next) + assert x0 is not None + return x0, exit_index, sigmas[exit_index], sigmas[exit_index + 1] + + def sample_score_sigma(self, generated: torch.Tensor, runtime: DistillationRoleRuntime) -> torch.Tensor: + """Sample discrete shifted timesteps or continuous unshifted sigma values.""" + generator = self.generator_for_stream("score_sigma", generated.device, runtime) + if self.score_discrete_steps > 0: + scheduler = runtime.scheduler_for_role("student") + scheduler_config = getattr(scheduler, "config", {}) + num_train_timesteps = int(scheduler_config.get("num_train_timesteps", 1000)) + if self.score_discrete_steps != num_train_timesteps: + raise ValueError( + "score_discrete_steps must equal the Qwen scheduler num_train_timesteps; " + f"got {self.score_discrete_steps} and {num_train_timesteps}." + ) + timestep = torch.randint( + 0, + num_train_timesteps, + (generated.shape[0],), + device=generated.device, + generator=generator, + ).float() + sigma = timestep_shift(timestep, num_train_timesteps, self.score_timestep_shift) + sigma = sigma / num_train_timesteps + else: + sigma = torch.rand( + (generated.shape[0],), + device=generated.device, + dtype=torch.float32, + generator=generator, + ) + sigma = self.score_sigma_min + (self.score_sigma_max - self.score_sigma_min) * sigma + return sigma.clamp(self.score_sigma_min, self.score_sigma_max) + + def score_batch( + self, + generated: torch.Tensor, + runtime: DistillationRoleRuntime, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Re-noise detached generated latents using independent score RNG streams.""" + sigma = self.sample_score_sigma(generated, runtime) + noise = self.sample_noise(tuple(generated.shape), generated.device, runtime, "score_noise") + expanded = self.expand_sigma(sigma, generated) + noisy = (1.0 - expanded) * generated.detach().float() + expanded * noise + return noisy, noise, sigma + + @staticmethod + def prepare_condition_for_sequence_parallel( + runtime: DistillationRoleRuntime, + condition: ConditionBundle, + ) -> ConditionBundle: + """Reuse engine padding for sequence-parallel prompt embeddings.""" + engine = runtime.engine_for_role("student") + if not getattr(engine, "use_ulysses_sp", False): + return condition + embeds, mask = engine._pad_embeds_for_sp( + condition.tensors["prompt_embeds"], + condition.masks.get("prompt_embeds"), + engine.ulysses_sequence_parallel_size, + ) + return QwenImageConditionProvider.make_condition(embeds, mask) + + def student_loss( + self, + batch: TensorDict, + runtime: DistillationRoleRuntime, + condition: ConditionBundle, + negative_condition: ConditionBundle, + *, + height: int, + width: int, + latent_shape: tuple[int, ...], + ): + """Build DMD and optional paired-regression losses for the student only.""" + initial_noise = self.sample_noise( + latent_shape, condition.tensors["prompt_embeds"].device, runtime, "initial_noise" + ) + packed_noise = self.pack_latents(initial_noise) + rollout_start = time.perf_counter() + generated, exit_index, sigma_from, sigma_to = self.rollout( + runtime, + condition, + packed_noise, + height=height, + width=width, + grad_enabled=True, + ) + rollout_duration = time.perf_counter() - rollout_start + noisy, _, score_sigma = self.score_batch(generated, runtime) + with torch.no_grad(): + fake_start = time.perf_counter() + fake_velocity = self.predict_velocity( + runtime, + "fake_score", + noisy, + score_sigma, + condition, + height=height, + width=width, + grad_enabled=False, + ) + fake_duration = time.perf_counter() - fake_start + teacher_start = time.perf_counter() + teacher_positive = self.predict_velocity( + runtime, + "teacher_score", + noisy, + score_sigma, + condition, + height=height, + width=width, + grad_enabled=False, + ) + teacher_negative = self.predict_velocity( + runtime, + "teacher_score", + noisy, + score_sigma, + negative_condition, + height=height, + width=width, + grad_enabled=False, + ) + teacher_velocity = standard_cfg( + teacher_positive, + teacher_negative, + self.guidance_scale, + self.cfg_norm, + ) + expanded_sigma = self.expand_sigma(score_sigma, noisy) + fake_x0 = velocity_to_x0(noisy, fake_velocity, expanded_sigma) + teacher_x0 = velocity_to_x0(noisy, teacher_velocity, expanded_sigma) + gradient, normalizer, nonfinite = dmd_gradient( + fake_x0, + teacher_x0, + generated, + normalization_epsilon=self.normalization_epsilon, + ) + teacher_duration = time.perf_counter() - teacher_start + dmd_loss, active = dmd_surrogate_loss(generated, gradient) + total = self.dmd_loss_weight * dmd_loss + metrics = { + "dmd/loss": float(dmd_loss.detach()), + "dmd/normalizer": float(normalizer.detach().mean()), + "dmd/nonfinite": float(nonfinite), + "dmd/active_elements": float(active), + "rollout/exit_index": float(exit_index), + "rollout/sigma_from": float(sigma_from), + "rollout/sigma_to": float(sigma_to), + "score/sigma": float(score_sigma.detach().mean()), + "perf/student_rollout_s": rollout_duration, + "perf/fake_score_model_s": fake_duration, + "perf/teacher_score_model_s": teacher_duration, + } + if self.plan.name == "dmd": + regression_start = time.perf_counter() + regression_loss = self.regression_loss(batch, runtime, condition, height=height, width=width) + total = total + self.regression_loss_weight * regression_loss + metrics["regression/loss"] = float(regression_loss.detach()) + metrics["perf/regression_s"] = time.perf_counter() - regression_start + return total, metrics + + def fake_loss( + self, + runtime: DistillationRoleRuntime, + condition: ConditionBundle, + *, + height: int, + width: int, + latent_shape: tuple[int, ...], + ): + """Train the fake score to denoise detached student samples.""" + initial_noise = self.sample_noise( + latent_shape, condition.tensors["prompt_embeds"].device, runtime, "initial_noise" + ) + packed_noise = self.pack_latents(initial_noise) + rollout_start = time.perf_counter() + with torch.no_grad(): + generated, exit_index, sigma_from, sigma_to = self.rollout( + runtime, + condition, + packed_noise, + height=height, + width=width, + grad_enabled=False, + ) + generated = generated.detach() + rollout_duration = time.perf_counter() - rollout_start + noisy, score_noise, score_sigma = self.score_batch(generated, runtime) + fake_start = time.perf_counter() + fake_velocity = self.predict_velocity( + runtime, + "fake_score", + noisy, + score_sigma, + condition, + height=height, + width=width, + grad_enabled=True, + ) + fake_duration = time.perf_counter() - fake_start + loss, active = fake_score_loss(fake_velocity, score_noise, generated) + return loss, { + "fake_score/denoising_loss": float(loss.detach()), + "fake_score/active_elements": float(active), + "rollout/exit_index": float(exit_index), + "rollout/sigma_from": float(sigma_from), + "rollout/sigma_to": float(sigma_to), + "score/sigma": float(score_sigma.detach().mean()), + "perf/student_rollout_s": rollout_duration, + "perf/fake_score_model_s": fake_duration, + } + + def coerce_packed(self, value: Any, expected_shape: torch.Size, key: str, device: torch.device) -> torch.Tensor: + """Validate and pack a reference noise or teacher-target tensor.""" + if not isinstance(value, torch.Tensor): + raise ValueError(f"Qwen DMD regression requires tensor batch field {key!r}.") + value = value.to(device=device, dtype=torch.float32) + if value.ndim == 4: + value = value.unsqueeze(2) + if value.ndim == 5: + value = self.pack_latents(value) + if value.shape != expected_shape: + raise ValueError( + f"Qwen DMD {key} shape {tuple(value.shape)} does not match generated latent shape " + f"{tuple(expected_shape)}." + ) + return value + + @torch.profiler.record_function("distillation/regression") + def regression_loss( + self, + batch: TensorDict, + runtime: DistillationRoleRuntime, + condition: ConditionBundle, + *, + height: int, + width: int, + ) -> torch.Tensor: + """Regress a paired student sample against its teacher target.""" + manifest = tu.get(batch, "teacher_sampling_manifest") + manifests = [manifest] if isinstance(manifest, Mapping) and batch.batch_size[0] == 1 else manifest + if ( + not isinstance(manifests, list) + or len(manifests) != batch.batch_size[0] + or any(not isinstance(item, Mapping) or not item for item in manifests) + ): + raise ValueError("Original DMD regression requires teacher_sampling_manifest provenance for every sample.") + reference_noise = tu.get(batch, "reference_noise") + target_latents = tu.get(batch, "teacher_target_latents") + target_pixels = tu.get(batch, "teacher_target_pixels") + if target_latents is None and target_pixels is None: + raise ValueError("Original DMD requires teacher_target_latents or teacher_target_pixels.") + if target_latents is not None and target_pixels is not None: + raise ValueError("Provide only one of teacher_target_latents and teacher_target_pixels.") + + with runtime.use_role("student", grad_enabled=True) as module: + _, _, in_channels, latent_shape = self.latent_geometry(batch, module) + expected_packed = torch.Size((latent_shape[0], (latent_shape[-2] // 2) * (latent_shape[-1] // 2), in_channels)) + reference_noise = self.coerce_packed( + reference_noise, + expected_packed, + "reference_noise", + condition.tensors["prompt_embeds"].device, + ) + prediction, _, _, _ = self.rollout( + runtime, + condition, + reference_noise, + height=height, + width=width, + grad_enabled=True, + ) + if self.regression_type == "latent_mse": + if target_pixels is not None: + raise ValueError("regression_type='latent_mse' requires teacher_target_latents.") + target = self.coerce_packed(target_latents, prediction.shape, "teacher_target_latents", prediction.device) + return torch.mean((prediction.float() - target.detach().float()) ** 2) + return self.decoded_lpips_loss(prediction, target_latents, target_pixels, height=height, width=width) + + def ensure_vae_and_lpips(self, device: torch.device): + """Load frozen VAE and perceptual-loss modules only for original DMD.""" + if self._vae is None: + from diffusers import AutoencoderKLQwenImage + + self._vae = AutoencoderKLQwenImage.from_pretrained( + self.model_config.local_path or self.model_config.path, + subfolder="vae", + torch_dtype=torch.float32, + local_files_only=os.path.isdir(self.model_config.local_path or self.model_config.path), + ).to(device) + self._vae.requires_grad_(False) + self._vae.eval() + if self._lpips is None: + try: + import piq + except ImportError as exc: + raise ImportError( + "regression_type='decoded_lpips' requires piq; install verl-omni with the distillation extra." + ) from exc + self._lpips = piq.LPIPS(replace_pooling=True, reduction="none").to(device) + self._lpips.requires_grad_(False) + self._lpips.eval() + return self._vae, self._lpips + + def decode_latents(self, packed: torch.Tensor, *, height: int, width: int) -> torch.Tensor: + """Undo Qwen latent normalization and decode differentiably to RGB.""" + from diffusers import QwenImagePipeline + + vae, _ = self.ensure_vae_and_lpips(packed.device) + latent = QwenImagePipeline._unpack_latents( + packed, + height=height, + width=width, + vae_scale_factor=QWEN_IMAGE_VAE_SCALE_FACTOR, + ).float() + shape = (1, vae.config.z_dim, 1, 1, 1) + mean = latent.new_tensor(vae.config.latents_mean).view(shape) + std = latent.new_tensor(vae.config.latents_std).view(shape) + decoded = vae.decode(latent * std + mean, return_dict=False)[0][:, :, 0] + return decoded.float().mul(0.5).add(0.5) + + def decoded_lpips_loss( + self, + prediction: torch.Tensor, + target_latents: Optional[torch.Tensor], + target_pixels: Optional[torch.Tensor], + *, + height: int, + width: int, + ) -> torch.Tensor: + """Apply perceptual regression with gradients only through the prediction.""" + _, lpips = self.ensure_vae_and_lpips(prediction.device) + prediction_pixels = self.decode_latents(prediction, height=height, width=width) + if target_pixels is not None: + target = target_pixels.to(device=prediction.device, dtype=torch.float32) + if target.ndim == 3: + target = target.unsqueeze(0) + if target.shape != prediction_pixels.shape: + raise ValueError( + f"teacher_target_pixels shape {tuple(target.shape)} does not match decoded prediction " + f"{tuple(prediction_pixels.shape)}." + ) + if torch.any((target < 0) | (target > 1)): + raise ValueError("teacher_target_pixels must be normalized to [0, 1].") + else: + target = self.coerce_packed( + target_latents, + prediction.shape, + "teacher_target_latents", + prediction.device, + ) + with torch.no_grad(): + target = self.decode_latents(target, height=height, width=width) + return lpips(prediction_pixels, target.detach()).mean() + + def compute_phase( + self, + request: PhaseRequest, + batch: TensorDict, + runtime: DistillationRoleRuntime, + ) -> DistillationPhaseComputation: + """Build the requested role loss and detached metrics for one micro-batch.""" + from verl_omni.workers.diffusion_distillation_worker import DistillationPhaseComputation + + if request.kind not in {"student", "fake_score"}: + raise ValueError(f"Unsupported Qwen DMD phase {request.kind!r}.") + with runtime.use_role("student", grad_enabled=False) as module: + height, width, _, latent_shape = self.latent_geometry(batch, module) + dtype = self.module_dtype(module) + device = next(module.parameters()).device + condition_start = time.perf_counter() + condition, negative_condition = self.condition_provider.encode( + batch, + device=device, + dtype=dtype, + require_negative=request.kind == "student", + ) + condition = self.prepare_condition_for_sequence_parallel(runtime, condition) + if negative_condition is not None: + negative_condition = self.prepare_condition_for_sequence_parallel(runtime, negative_condition) + condition_duration = time.perf_counter() - condition_start + if request.kind == "student": + if negative_condition is None: + raise ValueError("Qwen DMD student phases require negative teacher conditioning.") + loss, metrics = self.student_loss( + batch, + runtime, + condition, + negative_condition, + height=height, + width=width, + latent_shape=latent_shape, + ) + metrics["perf/condition_encode_s"] = condition_duration + return DistillationPhaseComputation(losses={"student": loss}, metrics=metrics) + loss, metrics = self.fake_loss( + runtime, + condition, + height=height, + width=width, + latent_shape=latent_shape, + ) + metrics["perf/condition_encode_s"] = condition_duration + return DistillationPhaseComputation(losses={"fake_score": loss}, metrics=metrics) + + def state_dict(self) -> dict: + """Return independent worker-local RNG stream states.""" + states = {name: generator.get_state().cpu() for name, generator in self._generators.items()} + states.update(self._pending_generator_states) + return {"version": 1, "rng_seed": self.rng_seed, "generator_states": states} + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore worker-local rollout and score-noise RNG streams.""" + if state.get("version") != 1 or state.get("rng_seed") != self.rng_seed: + raise ValueError("Qwen DMD phase-runner checkpoint is incompatible with the active RNG configuration.") + generator_states = state.get("generator_states") + if not isinstance(generator_states, Mapping) or set(generator_states) - set(self.STREAM_OFFSETS): + raise ValueError("Qwen DMD phase-runner checkpoint contains invalid RNG streams.") + self._generators.clear() + self._pending_generator_states.clear() + for name, generator_state in generator_states.items(): + if not isinstance(generator_state, torch.Tensor): + raise TypeError(f"Qwen DMD RNG state for {name!r} must be a tensor.") + self._pending_generator_states[name] = generator_state + + +@DiffusionModelBase.register("QwenImagePipeline", algorithm="dmd") +@DiffusionModelBase.register("QwenImagePipeline", algorithm="dmd2") +class QwenImageDistributionMatching(QwenImage, DistributionMatchingModelAdapter): + """Qwen-Image architecture adapter for DMD and DMD2.""" + + @classmethod + def build_distribution_matching_computer( + cls, + model_config: DiffusionModelConfig, + plan: DistillationPlan, + ) -> QwenImageDMDComputer: + """Build the architecture-owned Qwen DMD phase program.""" + return QwenImageDMDComputer(model_config, plan) diff --git a/verl_omni/pipelines/qwen_image_distillation/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/qwen_image_distillation/vllm_omni_rollout_adapter.py new file mode 100644 index 000000000..dffe6be3a --- /dev/null +++ b/verl_omni/pipelines/qwen_image_distillation/vllm_omni_rollout_adapter.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. +"""Qwen-Image inference adapter for DMD and DMD2 students.""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from verl_omni.pipelines.model_base import VllmOmniPipelineBase +from verl_omni.pipelines.qwen_image_distillation.diffusers_training_adapter import build_qwen_dmd_sigmas +from verl_omni.pipelines.qwen_image_flow_grpo.vllm_omni_rollout_adapter import QwenImagePipelineWithLogProb + +__all__ = ["QwenImageDMDPipeline"] + + +@VllmOmniPipelineBase.register("QwenImagePipeline", algorithm="dmd") +@VllmOmniPipelineBase.register("QwenImagePipeline", algorithm="dmd2") +class QwenImageDMDPipeline(QwenImagePipelineWithLogProb): + """Qwen-Image rollout with the same fixed-shift schedule used for DMD training.""" + + supports_request_batch = True + supports_step_execution = False + rollout_timestep_shift = 3.0 + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): + """Reuse Qwen packing and request-local generators with training-matched fp32 noise.""" + return super().prepare_latents( + batch_size, num_channels_latents, height, width, torch.float32, device, generator, latents + ) + + def prepare_encode(self, state, **kwargs): + """Reject step mode until its DMD defaults and per-request schedules are validated.""" + raise NotImplementedError("Qwen DMD request batching requires step_execution=false.") + + def forward(self, req: OmniDiffusionRequest | DiffusionRequestBatch, *args, **kwargs): + """Apply one batch-consistent DMD rollout shift before normal Qwen generation.""" + requests = req.requests if isinstance(req, DiffusionRequestBatch) else [req] + if not requests: + raise ValueError("Qwen DMD request batches cannot be empty.") + if len(requests) > 1: + key = build_request_batch_sampling_params_key(requests[0]) + if any(build_request_batch_sampling_params_key(request) != key for request in requests[1:]): + raise ValueError("Packed Qwen DMD requests must have compatible sampling parameters.") + extra_args = [request.sampling_params.extra_args or {} for request in requests] + shifts = { + 3.0 if values.get("rollout_timestep_shift") is None else float(values["rollout_timestep_shift"]) + for values in extra_args + } + if len(shifts) != 1: + raise ValueError("Packed Qwen DMD requests must use the same rollout_timestep_shift.") + shift = shifts.pop() + if not math.isfinite(shift) or shift < 1: + raise ValueError(f"rollout_timestep_shift must be finite and at least 1, got {shift}.") + default_noise_level = float(kwargs.get("noise_level", 0.0)) + noise_levels = { + default_noise_level if values.get("noise_level") is None else float(values["noise_level"]) + for values in extra_args + } + if noise_levels != {0.0}: + raise ValueError("Qwen DMD inference requires noise_level=0 for deterministic Euler sampling.") + kwargs.setdefault("noise_level", 0.0) + kwargs.setdefault("logprobs", False) + kwargs.setdefault("true_cfg_scale", 1.0) + previous_shift = self.rollout_timestep_shift + self.rollout_timestep_shift = shift + try: + return super().forward(req, *args, **kwargs) + finally: + self.rollout_timestep_shift = previous_shift + + def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): + """Build the fixed linear-shift schedule shared with the training phase runner.""" + del image_seq_len + if num_inference_steps <= 0: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}.") + if sigmas is None: + sigmas = build_qwen_dmd_sigmas(num_inference_steps, self.rollout_timestep_shift)[:-1].numpy() + else: + sigmas = np.asarray(sigmas, dtype=np.float32) + if sigmas.ndim != 1 or len(sigmas) != num_inference_steps: + raise ValueError(f"sigmas must contain exactly {num_inference_steps} values, got {sigmas.shape}.") + if not np.all(np.isfinite(sigmas)) or np.any(sigmas <= 0) or np.any(sigmas > 1): + raise ValueError("Qwen DMD sigmas must be finite and lie in (0, 1].") + if np.any(sigmas[:-1] < sigmas[1:]): + raise ValueError("Qwen DMD sigmas must be monotonically non-increasing.") + + device = self.device + sigma_tensor = torch.as_tensor(sigmas, device=device, dtype=torch.float32) + self.scheduler.num_inference_steps = num_inference_steps + self.scheduler.sigmas = torch.cat((sigma_tensor, sigma_tensor.new_zeros(1))) + self.scheduler.timesteps = sigma_tensor * self.scheduler.config.get("num_train_timesteps", 1000) + self.scheduler._step_index = None + self.scheduler._begin_index = None + return self.scheduler.timesteps, num_inference_steps diff --git a/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py b/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py index 3cc7bbbb7..c859163f2 100644 --- a/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen_image_flow_grpo/vllm_omni_rollout_adapter.py @@ -797,11 +797,12 @@ def forward( max_sequence_length = sampling_params.max_sequence_length or max_sequence_length output_type = sampling_params.output_type or output_type - noise_level = coalesce_not_none(sampling_params.extra_args.get("noise_level", None), noise_level) - sde_window_size = coalesce_not_none(sampling_params.extra_args.get("sde_window_size", None), sde_window_size) - sde_window_range = coalesce_not_none(sampling_params.extra_args.get("sde_window_range", None), sde_window_range) - sde_type = coalesce_not_none(sampling_params.extra_args.get("sde_type", None), sde_type) - logprobs = coalesce_not_none(sampling_params.extra_args.get("logprobs", None), logprobs) + extra_args = sampling_params.extra_args or {} + noise_level = coalesce_not_none(extra_args.get("noise_level", None), noise_level) + sde_window_size = coalesce_not_none(extra_args.get("sde_window_size", None), sde_window_size) + sde_window_range = coalesce_not_none(extra_args.get("sde_window_range", None), sde_window_range) + sde_type = coalesce_not_none(extra_args.get("sde_type", None), sde_type) + logprobs = coalesce_not_none(extra_args.get("logprobs", None), logprobs) for request in request_batch.requests: request_sampling_params = request.sampling_params @@ -870,7 +871,7 @@ def forward( generator, latents, ) - img_shapes = build_img_shapes(height, width, batch_size, self.vae_scale_factor) + img_shapes = build_img_shapes(height, width, latents.shape[0], self.vae_scale_factor) timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) self._num_timesteps = len(timesteps) diff --git a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml index 5828fdcd2..5f5fa7f8f 100644 --- a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml @@ -241,6 +241,7 @@ actor_rollout_ref: sde_window_size: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_size,null} sde_window_range: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_range,null} sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} algo: _target_: verl_omni.workers.config.diffusion.DiffusionRolloutAlgoConfig noise_level: 1.0 @@ -248,6 +249,7 @@ actor_rollout_ref: sde_window_size: null sde_window_range: null sde_contiguous: true + rollout_timestep_shift: null sample_strategy: random iters_per_group: 1 sde_window_seed: 0 @@ -345,6 +347,7 @@ actor_rollout_ref: sde_window_size: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_size,null} sde_window_range: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_range,null} sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} model_type: diffusion_model separate: false hybrid_engine: true @@ -527,6 +530,20 @@ distillation: override_optimizer_config: null ema_decay: 0.999 ema_start_step: 0 + conditioning_provider: local_frozen_encoder + negative_prompt: ' ' + teacher_guidance_scale: 4.0 + teacher_cfg_norm: layer_norm + rollout_timestep_shift: 3.0 + score_sigma_min: 0.02 + score_sigma_max: 0.98 + score_timestep_shift: 3.0 + score_discrete_steps: 1000 + normalization_epsilon: 1.0e-05 + dmd_loss_weight: 1.0 + regression_type: decoded_lpips + regression_loss_weight: 1.0 + rng_seed: 0 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 63afe5a97..b1c591924 100644 --- a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml @@ -282,6 +282,7 @@ actor_rollout_ref: sde_window_size: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_size,null} sde_window_range: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_range,null} sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} algo: _target_: verl_omni.workers.config.diffusion.DiffusionRolloutAlgoConfig noise_level: 1.0 @@ -289,6 +290,7 @@ actor_rollout_ref: sde_window_size: null sde_window_range: null sde_contiguous: true + rollout_timestep_shift: null sample_strategy: random iters_per_group: 1 sde_window_seed: 0 @@ -386,6 +388,7 @@ actor_rollout_ref: sde_window_size: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_size,null} sde_window_range: ${oc.select:actor_rollout_ref.rollout.algo.sde_window_range,null} sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} model_type: diffusion_model separate: false hybrid_engine: true @@ -568,6 +571,20 @@ distillation: override_optimizer_config: null ema_decay: 0.999 ema_start_step: 0 + conditioning_provider: local_frozen_encoder + negative_prompt: ' ' + teacher_guidance_scale: 4.0 + teacher_cfg_norm: layer_norm + rollout_timestep_shift: 3.0 + score_sigma_min: 0.02 + score_sigma_max: 0.98 + score_timestep_shift: 3.0 + score_discrete_steps: 1000 + normalization_epsilon: 1.0e-05 + dmd_loss_weight: 1.0 + regression_type: decoded_lpips + regression_loss_weight: 1.0 + rng_seed: 0 algorithm: _target_: verl_omni.trainer.config.DiffusionAlgoConfig trainer_type: policy_gradient diff --git a/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml b/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml index 02b5bd2d3..6f5a540ae 100644 --- a/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml +++ b/verl_omni/trainer/config/diffusion/distillation/diffusion_distillation.yaml @@ -137,3 +137,45 @@ distribution_matching: # First completed student step that updates EMA ema_start_step: 0 + + # Conditioning source: local_frozen_encoder or precomputed + conditioning_provider: local_frozen_encoder + + # Negative prompt used by the guided frozen teacher + negative_prompt: " " + + # Standard CFG scale for frozen-teacher scoring + teacher_guidance_scale: 4.0 + + # Teacher CFG normalization: none, layer_norm, or scalar + teacher_cfg_norm: layer_norm + + # Linear time shift used by the few-step student rollout schedule + rollout_timestep_shift: 3.0 + + # Minimum score-noising sigma + score_sigma_min: 0.02 + + # Maximum score-noising sigma + score_sigma_max: 0.98 + + # Rational shift used for score-noising samples + score_timestep_shift: 3.0 + + # Number of discrete score-noising samples; zero selects continuous sampling + score_discrete_steps: 1000 + + # Stabilizer for distribution-matching normalization + normalization_epsilon: 1.0e-5 + + # Distribution-matching loss weight + dmd_loss_weight: 1.0 + + # Original-DMD paired-regression distance; latent_mse is a non-paper diagnostic + regression_type: decoded_lpips + + # Original-DMD paired-regression loss weight + regression_loss_weight: 1.0 + + # Base seed for worker-local rollout and score-noise generators + rng_seed: 0 diff --git a/verl_omni/trainer/config/diffusion/model/diffusion_model.yaml b/verl_omni/trainer/config/diffusion/model/diffusion_model.yaml index 12fff623a..c0c886d7d 100644 --- a/verl_omni/trainer/config/diffusion/model/diffusion_model.yaml +++ b/verl_omni/trainer/config/diffusion/model/diffusion_model.yaml @@ -147,3 +147,6 @@ algo: # Whether selected SDE steps must form one consecutive window sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + + # Optional fixed linear time shift for distilled-model inference + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} diff --git a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml index 896832f5e..7a35d3ca1 100644 --- a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml +++ b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml @@ -218,6 +218,9 @@ val_kwargs: # Whether selected SDE steps must form one consecutive window sde_contiguous: ${oc.select:actor_rollout_ref.rollout.algo.sde_contiguous,true} + # Optional fixed linear time shift for distilled-model validation + rollout_timestep_shift: ${oc.select:actor_rollout_ref.rollout.algo.rollout_timestep_shift,null} + # Rollout Algorithm configuration algo: @@ -239,6 +242,9 @@ algo: # Whether selected SDE steps must form one consecutive window sde_contiguous: true + # Optional fixed linear time shift for distilled-model inference + rollout_timestep_shift: null + # MixGRPO sliding-window scheduler configs # Sliding strategy: "random" (default) draws a fresh window per step; diff --git a/verl_omni/trainer/diffusion/distillation/controller.py b/verl_omni/trainer/diffusion/distillation/controller.py index c6345e57c..c2e3ae1c6 100644 --- a/verl_omni/trainer/diffusion/distillation/controller.py +++ b/verl_omni/trainer/diffusion/distillation/controller.py @@ -120,6 +120,7 @@ def run_cycle(self) -> UpdateCycle: ) before_metrics = dict(self.phase_metrics) try: + self.phase_metrics = {} cycle = self.plan.update_schedule.next_cycle(self.counters) student_step_reported = self.drive_requests(cycle.requests) @@ -180,10 +181,12 @@ def validate_result(result: PhaseResult, request: PhaseRequest) -> None: 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.""" + """Record validated counters and every phase result in the current cycle.""" 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) + count = sum(name.split("/")[0] == request.kind for name in self.phase_metrics) + phase_name = request.kind if count == 0 else f"{request.kind}/{count}" + self.phase_metrics[phase_name] = dict(result.metrics) def assert_progress(self, before: TrainerCounters) -> None: """A cycle must advance global_step or at least one role optimizer counter.""" @@ -195,7 +198,7 @@ def assert_progress(self, before: TrainerCounters) -> None: @property def metrics(self) -> dict[str, dict]: - """Metrics recorded for the most recent phase of each kind.""" + """Metrics for every phase in the last completed cycle, including repeats.""" return self.phase_metrics def state_dict(self) -> dict[str, Any]: diff --git a/verl_omni/trainer/diffusion/distillation/ray_trainer.py b/verl_omni/trainer/diffusion/distillation/ray_trainer.py index 20df259c4..0b4da0fd8 100644 --- a/verl_omni/trainer/diffusion/distillation/ray_trainer.py +++ b/verl_omni/trainer/diffusion/distillation/ray_trainer.py @@ -36,6 +36,8 @@ from verl.trainer.ppo.ray_trainer import Role from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path from verl.utils.config import omega_conf_to_dataclass +from verl.utils.metric import Metric, reduce_metrics +from verl.utils.profiler import marked_timer from verl.utils.tracking import Tracking from verl_omni.pipelines.model_base import DiffusionModelBase, DistributionMatchingModelAdapter @@ -188,6 +190,13 @@ def __init__( def validate_runtime_config(self) -> None: """Reject unsupported role storage and distributed batch layouts.""" distribution_matching = self.config.distillation.distribution_matching + model_algorithm = OmegaConf.select(self.config, "actor_rollout_ref.model.algorithm") + recipe = OmegaConf.select(self.config, "distillation.distribution_matching.recipe") + if model_algorithm is not None and recipe is not None and model_algorithm != recipe: + raise ValueError( + "actor_rollout_ref.model.algorithm must match distillation.distribution_matching.recipe; " + f"got {model_algorithm!r} and {recipe!r}." + ) strategy = self.config.actor_rollout_ref.actor.strategy if strategy not in {"fsdp", "fsdp2"}: raise ValueError(f"Distillation role groups require strategy 'fsdp' or 'fsdp2', got {strategy!r}.") @@ -413,8 +422,26 @@ def _load_checkpoint(self) -> int: @staticmethod def flatten_metrics(metrics: dict[str, dict]) -> dict[str, float]: - """Flatten phase metrics for the existing tracking backends.""" - return {key: value for phase_metrics in metrics.values() for key, value in phase_metrics.items()} + """Sum cycle timings/counts, retain peak memory, and average other phase values.""" + collected: dict[str, Metric] = {} + result = {} + for phase, phase_metrics in metrics.items(): + for key, value in phase_metrics.items(): + result[f"phase/{phase}/{key}"] = float(value) + is_duration = ( + key.startswith(("perf/", "perf_max_rank/")) and key.endswith("_s") and not key.endswith("_per_s") + ) + if is_duration or key.endswith(("/active_elements", "/nonfinite", "_samples", "_micro_batches")): + aggregation = "sum" + elif key.startswith("memory/max_"): + aggregation = "max" + else: + aggregation = "mean" + if key not in collected: + collected[key] = Metric(aggregation=aggregation) + collected[key].append(value) + result.update({key: float(value) for key, value in reduce_metrics(collected).items()}) + return result def profile_workers(self, *, start: bool, step: int) -> None: """Start or stop the configured distributed profiler.""" @@ -455,18 +482,35 @@ def fit(self, num_cycles: Optional[int] = None) -> None: while controller.counters.global_step < target_steps: before_global_step = controller.counters.global_step profile_step = before_global_step + 1 - do_profile = profile_steps is not None and profile_step in profile_steps - self.profile_workers(start=do_profile, step=profile_step) + do_profile = ( + profile_steps is not None + and profile_step in profile_steps + and controller.counters.completed_cycles >= self.plan.update_schedule.warmup_cycles + ) + if do_profile: + self.profile_workers(start=True, step=profile_step) + timing_raw = {} try: - controller.run_cycle() + with marked_timer("perf/cycle_s", timing_raw): + controller.run_cycle() finally: if do_profile: self.profile_workers(start=False, step=profile_step) self.global_steps = controller.counters.global_step metrics = self.flatten_metrics(controller.metrics) + metrics.update(timing_raw) + for kind in ("student", "fake_score"): + samples = metrics.setdefault(f"training/{kind}_samples", 0.0) + metrics[f"perf/{kind}_samples_per_s"] = samples / timing_raw["perf/cycle_s"] + metrics.update( + { + f"training/{role}_optimizer_steps": float(steps) + for role, steps in controller.counters.optimizer_steps.items() + } + ) metrics["training/global_step"] = float(self.global_steps) metrics["training/completed_cycles"] = float(controller.counters.completed_cycles) - self._logger.log(data=metrics, step=self.global_steps) + self._logger.log(data=metrics, step=controller.counters.completed_cycles) if self.global_steps > before_global_step: progress_bar.update(1) if hasattr(self.train_dataset, "on_batch_end"): diff --git a/verl_omni/trainer/diffusion/distillation/recipes.py b/verl_omni/trainer/diffusion/distillation/recipes.py index c26e3bf77..52a4ec60a 100644 --- a/verl_omni/trainer/diffusion/distillation/recipes.py +++ b/verl_omni/trainer/diffusion/distillation/recipes.py @@ -385,6 +385,43 @@ def require_choice(field: str, value: str, valid_values: set[str]) -> str: return value +def dmd_data_requirements(config, mode: str) -> dict: + """Declare the selected conditioning and sample-data contract.""" + return { + "mode": mode, + "conditioning_provider": get_config_or_default(config, "conditioning_provider", "local_frozen_encoder"), + "negative_prompt": get_config_or_default(config, "negative_prompt", " "), + } + + +def dmd_objective(config, *, name: str, profile: str, **extra) -> dict: + """Declare guidance and objective weights without executing loss math.""" + return { + "name": name, + "profile": profile, + "teacher_guidance_scale": float(get_config_or_default(config, "teacher_guidance_scale", 4.0)), + "teacher_cfg_norm": get_config_or_default(config, "teacher_cfg_norm", "layer_norm"), + "normalization_epsilon": float(get_config_or_default(config, "normalization_epsilon", 1e-5)), + "dmd_loss_weight": float(get_config_or_default(config, "dmd_loss_weight", 1.0)), + "regression_type": get_config_or_default(config, "regression_type", "decoded_lpips"), + "regression_loss_weight": float(get_config_or_default(config, "regression_loss_weight", 1.0)), + **extra, + } + + +def dmd_rollout(config, strategy: str) -> dict: + """Declare deterministic rollout and score-noise sampling settings.""" + return { + "strategy": strategy, + "rollout_timestep_shift": float(get_config_or_default(config, "rollout_timestep_shift", 3.0)), + "score_sigma_min": float(get_config_or_default(config, "score_sigma_min", 0.02)), + "score_sigma_max": float(get_config_or_default(config, "score_sigma_max", 0.98)), + "score_timestep_shift": float(get_config_or_default(config, "score_timestep_shift", 3.0)), + "score_discrete_steps": int(get_config_or_default(config, "score_discrete_steps", 1000)), + "rng_seed": int(get_config_or_default(config, "rng_seed", 0)), + } + + def common_plan_kwargs( config, *, @@ -424,9 +461,9 @@ def build_plan(cls, config, capabilities) -> DistillationPlan: shared_base_layout(model_ref), get_config_or_default(config, "role_storage", "shared_base_adapters"), ), - data_requirements={"mode": data_mode}, - objective={"name": "dmd", "profile": profile}, - rollout={"strategy": rollout}, + data_requirements=dmd_data_requirements(config, data_mode), + objective=dmd_objective(config, name="dmd", profile=profile), + rollout=dmd_rollout(config, rollout), initialization={"stage": "base"}, required_capabilities=frozenset({"distribution_matching"}), **common_plan_kwargs(config, default_fake_repeats=1), @@ -461,9 +498,14 @@ def build_plan(cls, config, capabilities) -> DistillationPlan: shared_base_layout(model_ref, with_discriminator=adversarial), get_config_or_default(config, "role_storage", "shared_base_adapters"), ), - data_requirements={"mode": data_mode}, - objective={"name": "dmd2", "profile": profile, "adversarial": adversarial}, - rollout={"strategy": rollout}, + data_requirements=dmd_data_requirements(config, data_mode), + objective=dmd_objective( + config, + name="dmd2", + profile=profile, + adversarial=adversarial, + ), + rollout=dmd_rollout(config, rollout), initialization={"stage": "base"}, required_capabilities=frozenset({"distribution_matching"} | ({"adversarial"} if adversarial else set())), **common_plan_kwargs(config, with_discriminator=adversarial), @@ -578,7 +620,28 @@ def build_plan_from_config(config, capabilities) -> DistillationPlan: "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", "role_storage"): + optional_keys = ( + "profile", + "fake_update_ratio", + "rollout_strategy", + "data_mode", + "role_storage", + "conditioning_provider", + "negative_prompt", + "teacher_guidance_scale", + "teacher_cfg_norm", + "rollout_timestep_shift", + "score_sigma_min", + "score_sigma_max", + "score_timestep_shift", + "score_discrete_steps", + "normalization_epsilon", + "dmd_loss_weight", + "regression_type", + "regression_loss_weight", + "rng_seed", + ) + for optional_key in optional_keys: value = get_config_value(distribution_matching, optional_key) if value is not None: recipe_config[optional_key] = value diff --git a/verl_omni/utils/dataset/qwen_image_distillation_dataset.py b/verl_omni/utils/dataset/qwen_image_distillation_dataset.py new file mode 100644 index 000000000..4f10ecbb3 --- /dev/null +++ b/verl_omni/utils/dataset/qwen_image_distillation_dataset.py @@ -0,0 +1,79 @@ +# 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. +"""Dataset adapter for original-DMD Qwen-Image regression pairs.""" + +from __future__ import annotations + +import io +import os +from typing import Any + +import numpy as np +import torch +from verl.utils.dataset.rl_dataset import RLHFDataset + +__all__ = ["QwenImageDMDPairDataset"] + + +def load_float_tensor(value: Any, field: str) -> torch.Tensor: + """Load a non-empty detached fp32 tensor without arbitrary pickle execution.""" + if isinstance(value, torch.Tensor): + tensor = value + elif isinstance(value, bytes | bytearray | memoryview): + buffer = io.BytesIO(bytes(value)) + tensor = torch.load(buffer, map_location="cpu", weights_only=True) + elif isinstance(value, str): + path = os.path.expanduser(value) + if not os.path.isfile(path): + raise FileNotFoundError(f"DMD tensor path for {field!r} does not exist: {path}") + tensor = torch.load(path, map_location="cpu", weights_only=True) + else: + tensor = torch.as_tensor(np.asarray(value)) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"DMD field {field!r} must resolve to a tensor, got {type(tensor)}.") + if tensor.numel() == 0: + raise ValueError(f"DMD field {field!r} must not be empty.") + return tensor.detach().float() + + +def is_present(value: Any) -> bool: + """Treat missing parquet cells and NaN placeholders as absent targets.""" + if value is None: + return False + if isinstance(value, float) and np.isnan(value): + return False + return True + + +class QwenImageDMDPairDataset(RLHFDataset): + """Load prompt, reference-noise, target, and provenance fields for original DMD.""" + + def __getitem__(self, item: int) -> dict[str, Any]: + row = super().__getitem__(item) + if "reference_noise" not in row: + raise ValueError("Original-DMD rows require reference_noise.") + has_latents = is_present(row.get("teacher_target_latents")) + has_pixels = is_present(row.get("teacher_target_pixels")) + if has_latents == has_pixels: + raise ValueError("Original-DMD rows require exactly one teacher target: latents or pixels.") + manifest = row.get("teacher_sampling_manifest") + if not isinstance(manifest, dict) or not manifest: + raise ValueError("Original-DMD rows require a non-empty teacher_sampling_manifest mapping.") + + row["reference_noise"] = load_float_tensor(row["reference_noise"], "reference_noise") + target_key = "teacher_target_latents" if has_latents else "teacher_target_pixels" + row[target_key] = load_float_tensor(row[target_key], target_key) + row.pop("teacher_target_pixels" if has_latents else "teacher_target_latents", None) + row["pair_id"] = str(row.get("pair_id", row.get("index", item))) + return row diff --git a/verl_omni/workers/config/diffusion/distillation.py b/verl_omni/workers/config/diffusion/distillation.py index 3832d394b..ad37e1619 100644 --- a/verl_omni/workers/config/diffusion/distillation.py +++ b/verl_omni/workers/config/diffusion/distillation.py @@ -91,6 +91,33 @@ class DiffusionDistributionMatchingConfig(BaseConfig): ema_decay: float = 0.999 # First completed student step that updates EMA. ema_start_step: int = 0 + # Conditioning source used by architecture phase runners. + conditioning_provider: str = "local_frozen_encoder" + # Negative prompt used by the guided frozen teacher. + negative_prompt: str = " " + # Standard CFG scale for the frozen teacher score. + teacher_guidance_scale: float = 4.0 + # Teacher CFG normalization mode. + teacher_cfg_norm: str = "layer_norm" + # Linear time shift used by the few-step student rollout schedule. + rollout_timestep_shift: float = 3.0 + # Lower and upper score-noising sigma bounds. + score_sigma_min: float = 0.02 + score_sigma_max: float = 0.98 + # Rational time shift applied to score-noising samples. + score_timestep_shift: float = 3.0 + # Number of discrete score-noising samples; zero selects continuous sampling. + score_discrete_steps: int = 1000 + # Stabilizer for the DMD score-difference normalizer. + normalization_epsilon: float = 1e-5 + # Distribution-matching loss weight. + dmd_loss_weight: float = 1.0 + # DMD paired-regression distance. + regression_type: str = "decoded_lpips" + # DMD paired-regression loss weight. + regression_loss_weight: float = 1.0 + # Base seed for worker-local rollout and score-noise generators. + rng_seed: int = 0 def __post_init__(self): valid_recipes = {"dmd", "dmd2", "causvid", "self_forcing"} @@ -137,6 +164,52 @@ def __post_init__(self): raise ValueError(f"ema_decay must be in [0, 1], got {self.ema_decay}") if self.ema_start_step < 0: raise ValueError(f"ema_start_step must be non-negative, got {self.ema_start_step}") + valid_conditioning_providers = {"local_frozen_encoder", "precomputed"} + if self.conditioning_provider not in valid_conditioning_providers: + raise ValueError( + f"Invalid conditioning_provider: {self.conditioning_provider}. " + f"Must be one of {sorted(valid_conditioning_providers)}" + ) + if not isinstance(self.negative_prompt, str): + raise ValueError("negative_prompt must be a string") + if self.teacher_guidance_scale <= 1.0: + raise ValueError( + f"teacher_guidance_scale must be greater than 1 for guided teacher scoring, " + f"got {self.teacher_guidance_scale}" + ) + valid_cfg_norms = {"none", "layer_norm", "scalar"} + if self.teacher_cfg_norm not in valid_cfg_norms: + raise ValueError( + f"Invalid teacher_cfg_norm: {self.teacher_cfg_norm}. Must be one of {sorted(valid_cfg_norms)}" + ) + if self.rollout_timestep_shift < 1: + raise ValueError(f"rollout_timestep_shift must be at least 1, got {self.rollout_timestep_shift}") + if not 0.0 <= self.score_sigma_min < self.score_sigma_max <= 1.0: + raise ValueError( + "score sigma bounds must satisfy 0 <= score_sigma_min < score_sigma_max <= 1, " + f"got [{self.score_sigma_min}, {self.score_sigma_max}]" + ) + if self.score_timestep_shift < 1: + raise ValueError(f"score_timestep_shift must be at least 1, got {self.score_timestep_shift}") + if self.score_discrete_steps < 0: + raise ValueError(f"score_discrete_steps must be non-negative, got {self.score_discrete_steps}") + if self.normalization_epsilon <= 0: + raise ValueError(f"normalization_epsilon must be positive, got {self.normalization_epsilon}") + if self.dmd_loss_weight < 0: + raise ValueError(f"dmd_loss_weight must be non-negative, got {self.dmd_loss_weight}") + valid_regression_types = {"decoded_lpips", "latent_mse"} + if self.regression_type not in valid_regression_types: + raise ValueError( + f"Invalid regression_type: {self.regression_type}. Must be one of {sorted(valid_regression_types)}" + ) + if self.regression_loss_weight < 0: + raise ValueError(f"regression_loss_weight must be non-negative, got {self.regression_loss_weight}") + if self.recipe == "dmd2" and self.dmd_loss_weight == 0: + raise ValueError("DMD2 requires dmd_loss_weight > 0") + if self.recipe == "dmd" and self.dmd_loss_weight == 0 and self.regression_loss_weight == 0: + raise ValueError("DMD requires at least one positive objective weight") + if self.rng_seed < 0: + raise ValueError(f"rng_seed must be non-negative, got {self.rng_seed}") @dataclass diff --git a/verl_omni/workers/config/diffusion/rollout.py b/verl_omni/workers/config/diffusion/rollout.py index 7130d3bea..5cdf6d38d 100644 --- a/verl_omni/workers/config/diffusion/rollout.py +++ b/verl_omni/workers/config/diffusion/rollout.py @@ -43,6 +43,8 @@ class DiffusionRolloutAlgoConfig(BaseConfig): sde_window_size: Optional[int] = None sde_window_range: Optional[list[int]] = None sde_contiguous: bool = True + # Optional fixed linear time shift for few-step distilled-model inference. + rollout_timestep_shift: Optional[float] = None # MixGRPO-only configs sample_strategy: str = "random" @@ -50,6 +52,8 @@ class DiffusionRolloutAlgoConfig(BaseConfig): sde_window_seed: int = 0 def __post_init__(self): + if self.rollout_timestep_shift is not None and self.rollout_timestep_shift < 1: + raise ValueError(f"rollout_timestep_shift must be at least 1 when set, got {self.rollout_timestep_shift}.") if self.sample_strategy not in ("random", "progressive"): raise ValueError(f"Unknown sample_strategy: {self.sample_strategy!r}") if self.sample_strategy == "progressive" and self.iters_per_group <= 0: diff --git a/verl_omni/workers/diffusion_distillation_worker.py b/verl_omni/workers/diffusion_distillation_worker.py index 85fe47a9f..26669e528 100644 --- a/verl_omni/workers/diffusion_distillation_worker.py +++ b/verl_omni/workers/diffusion_distillation_worker.py @@ -210,6 +210,7 @@ def validate_computation( raise ValueError(f"Role loss for {role!r} must retain an autograd graph.") return role + @torch.profiler.record_function("distillation/backward") def backward_micro_batch( self, request: PhaseRequest, @@ -262,7 +263,8 @@ def step_phase(self, request: PhaseRequest) -> tuple[dict[str, int], dict[str, f engine = self.engine_for_role(role) self.normalize_role_gradients(role) optimizer_start = time.perf_counter() - stepped, grad_norm = engine.optimizer_step(role) + with torch.profiler.record_function(f"distillation/{role}_optimizer"): + stepped, grad_norm = engine.optimizer_step(role) metrics = { f"{role}/grad_norm": grad_norm, f"perf/{role}_optimizer_s": time.perf_counter() - optimizer_start, @@ -295,6 +297,7 @@ def initialize_ema(self) -> None: """Initialize the semantic EMA role exactly from the student role.""" self.update_ema_parameters(decay=0.0) + @torch.profiler.record_function("distillation/ema") def update_ema(self) -> None: """Update the semantic student EMA in shared or independent storage.""" self.update_ema_parameters(decay=self.ema_decay) @@ -309,7 +312,7 @@ def update_ema_parameters(self, decay: float) -> None: ema_engine.update_module_ema_from(student_engine, decay) def reduce_metrics(self, metrics: dict[str, float]) -> dict[str, float]: - """Average scalar metrics across data-parallel replicas.""" + """Reduce DP means plus peak memory and explicit slowest-rank host timings.""" if not metrics: return metrics first_engine = next(iter(self.engines.values())) @@ -318,8 +321,16 @@ def reduce_metrics(self, metrics: dict[str, float]) -> dict[str, float]: return metrics names = sorted(metrics) values = torch.tensor([metrics[name] for name in names], dtype=torch.float32, device=get_device_id()) + maxima = values.clone() torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.AVG, group=group) - return {name: value for name, value in zip(names, values.cpu().tolist(), strict=True)} + torch.distributed.all_reduce(maxima, op=torch.distributed.ReduceOp.MAX, group=group) + reduced = dict(zip(names, values.cpu().tolist(), strict=True)) + for name, maximum in zip(names, maxima.cpu().tolist(), strict=True): + if name.startswith("memory/max_"): + reduced[name] = maximum + elif name.startswith("perf/") and name.endswith("_s") and not name.endswith("_per_s"): + reduced[name.replace("perf/", "perf_max_rank/", 1)] = maximum + return reduced def group_metrics(self) -> dict[str, float]: """Return stable placement diagnostics for logging.""" @@ -503,7 +514,8 @@ def execute_phase(self, data: TensorDict) -> TensorDict: normalizer = computation.loss_normalizer loss_weight = weight if normalizer is None else normalizer for name, value in computation.metrics.items(): - metric_weight = weight + is_duration = name.startswith("perf/") and name.endswith("_s") and not name.endswith("_per_s") + metric_weight = 1.0 if is_duration or name.endswith(("/active_elements", "/nonfinite")) else weight if normalizer is not None and name.endswith("/loss"): metric_weight = normalizer loss_denominators[name] = loss_denominators.get(name, 0.0) + normalizer @@ -524,6 +536,12 @@ def execute_phase(self, data: TensorDict) -> TensorDict: metrics["memory/max_allocated_gb"] = device_module.max_memory_allocated() / (1024**3) metrics["memory/max_reserved_gb"] = device_module.max_memory_reserved() / (1024**3) metrics.update(self.runtime.group_metrics()) + engine = self.runtime.engine_for_role(request.trainable_roles[0]) + metrics[f"training/{request.kind}_samples"] = float(total_samples * engine.get_data_parallel_size()) + metrics[f"batch/{request.kind}_micro_batches"] = float( + (total_samples + micro_batch_size - 1) // micro_batch_size + ) + metrics[f"batch/{request.kind}_micro_batch_size"] = float(min(total_samples, micro_batch_size)) metrics = self.runtime.reduce_metrics(metrics) for name, denominator in self.runtime.reduce_metrics(loss_denominators).items(): metrics[name] /= denominator