fix(hi3): batched chat-template tokenization encodes only the first sample #520

System info

• UniRL: 7a530364 (upstream/main); found while validating fix(hi3): support batched trainside AR with dynamic KV cache #479 on its head 01d30dbe
• Checkpoint: tencent/HunyuanImage-3.0-Instruct, its bundled tokenization_hunyuan_image_3.py (identical batch_gen_infer to HF main)
• GPU observation: 8x H20, Python 3.13, torch 2.10.0-rc6, transformers 5.12.1
• CPU reproduction (tokenizer only): transformers 4.57.6

Rollout / inference engine

Train-side only

Domain

unified_model / diffusion — HunyuanImage3 in-process pipeline.generate (t2t, i2t, t2ti, t2i, it2i)

Recipe / config-name

N/A. The in-process HI3 recipes (unified_model/hi3_trainside_t2i, unified_model/hi3_it2i) avoid this by pinning rollout.forward_batch_size: 1.

How you ran it

• Modified recipe or my own script

Reproduction

Load only the checkpoint's tokenizer and call its chat template with three prompts:

from hi3.tokenization_hunyuan_image_3 import HunyuanImage3TokenizerFast
from hi3.configuration_hunyuan_image_3 import HunyuanImage3Config
from hi3.image_processor import HunyuanImage3ImageProcessor

cfg = HunyuanImage3Config.from_pretrained(CKPT)
tok = HunyuanImage3TokenizerFast.from_pretrained(CKPT, model_version="instruct")
prompts = ["What is the capital of France?", "Write a haiku about autumn, please, in English.", "Hi"]
common = dict(image_base_size=cfg.image_base_size, sequence_template="instruct", drop_think=False)

tok.apply_chat_template(batch_prompt=prompts, mode="gen_text", bot_task="auto", cfg_factor=1, **common)
info = HunyuanImage3ImageProcessor(cfg).build_gen_image_info("1024x1024")
for f in (1, 2):
    tok.apply_chat_template(batch_prompt=prompts, mode="gen_image", bot_task="image", cfg_factor=f,
                            batch_gen_image_info=[info] * 3, **common)

The same happens through UniRL on the real checkpoint: HunyuanImage3TextEmbedStage.embed_for_ar(Texts(texts=<4 prompts>)) returns fused.input_ids of shape [1, 1238], and i2t with two images then fails in instantiate_continuous_tokens with IndexError: index 1 is out of bounds for dimension 0 with size 1.

Root cause, in the checkpoint's tokenization_hunyuan_image_3.py: apply_general_template(batchify=True) calls

self.batch_gen_infer(
    infer_fn=self.apply_general_template,
    prompt_list=[[]],
    infer_fn_kwargs_list=[dict(message_list=message_list_i, ...) for message_list_i in message_list],
    ...
)

and batch_gen_infer iterates zip(prompt_list, infer_fn_kwargs_list), so only the first sample is encoded and the rest are dropped without an error.

Impact today:

• The tokenizer step is silent. At the mode level, t2t / i2t / t2i / t2ti then fail at Part.fill (Part.primitives[...] batch 1 != sample_ids count N), and it2i fails its explicit tokenizer row-count check. No wrong data reaches training.
• Both in-process recipes pin forward_batch_size: 1, one sample per generate call (see the comment added in feat(hi3): pure text-to-image trainside FlowGRPO recipe #342), so batched HI3 in-process rollout is unavailable rather than wrong.

Expected behavior

embed_for_ar and embed_for_gen_image return one row per prompt (times the CFG factor), so in-process HI3 generation can run with forward_batch_size > 1.

Suggested direction: encode each row separately in HunyuanImage3TextEmbedStage and merge with HunyuanImage3FusedMultimodalCondition.concat, stacking tokenizer_output.real_pos. That depends on #479, whose _pad_positions makes ragged position_ids (the KV-cache write indices) safe to concat. The diffusion side of a ragged concat (zero-padded rope_cache, fully masked pad query rows) needs its own GPU check before lifting forward_batch_size: 1. A cheaper interim guard is to raise in _apply_chat_template when the returned row count differs from len(prompts) * cfg_factor.