Add a 4/16/4 recurrent 1B architecture and 100-step trial results

README.md

@@ -2,34 +2,62 @@
 license: apache-2.0
 library_name: transformers
 tags:
-  - nano_dense
+  - nano_dense_loop
   - architecture-only
   - from-scratch
   - deepseek-v4
   - dense
+  - recurrent
 language:
   - en
   - zh
 ---
 
-# Nano-1B
+# Recurrent Nano-1B
 
-A **1,076,700,672-parameter dense decoder architecture for training from scratch**.
-This repository contains configuration, model code and a DeepSeek V4 tokenizer.
-**There are no model weights or trained checkpoints.** Initialize with
-`AutoModelForCausalLM.from_config`, rather than loading pretrained model weights.
+A **1,076,700,672-parameter dense decoder architecture for training from scratch**,
+with a shared core that can run multiple times. This revision contains configuration,
+model code and the same pinned DeepSeek V4 tokenizer. **No model weights or trained
+checkpoints are included.** Initialize with `AutoModelForCausalLM.from_config`.
 
-The dimensions follow [MiniCPM5-1B](https://huggingface.co/openbmb/MiniCPM5-1B).
-Normalization, dense SwiGLU and interleaved RoPE reuse Transformers' DeepSeek V4
-components. `NanoDenseForCausalLM` supplies the model entry point and standard
-global grouped-query attention. The tokenizer is pinned to
-[DeepSeek-V4-Flash-Base, revision 8855555](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Base/tree/8855555deef230a27a21a8d6f294b7b7497759b6).
+The base dimensions follow [MiniCPM5-1B](https://huggingface.co/openbmb/MiniCPM5-1B).
+Normalization, dense SwiGLU and interleaved RoPE use Transformers' DeepSeek V4
+components. `NanoDenseForCausalLM` supplies global grouped-query attention and
+the recurrent execution schedule. The model type is `nano_dense_loop`.
 
 ## Architecture
 
+```text
+DSV4 tokens → embedding → prelude (4 layers)
+                       → core (16 layers) → same core (16 layers)
+                       → coda (4 layers) → final RMSNorm → LM head
+```
+
+| Component | Physical layers | Executions at K=2 |
+| --- | ---: | ---: |
+| Prelude | 4 | 4 |
+| Shared core | 16 | 32 |
+| Coda | 4 | 4 |
+| Total | **24** | **40** |
+
+The prelude runs once, the core runs `num_core_iterations` times, and the coda
+runs once. Each core pass receives the previous pass's output and reuses the
+same 16 sets of parameters. Token positions stay the same between passes.
+Backpropagation traverses every pass. No extra boundary normalization, anchor
+injection, adaptive halting or truncated backpropagation is added.
+
+The default is **K=2**. K=1 is exactly the original 24-layer execution for the
+same weights. Physical layer names and weight shapes are unchanged; changing
+K does not add parameters. The split preserves the 1/6, 2/3, 1/6 proportions of
+the earlier 36-layer Nano loop candidate; it is an initial candidate, not a
+measured optimum or a complete reproduction of a paper's training recipe.
+
 | Setting | Value |
 | --- | ---: |
-| Decoder layers | 24 |
+| Physical decoder layers | 24 |
+| Prelude / core / coda | 4 / 16 / 4 |
+| Default core iterations | 2 |
+| Default effective depth | 40 |
 | Residual hidden size | 1,536 |
 | Dense SwiGLU intermediate size | 4,608 |
 | Query / key-value heads | 16 / 2 |
@@ -37,7 +65,7 @@ global grouped-query attention. The tokenizer is pinned to
 | Attention projection width | 2,048 |
 | Attention pattern | Global causal GQA in every layer |
 | Normalization | DeepseekV4RMSNorm, epsilon 1e-6 |
-| Position encoding | V4 interleaved RoPE over all 128 head dimensions |
+| Position encoding | V4 interleaved full-head RoPE |
 | RoPE theta | 5,000,000 |
 | Configured context target | 131,072 tokens |
 | Vocabulary | 129,280 |
@@ -46,36 +74,28 @@ global grouped-query attention. The tokenizer is pinned to
 | Total parameters | **1,076,700,672** |
 | Excluding embedding and output head | **679,552,512** |
 
-Each layer computes:
-
-```text
-h = x + GQA(RMSNorm(x))
-y = h + SwiGLU(RMSNorm(h))
-```
-
-The attention width is independent of the residual width: Q projects
-1536 -> 2048, K/V each project 1536 -> 256, and O projects 2048 -> 1536.
-All 24 layers have independent parameters. Linear and embedding weights use
-normal initialization with standard deviation 0.02; RMSNorm weights start at one.
+Each layer computes `h = x + GQA(RMSNorm(x))`, then
+`y = h + SwiGLU(RMSNorm(h))`. Q projects 1536 → 2048, K/V each project
+1536 → 256, and O projects 2048 → 1536. Linear and embedding weights use normal
+initialization with standard deviation 0.02; RMSNorm weights start at one.
 No depth-dependent initialization scaling is applied.
 
-128K is a configuration target, not a measured capability of an untrained model.
-Long-context performance requires a suitable training curriculum and evaluation.
-At batch size one, a complete 128K BF16 KV cache occupies **3 GiB**, excluding
-weights, temporary tensors and allocator overhead.
+Each execution has a separate KV cache slot, including the two executions of
+each shared core layer. At batch size one, a complete 128K BF16 KV cache occupies
+**5 GiB at K=2**, versus 3 GiB at K=1, excluding weights and temporary tensors.
+The 128K context is a configuration target, not a validated capability.
 
 ## Initialize for training
 
-```bash
-pip install -r requirements.txt
-```
+Install `requirements.txt`, then run this from a local checkout of this revision:
 
 ```python
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
-repo = "bowang0911/Nano-1B"
+repo = "."
 config = AutoConfig.from_pretrained(repo, trust_remote_code=True)
+config.num_core_iterations = 2
 tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
 model = AutoModelForCausalLM.from_config(
     config,
@@ -83,70 +103,97 @@ model = AutoModelForCausalLM.from_config(
     attn_implementation="sdpa",
     dtype=torch.bfloat16,
 )
-# Randomly initialized model; pass it to your training loop / Trainer.
-```
-
-To inspect the full architecture without allocating weight storage:
-
-```python
-with torch.device("meta"):
-    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
-print(sum(p.numel() for p in model.parameters()))  # 1076700672
 ```
 
-From a local copy, `python count_parameters.py --meta` checks the analytical
-count against a full-size meta-device model.
-
-## Tokenized data and packing
+When loading directly from the Hub before merge, pass the PR revision to
+`AutoConfig.from_pretrained` and `AutoTokenizer.from_pretrained`; `main` still
+contains the vanilla architecture. `num_core_iterations` can also be passed to
+`forward` or `generate` for a specific call. Use the same K across distributed
+training ranks. This implementation does not sample K randomly.
 
-`tokenizer.json` is byte-identical to the pinned DeepSeek V4 artifact. The only
-tokenizer-config change is `model_max_length=131072`; token IDs and encoding are
-unchanged. `tokenizer_provenance.json` records upstream and published checksums.
-Existing data encoded with this pinned tokenizer can be reused directly.
-The pretraining convention used here is raw text with `add_special_tokens=False`
-and one EOS token (ID 1) appended per document. No chat template is provided.
+`python count_parameters.py --meta` verifies the full physical parameter count
+without allocating weight storage and reports the configured execution depth
+and KV cache size. `--core-iterations 1` checks the single-pass configuration.
 
-For ordinary padded batches, pass a two-dimensional `attention_mask` and set
-padding labels to `-100`. Default positions are derived from that mask, including
-when continuing a KV cache. The tokenizer uses EOS as its padding token; the
-model embedding has no dedicated zeroed padding row, preserving the EOS embedding.
+## Data, packing and cache
 
-For padding-free document packing, use one flattened row (`batch_size=1`) and
-reset `position_ids` to zero at each document boundary. Omit `attention_mask`, or
-pass an all-ones mask. The model isolates documents and masks the first label of
-each subsequent document so the LM loss does not cross boundaries. Packing with
-padding, nonzero restart positions or multiple batch rows raises an error.
-KV caching is disabled during training and for packed sequences. Normal
-unpacked evaluation/generation supports dynamic KV caching.
-
-Eager and SDPA backends are validated on CPU. The FlashAttention adapter receives
-packed positions and dispatches separate document segments; running actual
-FlashAttention kernels requires a compatible GPU installation.
+The tokenizer remains pinned to
+[DeepSeek-V4-Flash-Base, revision 8855555](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Base/tree/8855555deef230a27a21a8d6f294b7b7497759b6).
+`tokenizer.json` is byte-identical to the original artifact. The existing
+`model_max_length=131072` tokenizer-config setting and token IDs are unchanged.
+`tokenizer_provenance.json` records the hashes. Pretraining uses raw text with
+`add_special_tokens=False` and one appended EOS (ID 1), without a chat template.
+
+For padded batches, pass a two-dimensional `attention_mask` and set padding
+labels to `-100`. Default positions are derived from the mask. For padding-free
+packing, use one flattened row and reset `position_ids` at document boundaries.
+Omit the mask or pass an all-ones mask. The model isolates documents and excludes
+cross-document next-token targets. Packed batches with padding or multiple rows
+are rejected. Training and packed evaluation disable KV caching.
+
+Unpacked evaluation/generation supports dynamic cache only. The returned cache
+records its depth and layer schedule; these cannot change during continuation.
+An empty vanilla `DynamicCache` is accepted and replaced with the loop cache.
+Populated untagged caches and other cache implementations are rejected.
+
+## 100-step trial
+
+The full-size trial ran on the same CoreWeave 4×GB200 node as the dense baseline.
+Both used seed 20260925, the same DSV4-tokenized English web stream and token
+order, 2048-token sequences, global batch 32, and **6,553,600 training tokens**.
+AdamW used LR 3e-4 with 10 warmup steps, betas 0.9/0.95, weight decay 0.1, and
+gradient clipping 1.0; parameters and optimizer states were FP32 with BF16 autocast.
+The loop trained at fixed K=2. Every 20 steps, both evaluation depths used the
+same fixed samples from six existing validation panels. Each panel contains
+32 windows of 2048 tokens; the overall panel overlaps the domain panels.
+
+| Metric at step 100 | Base 1B (Vanilla) | Recurrent 1B (K=2) |
+| --- | ---: | ---: |
+| Physical / effective layers | 24 / 24 | 24 / 40 |
+| Prelude / core / coda | — | 4 / 16 / 4 |
+| Final training CE | 6.9774 | 7.1601 |
+| Overall validation CE | **7.2836** | 7.4369 |
+| Overall validation BPB | **2.2104** | 2.2568 |
+| Mean seconds per step, steps 11–100 | 0.741 | 0.920 |
+| Peak allocated GiB/GPU | 29.97 | 33.13 |
+
+On the recurrent checkpoint, evaluating with K=1 gives CE **7.5839** and BPB
+**2.3012**. A second core pass improves that checkpoint, but it does not outperform
+the separately trained dense baseline in this trial. These runs match training
+tokens, not compute. One seed and 100 steps do not establish long-run quality.
+CE includes EOS; BPB uses text-token NLL and exact UTF-8 byte counts, excluding
+special tokens from both numerator and denominator.
+
+- [Dense baseline](https://wandb.ai/pplx-ai/nano-1b-ablation/runs/opurhf9m)
+- [Recurrent trial](https://wandb.ai/pplx-ai/nano-1b-ablation/runs/h73o4lvs)
+
+The trial completed with finite losses, gradients and weight-update probes in
+all three sections. Model/optimizer weights remain on the training cluster.
+W&B server readback verified 100 training points and all 72 per-panel evaluation
+records (six panels × six checkpoints × two depths).
 
 ## Validation
 
-Validated with Python 3.12.11, PyTorch 2.12.0 and Transformers 5.9.0:
+Eleven local behavioral tests passed with Python 3.12, PyTorch 2.12.0 and
+Transformers 5.9.0. Tests remain outside this model repository. They cover:
 
-- Full-size meta parameter count, including the wider attention projections.
-- An independent complex-number RoPE / attention reference, including large positions.
-- Training loss, finite gradients through every parameter, an AdamW update,
-  and gradient-checkpointing equivalence.
-- Eager/SDPA cached chunked decoding and cached/uncached greedy generation.
-- Left/right padding, packed-document isolation, boundary loss masking and
-  FlashAttention varlen dispatch with external GPU kernels stubbed.
-- FP32/BF16 save/reload and invalid configuration override rejection.
-- Tokenizer parity with the pinned source.
+- Full-size parameter count on meta device and the 40-layer execution schedule.
+- Exact K=1 initialization, logits and gradients versus the dense reference.
+- K=2 outputs versus an independently unrolled network and summed shared gradients.
+- Gradient checkpointing, model save/reload and optimizer continuation.
+- Separate cache slots, chunked decoding and cached/uncached greedy generation.
+- Depth-change rejection, invalid settings, packed-document isolation and padding.
+- FlashAttention varlen dispatch at every execution, with external kernels stubbed.
 
-The 11 local behavioral tests are kept outside this model repository. This is
-small-model CPU validation, not full-size training, GPU throughput or 128K quality
-validation. Transformers is pinned because the implementation imports its V4
-components and attention/cache interfaces.
+The full-size GPU trial used SDPA, PyTorch
+`2.12.0a0+0291f960b6.nv26.04.48445190`, CUDA 13.2 and Transformers 5.9.0.
+Actual FlashAttention kernels and 128K training were not evaluated.
 
 ## Sources and licenses
 
-- [MiniCPM5-1B configuration](https://huggingface.co/openbmb/MiniCPM5-1B/blob/87179e5c1f455ef22e6223592d2d61351b525bfc/config.json): architecture dimensions.
+- [MiniCPM5-1B configuration](https://huggingface.co/openbmb/MiniCPM5-1B/blob/87179e5c1f455ef22e6223592d2d61351b525bfc/config.json): base dimensions.
 - [Transformers 5.9.0 DeepSeek V4](https://github.com/huggingface/transformers/blob/v5.9.0/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py): RMSNorm, dense MLP and interleaved RoPE.
 - Model implementation: Apache-2.0, see `LICENSE`.
 - DeepSeek tokenizer assets: MIT, see `LICENSE-TOKENIZER`.
 
-No pretrained MiniCPM or DeepSeek model weights are included or converted.
+No pretrained MiniCPM or DeepSeek weights are included or converted.

config.json

@@ -19,7 +19,11 @@
   "intermediate_size": 4608,
   "max_position_embeddings": 131072,
   "mlp_bias": false,
-  "model_type": "nano_dense",
+  "model_type": "nano_dense_loop",
+  "prelude_num_layers": 4,
+  "core_num_layers": 16,
+  "coda_num_layers": 4,
+  "num_core_iterations": 2,
   "num_attention_heads": 16,
   "num_hidden_layers": 24,
   "num_key_value_heads": 2,

configuration_nano_dense.py

@@ -19,13 +19,17 @@ def _reject_unused(values):
 
 @strict
 class NanoDenseConfig(PreTrainedConfig):
-    model_type = "nano_dense"
+    model_type = "nano_dense_loop"
     keys_to_ignore_at_inference = ["past_key_values"]
 
     vocab_size: int = 129280
     hidden_size: int = 1536
     intermediate_size: int = 4608
     num_hidden_layers: int = 24
+    prelude_num_layers: int = 4
+    core_num_layers: int = 16
+    coda_num_layers: int = 4
+    num_core_iterations: int = 2
     num_attention_heads: int = 16
     num_key_value_heads: int = 2
     head_dim: int = 128
@@ -80,6 +84,11 @@ class NanoDenseConfig(PreTrainedConfig):
 
     def validate_architecture(self):
         _reject_unused(vars(self))
+        for name in ('prelude_num_layers', 'core_num_layers', 'coda_num_layers', 'num_core_iterations'):
+            if type(getattr(self, name)) is not int or getattr(self, name) <= 0:
+                raise ValueError(f'{name} must be a positive integer')
+        if self.prelude_num_layers + self.core_num_layers + self.coda_num_layers != self.num_hidden_layers:
+            raise ValueError('prelude + core + coda must equal num_hidden_layers')
         for name in ("vocab_size", "hidden_size", "intermediate_size", "num_hidden_layers",
                      "num_attention_heads", "num_key_value_heads", "head_dim", "max_position_embeddings"):
             value = getattr(self, name)
@@ -99,5 +108,13 @@ class NanoDenseConfig(PreTrainedConfig):
         if not math.isfinite(self.rms_norm_eps) or self.rms_norm_eps <= 0:
             raise ValueError("rms_norm_eps must be finite and positive")
 
+    def get_layer_schedule(self, num_core_iterations=None):
+        iterations = self.num_core_iterations if num_core_iterations is None else num_core_iterations
+        if type(iterations) is not int or iterations <= 0:
+            raise ValueError('num_core_iterations must be a positive integer')
+        start = self.prelude_num_layers
+        end = start + self.core_num_layers
+        return tuple(range(start)) + tuple(range(start, end)) * iterations + tuple(range(end, self.num_hidden_layers))
+
 
 NanoDenseConfig.register_for_auto_class()

count_parameters.py

@@ -20,18 +20,27 @@ def parameter_breakdown(config):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--meta", action="store_true")
+    parser.add_argument("--core-iterations", type=int, default=None)
     args = parser.parse_args()
     root = Path(__file__).resolve().parent
     config = json.loads((root / "config.json").read_text())
+    iterations = config["num_core_iterations"] if args.core_iterations is None else args.core_iterations
+    if iterations <= 0:
+        parser.error("--core-iterations must be positive")
+    effective_layers = config["prelude_num_layers"] + config["core_num_layers"] * iterations + config["coda_num_layers"]
     result = parameter_breakdown(config)
+    result["physical_layers"] = config["num_hidden_layers"]
+    result["core_iterations"] = iterations
+    result["effective_layers"] = effective_layers
     result["bf16_kv_bytes_at_context_limit_batch1"] = (
-        config["num_hidden_layers"] * 2 * config["num_key_value_heads"]
+        effective_layers * 2 * config["num_key_value_heads"]
         * config["head_dim"] * 2 * config["max_position_embeddings"]
     )
     if args.meta:
         import torch
         from transformers import AutoConfig, AutoModelForCausalLM
         cfg = AutoConfig.from_pretrained(root, trust_remote_code=True)
+        cfg.num_core_iterations = iterations
         with torch.device("meta"):
             model = AutoModelForCausalLM.from_config(cfg, trust_remote_code=True)
         actual = sum(p.numel() for p in model.parameters())

modeling_nano_dense.py

@@ -3,6 +3,8 @@
 # Licensed under the Apache License, Version 2.0.
 """Nano dense decoder using DeepSeek V4 normalization, MLP and interleaved RoPE."""
 
+from copy import copy
+
 import torch
 from torch import nn
 from transformers import initialization as init
@@ -22,6 +24,15 @@ from transformers.utils.output_capturing import capture_outputs
 from .configuration_nano_dense import NanoDenseConfig
 
 
+class NanoDenseLoopCache(DynamicCache):
+    def __init__(self, config, num_core_iterations):
+        self.num_core_iterations = num_core_iterations
+        self.physical_layer_indices = config.get_layer_schedule(num_core_iterations)
+        cache_config = copy(config)
+        cache_config.num_hidden_layers = len(self.physical_layer_indices)
+        super().__init__(config=cache_config)
+
+
 def _document_boundaries(position_ids, attention_mask):
     boundaries = position_ids[:, 1:] != position_ids[:, :-1] + 1
     if attention_mask is not None:
@@ -57,7 +68,7 @@ class NanoDenseAttention(nn.Module):
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
 
     def forward(self, hidden_states, position_embeddings, attention_mask=None,
-                past_key_values=None, **kwargs):
+                past_key_values=None, cache_layer_idx=None, **kwargs):
         batch, length, _ = hidden_states.shape
         shape = (batch, length, -1, self.head_dim)
         query = self.q_proj(hidden_states).view(shape).transpose(1, 2)
@@ -67,7 +78,7 @@ class NanoDenseAttention(nn.Module):
         query = apply_rotary_pos_emb(query, cos, sin)
         key = apply_rotary_pos_emb(key, cos, sin)
         if past_key_values is not None:
-            key, value = past_key_values.update(key, value, self.layer_idx)
+            key, value = past_key_values.update(key, value, cache_layer_idx)
         backend = ALL_ATTENTION_FUNCTIONS.get_interface(
             self.config._attn_implementation, eager_attention_forward
         )
@@ -132,12 +143,30 @@ class NanoDenseModel(NanoDensePreTrainedModel):
         self.gradient_checkpointing = False
         self.post_init()
 
+    def _prepare_loop_cache(self, cache, iterations, schedule):
+        if isinstance(cache, NanoDenseLoopCache):
+            if cache.num_core_iterations != iterations or cache.physical_layer_indices != schedule:
+                raise ValueError('Cannot change iterations or layer schedule of an existing cache')
+            return cache
+        if cache is not None:
+            if type(cache) is not DynamicCache or cache.offloading or any(
+                layer.get_seq_length() != 0 for layer in cache.layers
+            ):
+                raise ValueError('Use an empty DynamicCache or the returned loop cache')
+        return NanoDenseLoopCache(self.config, iterations)
+
     @merge_with_config_defaults
     @capture_outputs
     def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs):
+                past_key_values=None, inputs_embeds=None, use_cache=None,
+                num_core_iterations=None, **kwargs):
         if (input_ids is None) == (inputs_embeds is None):
             raise ValueError("Specify exactly one of input_ids and inputs_embeds")
+        if num_core_iterations is None:
+            num_core_iterations = (past_key_values.num_core_iterations
+                                   if isinstance(past_key_values, NanoDenseLoopCache)
+                                   else self.config.num_core_iterations)
+        schedule = self.config.get_layer_schedule(num_core_iterations)
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
         if self.training:
@@ -173,19 +202,20 @@ class NanoDenseModel(NanoDensePreTrainedModel):
                     raise ValueError("Packed documents with padding are unsupported; use padding-free packing")
                 attention_mask = None
             use_cache = False
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(config=self.config)
+        if use_cache:
+            past_key_values = self._prepare_loop_cache(past_key_values, num_core_iterations, schedule)
         mask = create_causal_mask(
             config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask,
             past_key_values=past_key_values, position_ids=position_ids,
         )
         position_embeddings = self.rotary_emb(inputs_embeds, position_ids, layer_type="main")
         hidden_states = inputs_embeds
-        for layer in self.layers:
+        for execution_index, physical_index in enumerate(schedule):
+            layer = self.layers[physical_index]
             hidden_states = layer(
                 hidden_states, attention_mask=mask, position_ids=position_ids,
                 position_embeddings=position_embeddings, past_key_values=past_key_values,
-                use_cache=use_cache, **kwargs,
+                cache_layer_idx=execution_index, use_cache=use_cache, **kwargs,
             )
         return BaseModelOutputWithPast(
             last_hidden_state=self.norm(hidden_states), past_key_values=past_key_values,
@@ -205,13 +235,13 @@ class NanoDenseForCausalLM(NanoDensePreTrainedModel, GenerationMixin):
     @can_return_tuple
     def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                 past_key_values=None, inputs_embeds=None, labels=None, use_cache=None,
-                logits_to_keep=0, **kwargs):
+                logits_to_keep=0, num_core_iterations=None, **kwargs):
         if labels is not None and not (isinstance(logits_to_keep, int) and logits_to_keep == 0):
             raise ValueError("Training with labels requires all logits (logits_to_keep=0)")
         outputs = self.model(
             input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
             past_key_values=past_key_values, inputs_embeds=inputs_embeds,
-            use_cache=use_cache, **kwargs,
+            use_cache=use_cache, num_core_iterations=num_core_iterations, **kwargs,
         )
         indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.lm_head(outputs.last_hidden_state[:, indices, :])
@@ -228,6 +258,11 @@ class NanoDenseForCausalLM(NanoDensePreTrainedModel, GenerationMixin):
             hidden_states=outputs.hidden_states, attentions=outputs.attentions,
         )
 
+    def _prepare_cache_for_generation(self, generation_config, *args, **kwargs):
+        if generation_config.cache_implementation not in (None, 'dynamic'):
+            raise ValueError("Loop generation supports only cache_implementation='dynamic'")
+        return super()._prepare_cache_for_generation(generation_config, *args, **kwargs)
+
 
 NanoDenseModel.register_for_auto_class("AutoModel")
 NanoDenseForCausalLM.register_for_auto_class("AutoModelForCausalLM")