Repository navigation
[reward,trainer] fix: gather prompt strings for groupwise rewards - #249
Merged
Merged
Conversation
RewardProcessor rebuilt cross-rank prompts for groupwise and global rewards by decoding the gathered prompt_ids with the adapter tokenizer. prompt_ids come from the same tokenization as prompt_embeds, so for chat-template models the decoded text carried the chat template, system message, and padding instead of the dataset prompt, and adapters without a tokenizer failed with "Cannot decode prompts: tokenizer not provided". gather_samples already transports Optional[str] fields as packed UTF-8, so the processor now gathers the prompt strings directly and no longer takes a tokenizer.
Jayce-Ping
marked this pull request as ready for review
October 17, 2026 14:36
Contributor
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
This was referenced Oct 17, 2026
Jayce-Ping
added a commit
that referenced
this pull request
Oct 18, 2026
* [model,test,docs] feat: add FLUX.3 adapter * [model] fix: default FLUX.3 encode_prompt to the official 8192-token limit Flow-Factory exposes no max_sequence_length setting, so dataset preprocessing encoded prompts with the encode_prompt default of 512 tokens, while the official FLUX.3 text encoder step defaults to 8192. A 999-token caption produced a 512-token window instead of the official 1040-token one: the prompt was truncated, and the window was no longer a multiple of the 80-token pad recipe. * [dataset,test] fix: register flux3 in the offline smoke catalog and README table check Two existing tests failed once flux3 was registered: - test_offline_smoke_profiles requires every registered model type to appear in an offline smoke profile. Add flux3 to text_to_image, image_to_image, and multi_image_to_image, matching flux2-klein and the adapter's within-type multi-reference contract. - test_sensenova_docs locates the combined T2I/I2I README table by its rowspan, which grew from 9 to 11 with the two FLUX.3 rows. * [model,docs,agents] fix: align Diffusers submodule references with the 7ecb1e1b pin The FLUX.3 change moves the shared Diffusers submodule from 80c7ed262 to 7ecb1e1b. That pin raises the submodule's floor to Python 3.11 and huggingface-hub 1.32, and Qwen-Image 2.1 still named the old commit in its install error. - Point QWEN_IMAGE_21_DIFFUSERS_COMMIT at 7ecb1e1b. The functions it documents as copied (calculate_dimensions, retrieve_latents, _get_qwen_prompt_embeds, _pack_latents, _unpack_latents) are unchanged between the two commits, and the released checkpoint sets no sample_sigmas, so Qwen-Image 2.1 sampling is unaffected. - State the Python 3.11 requirement in the README install note and in dependencies.md. * [model,test] fix: default every FLUX.3 entry point to the official 8192 tokens encode_prompt already defaults to the official 8192-token limit, but inference, _inference, and the Qwen3-VL helpers kept their own 512 default and pass it to encode_prompt explicitly, so raw-prompt rollouts (for example with enable_preprocess: false) truncated long captions to 512 tokens. Share one DEFAULT_MAX_SEQUENCE_LENGTH across every entry point, run the window tests on the default, and cover the inference path. * [model] docs: describe the FLUX.3 tokenizer without reward decoding #249 removes the reward-side prompt_ids decoding, so the tokenizer property no longer serves it. Match the wording #249 uses for the Qwen-Image 2.1 tokenizer. * [model,test,agents] fix: run every FLUX.3 sample at transformer batch size 1 for exact replay With per_device_batch_size > 1 the on-policy log-probs did not match the rollout: an exact probe found up to 7e-6 differences in 421 of 768 replays with FLUX.3-upsampled GenEval prompts and 140 of 768 with PickScore captions. inference() and forward() each ran a batch batched when all its prompts shared one 80-token window count and per sample otherwise. Group-contiguous rollout packs repeat one prompt and ran at batch size 8, while the shuffled optimize micro-batches mix window counts and replayed the same samples at batch size 1, and bf16 GEMMs of different shapes round differently. Both phases now run each sample's transformer forward at batch size 1, so a sample's result no longer depends on its batch, and _inference handles exactly one sample. A regression test regroups rollout packs into new micro-batches and requires bit-identical log-probs from batch size 1 transformer calls. The train-inference consistency notes record the failure mode. * [trainer,model,test,docs,agents] feat: shuffle whole rollout packs for pack-dependent adapters train.shuffle_samples permuted individual samples, so an adapter whose batched forward depends on batch composition replayed samples in micro-batches that differed from their rollout packs, and the on-policy ratio was no longer exactly 1. Bagel could only warn and ask for shuffle_samples: false, and FLUX.3 ran every sample alone, which cost about 13% rollout and 15% optimization time at per_device_batch_size 8. For adapters that declare requires_preserved_replay_batch_composition, the generation loop now records the rollout packs and the optimize loop permutes whole packs, keeping each pack's samples and order. It fails fast when a pack does not match per_device_batch_size or the samples changed after rollout; reward overlap shares that validation. Batch-invariant adapters still shuffle individual samples with the same permutation, and shuffle_samples: false still keeps the rollout order. FLUX.3 declares the flag and goes back to batched execution for packs whose prompts share one 80-token window count. Bagel's warning is dropped, and the help text, guidance, and train-inference consistency notes describe the semantics.
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
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
RewardProcessorrebuilt cross-rank prompts for groupwise and global rewards by decoding thegathered
prompt_idswith the adapter tokenizer.prompt_idscome from the same tokenization asprompt_embeds, so for chat-template models (FLUX.2, Qwen-Image, and others) the decoded textcontained the chat template, system message, and padding instead of the dataset prompt. Adapters
without a tokenizer failed with
ValueError: Cannot decode prompts: tokenizer not provided.gather_samplesalready transportsOptional[str]fields as packed UTF-8 (byte lengths ride theint64 metadata gather, bytes go through one uint8 all-gather), so the processor now gathers the
prompt strings directly.
prompt_idsstay aligned withprompt_embedsand are no longer read byreward code.
Changes
RewardProcessor: drop thetokenizerargument, theprompt_idsdecode path, and_decode_prompts.BaseTrainer: stop passing the adapter tokenizer to the train and eval reward processors.utils/dist.pyand the Qwen-Image 2.1 adapter.exact original prompts, including non-ASCII text, while
prompt_idshold template token ids. Theoffline trainer test adapter no longer needs a tokenizer.
GPU validation
scripts/agent_scope.pyclassifies this change as R4 because it touchesRewardProcessorandBaseTrainer. The change only removes a decode path in favor of an existing string transport, so themaintainers chose a targeted check over the 45-job framework-upgrade campaign: one SD3.5 and one
FLUX.3 GRPO job on a single backend.
Setup: 8x H20, one node, DDP, Python 3.12, torch 2.10.0+cu129. Each job ran two GRPO epochs on
dataset/pickscorewith the groupwisepickscore_rankreward,sampler_type: global_batch, andgroup_size: 8, so every group spans all eight ranks and its prompts go through the cross-rankgather. A temporary, uncommitted probe recorded the prompts each
PickScoreRankRewardModelcallreceived.
ec59059b)'user\na huge mess of pipes, digital art\nassistant\n'All three jobs completed (72-78 s) with
ratio = 1and finite losses and gradient norms.CPU evidence
collective packing, reward/optimization overlap, offline trainers, and rewards (128 passed).
pytest tests): 2984 passed, 50 skipped, 0 failed. Merged with [model,test,docs] feat: add FLUX.3 adapter #247 and [prompt_enhance,model,cli,docs] feat: move prompt enhancement into ff-prompt-enhance #248,the full suite gives 3021 passed, 50 skipped, 0 failed.
tests/trainers/test_group_coordinator_distributed.py, including thetorchrun-launchedtest_subgroup_coordinator_uses_only_its_two_rank_scope, passes on Linux. On macOS that test cannotresolve the hostname under
torchrun; launched as 4 gloo ranks withMASTER_ADDR=127.0.0.1, itpasses, and the same run fails on
mainwithCannot decode prompts: tokenizer not provided.scripts/check_docs_only_diff.pyconfirms the docstring edits are docs-only, andscripts/audit_docstrings.pyreports no findings.