【Bug已解决】Kosmos2.5 index error on long ocr input 解决方案一、现象长什么样Kosmos2.5 是一个面向 OCR/文档理解的多模态模型输入通常是一张文档图 一段提示文本。当文档较长高分辨率扫描、多行密集文字时generate或forward抛索引错误# 现象 A图像 patch 索引越界 IndexError: index 1024 is out of bounds for dimension 0 with size 1024 File .../models/kosmos2_5/modeling_kosmos2_5.py, line 210, in forward img_feat image_features[image_token_indices] # 现象 B截断后位置错位 RuntimeError: index -1 is out of bounds for dimension 0 with size 0 # 长输入被截断到 max_length但 image_token 的占位索引还指着被截掉的位置 # 现象 Cbatch 内长短不一拼 padding 后越界 ValueError: too many indices for tensor of dimension 1 # padding 把序列补齐到统一长度但 image_token_indices 仍是原始未 padding 的下标最典型的触发一张高分辨率文档图被切成很多 patch比如 1024 个加上 OCR 提示文本后总长度超过max_length2048截断逻辑只截了文本侧却忘了同步修正图像 token 的索引 → 索引越界。二、背景Kosmos2.5 的输入构造流程是图像经 vision encoder 切成若干 patch 特征image_features长度 patch 数 N。文本里用特殊imagetoken 占位 tokenizer 后这些占位被展开成 N 个 image token分布在序列的不同位置。forward 时模型根据image_token_indices这些 image token 在序列里的下标去image_features里取对应特征拼回序列。这套机制依赖一个不变式image_token_indices里的最大值 Npatch 总数且截断/ padding 后这些索引必须同步更新。当长 OCR 输入触发截断为了塞进max_length或 padding为了 batch这个不变式被打破就出现上面的索引错误。三、根因根因有三类截断只截断文本不修正 image_token_indices。processor在max_length超限时直接裁掉文本 token 尾部但 image token 的下标是相对裁剪前序列算的。裁剪后序列变短那些指向被裁区域的 image token 下标变成非法指向越界或负位置。padding 后索引未平移/未 mask。 batch 推理时短样本被 pad 到最长。padding 在序列前面或后面插入了 dummy token但image_token_indices仍是原始下标没有随 padding 偏移 → 在 padding 之后的位置取 image_features 时错位。patch 数 N 与 image token 数不一致。 长文档切的 patch 数超过image_features实际长度例如 image encoder 有自己内部的max_patches限制超出的 patch 被丢弃但 tokenizer 展开的 image token 数仍是按未限制算的 → 下标越界。四、最小可运行复现下面用纯 Python 模拟截断文本后 image_token_indices 越界的逻辑from typing import List def build_image_token_indices(seq_len: int, n_image_tokens: int) - List[int]: 模拟在序列末尾均匀放置 n_image_tokens 个 image token 的下标。 step max(1, seq_len // n_image_tokens) return list(range(0, seq_len, step))[:n_image_tokens] def truncate(seq_len: int, max_len: int) - int: return min(seq_len, max_len) # 正常短输入 seq_len 500 n_img 100 idx build_image_token_indices(seq_len, n_img) print(短输入 max idx:, max(idx), patch 数 N:, n_img) # 合法 # 长输入触发截断 max_len 300 new_len truncate(seq_len, max_len) new_idx build_image_token_indices(seq_len, n_img) # 索引仍按旧 seq_len 算 print(截断后序列长:, new_len, 但 image idx max:, max(new_idx)) assert max(new_idx) new_len, 复现成功截断后 image token 索引越界 # 修正版截断时同步裁剪 image_token_indices def truncate_with_indices(seq_len, max_len, idx): return [i for i in idx if i max_len] fixed truncate_with_indices(seq_len, max_len, idx) print(修正后 image idx:, fixed, max:, max(fixed) if fixed else None) assert all(i new_len for i in fixed), 修正失败运行后原new_idx的最大值~499超过了截断后的序列长度300触发越界修正函数把越界索引裁掉恢复不变式。五、解决方案第一层最小直接修复最快的止血在调用 processor / 截断前手动清洗 image_token_indices使其始终落在有效范围内import torch def sanitize_image_token_indices(image_token_indices, seq_len, image_features_len): 第一层修复保证索引在 [0, seq_len) 且 image_features_len。 valid [] for i in image_token_indices: if 0 i seq_len and i image_features_len: valid.append(i) # 若全部越界极端长输入退化为均匀取样 image_features if not valid and image_features_len 0: step max(1, image_features_len // seq_len) if seq_len else 1 valid list(range(0, image_features_len, max(step, 1)))[:seq_len] return valid # 使用示意在构造模型输入后、forward 前 inputs processor(imagesdoc_image, textprompt, return_tensorspt, truncationTrue, max_length2048) seq_len inputs[input_ids].shape[1] image_token_indices (inputs[input_ids][0] processor.image_token_id).nonzero().flatten().tolist() # 取出 image_features来自 vision encoder valid_idx sanitize_image_token_indices(image_token_indices, seq_len, image_features.shape[0]) # 用合法索引重建或传给模型一个已清洗的 indices 参数 assert max(valid_idx) image_features.shape[0], 仍有越界检查 image_features 长度第一层让用户立刻消除IndexError长 OCR 文档也能跑。六、解决方案第二层结构性改进把索引与序列同步做成KosmosIndexSync在 processor 和模型之间统一维护不变式from dataclasses import dataclass from typing import List dataclass class KosmosIndexSync: 维护 image_token_indices 与(截断后)序列长度、image_features 长度的一致。 max_patches: int 1024 def sync_after_truncation(self, indices: List[int], new_seq_len: int) - List[int]: kept [i for i in indices if 0 i new_seq_len] # 同时保证不超过 image_features 实际容量 kept [i for i in kept if i self.max_patches] # 若因截断丢失过多 image token从 image_features 均匀补回 if len(kept) max(1, len(indices) // 2) and self.max_patches 0: step max(1, self.max_patches // max(new_seq_len, 1)) kept list(range(0, self.max_patches, step))[:new_seq_len] return kept def sync_after_padding(self, indices: List[int], pad_left: int) - List[int]: # padding 在左侧插入 dummy 时所有索引右移 pad_left return [i pad_left for i in indices] # 使用 sync KosmosIndexSync(max_patches1024) inputs processor(imagesdoc_image, textprompt, return_tensorspt, truncationTrue, max_length2048, paddingmax_length) seq_len inputs[input_ids].shape[1] raw_idx (inputs[input_ids][0] processor.image_token_id).nonzero().flatten().tolist() valid sync.sync_after_truncation(raw_idx, seq_len) if inputs.get(attention_mask) is not None: pad_left int((inputs[attention_mask][0] 0).sum().item()) # 左 padding 数量 valid sync.sync_after_padding(valid, pad_left)KosmosIndexSync把截断同步 左 padding 平移 容量上限三件事集中处理保证image_token_indices永远落在合法区间。七、解决方案第三层断言 / CI 守护用 pytest 固化长输入不产生越界索引的契约import pytest def test_indices_within_bounds_after_truncation(): from index_sync import KosmosIndexSync sync KosmosIndexSync(max_patches1024) # 模拟长 OCR1024 个 image token序列被截到 300 indices list(range(0, 10000, 10))[:1024] new_len 300 valid sync.sync_after_truncation(indices, new_len) assert all(0 i new_len for i in valid), 截断后索引仍越界 assert all(i 1024 for i in valid), 索引超过 image_features 容量 def test_padding_shifts_indices(): from index_sync import KosmosIndexSync sync KosmosIndexSync() idx [5, 10, 15] shifted sync.sync_after_padding(idx, pad_left4) assert shifted [9, 14, 19], 左 padding 后索引应整体右移 def test_no_indexerror_on_long_ocr(): # 端到端长文档不应抛 IndexError import torch from unittest.mock import MagicMock image_features torch.randn(1024, 64, 64) # N1024 patch indices list(range(0, 10000, 10))[:1024] new_len 300 valid [i for i in indices if i new_len and i image_features.shape[0]] # 取特征不应越界 feats image_features[valid] assert feats.shape[0] len(valid)CI 跑pytest tests/test_kosmos2_5_long_ocr.py以后只要截断/padding 逻辑又忘了同步索引测试立刻红灯。八、排查清单当 Kosmos2.5 在长 OCR 输入上报索引错误按顺序查IndexError: index X is out of bounds for image_features→image_token_indices越界先用sanitize_image_token_indices清洗。index -1 is out of bounds→ 截断把 image token 全裁掉了需要同步裁剪或均匀补回。batch 推理报too many indices→ padding 后索引未平移用sync_after_padding右移。确认image_features实际 patch 数 N与 tokenizer 展开的 image token 数是否一致不一致要限制max_patches。长期方案把索引同步逻辑收进 processor返回已清洗的 indices而不是让模型侧去猜。九、小结Kosmos2.5: index error on long ocr input 的根因是图像 token 的下标image_token_indices在长输入触发截断/padding 后未同步更新破坏了下标 patch 数 序列长的不变式于是索引越界。第一层forward 前用手动sanitize_image_token_indices清洗越界索引长文档立即能跑。第二层用KosmosIndexSync统一处理截断裁剪 左 padding 平移 patch 容量上限结构性保证索引合法。第三层pytest 断言截断后索引在界内、padding 后正确平移、长 OCR 端到端不越界防止回归。记住多模态模型里跨模态的索引image token ↔ image_features必须在每次序列变换截断/padding后同步修正这个不变式一旦破坏就是索引错误。