v0.10.0: análise/prova/teste das críticas de engenharia Qwen3.8 lote 3 (r-35) — divisão SRP dos monólitos jepa/pensamento/moe/python em 26 submódulos (API estável via re-exports; 'Amdahl de compilação' refutada por experimento controlado, builds frios neutros), checkpoints atômicos SHA-256 puro FIPS 180-4 + fsync + rename + manifesto (CP-1..CP-3; correções da proposta: sem destruir estado anterior, manifesto com barreira, sem bincode/unwrap), TemperaturaInversa T-V3 (fronteira tipada da convexidade N-2; alegação compile-time corrigida), trait CamadaKhtst+CamadaHomeostatica (premissa do controlador refutada no fonte; Vec<Var>; adaptador MoeComPapel bit-a-bit), CLI Typer 75 (fachada aditiva com resume verificado por hash) — 123/123 cargo, 107/107 gates, treino real Wikipédia PT-BR com checkpoints íntegros
06e4b6f verified | //! `jepa::sonic_moe` — extraído do monólito na v0.10.0 | |
| //! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2). | |
| use super::*; | |
| // --------------------------------------------------------------------------- | |
| // Parte 5 — SonicMoE (IO/tile-aware) | |
| // --------------------------------------------------------------------------- | |
| /// MoE top-k com capacidade por especialista e perda de balanceamento. | |
| /// | |
| /// **Prova SM-1 (limite de capacidade):** o roteamento tile-aware mantém no | |
| /// máximo `C` tokens por especialista por ladrilho — o drop é determinístico | |
| /// (ordem de chegada) e o nº de tokens descartados é `Σ_e max(0, load_e − C)`. | |
| /// | |
| /// **Prova SM-2 (balanceamento):** `L_aux = E·Σ_e f_e·p_e`; com roteamento | |
| /// auto-consistente (`f = p`), `L_aux = E·Σ p_e² ≥ 1` com igualdade sse | |
| /// `p_e = 1/E` (Cauchy–Schchwarz) — o mínimo força o roteador uniforme. | |
| pub struct SonicMoe { | |
| pub roteador: Linear, // d → E | |
| pub especialistas: Vec<MlpJepa>, // E experts d → d | |
| pub n_top: usize, | |
| pub capacidade: usize, | |
| pub lambda_aux: f64, | |
| } | |
| pub struct SaidaMoe { | |
| pub saida: Tensor, | |
| pub perda_aux: Tensor, | |
| pub carga: Vec<usize>, // tokens efetivamente processados por especialista | |
| pub descartados: usize, // tokens com especialistas cortados pela capacidade | |
| } | |
| impl SonicMoe { | |
| pub fn novo(d: usize, n_expertos: usize, n_top: usize, capacidade: usize, dev: &Device) -> Result<Self> { | |
| let roteador = Linear::sem_vies(d, n_expertos, 0.1, dev)?; | |
| let especialistas = (0..n_expertos) | |
| .map(|_| MlpJepa::novo(d, 2 * d, d, dev)) | |
| .collect::<Result<Vec<_>>>()?; | |
| Ok(Self { roteador, especialistas, n_top, capacidade, lambda_aux: 0.01 }) | |
| } | |
| pub fn vars(&self) -> Vec<Var> { | |
| let mut v = self.roteador.vars(); | |
| for e in &self.especialistas { | |
| v.extend(e.vars()); | |
| } | |
| v | |
| } | |
| /// Forward tile-aware: (n, d) → (n, d). | |
| pub fn fwd(&self, x: &Tensor) -> Result<SaidaMoe> { | |
| let n = x.dim(0)?; | |
| let d = x.dim(1)?; | |
| let e = self.especialistas.len(); | |
| let dev = x.device(); | |
| let logits = self.roteador.fwd(x)?; // (n, E) | |
| let probs = candle_nn::ops::softmax(&logits, candle_core::D::Minus1)?; | |
| // top-k por token (não diferenciável — pesos do top-k são dif. via probs) | |
| let probs_v = probs.to_vec2::<f32>()?; | |
| let mut carga = vec![0usize; e]; | |
| let mut descartados = 0usize; | |
| // acumulador de saída (n, d) — construído por soma de contribuições | |
| let mut saida = Tensor::zeros((n, d), candle_core::DType::F32, &dev)?; | |
| let mut peso_total = Tensor::zeros((n, 1), candle_core::DType::F32, &dev)?; | |
| // por especialista (tile): tokens cujo top-k contém e e com carga < C | |
| for ei in 0..e { | |
| // candidatos: tokens ordenados por prob do especialista ei | |
| let mut cand: Vec<(usize, f32)> = (0..n) | |
| .map(|i| (i, probs_v[i][ei])) | |
| .collect(); | |
| cand.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); | |
| let mut escolhidos: Vec<usize> = Vec::new(); | |
| for (i, _) in cand { | |
| if carga[ei] >= self.capacidade { | |
| break; | |
| } | |
| // verifica se ei está no top-k do token i | |
| let mut ordem: Vec<(f32, usize)> = | |
| (0..e).map(|k| (probs_v[i][k], k)).collect(); | |
| ordem.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); | |
| let top: Vec<usize> = ordem.iter().rev().take(self.n_top).map(|(_, k)| *k).collect(); | |
| if top.contains(&ei) { | |
| escolhidos.push(i); | |
| carga[ei] += 1; | |
| } | |
| } | |
| if escolhidos.is_empty() { | |
| continue; | |
| } | |
| // tile do especialista: uma única passagem em lote | |
| let idx: Vec<u32> = escolhidos.iter().map(|&i| i as u32).collect(); | |
| let n_e = idx.len(); | |
| let idx_t = Tensor::from_vec(idx, (n_e,), &dev)?; | |
| let tile = x.index_select(&idx_t, 0)?; | |
| let y = self.especialistas[ei].fwd(&tile)?; // (m, d) | |
| // pesos: prob do especialista ei para cada token escolhido | |
| let w: Vec<f32> = escolhidos.iter().map(|&i| probs_v[i][ei]).collect(); | |
| let n_w = escolhidos.len(); | |
| let w_t = Tensor::from_vec(w.clone(), (n_w, 1), &dev)?; | |
| let contrib = y.broadcast_mul(&w_t)?; | |
| // scatter-add: acumula nas posições | |
| let mut buf = saida.to_vec2::<f32>()?; | |
| let contrib_v = contrib.to_vec2::<f32>()?; | |
| for (k, &i) in escolhidos.iter().enumerate() { | |
| for j in 0..d { | |
| buf[i][j] += contrib_v[k][j]; | |
| } | |
| } | |
| saida = Tensor::from_vec( | |
| buf.into_iter().flatten().collect::<Vec<f32>>(), | |
| (n, d), | |
| &dev, | |
| )?; | |
| let mut pt = peso_total.to_vec2::<f32>()?; | |
| for (k, &i) in escolhidos.iter().enumerate() { | |
| pt[i][0] += w[k]; | |
| } | |
| peso_total = Tensor::from_vec( | |
| pt.into_iter().flatten().collect::<Vec<f32>>(), | |
| (n, 1), | |
| &dev, | |
| )?; | |
| } | |
| // tokens sem especialista (todos cortados) → descartados; saída 0 | |
| let pt_v = peso_total.to_vec2::<f32>()?; | |
| for i in 0..n { | |
| if pt_v[i][0] == 0.0 { | |
| descartados += 1; | |
| } | |
| } | |
| // renormaliza pelo peso efetivo (evita viés da capacidade) | |
| let denom = peso_total.clamp(1e-6, f64::INFINITY)?; | |
| let saida = saida.broadcast_div(&denom)?; | |
| // perda auxiliar: E·Σ f_e·p_e (f = fração de tokens cujo top-1 é e) | |
| let mut f = vec![0f32; e]; | |
| for i in 0..n { | |
| let (mut melhor, mut mp) = (0usize, f32::NEG_INFINITY); | |
| for k in 0..e { | |
| if probs_v[i][k] > mp { | |
| mp = probs_v[i][k]; | |
| melhor = k; | |
| } | |
| } | |
| f[melhor] += 1.0 / n as f32; | |
| } | |
| let f_t = Tensor::from_vec(f, (e,), &dev)?; | |
| let p_media = probs.mean(0)?; // (E,) | |
| let perda_aux = ((&f_t * &p_media)?.sum_all()? * (e as f64))?; | |
| Ok(SaidaMoe { saida, perda_aux, carga, descartados }) | |
| } | |
| } | |