fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards

[CjhHa1]

Summary

Batch.concat could silently produce a wrong batch when shards disagreed. Several cases were already caught by per-trainer guards that each patched one symptom; others were not. This PR makes the container fail closed in one place, removes those guards, and fixes the producers that the stricter rule exposed.

• Shared fields must agree. Batch.concat used to keep values[0] for every SHARED field. It now requires structural equality across items and names the first differing path and values on mismatch (e.g. Part.sampling_params.sde_indices[1]: 1 != 2; tensors as path.flat[i]: x != y). A TensorRef compares by shape and dtype, because its rows live in the transport; the per-rollout SDE schedule is still caught exactly through Part.sampling_params.
• Per-sample fields must be present in every non-empty item. A CONCAT / PACKED field (and packed cu_seqlens) that is None or a missing dict key in some shards used to be dropped, leaving the merged lists shorter than batch_size. It may now be None only in zero-row items.
• A TensorRef in a CONCAT field must have one row per sample. TensorRef.concat joins spans regardless of row count, so a ref with any other row count (e.g. a 2×B CFG batch) used to merge into misaligned rows; it now raises.
• Sample rejects a part with a blank, non-string, or repeated sample id, replacing the checks in Part.fork and the SGLang seed derivation.
• Guards made redundant are removed: the Part.concat override (primitive_metadata equality), Part.validated_sample_ids and its two callers, and the async trainer's require_single_rollout_id / _rollout_id / _require_single_generation plus its separate mixed-version check. combine_rollout_prompts now concatenates first, then checks a single non-None output_version across generated parts.
• Producers emit a fixed key set. The Qwen3-Omni vLLM-Omni output adapter and the Qwen3-Omni / Qwen3.5 chat templates dropped a media key whenever their call had no such media. An image-only shard merged with a video-only shard then misaligned the media lists. They now always emit per-row lists with None for rows lacking that media; every consumer already filters None rows. AgenticTrainer clears the per-trajectory harness_status on train parts before concatenating, as it already did for primitive_metadata.
• HunyuanImage3 guided diffusion stores rows per sample. Guided t2i and t2ti stored the tokenizer's [cond; uncond] fused batch (2×B rows) unsplit, so a DP merge or select paired one sample's cond with another's uncond; it2i already split it, and that split is now split_cfg_fused, used by all three modes (_expand_cfg_for_forward restacks it bit-exactly). The raw tokenizer_output (also 2×B under CFG, read only by the in-process sampling KV cache) is no longer written to Part.conditions.

Related Issue

N/A

Test Plan

• SKIP=no-commit-to-branch pre-commit run --from-ref 0fe4f99d --to-ref HEAD → all hooks passed; re-run after the rebase onto 89ac5b03 as --from-ref 89ac5b03 on e3d7f0a5 → all hooks passed.
• One-off harness, not committed per AGENTS.md §5, on Python 3.12.13 + torch 2.14.1 CPU with ray stubbed → 32/32 checks passed. It covers:
  • Unchanged behaviour: Sample DP chunk/cat round trip, split / select / slice, packed TextSegment concat, TensorRef concat, and shared sigmas as TensorRefs from two workers.
  • New failures: output_version, sampling_params with different sde_indices, LatentSegment.sde_indices, primitive_metadata, and harness_status mismatches; None rewards, a missing condition key, and a None packed field in one shard; differing non-batched tensors; duplicate and empty sample ids.
  • Allowed cases: equal-but-distinct sampling_params, Broadcast(Sample), a mismatched TensorRef, and None in zero-row items.
  • Callers: combine_rollout_prompts with same, mixed, and missing versions; agentic train parts after clearing harness_status; Qwen3-Omni image + video shards and Qwen3.5 text-only + image shards stay row-aligned after concat and microbatch select.
• Follow-up commits up to c2899833: one-off harnesses on Python 3 + torch 2.13 CPU with ray stubbed → 27/27, plus two HunyuanImage3 harnesses (all pass). The HunyuanImage3 ones cover: split_cfg_fused + _expand_cfg_for_forward reproducing the tokenizer fused bit-exactly, cfg=False passthrough, row-count mismatch raising, merged and selected guided parts keeping each cond row with its own uncond (worker-side and driver-side TensorRef), and guided it2i parts with a 2×B tokenizer_output merging (the parent commit raised). The 27 general checks cover a TensorRef CONCAT row mismatch (raises), batched and zero-row TensorRef concat, a tensor mixed with a TensorRef (raises), equal vs differing non-batched tensors, path-named errors (sampling_params.sde_indices, output_version, dict keys, tensor shape, a middle-only tensor element, 2-D, NaN, numpy), all dataclass fields compared, shared TensorRef by shape/dtype, None only in zero-row items, the Sample DP chunk/cat round trip, and duplicate sample ids.
• Trajectory-IR CPU regression from the UniRL paper harness (experiments/trajectory_ir/run_cpu_evidence.py --allow-dirty against 07dec547) → PASS: 10 correctness checks. Not re-run on the follow-up commits; reason: harness not present in that checkout.
• e3d7f0a5 (drops the numpy branch and the non-batched CONCAT tensor equality check, which have no in-repo producer; non-batched tensors keep item 0 as on main): general harness → 26/26, both HunyuanImage3 harnesses pass.
• GPU smoke on c2899833 (8×H20, 2–4 rollouts each), all passed: SD3.5 FlowGRPO sync DP merge and accumulate_rollouts=2 (ratio 1.0000), BAGEL async with weight sync every rollout, Qwen3-Omni mixed image + video rows (single engine replica, so no cross-replica merge), HunyuanImage3 t2i guided + KV cache (ratio 0.9997) and it2i guided (ratio 0.9999), agentic default and tight. Not re-run after the rebase or on e3d7f0a5; reason: patches unchanged by the rebase, and e3d7f0a5 removes only branches no in-repo path reaches. Details in the PR comment.

Compatibility / Risk

• Behaviour change: inconsistent shards now raise instead of silently producing a wrong batch. The in-repo call paths were audited: the sync diffusion accumulation window passes a tuple and never concatenates rollouts; async diffusion keeps max_inflight=1; agentic now clears harness_status; every DP_SCATTER method takes a single Batch or a tuple of equal-size parts, so pytree_chunk is left unchanged.
• What is hardening vs. a live fix. Live today: agentic harness_status (now cleared) and guided HunyuanImage3 t2i/t2ti cond/uncond mis-pairing (fixed). Already guarded before, now unified: mixed output_version, primitive_metadata, multi-rollout async diffusion batches, root sample-id uniqueness. Data-dependent: producers dropping a media key per shard (Qwen3-Omni / Qwen3.5 fixed here). The remaining checks guard values from remote-code objects whose types cannot be verified in-repo.
• Interaction with open PR feat(tempflowgrpo): add TempFlow-GRPO with per-step branching #552 (TempFlow-GRPO / FlashGRPO). It concatenates step groups with different sde_indices and documents that the merged shared fields "are not per-sample truth". After this PR that concat raises, so feat(tempflowgrpo): add TempFlow-GRPO with per-step branching #552 should normalise those shared fields (per-sample truth already lives in its sde_index_per_sample CONCAT field) before merging.
• Compatible with feat(agentic): add bounded async turn-boundary rollout #558. It already moves the behavior version into row metadata and clears output_version before concatenating.
• Other shared fields. A SHARED field holding an object without a meaningful __eq__ would now raise after Ray deserialization. Today sampling_params carries only plain values (scheduler is set to None in requests, and sde_strategy is unset in every recipe).
• No config, checkpoint, or data-format changes.

Reviewer Notes

• AI-assisted (Cursor agent); the submitter reviewed the full diff.
• Duplicate-work check: searched open PRs and issues. Five open PRs touch the same files (perf(reward): send only scoring payload across slabs #313, feat(async): release rollout lanes during remote reward scoring #362, fix(async): make version metrics temporally consistent #540, feat(tempflowgrpo): add TempFlow-GRPO with per-step branching #552, feat(agentic): add bounded async turn-boundary rollout #558); none changes the same symbols. feat(tempflowgrpo): add TempFlow-GRPO with per-step branching #552 and feat(agentic): add bounded async turn-boundary rollout #558 are discussed above.
• Pre-PR review (.claude/skills/development/pr-review): two rounds with a fresh-context subagent on the same model as the author (no different model was available to request). Round 1 found that the trainside Qwen3-Omni / Qwen3.5 chat templates still dropped media keys; that is fixed here. Round 2 was clean.
• Follow-up reviews on 9ca249d0 and c2899833 with a fresh-context Grok 4.7 subagent (different model family), two rounds on the full branch diff. Fixed: unbatched TensorRef in a CONCAT field now raises; shared mismatches name the differing path and values (tensor diffs report the first differing element, since truncated reprs could print identical text for both sides); numpy compared by value; an unreachable torch.is_tensor fallback removed. Rejected: a finding that matching-row TensorRefs with different dtype or shape[1:] merge with item 0's metadata. cat_rows pads a ragged dim 1 on purpose, and the metadata comes from TensorRef._from_spans, which this PR does not touch; the base reached the same TensorRef.concat. The reviewer withdrew it. A later audit removed two guards no in-repo path can reach (pytree_chunk raising on a mismatched Batch; skipping compare=False dataclass fields) and found the HunyuanImage3 guided cond/uncond mis-pairing, fixed and re-reviewed clean. After the rebase onto 89ac5b03, a fresh-context Grok 4.7 review of the interaction with perf(trainer): stop driver-hydrating anchored rollout outputs before DP scatter #586 (driver-side TensorRefs) found nothing; a further audit dropped the numpy branch and the non-batched CONCAT tensor equality check, whose only justification (HunyuanImage3 diffusion tokenizer_output) left with 0777c210, reviewed clean by Grok 4.7.
• Start with unirl/distributed/tensor/batch.py: _first_difference, _shared_value, _any_present, _concat_value; then unirl/models/hunyuan_image3/conditions.py: split_cfg_fused.

Checklist

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

[CjhHa1]

Rebased onto main 89ac5b03; head is now 0777c210. No conflicts, and git range-diff shows all four patches unchanged.

Interaction with the new main commits. After #586, anchored AR and two-engine unified rollout outputs stay TensorRefs on the driver. The driver-side Sample.concat and DP scatter on those paths therefore see refs where they used to see hydrated tensors. On those paths every CONCAT ref has one row per sample: LatentSegment latents / sde_logp, and HunyuanImage3 fused / fused_uncond after split_cfg_fused. The SHARED refs (sigmas, indices, sde_indices) come from the same sampling params on every shard. So the new checks pass, and deep_hydrate is left only in the rollout image dump. #568, #583 and #554 touch none of this PR's files. A fresh-context Grok 4.7 review of git diff 89ac5b03 HEAD, focused on these interactions, found nothing.

Checks on 0777c210

• SKIP=no-commit-to-branch pre-commit run --from-ref 89ac5b03 --to-ref HEAD: all hooks passed.
• One-off CPU harnesses (torch CPU, ray stubbed): 27/27 general checks pass, and both HunyuanImage3 harnesses pass (cond/uncond pairing after merge and select, worker- and driver-side TensorRef).

GPU smoke (8×H20, 2–4 rollouts each). These ran on c2899833, the pre-rebase head with identical patches. I haven't re-run them on 0777c210.

Run	Result
SD3.5 FlowGRPO sync, DP merge	pass, ratio 1.0000
SD3.5 FlowGRPO sync, accumulate_rollouts=2 (tuple window)	pass, ratio 1.0000
BAGEL async, weight sync every rollout (combine_rollout_prompts across versions)	pass
Qwen3-Omni, mixed image + video rows	pass (see note)
HunyuanImage3 t2i, guidance_scale=5 + KV cache	pass, ratio 0.9997
HunyuanImage3 it2i, guidance_scale=5	pass, ratio 0.9999, clip 0
Agentic (calculator tool), default and tight	pass, ratio 0.9983–1.0013

Notes:

• Qwen3-Omni had a single engine replica, so image and video rows came from one call and no cross-replica merge happened. That merge is covered only by the CPU harness. Its first-update GSPO ratio was 1.0103 with every sequence clipped. The model inputs for that batch are the same before and after this PR, so I filed it separately as [Bug] Qwen3-Omni GSPO: first-update ratio != 1 with old_logp_source=replay (every sequence clipped) #590.
• HunyuanImage3 t2i clips 50–83% of samples. That is expected with old_logp_source=rollout: KV-cached sampling and full-prefill replay are not bit-exact, and the recipe's clip range is 2e-5. A cond/uncond mis-pairing would push the ratio far from 1, not to 0.9997.