modelscope / modelscope/ms-swift

`fix(gkd): defer full-vocabulary local teacher logits to forward_step`

Open
#10,097 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Checklist / 检查清单
  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues,确认这是一个新的 bug report。
Bug Description / Bug 描述

For a separate local teacher with exact full-vocabulary logits
(gkd_logits_topk is None), _compute_teacher_logits_local eagerly materializes
one [S, V] teacher logits tensor per rollout micro-batch at batch
preparation
, and all of them stay alive for the whole generation cycle. With
64 queries × n=4 (256 rollouts), up to 256 full [S, 151k] logits tensors are
retained at once, OOMing long teacher responses (observed in practice).

How to Reproduce / 如何复现
Change
  • Split the per-batch eager loop into _compute_teacher_output_local, a
    single-micro-batch helper that forwards one local teacher microbatch and
    returns its TeacherOutput.
  • _compute_teacher_logits skips the eager materialization when
    gkd_logits_topk is None (full-vocabulary mode): the encoded batches keep
    teacher_model_inputs instead of teacher_output.
  • forward_step computes the teacher output just in time for the micro-batch it
    is about to train on, so only one [S, V] logits tensor is alive at a
    time. Fails closed with RuntimeError if a batch has neither an eager
    teacher_output nor usable teacher_model_inputs.
  • Compressed top-k logprobs ([S, K]) remain eager — they are small; the
    deferral targets the exact full-vocabulary path only.
  • Self-distillation behavior is unchanged (still recomputed per train step via
    _on_train_step_batch).
Additional Information / 补充信息

No response

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Locate _compute_teacher_logits_local, _compute_teacher_output_local, _compute_teacher_logits, and forward_step, then trace how encoded batches carry teacher data. Verify that full-vocabulary local-teacher logits are computed only for the micro-batch being trained, while compressed top-k outputs remain eager and self-distillation is unchanged. Done means the full-vocabulary path retains only one logits tensor and fails clearly when neither teacher output nor inputs are available.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.