"""Minimal modeling.py file for HF compatibility and funny zero-shot experiments. Best used for inference, finetuning should work, but is untested with this implementation.""" import torch import math from torch import Tensor from dataclasses import dataclass from typing import Optional, Union, Any from .raven_config_minimal import RavenConfig from transformers.cache_utils import Cache, DynamicCache ###################### Huggingface Glue code I ################################################################## from transformers import PreTrainedModel, GenerationMixin from transformers.utils import ModelOutput from transformers.generation.utils import GenerateDecoderOnlyOutput import torch.nn.functional as F from transformers import GenerationConfig class RavenPreTrainedModel(PreTrainedModel): config_class = RavenConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["SandwichBlock"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_quantized_cache = False _supports_static_cache = False def _init_weights(self, module): if not torch.rand((1,)).is_meta: print("Random Initialization not implemented.") @dataclass class CausalLMOutputRecurrentLatents(ModelOutput): loss: Optional[torch.Tensor] = None log_ppl: Optional[torch.Tensor] = None logits: Optional[torch.Tensor] = None past_key_values: Optional[Cache] = None latent_states: Optional[torch.Tensor] = None hidden_states: Optional[torch.Tensor] = None attention_maps: Optional[dict[int, torch.Tensor]] = None stats: Optional[dict] = None ###################### Minimal implementation from here ############################################################ class RMSNorm(torch.nn.Module): """Saner dtype handling and slightly better for fusion""" def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = torch.nn.Parameter(torch.ones(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): with torch.autocast(enabled=False, device_type=x.device.type): return self._norm(x.float()).type_as(x) * self.weight def reset_parameters(self) -> None: torch.nn.init.ones_(self.weight) class HuginnDynamicCache(DynamicCache): def __init__(self, lookup_strategy: str = "full") -> None: super().__init__() self._seen_tokens = 0 self.key_cache: dict[int, dict[int, torch.Tensor]] = {} self.value_cache: dict[int, dict[int, torch.Tensor]] = {} # structure: cache[index_of_layer_or_recurrent_step][index_in_sequence] # the cache is held uncoalesced because certain recurrent steps may be missing for some sequence ids if using # per-token adaptive compute. In those cases, the "lookup_strategy" determines how to proceed # Also, It is critical that the head indices do not overlap with the recurrent iteration indices self.lookup_strategy = lookup_strategy def update( self, key_states: torch.Tensor, value_states: torch.Tensor, step_idx: int, lookup_strategy: Optional[str] = None, ) -> tuple[torch.Tensor, torch.Tensor]: lookup_strategy = self.lookup_strategy if lookup_strategy is None else lookup_strategy if "compress-" in self.lookup_strategy and step_idx > 1: # hardcode for current model! compression_stage = int(self.lookup_strategy.split("compress-")[1][1:]) if "compress-s" in self.lookup_strategy: new_step_idx = (step_idx - 2) % compression_stage + 2 else: new_step_idx = (step_idx - 2) // compression_stage + 2 # @ print(step_idx, new_step_idx, compression_stage) step_idx = new_step_idx # Init if step_idx not in self.key_cache: self.key_cache[step_idx] = {} self.value_cache[step_idx] = {} # Update the number of seen tokens, we assume that step_idx=0 (first prelude) is always hit if step_idx == 0: self._seen_tokens += key_states.shape[-2] # Add entries to cache for idx, entry in enumerate(key_states.unbind(dim=-2)): if "compress-" not in self.lookup_strategy: assert step_idx < 0 or self._seen_tokens - key_states.shape[-2] + idx not in self.key_cache[step_idx] # print(f"Overwrote cache entry for step_idx {step_idx}") # likely the head self.key_cache[step_idx][self._seen_tokens - key_states.shape[-2] + idx] = entry for idx, entry in enumerate(value_states.unbind(dim=-2)): self.value_cache[step_idx][self._seen_tokens - value_states.shape[-2] + idx] = entry # Materialize past state based on lookup strategy: if len(self.key_cache[step_idx]) == self._seen_tokens or self.lookup_strategy == "full": # All entries are present, materialize cache as normal return ( torch.stack(list(self.key_cache[step_idx].values()), dim=-2), torch.stack(list(self.value_cache[step_idx].values()), dim=-2), ) else: # some entries where not previously computed # if lookup_strategy.startswith("latest"): # latest_keys = [] # latest_values = [] # for token_pos in range(self._seen_tokens): # # Find the latest step that has this token position # max_step = max((s for s in range(step_idx + 1) if token_pos in self.key_cache[s]), default=None) # if max_step is None: # raise ValueError(f"No cache entry found for token position {token_pos}") # latest_keys.append(self.key_cache[max_step][token_pos]) # latest_values.append(self.value_cache[max_step][token_pos]) # return torch.stack(latest_keys, dim=-2), torch.stack(latest_values, dim=-2) if lookup_strategy.startswith("latest-m4"): latest_keys = [] latest_values = [] for token_pos in range(self._seen_tokens): # For steps >= 2, use modulo 4 if step_idx >= 2: # Find valid steps for this token position valid_steps = [s for s in range(step_idx + 1) if token_pos in self.key_cache[s]] max_step = max([s for s in valid_steps if s >= 2 and s % 4 == step_idx % 4]) else: max_step = step_idx if token_pos in self.key_cache[step_idx] else 0 if max_step is None: raise ValueError(f"No cache entry found for token position {token_pos}") latest_keys.append(self.key_cache[max_step][token_pos]) latest_values.append(self.value_cache[max_step][token_pos]) return torch.stack(latest_keys, dim=-2), torch.stack(latest_values, dim=-2) elif lookup_strategy.startswith("skip"): existing_keys = [] existing_values = [] for token_pos in range(self._seen_tokens): if token_pos in self.key_cache[step_idx]: existing_keys.append(self.key_cache[step_idx][token_pos]) existing_values.append(self.value_cache[step_idx][token_pos]) return torch.stack(existing_keys, dim=-2), torch.stack(existing_values, dim=-2) elif lookup_strategy.startswith("randomized"): # sanity check rand_keys = [] rand_values = [] for token_pos in range(self._seen_tokens): if step_idx < 2: # For prelude steps max_step = step_idx if token_pos in self.key_cache[step_idx] else 0 else: # Get all steps from same block position curr_modulo = (step_idx - 2) % 4 + 2 valid_steps = [ s for s in range(2, step_idx + 1) if (s - 2) % 4 + 2 == curr_modulo and token_pos in self.key_cache[s] ] max_step = valid_steps[torch.randint(len(valid_steps), (1,))] rand_keys.append(self.key_cache[max_step][token_pos]) rand_values.append(self.value_cache[max_step][token_pos]) return torch.stack(rand_keys, dim=-2), torch.stack(rand_values, dim=-2) else: raise ValueError(f"Unknown lookup strategy: {lookup_strategy}") def reset(self) -> None: """Reset the cache state.""" self._seen_tokens = 0 self.key_cache.clear() self.value_cache.clear() def get_seq_length(self, step_idx: int = 0) -> int: return self._seen_tokens def get_memory_usage(self) -> float: total_bytes = 0 # For each recurrent step/layer index for step_idx in self.key_cache: # Get the sequence cache for this step key_seq_cache = self.key_cache[step_idx] for seq_idx in key_seq_cache: key_tensor = key_seq_cache[seq_idx] # Add memory for of key tensors, assuming value is the same total_bytes += key_tensor.nelement() * key_tensor.element_size() return total_bytes * 2 / (1024 * 1024) class CausalSelfAttention(torch.nn.Module): def __init__(self, config: RavenConfig) -> None: super().__init__() self.config = config self.n_head = config.num_attention_heads self.n_kv_heads = config.num_key_value_heads self.head_dim = config.n_embd // self.n_head shape = (self.n_head + 2 * self.n_kv_heads) * self.head_dim self.chunks = [config.n_embd, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim] self.Wqkv = torch.nn.Linear(config.n_embd, shape, bias=False) if config.qk_bias: self.qk_bias = torch.nn.Parameter(torch.zeros(2, 1, self.n_head, self.head_dim)) self.proj = torch.nn.Linear(config.n_embd, config.n_embd, bias=False) def forward( self, x: Tensor, freqs_cis: Tensor, step_idx: int, mask: Optional[Tensor] = None, past_key_values: Optional[Cache] = None, return_attn: bool = False, ) -> tuple[Tensor, Optional[Tensor]]: B, S, E = x.shape # batch size, sequence length, embedding dimensionality (n_embd) q, k, v = self.Wqkv(x).split(self.chunks, dim=2) q = q.view(B, S, self.n_head, self.head_dim) k = k.view(B, S, self.n_kv_heads, self.head_dim) v = v.view(B, S, self.n_kv_heads, self.head_dim) # bias? if self.config.qk_bias: q_bias, k_bias = self.qk_bias.split(1, dim=0) q, k = (q + q_bias).to(q.dtype), (k + k_bias).to(q.dtype) # apply rotary q, k = apply_rotary_emb_complex_like(q, k, freqs_cis=freqs_cis) q = q.transpose(1, 2) # (B, nh, S, hs) k = k.transpose(1, 2) v = v.transpose(1, 2) if past_key_values is not None: k, v = past_key_values.update(k, v, step_idx) if return_attn: y, attention_map = self.compute_eager_sdpa(q, k, v, attn_mask=mask) else: y = torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=q.shape[2] > 1 ) y = y.transpose(1, 2).reshape(B, S, E).contiguous() # reshape is a view if possible (it mostly is) return self.proj(y), attention_map if return_attn else None def compute_eager_sdpa(self, q, k, v, attn_mask): scale = 1.0 / math.sqrt(self.head_dim) scores = torch.matmul(q, k.transpose(-2, -1)) * scale if attn_mask is not None: scores = scores + attn_mask if q.shape[2] > 1: causal_mask = torch.triu(torch.ones(q.shape[2], q.shape[2]), diagonal=1).bool() scores.masked_fill_(causal_mask.to(scores.device), float("-inf")) attention_weights = torch.nn.functional.softmax(scores, dim=-1) y = torch.matmul(attention_weights, v) return y, attention_weights.max(dim=1)[0] class GatedMLP(torch.nn.Module): def __init__(self, config: RavenConfig, in_features: int = 0) -> None: super().__init__() in_features = config.n_embd if in_features == 0 else in_features self.fc = torch.nn.Linear(in_features, config.intermediate_size * 2, bias=False) self.proj = torch.nn.Linear(config.intermediate_size, config.n_embd, bias=False) self.nonlin = torch.nn.SiLU() def forward(self, x: Tensor) -> Tensor: # modified to single FC layer to improve parallelism x_fc_1, x_fc_2 = self.fc(x).chunk(2, dim=-1) x = self.nonlin(x_fc_1) * x_fc_2 return self.proj(x) class SandwichBlock(torch.nn.Module): expanded = False def __init__(self, config: RavenConfig, layer_id: int) -> None: super().__init__() self.norm_1 = RMSNorm(config.n_embd, eps=config.norm_eps) self.attn = CausalSelfAttention(config) self.norm_2 = RMSNorm(config.n_embd, eps=config.norm_eps) self.mlp = GatedMLP(config) self.norm_3 = RMSNorm(config.n_embd, eps=config.norm_eps) self.norm_4 = RMSNorm(config.n_embd, eps=config.norm_eps) self.layer_id = layer_id def forward( self, x: Tensor, freqs_cis: Tensor, step_idx: int, mask: Optional[Tensor] = None, past_key_values: Optional[Cache] = None, return_attn: bool = False, ) -> tuple[Tensor, Optional[Tensor]]: attn_out, attn_map = self.attn(self.norm_1(x), freqs_cis, step_idx, mask, past_key_values, return_attn) x = self.norm_2(attn_out + x) x = self.norm_4(self.mlp(self.norm_3(x)) + x) return x, attn_map class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin): def __init__( self, config: RavenConfig, ) -> None: super().__init__(config) self.config = config # Transformer layers prelude = torch.nn.ModuleList(SandwichBlock(config, layer_id=i) for i in range(config.n_layers_in_prelude)) adapter = torch.nn.Linear(config.n_embd * 2, config.n_embd, bias=config.bias) core_block = torch.nn.ModuleList( SandwichBlock(config, layer_id=i + config.n_layers_in_prelude) for i in range(config.n_layers_in_recurrent_block) ) o = config.n_layers_in_prelude + config.n_layers_in_recurrent_block * config.mean_recurrence coda = torch.nn.ModuleList(SandwichBlock(config, layer_id=i + o) for i in range(config.n_layers_in_coda)) self.transformer = torch.nn.ModuleDict( dict( wte=torch.nn.Embedding(config.padded_vocab_size, config.n_embd), prelude=prelude, adapter=adapter, core_block=core_block, coda=coda, ln_f=RMSNorm(config.n_embd, eps=config.norm_eps), # used twice :> ) ) self.emb_scale = config.init_values["embed_scale"] # Head self.lm_head = torch.nn.Linear(config.n_embd, config.padded_vocab_size, bias=False) if self.config.tie_embeddings: self.lm_head.weight = self.transformer.wte.weight # rope self.register_buffer("freqs_cis", self._precompute_freqs_cis(), persistent=True) def _precompute_freqs_cis(self): # can actually be a buffer now, and remains in fp32! (at least in the settings I tested) freqs_cis = precompute_freqs_cis( self.config.n_embd // self.config.num_attention_heads, self.config.block_size, self.config.rope_base, 1 ) return freqs_cis def forward( self, input_ids: torch.Tensor, input_embeds: Optional[torch.Tensor] = None, input_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, num_steps: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = None, output_details: dict = { "return_logits": True, "return_latents": True, "return_attention": False, "return_head": False, "return_stats": False, }, use_cache: bool = False, cache_position: Optional[torch.Tensor] = None, **kwargs, ) -> CausalLMOutputRecurrentLatents: # Support multiple position formats: if position_ids is None and cache_position is None: freqs_cis = self.freqs_cis[:, : input_ids.shape[1]] elif position_ids is not None: freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze()) elif cache_position is not None: freqs_cis = self.freqs_cis[:, cache_position] if input_embeds is None: input_embeds = self.transformer.wte(input_ids) if self.emb_scale != 1: input_embeds = input_embeds * self.emb_scale # type: ignore if use_cache and past_key_values is None: past_key_values = HuginnDynamicCache() attn_maps = {} return_attn = output_details["return_attention"] # Non-recurrent prelude for block_idx, block in enumerate(self.transformer.prelude): input_embeds, attn_map = block( input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn=return_attn ) attn_maps[block_idx] = attn_map # Main recurrence x, num_steps_no_grad, num_steps_with_grad, xk, block_idx, attn_maps = self.iterate_forward( input_embeds, # type: ignore input_states, freqs_cis, block_idx, attention_mask, past_key_values, num_steps, attn_maps, return_attn=return_attn, ) latent_states = x.clone().detach() # Coda layers for block_idx, block in enumerate(self.transformer.coda, start=1): x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values, return_attn=return_attn) attn_maps[-block_idx] = attn_map x = self.transformer.ln_f(x) # Prediction head, assuming labels really are labels and not equal to input_ids if labels is not None: logits = self.lm_head(x).float() loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.shape[-1]), labels.view(-1)) log_ppl = loss.clone().detach() else: logits = self.lm_head(x).float() loss, log_ppl = torch.as_tensor(0.0), torch.as_tensor(0.0) return CausalLMOutputRecurrentLatents( loss=loss, log_ppl=log_ppl, logits=logits if output_details["return_logits"] else None, past_key_values=past_key_values, hidden_states=x if output_details["return_head"] else None, latent_states=latent_states if output_details["return_latents"] else None, attention_maps=attn_maps if output_details["return_attention"] else None, # type: ignore stats=self.get_stats(logits, x, latent_states, xk, input_embeds, num_steps_no_grad, num_steps_with_grad) if output_details["return_stats"] else None, ) @torch._dynamo.disable(recursive=False) # type: ignore def iterate_forward( self, input_embeds, input_states, freqs_cis, block_idx, mask, past_key_values: Optional[Cache] = None, num_steps: Optional[torch.Tensor] = None, attn_maps: dict = {}, return_attn: bool = False, ): x = xk = self.initialize_state(input_embeds) if input_states is None else input_states.clone() if num_steps is None: num_steps_no_grad, num_steps_with_grad = self.randomized_iteration_sampler() # type: ignore elif hasattr(num_steps, "__len__") and len(num_steps) > 1: num_steps_no_grad, num_steps_with_grad = num_steps else: num_steps_no_grad, num_steps_with_grad = num_steps, torch.tensor(0) with torch.no_grad(): # ultra annoying in ddp due to # https://discuss.pytorch.org/t/does-distributeddataparallel-work-with-torch-no-grad-and-find-unused-parameters-false/122594 # for now running with find_unused_params=True enabled even though the graph structure is (technically) clear # and all parameters are always used for step in range(num_steps_no_grad): xk = x x, block_idx, attn_maps = self.core_block_forward( xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps, return_attn ) for step in range(num_steps_with_grad): xk = x x, block_idx, attn_maps = self.core_block_forward( xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps, return_attn ) return self.transformer.ln_f(x), num_steps_no_grad, num_steps_with_grad, xk.detach(), block_idx, attn_maps def core_block_forward( self, x, input_embeds, freqs_cis, mask, past_key_values, block_idx: Union[torch.Tensor, int], attn_maps: dict = {}, return_attn: bool = False, ): x = self.transformer.adapter(torch.cat([x, input_embeds], dim=-1)) for idx, block in enumerate(self.transformer.core_block, start=1): x, attn_map = block(x, freqs_cis, block_idx + idx, mask, past_key_values, return_attn=return_attn) attn_maps[block_idx + idx] = attn_map return x, block_idx + idx, attn_maps @torch.no_grad() def iterate_one_step( self, input_embeds, input_states, position_ids: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, block_idx: Union[torch.Tensor, int] = 0, attention_mask: Optional[Tensor] = None, past_key_values: Optional[Cache] = None, attn_maps: dict = {}, ): if position_ids is None and cache_position is None: freqs_cis = self.freqs_cis[:, : input_embeds.shape[1]] elif position_ids is not None: freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze()) elif cache_position is not None: freqs_cis = self.freqs_cis[:, cache_position] x, block_idx, attn_maps = self.core_block_forward( input_states, input_embeds, freqs_cis, attention_mask, past_key_values, block_idx, attn_maps ) return x, block_idx, attn_maps def predict_from_latents( self, latents, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = None, return_attn: bool = False, attn_maps: dict = {}, ): if position_ids is None and cache_position is None: freqs_cis = self.freqs_cis[:, : latents.shape[1]] elif position_ids is not None: freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze()) elif cache_position is not None: freqs_cis = self.freqs_cis[:, cache_position] x = self.transformer.ln_f(latents) # Coda layers for block_idx, block in enumerate(self.transformer.coda, start=1): x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values) attn_maps[block_idx] = attn_map x = self.transformer.ln_f(x) logits = self.lm_head(x).float() return CausalLMOutputRecurrentLatents( loss=torch.as_tensor(0.0), log_ppl=torch.as_tensor(0.0), logits=logits, past_key_values=past_key_values, attention_maps=attn_maps if len(attn_maps) > 0 else None, ) def embed_inputs( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = None, use_cache: bool = False, cache_position: Optional[torch.Tensor] = None, return_attn: bool = False, **kwargs, ) -> tuple[torch.Tensor, int, dict[int, Tensor]]: # Support multiple position formats: if position_ids is None and cache_position is None: freqs_cis = self.freqs_cis[:, : input_ids.shape[1]] elif position_ids is not None: freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze()) elif cache_position is not None: freqs_cis = self.freqs_cis[:, cache_position] input_embeds = self.transformer.wte(input_ids) if self.emb_scale != 1: input_embeds = input_embeds * self.emb_scale # type: ignore if use_cache and past_key_values is None: past_key_values = HuginnDynamicCache() # Non-recurrent prelude attn_maps = {} for block_idx, block in enumerate(self.transformer.prelude): input_embeds, attn_maps = block( input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn ) return input_embeds, block_idx, attn_maps @torch._dynamo.disable(recursive=False) # type: ignore def randomized_iteration_sampler(self) -> tuple[torch.Tensor, torch.Tensor]: """Outputs are long tensors so that they can be passed through compiled functions""" t = max(self.config.mean_recurrence - self.config.mean_backprop_depth, 0) s = self.config.mean_backprop_depth if self.training: sigma = 0.5 mu = math.log(t + s) - (sigma**2 / 2) rate = torch.zeros((1,)).log_normal_(mean=mu, std=sigma) p = torch.poisson(torch.tensor([rate], dtype=torch.float)) + 1 n = torch.clamp(p - s, min=0) k = torch.as_tensor(torch.minimum(torch.as_tensor(s), p)) else: n, k = torch.as_tensor(self.config.mean_recurrence), torch.as_tensor(0) return n.to(dtype=torch.long), k.to(dtype=torch.long) def initialize_state(self, input_embeds, deterministic: bool = False): x = torch.randn_like(input_embeds) std = self.config.init_values["std"] torch.nn.init.trunc_normal_(x, mean=0.0, std=std, a=-3 * std, b=3 * std) if self.emb_scale != 1: x = x * self.emb_scale return x if not deterministic else x.zero_() def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Cache] = None, attention_mask: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ): model_inputs = {} model_inputs["cache_position"] = cache_position current_input_length = input_ids.shape[1] if past_key_values is not None: if type(past_key_values) != HuginnDynamicCache: # Need to use custom cache, detect and replace HF dynamic cache if generate injects it assert past_key_values.get_seq_length() == 0 past_key_values = HuginnDynamicCache() model_inputs["past_key_values"] = past_key_values if kwargs["use_cache"] else None input_ids = input_ids[:, cache_position] # type: ignore model_inputs["input_ids"] = input_ids.clone(memory_format=torch.contiguous_format) if cache_position is None: position_ids = torch.arange(current_input_length)[None, :].to(input_ids.device) model_inputs["position_ids"] = position_ids[:, -current_input_length:].clone( memory_format=torch.contiguous_format ) # some form of position_ids is a critical argument for the model to correctly apply rope! # forward all other entries for key, value in kwargs.items(): if key not in model_inputs: model_inputs[key] = value return model_inputs @torch.no_grad() def generate(self, *args, **kwargs): """Dispatcher - use HF generate in all normal cases.""" if any( k in kwargs for k in ("continuous_compute", "latent_dampening", "criterion", "exit_threshold", "cache_kwargs") ): print("Dispatching to custom generate function call") return self.generate_with_adaptive_compute(*args, **kwargs) else: return super().generate(*args, **kwargs) @torch.no_grad() def generate_minimal( self, input_ids: torch.LongTensor, generation_config: Optional[GenerationConfig] = None, # type: ignore tokenizer=None, streamer=None, continuous_compute=False, # warm-start state / continuous CoT cache_kwargs: dict = {}, **model_kwargs, ) -> Union[torch.Tensor, dict[str, Any]]: """Minimal single-sequence generation. Template for more complicated generate tasks""" # Setup if generation_config is None: generation_config: GenerationConfig = self.generation_config # type: ignore model_kwargs["past_key_values"] = HuginnDynamicCache(**cache_kwargs) model_kwargs["use_cache"] = True model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) stop_tokens = self._get_stops(generation_config, tokenizer).to(input_ids.device) if continuous_compute: embedded_inputs, _, _ = self.embed_inputs(input_ids) model_kwargs["input_states"] = self.initialize_state(embedded_inputs) # Generate tokens for _ in range(generation_config.max_length - input_ids.shape[1]): # Forward pass model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) outputs = self(**model_inputs) next_token_logits = outputs.logits[0, -1, :] if continuous_compute: current_last_latent = outputs.latent_states[:, -1:, :] # Sample or select next token if generation_config.do_sample: if generation_config.temperature: next_token_logits = next_token_logits / generation_config.temperature probs = F.softmax(next_token_logits, dim=-1) # Apply top_k if generation_config.top_k: top_k_probs, _ = torch.topk(probs, generation_config.top_k) probs[probs < top_k_probs[-1]] = 0 # Apply top_p if generation_config.top_p: sorted_probs = torch.sort(probs, descending=True)[0] cumsum = torch.cumsum(sorted_probs, dim=-1) probs[cumsum > generation_config.top_p] = 0 # Apply min_p if generation_config.min_p: probs[probs < generation_config.min_p * probs.max()] = 0 probs = probs / probs.sum() next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) input_ids = torch.cat([input_ids, next_token[None, :]], dim=-1) # type: ignore if streamer: streamer.put(next_token.cpu()) # Update model kwargs model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs) if continuous_compute: model_kwargs["input_states"] = current_last_latent # Check if we hit a stop token if stop_tokens is not None and next_token in stop_tokens: break if streamer: streamer.end() if generation_config.return_dict_in_generate: return GenerateDecoderOnlyOutput( sequences=input_ids, scores=None, logits=None, attentions=None, hidden_states=None, past_key_values=model_kwargs.get("past_key_values"), ) return input_ids @torch.no_grad() def generate_with_adaptive_compute( self, input_ids: torch.LongTensor, generation_config: Optional[GenerationConfig] = None, # type: ignore tokenizer=None, streamer=None, continuous_compute=False, # warm-start state / continuous CoT latent_dampening=False, criterion="entropy-diff", exit_threshold: Union[str, float, int] = "auto", cache_kwargs: dict = {}, **model_kwargs, ) -> Union[torch.Tensor, GenerateDecoderOnlyOutput]: """Minimal single-sequence generation. Template for more complicated generate tasks""" # Setup if generation_config is None: generation_config: GenerationConfig = self.generation_config # type: ignore model_kwargs["past_key_values"] = HuginnDynamicCache(**cache_kwargs) model_kwargs["use_cache"] = True model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) stop_tokens = self._get_stops(generation_config, tokenizer).to(input_ids.device) if continuous_compute: embedded_inputs, _, _ = self.embed_inputs(input_ids) current_last_latent = self.initialize_state(embedded_inputs) compute_steps = [] # Generate tokens for step in range(generation_config.max_length - input_ids.shape[1]): # Adaptive compute forward model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) aux_inputs = { k: model_inputs[k] for k in ["cache_position", "past_key_values", "attention_mask"] if k in model_inputs } embedded_inputs, block_idx, _ = self.embed_inputs(model_inputs["input_ids"], **aux_inputs) if not continuous_compute: current_latents = self.initialize_state(embedded_inputs, deterministic=False) else: current_latents = current_last_latent # Prep criterions: if criterion == "entropy-diff": entropy = torch.tensor(100.0, device=input_ids.device) exit_threshold = 1e-3 if exit_threshold == "auto" else float(exit_threshold) elif criterion in ["latent-diff", "none"]: exit_threshold = 0.03 if exit_threshold == "auto" else float(exit_threshold) elif "kl" in criterion: V = self.config.padded_vocab_size log_probs = (1 / V * torch.ones(V, device=input_ids.device)).log() if criterion == "minp-kl": exit_threshold = 1e-6 if exit_threshold == "auto" else float(exit_threshold) else: exit_threshold = 5e-4 if exit_threshold == "auto" else float(exit_threshold) elif criterion == "argmax-stability": stable_for_n_steps = 0 current_argmax = torch.tensor(-1, dtype=torch.long, device=input_ids.device) exit_threshold = 5 if exit_threshold == "auto" else int(exit_threshold) else: raise ValueError("Invalid adaptive compute strategy.") all_latents = [] exit_values = [] for compute_step in range(model_inputs["num_steps"]): prev_latents = current_latents.clone() current_latents, block_idx, _ = self.iterate_one_step( embedded_inputs, current_latents, block_idx=block_idx, **aux_inputs ) all_latents.append(current_latents if latent_dampening else None) if step > 0: # do not exit in prefill: if criterion == "entropy-diff": prev_entropy = entropy.clone() outputs = self.predict_from_latents(current_latents, **aux_inputs) probs = F.softmax(outputs.logits[:, -1, :], dim=-1) # type: ignore entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1).mean() entropy_diff = (entropy - prev_entropy).abs() exit_values.append(entropy_diff.item()) if entropy_diff < exit_threshold: break elif criterion == "latent-diff": norm_diff = (prev_latents - current_latents).norm() / current_latents.norm() exit_values.append(norm_diff.item()) if norm_diff < exit_threshold: break elif criterion == "kl": prev_log_probs = log_probs.clone() outputs = self.predict_from_latents(current_latents, **aux_inputs) log_probs = F.log_softmax(outputs.logits[:, -1, :], dim=-1) # type: ignore kl = F.kl_div(log_probs, prev_log_probs, reduction="none", log_target=True).sum(dim=-1) exit_values.append(kl.item()) if kl < exit_threshold: break elif criterion == "minp-kl": prev_log_probs = log_probs.clone() outputs = self.predict_from_latents(current_latents, **aux_inputs) probs = F.softmax(outputs.logits[:, -1, :], dim=-1) # type: ignore probs[probs < 0.1 * probs.max()] = 1 / V probs = probs / probs.sum() log_probs = probs.log() kl = F.kl_div(log_probs, prev_log_probs, reduction="none", log_target=True).sum(dim=-1) exit_values.append(kl.item()) if kl < exit_threshold: break elif criterion == "argmax-stability": prev_argmax = current_argmax.clone() outputs = self.predict_from_latents(current_latents, **aux_inputs) current_argmax = outputs.logits[0, -1, :].argmax(dim=-1) # type: ignore if current_argmax == prev_argmax: stable_for_n_steps += 1 else: stable_for_n_steps = 0 exit_values.append(stable_for_n_steps) if stable_for_n_steps >= exit_threshold: break elif criterion == "none": pass else: if not latent_dampening: outputs = self.predict_from_latents(current_latents, **aux_inputs) else: dampened_latents = torch.sum(torch.cat(all_latents, dim=0), dim=0, keepdim=True) outputs = self.predict_from_latents(dampened_latents, **aux_inputs) compute_steps.append([compute_step + 1, exit_values]) next_token_logits = outputs.logits[0, -1, :] # type: ignore if continuous_compute: # Save last latent current_last_latent = current_latents[:, -1:, :] # Sample or select next token if generation_config.do_sample: if generation_config.temperature: next_token_logits = next_token_logits / generation_config.temperature probs = F.softmax(next_token_logits, dim=-1) # Apply top_k if generation_config.top_k: top_k_probs, _ = torch.topk(probs, generation_config.top_k) probs[probs < top_k_probs[-1]] = 0 # Apply top_p if generation_config.top_p: sorted_probs = torch.sort(probs, descending=True)[0] cumsum = torch.cumsum(sorted_probs, dim=-1) probs[cumsum > generation_config.top_p] = 0 # Apply min_p if generation_config.min_p: probs[probs < generation_config.min_p * probs.max()] = 0 probs = probs / probs.sum() next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) input_ids = torch.cat([input_ids, next_token[None, :]], dim=-1) # type: ignore if streamer: streamer.put(next_token.cpu()) # Update model kwargs model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs) # Check if we hit a stop token if stop_tokens is not None and next_token in stop_tokens: break if streamer: streamer.end() if generation_config.return_dict_in_generate: return GenerateDecoderOnlyOutput( sequences=input_ids, scores=compute_steps, # type: ignore logits=None, attentions=None, hidden_states=None, past_key_values=model_kwargs.get("past_key_values"), ) return input_ids def _get_stops(self, generation_config, tokenizer): stop_tokens = set() if generation_config.eos_token_id is not None: stop_tokens.add(generation_config.eos_token_id) if hasattr(generation_config, "stop_strings") and tokenizer and generation_config.stop_strings: for s in generation_config.stop_strings: token_id = tokenizer(s, add_special_tokens=False)["input_ids"][0] stop_tokens.add(token_id) return torch.tensor(list(stop_tokens)) def get_stats(self, logits, x, latent_states, xk, input_embeds, num_steps_no_grad, num_steps_with_grad): probs = torch.softmax(logits.float(), dim=-1) prob_entropy = torch.where(probs > 0, -probs * probs.log(), 0).sum(dim=-1) residual_diff = (x - latent_states).norm(dim=-1) rel_residual = residual_diff / latent_states.norm(dim=-1) stats = { "entropy": prob_entropy, "residual_diff": residual_diff, "rel_residual": rel_residual, "num_steps_no_grad": num_steps_no_grad, "num_steps_with_grad": num_steps_with_grad, } return stats #################################### Utils ####################################################################### def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, condense_ratio: int = 1): with torch.autocast("cuda", enabled=False): inv_freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) t = torch.arange(end, dtype=torch.float32, device=inv_freqs.device) / condense_ratio freqs = torch.outer(t, inv_freqs).float() return torch.stack([torch.cos(freqs)[None, :, None, :], torch.sin(freqs)[None, :, None, :]], dim=4) # equivalent to # freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) def apply_rotary_emb_complex_like(q: Tensor, k: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: with torch.autocast("cuda", enabled=False): qk_r2 = torch.cat([q, k], dim=2).unflatten(dim=-1, sizes=(-1, 2)).float() # cast to float32 for smooth skin rotated_qk_r2 = torch.stack( [ qk_r2[..., 0] * freqs_cis[..., 0] - qk_r2[..., 1] * freqs_cis[..., 1], qk_r2[..., 1] * freqs_cis[..., 0] + qk_r2[..., 0] * freqs_cis[..., 1], ], -1, ).flatten(3) rotated_qk = rotated_qk_r2 return torch.split(rotated_qk.type_as(q), q.shape[2], dim=2) # type: ignore #################################### HF registration ############################################################ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM # New RavenConfig.register_for_auto_class() RavenForCausalLM.register_for_auto_class("AutoModel") RavenForCausalLM.register_for_auto_class("AutoModelForCausalLM") # Old? AutoConfig.register("huginn_raven", RavenConfig) AutoModel.register(RavenConfig, RavenForCausalLM) AutoModelForCausalLM.register(RavenConfig, RavenForCausalLM)