Skip to content

feat(model): add SenseNova U1.5 pixel-flow support - #389

Merged
leviking98z-rgb merged 13 commits into
Tencent-Hunyuan:mainfrom
CjhHa1:feat/sensenova-u1-5
Aug 28, 2026
Merged

feat(model): add SenseNova U1.5 pixel-flow support#389
leviking98z-rgb merged 13 commits into
Tencent-Hunyuan:mainfrom
CjhHa1:feat/sensenova-u1-5

Conversation

@CjhHa1

@CjhHa1 CjhHa1 commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator

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.

  • Vendor the pinned OpenSenseNova NEO-Unify definitions (a62fd54) and wrap them in a checkpoint-compatible bundle / FSDP-callable facade.
  • Implement packed RGB pixel-flow sampling that matches official 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.
  • Carry per-sample prefix KV caches as typed conditions, including device movement with the batch.
  • Add the 8B trainside Hydra recipe (diffusion/sensenova_u1_5/sensenova_u1_5_trainside), freeze the understanding path, and train fm_modules.* plus every *_mot_gen branch 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 current main.

  • python3.12 -m compileall -q unirl

  • Compose the trainside recipe:

    python -m unirl.train_diffusion \
      --config-name=diffusion/sensenova_u1_5/sensenova_u1_5_trainside \
      --cfg job --resolve
  • Load SenseNova-U1.5-8B-MoT-Preview on H20 and run packed pixel-flow inference

  • Full-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 a7352cc remains in the stacked vLLM-Omni follow-up #393.

Compatibility / Risk

  • The SenseNova package and recipe are additive; shared algorithm and ReplayResult interfaces remain unchanged.
  • Vendored neo_unify/ is excluded from ruff and pre-commit, same pattern as Bagel.
  • The documented recipe pins attention_backend: sdpa. The vendor flash path is Flash-Attention 2; UniRL's engine extras ship Flash-Attention 4. attention_backend="flash" is rejected. auto still follows the vendor dispatcher (FA2 if flash_attn imports).
  • Interleaved MoT layers mix frozen understanding and trainable generation params, so FSDP2 masters stay fp32 and compute is bf16 mixed precision. The recipe also sets root_wrap: true so fm_modules is covered.
  • Trainside rollout is per-sample (forward_batch_size: 1) because prefix caches and packed pixel trajectories are not batched.
  • Scope is T2I trainside FlowGRPO only. Image editing, LoRA, and engine rollout are out of this PR.

Reviewer Notes

  • Start with unirl/models/sensenova_u1/diffusion.py: UniRL σ-time ↔ vendor data-time (t = 1-σ), unit-noise SDE vs bf16 Euler, and CFG reduction.
  • Vendor pin and update steps are in unirl/models/sensenova_u1/vendor/VENDOR_COMMIT.txt. UniRL-owned seams outside the pristine subtree are the bundle/pipeline/stage, plus transformers_compat.py and the attention-backend guard.
  • Suggested reading order: recipe YAML → bundle.py trainability/FSDP facade → diffusion.pyconditions.py cache movement → pixels.py pack/unpack.
  • Duplicate-work check: the only overlapping open PR is feat(sensenova): add vLLM-Omni rollout #393, which stacks on this branch for vLLM-Omni rollout.
  • AI assistance: Codex reviewed the full diff and helped remove unused cross-model algorithm plumbing; the commands above were rerun after merging current main.

Checklist

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

Vendor the NEO-Unify model definitions and add trainside FlowGRPO rollout/replay so SenseNova U1.5 can be trained with deterministic pixel-space trajectories.
@github-actions github-actions Bot added the wip Draft / work in progress label Aug 24, 2026
CjhHa1 added 3 commits August 24, 2026 22:24
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.
Keep the draft focused on the model integration and validated training recipe while retaining the external GPU experiment record.
@CjhHa1
CjhHa1 marked this pull request as ready for review August 26, 2026 12:15
@github-actions github-actions Bot added need review Ready and waiting for review and removed wip Draft / work in progress labels Aug 26, 2026
CjhHa1 added 2 commits August 27, 2026 20:26
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

Copy link
Copy Markdown
Collaborator

LGTM

@leviking98z-rgb
leviking98z-rgb merged commit 10b6951 into Tencent-Hunyuan:main Aug 28, 2026
6 checks passed
@github-actions github-actions Bot removed the need review Ready and waiting for review label Aug 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants