[alwaysyiyu]

What does this PR do?

When pad_bshd_to_minibatch_max=True (introduced in #6901), the student forward pads every
micro-batch to the mini-batch's global max seqlen via forced_max_seqlen, but
compute_forward_kl_topk still splits the teacher top-k tensors with preprocess_bshd_engine
without passing forced_max_seqlen. With heterogeneous sequence lengths the teacher is
padded to the per-micro-batch max instead, so the two shapes diverge and the assertion

teacher_topk_log_probs_cp_split.shape[:2] == teacher_topk_ids_cp_split.shape[:2] == student_logits.shape[:2]

fails.

This threads the micro-batch's forced_max_seqlen (already stored as non-tensor data by the
engine) through compute_topk_loss into the BSHD teacher split. It picks up the earlier closed
attempt #7012 (which recomputed the value as student_logits.shape[1] * cp_size) and instead uses
the authoritative forced_max_seqlen recorded by the engine, matching the student path exactly.

Related: #6901 (introduced the padding), #7012 (prior closed attempt).

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here:
  • https://github.com/verl-project/verl/pulls?q=is%3Apr+pad_bshd_to_minibatch_max
  • https://github.com/verl-project/verl/pulls?q=is%3Apr+forced_max_seqlen
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • [trainer, megatron, fsdp] fix: forward-KL topk distillation shape mismatch with pad_bshd_to_minibatch_max

Test

Validated on an Ascend NPU cluster running on-policy distillation (OPD) with Qwen3.6-35B-A3B MoE,
TP=2 CP=2, pad_bshd_to_minibatch_max=True, and heterogeneous prompt/response lengths. Before the
fix the step aborts with the shape-mismatch assertion above; after the fix the step completes and
student_logits.shape[:2] == teacher_split.shape[:2] holds. The THD path and
pad_bshd_to_minibatch_max=False are unaffected (no forced_max_seqlen is stored).

API and Usage Example

No user-facing API/config change. compute_forward_kl_topk gains a backward-compatible
forced_max_seqlen: int | None = None parameter (internal), and compute_topk_loss reads
forced_max_seqlen from the micro-batch's non-tensor data. Callers need no change:

compute_forward_kl_topk(
    student_logits,
    teacher_topk_log_probs,
    teacher_topk_ids,
    config,
    data_format,
    forced_max_seqlen=tu.get_non_tensor_data(data, key="forced_max_seqlen", default=None),
)

Design & Code Changes

• verl/trainer/distillation/losses.py: read forced_max_seqlen from the micro batch and forward it to the distillation loss fn.
• verl/trainer/distillation/megatron/losses.py: accept forced_max_seqlen and pass it to both preprocess_bshd_engine calls in the BSHD branch (THD branch untouched).
• verl/trainer/distillation/fsdp/losses.py: accept the parameter for signature parity (unused; FSDP2 does not use pad_bshd_to_minibatch_max).

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: reproducing the mismatch requires a CP>1 distributed run; validated manually (see Test above). Open to adding a targeted unit test if reviewers prefer.
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.