Support emb/reranker/seq_cls padding_free (#6007)

Author: tastelikefeet

README.md

@@ -75,6 +75,7 @@ You can contact us and communicate with us by adding our group:
 
 
 ## 🎉 News
+- 🎁 2025.09.29: Support padding_free for embedding/reranker/seq_cls tasks, use `--padding_free true --task_type embedding/reranker/generative_reranker/seq_cls` to begin!
 - 🎁 2025.09.07: Added support for CHORD training algorithm. See the [documentation](./docs/source_en/Instruction/GRPO/AdvancedResearch/CHORD.md)
 - 🎁 2025.09.06: Ulysses can now be used with ring-attention, allowing sequences to be sharded into any number of chunks (no longer limited by the number of heads). The argument remains `--sequence_parallel_size N`.
 - 🎁 2025.09.02: Megatron-SWIFT now supports multimodal model training. Documentation can be found [here](./docs/source_en/Megatron-SWIFT/Multimodal-Model.md).

README_CN.md

@@ -71,6 +71,7 @@
 - **模型量化**：支持AWQ、GPTQ、FP8和BNB的量化导出，导出的模型支持使用vLLM/SGLang/LmDeploy推理加速，并支持继续训练。
 
 ## 🎉 新闻
+- 🎁 2025.09.29: 支持embedding/reranker/seq_cls任务的padding_free参数, 使用`--padding_free true --task_type embedding/reranker/generative_reranker/seq_cls`开始训练!
 - 🎁 2025.09.07: 支持CHORD训练算法，请查看[文档](docs/source/Instruction/GRPO/AdvancedResearch/CHORD.md)。
 - 🎁 2025.09.06: Ulysses现已支持与ring-attention结合使用，使得输入序列可以被切分成任意数量的块（不再受限于num_heads），命令参数仍然是`--sequence_parallel_size N`。
 - 🎁 2025.09.02: Megatron-SWIFT支持多模态模型训练。文档参考[这里](./docs/source/Megatron-SWIFT/多模态模型.md)。

examples/train/reranker/train_generative_reranker.sh

@@ -1,6 +1,7 @@
 nproc_per_node=4
 # 4*47G
 # losses: plugin/loss.py
+# only support --padding_side left
 NPROC_PER_NODE=$nproc_per_node \
 swift sft \
     --model Qwen/Qwen3-Reranker-4B \
@@ -11,6 +12,7 @@ swift sft \
     --load_from_cache_file true \
     --split_dataset_ratio 0.05 \
     --eval_strategy steps \
+    --padding_side left \
     --output_dir output \
     --eval_steps 100 \
     --num_train_epochs 1 \

examples/train/reranker/train_generative_reranker_listwise.sh

@@ -1,6 +1,7 @@
 nproc_per_node=4
 # 4*47G
 # losses: plugin/loss.py
+# only support --padding_side left
 NPROC_PER_NODE=$nproc_per_node \
 swift sft \
     --model Qwen/Qwen3-Reranker-4B \
@@ -10,6 +11,7 @@ swift sft \
     --dataset MTEB/scidocs-reranking \
     --load_from_cache_file true \
     --split_dataset_ratio 0.05 \
+    --padding_side left \
     --eval_strategy steps \
     --output_dir output \
     --eval_steps 100 \

swift/llm/model/patcher.py

@@ -4,22 +4,24 @@
 from contextlib import contextmanager
 from functools import wraps
 from types import MethodType
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import accelerate
 import torch
 import torch.nn as nn
 import transformers
 from accelerate.utils import find_device
 from packaging import version
+from peft import PeftModel
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from torch.nn.parallel import DistributedDataParallel as DDP
 from transformers import PreTrainedModel, dynamic_module_utils, trainer
 from transformers.modeling_outputs import SequenceClassifierOutputWithPast
 
 from swift.llm import deep_getattr, to_device, to_float_dtype
 from swift.utils import get_dist_setting, get_logger, is_mp, is_mp_ddp, safe_ddp_context
-from swift.utils.torch_utils import _get_max_memory, _sync_max_memory, get_device_count
+from swift.utils.torch_utils import (_get_max_memory, _sync_max_memory, get_cu_seqlens_from_position_ids,
+                                     get_device_count, get_position_ids_from_cu_seqlens)
 from .utils import HfConfigFactory
 
 logger = get_logger()
@@ -151,6 +153,8 @@ def _check_imports(filename) -> List[str]:
 
 
 def get_lm_head_model(model, model_meta=None, lm_heads=None):
+    if isinstance(model, PeftModel):
+        model = model.model
     model_meta = model_meta or model.model_meta
     lm_heads = lm_heads or ['lm_head']
     llm_prefix_list = getattr(model_meta.model_arch, 'language_model', None)
@@ -167,6 +171,81 @@ def get_lm_head_model(model, model_meta=None, lm_heads=None):
     return model
 
 
+def transformers_seq_cls_forward(self, *args, origin_forward, **kwargs):
+    labels = kwargs.pop('labels', None)
+    return_dict = kwargs.pop('return_dict', None)
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    input_ids = kwargs.get('input_ids')
+    inputs_embeds = kwargs.get('inputs_embeds')
+
+    output = origin_forward(*args, **kwargs)
+    if hasattr(output, 'logits'):
+        output.logits = output.logits.to(self.score.weight.dtype)
+    elif 'last_hidden_state' in output:
+        output.logits = output['last_hidden_state'].to(self.score.weight.dtype)
+    logits = self.score(output.logits)
+    if input_ids is not None:
+        batch_size = input_ids.shape[0]
+    else:
+        batch_size = inputs_embeds.shape[0]
+
+    if self.config.pad_token_id is None and batch_size != 1:
+        raise ValueError('Cannot handle batch sizes > 1 if no padding token is defined.')
+    if self.config.pad_token_id is None:
+        sequence_lengths = -1
+    else:
+        if output.get('attention_mask') is not None:
+            # When use padding_free in seq_cls tasks, `revert_padding_free` will add a attention_mask in the output
+            batch_size = output.get('attention_mask').shape[0]
+            sequence_lengths = output.get('attention_mask').sum(dim=1) - 1
+        elif input_ids is not None:
+            # 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]
+            sequence_lengths = sequence_lengths.to(logits.device)
+        elif kwargs.get('attention_mask') is not None:
+            sequence_lengths = kwargs['attention_mask'].sum(dim=1) - 1
+        else:
+            sequence_lengths = -1
+
+    pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+    loss = None
+    if labels is not None:
+        labels = labels.to(logits.device)
+        if self.config.problem_type is None:
+            if self.num_labels == 1:
+                self.config.problem_type = 'regression'
+            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                self.config.problem_type = 'single_label_classification'
+            else:
+                self.config.problem_type = 'multi_label_classification'
+
+        if self.config.problem_type == 'regression':
+            loss_fct = MSELoss()
+            if self.num_labels == 1:
+                loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+            else:
+                loss = loss_fct(pooled_logits, labels)
+        elif self.config.problem_type == 'single_label_classification':
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+        elif self.config.problem_type == 'multi_label_classification':
+            loss_fct = BCEWithLogitsLoss()
+            loss = loss_fct(pooled_logits, labels)
+    if not return_dict:
+        output = (pooled_logits, ) + output[1:]
+        return ((loss, ) + output) if loss is not None else output
+
+    return SequenceClassifierOutputWithPast(
+        loss=loss,
+        logits=pooled_logits,
+        past_key_values=output.past_key_values,
+        hidden_states=output.hidden_states,
+        attentions=output.attentions,
+    )
+
+
 def _patch_sequence_classification(model, model_meta):
     hidden_size = HfConfigFactory.get_config_attr(model.config, 'hidden_size')
     initializer_range = HfConfigFactory.get_config_attr(model.config, 'initializer_range')
@@ -183,73 +262,11 @@ def _patch_sequence_classification(model, model_meta):
             setattr(llm_model, lm_head, nn.Identity())
             break
 
-    origin_forward = llm_model.forward.__func__
+    origin_forward = llm_model.forward
 
-    @wraps(origin_forward)
+    @wraps(origin_forward.__func__)
     def new_forward(self, *args, **kwargs):
-        labels = kwargs.pop('labels', None)
-        return_dict = kwargs.pop('return_dict', None)
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        input_ids = kwargs.get('input_ids')
-        inputs_embeds = kwargs.get('inputs_embeds')
-
-        output = origin_forward(self, *args, **kwargs)
-        output.logits = output.logits.to(self.score.weight.dtype)
-        logits = self.score(output.logits)
-        if input_ids is not None:
-            batch_size = input_ids.shape[0]
-        else:
-            batch_size = inputs_embeds.shape[0]
-
-        if self.config.pad_token_id is None and batch_size != 1:
-            raise ValueError('Cannot handle batch sizes > 1 if no padding token is defined.')
-        if self.config.pad_token_id is None:
-            sequence_lengths = -1
-        else:
-            if input_ids is not None:
-                # 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]
-                sequence_lengths = sequence_lengths.to(logits.device)
-            else:
-                sequence_lengths = -1
-
-        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-
-        loss = None
-        if labels is not None:
-            labels = labels.to(logits.device)
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = 'regression'
-                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
-                    self.config.problem_type = 'single_label_classification'
-                else:
-                    self.config.problem_type = 'multi_label_classification'
-
-            if self.config.problem_type == 'regression':
-                loss_fct = MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(pooled_logits, labels)
-            elif self.config.problem_type == 'single_label_classification':
-                loss_fct = CrossEntropyLoss()
-                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
-            elif self.config.problem_type == 'multi_label_classification':
-                loss_fct = BCEWithLogitsLoss()
-                loss = loss_fct(pooled_logits, labels)
-        if not return_dict:
-            output = (pooled_logits, ) + output[1:]
-            return ((loss, ) + output) if loss is not None else output
-
-        return SequenceClassifierOutputWithPast(
-            loss=loss,
-            logits=pooled_logits,
-            past_key_values=output.past_key_values,
-            hidden_states=output.hidden_states,
-            attentions=output.attentions,
-        )
+        return transformers_seq_cls_forward(self, *args, origin_forward=origin_forward, **kwargs)
 
     llm_model.forward = MethodType(new_forward, llm_model)
 
@@ -454,6 +471,69 @@ def patch_tp_plan(load_model: bool):
     os.environ['WORLD_SIZE'] = WORLD_SIZE
 
 
+def revert_padding_free(outputs: Dict[str, Any], inputs: Dict[str, Any], padding_side='left'):
+    hidden_state_key = None
+    if 'last_hidden_state' in outputs:
+        hidden_state_key = 'last_hidden_state'
+    elif 'logits' in outputs:
+        hidden_state_key = 'logits'
+    elif 'token_embeddings' in outputs:
+        hidden_state_key = 'token_embeddings'
+
+    if hidden_state_key is None:
+        raise NotImplementedError()
+    last_hidden_state = outputs[hidden_state_key]
+    last_hidden_state = last_hidden_state.squeeze(dim=0)
+    if 'cu_seq_lens_q' in inputs:
+        position_ids = get_position_ids_from_cu_seqlens(inputs['cu_seq_lens_q'])
+    elif 'position_ids' in inputs and inputs['position_ids'].shape[0] == 1:
+        position_ids = inputs['position_ids']
+    else:
+        raise ValueError(
+            "revert_padding_free requires 'cu_seq_lens_q' or 'position_ids' in inputs, but neither was found.")
+
+    seq_lengths = []
+    pos = position_ids[0]
+    resets = torch.where(pos[1:] < pos[:-1])[0] + 1
+
+    if len(resets) == 0:
+        # Only one sequence in this batch item
+        seq_lengths = [pos.max().item() + 1]
+    else:
+        # Multiple sequences
+        start = 0
+        for end in resets:
+            seq_lengths.append(end - start)
+            start = end
+        seq_lengths.append(pos.shape[0] - start)
+
+    max_length = max(seq_lengths)
+    unpacked_logits = []
+    attention_mask = []
+
+    start = 0
+    for length in seq_lengths:
+        seq_state = last_hidden_state[start:start + length]
+        mask = torch.ones((seq_state.shape[0])).to(last_hidden_state.device)
+        padding = torch.zeros(
+            (max_length - length, last_hidden_state.shape[-1])).to(last_hidden_state.dtype).to(last_hidden_state.device)
+        attention_padding = torch.zeros((max_length - length)).to(last_hidden_state.device)
+        # re-padding
+        if padding_side == 'left':
+            seq_state = torch.cat((padding, seq_state), dim=0)
+            mask = torch.cat((attention_padding, mask), dim=0)
+        else:
+            seq_state = torch.cat((seq_state, padding), dim=0)
+            mask = torch.cat((mask, attention_padding), dim=0)
+        unpacked_logits.append(seq_state)
+        attention_mask.append(mask)
+        start += length
+    outputs[hidden_state_key] = torch.stack(unpacked_logits, dim=0)
+    inputs['attention_mask'] = torch.stack(attention_mask, dim=0).to(torch.int64)
+    outputs['attention_mask'] = inputs['attention_mask']
+    return outputs
+
+
 @contextmanager
 def patch_attach_align_device_hook_on_blocks():
     from accelerate import big_modeling

swift/llm/model/register.py

@@ -357,14 +357,6 @@ def make_inputs_require_grads(module, input, output):
 
         model.enable_input_require_grads = MethodType(enable_input_require_grads, model)
         tokenizer = model.tokenizer
-
-        def forward(self, **kwargs):
-            output = self._forward_origin(input=kwargs)
-            return {'last_hidden_state': output['sentence_embedding']}
-
-        if not hasattr(model, '_forward_origin'):
-            model._forward_origin = model.forward
-            model.forward = MethodType(forward, model)
     else:
         model = None
         tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)

swift/llm/template/base.py

@@ -521,7 +521,7 @@ def encode(self,
         if chosen.channel is not None:
             encoded['channel'] = chosen.channel
 
-        lengths = [0]
+        lengths = [0] if self.task_type not in {'reranker', 'generative_reranker'} else []
         for key in list(encoded.keys()):
             if encoded[key] is None:
                 encoded.pop(key)
@@ -532,7 +532,10 @@ def encode(self,
                 elif isinstance(value, (tuple, list)):
                     lengths += value
         if return_length:
-            encoded['length'] = sum(lengths)
+            if self.task_type in {'reranker', 'generative_reranker'}:
+                encoded['length'] = lengths
+            else:
+                encoded['length'] = sum(lengths)
         else:
             encoded.pop('length', None)
         if return_template_inputs:
@@ -1542,10 +1545,13 @@ def _reranker_data_collator(self,
                 max_positive = min(positive_num, max_positive_samples)
                 max_negative = min(negative_num, max_negative_samples)
                 for i in random.sample(range(positive_num), max_positive):
-                    new_batch.append({'input_ids': b['input_ids'][i]})
+                    new_batch.append({'input_ids': b['input_ids'][i], 'length': b['length'][i]})
                     labels_list.append(1)
                     for j in random.sample(range(negative_num), max_negative):
-                        new_batch.append({'input_ids': b['input_ids'][j + positive_num]})
+                        new_batch.append({
+                            'input_ids': b['input_ids'][j + positive_num],
+                            'length': b['length'][j + positive_num]
+                        })
                         labels_list.append(0)
 
             res = self._data_collator(new_batch, padding_to=padding_to)

swift/trainers/arguments.py

@@ -154,6 +154,9 @@ class SwiftArgumentsMixin(RLHFArgumentsMixin, TrainArgumentsMixin):
     train_type: Optional[str] = None
     local_repo_path: Optional[str] = None
     galore_config: Optional[GaLoreConfig] = None
+    padding_side: Optional[str] = None
+    padding_free: Optional[bool] = None
+    task_type: Optional[str] = None
 
     def __post_init__(self):
         if hasattr(self, 'output_dir'):