fix(target-lm-head/dflash): fail closed on verifier-norm load and honour last_hidden_states_prenorm in DFlash - #162
Merged
Conversation
…unavailable `TargetLMHead.__init__` kept only `config.text_config`, so multimodal targets could not be rebuilt by `AutoModelForCausalLM.from_config` and final-norm extraction raised. Both norm entry points swallowed that exception and left `self.norm = None`, so a run configured with `last_hidden_states_prenorm: true` silently trained the draft against unnormalised pre-norm hidden states — no error, just wrong targets. Retain the outer config for architecture introspection, and when introspection still fails, rebuild the norm from `hidden_size`/`rms_norm_eps`; every target seen so far ends in a plain RMSNorm. When even that is impossible, raise instead of continuing: `load_norm=True` is a correctness requirement, and `_init_norm_structure` silently returning without a norm additionally desynchronises the rank-0 parameter broadcast and deadlocks it. `Eagle3Trainer` now rejects a missing norm under `last_hidden_states_prenorm` both on rank 0 and at the post-broadcast status check, and logs the norm state it actually ended up with rather than `(with verifier norm)`, which only ever reflected the request. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
…is set `last_hidden_states_prenorm` defaults to True for the vLLM engine, but `DFlashTrainer` never passed `load_norm` through to `TargetLMHead` and never applied a norm in `_forward`. DFlash and DSpark runs against a vLLM target therefore trained on unnormalised pre-norm hidden states while the resolved config said otherwise. Mirror the Eagle3 path: load the norm on rank 0, build the matching structure on the other ranks so the parameter broadcast stays in sync, reconcile the status across ranks before that broadcast, and normalise `last_hidden_states` under `no_grad` in `_forward`. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 773f9fa435
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…asting Rank 0's fail-closed norm load raised inside `TargetLMHead.from_pretrained`, before the status broadcast that the other ranks reach after building their own norm structure. Those ranks then sat in that broadcast until the process-group timeout instead of failing with the intended clean error. Deduplicate `_init_target_lm_head` — already a near-verbatim copy in the Eagle3 and DFlash trainers, and the preceding commit added the norm handling to both — into `load_synced_target_lm_head`, so this failure handling exists in one place. Every rank now reports its local outcome through an `all_reduce(MIN)` before the parameter broadcast; the rank that failed re-raises its own exception so the type and traceback survive, and the others report which side failed. `MIN` rather than a broadcast from rank 0 so that a follower failing in `_init_norm_structure` or `AutoConfig.from_pretrained` aborts the job too, instead of stranding rank 0 in the same way. Drop the `has_norm` broadcast and its two recovery branches, which the new all-reduce subsumes. With initialization failing closed, `norm is not None` holds exactly when `load_norm` does, on every rank: rank 0 either loads the norm or raises, followers either build the structure or raise, and `load_norm` comes from the same resolved config everywhere. The broadcast could therefore only confirm what each rank already knew, and the mismatch-recovery branches were unreachable. The number of collectives before the parameter broadcast is unchanged. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cf2ec4f1d1
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.target_lm_head = load_synced_target_lm_head( | ||
| target_model_path, | ||
| load_norm=self._last_hs_prenorm, | ||
| lm_head_key=getattr(self.args, "lm_head_key", "lm_head.weight"), |
There was a problem hiding this comment.
For DFlash, last_hidden_states_prenorm defaults to true with the vLLM backend, but the DFlash model only consumes last_hidden_states when dflash_l1_loss_alpha > 0 (torchspec/models/dflash.py:398). This line now forces verifier-norm loading for ordinary CE-only DFlash startup, so a missing or unsupported norm_key can abort a run even though the norm is never used; gate load_norm on actually using last-hidden-state distillation rather than the prenorm flag alone.
Useful? React with 👍 / 👎.
…the target `last_hidden_states_prenorm` defaults to True for the vLLM engine, so the preceding commit made every DFlash run load the target's final norm. But DFlash trains on hard-label cross entropy against `input_ids` and only reads `last_hidden_states` for the L1 distribution term — and `dflash_l1_loss_alpha` defaults to 0. An ordinary CE-only run therefore aborted on an unsupported `norm_key` for a norm it would never apply. Gate the load on the objective actually consuming the target's hidden states. The predicate lives on the loss module, next to the branches it describes, so it cannot drift from them: `DFlashModel` reports the L1 weight, `DSparkModel` adds the confidence head, whose accept-rate target also comes from the verifier distribution. Both now guard their own branches with it, and the trainer reads it off the unwrapped model it already keeps in `self.dflash`. `_forward` drops `last_hidden_states` under the same predicate rather than forwarding them unnormalised. The models discard them anyway when it is false, so no loss changes; but should a future branch read them without widening the predicate, it meets the existing explicit `ValueError` instead of silently training against unnormalised pre-norm activations — the failure this series exists to close. It also skips a full-size host-to-device copy and an RMSNorm per step on CE-only runs that still carry the tensor. Log when a prenorm run skips the norm, since `load_synced_target_lm_head` would otherwise report `not requested` and leave no way to tell that from a config that never asked for it. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
torchspec-bot
force-pushed
the
export/target-norm-failfast
branch
from
August 9, 2026 11:23
3b6c73b to
8f270a9
Compare
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
2 participants
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
Two related fixes to how the target model's final norm is loaded and applied when
last_hidden_states_prenormis set.1.
TargetLMHeadfails closed instead of silently dropping the normTargetLMHead.__init__kept onlyconfig.text_config, so multimodal targets could not berebuilt by
AutoModelForCausalLM.from_configand final-norm extraction raised. Both_init_and_load_normand_init_norm_structureswallowed that exception and leftself.norm = None, so a run configured withlast_hidden_states_prenorm: truetrained thedraft against unnormalised pre-norm hidden states with nothing louder than a warning.
model_config) for architecture introspection, alongside the textsub-config the module allocates against.
hidden_size/rms_norm_eps— everytarget seen so far ends in a plain RMSNorm.
load_norm=Trueis a correctness requirement, and_init_norm_structurereturning without a norm additionally desynchronises the rank-0parameter broadcast and deadlocks it.
Eagle3Trainerrejects a missing norm underlast_hidden_states_prenormon rank 0 and atthe post-broadcast status check, and logs the norm state it actually ended up with rather
than
(with verifier norm), which only ever reflected the request.2. DFlash honours
last_hidden_states_prenormresolve_configdefaultslast_hidden_states_prenormto True forinference_engine_type: vllm, butDFlashTrainernever passedload_normthrough toTargetLMHeadand neverapplied a norm in
_forward— so DFlash and DSpark runs against a vLLM target trained onunnormalised hidden states while the resolved config said otherwise. This mirrors the Eagle3
path: load on rank 0, build the matching structure elsewhere, reconcile the status across
ranks before the parameter broadcast, and normalise under
no_gradin_forward.Tests