feat(janus-pro): add i2t and t2i AR training support

[HaitaoWuTJU]

Summary

Adds Janus-Pro support for both Text+Image -> Text and Text -> Image AR training in UniRL.

This includes Janus-Pro bundle/config/pipeline wiring, multimodal I2T conditions, AR image-token prompt/rollout/replay/decode support for T2I, vendored Janus runtime code, LoRA training recipes for Geo3K-MC I2T and PickScore T2I, and the direct Janus runtime dependencies.

The branch is rebased onto the current UniRL API. Both I2T and T2I now store old-policy log-probabilities using the same full-sequence teacher-forced geometry used by training. T2I replay is a single batched forward rather than one cached forward per image token, and the recipes enable gradient checkpointing.

Related Issue

Related to #25

Test Plan

• Repository validation:
  • 25 current tests/ tests passed.
  • Full pre-commit suite passed.
  • Both Janus-Pro Hydra recipes resolve.
  • The final PR diff contains no test files.
• Source/runtime validation:
  • Vendored Janus source matches DeepSeek Janus commit 1daa72fa except for declared integration compatibility changes; the upstream MIT license is preserved.
  • The real Janus-Pro-1B processor loaded two images into a (2, 626) token batch with 1,152 image slots.
  • The real 1B model constructed under Transformers 5.6.2.
• Single-H20 GPU smoke with the real Janus-Pro-1B checkpoint:
  • I2T processor + vision + generate + replay + backward + optimizer passed.
  • T2I generated 16 image tokens, decoded a (1, 3, 64, 64) image, replayed, backpropagated, and optimized.
  • Full 576-token T2I teacher-forced replay + backward + optimizer passed with 4.23 GiB peak allocated memory.
  • T2I old/new log-probability drift was reduced from 0.6168 to 0.0 by replay-geometry anchoring.
• Two-H20 end-to-end unirl.train_ar smoke with Ray, DP2 FSDP, LoRA, local MC reward, and one complete rollout:
  • reward=0.5000
  • grad_norm=0.3854
  • ratio=1.0000±0.0000
  • clip_fraction=0.00
  • Process exited successfully after reward → advantage → backward → optimizer.

Compatibility / Risk

This is a new model package and two new recipes, so the changes are isolated from existing model paths. The old-policy anchoring adds one graph-free full-sequence forward per rollout, but removes the previous 576-forward T2I replay pattern and prevents nominally on-policy samples from being clipped solely because cached BF16 decode and full-sequence replay use different numerical kernels.

The remaining validation gap is a sustained multi-step training curve, especially with the 7B checkpoint and the full T2I trainer path. The PR remains draft until that acceptance criterion is agreed or supplied.

Reviewer Notes

AI assistance was used. No AI co-author trailer is present. Local validation did not add tests to the PR. Please review the old-policy anchoring in both AR stages, Janus-Pro pipeline routing, and the two recipe resource settings first.

Checklist

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

[haonan3]

Could you provide the training curve (reward / loss)?

[CjhHa1]

Reviewed this and pushed five commits to the branch (9b684dc4..7378daa8), plus the training curve that has been outstanding since #141 (comment) on Jul 10.

Blocking issue found: the package did not import

>>> import unirl.models.janus_pro
ModuleNotFoundError: No module named 'unirl.types.rollout_req'

#214 replaced RolloutReq / RolloutResp with the sample-native Sample / Part API. The Merge branch 'main' into this branch did not conflict, because unirl/models/janus_pro/pipeline.py only exists on the feature side — so it kept importing deleted modules. Every validation number in the PR description predates that merge and no longer holds.

7378daa8 ports it. generate is now Sample -> Sample; task selection reads parts[0].control["task"] and otherwise infers from Sample.has_image_input(); i2t pulls turns from vision_conditioning() and t2i from turns(), both failing loud on a multi-turn trajectory the single-user-turn Janus template cannot encode; sampling params come off the frontier gen shell, and t2i requires JanusProImageARSamplingParams since the image grid / CFG weight / token count all ride on it. The stages, conditions, bundle, config, chat_template and image_prompt were already API-agnostic and are untouched.

Training curve

Janus-Pro-1B, examples/ar/janus_pro_grpo_geo3k_mc_lora.yaml unmodified except num_rollouts=40 and W&B on. One 8x H20 node. 40/40 rollouts, exit=0, 44.3 min.

metric	first	last	notes
reward (avg@8)	0.0078	0.4062	max 0.4570, last-5 mean 0.3289, crosses the 0.25 4-way baseline at rollout 6
grad_norm	0.3405	0.0211	peaks 1.1775, clipped at 1.0 during rollouts 4-6
step time	268 s	49 s	responses collapse toward a single letter

The near-zero start is 1B ignoring "answer with the letter only"; GRPO teaches the format first, which is also why step time drops 5.5x and grad_norm falls two orders of magnitude once it is learned. W&B: https://wandb.ai/hfmoe-hfmoe/unirl/runs/i5jx76la

Not covered: 7B (1B was the fastest route to a green end-to-end curve; the code path is identical but 7B throughput/memory are unmeasured), t2i, and anything past the format-learning phase.

Review fixes

257ccd6b + 912e860e — anchoring moved to the framework. Both autoregress paths overwrote segment.log_probs with a replay forward. The repo already has this as a first-class knob (old_logp_source: rollout|replay with prepare_segment, used by FlowGRPO / DRPO / CPPO / DPPO); GRPO was simply the one algorithm missing it. Doing it in the stage could not be switched off, permanently zeroed the rollout/replay K3 metric from #187, and only matched training geometry while forward_batch_size == micro_batch_size == 1. GRPO now declares anchor_fields=("log_probs",) with recomputes_anchor() in replay mode, so TrainStack drives the hook per micro-slice at the exact geometry training replays at. Both recipes opt in.

Worth flagging for the description: ratio=1.0000±0.0000 and clip_fraction=0.00 are true by construction under replay anchoring — they show the anchor is wired correctly, not that anything converged. In the run above they held on all 40 rollouts; the reward and grad_norm columns are what carry signal.

912e860e — ar.py::replay no longer materializes full-sequence logits. It ran the fused LlamaForCausalLM forward and kept logits over the whole prompt+response window to read T rows out. At vocab 102400 and max_prompt_length: 2048 that is ~430 MB of bf16 logits per sequence, retained by autograd, to use ~3% of it. Now takes last_hidden_state from the decoder body and runs lm_head only at each row's predict positions — the same hidden-then-head split image_ar.py already used. Verified numerically equivalent (max abs diff 4.8e-7 on a causal toy model).

94243cc0 — both stages carried their own copy of the left-repack loop that unirl.models.types.ar.left_pad_prompt provides (qwen3 and qwen_vl use it); collapsed onto one helper. ar.autoregress called .item() per row per step, costing 2*B host syncs per decode step; now accumulates on device. cfg_weight now comes only from conditions — replay runs from GRPO with no sampling_params, so sampling under a different value would bias every ratio; a disagreement is an error rather than a silent skew. Also drops the stages' own .eval() calls, since TrainsideRolloutEngine already saves and restores training mode and model.model.eval() flipped the frozen towers with nothing restoring them.

d9b984e9 — trust_remote_code=true lets a checkpoint ship modeling code that wins over the vendored registration, silently dropping the transformers-5.x patches in VENDOR_COMMIT.txt, including the grad-safety clone in prepare_inputs_embeds that replay backprop depends on. Now asserts the resolved class.

Still open

A 7B curve, and t2i. Happy to run either — PickScore_v1 and CLIP-ViT-H-14-laion2B are already staged, so t2i only needs a prompt list for datasets/pickscore/train.txt.

[CjhHa1]

7B curve, closing the gap I left open above. Same recipe, same overrides, same 8x H20 node — only JANUS_PRO_PATH differs, so the two runs are directly comparable.

run	first	last	max	last-5 mean	crosses 0.25	grad_norm max -> last	step time first -> last	wall	anchor
Janus-Pro-1B	0.0078	0.4062	0.4570	0.3289	rollout 6	1.1775 -> 0.0211	268 s -> 49 s	44.3 min	ratio=1, |Δlogp|=0, clip=0
Janus-Pro-7B	0.0234	0.4062	0.4375	0.3313	rollout 5	1.0149 -> 0.0057	353 s -> 113 s	91.4 min	ratio=1, |Δlogp|=0, clip=0

Both 40/40, exit=0. W&B: 1B https://wandb.ai/hfmoe-hfmoe/unirl/runs/i5jx76la · 7B https://wandb.ai/hfmoe-hfmoe/unirl/runs/cvfuwm0p

7B does not beat 1B here, and that is worth being explicit about. Last-5 means are within noise of each other (0.3313 vs 0.3289) and both sit only a little above the 0.25 four-way random baseline. What 40 rollouts of LoRA GRPO buys on this task is output-format compliance — the model learns to emit a bare letter, which is also why step time drops ~3x and grad_norm falls two orders of magnitude — not geometry reasoning. So these curves demonstrate that the training path runs end to end and stays numerically well-behaved; they are not evidence that Janus-Pro learns Geometry3K. A longer run, a harder reward, or full fine-tuning instead of LoRA r16 would be needed for that.

What the 7B run does confirm that 1B could not:

• No OOM in replay. This is the case 912e860e exists for: at vocab 102400 with max_prompt_length: 2048, the original fused-logits path retained roughly 430 MB of bf16 logits per sequence in the autograd graph to read 64 rows out of it. With lm_head applied only at the predict positions, peak per-GPU memory stayed at 13.7 GB through all 40 rollouts on 96 GB cards.
• The anchor holds at 7B. ratio is exactly 1.0 and |Δlogp| exactly 0.0 on all 40 first-updates, same as 1B — recomputes_anchor() driving prepare_segment per micro-slice behaves identically at both scales.

One operational note for anyone reproducing: startup was 31 minutes before the first rollout, almost entirely eight Ray workers each cold-reading the 14.85 GB checkpoint from CephFS. references/startup_optimization.md-style local staging would cut most of that; it does not affect the training numbers above.

Still not covered: t2i. PickScore_v1 and CLIP-ViT-H-14-laion2B are staged already, so it only needs a prompt list for datasets/pickscore/train.txt — happy to run it if you want it before merge.

[CjhHa1]

Follow-up on the T2I half, which my earlier comments left uncovered. Posting for the record since the PR is closed — nothing here asks for it to be reopened.

A real bug: every T2I reward was silently zeroed

RewardService.score_and_attach zeroes the reward of any AR generation whose segment reaches max_new_tokens, on the assumption that a trace which never emitted a stop token rambled to the cap (truncated_reward defaults to "zero").

Janus-Pro T2I emits one image token per grid cell, so its length is max_new_tokens on every sample by construction — _resolve_image_grid refuses any other value. Because JanusProImageARSamplingParams subclasses ARSamplingParams, every T2I sample matched the truncation test, so every PickScore reward was replaced with 0. Advantages went flat, grad_norm was exactly 0, and the run trained on nothing while looking healthy.

Localised by instrumenting the backend: PickScore returned [0.7807, 0.7724, 0.7623, 0.8192] etc. while the trainer logged reward=0.0000.

Fix: ARSamplingParams now declares whether max_new_tokens is a cap or an exact count, the image params override it, and the shaping consults it. On 8x H20 reward went 0.0000 -> 0.7905 and grad_norm 0.0000 -> 0.1870 on the first rollout. Text behaviour is unchanged — traces at the cap are still zeroed, shorter ones and keep mode untouched.

Commit d4be7e6b, on CjhHa1/UniRL@janus-pro-i2t-t2i-reviewfixes (the branch also carries the five review commits from 9b684dc4..7378daa8); I could not push it here once the PR closed.

Two more T2I findings

rollout.forward_batch_size > 1 kills the worker. With forward_batch_size=4 all eight bundle actors die during construction — ActorUnavailableError: Socket closed, has_creation_task_exception=false, no Python traceback. num_devices=1 runs the same path fine, so it is specific to concurrent actors. forward_batch_size=1 (as authored) is stable; I ran 280+ rollouts on it without a failure.

PickScore with base_device: cuda deadlocks at 8 actors. Every rank parks in Module._apply/convert during the CLIP host-to-device copy — py-spy shows all eight at the same line, one actor per GPU, 0% util, memory static, no progress in 20+ minutes. base_device: cpu sidesteps it. This one is in the framework's reward code, not this PR; the T2I recipe is just the first config to exercise it (the i2t recipe's MCExactMatchRewardScorer loads no model at all, which is why i2t never hit it).

T2I trains, but does not learn on PickScore

With the fix in, 100 rollouts at the recipe defaults produce a flat curve: least-squares drift +0.0021 against per-rollout noise of 0.0217, i.e. one tenth of a single rollout's noise. Not slow learning — flat.

Four-arm ablation, 8x H20 each:

arm	n	first10	last10	drift	noise sd	drift/noise
A baseline (scope=global)	100	0.7616	0.7625	+0.0021	0.0217	0.10
B scope=group	60	0.7550	0.7616	+0.0117	0.0230	0.51
C scope=group + spp=8	60	0.7568	0.7641	+0.0125	0.0230	0.54
D scope=global + lr 3e-4	60	0.7571	0.7561	+0.0016	0.0225	0.07

No arm clears a drift-to-noise of 1, so none of them is a demonstrated effect. But the split is clean: the two scope=group arms sit 5-7x above the two scope=global arms, independent of learning rate. That matches a variance decomposition over 95 real sibling groups — between-prompt reward std 0.0530 vs within-prompt 0.0186, so under adv_normalization_scope: global roughly 89% of the advantage variance is prompt difficulty the policy cannot act on. The T2I recipe ships global.

Two things this rules out: raising the learning rate does nothing under global (arm D came in below baseline), and clip_range is inert — clip_fraction was exactly 0.00 across all 280 ablation rollouts, because old_logp_source: replay plus one update per batch pins the ratio at 1 and reduces the objective to plain REINFORCE.

A 500-rollout scope=group run is going now; extrapolating arm B's slope it should reach a drift of ~4x noise, which would settle whether normalisation scope is the cause. Also worth suspecting, and untested: PickScore's within-prompt spread is only 0.0186, so it may simply not discriminate at this model's quality level.

To be explicit about what is and is not shown: the T2I path runs end to end and stays numerically well-behaved (ratio exactly 1, |Δlogp| exactly 0, no clipping, 280+ rollouts without a failure). It has not been shown to improve PickScore at the recipe's defaults, and on current evidence it does not.