| """Text-only ModilifyMk2 model with recurrent latent deliberation.""" |
|
|
| from __future__ import annotations |
|
|
| import copy |
| from collections.abc import Sequence |
| from dataclasses import dataclass, replace |
| from typing import Any |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.checkpoint import checkpoint |
| from transformers.cache_utils import Cache |
| from transformers.modeling_outputs import BaseModelOutputWithPast |
|
|
| from transformers.masking_utils import ( |
| ALL_MASK_ATTENTION_FUNCTIONS, |
| bidirectional_mask_function, |
| ) |
| from transformers.models.diffusion_gemma import ( |
| DiffusionGemmaDecoderModel, |
| DiffusionGemmaEncoderTextModel, |
| DiffusionGemmaPreTrainedModel, |
| ) |
| from transformers.models.diffusion_gemma.modeling_diffusion_gemma import ( |
| DiffusionGemmaRMSNorm, |
| DiffusionGemmaTextRouter, |
| ) |
|
|
| from .mps_ops import register_mps_backends |
| from .adapters import LoRAConfig, inject_lora |
| from .configuration_modilify_mk2 import ModilifyMk2Config, DENOISE_TEMPERATURE |
| from .generation_modilify_mk2 import ModilifyMk2GenerationConfig, ModilifyMk2GenerationMixin |
| from .latent_deliberation import LatentDeliberationState, LatentProcessorOutput, _sdpa_mask_value |
| from .latent_deliberation import LatentDeliberationTransformer |
| from .vocab_ops import ( |
| chunked_vocab_statistics, |
| chunked_vocab_training_statistics, |
| ) |
|
|
|
|
| register_mps_backends() |
|
|
|
|
| @dataclass |
| class ModilifyMk2DecoderOutput(BaseModelOutputWithPast): |
| token_embeddings: torch.FloatTensor | None = None |
| latent_residual_diagnostics: dict[str, torch.Tensor] | None = None |
|
|
|
|
| @dataclass |
| class ModilifyMk2ModelOutput(BaseModelOutputWithPast): |
| token_embeddings: torch.FloatTensor | None = None |
| encoder_last_hidden_state: torch.FloatTensor | None = None |
| latent_residual_diagnostics: dict[str, torch.Tensor] | None = None |
|
|
|
|
| @dataclass |
| class ModilifyMk2BlockDiffusionOutput: |
| logits: torch.FloatTensor | None |
| heavy_hidden_state: torch.FloatTensor |
| next_latent_state: LatentDeliberationState |
| gold_nll: torch.FloatTensor | None = None |
| loss: torch.FloatTensor | None = None |
| past_key_values: Cache | None = None |
| hidden_states: tuple[torch.FloatTensor, ...] | None = None |
| attentions: tuple[torch.FloatTensor, ...] | None = None |
| encoder_last_hidden_state: torch.FloatTensor | None = None |
| temporal_context: torch.FloatTensor | None = None |
| heavy_context_leaf: torch.FloatTensor | None = None |
| temporal_memory_slots: torch.FloatTensor | None = None |
| heavy_memory_leaf: torch.FloatTensor | None = None |
| working_state: torch.FloatTensor | None = None |
| latent_residual_diagnostics: dict[str, torch.Tensor] | None = None |
| temperature_gold_nll: torch.FloatTensor | None = None |
| proposal: torch.LongTensor | None = None |
| proposal_confidence: torch.FloatTensor | None = None |
| sampled_nll: torch.FloatTensor | None = None |
| token_entropy: torch.FloatTensor | None = None |
| greedy_proposal: torch.LongTensor | None = None |
| greedy_confidence: torch.FloatTensor | None = None |
|
|
| def to_tuple(self) -> tuple[Any, ...]: |
| return tuple(value for value in ( |
| self.loss, self.logits, self.past_key_values, self.hidden_states, |
| self.attentions, self.encoder_last_hidden_state, |
| self.heavy_hidden_state, self.next_latent_state, |
| ) if value is not None) |
|
|
| def __getitem__(self, key: str | int) -> Any: |
| return getattr(self, key) if isinstance(key, str) else self.to_tuple()[key] |
|
|
|
|
| class ModilifyMk2TextRouter(DiffusionGemmaTextRouter): |
| """Official router parameters with fp32 softmax so top-k weights stay finite.""" |
|
|
| def __init__(self, config: Any) -> None: |
| super().__init__(config) |
| self.norm = DiffusionGemmaRMSNorm(self.hidden_size, eps=self.eps, with_scale=False) |
|
|
| def forward( |
| self, hidden_states: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| hidden_states = self.norm(hidden_states) |
| hidden_states = hidden_states * self.scale * self.scalar_root_size |
| expert_scores = self.proj(hidden_states) |
| |
| |
| router_probabilities = F.softmax(expert_scores, dim=-1, dtype=torch.float32) |
| top_k_weights, top_k_index = torch.topk( |
| router_probabilities, |
| k=self.config.top_k_experts, |
| dim=-1, |
| ) |
| top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True).clamp_min( |
| torch.finfo(torch.float32).tiny |
| ) |
| top_k_weights = top_k_weights * self.per_expert_scale[top_k_index] |
| return router_probabilities, top_k_weights, top_k_index |
|
|
|
|
| def _finite_decoder_attention_mask( |
| mask: torch.Tensor | None, |
| *, |
| dtype: torch.dtype, |
| ) -> torch.Tensor | None: |
| """Keep SDPA masks boolean and unmask fully-blocked query rows. |
| |
| Transformers' SDPA mask is True=attend. A fully-False query row makes |
| softmax(-inf) in the backward. Converted additive 8k masks also double |
| decoder workspace on long prefixes, so stay in bool and only apply the |
| torch<2.5 empty-row workaround (pytorch/pytorch#110213). |
| """ |
|
|
| del dtype |
| if mask is None: |
| return None |
| if mask.dtype == torch.bool: |
| return mask | ~mask.any(dim=-1, keepdim=True) |
| return mask.clamp_min(_sdpa_mask_value(mask.dtype)) |
|
|
|
|
| def install_modilify_mk2_trunk_semantics(module: nn.Module) -> None: |
| """Swap official leaf modules on this instance. Never patch Transformers classes.""" |
|
|
| for name, child in list(module.named_children()): |
| if type(child) is DiffusionGemmaTextRouter: |
| replacement = ModilifyMk2TextRouter(child.config) |
| replacement.load_state_dict(child.state_dict()) |
| setattr(module, name, replacement) |
| else: |
| install_modilify_mk2_trunk_semantics(child) |
|
|
|
|
| class ModilifyMk2EncoderModel(DiffusionGemmaPreTrainedModel): |
| """Checkpoint-compatible wrapper around the text encoder trunk.""" |
|
|
| config_class = ModilifyMk2Config |
|
|
| def __init__(self, config: ModilifyMk2Config): |
| super().__init__(config) |
| self.language_model = DiffusionGemmaEncoderTextModel(config.text_config) |
| install_modilify_mk2_trunk_semantics(self) |
| self.post_init() |
|
|
| def get_input_embeddings(self) -> nn.Module: |
| return self.language_model.embed_tokens |
|
|
| def set_input_embeddings(self, value: nn.Module) -> None: |
| self.language_model.embed_tokens = value |
|
|
| def forward(self, **kwargs: Any) -> BaseModelOutputWithPast: |
| return self.language_model(**kwargs) |
|
|
|
|
| class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel): |
| """Diffusion decoder accepting compact self-conditioning embeddings.""" |
|
|
| config_class = ModilifyMk2Config |
| latent_residual_rms_ratio_cap = 0.5 |
|
|
| def __init__(self, config: ModilifyMk2Config): |
| super().__init__(config) |
| self.gradient_checkpointing = False |
| install_modilify_mk2_trunk_semantics(self) |
|
|
| @staticmethod |
| def create_diffusion_decoder_attention_mask( |
| config: Any, |
| inputs_embeds: torch.Tensor, |
| past_key_values: Cache, |
| decoder_attention_mask: torch.Tensor | dict | None = None, |
| ) -> dict[str, torch.Tensor | None]: |
| """Official mask builder without the all-True sliding-window skip.""" |
|
|
| if past_key_values is None: |
| raise ValueError( |
| "The diffusion mask requires `past_key_values` to construct the next attention mask correctly" |
| ) |
| if ( |
| decoder_attention_mask is None |
| or config._attn_implementation |
| not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping |
| ): |
| return {"full_attention": None, "sliding_attention": None} |
| if isinstance(decoder_attention_mask, dict) and all( |
| mask.ndim == 4 for mask in decoder_attention_mask.values() |
| ): |
| return { |
| name: _finite_decoder_attention_mask( |
| mask, dtype=inputs_embeds.dtype |
| ) |
| for name, mask in decoder_attention_mask.items() |
| } |
|
|
| text_config = config.get_text_config() if hasattr(config, "get_text_config") else config |
| q_length = inputs_embeds.shape[1] |
| q_offset = past_key_values.get_seq_length() |
| if isinstance(q_offset, torch.Tensor): |
| q_offset = q_offset.to(inputs_embeds.device) |
| additional_kv_length = ( |
| getattr(config, "canvas_length", 0) if past_key_values.is_compileable else 0 |
| ) |
| mask_mapping: dict[str, torch.Tensor | None] = {} |
| for layer_pattern in set(text_config.layer_types): |
| layer_idx = past_key_values.is_sliding.index( |
| layer_pattern == "sliding_attention" |
| ) |
| kv_length, kv_offset = past_key_values.get_mask_sizes(q_length, layer_idx) |
| kv_length += additional_kv_length |
| if layer_pattern == "sliding_attention" and past_key_values.is_compileable: |
| sliding_layer = past_key_values.layers[layer_idx] |
| max_length = sliding_layer.get_max_length() + additional_kv_length |
| if kv_length >= max_length: |
| kv_length = max_length |
| mask_mapping[layer_pattern] = _finite_decoder_attention_mask( |
| ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]( |
| batch_size=inputs_embeds.shape[0], |
| q_length=q_length, |
| kv_length=kv_length, |
| q_offset=q_offset, |
| kv_offset=kv_offset, |
| mask_function=bidirectional_mask_function, |
| attention_mask=decoder_attention_mask, |
| allow_is_causal_skip=False, |
| allow_is_bidirectional_skip=True, |
| local_size=getattr(text_config, "sliding_window", None), |
| dtype=inputs_embeds.dtype, |
| config=text_config, |
| use_vmap=False, |
| device=inputs_embeds.device, |
| ), |
| dtype=inputs_embeds.dtype, |
| ) |
| return mask_mapping |
|
|
| def merge_latent_context( |
| self, |
| token_embeddings: torch.Tensor, |
| latent_context: torch.Tensor | None, |
| *, |
| collect_diagnostics: bool = True, |
| ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| """Map latent context through the frozen native self-conditioning bridge.""" |
|
|
| if latent_context is None: |
| context = torch.zeros_like(token_embeddings) |
| else: |
| if latent_context.shape != token_embeddings.shape: |
| raise ValueError("Temporal context must match canvas hidden-state shape.") |
| context = latent_context.to(token_embeddings) |
| mapper = self.self_conditioning |
| normalized_context = mapper.pre_norm(context) |
| mapped_context = mapper.down_proj( |
| mapper.act_fn(mapper.gate_proj(normalized_context)) |
| * mapper.up_proj(normalized_context) |
| ) |
| mapped_fp32 = mapped_context.float() |
| |
| |
| |
| |
| mapped_token_energy = mapped_fp32.square().mean(dim=-1, keepdim=True) |
| token_rms_per_token = ( |
| token_embeddings.float().square().mean(dim=-1, keepdim=True).sqrt() |
| ) |
| residual_cap = self.latent_residual_rms_ratio_cap * token_rms_per_token |
| soft_cap_scale = residual_cap / torch.sqrt( |
| mapped_token_energy + residual_cap.square() + 1.0e-12 |
| ) |
| mapped_context = (mapped_fp32 * soft_cap_scale).to(dtype=mapped_context.dtype) |
| combined = token_embeddings + mapped_context |
| diagnostics: dict[str, torch.Tensor] = {} |
| if collect_diagnostics: |
| token_rms = token_embeddings.detach().float().square().mean().sqrt() |
| mapped_rms = mapped_context.detach().float().square().mean().sqrt() |
| diagnostics = { |
| "token_embedding_rms": token_rms, |
| "latent_context_rms": context.detach().float().square().mean().sqrt(), |
| "mapped_sc_rms": mapped_rms, |
| "actual_residual_rms": mapped_rms, |
| "latent_to_embedding_rms_ratio": ( |
| mapped_rms / token_rms.clamp_min(1.0e-12) |
| ), |
| "merged_input_rms": combined.detach().float().square().mean().sqrt(), |
| } |
| return mapper.post_norm(combined), diagnostics |
|
|
| def _run_stack( |
| self, |
| inputs_embeds: torch.Tensor, |
| *, |
| past_key_values: Cache | None, |
| decoder_attention_mask: torch.Tensor | dict | None, |
| decoder_position_ids: torch.LongTensor | None, |
| **kwargs: Any, |
| ) -> torch.Tensor: |
| if decoder_position_ids is None: |
| prefix = past_key_values.get_seq_length(layer_idx=0) if past_key_values is not None else 0 |
| decoder_position_ids = torch.arange( |
| prefix, prefix + inputs_embeds.shape[1], device=inputs_embeds.device |
| ).unsqueeze(0) |
| if not isinstance(mask_mapping := decoder_attention_mask, dict): |
| mask_mapping = self.create_diffusion_decoder_attention_mask( |
| config=self.text_config, |
| inputs_embeds=inputs_embeds, |
| past_key_values=past_key_values, |
| decoder_attention_mask=decoder_attention_mask, |
| ) |
| working_bus = kwargs.pop("working_bus", None) |
| working_state = kwargs.pop("working_state", None) |
| working_canvas_head = kwargs.pop("working_canvas_head", None) |
| persistent_bus = kwargs.pop("persistent_bus", None) or kwargs.pop("memory_bus", None) |
| memory_slots = kwargs.pop("memory_slots", None) |
| slot_identity = kwargs.pop("slot_identity", None) |
| cache = past_key_values |
| if self.gradient_checkpointing and self.training and cache is not None: |
| cache = copy.copy(cache) |
| cache.layers = [copy.copy(layer) for layer in cache.layers] |
| hidden = inputs_embeds |
| positions = { |
| layer_type: self.rotary_emb(hidden, decoder_position_ids, layer_type) |
| for layer_type in self.unique_layer_types |
| } |
| working_kv = None |
| if working_bus is not None and working_state is not None: |
| working_kv = working_bus.prepare_kv(working_state) |
| memory_kv = None |
| if persistent_bus is not None and memory_slots is not None: |
| memory_kv = persistent_bus.prepare_kv(memory_slots, slot_identity) |
| working_reader = 0 |
| reader_index = 0 |
| for index in range(self.text_config.num_hidden_layers): |
| layer = self.layers[index] |
| layer_type = self.text_config.layer_types[index] |
| if self.gradient_checkpointing and self.training: |
| def layer_forward( |
| values: torch.Tensor, |
| selected: nn.Module = layer, |
| selected_type: str = layer_type, |
| ) -> torch.Tensor: |
| return selected( |
| values, |
| position_embeddings=positions[selected_type], |
| attention_mask=mask_mapping[selected_type], |
| position_ids=decoder_position_ids, |
| past_key_values=cache, |
| **kwargs, |
| ) |
|
|
| hidden = self._gradient_checkpointing_func(layer_forward, hidden) |
| else: |
| hidden = layer( |
| hidden, |
| position_embeddings=positions[layer_type], |
| attention_mask=mask_mapping[layer_type], |
| position_ids=decoder_position_ids, |
| past_key_values=cache, |
| **kwargs, |
| ) |
| if layer_type == "full_attention": |
| if ( |
| working_kv is not None |
| and working_reader < working_bus.num_readers |
| ): |
| hidden = working_bus.read( |
| hidden, working_reader, *working_kv, |
| positions=( |
| decoder_position_ids if working_canvas_head is not None else None |
| ), |
| ) |
| working_reader += 1 |
| if ( |
| memory_kv is not None |
| and reader_index < persistent_bus.num_readers |
| ): |
| hidden = persistent_bus.read( |
| hidden, reader_index, memory_kv[0], memory_kv[1] |
| ) |
| reader_index += 1 |
| return self.norm(hidden) |
|
|
| def forward( |
| self, |
| decoder_input_ids: torch.LongTensor, |
| past_key_values: Cache | None = None, |
| decoder_token_embeddings: torch.FloatTensor | None = None, |
| temporal_context_embeddings: torch.FloatTensor | None = None, |
| decoder_attention_mask: torch.Tensor | dict | None = None, |
| decoder_position_ids: torch.LongTensor | None = None, |
| collect_latent_diagnostics: bool = True, |
| memory_slots: torch.Tensor | None = None, |
| working_bus: Any | None = None, |
| working_state: torch.Tensor | None = None, |
| persistent_bus: Any | None = None, |
| slot_identity: torch.Tensor | None = None, |
| **kwargs: Any, |
| ) -> ModilifyMk2DecoderOutput: |
| if "use_cache" in kwargs: |
| raise ValueError("The diffusion decoder always reads the supplied cache.") |
| if decoder_token_embeddings is None: |
| token_embeddings = self.embed_tokens(decoder_input_ids) |
| else: |
| token_embeddings = decoder_token_embeddings |
| expected = (*decoder_input_ids.shape, self.text_config.hidden_size) |
| if token_embeddings.shape != expected: |
| raise ValueError("Precomputed decoder embeddings have the wrong shape.") |
| context_embeddings = ( |
| torch.zeros_like(token_embeddings) |
| if temporal_context_embeddings is None |
| else temporal_context_embeddings.to(token_embeddings) |
| ) |
| inputs_embeds, diagnostics = self.merge_latent_context( |
| token_embeddings, |
| context_embeddings if temporal_context_embeddings is not None else None, |
| collect_diagnostics=collect_latent_diagnostics, |
| ) |
| hidden = self._run_stack( |
| inputs_embeds, |
| past_key_values=past_key_values, |
| decoder_attention_mask=decoder_attention_mask, |
| decoder_position_ids=decoder_position_ids, |
| memory_slots=memory_slots, |
| working_bus=working_bus, |
| working_state=working_state, |
| persistent_bus=persistent_bus, |
| slot_identity=slot_identity, |
| **kwargs, |
| ) |
| return ModilifyMk2DecoderOutput( |
| last_hidden_state=hidden, |
| past_key_values=past_key_values, |
| token_embeddings=token_embeddings, |
| latent_residual_diagnostics=diagnostics, |
| ) |
|
|
|
|
| class ModilifyMk2Model(DiffusionGemmaPreTrainedModel): |
| config_class = ModilifyMk2Config |
| _tied_weights_keys = { |
| "encoder.language_model.norm.weight": "decoder.norm.weight", |
| r"encoder.language_model.layers\.(?:[^.]+\.)*weight": r"decoder.layers\.(?:[^.]+\.)*weight", |
| r"encoder.language_model.layers\.(?:[^.]+\.)*scale": r"decoder.layers\.(?:[^.]+\.)*scale", |
| r"encoder.language_model.layers\.(?:[^.]+\.)*per_expert_scale": r"decoder.layers\.(?:[^.]+\.)*per_expert_scale", |
| r"encoder.language_model.layers\.(?:[^.]+\.)*gate_up_proj": r"decoder.layers\.(?:[^.]+\.)*gate_up_proj", |
| r"encoder.language_model.layers\.(?:[^.]+\.)*down_proj": r"decoder.layers\.(?:[^.]+\.)*down_proj", |
| "encoder.language_model.embed_tokens.weight": "decoder.embed_tokens.weight", |
| } |
|
|
| def __init__(self, config: ModilifyMk2Config): |
| super().__init__(config) |
| self.encoder = ModilifyMk2EncoderModel(config) |
| self.decoder = ModilifyMk2DecoderModel(config) |
| install_modilify_mk2_trunk_semantics(self) |
| self.post_init() |
|
|
| def get_encoder(self): |
| return self.encoder |
|
|
| def get_decoder(self): |
| return self.decoder |
|
|
| def get_input_embeddings(self): |
| return self.encoder.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value): |
| self.encoder.set_input_embeddings(value) |
| self.decoder.embed_tokens = value |
|
|
| def forward( |
| self, |
| *, |
| input_ids: torch.LongTensor | None = None, |
| attention_mask: torch.Tensor | dict | None = None, |
| past_key_values: Cache | None = None, |
| position_ids: torch.LongTensor | None = None, |
| decoder_input_ids: torch.LongTensor, |
| decoder_token_embeddings: torch.FloatTensor | None = None, |
| temporal_context_embeddings: torch.FloatTensor | None = None, |
| decoder_attention_mask: torch.Tensor | dict | None = None, |
| decoder_position_ids: torch.LongTensor | None = None, |
| return_encoder_outputs: bool = True, |
| collect_latent_diagnostics: bool = True, |
| memory_slots: torch.Tensor | None = None, |
| working_bus: Any | None = None, |
| working_state: torch.Tensor | None = None, |
| persistent_bus: Any | None = None, |
| slot_identity: torch.Tensor | None = None, |
| **kwargs: Any, |
| ) -> ModilifyMk2ModelOutput: |
| encoder_hidden = None |
| if input_ids is not None: |
| encoded = self.encoder( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| past_key_values=past_key_values, |
| position_ids=position_ids, |
| **kwargs, |
| ) |
| past_key_values = encoded.past_key_values |
| if return_encoder_outputs: |
| encoder_hidden = encoded.last_hidden_state |
| elif past_key_values is None: |
| raise ValueError("Either `input_ids` or `past_key_values` is required.") |
| decoded = self.decoder( |
| decoder_input_ids=decoder_input_ids, |
| decoder_token_embeddings=decoder_token_embeddings, |
| past_key_values=past_key_values, |
| temporal_context_embeddings=temporal_context_embeddings, |
| decoder_attention_mask=decoder_attention_mask, |
| decoder_position_ids=decoder_position_ids, |
| collect_latent_diagnostics=collect_latent_diagnostics, |
| memory_slots=memory_slots, |
| working_bus=working_bus, |
| working_state=working_state, |
| persistent_bus=persistent_bus, |
| slot_identity=slot_identity, |
| **kwargs, |
| ) |
| return ModilifyMk2ModelOutput( |
| last_hidden_state=decoded.last_hidden_state, |
| past_key_values=past_key_values, |
| token_embeddings=decoded.token_embeddings, |
| encoder_last_hidden_state=encoder_hidden, |
| latent_residual_diagnostics=decoded.latent_residual_diagnostics, |
| ) |
|
|
|
|
| class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2GenerationMixin): |
| config_class = ModilifyMk2Config |
| _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"} |
| generation_config_class = ModilifyMk2GenerationConfig |
| _no_split_modules = ["DiffusionGemmaDecoderLayer", "LatentDeliberationTransformer"] |
| supports_gradient_checkpointing = True |
|
|
| @torch.no_grad() |
| def _init_weights(self, module: nn.Module) -> None: |
| super()._init_weights(module) |
| if isinstance(module, LatentDeliberationTransformer): |
| module.reset_identity_parameters() |
| module.working_memory_bus.freeze() |
| module.persistent_memory_bus.freeze() |
|
|
| def __init__(self, config: ModilifyMk2Config): |
| super().__init__(config) |
| self.model = ModilifyMk2Model(config) |
| layer_types = tuple(getattr(config.text_config, "layer_types", None) or ()) |
| self.latent_deliberation = LatentDeliberationTransformer( |
| hidden_size=config.text_config.hidden_size, |
| latent_dim=config.latent_dim, |
| ffn_dim=config.latent_ffn_dim, |
| num_layers=config.latent_num_layers, |
| num_heads=config.latent_num_heads, |
| local_attention_window=config.latent_local_attention_window, |
| tape_probes=config.latent_tape_probes, |
| history_kv_rank=config.latent_history_kv_rank, |
| num_memory_readers=sum(layer_type == "full_attention" for layer_type in layer_types), |
| num_working_readers=( |
| sum(layer_type == "full_attention" for layer_type in layer_types) |
| if config.working_memory_bus else 0 |
| ), |
| num_persistent_readers=( |
| sum(layer_type == "full_attention" for layer_type in layer_types) |
| if config.persistent_memory_bus else 0 |
| ), |
| working_last_block_global=config.latent_working_last_block_global, |
| commit_sequence_dim=config.commit_sequence_dim, |
| max_canvas_length=config.canvas_length, |
| ) |
| self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) |
| self.final_logit_softcapping = config.text_config.final_logit_softcapping |
| self.post_init() |
| install_modilify_mk2_trunk_semantics(self) |
| if getattr(config, "lora_config", None): |
| payload = dict(config.lora_config) |
| payload["target_modules"] = tuple(payload["target_modules"]) |
| inject_lora(self, LoRAConfig(**payload)) |
| |
| |
| self._tied_weights_keys = dict(type(self)._tied_weights_keys) |
| self._tied_weights_keys.update({ |
| name: name.replace("model.encoder.language_model.", "model.decoder.", 1) |
| for name, _ in self.named_parameters(remove_duplicate=False) |
| if name.startswith("model.encoder.language_model.") and ".lora_" in name |
| }) |
| self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=True) |
| self.tie_weights(recompute_mapping=False) |
| config.text_config._experts_implementation = "modilify_mk2_lora" |
| |
| |
| self._keep_in_fp32_modules_strict = { |
| name for name, _ in self.named_parameters() |
| if name.startswith("latent_deliberation.") and ( |
| name.endswith((".dt_bias", ".a_log")) or |
| (name.endswith(".weight") and "norm" in name.rsplit(".", 2)[-2]) |
| ) |
| } |
|
|
| def _finalize_logits(self, hidden: torch.Tensor) -> torch.Tensor: |
| logits = self.lm_head(hidden) |
| return torch.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping |
|
|
| def _compute_gold_nll( |
| self, logits: torch.Tensor, labels: torch.LongTensor | None |
| ) -> torch.Tensor | None: |
| if labels is None: |
| return None |
| labels = labels.to(logits.device) |
| return F.cross_entropy( |
| logits.flatten(0, 1), labels.flatten(), |
| ignore_index=-100, reduction="none", |
| ).view_as(labels) |
|
|
| def _prepare_latent_context( |
| self, |
| decoder_input_ids: torch.LongTensor, |
| *, |
| |
| |
| confidence: torch.Tensor | None, |
| entropy: torch.Tensor | None, |
| age: torch.Tensor | None, |
| latent_state: LatentDeliberationState | None, |
| canvas_head: torch.Tensor | None, |
| ) -> tuple[torch.Tensor, LatentDeliberationState, torch.Tensor]: |
| batch, canvas = decoder_input_ids.shape |
| dtype = self.model.decoder.embed_tokens.weight.dtype |
| if latent_state is None: |
| latent_state = LatentDeliberationState.empty( |
| batch_size=batch, canvas_length=canvas, |
| latent_dim=self.config.latent_dim, memory_slots=self.config.latent_memory_slots, |
| device=decoder_input_ids.device, dtype=dtype, enable_gdn2=True, |
| ) |
| confidence = latent_state.confidence if confidence is None else confidence.squeeze(-1).float() |
| entropy = latent_state.entropy if entropy is None else entropy.squeeze(-1).float() |
| if age is not None: |
| latent_state = replace( |
| latent_state, |
| age=age.to(device=decoder_input_ids.device, dtype=torch.int32), |
| ) |
| token_embeddings = self.model.decoder.embed_tokens(decoder_input_ids) |
| use_outer_checkpoint = ( |
| self.training |
| and torch.is_grad_enabled() |
| and self.latent_deliberation.activation_checkpointing |
| ) |
| if use_outer_checkpoint: |
| def run_latent_processor( |
| embeddings: torch.Tensor, |
| memory_slots: torch.Tensor, |
| working_cells: torch.Tensor, |
| working_row: torch.Tensor, |
| state_confidence: torch.Tensor, |
| state_entropy: torch.Tensor, |
| state_age: torch.Tensor, |
| state_token_changed: torch.Tensor, |
| state_confidence_delta: torch.Tensor, |
| state_entropy_delta: torch.Tensor, |
| state_ponder_steps: torch.Tensor, |
| state_stagnation_steps: torch.Tensor, |
| |
| |
| ) -> torch.Tensor: |
| checkpoint_state = LatentDeliberationState( |
| memory_slots=memory_slots, |
| confidence=state_confidence, |
| entropy=state_entropy, |
| age=state_age, |
| token_changed=state_token_changed, |
| confidence_delta=state_confidence_delta, |
| entropy_delta=state_entropy_delta, |
| ponder_steps=state_ponder_steps, |
| stagnation_steps=state_stagnation_steps, |
| gdn2=replace( |
| latent_state.gdn2, cells=working_cells, |
| row=working_row, persistent=memory_slots, |
| ), |
| ) |
| result = self.latent_deliberation( |
| token_embeddings=embeddings, |
| confidence=confidence, |
| entropy=entropy, |
| state=checkpoint_state, |
|
|
|
|
| canvas_head=canvas_head, |
| ) |
| return result.context |
|
|
| latent_context = checkpoint( |
| run_latent_processor, |
| token_embeddings, |
| latent_state.memory_slots, |
| latent_state.gdn2.cells, |
| latent_state.gdn2.row, |
| latent_state.confidence, |
| latent_state.entropy, |
| latent_state.age, |
| latent_state.token_changed, |
| latent_state.confidence_delta, |
| latent_state.entropy_delta, |
| latent_state.ponder_steps, |
| latent_state.stagnation_steps, |
|
|
|
|
| use_reentrant=False, |
| preserve_rng_state=True, |
| ) |
| next_state = LatentDeliberationState( |
| memory_slots=latent_state.memory_slots, |
| confidence=confidence.to(dtype=torch.float32), |
| entropy=entropy.to(dtype=torch.float32), |
| age=latent_state.age, |
| token_changed=latent_state.token_changed, |
| confidence_delta=latent_state.confidence_delta, |
| entropy_delta=latent_state.entropy_delta, |
| ponder_steps=latent_state.ponder_steps, |
| stagnation_steps=latent_state.stagnation_steps, |
| gdn2=latent_state.gdn2, |
| ) |
| else: |
| processed: LatentProcessorOutput = self.latent_deliberation( |
| token_embeddings=token_embeddings, |
| confidence=confidence, |
| entropy=entropy, |
| state=latent_state, |
|
|
|
|
| canvas_head=canvas_head, |
| ) |
| latent_context = processed.context |
| next_state = processed.state |
| return ( |
| latent_context, |
| next_state, |
| token_embeddings, |
| ) |
|
|
| def forward( |
| self, |
| *, |
| input_ids: torch.LongTensor | None = None, |
| attention_mask: torch.Tensor | dict | None = None, |
| past_key_values: Cache | None = None, |
| position_ids: torch.LongTensor | None = None, |
| decoder_input_ids: torch.LongTensor, |
| previous_confidence: torch.FloatTensor | None = None, |
| previous_entropy: torch.FloatTensor | None = None, |
| token_age: torch.Tensor | None = None, |
| latent_state: LatentDeliberationState | None = None, |
| |
| |
| canvas_head: torch.Tensor | None = None, |
| |
| decoder_attention_mask: torch.Tensor | dict | None = None, |
| decoder_position_ids: torch.LongTensor | None = None, |
| labels: torch.LongTensor | None = None, |
| return_encoder_outputs: bool = True, |
| detach_temporal_context_for_heavy: bool = False, |
| compact_vocab: bool = False, |
| repetition_token_mask: torch.BoolTensor | None = None, |
| repetition_penalty: float = 1.0, |
| sampling_generators: Sequence[torch.Generator] | None = None, |
| collect_latent_diagnostics: bool = True, |
| **kwargs: Any, |
| ) -> ModilifyMk2BlockDiffusionOutput: |
| ( |
| latent_context, |
| next_state, |
| decoder_token_embeddings, |
| ) = self._prepare_latent_context( |
| decoder_input_ids, |
|
|
|
|
| confidence=previous_confidence, |
| entropy=previous_entropy, |
| age=token_age, |
| latent_state=latent_state, |
| canvas_head=canvas_head, |
| ) |
| heavy_context = latent_context |
| heavy_context_leaf = None |
| temporal_memory_slots = next_state.memory_slots |
| heavy_memory_slots = temporal_memory_slots |
| heavy_memory_leaf = None |
| working_bus = self.latent_deliberation.working_memory_bus |
| persistent_bus = self.latent_deliberation.persistent_memory_bus |
| working_readers_active = ( |
| working_bus.num_readers > 0 |
| and (not working_bus.training or working_bus.enabled) |
| ) |
| memory_readers_active = ( |
| persistent_bus.num_readers > 0 |
| and (not persistent_bus.training or persistent_bus.enabled) |
| ) |
| if detach_temporal_context_for_heavy and latent_context.requires_grad: |
| |
| |
| |
| heavy_context_leaf = latent_context.detach().clone().requires_grad_(True) |
| heavy_context = heavy_context_leaf |
| if ( |
| detach_temporal_context_for_heavy |
| and memory_readers_active |
| and temporal_memory_slots.requires_grad |
| ): |
| |
| |
| |
| |
| heavy_memory_leaf = ( |
| temporal_memory_slots.detach().clone().requires_grad_(True) |
| ) |
| heavy_memory_slots = heavy_memory_leaf |
| |
| |
| |
| reusable_decoder_embeddings = ( |
| None |
| if self.model.decoder.embed_tokens.weight.requires_grad |
| else decoder_token_embeddings |
| ) |
| outputs = self.model( |
| input_ids=input_ids, attention_mask=attention_mask, |
| past_key_values=past_key_values, position_ids=position_ids, |
| decoder_input_ids=decoder_input_ids, |
| decoder_token_embeddings=reusable_decoder_embeddings, |
| temporal_context_embeddings=heavy_context, |
| decoder_attention_mask=decoder_attention_mask, |
| decoder_position_ids=decoder_position_ids, |
| return_encoder_outputs=return_encoder_outputs, |
| collect_latent_diagnostics=collect_latent_diagnostics, |
| working_bus=working_bus if working_readers_active else None, |
| working_state=( |
| heavy_context, |
| |
| |
| next_state.gdn2.seen, |
| ) if working_readers_active else None, |
| working_canvas_head=canvas_head, |
| persistent_bus=persistent_bus if memory_readers_active else None, |
| memory_slots=heavy_memory_slots if memory_readers_active else None, |
| slot_identity=( |
| next_state.gdn2.seen |
| if memory_readers_active else None |
| ), |
| **kwargs, |
| ) |
| latent_residual_diagnostics = outputs.latent_residual_diagnostics |
| proposal = proposal_confidence = token_entropy = None |
| greedy_proposal = greedy_confidence = None |
| temperature_gold_nll = None |
| sampled_nll = None |
| if compact_vocab: |
| logits = None |
| if labels is None: |
| ( |
| proposal, |
| proposal_confidence, |
| token_entropy, |
| greedy_proposal, |
| greedy_confidence, |
| ) = ( |
| chunked_vocab_statistics( |
| outputs.last_hidden_state.detach(), |
| self.lm_head.weight.detach(), |
| softcap=self.final_logit_softcapping, |
| temperature=DENOISE_TEMPERATURE, |
| chunk_size=self.config.vocab_chunk_size, |
| repetition_token_mask=repetition_token_mask, |
| repetition_penalty=repetition_penalty, |
| sampling_generators=sampling_generators, |
| top_k=getattr(self.config, "commit_top_k", 40), |
| min_p=getattr(self.config, "commit_min_p", 0.05), |
| ) |
| ) |
| gold_nll = None |
| else: |
| ( |
| gold_nll, |
| temperature_gold_nll, |
| proposal, |
| proposal_confidence, |
| token_entropy, |
| greedy_proposal, |
| greedy_confidence, |
| sampled_nll, |
| ) = chunked_vocab_training_statistics( |
| outputs.last_hidden_state, |
| self.lm_head.weight, |
| labels, |
| return_sampled_nll=True, |
| softcap=self.final_logit_softcapping, |
| temperature=DENOISE_TEMPERATURE, |
| chunk_size=self.config.vocab_chunk_size, |
| top_k=getattr(self.config, "commit_top_k", 40), |
| min_p=getattr(self.config, "commit_min_p", 0.05), |
| ) |
| else: |
| logits = self._finalize_logits(outputs.last_hidden_state) |
| gold_nll = self._compute_gold_nll(logits, labels) |
| loss = None |
| if gold_nll is not None and labels is not None: |
| valid = labels.to(gold_nll.device).ne(-100) |
| loss = gold_nll.masked_fill(~valid, 0.0).sum() / valid.sum().clamp_min(1) |
| return ModilifyMk2BlockDiffusionOutput( |
| logits=logits, loss=loss, |
| past_key_values=outputs.past_key_values, |
| encoder_last_hidden_state=outputs.encoder_last_hidden_state, |
| heavy_hidden_state=outputs.last_hidden_state, |
| next_latent_state=next_state, |
| gold_nll=gold_nll, |
| temporal_context=latent_context, |
| heavy_context_leaf=heavy_context_leaf, |
| temporal_memory_slots=temporal_memory_slots, |
| heavy_memory_leaf=heavy_memory_leaf, |
|
|
| working_state=latent_context, |
| latent_residual_diagnostics=latent_residual_diagnostics, |
| temperature_gold_nll=temperature_gold_nll, |
| proposal=proposal, |
| proposal_confidence=proposal_confidence, |
| sampled_nll=sampled_nll, |
| token_entropy=token_entropy, |
| greedy_proposal=greedy_proposal, |
| greedy_confidence=greedy_confidence, |
| ) |
|
|
|
|
| __all__ = [ |
| "ModilifyMk2BlockDiffusionOutput", "ModilifyMk2Config", "ModilifyMk2DecoderModel", |
| "ModilifyMk2EncoderModel", "ModilifyMk2ForBlockDiffusion", "ModilifyMk2Model", |
| "ModilifyMk2TextRouter", "install_modilify_mk2_trunk_semantics", |
| ] |
|
|
