Skip to content

feat(dyref): add multi-reference SFT and RL training - #478

Draft
Ideny42 wants to merge 14 commits into
Tencent-Hunyuan:mainfrom
Ideny42:feat/dyref-multi-reference
Draft

Ideny42 wants to merge 14 commits into
Tencent-Hunyuan:mainfrom
Ideny42:feat/dyref-multi-reference

Conversation

@Ideny42

@Ideny42 Ideny42 commented Sep 17, 2026

Copy link
Copy Markdown
Collaborator

Summary

Add the DyRef two-stage training path on top of #417's ordered ImageSet / ImageSets foundation.

  • Convert official DyRef metadata into one raw-media manifest shared by SFT and RL, preserving ordered references while keeping the clean target out of model conditioning.
  • Train Qwen-Image-Edit-2511 and fixed-N FLUX.2-Klein with multi-reference flow-matching SFT, including Qwen N=1 compatibility and SGLang rollout/replay condition parity.
  • Add SigLIP2 semantic similarity, CSD style similarity, discriminative reward scaling, and difficulty-aware advantage reweighting on top of FlowGRPO.
  • Warm-start Stage II from an exported Stage-I PEFT adapter and synchronize that adapter before rollout 0.

This is a stacked draft. It depends on #417 and should be rebased onto main after #417 lands.

Related Issue

Refs #470

Depends on #417.

Test Plan

  • Full repository hooks:
    SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure — passed.
  • Python syntax:
    python3 -m compileall -q unirl datasets/dyref — passed.
  • Hydra composition: all four new Qwen/FLUX SFT and RL configs composed and resolved with local placeholder paths — passed.
  • CPU contract harnesses with synthetic N=1/N=2/N=3 data:
    • DyRef conversion, ordered reference typing, target isolation, and fixed-N filtering — passed.
    • Multi-reference diffusion SFT track construction and clean-target VAE shapes — passed.
    • Qwen text placeholders, source-latent packing, img_shapes, SGLang capture splitting, replay reconstruction, and legacy N=1 wire compatibility — passed.
    • SigLIP2 DRS F2, CSD style-reference selection, DAR weighting, and first-rollout weight sync — passed.
    • PEFT Stage-I-to-Stage-II adapter loading plus rank/alpha mismatch rejection — passed.
    • FLUX legacy stretch parity, DyRef aspect-preserving crop, and one-of-{1,2,3} SDE scheduling — passed.
  • Not run: real-checkpoint GPU SFT/RL. This draft is being opened before the planned 8×H20 validation (SFT on 8 GPUs; RL on 7 train/rollout GPUs + 1 reward GPU).

Compatibility / Risk

  • Existing manifest fields remain valid. DyRef rows add one role="target" image that is projected only into reward metadata; it never becomes a model input.
  • LoraConfig.initial_adapter_path is optional. It accepts standard PEFT adapter directories and intentionally rejects rank/alpha mismatch and meta-initialized loading.
  • Legacy flat Qwen N=1 latent conditions and flat one-image-per-prompt SGLang requests are normalized without changing N=1 model numerics.
  • FLUX keeps stretch as the default condition resize mode; only the DyRef recipes opt into aspect-preserving center crop.
  • CSD requires its external ViT-L checkpoint. The adapted model wrapper retains the upstream MIT notice; checkpoint licensing remains governed by its model card.
  • The public DyRef code's dar_min_weight=0.5 collapses typical published similarity values to uniform weights. These recipes use 0.01 so DAR remains active; GPU validation must report the resulting weight distribution.
  • The Qwen SGLang path may redundantly encode shared references across expanded outputs; GPU profiling will determine whether a follow-up deduplication patch is needed.

Reviewer Notes

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

celve and others added 14 commits September 14, 2026 19:11
Reference images travel as a slot-major input Part chain. The driver loads
N per prompt, transposes prompt-major to slot-major, and chains one input
Part per slot, so every Part stays row-aligned to the prompt batch --
Part.__post_init__ requires len(primitive) == len(sample_ids), and
primitives are keyed by canonical modality, so N references cannot share
one Part.

FLUX consumes them as a longer token sequence. Each reference is already
resized to the generation canvas, so the condition tensors stay dense and
Flux2KleinConditions / predict_noise need no change -- they concatenate
tokens and ids blindly. The only model-side change is the RoPE time
offset, now scale * (slot + 1), replicating upstream's
t_coords = [scale + scale * k for k in range(N)] so the trainside
(recomputed) and SGLang (captured) ids stay identical. They agreed at N=1
only because both produced t=10.

On the SGLang path one reference keeps the flat per-prompt condition_image
shape; two or more nest one list per prompt, which is what upstream's
preprocess_condition_image and ImageVAEEncodingStage already iterate. The
per-prompt stash now keys on nesting rather than length, because a single
prompt with N references is a length-1 outer list.

Batches must be uniform in reference count: the lineage walk resolves every
ancestor id by name, so a Part holding only some samples is invalid. Ragged
batches are rejected with a message pointing at bucketing.

Models that do not support multiple references (Qwen-Image-Edit-Plus,
HunyuanImage 3.0, BAGEL) keep their own fail-fast guards, which now surface
on N>1 instead of the loader truncating first.
Keep ordered reference slots for reward requests so EditReward scores against the primary source, and reject extra references in single-image model paths instead of silently ignoring them.
Close condition-image handles deterministically and describe both supported per-prompt wire shapes accurately.
Carry every image through the VLM encoding and emit the SGLang image_data shape required for both single-candidate and parallel-sampling requests.
@github-actions github-actions Bot added the wip Draft / work in progress label Sep 17, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

wip Draft / work in progress

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants