methods: sdft + sdpo — self-distillation trainers for SFT and RL - #13
Open
ghanshyam-lyzr wants to merge 1 commit into
Open
methods: sdft + sdpo — self-distillation trainers for SFT and RL#13ghanshyam-lyzr wants to merge 1 commit into
ghanshyam-lyzr wants to merge 1 commit into
Conversation
Two new methods on both local backends (mlx + torch), bringing the count to 15. Both make the model its own teacher: instead of cross-entropy on golden tokens or a scalar reward per rollout, the student's on-policy samples are pushed toward the same model reading richer context. - sdft (arXiv 2601.19897): on-policy learning from demonstrations. Each chat row's golden response is placed in-context for the teacher pass; the student's sampled completion is distilled toward the demonstration-conditioned distributions. Learns the task like SFT with far less catastrophic forgetting. Knobs: sdft_alpha (forward KL ↔ reverse KL ↔ JSD), completion length, temperature, teacher template. - sdpo (arXiv 2601.20802): RL via self-distillation on the grpo data surface. Rollout groups are scored by reward_fns — a fn may return (score, feedback) pairs — then each rollout is distilled toward the self-teacher reading a successful sibling solution and/or the textual feedback. Dense per-token advantages; all-fail groups still teach. The teacher trails the student as an EMA over adapter weights. Shared prompt-building and group-scoring helpers live in shadowlm/sdft.py and shadowlm/sdpo.py (pure stdlib) so both backends and the CPU tests use one implementation. Ships CPU tests, GPU matrix cells, mlx/torch/remote examples, and Train-page params (grouped under Preference · RL · distillation; no held-out eval, like grpo). The remote backend's reward_fns error now covers sdpo. pyproject: pin mlx-lm<0.30 and mlx-lm-lora==1.0.0 — newer releases require transformers 5.x / mlx-lm 0.30.6+, past our transformers<4.57 ceiling. Frontend static bundle rebuilt.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two new methods on both local backends (mlx + torch), bringing the count to 15. Both make the model its own teacher: instead of cross-entropy on golden tokens or a scalar reward per rollout, the student's on-policy samples are pushed toward the same model reading richer context.
sdft (arXiv 2601.19897): on-policy learning from demonstrations. Each chat row's golden response is placed in-context for the teacher pass; the student's sampled completion is distilled toward the demonstration-conditioned distributions. Learns the task like SFT with far less catastrophic forgetting. Knobs: sdft_alpha (forward KL ↔ reverse KL ↔ JSD), completion length, temperature, teacher template.
sdpo (arXiv 2601.20802): RL via self-distillation on the grpo data surface. Rollout groups are scored by reward_fns — a fn may return (score, feedback) pairs — then each rollout is distilled toward the self-teacher reading a successful sibling solution and/or the textual feedback. Dense per-token advantages; all-fail groups still teach. The teacher trails the student as an EMA over adapter weights.
Shared prompt-building and group-scoring helpers live in shadowlm/sdft.py and shadowlm/sdpo.py (pure stdlib) so both backends and the CPU tests use one implementation. Ships CPU tests, GPU matrix cells, mlx/torch/remote examples, and Train-page params (grouped under Preference · RL · distillation; no held-out eval, like grpo). The remote backend's reward_fns error now covers sdpo.
pyproject: pin mlx-lm<0.30 and mlx-lm-lora==1.0.0 — newer releases require transformers 5.x / mlx-lm 0.30.6+, past our transformers<4.57 ceiling. Frontend static bundle rebuilt.