[Arthur031221]

Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install

Issue number

Refs #6457

What does this PR do?

With pipeline parallelism, LlamaForSequenceClassification and Qwen2ForSequenceClassification pool the wrong token on left-padded batches, so results differ from the unpipelined model.

The pipeline forwards in _llama_tf4.py, _llama_tf5.py and _qwen2_tf4.py use torch.ne(input_ids, pad_token_id).sum(-1) - 1, which only gives the last real token with right padding. This now takes the rightmost non-pad token, like transformers and the Qwen3 forward do.

opt.py has the same line, but its pipeline forward raises a ValueError on the pinned transformers 5.17.0, so I left it alone and used Refs. The eq(...).argmax(-1) - 1 poolers in gpt2, bloom, falcon, gptj and mistral are different and untouched.

The new test compares both pipeline forwards with the model's own forward on left and right padded batches. Without the fix Llama fails on 5.17.0 and Qwen2 on 4.51.3.

Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

Do you enjoy contributing to Colossal-AI?

• Yes, I do.
• No, I don't.

Assisted by Claude/Codex.

[mira687]

Grep says ne(input_ids, pad).sum(-1) - 1 appears in exactly these three files, so the scope is right for that pattern. But five of the other last-token poolers in shardformer/modeling are on a different legacy pattern and this PR doesn't reach them — falcon.py:517, gpt2.py:764, gptj.py:427, bloom.py:488, mistral.py:331:

# if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
sequence_lengths = sequence_lengths % input_ids.shape[-1]   # mistral.py omits this second line

Those five are not hit by the left-padding bug: on [0, 0, 5, 6, 7, 8] the argmax is 0 and -1 wraps round to 5, which is what last_non_pad_token gives. They diverge somewhere else — when a pad id sits before the last real token, which is the ordinary case once pad_token_id == eos_token_id. Indices for pad 0, L = 6:

row	ne.sum-1	eq.argmax-1 % L	last_non_pad
[0,0,5,6,7,8]	3	5	5
[5,6,7,8,0,0]	3	3	3
[0,0,5,0,7,8]	2	5	5
[5,0,7,8,0,0]	2	0	3

Only the last row separates those five from last_non_pad_token. qwen3.py:434-436 is already on it, so this PR takes the directory to 4 of 9 aligned with the pinned transformers==4.51.3.

That is the index arithmetic computed on its own (torch 2.2.2, CPU) — I have not run those five pipeline forwards. And the modulo is there deliberately for ONNX, so it isn't a free swap; happy to leave the five if that's the intended scope.

[Arthur031221]

Thanks. Your last row is right, and it disproves the sentence in the PR body that said the eq(...).argmax(-1) - 1 expression in the other five files was correct for left-padded, right-padded and unpadded rows. It points one position before the first pad_token_id in the row, so a row that contains pad_token_id among its real tokens can be pooled at the wrong position. I reproduced your four rows in plain Python and corrected that sentence in the body. I have only checked the index arithmetic, not those five forwards, so this PR stays on the three files with the left-padding bug.