# coding=utf-8 """Swarm-MoE — a real, trainable, HuggingFace-native sparse Mixture-of-Experts LLM. Design goals (kept faithful to the project identity, but made *actually trainable*): * Swarm Mixture-of-Experts: a pool of experts per layer, only top-k routed per token (the "swarm" sparsity) — vectorized dispatch, fully differentiable. * Shared experts: one (or more) always-on FFN per layer captures common knowledge (DeepSeek/Qwen-MoE style "shared expert isolation") -> better quality + better parameter sharing. * Grouped-Query Attention (GQA) + Rotary embeddings (RoPE) + RMSNorm + SwiGLU. * Optional BitNet b1.58 ternary weights (Straight-Through Estimator) via ``config.quantization="bitnet"``. * Optional cross-layer weight sharing (ALBERT-style) via ``config.num_unique_layers``. * QOL speed: torch SDPA (memory-efficient/flash attention), KV cache for generation, gradient checkpointing, weight tying. Self-contained: only depends on ``torch`` and ``transformers`` base classes, so this file can be shipped inside an ``hf_upload/`` folder and loaded with ``trust_remote_code=True``. """ from __future__ import annotations from pathlib import Path from typing import Optional, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers.modeling_utils import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import MoeModelOutputWithPast, MoeCausalLMOutputWithPast from transformers.cache_utils import Cache, DynamicCache from .configuration_swarm_moe import SwarmMoEConfig # ----------------------------------------------------------------------------- # # BitNet b1.58 (optional ternary weights, trained with a Straight-Through Est.) # # ----------------------------------------------------------------------------- # def _activation_quant(x: torch.Tensor) -> torch.Tensor: """Per-token absmax INT8 quantization of activations (BitNet b1.58).""" scale = 127.0 / x.abs().amax(dim=-1, keepdim=True).clamp_(min=1e-5) return (x * scale).round().clamp_(-128, 127) / scale def _weight_quant(w: torch.Tensor) -> torch.Tensor: """Per-tensor absmean ternary {-1,0,1} quantization of weights (BitNet b1.58).""" scale = 1.0 / w.abs().mean().clamp_(min=1e-5) return (w * scale).round().clamp_(-1, 1) / scale class BitLinear(nn.Linear): """Linear layer with BitNet b1.58 ternary weights + INT8 activations. Uses the Straight-Through Estimator: the quantized values are used in the forward pass while gradients flow to the latent full-precision weights. This is the real BitNet training recipe (not a simulation) and converges with standard optimizers. """ def forward(self, x: torch.Tensor) -> torch.Tensor: w = self.weight # STE: forward uses quantized values, backward sees identity. x_q = x + (_activation_quant(x) - x).detach() w_q = w + (_weight_quant(w) - w).detach() return F.linear(x_q, w_q, self.bias) def make_linear(in_f: int, out_f: int, bias: bool, config: SwarmMoEConfig) -> nn.Linear: if config.quantization == "bitnet": return BitLinear(in_f, out_f, bias=bias) return nn.Linear(in_f, out_f, bias=bias) class ZeroLinear(nn.Linear): """A Linear that is initialized to all-zeros (see ``_init_weights``). Used for residual "depth" add-ons (deeper router) so that at initialization the new module is the identity — letting us add capacity while *retaining* a model's existing behavior and weights exactly, then learn the correction during finetuning. """ # ----------------------------------------------------------------------------- # # Norm / RoPE # # ----------------------------------------------------------------------------- # class SwarmRMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (self.weight * x.to(dtype)) class SwarmRotaryEmbedding(nn.Module): def __init__(self, dim: int, base: float = 10000.0): super().__init__() # NOTE: computed lazily (NOT a registered buffer) so the module never holds a # meta tensor — keeps from_pretrained(...).to(device) working under lazy loading. self.dim = dim self.base = base self._inv_freq = None def _inv(self, device): if self._inv_freq is None or self._inv_freq.device != device: self._inv_freq = 1.0 / (self.base ** ( torch.arange(0, self.dim, 2, dtype=torch.int64, device=device).float() / self.dim)) return self._inv_freq @torch.no_grad() def forward(self, x: torch.Tensor, position_ids: torch.Tensor): # position_ids: [batch, seq] inv_freq = self._inv(x.device)[None, :, None].float().expand(position_ids.shape[0], -1, 1) pos = position_ids[:, None, :].float() freqs = (inv_freq @ pos).transpose(1, 2) # [batch, seq, dim/2] emb = torch.cat((freqs, freqs), dim=-1) return emb.cos().to(x.dtype), emb.sin().to(x.dtype) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rotary(q, k, cos, sin): cos = cos.unsqueeze(1) # [b, 1, seq, dim] sin = sin.unsqueeze(1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: if n_rep == 1: return x b, kvh, s, d = x.shape return x[:, :, None, :, :].expand(b, kvh, n_rep, s, d).reshape(b, kvh * n_rep, s, d) # ----------------------------------------------------------------------------- # # Attention (GQA + RoPE + SDPA + KV cache) # # ----------------------------------------------------------------------------- # class SwarmAttention(nn.Module): def __init__(self, config: SwarmMoEConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.n_rep = self.num_heads // self.num_kv_heads self.attention_dropout = config.attention_dropout self.q_proj = make_linear(config.hidden_size, self.num_heads * self.head_dim, config.attention_bias, config) self.k_proj = make_linear(config.hidden_size, self.num_kv_heads * self.head_dim, config.attention_bias, config) self.v_proj = make_linear(config.hidden_size, self.num_kv_heads * self.head_dim, config.attention_bias, config) self.o_proj = make_linear(self.num_heads * self.head_dim, config.hidden_size, False, config) # Aligned Q/K/V RMSNorm (per head). Q/K norms are active (GPT-OSS/Qwen3 style); # the V norm is gated by `v_gate` (zero-init) so it is the identity at start and # blends in during finetuning -> a fuller, better-conditioned QKV-norm setup that # preserves a loaded model's behavior exactly. self.q_norm = SwarmRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = SwarmRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.use_v_norm = getattr(config, "use_v_norm", False) if self.use_v_norm: self.v_norm = SwarmRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.v_gate = nn.Parameter(torch.zeros(1)) def forward(self, hidden_states, cos, sin, attention_mask=None, past_key_values=None, cache_position=None): b, s, _ = hidden_states.shape q = self.q_proj(hidden_states).view(b, s, self.num_heads, self.head_dim) k = self.k_proj(hidden_states).view(b, s, self.num_kv_heads, self.head_dim) v = self.v_proj(hidden_states).view(b, s, self.num_kv_heads, self.head_dim) q = self.q_norm(q).transpose(1, 2) k = self.k_norm(k).transpose(1, 2) if self.use_v_norm: v = v + self.v_gate * (self.v_norm(v) - v) # gated; identity when v_gate=0 v = v.transpose(1, 2) q, k = apply_rotary(q, k, cos, sin) if past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx, {"cache_position": cache_position}) k = repeat_kv(k, self.n_rep) v = repeat_kv(v, self.n_rep) is_causal = attention_mask is None and s > 1 attn = F.scaled_dot_product_attention( q, k, v, attn_mask=attention_mask, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=is_causal, ) attn = attn.transpose(1, 2).contiguous().view(b, s, -1) return self.o_proj(attn) # ----------------------------------------------------------------------------- # # Experts + Swarm-MoE block # # ----------------------------------------------------------------------------- # class SwiGLUExpert(nn.Module): """A single SwiGLU feed-forward expert.""" def __init__(self, config: SwarmMoEConfig): super().__init__() self.gate_proj = make_linear(config.hidden_size, config.intermediate_size, False, config) self.up_proj = make_linear(config.hidden_size, config.intermediate_size, False, config) self.down_proj = make_linear(config.intermediate_size, config.hidden_size, False, config) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class SwarmMoE(nn.Module): """Swarm Mixture-of-Experts: top-k routed experts (sparse) + shared expert(s) (dense). Returns the mixed hidden states and the load-balancing auxiliary loss. """ def __init__(self, config: SwarmMoEConfig): super().__init__() self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.norm_topk_prob = config.norm_topk_prob self.jitter = config.router_jitter_noise self.aux_coef = config.router_aux_loss_coef self.z_coef = getattr(config, "router_z_loss_coef", 0.0) # Deeper router: a zero-init residual pre-layer makes routing input-adaptive with # extra depth, while being the identity at init (retains a loaded gate exactly). self.deeper_router = getattr(config, "deeper_router", False) self.has_routed_experts = self.num_experts > 0 and self.top_k > 0 if self.deeper_router and self.has_routed_experts: self.router_pre = ZeroLinear(config.hidden_size, config.hidden_size, bias=False) else: self.router_pre = None if self.has_routed_experts: self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) self.experts = nn.ModuleList([SwiGLUExpert(config) for _ in range(config.num_experts)]) else: self.gate = None self.experts = nn.ModuleList() self.num_shared_experts = getattr(config, "num_shared_experts", 1) if self.num_shared_experts > 0: self.shared_expert = SwiGLUExpert(config) else: self.shared_expert = None def forward(self, hidden_states: torch.Tensor, router_bias: Optional[torch.Tensor] = None): b, s, d = hidden_states.shape x = hidden_states.view(-1, d) # [N, d] n_tokens = x.shape[0] if self.training and self.jitter > 0: x = x * torch.empty_like(x).uniform_(1.0 - self.jitter, 1.0 + self.jitter) out = torch.zeros_like(x) if self.has_routed_experts: route_in = x + F.silu(self.router_pre(x)) if self.deeper_router else x router_logits = self.gate(route_in) # [N, E] if router_bias is not None: router_logits = router_logits + router_bias.to(device=router_logits.device, dtype=router_logits.dtype) routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) top_w, top_i = torch.topk(routing_weights, self.top_k, dim=-1) # [N, k] if self.norm_topk_prob: top_w = top_w / top_w.sum(dim=-1, keepdim=True) top_w = top_w.to(x.dtype) # one-hot expert assignment for vectorized per-expert gather: [E, k, N] expert_mask = F.one_hot(top_i, num_classes=self.num_experts).permute(2, 1, 0) for e in range(self.num_experts): idx_k, idx_tok = torch.where(expert_mask[e]) if idx_tok.numel() == 0: continue cur = x[idx_tok] y = self.experts[e](cur) * top_w[idx_tok, idx_k, None] out.index_add_(0, idx_tok, y.to(out.dtype)) if self.shared_expert is not None: out = out + self.shared_expert(x) if not self.has_routed_experts: return out.view(b, s, d), out.new_zeros(()) # Two router regularizers, summed and pre-scaled: # 1) Load-balancing (Switch/Mixtral): keeps all experts utilized so capacity isn't # wasted and routing doesn't collapse onto a few experts. # 2) Router z-loss (ST-MoE): keeps the gating logits small/well-conditioned, which # stabilizes training and prevents the softmax from saturating. lb_loss = self._aux_loss(routing_weights, top_i, n_tokens) z_loss = (torch.logsumexp(router_logits, dim=-1) ** 2).mean() aux = self.aux_coef * lb_loss + self.z_coef * z_loss return out.view(b, s, d), aux def _aux_loss(self, routing_weights, top_i, n_tokens): """Switch-Transformer load-balancing loss L = E * sum_i f_i * P_i. * f_i = fraction of routed (token,slot) assignments to expert i — a hard count, so computed under no_grad (the *actual* load; non-differentiable). * P_i = mean gating probability for expert i (the router's *intended* load; gradients flow here). * Minimizing f.P pushes probability mass away from overloaded experts toward idle ones. At perfect balance f_i = P_i = 1/E, so L = 1.0 (its minimum) — a clean target to watch in the logs. """ with torch.no_grad(): counts = torch.zeros(self.num_experts, device=top_i.device, dtype=routing_weights.dtype) ones = torch.ones(top_i.numel(), device=top_i.device, dtype=routing_weights.dtype) counts.scatter_add_(0, top_i.reshape(-1), ones) f = counts / (n_tokens * self.top_k) # actual load (detached) p = routing_weights.mean(dim=0) # intended load (differentiable) return (f * p).sum() * self.num_experts # ----------------------------------------------------------------------------- # # Decoder layer # # ----------------------------------------------------------------------------- # class SwarmDecoderLayer(nn.Module): def __init__(self, config: SwarmMoEConfig, layer_idx: int): super().__init__() self.input_layernorm = SwarmRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.self_attn = SwarmAttention(config, layer_idx) self.post_attention_layernorm = SwarmRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp = SwarmMoE(config) def forward(self, hidden_states, cos, sin, attention_mask=None, past_key_values=None, cache_position=None, router_bias: Optional[torch.Tensor] = None): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, cos, sin, attention_mask, past_key_values, cache_position) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states, aux_loss = self.mlp(hidden_states, router_bias=router_bias) hidden_states = residual + hidden_states return hidden_states, aux_loss # ----------------------------------------------------------------------------- # # Base / model / causal-LM # # ----------------------------------------------------------------------------- # class SwarmMoEPreTrainedModel(PreTrainedModel): config_class = SwarmMoEConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["SwarmDecoderLayer"] _supports_sdpa = True _supports_cache_class = True _supports_flash_attn = False def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, ZeroLinear): # identity-at-init add-ons module.weight.data.zero_() if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) elif isinstance(module, SwarmRMSNorm): module.weight.data.fill_(1.0) if isinstance(module, SwarmAttention) and getattr(module, "v_gate", None) is not None: module.v_gate.data.zero_() # gated V-norm starts as identity class SwarmMoEModel(SwarmMoEPreTrainedModel): def __init__(self, config: SwarmMoEConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) # Optional ALBERT-style cross-layer sharing: build `num_unique_layers` real # layers and reuse them to fill `num_hidden_layers` (better parameter sharing). num_unique = getattr(config, "num_unique_layers", None) or config.num_hidden_layers self.num_unique_layers = num_unique unique_layers = [SwarmDecoderLayer(config, i) for i in range(num_unique)] self._unique = nn.ModuleList(unique_layers) # layer_idx -> unique module index, evenly grouped group = max(1, round(config.num_hidden_layers / num_unique)) self.layer_plan = [min(i // group, num_unique - 1) for i in range(config.num_hidden_layers)] self.norm = SwarmRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = SwarmRotaryEmbedding(config.head_dim, base=config.rope_theta) self.gradient_checkpointing = False self.post_init() @property def layers(self): # exposes the *effective* (possibly shared) layer sequence return [self._unique[p] for p in self.layer_plan] def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, cache_position=None, output_router_logits=False, thinking_steps=None, memory_embeds=None, router_bias=None, plasticity_state=None, **kwargs, ): use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training) if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if _dynamic_cache_takes_config() else DynamicCache() past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 if cache_position is None: cache_position = torch.arange(past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device) if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = self._build_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values) cos, sin = self.rotary_emb(inputs_embeds, position_ids) hidden_states = inputs_embeds hidden_states = self._inject_memory(hidden_states, memory_embeds) runtime_router_bias = self._runtime_router_bias( hidden_states, router_bias=router_bias, plasticity_state=plasticity_state, ) total_aux = 0.0 layers = self.layers for layer_idx, layer in enumerate(layers): layer_router_bias = self._select_router_bias(runtime_router_bias, layer_idx, hidden_states) if self.gradient_checkpointing and self.training: hidden_states, aux = self._gradient_checkpointing_func( layer.__call__, hidden_states, cos, sin, causal_mask, past_key_values, cache_position, layer_router_bias, ) else: hidden_states, aux = layer( hidden_states, cos, sin, causal_mask, past_key_values, cache_position, router_bias=layer_router_bias, ) total_aux = total_aux + aux extra_steps = self._thinking_steps(thinking_steps, past_key_values) if extra_steps > 0 and layers: think_idx = getattr(self.config, "test_time_compute_layer", -1) think_layer = layers[think_idx] think_bias = self._select_router_bias(runtime_router_bias, think_idx, hidden_states) for _ in range(extra_steps): hidden_states, aux = think_layer( hidden_states, cos, sin, causal_mask, None, cache_position, router_bias=think_bias, ) total_aux = total_aux + aux hidden_states = self.norm(hidden_states) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, router_logits=(total_aux,) if output_router_logits else None, ) def _inject_memory(self, hidden_states, memory_embeds): """Zero-parameter memory influence aligned to the model hidden space.""" strength = float(getattr(self.config, "memory_injection_strength", 0.0) or 0.0) if memory_embeds is None or strength == 0.0: return hidden_states mem = memory_embeds.to(device=hidden_states.device, dtype=hidden_states.dtype) if mem.dim() == 2: mem = mem[:, None, :] if mem.shape[-1] != hidden_states.shape[-1]: return hidden_states mem = mem.mean(dim=1, keepdim=True) mem = F.normalize(mem.float(), dim=-1).to(hidden_states.dtype) gate = torch.sigmoid((hidden_states.float() * mem.float()).mean(dim=-1, keepdim=True)).to(hidden_states.dtype) return hidden_states + (strength * gate * mem) def _runtime_router_bias(self, hidden_states, router_bias=None, plasticity_state=None): """Build a per-expert bias without adding parameters or copying model weights.""" bias = router_bias if plasticity_state is not None: state_bias = plasticity_state.get("router_bias") if isinstance(plasticity_state, dict) else None if state_bias is not None: bias = state_bias if bias is None else bias + state_bias.to(bias.device) if bias is None: return None strength = float(getattr(self.config, "router_bias_strength", 0.0) or 0.0) plastic = float(getattr(self.config, "plasticity_strength", 0.0) or 0.0) scale = strength + plastic if scale == 0.0: return None return bias.to(device=hidden_states.device, dtype=hidden_states.dtype) * scale def _select_router_bias(self, bias, layer_idx, hidden_states): if bias is None: return None if bias.dim() == 1: return bias if bias.dim() == 2: idx = layer_idx % bias.shape[0] return bias[idx] return None def _thinking_steps(self, thinking_steps, past_key_values): steps = getattr(self.config, "test_time_compute_steps", 0) if thinking_steps is None else thinking_steps steps = int(steps or 0) if steps <= 0: return 0 # Cache mutation and repeated internal passes do not mix; use prompt/full-pass thinking. if past_key_values is not None and past_key_values.get_seq_length() > 0: return 0 return min(steps, 8) def _build_causal_mask(self, attention_mask, inputs_embeds, cache_position, past_key_values): # Returns None to let SDPA use its fast is_causal path when there is no padding # and no cache; otherwise builds an additive float mask. b, q_len, _ = inputs_embeds.shape if attention_mask is None and (past_key_values is None or past_key_values.get_seq_length() == 0): return None kv_len = (past_key_values.get_seq_length() if past_key_values is not None else 0) + q_len dtype = inputs_embeds.dtype min_val = torch.finfo(dtype).min causal = torch.full((q_len, kv_len), min_val, dtype=dtype, device=inputs_embeds.device) positions = cache_position.reshape(-1, 1) kv_idx = torch.arange(kv_len, device=inputs_embeds.device).reshape(1, -1) causal = causal.masked_fill(kv_idx <= positions, 0.0) causal = causal[None, None, :, :].expand(b, 1, q_len, kv_len).clone() if attention_mask is not None: pad = attention_mask[:, None, None, :].to(dtype) causal = causal.masked_fill(pad == 0, min_val) return causal def _dynamic_cache_takes_config() -> bool: import inspect try: return "config" in inspect.signature(DynamicCache.__init__).parameters except (ValueError, TypeError): return False class SwarmMoEForCausalLM(SwarmMoEPreTrainedModel, GenerationMixin): # When tie_word_embeddings is True we reuse the embedding matrix *functionally* (no # separate lm_head parameter), which avoids the meta-tensor / tie pitfalls of lazy # loading entirely — so the exported model loads cleanly for everyone. When untied, a # real lm_head Linear is created and saved (also meta-safe, since it is in the ckpt). _tied_weights_keys = None def __init__(self, config: SwarmMoEConfig): super().__init__(config) self.model = SwarmMoEModel(config) self.vocab_size = config.vocab_size self.lm_head = None if config.tie_word_embeddings else \ nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.router_aux_loss_coef = config.router_aux_loss_coef self._vision_sidecar = None self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head # None when weights are tied (functional head) def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def _compute_logits(self, hidden): if self.lm_head is not None: return self.lm_head(hidden) return F.linear(hidden, self.model.embed_tokens.weight) def set_decoder(self, decoder): self.model = decoder def _get_vision_sidecar(self, device: torch.device, dtype: torch.dtype): if not bool(getattr(self.config, "vision_sidecar_enabled", False)): return None sidecar_path = getattr(self.config, "vision_sidecar_path", None) if not sidecar_path: return None path = Path(sidecar_path).expanduser() if not path.is_absolute(): base = Path(str(getattr(self.config, "_name_or_path", "") or ".")).expanduser() path = base / path manifest = path / "manifest.json" weights = path / "model.safetensors" if not manifest.exists() or not weights.exists(): return None if self._vision_sidecar is not None and getattr(self._vision_sidecar, "sidecar_dir", None) == path: return self._vision_sidecar from .vision_sidecar import Qwen35VisionSidecar self._vision_sidecar = Qwen35VisionSidecar( path, device=device, dtype=dtype if dtype in (torch.float16, torch.bfloat16, torch.float32) else torch.float32, max_blocks=int(getattr(self.config, "vision_sidecar_max_blocks", 24) or 24), ) return self._vision_sidecar @torch.no_grad() def route_vla_action(self, vision_tensor: torch.Tensor, command: str = "", action_dim: int = 8): """Lightweight browser VLA routing hook. This is intentionally runtime-only and weight-free: it lets the browser tool pass a screenshot tensor through a stable HF-compatible method and receive grounded control priors. The learned language model still consumes the text/OCR/browser observation; this hook provides visual-state routing metadata for the external Playwright actuator without changing checkpoint architecture or parameters. """ if vision_tensor is None: return {"status": "no_vision_tensor", "controls": {}, "route": {}} x = vision_tensor.detach().float() if x.ndim == 3: x = x.unsqueeze(0) if x.ndim != 4: return {"status": "bad_vision_tensor", "shape": list(x.shape), "controls": {}, "route": {}} # Normalize expected BCHW/BHWC inputs into BCHW. if x.shape[1] not in {1, 3, 4} and x.shape[-1] in {1, 3, 4}: x = x.permute(0, 3, 1, 2).contiguous() if x.max() > 2.0: x = x / 255.0 x = x.clamp(0.0, 1.0) gray = x[:, :3].mean(dim=1, keepdim=True) if x.shape[1] >= 3 else x[:, :1] brightness = float(gray.mean().item()) contrast = float(gray.std().item()) h_edges = torch.mean(torch.abs(gray[:, :, :, 1:] - gray[:, :, :, :-1])).item() if gray.shape[-1] > 1 else 0.0 v_edges = torch.mean(torch.abs(gray[:, :, 1:, :] - gray[:, :, :-1, :])).item() if gray.shape[-2] > 1 else 0.0 edge_density = float((h_edges + v_edges) * 0.5) lower = (command or "").lower() wants_click = any(token in lower for token in ("click", "press", "select", "open")) wants_type = any(token in lower for token in ("type", "enter", "fill", "search")) wants_scroll = any(token in lower for token in ("scroll", "more", "lower", "down")) wants_observe = any(token in lower for token in ("observe", "inspect", "verify", "read", "summarize")) controls = { "mouse": { "click": float(0.85 if wants_click else 0.15), "scroll": float(0.85 if wants_scroll else (0.35 if edge_density > 0.08 else 0.10)), "center_x": 0.5, "center_y": 0.52, }, "keyboard": { "type": float(0.85 if wants_type else 0.05), "enter": float(0.45 if wants_type else 0.05), }, "observe": float(0.90 if wants_observe or not (wants_click or wants_type or wants_scroll) else 0.35), } sidecar_packet = None try: sidecar = self._get_vision_sidecar(x.device, x.dtype) if sidecar is not None: sidecar_packet = sidecar.route(x, command=command, action_dim=action_dim) sidecar_controls = sidecar_packet.get("controls") or {} if sidecar_packet.get("status") == "ok" and sidecar_controls: # Blend sidecar priors with cheap visual priors; keep deterministic routing. for group, values in sidecar_controls.items(): if isinstance(values, dict): controls.setdefault(group, {}) for name, value in values.items(): if isinstance(value, (int, float)): controls[group][name] = float((controls[group].get(name, value) + value) * 0.5) elif isinstance(value := values, (int, float)): controls[group] = float((controls.get(group, value) + value) * 0.5) except Exception as exc: sidecar_packet = {"status": "error", "error": f"{type(exc).__name__}: {exc}", "route": {}} return { "status": "ok", "controls": controls, "route": { "source": "SwarmMoEForCausalLM.route_vla_action", "weight_free": sidecar_packet is None, "runtime_only": True, "visual_features": { "brightness": round(brightness, 4), "contrast": round(contrast, 4), "edge_density": round(edge_density, 4), }, "intent": { "click": wants_click, "type": wants_type, "scroll": wants_scroll, "observe": wants_observe, }, "policy": "browser screenshot -> visual priors -> Playwright action; text/OCR remains in model prompt", "action_dim": int(action_dim or 0), "vision_sidecar": sidecar_packet.get("route", sidecar_packet) if isinstance(sidecar_packet, dict) else None, }, } def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, cache_position=None, logits_to_keep: Union[int, torch.Tensor] = 0, thinking_steps=None, memory_embeds=None, router_bias=None, plasticity_state=None, **kwargs, ): outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, output_router_logits=labels is not None, thinking_steps=thinking_steps, memory_embeds=memory_embeds, router_bias=router_bias, plasticity_state=plasticity_state, ) hidden_states = outputs.last_hidden_state slice_idx = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self._compute_logits(hidden_states[:, slice_idx, :]) loss = None aux_loss = None if labels is not None: logits_f = logits.float() shift_logits = logits_f[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, self.vocab_size), shift_labels.view(-1), ignore_index=-100, ) if outputs.router_logits is not None: aux_loss = outputs.router_logits[0] # already scaled (aux_coef*lb + z_coef*z) loss = loss + aux_loss.to(loss.device) return MoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, ) __all__ = [ "SwarmMoEPreTrainedModel", "SwarmMoEModel", "SwarmMoEForCausalLM", "BitLinear", "ZeroLinear", ]