Sync modeling_domynedge.py, configuration_domynedge.py from domyn-forge@9b691e2

configuration_domynedge.py

@@ -1,6 +1,10 @@
+# SPDX-FileCopyrightText: 2026 Domyn
+# SPDX-License-Identifier: Apache-2.0
+
 import transformers
 
-class RopeParameters():
+
+class RopeParameters:
     def __init__(
         self,
         rope_type: str = "default",
@@ -9,10 +13,18 @@ class RopeParameters():
         self.rope_type = rope_type
         self.rope_theta = rope_theta
 
+
 # to support both transformers > 5 and older
-PretrainedConfig = transformers.modeling_utils.PreTrainedConfig if hasattr(transformers.modeling_utils, "PreTrainedConfig") else transformers.configuration_utils.PretrainedConfig
+PretrainedConfig = (
+    transformers.modeling_utils.PreTrainedConfig
+    if hasattr(transformers.modeling_utils, "PreTrainedConfig")
+    else transformers.configuration_utils.PretrainedConfig
+)
+
 
 class DomynEdgeConfig(PretrainedConfig):
+    """Configuration of a DomynEdge model, with per-layer position encoding, sliding window and head size."""
+
     model_type: str = "domynedge"
 
     def __init__(
@@ -26,6 +38,7 @@ class DomynEdgeConfig(PretrainedConfig):
         max_position_embeddings: int = 8192,
         initializer_range: float = 0.02,
         rms_norm_eps: float = 1e-6,
+        *,
         use_cache: bool = True,
         tie_word_embeddings: bool = True,
         rope_parameters: RopeParameters | None = None,
@@ -33,13 +46,16 @@ class DomynEdgeConfig(PretrainedConfig):
         attention_bias: bool = False,
         attention_dropout: float = 0.0,
         position_layers: list[str] = ["nope", "nope"] + ["rope", "rope", "rope", "nope"] * 8 + ["nope", "nope"],
-        window_sizes: list[tuple[int,int]] = [(-1,0), (-1,0)] + [(1024,0), (1024,0), (1024,0), (-1,0)] * 8 + [(-1,0), (-1,0)],
-        head_sizes: list[int] = [256,256] + [64, 64, 64, 256] * 8 + [256, 256],
+        window_sizes: list[tuple[int, int]] = [(-1, 0), (-1, 0)]
+        + [(1024, 0), (1024, 0), (1024, 0), (-1, 0)] * 8
+        + [(-1, 0), (-1, 0)],
+        head_sizes: list[int] = [256, 256] + [64, 64, 64, 256] * 8 + [256, 256],
         pad_token_id: int | None = None,
         bos_token_id: int | None = None,
         eos_token_id: int | None = None,
         **kwargs,
     ):
+        """Initialize and validate the per-layer settings, which must all have `num_hidden_layers` entries."""
         self.vocab_size = vocab_size
         self.max_position_embeddings = max_position_embeddings
         self.hidden_size = hidden_size
@@ -58,35 +74,50 @@ class DomynEdgeConfig(PretrainedConfig):
         self.use_cache = use_cache
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout
-        
+
         self.position_layers = position_layers
         self.window_sizes = window_sizes
         self.head_sizes = head_sizes
 
-        # validate position_layers
+        # Validate position_layers
         if len(self.position_layers) != self.num_hidden_layers:
-            raise ValueError(f"Length of position_layers ({len(self.position_layers)}) must be equal to num_hidden_layers ({self.num_hidden_layers})")
+            msg = (
+                f"Length of position_layers ({len(self.position_layers)}) "
+                f"must be equal to num_hidden_layers ({self.num_hidden_layers})"
+            )
+            raise ValueError(msg)
         if any(pl not in ["rope", "nope"] for pl in self.position_layers):
-            raise ValueError(f"position_layers must be one of 'rope', 'nope'. Got {self.position_layers}")
-        # validate window_sizes
+            msg = f"position_layers must be one of 'rope', 'nope'. Got {self.position_layers}"
+            raise ValueError(msg)
+        # Validate window_sizes
         if len(self.window_sizes) != self.num_hidden_layers:
-            raise ValueError(f"Length of window_sizes ({len(self.window_sizes)}) must be equal to num_hidden_layers ({self.num_hidden_layers})")
-        # validate head_sizes
+            msg = (
+                f"Length of window_sizes ({len(self.window_sizes)}) "
+                f"must be equal to num_hidden_layers ({self.num_hidden_layers})"
+            )
+            raise ValueError(msg)
+        # Validate head_sizes
         if len(head_sizes) != self.num_hidden_layers:
-            raise ValueError(f"Length of head_sizes ({len(head_sizes)}) must be equal to num_hidden_layers ({self.num_hidden_layers})")
+            msg = (
+                f"Length of head_sizes ({len(head_sizes)}) "
+                f"must be equal to num_hidden_layers ({self.num_hidden_layers})"
+            )
+            raise ValueError(msg)
         if any(not isinstance(hs, int) for hs in head_sizes):
-            raise ValueError(f"head_sizes must be a list of integers. Got {head_sizes}")
-        # only one rope dim
-        rope_dims = [head_size for rope,head_size in zip(position_layers, head_sizes) if rope == "rope"]
+            msg = f"head_sizes must be a list of integers. Got {head_sizes}"
+            raise ValueError(msg)
+        # Only one rope dim
+        rope_dims = [head_size for rope, head_size in zip(position_layers, head_sizes, strict=True) if rope == "rope"]
         if len(set(rope_dims)) > 1:
-            raise ValueError(f"All rope layers must have the same head size. Got {rope_dims}")
+            msg = f"All rope layers must have the same head size. Got {rope_dims}"
+            raise ValueError(msg)
         self.rope_dim = rope_dims[0]
 
         # Set default rope_parameters if not provided
         if rope_parameters is None:
             self.rope_parameters = RopeParameters(
-                rope_type = "default",
-                rope_theta = 10_000,
+                rope_type="default",
+                rope_theta=10_000,
             )
 
         self.pad_token_id = pad_token_id

modeling_domynedge.py

@@ -1,13 +1,59 @@
-from .configuration_domynedge import DomynEdgeConfig
-import transformers.modeling_outputs
-import transformers.modeling_utils
-import torch.nn.attention.varlen
-import transformers.utils
-import transformers
+# SPDX-FileCopyrightText: 2026 Domyn
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file includes third-party code:
+#
+# - `IndexFirstAxis`, `IndexPutFirstAxis`, `unpad_input` and `pad_input` are adapted from
+#   https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
+#   and are distributed under the following license:
+#
+#   BSD 3-Clause License
+#
+#   Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+#   All rights reserved.
+#
+#   Redistribution and use in source and binary forms, with or without
+#   modification, are permitted provided that the following conditions are met:
+#
+#   * Redistributions of source code must retain the above copyright notice, this
+#     list of conditions and the following disclaimer.
+#
+#   * Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+#   * Neither the name of the copyright holder nor the names of its
+#     contributors may be used to endorse or promote products derived from
+#     this software without specific prior written permission.
+#
+#   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#   AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#   IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#   DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#   FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#   DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#   SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#   CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#   OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#   OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# - `DomynEdgeRotaryEmbedding.compute_default_rope_parameters`, `rotate_half` and
+#   `apply_rotary_pos_emb` are adapted from HuggingFace transformers
+#   (https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py),
+#   Copyright The HuggingFace Team, licensed under the Apache License, Version 2.0.
+#   Modified for DomynEdge.
+
 import dataclasses
-import einops
 import typing
+
+import einops
 import torch
+import torch.nn.attention.varlen
+import transformers
+import transformers.modeling_outputs
+
+from .configuration_domynedge import DomynEdgeConfig
+
 
 @dataclasses.dataclass
 class VarlenMetadataMask:
@@ -16,17 +62,18 @@ class VarlenMetadataMask:
     max_q: int
     max_k: int
 
+
 class IndexFirstAxis(torch.autograd.Function):
-    """ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py """
+    """Gather the rows of the first axis at `indices`, scattering the gradient back in backward."""
 
     @staticmethod
-    def forward(ctx, input, indices):
+    def forward(ctx, input_, indices):
         ctx.save_for_backward(indices)
-        assert input.ndim >= 2
-        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        assert input_.ndim >= 2
+        ctx.first_axis_dim, other_shape = input_.shape[0], input_.shape[1:]
         second_dim = other_shape.numel()
         return torch.gather(
-            einops.rearrange(input, "b ... -> b (...)"), 0, einops.repeat(indices, "z -> z d", d=second_dim)
+            einops.rearrange(input_, "b ... -> b (...)"), 0, einops.repeat(indices, "z -> z d", d=second_dim)
         ).reshape(-1, *other_shape)
 
     @staticmethod
@@ -43,17 +90,16 @@ class IndexFirstAxis(torch.autograd.Function):
         grad_input.scatter_(0, einops.repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
         return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
 
+
 class IndexPutFirstAxis(torch.autograd.Function):
-    """ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py """
+    """Scatter `values` into a zero tensor at the first-axis `indices`, gathering the gradient back in backward."""
 
     @staticmethod
     def forward(ctx, values, indices, first_axis_dim):
         ctx.save_for_backward(indices)
         assert indices.ndim == 1
         assert values.ndim >= 2
-        output = torch.zeros(
-            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
-        )
+        output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
         output[indices] = values
         return output
 
@@ -64,12 +110,15 @@ class IndexPutFirstAxis(torch.autograd.Function):
         return grad_values, None, None
 
 
-def unpad_input(hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py 
+def unpad_input(
+    hidden_states: torch.Tensor, attention_mask: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Remove the padding tokens, packing the valid tokens of all sequences along one axis.
 
     Arguments:
         hidden_states: (batch, seqlen, ...)
         attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+
     Return:
         hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask.
         indices: (total_nnz), the indices of masked tokens from the flattened input sequence.
@@ -88,14 +137,16 @@ def unpad_input(hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> tu
         max_seqlen,
     )
 
+
 def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch_size: int, seqlen: int) -> torch.Tensor:
-    """ https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py 
+    """Restore the padded (batch, seqlen, ...) layout from packed tokens, the inverse of `unpad_input`.
 
     Arguments:
         hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
         indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.
-        batch: int, batch size for the padded sequence.
+        batch_size: int, batch size for the padded sequence.
         seqlen: int, maximum sequence length for the padded sequence.
+
     Return:
         hidden_states: (batch, seqlen, ...)
 
@@ -103,6 +154,7 @@ def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch_size: in
     output = IndexPutFirstAxis.apply(hidden_states, indices, batch_size * seqlen)
     return einops.rearrange(output, "(b s) ... -> b s ...", b=batch_size)
 
+
 class DomynEdgeFeedForward(torch.nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -116,6 +168,7 @@ class DomynEdgeFeedForward(torch.nn.Module):
     def forward(self, x):
         return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
 
+
 class DomynEdgeRotaryEmbedding(torch.nn.Module):
     inv_freq: torch.Tensor  # fix linting for `register_buffer`
 
@@ -137,17 +190,15 @@ class DomynEdgeRotaryEmbedding(torch.nn.Module):
     def compute_default_rope_parameters(
         config: DomynEdgeConfig | None = None,
         device: typing.Optional["torch.device"] = None,
-        seq_len: int | None = None,
     ) -> tuple["torch.Tensor", float]:
-        """
-        Computes the inverse frequencies according to the original RoPE implementation
+        """Computes the inverse frequencies according to the original RoPE implementation.
+
         Args:
             config ([`~transformers.PreTrainedConfig`]):
                 The model configuration.
             device (`torch.device`):
                 The device to use for initialization of the inverse frequencies.
-            seq_len (`int`, *optional*):
-                The current sequence length. Unused for this type of RoPE.
+
         Returns:
             Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
             post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
@@ -168,6 +219,11 @@ class DomynEdgeRotaryEmbedding(torch.nn.Module):
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
         position_ids_expanded = position_ids[:, None, :].float()
 
+        # TODO: the angles lose precision in two cases, both wrong at large positions:
+        # - under torch.autocast(bfloat16) the matmul below runs in bf16 despite the `.float()` casts;
+        # - after model.to(torch.bfloat16) / model.bfloat16() the `inv_freq` buffer itself is cast to bf16.
+        # Loading with `from_pretrained(dtype=torch.bfloat16)` and no autocast is not affected. Candidate fix:
+        # recompute `inv_freq` in fp32 here, with autocast disabled, instead of reading the buffer.
         freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
         emb = torch.cat((freqs, freqs), dim=-1)
         cos = emb.cos() * self.attention_scaling
@@ -198,6 +254,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
             k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
             cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
             the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
@@ -208,23 +265,36 @@ def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     k_embed = (kf * cos) + (rotate_half(kf) * sin)
     return q_embed.to(q.dtype), k_embed.to(k.dtype)
 
+
 class DomynEdgeAttention(torch.nn.Module):
-    """ Grouped Query Attention with support for sliding window """
+    """Grouped Query Attention with support for sliding window."""
 
     def __init__(self, config: DomynEdgeConfig, layer_idx: int):
         super().__init__()
         self.config = config
         self.head_dim = config.head_sizes[layer_idx]
-        self.scaling = self.head_dim ** -0.5
+        self.scaling = self.head_dim**-0.5
         self.window_size = config.window_sizes[layer_idx]
-        self.apply_rope = apply_rotary_pos_emb if config.position_layers[layer_idx] == "rope" else lambda q, k, cos, sin, unsqueeze_dim: (q, k)
+        self.apply_rope = (
+            apply_rotary_pos_emb
+            if config.position_layers[layer_idx] == "rope"
+            else lambda q, k, cos, sin, unsqueeze_dim: (q, k)  # noqa: ARG005
+        )
         self.layer_idx = layer_idx
         self.enable_gqa = config.num_attention_heads > config.num_key_value_heads
 
-        self.q_proj = torch.nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias)
-        self.k_proj = torch.nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
-        self.v_proj = torch.nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
-        self.o_proj = torch.nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
+        self.q_proj = torch.nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = torch.nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = torch.nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = torch.nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
         self.q_norm = torch.nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps, elementwise_affine=False)
         self.k_norm = torch.nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps, elementwise_affine=True)
 
@@ -234,13 +304,17 @@ class DomynEdgeAttention(torch.nn.Module):
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor,
         past_key_values: transformers.Cache | None = None,
-        **kwargs,
+        **_kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
 
-        batch_size, sequence_length, hidden_size = hidden_states.shape
+        batch_size, sequence_length, _ = hidden_states.shape
 
-        q = self.q_norm(self.q_proj(hidden_states).view(batch_size, sequence_length, self.config.num_attention_heads, self.head_dim))
-        k = self.k_norm(self.k_proj(hidden_states).view(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim))
+        q = self.q_norm(
+            self.q_proj(hidden_states).view(batch_size, sequence_length, self.config.num_attention_heads, self.head_dim)
+        )
+        k = self.k_norm(
+            self.k_proj(hidden_states).view(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
+        )
         v = self.v_proj(hidden_states).view(batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim)
 
         cos, sin = position_embeddings
@@ -252,33 +326,33 @@ class DomynEdgeAttention(torch.nn.Module):
             k, v = k.transpose(1, 2), v.transpose(1, 2)
 
         if q.size(1) > 1:
-            q, q_indices, q_cu_seqlen, q_maxseq = unpad_input(q, attention_mask)
+            q, q_indices, q_cu_seqlen, q_maxseq = unpad_input(q, attention_mask[:, -sequence_length:])
         else:
             q_maxseq = 1
             q_cu_seqlen = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device)
             q_indices = q_cu_seqlen[:-1]
-            q = q.squeeze(1) 
+            q = q.squeeze(1)
 
-        k, k_indices, k_cu_seqlen, k_maxseq = unpad_input(k, attention_mask)
-        v, v_indices, v_cu_seqlen, v_maxseq = unpad_input(v, attention_mask)
+        k, _, k_cu_seqlen, k_maxseq = unpad_input(k, attention_mask)
+        v, *_ = unpad_input(v, attention_mask)
 
         attn_output = torch.nn.attention.varlen.varlen_attn(
             q if q.dtype in [torch.float16, torch.bfloat16] else q.to(torch.bfloat16),
             k if k.dtype in [torch.float16, torch.bfloat16] else k.to(torch.bfloat16),
             v if v.dtype in [torch.float16, torch.bfloat16] else v.to(torch.bfloat16),
-            cu_seq_q = q_cu_seqlen,
-            cu_seq_k = k_cu_seqlen,
-            max_q = q_maxseq,
-            max_k = k_maxseq,
-            scale = self.scaling,
-            window_size = self.window_size,
-            enable_gqa = self.enable_gqa,
+            cu_seq_q=q_cu_seqlen,
+            cu_seq_k=k_cu_seqlen,
+            max_q=q_maxseq,
+            max_k=k_maxseq,
+            scale=self.scaling,
+            window_size=self.window_size,
+            enable_gqa=self.enable_gqa,
         ).to(q.dtype)
 
         attn_output = pad_input(attn_output, q_indices, batch_size, sequence_length)
         attn_output = attn_output.view(batch_size, sequence_length, self.config.num_attention_heads * self.head_dim)
-        attn_output = self.o_proj(attn_output)
-        return attn_output
+        return self.o_proj(attn_output)
+
 
 class DomynEdgeDecoderLayer(torch.nn.Module):
     def __init__(self, config: DomynEdgeConfig, layer_idx: int):
@@ -305,24 +379,21 @@ class DomynEdgeDecoderLayer(torch.nn.Module):
         # Self Attention
         hidden_states = hidden_states + self.post_attention_norm(
             self.attention(
-                hidden_states = self.pre_attention_norm(hidden_states),
-                attention_mask = attention_mask,
-                past_key_values = past_key_values,
-                position_embeddings = position_embeddings,
+                hidden_states=self.pre_attention_norm(hidden_states),
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
                 **kwargs,
             )
         )
 
         ## Fully Connected
-        hidden_states = hidden_states + self.post_feedforward_norm(
-            self.feedforward(
-                self.pre_feedforward_norm(hidden_states)
-            )
-        )
+        return hidden_states + self.post_feedforward_norm(self.feedforward(self.pre_feedforward_norm(hidden_states)))
 
-        return hidden_states
 
 class DomynEdgePreTrainedModel(transformers.PreTrainedModel):
+    """Base class wiring DomynEdge into the HuggingFace `PreTrainedModel` machinery."""
+
     config: DomynEdgeConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = False
@@ -338,13 +409,18 @@ class DomynEdgePreTrainedModel(transformers.PreTrainedModel):
 
 
 class DomynEdgeModel(DomynEdgePreTrainedModel):
+    """DomynEdge decoder stack without the language modeling head."""
+
     def __init__(self, config: DomynEdgeConfig):
+        """Build the embeddings, decoder layers, final norm and rotary embedding from `config`."""
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
         self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = torch.nn.ModuleList([DomynEdgeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.layers = torch.nn.ModuleList(
+            [DomynEdgeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
         self.norm = torch.nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = DomynEdgeRotaryEmbedding(config=config)
 
@@ -361,9 +437,10 @@ class DomynEdgeModel(DomynEdgePreTrainedModel):
         use_cache: bool | None = None,
         **kwargs,
     ) -> transformers.modeling_outputs.BaseModelOutputWithPast:
-
+        """Run the decoder stack and return the final hidden states (and the cache when `use_cache` is set)."""
         if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+            msg = "You must specify exactly one of input_ids or inputs_embeds"
+            raise ValueError(msg)
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
@@ -371,8 +448,9 @@ class DomynEdgeModel(DomynEdgePreTrainedModel):
         if use_cache and past_key_values is None:
             past_key_values = transformers.DynamicCache(config=self.config)
 
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+
         if position_ids is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
             position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
             position_ids = position_ids.unsqueeze(0)
 
@@ -380,7 +458,11 @@ class DomynEdgeModel(DomynEdgePreTrainedModel):
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
         if attention_mask is None:
-            attention_mask = torch.ones(hidden_states.shape[:2], device=hidden_states.device, dtype=torch.bool)
+            attention_mask = torch.ones(
+                (hidden_states.shape[0], past_seen_tokens + hidden_states.shape[1]),
+                device=hidden_states.device,
+                dtype=torch.bool,
+            )
 
         for decoder_layer in self.layers:
             hidden_states = decoder_layer(
@@ -397,10 +479,14 @@ class DomynEdgeModel(DomynEdgePreTrainedModel):
             past_key_values=past_key_values if use_cache else None,
         )
 
+
 class DomynEdgeForCausalLM(DomynEdgePreTrainedModel, transformers.GenerationMixin):
+    """DomynEdge model with a language modeling head tied to the input embeddings."""
+
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
 
     def __init__(self, config):
+        """Build the decoder stack and the language modeling head from `config`."""
         super().__init__(config)
         self.model = DomynEdgeModel(config)
         self.vocab_size = config.vocab_size
@@ -419,9 +505,9 @@ class DomynEdgeForCausalLM(DomynEdgePreTrainedModel, transformers.GenerationMixi
         labels: torch.LongTensor | None = None,
         use_cache: bool | None = None,
         logits_to_keep: int | torch.Tensor = 0,
-        **kwargs
+        **kwargs,
     ) -> transformers.modeling_outputs.CausalLMOutputWithPast:
-
+        """Compute logits for the last `logits_to_keep` positions, plus the loss when `labels` are given."""
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -436,7 +522,11 @@ class DomynEdgeForCausalLM(DomynEdgePreTrainedModel, transformers.GenerationMixi
         # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.lm_head(hidden_states[:, slice_indices, :])
-        loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) if labels is not None else None
+        loss = (
+            self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+            if labels is not None
+            else None
+        )
 
         return transformers.modeling_outputs.CausalLMOutputWithPast(
             loss=loss,
@@ -446,6 +536,7 @@ class DomynEdgeForCausalLM(DomynEdgePreTrainedModel, transformers.GenerationMixi
             attentions=outputs.attentions,
         )
 
+
 __all__ = [
     "DomynEdgeForCausalLM",
     "DomynEdgeModel",