feat(model): add SenseNova U1.5 pixel-flow support - #389
Merged
leviking98z-rgb merged 13 commits intoAug 28, 2026
Conversation
Vendor the NEO-Unify model definitions and add trainside FlowGRPO rollout/replay so SenseNova U1.5 can be trained with deterministic pixel-space trajectories.
Load the complete interleaved MoT layer in fp32 so frozen understanding and trainable generation parameters satisfy FSDP2's per-group dtype contract while compute remains bf16.
Keep deterministic Euler updates in the BF16 trajectory dtype so 512px CFG rollouts remain tensor-identical to the official 50-step implementation while stochastic log-probability math stays fp32.
Normalize stochastic transitions into unit-noise coordinates, match official CFG reduction semantics, move prefix caches with conditions, and remove redundant per-branch and batch work while expanding contract tests.
6 tasks
Keep the draft focused on the model integration and validated training recipe while retaining the external GPU experiment record.
CjhHa1
marked this pull request as ready for review
August 26, 2026 12:15
CjhHa1
requested review from
celve,
haonan3 and
leviking98z-rgb
as code owners
August 26, 2026 12:15
Keep the default recipe on SDPA and reject forced Flash-Attention 2 usage in UniRL's Flash-Attention 4 stack.
Make diffusion stages return standard deviations in the same coordinates as transition means so KL consumers avoid implicit shape and scaling assumptions.
Bring the review branch onto the current main baseline so the SenseNova integration is checked against the latest shared model and algorithm code. Resolve the examples/README.md conflict by retaining both the upstream entries and the new SenseNova recipe.
Keep replay producers compatible by allowing means without an accompanying standard deviation, while validating the shape whenever a producer supplies one. Variance-normalized algorithms remain responsible for requiring the value at the point of use. Remove the temporary in-tree harness because repository policy keeps one-off verification scripts out of the source tree.
Only SenseNova needs to report a transition standard deviation in replay-mean coordinates because its pixel-space transitions include a resolution-dependent noise scale. Requiring every diffusion stage to calculate the same value would duplicate the existing schedule-based path and broaden this model PR. Prefer a producer override when present, but retain the original schedule/strategy calculation for all existing models and restore their replay implementations unchanged.
The SenseNova recipe uses FlowGRPO only. Passing its transition override into FlowDPPO and DiffusionOPD changed unrelated algorithms and implied support that this PR does not provide, so restore both algorithms to their existing paths.
The shipped SenseNova recipe keeps beta at its default value of zero, and LoRA reference-policy training is explicitly outside this PR's scope. Therefore the FlowGRPO reference-KL branch never consumes the added transition standard deviation. Restore the shared algorithm and ReplayResult interfaces, and stop computing an unused standard deviation during SenseNova replay. A future SenseNova LoRA plus beta>0 change should introduce that contract together with its recipe and tests.
leviking98z-rgb
force-pushed
the
feat/sensenova-u1-5
branch
from
August 28, 2026 17:26
3a52dd1 to
251c5e7
Compare
Collaborator
|
LGTM |
leviking98z-rgb
approved these changes
Aug 28, 2026
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
Add trainside FlowGRPO support for SenseNova-U1.5 pixel flow so the released 8B MoT checkpoint can roll out, replay, and train inside UniRL without a VAE.
a62fd54) and wrap them in a checkpoint-compatible bundle / FSDP-callable facade.t2i_generate: resolution-dependent noise scale, CFG (none/global/channel/cfg_zero_star), bf16 deterministic Euler, and unit-noise FlowSDE transitions for stochastic log-probs.diffusion/sensenova_u1_5/sensenova_u1_5_trainside), freeze the understanding path, and trainfm_modules.*plus every*_mot_genbranch under FSDP2.vLLM-Omni rollout is the stacked follow-up in #393.
Related Issue
N/A
Test Plan
SKIP=no-commit-to-branch pre-commit run --files $(git diff --name-only --diff-filter=ACMR origin/main --) --show-diff-on-failure— passed after merging currentmain.python3.12 -m compileall -q unirlCompose the trainside recipe:
Load
SenseNova-U1.5-8B-MoT-Previewon H20 and run packed pixel-flow inferenceFull-checkpoint SDE replay/backward with finite gradients
32-GPU / 4-node FSDP optimizer step and checkpoint-resume smoke
Official 512×512 CFG=4 50-step output matches tensor-for-tensor
Verification harnesses are intentionally not committed per repository policy. The larger packing, schedule-conversion, CFG, and cache-movement harness removed in
a7352ccremains in the stacked vLLM-Omni follow-up #393.Compatibility / Risk
ReplayResultinterfaces remain unchanged.neo_unify/is excluded from ruff and pre-commit, same pattern as Bagel.attention_backend: sdpa. The vendorflashpath is Flash-Attention 2; UniRL's engine extras ship Flash-Attention 4.attention_backend="flash"is rejected.autostill follows the vendor dispatcher (FA2 ifflash_attnimports).root_wrap: truesofm_modulesis covered.forward_batch_size: 1) because prefix caches and packed pixel trajectories are not batched.Reviewer Notes
unirl/models/sensenova_u1/diffusion.py: UniRL σ-time ↔ vendor data-time (t = 1-σ), unit-noise SDE vs bf16 Euler, and CFG reduction.unirl/models/sensenova_u1/vendor/VENDOR_COMMIT.txt. UniRL-owned seams outside the pristine subtree are the bundle/pipeline/stage, plustransformers_compat.pyand the attention-backend guard.bundle.pytrainability/FSDP facade →diffusion.py→conditions.pycache movement →pixels.pypack/unpack.main.Checklist