微调LLaMA做Token分类遇CUDA错误及SFTTrainer异常求助

一、默认Trainer的CUDA错误原因 #

1. 任务架构不匹配导致张量维度冲突 #

你从Roberta迁移的多标签Token分类逻辑（独热编码输出）属于序列标注任务，而LLaMA本质是自回归语言模型，二者的输出结构、损失逻辑完全不同：

• Roberta Token分类输出维度为 (batch_size, seq_len, num_labels)，损失直接基于每个token的多标签概率计算；
• LLaMA默认输出维度是 (batch_size, seq_len, vocab_size)，对应下一个token的预测分布。若强行替换输出头却未同步调整损失函数、标签维度，会触发张量形状不匹配，导致CUDA运算时索引越界，进而引发断言失败和CUBLAS初始化错误（CUBLAS对张量合法性要求极高）。

2. 标签数据格式错误 #

多标签Token分类的标签需是 (batch_size, seq_len, num_labels) 的独热张量，若你仍沿用CasualLM微调时的 (batch_size, seq_len) 单标签格式，损失计算时会出现维度广播错误，触发CUDA层面的异常。

二、SFTTrainer抛出NotImplemented异常的原因 #

SFTTrainer是专为**自回归语言模型的监督微调（SFT）**设计的，仅支持标准CausalLMOutput类型的输出，不兼容Token分类任务的TokenClassifierOutput。当模型输出不符合CasualLM格式时，SFTTrainer内部的损失计算、日志跟踪逻辑会因找不到对应处理分支，直接抛出NotImplementedError。

三、修复步骤 #

• 重构LLaMA输出头：将默认的语言建模头替换为 Linear(in_features=llama_hidden_dim, out_features=num_labels)，确保输出维度为 (batch_size, seq_len, num_labels)。
• 自定义损失函数：多标签任务需使用BCEWithLogitsLoss（而非默认交叉熵损失），设置reduction='none'后按token维度计算损失，适配序列标注场景。
• 修正标签数据格式：确保数据加载器输出的标签是(batch_size, seq_len, num_labels)的独热编码张量，而非单token标签。
• 调试CUDA问题：临时禁用自动混合精度（AMP），排查是否因AMP对非标准张量的兼容问题导致CUBLAS错误；同时检查显卡显存占用，3090虽显存充足，但张量维度不匹配可能导致异常显存分配。
• 弃用SFTTrainer：Token分类不属于SFT任务范畴，直接使用默认Trainer配合自定义训练参数即可。

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