使用Seq2SeqTrainer在GPU训练时Assertion `srcIndex < srcSelectDimSize`失败求助

问题 #

使用HuggingFace Transformer框架的Seq2SeqTrainer微调预训练模型时，CPU上运行正常，但在NVIDIA GeForce 4090/3090 GPU上触发错误：

Assertion srcIndex < srcSelectDimSize failed

自定义Seq2SeqTrainer调试张量信息后，运行trainer.train()时前14步输出正常，随后抛出：

RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

执行resize_token_embeddings(len(tokenizer) + 10000)后，错误变为：

RuntimeError: shape '[-1, 58685]' is invalid for input of size 109896000

解决方案 #

1. 定位索引越界根源 #

• 启用CUDA设备端断言：设置环境变量export TORCH_USE_CUDA_DSA=1后重新运行训练，会输出更具体的越界位置，明确是哪个张量的索引超出范围。
• 核对输入张量的索引范围：重点检查labels、input_ids等张量的取值，确保所有索引值都小于对应词典的大小（比如labels中的值必须小于tokenizer.vocab_size），GPU运行时异步操作可能放大数据加载时的索引异常。
• 验证数据加载逻辑：对比CPU和GPU训练时的同一批输入数据，检查collate_fn在GPU环境下是否出现padding错误或张量形状异常。

2. 修复嵌入层形状不匹配问题 #

• 精准调整嵌入层尺寸：resize_token_embeddings的参数需与实际新增token数量匹配，不能随意添加固定值。先统计新增token数（如new_tokens = len(custom_vocab) - len(tokenizer)），再执行model.resize_token_embeddings(len(tokenizer) + new_tokens)。
• 校验索引与嵌入层的匹配：确保input_ids、labels中的最大索引值不超过调整后的model.config.vocab_size，存在无效索引时直接过滤或替换。

3. GPU环境专属排查 #

• 统一张量设备：训练前执行model = model.to('cuda')，同时在DataLoader中设置pin_memory=True，避免部分张量在CPU、部分在GPU的混合设备状态。
• 降低训练批量：GPU大批次可能触发隐性内存或形状问题，尝试减小per_device_train_batch_size后重新运行，排查是否为批量过大导致的异常。
• 升级框架版本：旧版本PyTorch/Transformers在4090等新GPU上存在兼容性bug，升级到PyTorch>=2.0、Transformers>=4.30的稳定版，可解决底层CUDA操作问题。

内容的提问来源于stack exchange，提问作者Raptor