Support use_logits_to_keep with sequence parallelism for long-context SFT

[JarvisPei]

Checklist / 检查清单

• I have searched existing issues, and this is a new feature request.

Feature Request Description / Feature Request 描述

Please support use_logits_to_keep together with sequence parallelism (SP) for long-context SFT.

These two optimizations reduce different sources of memory usage:

• sequence parallelism shards sequence-related transformer computation;
• logits_to_keep avoids materializing vocabulary logits for tokens excluded from the loss.

Current behavior

In Seq2SeqTrainer._prepare_inputs, ms-swift only enables
use_logits_to_keep when sequence_parallel_size == 1:

https://github.com/modelscope/ms-swift/blob/main/swift/trainers/seq2seq_trainer.py

use_logits_to_keep = self.get_use_logits_to_keep(
    self.template.sequence_parallel_size == 1
)

Therefore, enabling SP makes the effective value False by default.

If --use_logits_to_keep true is explicitly specified,
prepare_logits_to_keep() raises immediately:

https://github.com/modelscope/ms-swift/blob/main/swift/trainers/mixin.py#L1179-L1184

if self.template.sequence_parallel_size > 1:
    raise NotImplementedError()

Reproduction environment

• model: Qwen3.6-35B-A3B
• hardware: 8 × NVIDIA H20, 95.08 GiB per GPU
• ms-swift: 4.4.1
• transformers: 5.12.1
• PyTorch: 2.13.0+cu130
• CUDA: 13.0
• bf16, SDPA, DeepSpeed ZeRO-3, gradient checkpointing
• per-device batch size: 1

Representative command:

NPROC_PER_NODE=8 swift sft \
  --model /path/to/Qwen3.6-35B-A3B \
  --dataset /path/to/long-context-sft.jsonl \
  --tuner_type lora \
  --torch_dtype bfloat16 \
  --attn_impl sdpa \
  --max_length 131072 \
  --sequence_parallel_size 2 \
  --deepspeed zero3 \
  --gradient_checkpointing true \
  --per_device_train_batch_size 1

The training log reports:

[INFO:swift] use_logits_to_keep: False
...
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 22.00 GiB.
GPU 2 has a total capacity of 95.08 GiB of which 20.06 GiB is free.

The dense vocabulary logits can still dominate memory after the sequence
has been sharded. Adding --use_logits_to_keep true is currently not a
workaround because it reaches the NotImplementedError above.

Desired behavior

Allow use_logits_to_keep=true when sequence_parallel_size > 1, while
preserving correct causal-label alignment and loss normalization across
the SP group.

A possible implementation could:

1. align and shard labels using the existing SP path;
2. derive the supervised-logit selection after accounting for causal
   shifting and shard boundaries;
3. apply lm_head only to selected local hidden states;
4. normalize the loss correctly across the SP group;
5. preserve compatibility with loss_scale and padding-free/packed data.

Even an initial implementation limited to batch size 1 and non-packed
SFT would address an important long-context use case.

Related issues:

• DeepSpeed-Ulysses not work for Qwen-3.5 #8468 mentioned the incompatibility, but was closed after the separate
  Qwen3.5 SP/MRoPE fix in [bugfix] Qwen3.5 SP compat transformers 5.9.0 #9434.
• 训练时，logits.float()新申请的显存数过大导致oom #6966 and 4*A800 LoRA 微调 Qwen3.5-4B（图片文本识别，seq_len=32k）：训练几步后单卡显存占满 OOM，同配置 Qwen3-VL 无此问题 #9055 discuss logits-related memory pressure but do not
  directly track SP + use_logits_to_keep compatibility.

Pull Request / Pull Request 信息

No PR yet. I can help test a proposed implementation on long-context SFT.