Repository navigation
Conversation
The online teacher depended on the variant: DFlash, DFlash2 and LiLiCorr ran the CausalLM forward and gathered their teacher rows out of its full [B, seq, vocab] logits, DSpark ran the inner model and projected only the rows its loss reads, and Domino kept no teacher at all. Offline and streaming training already used the row projection for everyone. _online_base_outputs is no longer a variant hook, and every batch gets one teacher: _teacher_logits projects the base's final hidden state, the inner model's output, at just the rows a loss reads, as it does offline. - A text batch runs the inner model once and keeps its final hidden state next to the target hidden states. - A multimodal batch keeps the top-level forward, which inserts the image/video features, and a forward hook on the inner model reads its output during it. The hook exists for that forward only and is removed in a finally. It reads where a text batch reads, not at lm_head's input, which a CausalLM may rescale first (Inkling divides it by a multiplier), so the two batches cannot get different teachers. The inner model must run once over the whole sequence, or the step raises; logits_to_keep narrows only lm_head's input. - DSpark's and Domino's overrides go, so their image batches reach the multimodal forward too. Theirs ran the language model on the raw token ids and dropped the image features. - The multimodal forward now also gets image_position_ids and video_position_ids, without which Gemma 4's vision tower fails: DFlash's Gemma 4 image and video batches failed there before, and DSpark's and Domino's would have started to. It still gets no token_type_ids, from which Gemma 3 and PaliGemma build their bidirectional attention over the image, so on transformers 5.5 their image batches raise, now for DSpark and Domino too, and on 5.12 and later they attend causally over the image, as DFlash's did. Behavior change, for online DFlash, DFlash2 and LiLiCorr only: their teacher is lm_head on the final hidden state instead of the CausalLM's logits. That equals the old teacher up to matmul rounding, except on a base whose CausalLM rescales its logits around lm_head, which the teacher now leaves out: Gemma's final_logit_softcapping, for which the next commit adds a switch, and the scaling of Granite (logits_scaling), Cohere (logit_scale), Falcon-H1 (lm_head_multiplier) and Inkling (logits_mup_width_multiplier), which changes the teacher's temperature by that factor. Text batches no longer build full-sequence logits, an image batch's are dropped, and the training output's logits are None. Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py 533 passed on transformers 5.5.4. The online case of the full-sequence test fails on the base tree, where the CausalLM forward fed lm_head the whole sequence; a new test pins the projected online teacher to the CausalLM's logits. On a tiny Qwen3-VL, new tests check that an image batch's teacher equals the model's own logits at the rows read, that an image batch and a text batch read the same final hidden state before the image even when the CausalLM rescales lm_head's input, that another image changes DFlash's, DSpark's and Domino's base rows from the image on (on the base tree it left DSpark's and Domino's unchanged), that the hook never exists for a text batch and is gone after an image batch, also one whose forward fails, that logits_to_keep leaves the teacher whole, and that an inner model run twice or on fewer rows raises. On a tiny Gemma 4 VLM, a new test checks that an image and a video reach the base. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…ly teacher The teacher is now lm_head on the base's final hidden state for every batch, which leaves out what a base's CausalLM does to its logits around lm_head. Gemma-family models cap them, and some bases need that cap in the teacher, so it comes back as a switch, off by default so that no other run changes. The logit scaling of other bases stays out, and conversion says so. - dflash_teacher_softcap: true makes _teacher_logits apply cap * tanh(logits / cap), with cap the base config's final_logit_softcapping, on the full rows and on the token_ids path LiLiCorr reads, for text and image batches alike. It runs in the base model's own op order. Logits handed over in an offline batch are the model's own and are still only gathered. Converting raises when the switch is on and the base config sets no final_logit_softcapping. - The switch alone decides the cap. A multimodal base gets the same rule as a text-only one, and restore reads the switch from the saved config like any field, so a checkpoint saved before the switch existed restores with it off. - One table, next to the base-model path constants in modeling_fakebase.py, lists the base config fields by which a CausalLM scales its logits at lm_head: logits_scaling (Granite family, HyperCLOVA X, MiniCPM3), logit_scale (Cohere family), lm_head_multiplier (Falcon-H1), output_multiplier (MuseGlimmer) and logits_mup_width_multiplier (Inkling). When a loss term reads the teacher and the base config sets one of them to other than 1.0, conversion warns once, on rank 0, that the teacher is lm_head's raw projection and does not apply it. Nothing else changes. - Both are decided at convert time, so whether a loss reads the teacher is decided there too: _reads_teacher(config) is stored as _needs_base_logits, and LiLiCorr declares its penalty and calibration terms from the config instead of from attributes it sets later. - FakeBaseConfig carries final_logit_softcapping and whichever of those scaling fields the base config sets. Offline training on a fake base, which every offline VLM uses, can then use the switch and gets the warning. Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py 558 passed on transformers 5.5.4. New tests pin the capped teacher on both projection paths and on a tiny Qwen3-VL image batch, that handed-over logits are not capped again, the error for a switch on a base without a cap, that a capping base converts with the cap the switch sets whether multimodal or not, that restore takes the switch from the saved config and leaves it off for a config saved without it, the fake base carrying the cap and each scaling field, and the warning: once for a tiny Granite with logits_scaling=8, also offline on a fake base loaded from it, and not for Llama, a scaling of 1.0, a CE-only objective, or a rank other than 0. A drafter trained online on a tiny capping Gemma-4 VLM, by the base branch or by 0.47.0, loads through load_vlm_or_llm with the cap off: a resumed DSpark step matches the step of the tree that trained it bitwise, as does a DFlash CE step (base branch), while a DFlash KD step distills from the uncapped teacher. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…ss_terms
Which divergence a DFlash-family draft trains on was tied to its variant: DFlash, DFlash2 and
LiLiCorr picked KD or CE by dflash_self_logit_distillation (DFlash2 also tv or lambda by
dflash_lk_loss_type), DSpark always blended CE and TVD, and Domino always trained CE. The terms
already live in one registry that names no variant, so the objective becomes a config choice.
- DFlashConfig.dflash_loss_terms takes {term: weight} over the registry's names (ce, kd, tvd,
tv, lambda), summed in the order given; lambda keeps dflash_lk_ce_scale/_decay, whose
descriptions stop saying DFlash2 only. Unknown names, negative or non-finite weights and an
all-zero objective fail config validation, as does setting it together with a legacy
selector moved off its default. The legacy check compares against the field defaults
rather than model_fields_set, because main.py converts from model_dump() and modelopt_state
restores from it, and both mark every field as set.
- No DSpark or Domino recipe pins dflash_self_logit_distillation: false any more: not
dspark.yaml or domino.yaml, the Nemotron-3.5 DSpark warm-start recipe, or the Qwen3-8B
Domino launcher example. Neither variant reads it: each states its own default objective,
and only DFlash's default objective, which DFlash2 and LiLiCorr inherit, and DFlash2's
modify() read the selector. But next to dflash_loss_terms the pin had to be flipped back
first. The DFlash2 and LiLiCorr recipes read it and keep theirs; the DSpark and Domino
skill notes now say the two ignore it.
- Unset, each variant keeps its objective exactly: a variant now states its default as
_default_loss_terms(config), and _build_objective turns either source into the same
[(name, weight, term)] list it built before.
- Any term works for any variant, online and offline: online every variant now has the
teacher, and the offline entry check still follows the terms' needs_teacher declarations.
DSpark declares its confidence head there, whose target reads the TVD, now that its
objective may contain no TVD term.
- Every variant reports each term as loss_<term>, a term of weight 0 included. Domino built
its metrics from scratch, harmless while its objective was a fixed CE equal to loss_final;
it now keeps the objective's terms next to loss_final and loss_backbone.
- DIVERGENCES is keyed by the config's DFlashLossTerm, so mypy checks its names against the
config, and a test checks that every configurable name has a term.
Variant-owned terms keep their own knobs: the DFlash2 selector, the LiLiCorr lattice, the DSpark
confidence head, and Domino's lambda schedule, which now blends the configured objective on its
corrected and backbone logits.
Tested: tests/unit/torch/speculative and tests/unit/torch/export/test_hf_spec_rope_export.py
583 passed, and tests/unit/recipe 540 passed, on transformers 5.5.4; the new
tests take at most 0.45 s each on 4 CPUs. They cover the config errors, that the unset
objective of each variant is the one it built before and equals the same terms set explicitly
(lambda's bound parameters included), the terms surviving a model_dump() convert and a
modelopt_state restore in order, a composed objective training every variant online with a
gradient on every draft parameter and every term reported, which configs make offline
batches carry the base hidden state, and that every shipped DSpark and Domino recipe takes a
dflash_loss_terms override alone through load_recipe and main.py's model_dump() round trip.
Without one, dspark.yaml and domino.yaml convert to the objective they built before, and a
DSpark or Domino step is bitwise the same with the selector on or off.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
CHANGELOG: one New Features (Speculative Decoding) entry per field, two Backward Breaking Changes entries for what the one teacher changes for existing users, and two Bug Fixes entries. - Online DFlash, DFlash2 and LiLiCorr now distill text and multimodal batches alike from lm_head on the base's final hidden state, which leaves out the logit scaling of Granite, Cohere, Falcon-H1 and others; conversion warns about it. - On a base that caps its logits, Gemma 3n and Gemma 4 VLMs included, the teacher is capped only with dflash_teacher_softcap set. A run resumed from an HF-format checkpoint saved before this release restores that checkpoint's config, in which the switch is off; one resumed from an FSDP-sharded checkpoint converts from the recipe and takes its switch. - The bug fixes are what the one online forward brings. DSpark's and Domino's image and video batches now reach the base model's multimodal forward, where since 0.46 they ran its language model on the raw token ids. Like DFlash's, their Gemma 3 and PaliGemma image batches reach it without token_type_ids, so on transformers 5.5 it raises and on 5.12 and later it attends causally over the image. Every variant's Gemma 4 image and video batches, which failed in the vision tower, now pass it their patch positions. The DFlash recipe table and the agent skill's knob table gain both fields, and the skill's per-model table says when a Gemma-family base needs the switch, DFlash's default KD included. No shipped recipe or launcher YAML trains a Gemma-family base (the only Gemma launcher example benchmarks vLLM MTP), so none sets the switch. Tested: pre-commit on the three files. On a tiny Qwen3-VL (transformers 5.5.4, 5.12.1 and 5.18.0), a tiny Gemma 3 VLM (5.12.1 and 5.18.0) and a tiny Gemma 4 VLM (5.5.4, 5.12.1 and 5.18.0), two different images give online DSpark and Domino the same target hidden states on the base branch and different ones here, as they do for DFlash. The image batches of a tiny Gemma 3 and a tiny PaliGemma VLM, with the token_type_ids their processors return, raise for every variant on transformers 5.5.4; on 5.12.1 and 5.18.0 they train, and the final hidden state is bitwise the model's own without those token_type_ids. On the base branch DSpark and Domino trained on them at 5.5.4 without the image. Resumed through main.py, a checkpoint that the base branch saved from online DFlash KD on a tiny capping Gemma2 continues uncapped from its HF-format copy whatever the recipe says, and takes the recipe's switch from its FSDP-sharded copy (transformers 5.5.4, 5.12.1 and 5.18.0). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Comment |
This was referenced Oct 24, 2026
This branch has not been deployed
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.
What does this PR do?
Type of change: New feature and two bug fixes, with two backward-breaking changes to online training (listed below and in the CHANGELOG).
This PR is stacked on #2990 ([2/2], branch
haoguo/dflash-shared-forward, head 28f4f45), whose branch is its base; merge it after that one. It does not depend on the other PRs stacked on #2990. It conflicts with two of them: with #2991 (DSpark and Domino followdflash_loss_objective) inCHANGELOG.rst,hf_dflash.py,hf_dspark.py,hf_domino.py,dspark.yamlanddomino.yaml(both change how the objective is configured), and with thedraft_metricslogging PR (#2992) inCHANGELOG.rstonly. Whichever of each pair merges second has to be rebased on the other. It merges cleanly with the AR-decode helper (#2994). #2990 gave the DFlash family one training forward and a registry of loss terms that names no variant. This PR makes the objective a config choice and gives the family one teacher.Loss terms in config. New field
DFlashConfig.dflash_loss_terms: a{term: weight}dict over the registry (ce,kd,tvd,tv,lambda), summed in the order given. It works for anyprojector_type, online or offline.lambdakeeps its parametersdflash_lk_ce_scaleanddflash_lk_ce_decay, whose descriptions no longer say "DFlash2 only".kdorce, perdflash_self_logit_distillation.tv/lambdaperdflash_lk_loss_type.ce_alpha * ce + l1_alpha * tvd.ce._default_loss_terms(config).dflash_loss_termstogether with any of the four old selectors (dflash_self_logit_distillation,dflash_lk_loss_type,dflash_ce_loss_alpha,dflash_l1_loss_alpha) set off its default.dflash_self_logit_distillation: falseany more: notdspark.yaml,domino.yaml, the Nemotron-3.5 DSpark warm-start recipe, or the Qwen3-8B Domino launcher example. Neither variant reads that selector, so they train as before and now takedflash_loss_termswith no other override. The DFlash2 and LiLiCorr recipes read it and keep it.loss_<term>indraft_metrics, including a term of weight 0. Domino now does too, next toloss_finalandloss_backbone._reads_teacher(config): LiLiCorr's penalty and calibration terms, and now DSpark's confidence head, whose target is 1 − TVD/2. For example, DSpark with{ce: 1}and no confidence head now trains offline withoutbase_model_hidden_states.One teacher. Before, the online teacher depended on the variant:
[B, seq, vocab]logits._online_base_outputsis no longer a hook a variant overrides, and every batch gets the same teacher:_teacher_logitsprojects the base's final hidden state (the inner model's output) at only the rows the loss reads, which is what offline training already did.lm_head).finally. A text batch never gets one.logits_to_keepnarrows onlylm_head's input, so it leaves the teacher whole.logitsisNone.Opt-in softcap. New field
DFlashConfig.dflash_teacher_softcap, defaultfalse._teacher_logitsappliescap * tanh(logits / cap), withcapthe base config'sfinal_logit_softcapping. It does so on the full rows and on thetoken_idspath LiLiCorr reads, for text and image batches alike, in the base model's own order of operations.final_logit_softcapping. A multimodal base follows the same rule as a text-only one, and restore reads the switch from the saved config, so a checkpoint saved before this PR restores with the cap off.A warning for logit scaling. Other scaling a base's CausalLM applies at
lm_headstays out of the teacher, and conversion now says so.modeling_fakebase.py, lists the config fields:logits_scaling(Granite family, HyperCLOVA X, MiniCPM3),logit_scale(Cohere family),lm_head_multiplier(Falcon-H1),output_multiplier(MuseGlimmer) andlogits_mup_width_multiplier(Inkling).The base config sets logits_scaling=8.0, which the distillation teacher does not apply: the teacher is lm_head's raw projection of the base model's final hidden state.Nothing else changes.FakeBaseConfignow carriesfinal_logit_softcappingand whichever of these fields the base config sets. Offline training on a fake base, which every offline VLM uses, can then use the switch and gets the warning.Changes for existing users (both are in the CHANGELOG under Backward Breaking Changes):
lm_headon the base's final hidden state, which offline training and DSpark already used. It leaves out whatever the base's CausalLM does to its logits aroundlm_head:final_logit_softcapping(Gemma family): see item 2.logits_scaling), Cohere and Cohere2 (logit_scale), and Falcon-H1 (lm_head_multiplier). With transformers 5.18 this also covers MiniCPM3, HyperCLOVA X, MuseGlimmer and Inkling.logits_scaling=8gives a teacher exactly 8x the old one (loss 4.165921688 → 4.166797638). Cohere withlogit_scale=0.0625gives one 16x the old one (loss 4.165035248 → 4.165026188).lm_headoutput.dflash_teacher_softcap: true.ar_validate.pyandmain.py's resume from an HF-format checkpoint restore the checkpoint's saved config, not the recipe's. A checkpoint saved before this PR therefore restores with the switch off and continues uncapped, and the recipe cannot turn the switch on.main.py's resume from an FSDP-sharded checkpoint converts from the recipe, so there the recipe's switch keeps the cap.Bug fixes (both in the CHANGELOG under Bug Fixes):
image_position_idsandvideo_position_ids. Without them Gemma 4's vision tower fails: DFlash's Gemma 4 image and video batches failed there before this PR, and DSpark's and Domino's would have started to with the first fix.Other notes:
token_type_ids(see Known gaps), now for DSpark and Domino too.Usage
Testing
Rebased again on 2026-10-24, onto #2990 at 28f4f45, which now sits on main 0638a27. Of what main brought, #2902 (DSpark serving export), #2897 (GPT-OSS final norm in the fake base) and #2839 (launcher regression gate) come closest to this PR, but none of them changes a line this PR changes. The rebase applied without conflicts and leaves every changed line of this PR as it was; the six CHANGELOG entries stay under 0.49.0, next to main's. After the rebase these were run again and give the results below: the unit tests (new counts), the GPU, launcher and skill gate tests, the bitwise A/B, the logit-scaling A/B, the Gemma 2 softcap A/B, the tiny Qwen3-VL A/B (5.5.4 and 5.18.0), the image-reach and capped-teacher checks on tiny Gemma 3 and Gemma 4 VLMs, the restore of drafter checkpoints that the base branch and the 0.47.0 release saved, the FSDP2 probe and pre-commit. Apart from the test counts and the launcher tests' environment (see below), each gives the same result as before the rebase. The other results are from before the rebase.
test_hf_dflash.py,test_hf_dflash2.py,test_hf_dspark.py,test_hf_domino.py,test_hf_lilicorr.pyandtest_fakebase.pyundertests/unit/torch/speculative/plugins,tests/unit/torch/export/test_hf_spec_rope_export.pyandtests/unit/recipe/test_loader.py): 890 passed on transformers 5.5.4 and on 5.18.0, CI'stf_minandtf_latestpins.tests/unit/torch/speculativeplustest_hf_spec_rope_export.py(583) also pass on 5.5.4, 5.12.1 and 5.18.0, andtests/unit/recipe(540) on 5.5.4 and 5.12.1. These include main's GPT-OSS final-norm tests (Fix GPT-OSS final norm reconstruction in FakeBaseModel #2897), which load the fake base whose config this PR extends. The commit messages give each commit's counts on 5.5.4.tools/launcher/tests, which parse every launcher example (this PR edits the Qwen3-8B Domino one): 191 passed, as on the base branch. This run used an existing launcher environment (nemo-run 0.8.1) in which the launcher package itself is not installed, so the one remaining test,test_package_install.py::test_package_dir_resolves_to_checkout, fails on both branches; CI'suv runinstalls the package. Before this rebase, installed as CI's launcher job does, all 165 tests then present passed.plugins/modelopt/skills), installed as CI's skills job does (pytest only): 110 passed and 3 skipped, namelytest_benchmark_model.py, which needs torch, accelerate and transformers, and two tests that need jinja2. With those installed (transformers 5.12.1), all 141 pass.tests/gpu/torch/speculative/plugins: 20 passed, 1 skipped (it needs 2 GPUs).train_acc, everydraft_metricsvalue and the generated tokens.dflash_loss_terms, with the old selectors left at their defaults, against the base running the old selectors. All 16 are identical.loss_ce. In all 14 Domino runs it equalsloss_finalbitwise, and everything else is identical.lilicorr-ce-dpace-on,lilicorr-kd-decay-on):lilicorr_calibrationis 0.00758991158 vs 0.00758990459, and all 71 gradients differ by at most 6.2e-5 relative to each gradient's max.Gemma3ForConditionalGenerationdoes not apply the cap): it converts with the switch off, and DFlash KD is bitwise identical on both batches. With the switch on, both batches' teacher is capped.PaliGemmaForConditionalGenerationdoes not apply the cap either): DFlash KD is bitwise identical on the image batch. On the text batch the loss is equal to 9 digits, and the target hidden states and the teacher differ by rounding, as in item 1.lm_head's input; that the switch caps it; that the image reaches the base rows of DFlash, DSpark and Domino; and the hook's lifetime: never there for a text batch, gone after an image batch, also one whose forward fails. They also pin thatlogits_to_keepleaves the teacher whole and that an inner model run twice or on fewer rows raises. On a tiny Gemma 4 VLM, unit tests pin that an image and a video reach the base.lm_head's input, in each image-text-to-text class that a generic tiny build could run and whose inner model DFlash finds (28 classes on transformers 5.5.4, 32 on 5.12.1 and on 5.18.0), and in Qwen2-VL, Qwen2.5-VL, Qwen3.5, GLM-4V, Llama 4, Mllama, SmolVLM and Idefics3 built by hand. The one exception is Florence2, an encoder-decoder, whose forward never runs the module DFlash takes as its inner model.token_type_idstheir processors return (Known gaps):token_type_ids. Gemma 3's own forward with them differs from the first image token on (by up to 1.15 on 5.12.1 and 0.43 on 5.18.0).Gemma4ForConditionalGeneration(cap 0.5), saved withsave_pretrained, then loaded throughload_vlm_or_llmas export,ar_validate.pyandmain.py's HF-format resume do:main.pytrains online DFlash KD on a tiny capping Gemma2 (cap 0.5) and writes the same weights both as an HF-format checkpoint and as FSDP-sharded DCP shards. This branch'smain.pythen resumes each (transformers 5.5.4, 5.12.1 and 5.18.0):fully_shardunit, one rank on the zero-loss branch, online DFlash with CE and with KD. Both ranks finish OK with the same losses as the base: 3.47548 and 3.48123 for CE, 3.47471 and 3.47468 for KD.lm_head's input instead of the inner model's output;image_position_ids,video_position_idsorpixel_values_videos(three changes);domino.yamlpins the selector again (two changes);model_fields_set;dflash_loss_termsis ignored;lambdaloses its parameters;dflash.dflash_loss_terms={ce: 0.5, tvd: 0.5}and checks it survives main.py'smodel_dump()round trip. With the selector pinned again, it fails.dflash.yamlwithdflash.dflash_loss_terms={ce: 0.5, kd: 0.5}loads the same way, on transformers 5.5.4, 5.12.1 and 5.18.0.dspark.yamlanddomino.yamlof the base branch and of this branch differ only indflash_self_logit_distillation, and each pair converts to the same objective and the same need for the teacher.cetokd.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).Is this change backward compatible?: ❌ Two changes for existing users, both described above and listed in the CHANGELOG under Backward Breaking Changes:
lm_head, on text and image batches; conversion warns about it.Every other config builds the same objective. Every checkpoint loads, and every config converts: no base is refused over its cap, Gemma 3 and PaliGemma VLMs included.
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in
CONTRIBUTING.md: N/ADid you write any new necessary tests?: ✅
Did you update Changelog?: ✅ Under 0.49.0: one New Features entry per field and two Backward Breaking Changes entries (Speculative Decoding), and two Bug Fixes entries.
Did you get Claude approval on this PR?: ❌ Not run yet.
Additional Information
git diff --shortstat <base>...HEAD -- . ':!tests' ':!docs', of whichmodelopt/is +224/−97.lm_head's input.lm_head: Inkling (transformers 5.18) divides it bylogits_mup_width_multiplier.lm_head's input would give an Inkling image batch a different teacher from a text batch.logits_to_keep, which slices onlylm_head's input.model_fields_set.examples/speculative_decoding/main.pyconverts fromrecipe.dflash.model_dump(), andmodelopt_staterestores from a dump; both mark every field as set.dflash.yaml'sdflash_self_logit_distillation: true, is accepted next todflash_loss_terms, where it has no effect.token_type_ids, from which Gemma 3 and PaliGemma build their bidirectional attention over the image:token_type_ids, attends bidirectionally within each image. DFlash already trained Gemma 3 that way.token_type_idson would fix Gemma 3 but not PaliGemma: with no suffix, as in this pipeline, PaliGemma's mark the whole prompt as one bidirectional prefix, which would let the base see future tokens. This PR leaves it as is.PaliGemmaForConditionalGenerationdoes not apply the cap. With the switch on, its teacher is capped anyway.Literalthere becauseconfig.pycannot import the plugin.DIVERGENCESis typed with thatLiteral, and a test checks the two hold the same set of names.🤖 Generated with Claude Code