[diffusion, trainer, worker, data, cfg, tests] feat: add Qwen Image DMD and DMD2

[NancyFyong]

Summary

Implements the Qwen-Image architecture stage of the distribution-matching RFC as a stacked change on PR fork #8:

• registers Qwen-Image training adapters for dmd and dmd2
• adds differentiable one-step, deterministic Euler, consistency re-noise, and backward-simulated student rollout paths
• keeps only the sampled exit-step student forward in the autograd graph and runs every denoiser role in evaluation mode
• adds conditional fake-score training and positive/negative frozen-teacher scoring with standard CFG normalization
• preserves Qwen's packed latent geometry, noise - x0 velocity convention, and fp32 scheduler/objective arithmetic
• adds separate checkpointed RNG streams for initial noise, rollout decisions, transition noise, score sigma, and score noise
• supports local frozen Qwen prompt encoding and precomputed conditioning
• implements original-DMD paired regression with paper-oriented decoded LPIPS and an explicit non-paper latent-MSE mode
• adds a regression-pair dataset adapter with provenance checks
• registers a deterministic vLLM-Omni inference adapter using the exact training sigma schedule
• validates Qwen DMD/DMD2 under FSDP1/FSDP2, including semantic EMA LoRA export and eight-rank physical-batching tests
• adds a runnable DMD2 LoRA recipe, algorithm documentation, and a GPU smoke entry

The DMD2 adversarial classifier profile remains intentionally out of scope and fails closed; it is the next dedicated stage in RFC #519.

Review update — 2026-09-07

Current validated head: a0610326de063e5c180e74d6235996023b52de8c, stacked on runtime 1a0edbb. Parent fixes were merged without rewriting history; the final code tree matches the tested repair snapshot.

• Remove inactive nullable teacher_target_latents/teacher_target_pixels columns before collation, so valid None/NaN cells do not become conflicting worker inputs.
• Add dataset → collate → worker-batch regressions for both target representations.
• Inherit the shared mask expansion and optional runtime loss-denominator fixes; existing unmasked Qwen sample-mean behavior is unchanged.
• Register the Qwen example with its docs symlink, toctree entry and required timestamp.
• Synchronize public spec #543: same-resolution physical batching is supported, and collective-dependent rollout exit selection is broadcast globally rather than chosen independently per DP sample.

Stack

• Parent RFC: verl-project/verl-omni#519
• PR 1 implementation issue: verl-project/verl-omni#520
• PR 2 implementation issue: verl-project/verl-omni#535
• This implementation issue: verl-project/verl-omni#543
• Base PR: fork #8

Why this is not duplicate work

Before opening this PR, I checked:

gh issue view 543 --repo verl-project/verl-omni --comments
gh pr list --repo verl-project/verl-omni --state open --search '543 in:body'
gh pr list --repo verl-project/verl-omni --state open --search 'Qwen Image DMD DMD2'
gh pr list --repo NancyFyong/verl-omni --state open --search 'Qwen Image DMD DMD2'

The pre-creation search found no competing Qwen-Image DMD/DMD2 architecture adapter PR. PR fork #8 is the required architecture-neutral runtime underneath this stacked change; existing OPD, DPO, and DiffusionNFT implementations use different optimization and data-flow contracts.

The follow-up search for 543 in:body and Qwen DMD found no competing upstream implementation; upstream #546 is the required parent runtime, not another Qwen adapter. This update extends this existing PR rather than creating another one.

Algorithm and boundary details

• Four-step reference schedule with linear shift 3.0: [1.0, 0.9, 0.75, 0.5, 0.0].
• Discrete score timesteps use the scheduler's 1,000-step domain, one shift application, and [0.02, 0.98] clamping.
• Teacher score uses negative + scale * (positive - negative) and configurable none, layer_norm, or scalar normalization.
• Fake-score targets detach generated latents before epsilon - x_g construction.
• The DMD normalizer spans all non-batch dimensions and never uses a gradient mask.
• vLLM-Omni inference rejects stochastic SDE noise and shares the training schedule builder.
• Only student or student_ema is exportable; teacher and fake-score state remains training-only.

Batching and profiling follow-up

• Reuses the existing worker micro-batch/accumulation loop for same-resolution physical batches; checks every original-DMD sample's provenance and expands cached negative conditioning for the current batch.
• Explicitly enables native vLLM-Omni request batching, checks compatibility before execution, preserves per-request seeds and multiple outputs, and rejects unvalidated step-execution mode.
• Preserves all repeated fake phases, clears stale checkpoint timings, distinguishes summed durations/counts from averaged losses/ratios, and emits ordinary numeric scalars for console/TensorBoard logging.
• Reuses Metric, DistProfiler and Tracking; adds component trace ranges, slowest-DP-rank timing, cycle latency, sample counts, and throughput.
• Documents the existing gradient-checkpointing, resharding and FA3 Hub switches. No performance defaults changed; checkpoint-layout work is excluded.

Validation

Current-head CPU and sanity verification

CPU_PYTEST_INI=$(mktemp)
printf '[pytest]\npython_files = *_on_cpu.py\nasyncio_mode = auto\n' > "$CPU_PYTEST_INI"
PYTHONPATH=. TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1 \
  python -m pytest -c "$CPU_PYTEST_INI" -p no:cacheprovider --no-cov -q \
  tests/trainer/diffusion \
  tests/workers/test_diffusion_distillation_runtime_on_cpu.py \
  tests/workers/test_diffusion_distillation_lora_on_cpu.py \
  tests/utils/dataset/test_qwen_image_distillation_dataset_on_cpu.py \
  tests/pipelines/test_qwen_image_distillation_adapter_on_cpu.py \
  tests/pipelines/test_qwen_image_dmd_request_batch_on_cpu.py
# 571 passed

pre-commit run --files $(git diff --name-only distillation-pr2-runtime...HEAD)
# applicable hooks passed, including mypy

python tests/special_sanity/check_example_docs_symlinks.py
bash scripts/generate_trainer_config.sh
git diff --check distillation-pr2-runtime...HEAD
# docs registration passed; generated configs unchanged; diff clean

The dependent PR 5 stack passed 1,521 full L1 CPU tests. Current CPU verification used Python 3.11 / PyTorch 2.13.0+cu129 and repository-matching verl/vLLM-Omni pins. No new GPU or real multi-rank denominator validation was run for this repair. Historical GPU evidence below is retained separately, not relabeled as current-head validation.

This stacked PR targets distillation-pr2-runtime, which is outside the CPU workflow's main/v0.* base filter. The CPU/GPU evidence here is local validation, not remote CPU CI coverage.

GPU evidence carried forward from the profiling worktree

These historical GPU runs preceded the previous branch-only fast-forward and were not repeated for it or for the current review fixes:

# QWEN_IMAGE_MODEL_PATH points to a prepared tiny Qwen-Image checkpoint.
PYTHONPATH=. python -m torch.distributed.run --standalone --nproc_per_node=8 \
  --module pytest -q tests/workers/test_distillation_fsdp_roles.py \
  -k qwen_image_distillation

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python -m pytest -q \
  tests/pipelines/test_qwen_image_dmd_request_batch.py

• 8 cases passed on every rank: DMD/DMD2 × physical batch 1/2 × FSDP1/FSDP2, with repeated student/fake updates and role/EMA checks. Original-DMD cases use decoded-LPIPS regression.
• 4 passed: native scheduled serial-versus-packed requests with one/two outputs per prompt, plus masked/unmasked FA3 attention forward/backward parity. The test uses native enrich_config() and step_streaming() APIs.
• An eight-GPU production tiny-Qwen Ray/FSDP2 FA3 smoke completed three cycles at global batch 16 and physical batch 2, with two fake updates per cycle.
• The real Qwen-Image DMD2 distribution-only PickScore-SFW run completed 1,000 steps. Separately, baseline, accumulation, physical batching, no-checkpointing, no-resharding and FA3 configurations each completed six real-model cycles without reported DMD nonfinite elements.
• TensorBoard was reloaded to verify both fake phases, cycle timing and sample counts. Four native CUDA traces cover two ranks and two post-warmup cycles.

Exploratory performance, not a quality or isolated-speedup claim

Eight H20 GPUs, 1024×1024, four denoising steps and two fake updates. Means exclude the first two of six cycles; other GPU workloads remained running.

Configuration	Global batch	Mean cycle (s)	Peak allocated GiB
Native baseline	8	19.49	23.96
FA3 Hub	8	16.89	23.96
No gradient checkpointing	8	16.78	47.12
No reshard-after-forward	8	18.56	61.57
Micro-batch 1 accumulation	16	37.42	23.96
Physical micro-batch 2	16	35.80	26.04

Physical batching and accumulation sampled different mean rollout exit depths. Sampling seeds were fixed, but fresh adapter initialization was not guaranteed identical. Combinations, statistical quality gains and isolated latency were not established. Original-DMD real-model throughput is not claimed. The installed vLLM-Omni Git revision matches .github/vllm_omni_pin.txt; the alternative rc1 wheel was not independently validated.

API and Usage Example

At eight data-parallel ranks, compare the existing accumulation path against physical batching at fixed effective batch 16:

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

Use micro-batch size 1 for the accumulation comparison. Inference request batching is separate from offline differentiable student rollout; start with a small request limit and measure memory.

Scope notes

• No adversarial DMD2 classifier or real-data GAN phase.
• No causal Wan, CausVid, Self-Forcing, or MiniMax-H3 adapter.
• No standalone score transport.
• Validation-replica/CheckpointEngine orchestration remains separate: semantic export APIs alone do not complete inference integration. Tiny request-batch parity is not a claim of real-model batched LoRA inference throughput or complete training/backend latent parity.
• No gradient through vLLM-Omni, the text encoder, or VAE parameters.
• Existing policy-gradient, direct-preference, and OPD paths remain unchanged.

AI assistance (OpenAI via pi) was used for implementation, review, and test construction. The human submitter confirmed reviewing every changed line, personally running relevant tests, and accepting responsibility for explaining and defending the implementation.

Review follow-up: naming propagation

Rebased onto the PR 1/PR 2 naming cleanup: stale distillation.equations / distillation.control_plane imports and trainer.control_plane accessors were converted to distillation.utils / distillation.controller and trainer.controller; the fake-executor test field _fail_on became public fail_on. Newly introduced state uses descriptive snake_case names; framework overrides and protocol dunders are unchanged.