refactor(diffusion): extract Flow-GRPO prepare/formula hooks and add DiffusionNFT - #63
Merged
Merged
Conversation
niehen6174
force-pushed
the
refactor/flow-grpo-loss-hub
branch
2 times, most recently
from
July 27, 2026 12:30
a62e1b5 to
a7e49fc
Compare
niehen6174
changed the title
niehen6174
added a commit
to niehen6174/miles_diffusion
that referenced
this pull request
Jul 28, 2026
No code changes — triggers stage-c-5 against the PR radixark#63 merge base (0533c99) to verify the committed LTX standard still passes without the loss_hub refactor.
…formula Keep DiT forward in the FSDP actor and swap convert/prepare/formula hooks for Flow-GRPO vs NFT, instead of dispatching a whole-package loss_fn.
niehen6174
force-pushed
the
refactor/flow-grpo-loss-hub
branch
from
July 28, 2026 07:19
36c1868 to
012dc10
Compare
Separate DiT forward input preparation from loss formulas so Flow-GRPO and NFT prepare paths live in one module while losses.py stays formula-only.
Rockdu
reviewed
Jul 29, 2026
…mode Keep custom_convert as a full override, route NFT through a dedicated converter after post_process_rewards, and resolve reference selection in arguments.py so the actor only consumes args.ref_mode.
Rename lora_ema → ema / --ema-shadow and drop the --use-lora gate so full finetune can share the same pi_old path as LoRA.
Collaborator
|
TODO: add an e2e 2GPU ci for SD3.5 NFT |
Keep rollout orchestration generic while allowing NFT to provide its own sample expansion without a converter class. Co-authored-by: Cursor <cursoragent@cursor.com>
Make Flow-GRPO the explicit actor fallback and keep NFT-specific preparation and loss behavior together without resolver helpers. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Collaborator
Follow-up changes
|
Collaborator
|
Thanks, the pr is now merged! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
prepare → DiT forward → [optional ref] → loss formula → backward.prepare_flow_grpo_batch/flow_grpo_loss_formula) instead of a whole-packageloss_fn.pi_old.--loss-type nftonly auto-fills convert / prepare / formula paths; DiT forward always stays in the actor.Motivation
We need a second diffusion RL objective (DiffusionNFT) without forking the actor train loop or swapping an entire forward+loss bundle. Customization should replace parts—primarily the loss formula, plus convert/prepare when needed—while schedule, DiT forward, and backward remain shared.
Design
--loss-type nft)(x0, t)xt, CFG offlora_baseifkl_beta>0Test
PickScore reward during NFT training
Qualitative comparison on eval prompts — baseline vs. NFT checkpoint
