Update raven_modeling_minimal.py

raven_modeling_minimal.py

@@ -11,7 +11,7 @@ from .raven_config_minimal import RavenConfig
 from transformers.cache_utils import Cache, DynamicCache
 
 ###################### Huggingface Glue code I ##################################################################
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, GenerationMixin
 from transformers.utils import ModelOutput
 from transformers.generation.utils import GenerateDecoderOnlyOutput
 
@@ -32,7 +32,8 @@ class RavenPreTrainedModel(PreTrainedModel):
     _supports_static_cache = False
 
     def _init_weights(self, module):
-        print("Random Initialization not implemented.")
+        if not torch.rand((1,)).is_meta:
+            print("Random Initialization not implemented.")
 
 
 @dataclass
@@ -70,7 +71,7 @@ class RMSNorm(torch.nn.Module):
 
 
 class HuginnDynamicCache(DynamicCache):
-    def __init__(self, lookup_strategy: str = "latest") -> None:
+    def __init__(self, lookup_strategy: str = "full") -> None:
         super().__init__()
         self._seen_tokens = 0
         self.key_cache: dict[int, dict[int, torch.Tensor]] = {}
@@ -89,6 +90,14 @@ class HuginnDynamicCache(DynamicCache):
         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] = {}
@@ -98,32 +107,49 @@ class HuginnDynamicCache(DynamicCache):
             self._seen_tokens += key_states.shape[-2]
         # Add entries to cache
         for idx, entry in enumerate(key_states.unbind(dim=-2)):
-            assert step_idx < 0 or self._seen_tokens - key_states.shape[-2] + idx not in self.key_cache[step_idx]
+            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:
+        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 == "latest":
+            # 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):
-                    # 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)
+                    # 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 == "skip":
+            elif lookup_strategy.startswith("skip"):
                 existing_keys = []
                 existing_values = []
                 for token_pos in range(self._seen_tokens):
@@ -131,15 +157,22 @@ class HuginnDynamicCache(DynamicCache):
                         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 == "randomized":  # sanity check
+            elif lookup_strategy.startswith("randomized"):  # sanity check
                 rand_keys = []
                 rand_values = []
                 for token_pos in range(self._seen_tokens):
-                    # Find steps that have this token position
-                    steps = [s for s in range(step_idx + 1) if token_pos in self.key_cache[s]]
-                    rand_step = steps[torch.randint(len(steps), (1,))]
-                    rand_keys.append(self.key_cache[rand_step][token_pos])
-                    rand_values.append(self.value_cache[rand_step][token_pos])
+                    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}")
@@ -153,6 +186,18 @@ class HuginnDynamicCache(DynamicCache):
     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:
@@ -265,7 +310,7 @@ class SandwichBlock(torch.nn.Module):
         return x, attn_map
 
 
-class RavenForCausalLM(RavenPreTrainedModel):
+class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
     def __init__(
         self,
         config: RavenConfig,
@@ -323,7 +368,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
             "return_latents": True,
             "return_attention": False,
             "return_head": False,
-            "return_stats": True,
+            "return_stats": False,
         },
         use_cache: bool = False,
         cache_position: Optional[torch.Tensor] = None,
@@ -351,7 +396,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
         # 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
+                input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn=return_attn
             )
             attn_maps[block_idx] = attn_map
 
@@ -365,12 +410,13 @@ class RavenForCausalLM(RavenPreTrainedModel):
             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)
+            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)
 
@@ -407,6 +453,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
         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:
@@ -424,13 +471,13 @@ class RavenForCausalLM(RavenPreTrainedModel):
             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
+                    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
+                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
 
@@ -443,10 +490,11 @@ class RavenForCausalLM(RavenPreTrainedModel):
         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=len(attn_maps) > 0)
+            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
 
@@ -579,7 +627,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
         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) == DynamicCache:
+            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()
@@ -599,6 +647,18 @@ class RavenForCausalLM(RavenPreTrainedModel):
                 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,
@@ -693,6 +753,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
         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]:
@@ -725,46 +786,64 @@ class RavenForCausalLM(RavenPreTrainedModel):
             # 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"]:
-                pass
-            elif criterion == "kl":
+                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 = []
-            for compute_step in range(1, model_inputs["num_steps"]):
+            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 compute_step > 1 and step > 0:  # do not exit in prefill:
+                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()
-                        if entropy_diff < 1e-3:
-                            compute_steps.append([compute_step, entropy_diff.item()])
+                        exit_values.append(entropy_diff.item())
+                        if entropy_diff < exit_threshold:
                             break
                     elif criterion == "latent-diff":
-                        norm_diff = (prev_latents - current_latents).norm()
-                        if norm_diff < 1:
-                            compute_steps.append([compute_step, norm_diff.item()])
+                        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)
-                        if kl < 2e-4:
-                            compute_steps.append([compute_step, kl.item()])
+                        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()
@@ -774,19 +853,19 @@ class RavenForCausalLM(RavenPreTrainedModel):
                             stable_for_n_steps += 1
                         else:
                             stable_for_n_steps = 0
-                        if stable_for_n_steps >= 10:
-                            compute_steps.append([compute_step, stable_for_n_steps])
+                        exit_values.append(stable_for_n_steps)
+                        if stable_for_n_steps >= exit_threshold:
                             break
                     elif criterion == "none":
                         pass
 
             else:
-                compute_steps.append([compute_step, float("NaN")])
                 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