Skip to content

[Train] Reduce padding-free embedding output memory - #9893

Open
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:agent/fix-qwen3-vl-embedding-memory
Open

[Train] Reduce padding-free embedding output memory#9893
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:agent/fix-qwen3-vl-embedding-memory

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

What changed

  • Select the last hidden state of every sequence directly from the packed embedding output.
  • Keep the existing re-padding behavior for sequence classification and reranking tasks.
  • Support both cu_seq_lens_q and packed position_ids, including the sequence-parallel path.
  • Add regression tests for output selection and gradient propagation.

Why

With padding_free, the embedding path restored the full output to a dense
[num_sequences, max_sequence_length, hidden_size] tensor before the output
normalizer immediately selected only the last token of each sequence. For
multimodal batches with uneven sequence lengths, this creates a large temporary
allocation and retains its autograd graph until backward.

The embedding loss only needs one vector per sequence, so selecting those packed
positions directly avoids the re-padding allocation without changing the
resulting embeddings.

Fixes #9885.

Validation

  • pytest -q tests/models/test_patcher.py
  • pre-commit run --files swift/model/patcher.py swift/trainers/mixin.py tests/models/test_patcher.py
  • Tiny Qwen3-VL packed forward, normalized output, and backward smoke test

Signed-off-by: taking-lying-flat <1615405@qq.com>
@taking-lying-flat
taking-lying-flat marked this pull request as ready for review August 11, 2026 12:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Qwen3-vl-embedding-2B训练爆显存

1 participant