From be3a6c1178c76f2418383d4231e537fd686e4cf4 Mon Sep 17 00:00:00 2001 From: Arthur031221 <124417490+Arthur031221@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:00:27 +0800 Subject: [PATCH] [shardformer]: pick the last non-pad token in pipeline sequence classification The pipeline forwards of LlamaForSequenceClassification and Qwen2ForSequenceClassification located the pooled token with ne(input_ids, pad_token_id).sum(-1) - 1. That is the last real token only for right-padded rows. For a left-padded row it points into the middle of the sequence, so the logits and loss differ from the model without pipeline parallelism. Use the rightmost non-pad token, as transformers and the Qwen3 pipeline forward already do. Signed-off-by: Arthur031221 <124417490+Arthur031221@users.noreply.github.com> --- colossalai/shardformer/modeling/_llama_tf4.py | 5 +- colossalai/shardformer/modeling/_llama_tf5.py | 5 +- colossalai/shardformer/modeling/_qwen2_tf4.py | 5 +- .../test_model/test_seq_cls_padding.py | 86 +++++++++++++++++++ 4 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 tests/test_shardformer/test_model/test_seq_cls_padding.py diff --git a/colossalai/shardformer/modeling/_llama_tf4.py b/colossalai/shardformer/modeling/_llama_tf4.py index 3f9b7b12880c..9f334fb9114d 100644 --- a/colossalai/shardformer/modeling/_llama_tf4.py +++ b/colossalai/shardformer/modeling/_llama_tf4.py @@ -444,7 +444,10 @@ def llama_for_sequence_classification_forward( sequence_lengths = -1 else: if input_ids is not None: - sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + # Take the rightmost non-pad token so that both left and right padding work + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + sequence_lengths = (token_indices * non_pad_mask).argmax(-1) else: sequence_lengths = -1 diff --git a/colossalai/shardformer/modeling/_llama_tf5.py b/colossalai/shardformer/modeling/_llama_tf5.py index 01013a310592..0dc82f4ca0c5 100644 --- a/colossalai/shardformer/modeling/_llama_tf5.py +++ b/colossalai/shardformer/modeling/_llama_tf5.py @@ -486,7 +486,10 @@ def llama_for_sequence_classification_forward( sequence_lengths = -1 else: if input_ids is not None: - sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + # Take the rightmost non-pad token so that both left and right padding work + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + sequence_lengths = (token_indices * non_pad_mask).argmax(-1) else: sequence_lengths = -1 diff --git a/colossalai/shardformer/modeling/_qwen2_tf4.py b/colossalai/shardformer/modeling/_qwen2_tf4.py index 80dfb003dccf..85b8d33408f3 100644 --- a/colossalai/shardformer/modeling/_qwen2_tf4.py +++ b/colossalai/shardformer/modeling/_qwen2_tf4.py @@ -434,7 +434,10 @@ def qwen2_for_sequence_classification_forward( sequence_lengths = -1 else: if input_ids is not None: - sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + # Take the rightmost non-pad token so that both left and right padding work + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + sequence_lengths = (token_indices * non_pad_mask).argmax(-1) else: sequence_lengths = -1 diff --git a/tests/test_shardformer/test_model/test_seq_cls_padding.py b/tests/test_shardformer/test_model/test_seq_cls_padding.py new file mode 100644 index 000000000000..17d8788f17e4 --- /dev/null +++ b/tests/test_shardformer/test_model/test_seq_cls_padding.py @@ -0,0 +1,86 @@ +import pytest +import torch +from transformers import LlamaConfig, LlamaForSequenceClassification, Qwen2Config, Qwen2ForSequenceClassification + +import colossalai +from colossalai.cluster import ProcessGroupMesh +from colossalai.logging import disable_existing_loggers +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer import ShardConfig +from colossalai.shardformer.modeling.llama import LlamaPipelineForwards +from colossalai.shardformer.modeling.qwen2 import Qwen2PipelineForwards +from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn + +NUM_LAYERS = 2 +COMMON_CONFIG = dict( + vocab_size=64, hidden_size=32, num_hidden_layers=NUM_LAYERS, num_attention_heads=4, pad_token_id=0, num_labels=3 +) + +MODELS = { + "llama": ( + LlamaConfig(intermediate_size=64, num_key_value_heads=4, **COMMON_CONFIG), + LlamaForSequenceClassification, + LlamaPipelineForwards.llama_for_sequence_classification_forward, + ), + "qwen2": ( + Qwen2Config(intermediate_size=64, num_key_value_heads=4, **COMMON_CONFIG), + Qwen2ForSequenceClassification, + Qwen2PipelineForwards.qwen2_for_sequence_classification_forward, + ), +} + +# The first row of each batch is padded, the second one is not. +INPUT_IDS = { + "left": [[0, 0, 5, 6, 7, 8], [9, 10, 11, 12, 13, 14]], + "right": [[5, 6, 7, 8, 0, 0], [9, 10, 11, 12, 13, 14]], +} + + +def run_seq_cls_padding(name): + device = torch.device("cuda") + stage_manager = PipelineStageManager(ProcessGroupMesh(1), pipeline_axis=0) + assert stage_manager.is_first_stage() and stage_manager.is_last_stage() + shard_config = ShardConfig(enable_tensor_parallelism=False, pipeline_stage_manager=stage_manager) + + config, model_cls, pipeline_forward = MODELS[name] + torch.manual_seed(0) + model = model_cls(config).to(device).eval() + for padding_side, ids in INPUT_IDS.items(): + input_ids = torch.tensor(ids, device=device) + 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 = pipeline_forward( + model, + input_ids=input_ids, + attention_mask=attention_mask, + stage_manager=stage_manager, + stage_index=[0, NUM_LAYERS], + shard_config=shard_config, + ).logits + assert torch.allclose(output, expected, atol=1e-5), f"{name} with {padding_side} padding" + + +def check_seq_cls_padding(rank, world_size, port, name): + disable_existing_loggers() + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + run_seq_cls_padding(name) + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_llama_seq_cls_padding(): + spawn(check_seq_cls_padding, 1, name="llama") + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +@clear_cache_before_run() +def test_qwen2_seq_cls_padding(): + spawn(check_seq_cls_padding, 1, name="qwen2") + + +if __name__ == "__main__": + test_llama_seq_cls_padding() + test_qwen2_seq_cls_padding()