Skip to content

feat(speculative): DFlash-family loss terms in config; one teacher with an opt-in softcap - #2993

Draft
h-guo18 wants to merge 4 commits into
haoguo/dflash-shared-forwardfrom
haoguo/dflash-loss-terms-config
Draft

h-guo18 wants to merge 4 commits into
haoguo/dflash-shared-forwardfrom
haoguo/dflash-loss-terms-config

Conversation

@h-guo18

@h-guo18 h-guo18 commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

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 follow dflash_loss_objective) in CHANGELOG.rst, hf_dflash.py, hf_dspark.py, hf_domino.py, dspark.yaml and domino.yaml (both change how the objective is configured), and with the draft_metrics logging PR (#2992) in CHANGELOG.rst only. 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 any projector_type, online or offline. lambda keeps its parameters dflash_lk_ce_scale and dflash_lk_ce_decay, whose descriptions no longer say "DFlash2 only".

    • Config validation rejects unknown names (the message lists the valid ones), negative or non-finite weights, and an objective whose weights are all zero.
    • Left unset, each variant builds exactly the objective it built before:
      • DFlash and LiLiCorr: kd or ce, per dflash_self_logit_distillation.
      • DFlash2: the same, or tv/lambda per dflash_lk_loss_type.
      • DSpark: ce_alpha * ce + l1_alpha * tvd.
      • Domino: ce.
      • Each variant now states this default as _default_loss_terms(config).
    • One source for the objective: config validation rejects dflash_loss_terms together 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.
    • No DSpark or Domino recipe pins dflash_self_logit_distillation: false any more: not dspark.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 take dflash_loss_terms with no other override. The DFlash2 and LiLiCorr recipes read it and keep it.
    • Terms a variant adds itself keep their own settings: the DFlash2 selector, the LiLiCorr lattice and the DSpark confidence head. Domino's λ schedule now blends the configured objective on its corrected and backbone logits.
    • Every variant reports each term as loss_<term> in draft_metrics, including a term of weight 0. Domino now does too, next to loss_final and loss_backbone.
    • Whether offline batches must carry the base model's final hidden state now depends only on what reads the teacher. That is the objective's terms, plus terms a variant declares in _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 without base_model_hidden_states.
  • One teacher. Before, the online teacher depended on the variant:

    • DFlash, DFlash2 and LiLiCorr gathered their rows from the CausalLM's full [B, seq, vocab] logits.
    • DSpark projected rows of the inner model's final hidden state.
    • Domino had no teacher.

    _online_base_outputs is no longer a hook a variant overrides, and every batch gets the same teacher: _teacher_logits projects 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.

    • A text batch runs the inner model once and keeps the target hidden states plus the final hidden state.
    • A multimodal batch keeps the top-level forward, which inserts the image and video features. A forward hook on the inner model reads its output during that forward: the same tensor a text batch reads, so an image batch and a text batch get the same teacher (see Additional Information for why the hook is not on lm_head).
    • The hook exists only during that forward and is removed in a finally. A text batch never gets one.
    • The inner model must run once over the whole sequence during that forward; otherwise the step raises rather than hand the teacher the wrong rows. logits_to_keep narrows only lm_head's input, so it leaves the teacher whole.
    • Text batches no longer build full-sequence logits, an image batch's logits are dropped, and the training output's logits is None.
  • Opt-in softcap. New field DFlashConfig.dflash_teacher_softcap, default false.

    • When set, _teacher_logits applies cap * tanh(logits / cap), with cap the base config's final_logit_softcapping. It does so on the full rows and on the token_ids path LiLiCorr reads, for text and image batches alike, in the base model's own order of operations.
    • Logits handed over in an offline batch are the model's own and are not capped again. The draft's own logits are never capped.
    • The switch alone decides the cap. Conversion raises only when the switch is on and the base config sets no 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_head stays out of the teacher, and conversion now says so.

    • One table, next to the base-model path constants in 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) 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: 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.
    • FakeBaseConfig now carries final_logit_softcapping and 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):

  1. The online teacher of DFlash, DFlash2 and LiLiCorr changes, on text and image batches alike. It is now lm_head on 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 around lm_head:
    • final_logit_softcapping (Gemma family): see item 2.
    • Logit scaling: Granite and its MoE variants (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.
      • The teacher's temperature changes by that factor, and no setting brings the old one back. Conversion warns about it.
      • On tiny bases: Granite with logits_scaling=8 gives a teacher exactly 8x the old one (loss 4.165921688 → 4.166797638). Cohere with logit_scale=0.0625 gives one 16x the old one (loss 4.165035248 → 4.165026188).
      • The new teacher is on the same scale as the draft's own logits, which are also the raw lm_head output.
    • CausalLMs that upcast their logits to fp32 (Mamba, Mamba2, FalconMamba, NemotronH): KD's teacher softmax now runs in the model dtype. A tiny Mamba2 in fp32 is bitwise unchanged. In bf16, none of these bases trains online DFlash either before or after this PR: Mamba2 hands the bf16 draft fp32 hidden states, and online conversion does not find NemotronH's embeddings.
    • Every other base: the teacher equals the old one up to rounding. It is bitwise identical in every A/B case below, image batches of tiny VLMs included, except two:
      • LiLiCorr online in fp32.
      • A tiny PaliGemma's text batch, whose language model now runs on its own rather than inside PaliGemma's forward. Its target hidden states differ by up to 9.5e-7 and its teacher by up to 9.7e-8.
  2. On a base that caps its logits, the teacher is capped only with dflash_teacher_softcap: true.
    • That holds for the Gemma 3n and Gemma 4 VLMs too, whose forward applies the cap. Their teacher used to be the model's own, capped logits on every batch, and is now the uncapped projection unless the switch is set. With the switch on, a tiny Gemma 4 VLM's DFlash KD step matches the base branch's bitwise on a text batch and on an image batch (the base branch needs the patch positions passed on to run an image batch at all).
    • Resumed runs: export, ar_validate.py and main.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.
    • No base is refused at conversion over its cap. Gemma 3 and PaliGemma VLMs, whose forward does not apply the cap their text config may set, keep the teacher they had. On tiny ones whose text config caps at 0.5, a DFlash KD step on an image batch is bitwise identical to the base branch's; so is Gemma 3's text batch, and PaliGemma's differs as in item 1.

Bug fixes (both in the CHANGELOG under Bug Fixes):

  • Online DSpark and Domino now send image and video batches through the base's multimodal forward, as DFlash does. Since 0.46 their own base forward 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 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:

  • No shipped recipe or launcher YAML trains a Gemma-family base. The only Gemma launcher example benchmarks vLLM MTP. So none of them sets the switch.
  • Gemma 3 and PaliGemma image batches still get no token_type_ids (see Known gaps), now for DSpark and Domino too.

Usage

dflash:
  # Any projector_type. Leave dflash_self_logit_distillation, dflash_lk_loss_type,
  # dflash_ce_loss_alpha and dflash_l1_loss_alpha at their defaults alongside it.
  dflash_loss_terms: {ce: 0.1, tvd: 0.9}
  # Gemma-family base: distill from the capped distribution the model serves.
  dflash_teacher_softcap: true
mtsp.convert(model, [("dflash", {**cfg, "dflash_loss_terms": {"kd": 0.5, "ce": 0.5}})])

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.

  • Unit tests.
    • The test files this PR changes or exercises (test_hf_dflash.py, test_hf_dflash2.py, test_hf_dspark.py, test_hf_domino.py, test_hf_lilicorr.py and test_fakebase.py under tests/unit/torch/speculative/plugins, tests/unit/torch/export/test_hf_spec_rope_export.py and tests/unit/recipe/test_loader.py): 890 passed on transformers 5.5.4 and on 5.18.0, CI's tf_min and tf_latest pins.
    • Each of the four commits passes them on 5.5.4 on its own (839, 864, 890 and 890).
    • tests/unit/torch/speculative plus test_hf_spec_rope_export.py (583) also pass on 5.5.4, 5.12.1 and 5.18.0, and tests/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.
    • Each test this PR adds or changes takes at most 0.7 s on 4 CPUs.
    • 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's uv run installs the package. Before this rebase, installed as CI's launcher job does, all 165 tests then present passed.
    • The skill gate tests (plugins/modelopt/skills), installed as CI's skills job does (pytest only): 110 passed and 3 skipped, namely test_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.
  • GPU tests. tests/gpu/torch/speculative/plugins: 20 passed, 1 skipped (it needs 2 GPUs).
  • Bitwise A/B against the base branch's tree (28f4f45).
    • Setup: each run is one training step, plus three generation steps for online cases. It compares the loss, every gradient, train_acc, every draft_metrics value and the generated tokens.
    • 48 cases × fp32/bf16 = 96 runs, of which 94 are identical.
    • The cases cover every variant offline and online:
      • DFlash with KD and CE;
      • DFlash2 with ce, kd, tv and lambda, plus its selector;
      • LiLiCorr with penalty and calibration;
      • DSpark with vanilla and rnn heads, plus the confidence head;
      • Domino.
      • All run with both D-PACE and decay, and include zero-loss steps and labels with -100.
    • 16 of the cases state each variant's old objective through dflash_loss_terms, with the old selectors left at their defaults, against the base running the old selectors. All 16 are identical.
    • Domino now also reports loss_ce. In all 14 Domino runs it equals loss_final bitwise, and everything else is identical.
    • The 2 runs that differ are LiLiCorr online in fp32 (lilicorr-ce-dpace-on, lilicorr-kd-decay-on):
      • Their penalty and calibration teacher is now projected per candidate instead of gathered from the full logits.
      • The loss is identical. lilicorr_calibration is 0.00758991158 vs 0.00758990459, and all 71 gradients differ by at most 6.2e-5 relative to each gradient's max.
      • In bf16 they are identical.
  • Bases that post-process their logits. One online DFlash KD step against the base branch, with the same weights, on transformers 5.5.4 and 5.18.0:
    • Granite and Cohere give the numbers in item 1 above, and conversion warns once on each.
    • Llama and Mamba2 (fp32) are bitwise identical, and conversion does not warn.
    • Mamba2 in bf16 and NemotronH fail the same way on both trees (item 1).
  • Softcap A/B on a tiny Gemma2 (cap 0.5), one online DFlash KD step:
    • Switch off: the teacher rows change by up to 0.064, and the loss goes from 3.47262263 to 3.47232127.
    • Switch on: the loss, every gradient and the teacher rows are bitwise identical to the base.
  • Multimodal. One online step on an image batch and one on a text batch, against the base branch with the same weights:
    • Tiny Qwen3-VL (transformers 5.5.4, 5.12.1 and 5.18.0): DFlash KD and DFlash CE are bitwise identical on both batches. The image batch's teacher, gathered from the model's logits before and projected from the inner model's output now, is bitwise the same.
    • Tiny Gemma 3 VLM whose text config caps at 0.5 (5.12.1 and 5.18.0; Gemma3ForConditionalGeneration does 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.
    • Tiny PaliGemma whose Gemma 2 text config caps at 0.5 (5.12.1 and 5.18.0; PaliGemmaForConditionalGeneration does 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.
    • Tiny Gemma 4 VLM whose forward caps at 0.5 (5.5.4, 5.12.1 and 5.18.0): with the switch on, DFlash KD is bitwise identical on both batches. With it off, the teacher is uncapped (5.18.0: image loss 4.16426754 → 4.16428375, text loss 4.16597319 → 4.16605473). DFlash CE is bitwise identical either way. The base branch cannot run a Gemma 4 image batch at all, so for this comparison its allow-list gets the two position fields.
    • On all four, DSpark and Domino are bitwise identical on the text batch with the switch off (Domino, which reads no teacher, with it on too). Two different images give them the same target hidden states on the base branch and different ones on this branch, while DFlash's differ on both.
    • Unit tests on a tiny Qwen3-VL pin that an image batch's teacher equals the model's own logits at the rows read; that it equals a text batch's before the image even when the CausalLM rescales 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 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, unit tests pin that an image and a video reach the base.
    • A text-only top-level forward runs the inner model exactly once, and its output is exactly 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.
  • Gemma 3 and PaliGemma image batches carrying the token_type_ids their processors return (Known gaps):
    • On transformers 5.5.4 the step raises for every variant. On the base branch DFlash raised too, while DSpark and Domino trained on them without the image.
    • On 5.12.1 and 5.18.0 every variant trains, and the final hidden state is bitwise the model's own forward without 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).
  • Checkpoints saved before this PR.
    • A drafter trained online on a tiny capping Gemma4ForConditionalGeneration (cap 0.5), saved with save_pretrained, then loaded through load_vlm_or_llm as export, ar_validate.py and main.py's HF-format resume do:
      • Trained by the base branch (transformers 5.12.1 and 5.18.0): it loads with the cap off. One resumed step matches the base branch's bitwise for DFlash (CE) and DSpark. DFlash (KD) now distills from the uncapped teacher: loss 4.167181969 → 4.167268276.
      • Trained by the 0.47.0 release (5.18.0): it loads, the resumed DSpark step matches the release's bitwise, and DFlash (KD) goes from 4.167181969 to the same 4.167268276.
    • The base branch's main.py trains 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's main.py then resumes each (transformers 5.5.4, 5.12.1 and 5.18.0):
      • FSDP-sharded, recipe switch on: cap 0.5, and the step is bitwise the base branch's (loss 4.856307030).
      • FSDP-sharded, switch off: uncapped (loss 4.856229305).
      • HF-format: uncapped whether or not the recipe sets the switch (loss 4.856229305).
  • FSDP2 probe. 2 gloo ranks, the draft as its own fully_shard unit, 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.
  • Mutation checks. Each of these 30 source changes fails at least one test in the suite:
    • an image batch reads lm_head's input instead of the inner model's output;
    • an image batch hands over the model's logits again;
    • the check that the inner model ran once, the check that it ran over the whole sequence, or both are dropped (three changes);
    • the hook is not removed when the forward raises, or never removed;
    • a text batch registers a hook on the inner model;
    • the multimodal forward drops image_position_ids, video_position_ids or pixel_values_videos (three changes);
    • the fake base drops the scaling fields, or the cap;
    • the multimodal refusal comes back;
    • the cap is always on for a capping base;
    • the cap is not applied;
    • handed-over logits are capped again;
    • the warning ignores, in turn, whether a term reads the teacher, the identity value 1.0, and the rank, or never fires (four changes);
    • the Nemotron-3.5 DSpark warm-start recipe or domino.yaml pins the selector again (two changes);
    • online text batches go back to the CausalLM forward;
    • online text batches keep no final hidden state;
    • the old-selector check reads model_fields_set;
    • DSpark's confidence head is not declared as reading the teacher;
    • dflash_loss_terms is ignored;
    • lambda loses its parameters;
    • Domino drops its per-term metrics.
  • Recipes.
    • A new unit test loads every shipped DSpark and Domino recipe with the single override dflash.dflash_loss_terms={ce: 0.5, tvd: 0.5} and checks it survives main.py's model_dump() round trip. With the selector pinned again, it fails.
    • dflash.yaml with dflash.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.
    • Without the override, the dspark.yaml and domino.yaml of the base branch and of this branch differ only in dflash_self_logit_distillation, and each pair converts to the same objective and the same need for the teacher.
    • A DSpark step (vanilla and rnn heads, with the confidence head) or a Domino step is bitwise the same with that selector true or false: the loss, every gradient, every metric and the generated tokens, offline and online, in fp32 and bf16 (24 runs, transformers 5.12.1). The same flip changes DFlash's objective from ce to kd.
  • pre-commit is clean on every file this PR changes.

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:

    • The online DFlash, DFlash2 and LiLiCorr teacher leaves out the logit scaling a base applies at lm_head, on text and image batches; conversion warns about it.
    • On a base that caps its logits, the Gemma 3n and Gemma 4 VLMs included, the teacher is capped only with the switch. A run resumed from an HF-format checkpoint saved before this PR continues uncapped; one resumed from an FSDP-sharded checkpoint takes the recipe's switch.

    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/A

  • Did 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

  • Source changes are +236/−106 lines by git diff --shortstat <base>...HEAD -- . ':!tests' ':!docs', of which modelopt/ is +224/−97.
  • Why the teacher reads the inner model's output, not lm_head's input.
    • A CausalLM may rescale the final hidden state before lm_head: Inkling (transformers 5.18) divides it by logits_mup_width_multiplier.
    • A text batch runs only the inner model, so its teacher projects the inner model's output. A hook on lm_head's input would give an Inkling image batch a different teacher from a text batch.
    • The hook on the inner model reads the very tensor a text batch reads, so the two cannot differ. It also cannot be narrowed by logits_to_keep, which slices only lm_head's input.
    • For every other VLM checked (see Testing) the two tensors are identical.
  • Why the old-selector check compares against defaults.
    • It compares each selector with its field default rather than reading model_fields_set.
    • examples/speculative_decoding/main.py converts from recipe.dflash.model_dump(), and modelopt_state restores from a dump; both mark every field as set.
    • As a result, a selector stated explicitly at its default value, such as dflash.yaml's dflash_self_logit_distillation: true, is accepted next to dflash_loss_terms, where it has no effect.
  • Known gaps.
    • The family's multimodal forward does not pass on token_type_ids, from which Gemma 3 and PaliGemma build their bidirectional attention over the image:
      • On transformers 5.5 their training forward requires them, so an online Gemma 3 or PaliGemma image batch raises. It did for DFlash before this PR; DSpark and Domino, which used to train on these batches without the image, now raise too.
      • On 5.12 and later these batches train with causal attention over the image, while Gemma 3's own forward, given its processor's token_type_ids, attends bidirectionally within each image. DFlash already trained Gemma 3 that way.
      • Passing token_type_ids on 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.
    • PaliGemma 2's Gemma2 text config caps at 30, but PaliGemmaForConditionalGeneration does not apply the cap. With the switch on, its teacher is capped anyway.
  • Why the term names live in config.py. They are a Literal there because config.py cannot import the plugin. DIVERGENCES is typed with that Literal, and a test checks the two hold the same set of names.

🤖 Generated with Claude Code

h-guo18 and others added 4 commits October 24, 2026 13:00
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>
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2026

Copy link
Copy Markdown

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.

@coderabbitai

coderabbitai Bot commented Oct 24, 2026

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant