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) 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 (<latent>...</latent>) via a lightweight residual HybridLink bridge (adapter.pt / hybridlink.pt), and acts through text (<think>...</think>, <python>...</python>, <search>...</search>), 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 <latent> 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 </latent> (token_id=262145). Use HybridEngine or HybridGenerationEngine below.

---

🛠️ Step 1: Installation

Install the Gemma 4–adapted HybridAgent package:

# 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 (<search>...</search>), set your Tavily API key:

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, 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 (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):

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)

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)

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 <python> interpreter (and <search> 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 <latent> 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 <latent> 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:

# 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:

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) 混合推理框架遷移至 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. 隱空間思考需要專用狀態機引擎：當模型在文字模式吐出 <latent> (262144) 時，推論引擎必須接手將最後一層隱藏態 $h_t$ 送入 HybridLink 轉成連續嵌入向量（並乘上 Gemma 4 特有的縮放係數 embed_scale = $\sqrt{2816} \approx 53.0660$），在隱空間高速遞迴推演直到最近鄰餘弦相似度對齊 </latent> (262145) 後再切回文字模式。

三階段訓練流程 (3-Stage HybridSFT)

1. **Stage 1 (format_activation)**：訓練主幹與特殊錨點（<latent>、</latent>、<think>、</think>），學會何時切換進入隱空間與呼叫 <python> / <search> 工具（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）將序列中 <latent> 區間的真實詞嵌入替換為 Pass 1 的預測向量，並同時開放主幹（lr = 7e-6）與 HybridLink（lr = 5e-5）進行跨趟端到端反向傳播，消除推論時的曝露偏差（Exposure Bias）。

---

📚 Citation

@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}, 
}