# Gemma-4-26B-A4B-HybridAgent **`Gemma-4-26B-A4B-HybridAgent`** is a hybrid reasoning agent built upon **`google/gemma-4-26B-A4B` (25.81B Total Parameters / ~4.0B Active Parameters per token MoE)**, adapting the [Stanford / Together AI `HybridAgent` (`arXiv:2610.18861`)](https://arxiv.org/abs/2610.18861) architecture to the Gemma 4 Mixture-of-Experts (MoE) family. Instead of verbalizing every intermediate reasoning token in text, **`Gemma-4-26B-A4B-HybridAgent` thinks in continuous latent space (`...`) via a lightweight residual `HybridLink` bridge (`adapter.pt` / `hybridlink.pt`), and acts through text (`...`, `...`, `...`)**, preserving multi-turn KV working memory and full white-box latent interpretability. --- ## ⚠️ Important Usage Notice (Read Before Loading!) > **Why you MUST NOT use standard `AutoModelForCausalLM.from_pretrained(...).generate()` directly:** > 1. **`adapter.pt` & `adapter_config.json` are `HybridLink`, NOT a PEFT LoRA adapter!** If `peft` is installed in your environment, calling `AutoModelForCausalLM.from_pretrained()` directly will mistakenly treat `adapter_config.json` (`{"adapter_type": "ln_res_adapter"}`) as a LoRA config and raise a `ValueError`. Always load via `load_hybrid_agent(...)` (which disables PEFT auto-detection and attaches `HybridLink`). > 2. **Standard `.generate()` does not execute latent recurrence:** When the model emits the `` entry anchor (`token_id=262144`), the generation loop must route last-layer hidden states $h_t$ through `HybridLink` ($\mathcal{R}(h_t)$), scale by Gemma 4's `embed_scale` ($\sqrt{2816} \approx 53.0660$), and feed the continuous vector back into the decoder until the nearest-token cosine readout hits `` (`token_id=262145`). Use `HybridEngine` or `HybridGenerationEngine` below. --- ## 🛠️ Step 1: Installation Install the Gemma 4–adapted `HybridAgent` package: ```bash # Method 1: Clone from GitHub (gemma4 branch) git clone -b gemma4 https://github.com/milkshake0721/HybridAgent.git cd HybridAgent pip install -r requirements.txt pip install -e . # Method 2: Or download the bundled hybrid_agent package directly from Hugging Face python3 -c "from huggingface_hub import snapshot_download; p = snapshot_download('milkshake721/Gemma-4-26B-A4B-HybridAgent', allow_patterns=['hybrid_agent/**', 'interpretable_latent_mode/**', 'hybrid_inference/**', 'pyproject.toml', 'requirements.txt']); print('Downloaded package to:', p)" ``` *(Optional)* If you want to enable the live web search tool (`...`), set your Tavily API key: ```bash export TAVILY_API_KEY="tvly-..." ``` --- ## ⚡ TensorRT & Zero-Sync MoE Acceleration (`backend="tensorrt"`) Passing `backend="tensorrt"` to `load_hybrid_agent` (or `--backend tensorrt` to `sglang_tensorrt_runner.py`): 1. Compiles `HybridLink` (`adapter.pt` / `hybridlink.pt`) into a native **NVIDIA TensorRT `ICudaEngine`** via `torch_tensorrt.compile`. 2. Replaces standard loop-based `Gemma4TextExperts.forward` with a **Zero-Sync Batched Fast Path (`patch_bf16_experts_fastpath`)** across all 30 MoE layers (`4.60 ms / token`). 3. Replaces all 330 `Gemma4RMSNorm` calls per step with fused C++ `F.rms_norm`. *(Running on a 24GB GPU such as RTX 4090? Use the 4-bit edition [`milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit`](https://huggingface.co/milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit), which fits in `19.00 GB VRAM` and reaches `60.2 tok/s` via our Register-Fused Triton W4A16 MoE kernel.)* | Configuration | Allocated VRAM | Load Time | Prefill (180 tok) | Single-Token Decode | 10-Q Total Time | 10-Q Accuracy | | :--- | :---: | :---: | :---: | :---: | :---: | :---: | | **`Gemma-4-26B-A4B-HybridAgent` (BF16, `backend="tensorrt"`)** | **`54.73 GB`** | **`6.0s ~ 8.8s`** | **`313.03 ms`** | **`26.01 ms` (`38.5 tok/s`)** | **`106.80s` (`10.68s/q`)** | **`10/10 (100%)`** | | **[`Gemma-4-26B-A4B-HybridAgent-4bit`](https://huggingface.co/milkshake721/Gemma-4-26B-A4B-HybridAgent-4bit) (`backend="tensorrt"`)** | **`19.00 GB`** | `29.69s` | **`187.51 ms`** | **`16.62 ms` (`60.2 tok/s`)** | **`83.65s` (`8.36s/q`)** | **`10/10 (100%)`** | --- ## 🚀 Step 2: Quick Start ### CLI Runner (`sglang_tensorrt_runner.py`): ```bash python3 sglang_tensorrt_runner.py \ --backend tensorrt \ --model milkshake721/Gemma-4-26B-A4B-HybridAgent \ --question "A rectangle has a perimeter of 30 cm and an area of 54 cm^2. What are the lengths of its sides?" ``` ### Option A — High-Level Single-Question API (`HybridEngine`) ```python import torch from hybrid_agent.model import load_hybrid_agent from hybrid_agent.engine import HybridEngine from hybrid_agent.prompts import SYSTEM_PROMPT MODEL_ID = "milkshake721/Gemma-4-26B-A4B-HybridAgent" model, tokenizer = load_hybrid_agent( MODEL_ID, dtype=torch.bfloat16, device="cuda", backend="tensorrt", ) engine = HybridEngine( model, tokenizer, max_new_tokens=1024, max_latent_steps=64, temperature=0.0, top_p=1.0, ) traj = engine.run( question="If 3x + 7 = 22, what is the value of 5x - 2?", system_prompt=SYSTEM_PROMPT, ) print("Final Answer:", traj.final_answer) print("Total Latent Steps:", traj.latent_steps, "| Latent Spans:", traj.latent_spans) print("\n--- Rendered Trajectory (with Latent Pseudo-Text Readout) ---") print(traj.output_rendered) ``` --- ### Option B — Full Multi-Turn Lockstep Engine (`HybridGenerationEngine` with Python & Search Tools) ```python import torch from hybrid_agent import ( HybridGenerationEngine, SYSTEM_PROMPT, build_tool_executor, extract_answer, load_hybrid_agent, ) MODEL_ID = "milkshake721/Gemma-4-26B-A4B-HybridAgent" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, tokenizer = load_hybrid_agent(MODEL_ID, dtype=torch.bfloat16, device=device, backend="tensorrt") # Registers the sandboxed interpreter (and if TAVILY_API_KEY is set) tools = build_tool_executor(require_search=False) question = "How many trailing zeros does 100! have?" engine = HybridGenerationEngine( model, tokenizer, device, tools, SYSTEM_PROMPT, question, batch_size=1, max_response_tokens=1024, min_latent_steps=4, max_latent_steps=64, temperature=0.0, top_p=1.0, record_latent_readout=True, # Decodes vectors into readable pseudo-text ) engine.run(torch.Generator(device=device).manual_seed(0)) res = engine.row_result(engine.rows[0]) print("Extracted Answer:", extract_answer(res["output"])) print(f"Turns: {res['turns']} | Python calls: {res['python_calls']} | Latent steps: {res['latent_steps']}") print("\n--- Full Hybrid Trace ---") print(res["output_rendered"]) ``` --- ### Option C — Interactive White-Box Latent Inspector (`interpretable_latent_mode`) Every continuous `` thought vector produced by `HybridLink` can be projected onto its nearest vocabulary token embedding via cosine similarity ($\arg\max_{v \in \mathcal{V}} \cos(\mathcal{R}(h_t), \text{Emb}(v))$), rendering an interactive HTML report that shades each latent step by readability: ```bash # Generate an interactive self-contained HTML report for a question python -m interpretable_latent_mode \ --model milkshake721/Gemma-4-26B-A4B-HybridAgent \ --question "At the speed of sound in air at 20 degrees Celsius, how many seconds does sound take to travel 10 kilometers? Round to the nearest second." \ --readout_top_k 5 \ --output latent_report_26b.html # Or view step-by-step nearest tokens directly in the terminal python -m interpretable_latent_mode \ --model milkshake721/Gemma-4-26B-A4B-HybridAgent \ --question "If f(x) = x^2 - 3x + 4, what is f(f(2))?" \ --show_steps ``` --- ### Option D — Zero-Clone Standalone Loader (`transformers` + `huggingface_hub`) If you do not want to clone the repository, you can load the model and `HybridLink` (`adapter.pt`) directly in pure PyTorch with this self-contained snippet: ```python import torch import torch.nn as nn from unittest import mock from pathlib import Path from huggingface_hub import snapshot_download from transformers import AutoModelForCausalLM, AutoTokenizer import transformers.modeling_utils as modeling_utils class HybridLink(nn.Module): def __init__(self, hidden_size: int): super().__init__() self.proj1 = nn.Linear(hidden_size, hidden_size) self.act = nn.GELU() self.proj2 = nn.Linear(hidden_size, hidden_size) self.pre_ln = nn.LayerNorm(hidden_size) self.post_ln = nn.LayerNorm(hidden_size) def forward(self, h: torch.Tensor) -> torch.Tensor: return self.post_ln(h + self.proj2(self.act(self.proj1(self.pre_ln(h))))) repo_dir = Path(snapshot_download("milkshake721/Gemma-4-26B-A4B-HybridAgent")) tokenizer = AutoTokenizer.from_pretrained(str(repo_dir)) # Bypass PEFT auto-detection on adapter_config.json with mock.patch.object( modeling_utils, "maybe_load_adapters", lambda path, download_kwargs, **kw: (None, path, kw), ): backbone = AutoModelForCausalLM.from_pretrained( str(repo_dir), dtype=torch.bfloat16, device_map="cuda" ) hidden_size = backbone.get_input_embeddings().weight.shape[-1] # 2816 hybridlink = HybridLink(hidden_size).to(device="cuda", dtype=torch.bfloat16) hybridlink.load_state_dict(torch.load(repo_dir / "adapter.pt", map_location="cuda", weights_only=True)) backbone.eval() hybridlink.eval() print("Loaded Gemma-4-26B-A4B backbone + HybridLink successfully!") ``` --- ## 🔬 Architecture & MoE Parameter Breakdown `google/gemma-4-26B-A4B` uses a **Parallel Shared Dense MLP + Routed MoE** decoder layer (`hidden_states = hidden_states_shared + hidden_states_routed`): | Component | Parameters | Pure BF16 Size | Share | Training Role | | :--- | ---: | ---: | ---: | :--- | | **Total Model (`Gemma4ForConditionalGeneration`)** | `25,805,933,872` | `51.61 GB` | 100% | Pure `bfloat16` weights | | ↳ Vision Encoder (`vision_tower`) | `572,792,112` | `1.14 GB` | 2.2% | Frozen (`requires_grad=False`) | | **Language Model (`language_model`, 30 layers)** | `25,233,141,760` | `50.47 GB` | 97.8% | `d_model = 2816`, vocab `262,148` | | ↳ **Active Shared Trunk (Non-Expert)** | **`2,395,156,480`** | **`4.79 GB`** | **9.3%** | Trained in Stage 1 & 3: All 30 layers of Shared Attention (`self_attn`), Shared Dense MLP (`2112`), **MoE Routers (`router`, `10.9M`)**, RMSNorms, and Tied Embeddings (`738M`) | | ↳ **Routed MoE Experts (128 experts/layer, Top-8)** | **`22,837,985,280`** | **`45.68 GB`** | **88.5%** | • **Mode A (`--freeze_experts`)**: Preserves pre-trained expert weights while flowing input gradients $\frac{\partial L}{\partial x}$ through all 30 layers
• **Mode B (`--offload_expert_grads`)**: Streams expert gradients to CPU RAM | | **`HybridLink` Adapter (`adapter.pt`)** | **`15,865,344`** | **`0.03 GB`** | **0.1%** | Trained in Stage 2 (`lr=2e-4`) & Stage 3 (`lr=5e-5`): 2-layer GELU Residual Adapter + LayerNorm (`2816 -> 2816`) | --- ## 🇹🇼 繁體中文使用說明 (Traditional Chinese Guide) **`Gemma-4-26B-A4B-HybridAgent`** 是將 Stanford 與 Together AI 提出的 [`HybridAgent` (`arXiv:2610.18861`)](https://arxiv.org/abs/2610.18861) 混合推理框架遷移至 **`google/gemma-4-26B-A4B`（26B 總參數 / 每 Token 僅啟用約 4B 活躍參數之 MoE 架構）** 的模型。 ### 為什麼不能直接用 `AutoModelForCausalLM` + `model.generate()`？ 1. **`adapter.pt` 與 `adapter_config.json` 是外掛在模型頭尾之間的 `HybridLink` 空間橋接器，不是 LoRA！** 若直接呼叫 `AutoModelForCausalLM.from_pretrained()`，環境中的 `peft` 套件會誤將 `adapter_config.json` 當成 LoRA 設定檔而報錯。 2. **隱空間思考需要專用狀態機引擎**：當模型在文字模式吐出 `` (`262144`) 時，推論引擎必須接手將最後一層隱藏態 $h_t$ 送入 `HybridLink` 轉成連續嵌入向量（並乘上 `Gemma 4` 特有的縮放係數 `embed_scale` = $\sqrt{2816} \approx 53.0660$），在隱空間高速遞迴推演直到最近鄰餘弦相似度對齊 `` (`262145`) 後再切回文字模式。 ### 三階段訓練流程 (3-Stage HybridSFT) 1. **Stage 1 (`format_activation`)**：訓練主幹與特殊錨點（``、``、``、``），學會何時切換進入隱空間與呼叫 `` / `` 工具（`lr = 7e-6`）。 2. **Stage 2 (`hybridlink_activation`)**：凍結主幹，單獨以 `lr = 2e-4` 訓練 `HybridLink`（`15.87M` 參數），透過餘弦相似度與 MSE 損失將隱空間隱藏態對齊至下一步的輸入詞嵌入空間。 3. **Stage 3 (`self_conditioned`)**：**Two-Pass 雙趟自我條件化聯合訓練**——第一趟（Pass 1）由 `HybridLink` 產生模型自己的隱空間預測向量，第二趟（Pass 2）將序列中 `` 區間的真實詞嵌入替換為 Pass 1 的預測向量，並同時開放主幹（`lr = 7e-6`）與 `HybridLink`（`lr = 5e-5`）進行跨趟端到端反向傳播，消除推論時的曝露偏差（Exposure Bias）。 --- ## 📚 Citation ```bibtex @misc{zou2026hybridagent, title={Hybrid Reasoning Agents}, author={Jiaru Zou and Yunzhe Qi and Sirui Chen and Federico Bianchi and Shang Zhu and Pan Lu and Jingrui He and James Zou}, year={2026}, eprint={2610.18861}, archivePrefix={arXiv}, primaryClass={cs.CL}, url={https://arxiv.org/abs/2610.18861}, } ```

Full Model Card (README)