"""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) # Match the official DiffusionGemma router: bf16 softmax underflows to # all-zero top-k rows, then 0/0 NaNs the MoE residual and every loss. 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() # Use the energy directly in the cap denominator. Computing # sqrt(E[x²]) and immediately squaring it again has an undefined # backward at the identity-init point x=0 (0/0 in d(sqrt)/dx), which # poisoned the detached temporal-context VJP on every denoise. 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)) # The same adapters serve the encoder and decoder, just like the # tied backbone. Tell from_pretrained/save_pretrained about them. 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" # These 26 small parameters were trained and saved in FP32. Preserve # their dtype even when from_pretrained is called with BF16/FP16. 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: # Clone so the heavy decoder cannot alias the temporal graph storage. # `detach()` alone shares storage; an in-place decoder write would # corrupt the later latent backward and can produce non-finite grads. 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 ): # Memory readers are part of the heavy decoder too. Keep their # backward separate from the recurrent temporal graph for the same # reason as the canvas context above; both gradients are replayed # together at the TBPTT boundary. heavy_memory_leaf = ( temporal_memory_slots.detach().clone().requires_grad_(True) ) heavy_memory_slots = heavy_memory_leaf # The production LoRA path freezes the tied embedding table. Reuse its # lookup there; a trainable table needs two independent autograd paths # because heavy and temporal backward execute at different times. 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, # The Torch decoder keeps physical ring order and supplies # logical positions through RoPE; its mask must stay physical. 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", ]