[Train] Reduce padding-free embedding output memory

swift/model/patcher.py

@@ -572,6 +572,45 @@ def revert_padding_free(outputs: Dict[str, Any], inputs: Dict[str, Any], padding
     return outputs
 
 
+def select_last_packed_states(outputs: Dict[str, Any], inputs: Dict[str, Any]):
+    """Keep only the last hidden state of each sequence in a packed batch.
+
+    Embedding models only need one hidden state per sequence. Re-padding the
+    complete packed output before pooling can allocate a large temporary tensor,
+    especially for multimodal batches with uneven sequence lengths.
+    """
+    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()
+
+    hidden_states = outputs[hidden_state_key]
+    if hidden_states.shape[0] != 1:
+        raise ValueError(f'Expected a padding-free batch dimension of 1, but got shape {hidden_states.shape}.')
+
+    if 'cu_seq_lens_q' in inputs:
+        cu_seq_lens_q = inputs['cu_seq_lens_q']
+        sequence_end_indices = cu_seq_lens_q[1:].to(device=hidden_states.device, dtype=torch.long) - 1
+    elif 'position_ids' in inputs and inputs['position_ids'].shape[0] == 1:
+        position_ids = inputs['position_ids'][0]
+        resets = torch.where(position_ids[1:] < position_ids[:-1])[0] + 1
+        sequence_end_indices = torch.cat([resets - 1, position_ids.new_tensor([position_ids.shape[0] - 1])])
+        sequence_end_indices = sequence_end_indices.to(device=hidden_states.device, dtype=torch.long)
+    else:
+        raise ValueError("select_last_packed_states requires 'cu_seq_lens_q' or 'position_ids' in inputs, "
+                         'but neither was found.')
+
+    # Preserve a singleton sequence dimension for patch_output_normalizer.
+    outputs[hidden_state_key] = hidden_states[0].index_select(0, sequence_end_indices).unsqueeze(1)
+    return outputs
+
+
 def gather_sequence_parallel_outputs(
     outputs: Dict[str, Any],
     tensor_keys: Optional[List[str]] = None,

swift/trainers/mixin.py

@@ -47,7 +47,8 @@
 from swift.loss import loss_map
 from swift.metrics import MeanMetric, compute_acc, eval_metrics_map
 from swift.model import get_llm_model, get_lm_head_model, save_checkpoint
-from swift.model.patcher import gather_sequence_parallel_outputs, revert_padding_free, transformers_seq_cls_forward
+from swift.model.patcher import (gather_sequence_parallel_outputs, revert_padding_free, select_last_packed_states,
+                                 transformers_seq_cls_forward)
 from swift.optimizers import OptimizerCallback, optimizers_map
 from swift.sequence_parallel import SequenceParallelDispatcher, SequenceParallelSampler, sequence_parallel
 from swift.template import Template, update_generation_config_eos_token
@@ -836,17 +837,21 @@ def sp_gather_hook(module, args, input, output):
                 if pf_enabled:
                     if sp_enabled:
 
-                        def revert_padding_free_hook(module, args, input, output):
+                        def padding_free_hook(module, args, input, output):
                             # Use full packed position ids cached by sequence_parallel.prepare_inputs
                             position_ids = sequence_parallel.real_position_ids
                             tmp_input = {'position_ids': position_ids}
+                            if task_type == 'embedding':
+                                return select_last_packed_states(output, tmp_input)
                             return revert_padding_free(output, tmp_input, padding_side)
                     else:
 
-                        def revert_padding_free_hook(module, args, input, output):
+                        def padding_free_hook(module, args, input, output):
+                            if task_type == 'embedding':
+                                return select_last_packed_states(output, input)
                             return revert_padding_free(output, input, padding_side)
 
-                    hooks.append(revert_padding_free_hook)
+                    hooks.append(padding_free_hook)
 
                 if hooks:
                     _register_llm_hooks_in_order(llm_model, hooks)

tests/models/test_patcher.py

@@ -0,0 +1,29 @@
+import torch
+
+from swift.model.patcher import select_last_packed_states
+
+
+def test_select_last_packed_states_with_cu_seq_lens():
+    hidden_states = torch.arange(16, dtype=torch.float32).view(1, 8, 2).requires_grad_()
+    outputs = {'logits': hidden_states}
+    inputs = {'cu_seq_lens_q': torch.tensor([0, 3, 7, 8], dtype=torch.int32)}
+
+    result = select_last_packed_states(outputs, inputs)['logits']
+
+    torch.testing.assert_close(result[:, 0], hidden_states[0, [2, 6, 7]])
+    assert result.shape == (3, 1, 2)
+    result.sum().backward()
+    expected_grad = torch.zeros_like(hidden_states)
+    expected_grad[0, [2, 6, 7]] = 1
+    torch.testing.assert_close(hidden_states.grad, expected_grad)
+
+
+def test_select_last_packed_states_with_position_ids():
+    hidden_states = torch.arange(12, dtype=torch.float32).view(1, 6, 2)
+    outputs = {'last_hidden_state': hidden_states}
+    inputs = {'position_ids': torch.tensor([[0, 1, 2, 0, 1, 0]])}
+
+    result = select_last_packed_states(outputs, inputs)['last_hidden_state']
+
+    torch.testing.assert_close(result[:, 0], hidden_states[0, [2, 4, 5]])
+    assert result.shape == (3, 1, 2)