modelscope / modelscope/ms-swift
`fix(gkd): defer full-vocabulary local teacher logits to forward_step`
Nobody has claimed this yet.
- 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 itsTeacherOutput. _compute_teacher_logitsskips the eager materialization when
gkd_logits_topk is None(full-vocabulary mode): the encoded batches keep
teacher_model_inputsinstead ofteacher_output.forward_stepcomputes 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 withRuntimeErrorif a batch has neither an eager
teacher_outputnor usableteacher_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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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