"""arch_replace — architectural replacement via dual-path QAT distillation. Implements the "Jamba-conversion" idea (roadmap Stage 19e): replace attention layers with Mamba (selective SSM) layers using the SAME dual-path infrastructure as block-wise ternarization (Stage 18b), but for architectural substitution instead of weight quantization. Key insight: QuantizedModule(dual_path=True) + dual_path_loss is a generic mechanism. It can ternarize (18b), Jamba-convert (19e), or train MTP (Stage 24). One engine — three applications. Components: - MambaLayer — selective SSM (fixed recurrent state, input-dependent A/B/C/dt) - ReplaceModule — generic dual-path wrapper for architectural substitution - jamba_replace — block-wise progressive attention->Mamba replacement """ from __future__ import annotations import logging import math from typing import Any import torch import torch.nn as nn import torch.nn.functional as F logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # MambaLayer — selective State-Space Model (Mamba-style) # --------------------------------------------------------------------------- class MambaLayer(nn.Module): """Selective SSM layer with fixed recurrent state (Mamba architecture). Unlike attention, the recurrent state is fixed-size (d_state x d_inner) and does NOT grow with sequence length. This makes long context cheap in memory. Parameters (all learnable via standard backprop, no STE needed — these are continuous parameters, not quantized): in_proj — input projection: d_model -> 2*d_inner (gate + value branches) out_proj — output projection: d_inner -> d_model A_log — log-space diagonal of state matrix [d_inner, d_state] D — skip-connection scale [d_inner] dt_proj — timestep projection: d_inner -> d_inner (input-dependent dt) Forward (selective scan, simplified): x [B, L, d_model] -> in_proj -> (x_gate, x_value) [B, L, d_inner] -> dt = softplus(dt_proj(x_gate)) [B, L, d_inner] -> A = -exp(A_log) [d_inner, d_state] (constant per channel) -> selective_scan: for each t, state[t+1] = state[t] * exp(A*dt[t]) + B(t) * x_value(t); output[t] = C(t) . state[t] -> out = output * silu(x_gate) (gated) -> out_proj -> [B, L, d_model] The selective_scan is implemented as a sequential loop over time (Python- level). This is NOT optimized for speed (no CUDA kernel) — it is a reference implementation for distillation training and correctness. For production inference, a fused scan kernel would be needed. """ def __init__( self, d_model: int, d_state: int = 16, d_inner: int | None = None, dt_rank: int | None = None, bias: bool = False, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32, ): super().__init__() self.d_model = d_model self.d_state = d_state # d_inner defaults to 2*d_model (standard Mamba expansion). self.d_inner = d_inner if d_inner is not None else 2 * d_model # dt_rank defaults to d_model/16 (low-rank projection for dt). self.dt_rank = dt_rank if dt_rank is not None else max(1, d_model // 16) _dt = {"device": device, "dtype": dtype} # Input projection: x -> (gate, value, dt_base) of total 2*d_inner + dt_rank. self.in_proj = nn.Linear(d_model, 2 * self.d_inner + self.dt_rank, bias=bias, **_dt) # Timestep projection: dt_base -> d_inner (produces input-dependent dt). self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True, **_dt) # A_log: learnable log-space diagonal of the state matrix. # Initialized as in Mamba paper: A = -arange(1, d_state+1) repeated. A_init = torch.arange(1, d_state + 1, dtype=torch.float32).expand(self.d_inner, d_state).contiguous() self.A_log = nn.Parameter(torch.log(A_init).to(**_dt)) # D: skip-connection scale, initialized to 1. self.D = nn.Parameter(torch.ones(self.d_inner, **_dt)) # Output projection: d_inner -> d_model. self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias, **_dt) def _selective_scan( self, x_value: torch.Tensor, dt: torch.Tensor, B: torch.Tensor, C: torch.Tensor, ) -> torch.Tensor: """Sequential selective scan over time dimension. Args: x_value: [B, L, d_inner] — input values. dt: [B, L, d_inner] — input-dependent timesteps. B: [B, L, d_state] — input-dependent B matrix (interaction weights). C: [B, L, d_state] — input-dependent C matrix (readout weights). Returns: y: [B, L, d_inner] — scanned output. """ batch, seq_len, d_inner = x_value.shape d_state = self.d_state device = x_value.device dtype = x_value.dtype # A is constant (from A_log): [d_inner, d_state], negative. A = -torch.exp(self.A_log.to(dtype)) # [d_inner, d_state] # Discretized A: dA = exp(A * dt) [B, L, d_inner, d_state] # dt: [B, L, d_inner] -> [B, L, d_inner, 1] dt_exp = dt.unsqueeze(-1) # [B, L, d_inner, 1] dA = torch.exp(A.unsqueeze(0).unsqueeze(0) * dt_exp) # [B, L, d_inner, d_state] # Discretized B: dB = dt * B [B, L, d_inner, d_state] # B: [B, L, d_state] -> [B, L, 1, d_state] -> broadcast to d_inner dB = dt_exp * B.unsqueeze(2) # [B, L, d_inner, d_state] # Sequential scan (Python loop — reference implementation). h = torch.zeros(batch, d_inner, d_state, device=device, dtype=dtype) ys = [] for t in range(seq_len): # state update: h = dA * h + dB * x h = dA[:, t] * h + dB[:, t] * x_value[:, t].unsqueeze(-1) # output: y = C . h (dot product over d_state) y_t = (h * C[:, t].unsqueeze(1)).sum(dim=-1) # [B, d_inner] ys.append(y_t) y = torch.stack(ys, dim=1) # [B, L, d_inner] return y def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward pass. Args: x: [B, L, d_model] or [L, d_model] (auto-unsqueeze batch). Returns: [B, L, d_model] or [L, d_model] (matches input batch dim). """ squeeze_batch = False if x.dim() == 2: x = x.unsqueeze(0) # [1, L, d_model] squeeze_batch = True batch, seq_len, _ = x.shape dtype = x.dtype # Input projection: x -> (gate, value, dt_base) xzdr = self.in_proj(x) # [B, L, 2*d_inner + dt_rank] gate, value, dt_base = torch.split( xzdr, [self.d_inner, self.d_inner, self.dt_rank], dim=-1 ) # dt: input-dependent timestep via softplus(dt_proj(dt_base)). dt = F.softplus(self.dt_proj(dt_base)) # [B, L, d_inner] # Clamp dt to avoid numerical issues (exp(A*dt) overflow). dt = dt.clamp(max=50.0) # B and C matrices: derived from gate via a linear split. # In Mamba, B/C are input-dependent. Here we use a simplified approach: # project gate into d_state dims for B and C. # Use first d_state and second d_state slices of value for B, C. # This is a lightweight variant — full Mamba uses separate projections. # For distillation purposes (learning to mimic attention output), # this simplified B/C is sufficient as the layer will be trained. B = gate[..., :self.d_state] # [B, L, d_state] C = value[..., :self.d_state] # [B, L, d_state] # Selective scan. y = self._selective_scan(value, dt, B, C) # [B, L, d_inner] # Gated output: y * silu(gate) + D * value (skip connection). y = y * F.silu(gate) y = y + self.D.unsqueeze(0).unsqueeze(0) * value # Output projection. out = self.out_proj(y) # [B, L, d_model] if squeeze_batch: out = out.squeeze(0) return out def reset_state(self) -> None: """No-op — recurrent state is per-forward (ephemeral), not stored.""" pass # --------------------------------------------------------------------------- # ReplaceModule — generic dual-path wrapper for architectural substitution # --------------------------------------------------------------------------- class ReplaceModule(nn.Module): """Generic dual-path wrapper: student (new architecture) vs teacher (frozen original). This is the architectural analogue of QuantizedModule(dual_path=True). Instead of quantizing weights, it replaces the module TYPE entirely: - student: new architecture (e.g. MambaLayer) — learnable - teacher: frozen original (e.g. nn.MultiheadAttention) — reference forward(x, path): - "student" -> student.forward(x) (for inference after training) - "teacher" -> teacher.forward(x) [no_grad] (for validation) - "both" -> (student_out, teacher_out) (for dual_path_loss) distillation: dual_path_loss(student_out, teacher_out, "mse") uses the existing training_unified.dual_path_loss function. """ def __init__( self, student: nn.Module, teacher: nn.Module | None = None, dual_path: bool = True, teacher_forward_fn: Any = None, student_forward_fn: Any = None, ): super().__init__() self.student = student self._dual_path = dual_path self._has_teacher = teacher is not None # Optional custom forward functions for non-standard call signatures. # teacher_forward_fn(module, x) -> output (e.g. for MultiheadAttention: # lambda m, x: m(x, x, x)) # student_forward_fn(module, x) -> output (e.g. for MambaLayer that # needs extra args) # If None, defaults to module(x). self._teacher_fn = teacher_forward_fn self._student_fn = student_forward_fn if teacher is not None: # Register teacher as a non-trainable submodule. self.teacher = teacher for p in self.teacher.parameters(): p.requires_grad = False else: # No teacher — single-path mode (student only). self._dual_path = False def _call_student(self, x: torch.Tensor) -> torch.Tensor: if self._student_fn is not None: return self._student_fn(self.student, x) return self.student(x) def _call_teacher(self, x: torch.Tensor) -> torch.Tensor: if self._teacher_fn is not None: return self._teacher_fn(self.teacher, x) return self.teacher(x) def forward( self, x: torch.Tensor, path: str = "student" ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if path == "teacher": if not self._has_teacher: raise RuntimeError("teacher_forward called but no teacher set") with torch.no_grad(): return self._call_teacher(x) if path == "both": if not self._has_teacher: raise RuntimeError("dual_path='both' requires a teacher") student_out = self._call_student(x) with torch.no_grad(): teacher_out = self._call_teacher(x) return student_out, teacher_out # Default: student only. return self._call_student(x) def distillation_loss( self, x: torch.Tensor, loss_type: str = "mse" ) -> torch.Tensor: """Compute dual-path distillation loss (student vs teacher).""" from agiws_neural_quant.training_unified import dual_path_loss student_out, teacher_out = self.forward(x, path="both") return dual_path_loss(student_out, teacher_out, loss_type) def freeze_student(self) -> None: """Freeze student parameters (after training, for inference).""" for p in self.student.parameters(): p.requires_grad = False def unfreeze_student(self) -> None: """Unfreeze student parameters (for further training).""" for p in self.student.parameters(): p.requires_grad = True # --------------------------------------------------------------------------- # jamba_replace — block-wise progressive attention->Mamba replacement # --------------------------------------------------------------------------- def _count_parameters(module: nn.Module) -> int: return sum(p.numel() for p in module.parameters()) def _train_replace_block( replace_module: ReplaceModule, inputs: list[torch.Tensor], epochs: int = 100, lr: float = 1e-4, loss_type: str = "mse", verbose: bool = True, ) -> dict[str, float]: """Train a single ReplaceModule's student to mimic its teacher. Args: replace_module: ReplaceModule with student + teacher. inputs: list of calibration input tensors (each [B, L, d_model]). epochs: number of training epochs. lr: learning rate. loss_type: "mse", "cosine", or "kl". verbose: print progress. Returns: dict with 'loss_start', 'loss_end', 'cosine_end'. """ from agiws_neural_quant.training_unified import dual_path_loss optimizer = torch.optim.Adam( [p for p in replace_module.student.parameters() if p.requires_grad], lr=lr, ) loss_start = None loss_end = None cosine_end = None for epoch in range(epochs): epoch_loss = 0.0 n_batches = 0 for x in inputs: optimizer.zero_grad() student_out, teacher_out = replace_module(x, path="both") loss = dual_path_loss(student_out, teacher_out, loss_type) loss.backward() optimizer.step() epoch_loss += loss.item() n_batches += 1 avg_loss = epoch_loss / max(n_batches, 1) if loss_start is None: loss_start = avg_loss loss_end = avg_loss if verbose and (epoch % max(1, epochs // 10) == 0 or epoch == epochs - 1): logger.info(f" epoch {epoch}/{epochs}: {loss_type}={avg_loss:.6f}") # Final cosine similarity. if inputs: with torch.no_grad(): s_out, t_out = replace_module(inputs[0], path="both") cos = F.cosine_similarity( s_out.flatten().unsqueeze(0), t_out.flatten().unsqueeze(0), ).item() cosine_end = cos return { "loss_start": loss_start or 0.0, "loss_end": loss_end or 0.0, "cosine_end": cosine_end or 0.0, } def jamba_replace( model: nn.Module, calib_inputs: list[torch.Tensor], mamba_ratio: float = 0.85, d_state: int = 16, epochs_per_block: int = 100, lr: float = 1e-4, loss_type: str = "mse", teacher_call_fn: Any = None, verbose: bool = True, ) -> tuple[nn.Module, list[dict[str, Any]]]: """Replace attention layers with Mamba via block-wise dual-path distillation. Walks the model tree, finds attention layers (by type name containing "Attention" or "MultiheadAttention"), replaces mamba_ratio fraction of them with MambaLayer, and trains each replacement via dual-path distillation against the frozen original attention layer. Placement: attention layers are kept periodically (every N-th) so that ~15% remain for recall. The rest are replaced with Mamba. Args: model: source model with attention layers. calib_inputs: list of calibration tensors [B, L, d_model] for distillation. mamba_ratio: fraction of attention layers to replace (0.85 = 85% -> Mamba). d_state: Mamba recurrent state size. epochs_per_block: distillation epochs per replaced block. lr: learning rate for distillation. loss_type: "mse", "cosine", or "kl". teacher_call_fn: callable(module, x) -> output for non-standard teacher forward signatures. Default: None — auto-detects: - nn.MultiheadAttention -> lambda m, x: m(x, x, x)[0] (self-attention: q=k=v=x, extract output from tuple) - everything else -> module(x) For cross-attention, pass your own: lambda m, x: m(decoder_h, encoder_h, encoder_h) verbose: print progress. Returns: (modified_model, report) where report is a list of per-block dicts. """ # Collect all attention modules by path (top-level only — skip nested # attention submodules like nn.MultiheadAttention inside a wrapper). all_named = list(model.named_modules()) attn_candidates: list[str] = [] for name, module in all_named: if not name: # skip root module continue type_name = type(module).__name__ if "Attention" in type_name or "MultiheadAttention" in type_name: # Skip if it's inside an already-replaced ReplaceModule. if not any(p in name for p in (".student.", ".teacher.")): attn_candidates.append(name) # Keep only top-level (no parent in attn_candidates). attn_paths: list[str] = [] for n in attn_candidates: if not any(n != other and n.startswith(other + ".") for other in attn_candidates): attn_paths.append(n) if not attn_paths: if verbose: logger.info("jamba_replace: no attention layers found, nothing to replace") return model, [] n_attn = len(attn_paths) n_keep = max(1, int(n_attn * (1.0 - mamba_ratio))) n_replace = n_attn - n_keep if verbose: logger.info( f"jamba_replace: {n_attn} attention layers, " f"keep {n_keep}, replace {n_replace} with Mamba" ) # Placement: keep attention layers spread evenly (every N-th). # Indices to keep: evenly spaced across the list. keep_indices = set() if n_keep > 0: step = n_attn / n_keep for i in range(n_keep): keep_indices.add(int(i * step)) report: list[dict[str, Any]] = [] for idx, attn_path in enumerate(attn_paths): # Navigate to the parent module and get the attention submodule. parts = attn_path.split(".") parent = model for p in parts[:-1]: parent = getattr(parent, p) attr_name = parts[-1] attn_module = getattr(parent, attr_name) if idx in keep_indices: report.append({ "path": attn_path, "action": "keep_attention", "index": idx, }) if verbose: logger.info(f" [{idx}] KEEP attention: {attn_path}") continue # Determine d_model from the attention module. # Try common attributes; fallback to calib_inputs shape. d_model = getattr(attn_module, "embed_dim", None) if d_model is None: d_model = getattr(attn_module, "d_model", None) if d_model is None and calib_inputs: d_model = calib_inputs[0].shape[-1] if d_model is None: report.append({ "path": attn_path, "action": "skip_unknown_dim", "index": idx, "reason": "cannot determine d_model", }) if verbose: logger.info(f" [{idx}] SKIP (unknown d_model): {attn_path}") continue # Create Mamba student. device = next(attn_module.parameters()).device dtype = next(attn_module.parameters()).dtype mamba_student = MambaLayer( d_model=d_model, d_state=d_state, device=device, dtype=dtype, ).to(device=device, dtype=dtype) # Auto-detect teacher call function for nn.MultiheadAttention. effective_fn = teacher_call_fn if effective_fn is None: if isinstance(attn_module, nn.MultiheadAttention): # Self-attention: q=k=v=x, MHA returns (output, weights). effective_fn = lambda m, x: m(x, x, x)[0] # Create ReplaceModule with teacher = frozen attention. replace_mod = ReplaceModule( student=mamba_student, teacher=attn_module, dual_path=True, teacher_forward_fn=effective_fn, ) # Distillation training. if verbose: logger.info( f" [{idx}] REPLACE attention->Mamba: {attn_path} " f"(d_model={d_model}, d_state={d_state})" ) block_report = _train_replace_block( replace_mod, calib_inputs, epochs=epochs_per_block, lr=lr, loss_type=loss_type, verbose=verbose, ) # Freeze student and install. replace_mod.freeze_student() setattr(parent, attr_name, replace_mod) report.append({ "path": attn_path, "action": "replace_mamba", "index": idx, "d_model": d_model, "d_state": d_state, "teacher_params": _count_parameters(attn_module), "student_params": _count_parameters(mamba_student), **block_report, }) if verbose: logger.info( f" [{idx}] DONE: cosine={block_report['cosine_end']:.4f}, " f"loss {block_report['loss_start']:.6f} -> {block_report['loss_end']:.6f}" ) return model, report __all__: list[str] = []