[lincq2000]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

In RL training, every step waits for the slowest DP rank. Rollout lengths vary a lot, so the contiguous split slice_dp makes is often uneven. This PR adds an opt-in planner that reorders a mini-batch on the driver, before forward_backward, so that each rank gets a similar workload.

Feature

• plan_dp_balance(...) returns a DPBalancePlan; apply it with plan.apply() to every sample-aligned field (inputs, old_logps, advantages, ...).
• Each rank's cost is simulated with the model's own micro-batch planner, so it works for padded and packed inputs, with fixed or dynamic batching. The plan picks the cheapest of three orders (original, equal-size Karmarkar–Karp, snake by length) and reorders only if the slowest rank gets lighter.
• Per-rank sample counts and slice_dp stay unchanged, so the gradient is unchanged for per-sequence-mean losses such as GRPOLoss. Adds dp_balance/* metrics.
• Adds the cookbook/rl/grpo/dp_balance_grpo.{py,sh} example.

Bug fixes

• TypeError: not a sequence with Qwen3.5 / Qwen-VL rollouts. dataset.encode stores mm_token_type_ids as a nested [1, L] list, and swift's append_and_finalize extends it flat, so its length no longer matches input_ids. The twinkle wrapper now converts it to an array first.
• Packed samples attend to each other in forward_only. In eval mode use_cache is True and Transformers skips packed-sequence detection. forward_only now passes use_cache=False.

Experiment results

Environment: 4 × Ascend 910B3 (2 for training, 2 for vLLM), CANN 9.1.0, Python 3.12, torch 2.10.0 / torch_npu 2.10.0.post4, transformers 5.14.1, vLLM 0.26.0 / vllm-ascend 0.26.0rc2, flash-linear-attention 0.5.2.

Setup: Qwen3.5-4B GRPO on GSM8K, LoRA, FSDP2 with DP=2, dynamic batching, 50 steps. Each layout runs with balance_batch off and on, using the same seed.

layout	balance_batch	train ms / 1k tokens (median)	iteration time	mean reward	mean loss	mean grad norm
padded	off	1155	161.9 s	0.556	0.027	0.063
padded	on	913.5 (1.26x)	141.1 s (1.15x)	0.575	0.065	0.066
packed	off	1395	194.9 s	0.552	0.017	0.063
packed	on	923.4 (1.51x)	144.3 s (1.35x)	0.567	0.022	0.060

[meichangsu1]

batch 不能被 DP 数整除时，重排会改变 GRPO 梯度权重

[lincq2000]

不能被 DP 数整除时，已改为不重排，保持原顺序

[meichangsu1]

这里为什么要use_cache=False

[lincq2000]

padding-free 时多条样本拼成一行，没有 attention_mask，样本边界靠 position_ids 区分，transformers 只有在 past_key_values is None 时才按 position_ids 切分。
forward_only 在 eval 模式下，默认use_cache=True，模型内部会 past_key_values = DynamicCache（上面的条件不成立）导致整行按一条长序列的普通 causal mask 计算。forward_only 用不到 KV cache。