import json import math from pathlib import Path from typing import Optional, Tuple import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint as _torch_checkpoint def checkpoint(fn, *args, use_reentrant=False, **kwargs): # DeepSpeed's activation checkpoint runs the forward under no_grad # (reentrant); with a frozen self-decoder the loss comes out with no # grad_fn. Torch's non-reentrant checkpoint tracks internal params. return _torch_checkpoint(fn, *args, use_reentrant=use_reentrant, **kwargs) from .config import SharonaConfig from .layers import RMSNorm, RotaryEmbedding, SwiGLU from .attention import GQAAttention, CrossAttention, KVState from .kv_cache import KVCache class SelfDecoderLayer(nn.Module): def __init__(self, config: SharonaConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = GQAAttention( config, layer_idx, sliding_window=config.sliding_window_size ) self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.ffn = SwiGLU(config) def forward( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, rope: RotaryEmbedding, kv_cache: Optional[KVState] = None, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, KVState]: residual = hidden_states normed = self.attn_norm(hidden_states) attn_out, new_kv = self.attn(normed, position_ids, rope, kv_cache, attention_mask) hidden_states = residual + attn_out residual = hidden_states hidden_states = residual + self.ffn(self.ffn_norm(hidden_states)) return hidden_states, new_kv class CrossDecoderLayer(nn.Module): def __init__(self, config: SharonaConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.num_persistent_contexts = config.num_persistent_contexts self.self_attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.self_attn = GQAAttention(config, layer_idx, sliding_window=None) self.cross_attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.cross_attn = CrossAttention(config, layer_idx) if self.num_persistent_contexts >= 2: self.cross_attn_m3_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.cross_attn_m3 = CrossAttention(config, layer_idx) self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.ffn = SwiGLU(config) def forward( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, rope: RotaryEmbedding, shared_kv: KVState, self_kv_cache: Optional[KVState] = None, self_attn_mask: Optional[torch.Tensor] = None, cross_attn_mask: Optional[torch.Tensor] = None, m3_kv: Optional[KVState] = None, m3_attn_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, KVState]: residual = hidden_states normed = self.self_attn_norm(hidden_states) self_attn_out, new_self_kv = self.self_attn( normed, position_ids, rope, self_kv_cache, self_attn_mask ) hidden_states = residual + self_attn_out residual = hidden_states normed = self.cross_attn_norm(hidden_states) cross_out = self.cross_attn(normed, shared_kv, position_ids, rope, cross_attn_mask) hidden_states = residual + cross_out if self.num_persistent_contexts >= 2 and m3_kv is not None: residual = hidden_states normed = self.cross_attn_m3_norm(hidden_states) m3_out = self.cross_attn_m3(normed, m3_kv, position_ids, rope, m3_attn_mask) hidden_states = residual + m3_out residual = hidden_states hidden_states = residual + self.ffn(self.ffn_norm(hidden_states)) return hidden_states, new_self_kv class SharonaModel(nn.Module): def __init__(self, config: SharonaConfig): super().__init__() self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.rope = RotaryEmbedding( config.head_dim, max(config.max_context1_length, config.max_context2_length), theta=config.rope_theta, scaling_factor=config.rope_scaling_factor, scaling_type=config.rope_scaling_type, low_freq_factor=config.rope_low_freq_factor, high_freq_factor=config.rope_high_freq_factor, original_max_position_embeddings=config.rope_original_max_position_embeddings, ) self.self_decoder = nn.ModuleList([ SelfDecoderLayer(config, layer_idx=i) for i in range(config.num_self_decoder_layers) ]) self.cross_decoder = nn.ModuleList([ CrossDecoderLayer(config, layer_idx=config.num_self_decoder_layers + i) for i in range(config.num_cross_decoder_layers) ]) self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.embed_tokens.weight self.apply(self._init_weights) self._init_residual_projections() def gradient_checkpointing_enable(self): self.config.use_gradient_checkpointing = True def gradient_checkpointing_disable(self): self.config.use_gradient_checkpointing = False def _init_weights(self, module: nn.Module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) def _init_residual_projections(self): factor = 1.0 / math.sqrt(2.0 * self.config.total_layers) for layer in self.self_decoder: nn.init.normal_(layer.attn.o_proj.weight, mean=0.0, std=self.config.initializer_range * factor) nn.init.normal_(layer.ffn.down_proj.weight, mean=0.0, std=self.config.initializer_range * factor) for layer in self.cross_decoder: nn.init.normal_(layer.self_attn.o_proj.weight, mean=0.0, std=self.config.initializer_range * factor) nn.init.normal_(layer.cross_attn.o_proj.weight, mean=0.0, std=self.config.initializer_range * factor) if hasattr(layer, 'cross_attn_m3'): nn.init.normal_(layer.cross_attn_m3.o_proj.weight, mean=0.0, std=self.config.initializer_range * factor) nn.init.normal_(layer.ffn.down_proj.weight, mean=0.0, std=self.config.initializer_range * factor) def _run_self_decoder( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, kv_cache: Optional[KVCache] = None, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, KVCache]: if kv_cache is None: kv_cache = KVCache(self.config.num_self_decoder_layers) for i, layer in enumerate(self.self_decoder): layer_kv = kv_cache.get(i) if self.config.use_gradient_checkpointing and self.training: hidden_states, new_kv = checkpoint( layer, hidden_states, position_ids, self.rope, layer_kv, attention_mask, use_reentrant=False, ) else: hidden_states, new_kv = layer( hidden_states, position_ids, self.rope, layer_kv, attention_mask ) kv_cache.set(i, new_kv) return hidden_states, kv_cache def _run_cross_decoder( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, shared_kv: KVState, kv_cache: Optional[KVCache] = None, self_attn_mask: Optional[torch.Tensor] = None, cross_attn_mask: Optional[torch.Tensor] = None, m3_kv: Optional[KVState] = None, m3_attn_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, KVCache]: if kv_cache is None: kv_cache = KVCache(self.config.num_cross_decoder_layers) for i, layer in enumerate(self.cross_decoder): cache_idx = self.config.num_self_decoder_layers + i layer_kv = kv_cache.get(cache_idx) if self.config.use_gradient_checkpointing and self.training: hidden_states, new_kv = checkpoint( layer, hidden_states, position_ids, self.rope, shared_kv, layer_kv, self_attn_mask, cross_attn_mask, m3_kv, m3_attn_mask, use_reentrant=False, ) else: hidden_states, new_kv = layer( hidden_states, position_ids, self.rope, shared_kv, layer_kv, self_attn_mask, cross_attn_mask, m3_kv=m3_kv, m3_attn_mask=m3_attn_mask, ) kv_cache.set(cache_idx, new_kv) return hidden_states, kv_cache def forward( self, context1_ids: torch.Tensor, context1_position_ids: Optional[torch.Tensor] = None, context2_ids: Optional[torch.Tensor] = None, context2_position_ids: Optional[torch.Tensor] = None, context2_kv_cache: Optional[KVCache] = None, cross_decoder_kv_cache: Optional[KVCache] = None, labels: Optional[torch.Tensor] = None, m1_ids: Optional[torch.Tensor] = None, m3_ids: Optional[torch.Tensor] = None, m1_kv_cache: Optional[KVCache] = None, m3_kv_cache: Optional[KVCache] = None, ) -> dict: bsz, ctx1_len = context1_ids.shape device = context1_ids.device triple_context = m1_ids is not None or m3_ids is not None if context1_position_ids is None: context1_position_ids = torch.arange(ctx1_len, device=device).unsqueeze(0).expand(bsz, -1) if triple_context: return self._forward_triple( context1_ids, context1_position_ids, m1_ids, m3_ids, m1_kv_cache, m3_kv_cache, cross_decoder_kv_cache, labels, ) if context2_ids is not None and context2_kv_cache is None: ctx2_len = context2_ids.shape[1] if context2_position_ids is None: context2_position_ids = torch.arange(ctx2_len, device=device).unsqueeze(0).expand(bsz, -1) ctx2_hidden = self.embed_tokens(context2_ids) ctx2_hidden, context2_kv_cache = self._run_self_decoder(ctx2_hidden, context2_position_ids) self_decoder_cache = None if context2_kv_cache is not None: self_decoder_cache = KVCache(self.config.num_self_decoder_layers) for i in range(self.config.num_self_decoder_layers): cached = context2_kv_cache.get(i) if cached is not None: self_decoder_cache.set(i, cached) ctx1_hidden = self.embed_tokens(context1_ids) ctx1_hidden, self_decoder_final = self._run_self_decoder( ctx1_hidden, context1_position_ids, kv_cache=self_decoder_cache ) shared_kv = self_decoder_final.get_shared_kv() if self.training and shared_kv is not None: total_kv_len = shared_kv.key.shape[2] ctx2_kv_len = total_kv_len - ctx1_len if ctx2_kv_len > 0: shared_kv = KVState( shared_kv.key[:, :, :ctx2_kv_len, :], shared_kv.value[:, :, :ctx2_kv_len, :], ) ctx1_hidden, cross_decoder_cache = self._run_cross_decoder( ctx1_hidden, context1_position_ids, shared_kv, kv_cache=cross_decoder_kv_cache, ) ctx1_hidden = self.final_norm(ctx1_hidden) logits = self.lm_head(ctx1_hidden) result = { "logits": logits, "context2_kv_cache": context2_kv_cache, "cross_decoder_kv_cache": cross_decoder_cache, } if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = nn.functional.cross_entropy( shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100, ) result["loss"] = loss return result def _forward_triple( self, m2_ids: torch.Tensor, m2_position_ids: torch.Tensor, m1_ids: Optional[torch.Tensor], m3_ids: Optional[torch.Tensor], m1_kv_cache: Optional[KVCache], m3_kv_cache: Optional[KVCache], cross_decoder_kv_cache: Optional[KVCache], labels: Optional[torch.Tensor], ) -> dict: bsz, m2_len = m2_ids.shape device = m2_ids.device if m1_ids is not None and m1_kv_cache is None: m1_len = m1_ids.shape[1] m1_pos = torch.arange(m1_len, device=device).unsqueeze(0).expand(bsz, -1) m1_hidden = self.embed_tokens(m1_ids) m1_hidden, m1_kv_cache = self._run_self_decoder(m1_hidden, m1_pos) if m3_ids is not None and m3_kv_cache is None: m3_len = m3_ids.shape[1] m1_offset = m1_kv_cache.get(0).key.shape[2] if m1_kv_cache is not None else (m1_ids.shape[1] if m1_ids is not None else 0) m3_pos = torch.arange(m1_offset, m1_offset + m3_len, device=device).unsqueeze(0).expand(bsz, -1) m3_hidden = self.embed_tokens(m3_ids) m3_hidden, m3_kv_cache = self._run_self_decoder(m3_hidden, m3_pos) self_decoder_cache = KVCache(self.config.num_self_decoder_layers) m1_kv_len = 0 m3_kv_len = 0 for i in range(self.config.num_self_decoder_layers): parts_k, parts_v = [], [] if m1_kv_cache is not None: m1_layer = m1_kv_cache.get(i) if m1_layer is not None: parts_k.append(m1_layer.key) parts_v.append(m1_layer.value) m1_kv_len = m1_layer.key.shape[2] if m3_kv_cache is not None: m3_layer = m3_kv_cache.get(i) if m3_layer is not None: parts_k.append(m3_layer.key) parts_v.append(m3_layer.value) m3_kv_len = m3_layer.key.shape[2] if parts_k: combined_k = torch.cat(parts_k, dim=2) combined_v = torch.cat(parts_v, dim=2) self_decoder_cache.set(i, KVState(combined_k, combined_v)) prefix_len = m1_kv_len + m3_kv_len m2_pos_offset = torch.arange(prefix_len, prefix_len + m2_len, device=device).unsqueeze(0).expand(bsz, -1) # Key-padding masks so M2 never attends to <|pad|> positions in M1/M3. pad_id = self.config.pad_token_id dtype = self.embed_tokens.weight.dtype neg_inf = torch.finfo(dtype).min m1_pad_bias = None m3_pad_bias = None if m1_ids is not None and m1_kv_len > 0: m1_valid = (m1_ids != pad_id) # (bsz, m1_len) # All-pad rows would make SDPA softmax NaN; let them attend freely. m1_valid = m1_valid | ~m1_valid.any(dim=1, keepdim=True) m1_pad_bias = torch.where(m1_valid, 0.0, neg_inf).to(dtype) if m3_ids is not None and m3_kv_len > 0: m3_valid = (m3_ids != pad_id) m3_valid = m3_valid | ~m3_valid.any(dim=1, keepdim=True) m3_pad_bias = torch.where(m3_valid, 0.0, neg_inf).to(dtype) # Self-decoder mask for M2 over [M1 | M3 | M2]: bottom-right causal # plus pad bias on the prefix segment. self_attn_mask = None if prefix_len > 0: total_len = prefix_len + m2_len q_idx = torch.arange(prefix_len, total_len, device=device).unsqueeze(1) kv_idx = torch.arange(total_len, device=device).unsqueeze(0) causal = (kv_idx <= q_idx) self_attn_mask = torch.where(causal, 0.0, neg_inf).to(dtype) self_attn_mask = self_attn_mask.unsqueeze(0).unsqueeze(0).expand(bsz, 1, m2_len, total_len).clone() prefix_parts = [] if m1_pad_bias is not None: prefix_parts.append(m1_pad_bias) if m3_pad_bias is not None: prefix_parts.append(m3_pad_bias) if prefix_parts: prefix_bias = torch.cat(prefix_parts, dim=1) # (bsz, prefix_len) self_attn_mask[:, :, :, :prefix_len] += prefix_bias[:, None, None, :] m2_hidden = self.embed_tokens(m2_ids) m2_hidden, self_decoder_final = self._run_self_decoder( m2_hidden, m2_pos_offset, kv_cache=self_decoder_cache, attention_mask=self_attn_mask, ) m1_shared_kv = None if m1_kv_cache is not None: last_layer = m1_kv_cache.get_shared_kv() if last_layer is not None: m1_shared_kv = last_layer m3_shared_kv = None if m3_kv_cache is not None: last_layer = m3_kv_cache.get_shared_kv() if last_layer is not None: m3_shared_kv = last_layer if m1_shared_kv is None: sd_final = self_decoder_final.get_shared_kv() m1_shared_kv = sd_final cross_attn_mask = m1_pad_bias[:, None, None, :] if m1_pad_bias is not None else None m3_attn_mask = m3_pad_bias[:, None, None, :] if m3_pad_bias is not None else None cross_pos = m2_pos_offset # v4 fix: use real positions for consistent RoPE m2_hidden, cross_decoder_cache = self._run_cross_decoder( m2_hidden, cross_pos, m1_shared_kv, kv_cache=cross_decoder_kv_cache, cross_attn_mask=cross_attn_mask, m3_kv=m3_shared_kv, m3_attn_mask=m3_attn_mask, ) m2_hidden = self.final_norm(m2_hidden) logits = self.lm_head(m2_hidden) result = { "logits": logits, "m1_kv_cache": m1_kv_cache, "m3_kv_cache": m3_kv_cache, "cross_decoder_kv_cache": cross_decoder_cache, } if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = nn.functional.cross_entropy( shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100, ) result["loss"] = loss return result @torch.no_grad() def precompute_context2( self, context2_ids: torch.Tensor, position_ids: Optional[torch.Tensor] = None, chunk_size: int = 4096, ) -> KVCache: bsz, total_len = context2_ids.shape device = context2_ids.device if position_ids is None: position_ids = torch.arange(total_len, device=device).unsqueeze(0).expand(bsz, -1) kv_cache = KVCache(self.config.num_self_decoder_layers) for start in range(0, total_len, chunk_size): end = min(start + chunk_size, total_len) chunk_ids = context2_ids[:, start:end] chunk_pos = position_ids[:, start:end] chunk_hidden = self.embed_tokens(chunk_ids) chunk_hidden, kv_cache = self._run_self_decoder( chunk_hidden, chunk_pos, kv_cache=kv_cache ) return kv_cache @torch.no_grad() def generate( self, context1_ids: torch.Tensor, context2_kv_cache: Optional[KVCache] = None, max_new_tokens: int = 256, temperature: float = 1.0, top_k: int = 50, top_p: float = 0.9, eos_token_id: Optional[int] = None, ) -> torch.Tensor: bsz = context1_ids.shape[0] device = context1_ids.device generated = context1_ids.clone() self_decoder_cache = None if context2_kv_cache is not None: self_decoder_cache = KVCache(self.config.num_self_decoder_layers) for i in range(self.config.num_self_decoder_layers): cached = context2_kv_cache.get(i) if cached is not None: self_decoder_cache.set(i, cached) cross_decoder_cache = None ctx1_len = context1_ids.shape[1] position_ids = torch.arange(ctx1_len, device=device).unsqueeze(0).expand(bsz, -1) ctx1_hidden = self.embed_tokens(context1_ids) ctx1_hidden, self_decoder_cache = self._run_self_decoder( ctx1_hidden, position_ids, kv_cache=self_decoder_cache ) shared_kv = self_decoder_cache.get_shared_kv() ctx1_hidden, cross_decoder_cache = self._run_cross_decoder( ctx1_hidden, position_ids, shared_kv, kv_cache=cross_decoder_cache ) ctx1_hidden = self.final_norm(ctx1_hidden) logits = self.lm_head(ctx1_hidden[:, -1:, :]) for step in range(max_new_tokens): next_logits = logits[:, -1, :] if temperature != 1.0: next_logits = next_logits / temperature if top_k > 0: top_k_vals, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1))) threshold = top_k_vals[:, -1].unsqueeze(-1) next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf")) if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(next_logits, descending=True) cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) remove_mask = cumulative_probs - sorted_logits.softmax(dim=-1) >= top_p sorted_logits[remove_mask] = float("-inf") next_logits = sorted_logits.scatter(1, sorted_indices, sorted_logits) probs = next_logits.softmax(dim=-1) next_token = torch.multinomial(probs, num_samples=1) generated = torch.cat([generated, next_token], dim=1) if eos_token_id is not None and (next_token == eos_token_id).all(): break cur_pos = generated.shape[1] - 1 step_pos = torch.tensor([[cur_pos]], device=device).expand(bsz, -1) step_hidden = self.embed_tokens(next_token) step_hidden, self_decoder_cache = self._run_self_decoder( step_hidden, step_pos, kv_cache=self_decoder_cache ) shared_kv = self_decoder_cache.get_shared_kv() step_hidden, cross_decoder_cache = self._run_cross_decoder( step_hidden, step_pos, shared_kv, kv_cache=cross_decoder_cache ) step_hidden = self.final_norm(step_hidden) logits = self.lm_head(step_hidden) return generated def save_pretrained(self, save_dir: str): save_path = Path(save_dir) save_path.mkdir(parents=True, exist_ok=True) torch.save(self.state_dict(), save_path / "model.pt") config_dict = {k: v for k, v in self.config.__dict__.items() if not k.startswith("_")} with open(save_path / "config.json", "w") as f: json.dump(config_dict, f, indent=2) @classmethod def from_pretrained(cls, load_dir: str, device: str = "cpu"): load_path = Path(load_dir) with open(load_path / "config.json") as f: config_dict = json.load(f) config = SharonaConfig(**config_dict) model = cls(config) state_dict = torch.load(load_path / "model.pt", map_location=device, weights_only=False) model.load_state_dict(state_dict) return model