Skip to main content

ferrox_models/decoder/
ffn_act.rs

1//! The one place `ModelConfig::ffn_activation` becomes the gated
2//! activation the expert FFN actually runs.
3//!
4//! It is a separate module for one reason: this mapping used to be
5//! written out only in `run_ffn_block`'s DENSE arm, and every routed
6//! path hardcoded SwiGLU. That is not a drift between copies, it is a
7//! decision that was made in one place and never made in the others, and
8//! the only thing keeping it from producing wrong logits was that
9//! `loader.rs` hands out [`FfnActivation::Gelu`] for
10//! `DecoderFamily::GemmaFamily` alone and every GemmaFamily row on
11//! `ArchPath::GenericGqa` is dense.
12//!
13//! With one conversion and a [`GluAct`] argument the routed paths cannot
14//! refuse to pass, there is nothing left to forget: adding a third
15//! activation is a non-exhaustive `match` here and a compile error at
16//! every call site, not a silent SwiGLU.
17
18use crate::config::FfnActivation;
19use ferrox_moe::GluAct;
20
21impl From<FfnActivation> for GluAct {
22    fn from(a: FfnActivation) -> Self {
23        match a {
24            // `SwigluFused` is the same activation as `Swiglu`; it only
25            // says gate and up arrive as one on-disk tensor (Phi), which
26            // the loader has already split by the time a `WeightMatrix`
27            // exists.
28            FfnActivation::Swiglu | FfnActivation::SwigluFused => GluAct::Swiglu,
29            FfnActivation::Gelu => GluAct::Geglu,
30        }
31    }
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    /// The mapping itself, pinned. If a later edit makes `Gelu` map to
39    /// `Swiglu` -- which is exactly what every routed path did before
40    /// this existed -- this goes red without needing a model.
41    #[test]
42    fn gelu_maps_to_geglu_and_never_to_swiglu() {
43        assert_eq!(GluAct::from(FfnActivation::Gelu), GluAct::Geglu);
44        assert_eq!(GluAct::from(FfnActivation::Swiglu), GluAct::Swiglu);
45        assert_eq!(GluAct::from(FfnActivation::SwigluFused), GluAct::Swiglu);
46    }
47}