[BUG]: 工作流序列分类对左边填充的输入选择了错误的标记（LLaMA、Qwen2、OPT）

import torch from transformers import LlamaConfig, LlamaForSequenceClassification import colossalai from colossalai.cluster import ProcessGroupMesh from colossalai.pipeline.stage_manager import PipelineStageManager from colossalai.shardformer import ShardConfig from colossalai.shardformer.modeling.LLaMA import LlamaPipelineForwards colossalai.launch(rank=0, world_size=1, host="localhost", port=29512, backend="nccl") stage_manager = PipelineStageManager(ProcessGroupMesh(1), pipeline_axis=0) shard_config = ShardConfig(enable_tensor_parallelism=False, pipeline_stage_manager=stage_manager) config = LlamaConfig( vocab_size=64, hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, pad_token_id=0, num_labels=3, ) torch.manual_seed(0) model = LlamaForSequenceClassification(config).cuda().eval()

first row is left-padded, second row has no padding

input_ids = torch.tensor([[0, 0, 5, 6, 7, 8], [9, 10, 11, 12, 13, 14]]).cuda() attention_mask = (input_ids != config.pad_token_id).long() with torch.no_grad(): expected = model(input_ids=input_ids, attention_mask=attention_mask).logits output = LlamaPipelineForwards.llama_for_sequence_classification_forward( model, input_ids=input_ids, attention_mask=attention_mask, stage_manager=stage_manager, stage_index=[0, config.num_hidden_layers], shard_config=shard_config, ).logits print("row matches model.forward:", torch.isclose(output, expected, rtol=1e-05, atol=1e-05))