Add GraniteMoeHybrid architecture to mlx-lm
Session from traces.claude-code.mlx-lm-granitemoehybrid on Hugging Face by gabegoodhart · Apache-2.0 · condensed by Coders Talk
Implement a new GraniteMoeHybrid model module in mlx_lm/models, mirroring nemotron_h and granitemoe patterns, so mlx-lm can load and run GraniteMoeHybrid checkpoints with correct generation and caching.
- +0Prompt · +0Add GraniteMoeHybrid architecture support to mlx-lm as a new granitemoehybrid.py module, hybrid of mamba2 and granitemoe attention+MoE layers, with detailed pointers on config and scalar multipliers.
- +0mAgent did · +0mSet up dev env, studied nemotron_h/granitemoe/mamba implementations and HF transformers source, then wrote the initial granitemoehybrid.py module (~483 lines).
- +7mFail · +7mModel loaded but errors piled up: mismatched weight names, wrong MoE tensor shapes, dimension mismatches in Mamba conv cache, and shape-unpacking errors during generation.
- +7mAgent did · +7mIteratively rewrote weight naming, MoE structure using SwitchGLU/SwitchLinear to handle 3D expert weight tensors, and fixed conv/cache shape handling until the model loaded and generated text without crashing.
- +21mPrompt · +21mShow me the current inference results compared side-by-side with transformers using greedy decoding, and make sure they line up.
- +23mFail · +23mSide-by-side comparison showed MLX generating nonsense ('lylyly Usesfinite') with completely different logits statistics and zero token overlap with transformers.
- +27mIntervention · +27mDeveloper pointed out RoPE is usually disabled for GraniteMoeHybrid, controlled by position_embedding_type in config, which the agent had missed.Why: I remembered a key architectural detail the agent couldn't have guessed from the code alone: RoPE is optional here and off by default, so I flagged it directly instead of letting it keep guessing.
- +27mFail · +27mApplied the RoPE fix but predictions were still identical and wrong, showing the RoPE issue wasn't the root cause.
- +31mAgent did · +31mWrote side-by-side debugging scripts comparing embeddings, layer outputs, and full layer calls between MLX and transformers, eventually discovering that every layer (mamba or attention) also applies a second MoE/shared_mlp block with its own residual, which the MLX implementation had skipped.
- +40mAgent did · +40mFixed the layer forward pass to apply the second MoE+shared_mlp block after the mamba/attention block; predictions then matched transformers almost exactly and generations became coherent.
- +51mOutcome · +51mCache validation completed: MLX hybrid cache (Mamba states + KV cache) produced 0.00000000 max difference between cached and full inference; transformers' own HybridMambaAttentionDynamicCache setup remained broken during testing. Session ended with working GraniteMoeHybrid support.
Next time I'd front-load the transformers layer forward-pass source (not just the module files) so the agent sees the double-residual MoE/shared_mlp structure per layer up front instead of discovering it through hours of logit-diffing. My correction about RoPE being disabled via position_embedding_type didn't actually fix the bug, but it was still the right instinct to hand over architecture-specific knowledge the agent couldn't infer from code alone; the real unlock was pushing for a strict side-by-side comparison against transformers, which is what eventually surfaced the missing MoE block.