Repository navigation
Add NVIDIA Dory 1.0 (looped nGPT) model support - #58
Merged
Merged
Conversation
Ports the looped (input/recurrent/output) nGPT-style architecture from the upstream mlx-lm PR (ml-explore#1988, by pmichalas, MIT): input layers -> recurrent block run n_recurrent_loops times -> output layers, with per-loop or last-loop KV caching. Model code and tests carried over as-is (numerically verified there against an independent fp32 PyTorch implementation); extended this fork's existing AutoTokenizer fallback (already added for DSV4) to also catch StrictDataclassError/KeyError, since transformers doesn't recognize Dory's layer_types values yet and this fork's fallback path (stub PretrainedConfig) already covers the same class of problem. Verified: the 4 new dory tests pass; full suite shows no regressions (remaining failures are pre-existing, unrelated flakes on other architectures, confirmed against unmodified HEAD). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… and training paths Found by a codex review pass before merging: - Attention called mx.fast.scaled_dot_product_attention directly instead of the shared base.scaled_dot_product_attention dispatcher, so enabling --kv-bits broke generation (QuantizedKVCache's update_and_fetch returns packed tuples mx.fast's kernel can't accept). - last_loop cache rewind set cache.offset directly, which BatchKVCache (what BatchGenerator substitutes for batched decoding) doesn't use as its write cursor (it uses a separate _idx) -- each loop appended instead of overwriting. Moved the rewind to the model's loop driver, using each cache's own trim() so every cache type's bookkeeping stays consistent. Added test_dory_last_loop_batch_cache (confirmed it fails without the fix, passes with it). - _scales/_alpha/_ls were memoized on first forward and never recomputed, so full fine-tuning never saw the current parameter values after the first step. Guarded recomputation with self.training, matching the existing pattern in kimi_k3.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The self.training guard from the previous commit was still wrong: once training switches back to eval (e.g. a validation pass after an optimizer step), self.training is False and the cache is non-None, so it never recomputes and silently keeps pre-update sqk/s_gate/suv_*/ attn_alpha/mlp_alpha/logit_scale values. Verified with a probe (one AdamW step, then eval vs. a fresh model with identical parameters: max |delta logits| = 14.27). These are tiny vectors, so just stopped caching them and recompute on every forward instead. Also, from the same review: assert n_input_layers + n_recurrent_layers + n_output_layers == num_hidden_layers in ModelArgs.__post_init__ (an inconsistent checkpoint config silently shifted the loop boundaries instead of failing), and add a sanitize() round-trip test -- every other test builds the model with already-fused qkvg_proj/gate_up_proj weights, so sanitize() (the only path a real split checkpoint takes) had no coverage at all. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
3 tasks done
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.
Summary
StrictDataclassError/KeyError, sincetransformersdoesn't recognize Dory'slayer_typesvalues yet and our fallback (stubPretrainedConfig) already covers exactly this class of problem — kept one fallback path instead of adding the upstream PR's separate one.Test plan
test_dory,test_dory_sliding_window_cache,test_dory_last_loop_cache_types,test_dory_last_loop_past_window)tests/test_models.pysuite run; remaining failures confirmed pre-existing on unmodified HEAD (unrelated architectures/flaky tolerances), not introduced by this change