Update modeling_alinlight.py

modeling_alinlight.py

@@ -13,29 +13,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Optional, Tuple, List, Union
+from torch.utils.checkpoint import checkpoint
+
 from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
 from configuration_alinlight import AlinlightConfig
 
+# ==========================================
+# 1. BASE COMPONENTS
+# ==========================================
+
 class AlinlightRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.eps = eps
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor):
         input_dtype = x.dtype
         x = x.to(torch.float32)
         variance = x.pow(2).mean(-1, keepdim=True)
         x = x * torch.rsqrt(variance + self.eps)
         return self.weight * x.to(input_dtype)
 
+
 class AlinlightRotaryEmbedding(nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
         super().__init__()
@@ -43,16 +49,18 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.base = base
         self.max_position_embeddings = max_position_embeddings
         self.scaling_factor = scaling_factor
-        
 
         inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32).to(device) / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        
- 
         self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
- 
+        if (hasattr(self, 'cos_cached') and 
+            self.cos_cached.device == device and 
+            self.cos_cached.dtype == dtype and 
+            self.cos_cached.shape[0] >= seq_len):
+            return
+
         t = torch.arange(seq_len, device=device, dtype=torch.int64).type_as(self.inv_freq)
         t = t / self.scaling_factor
         freqs = torch.outer(t, self.inv_freq)
@@ -61,20 +69,20 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
-
         if seq_len > self.cos_cached.shape[0]:
-             self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-        
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
         return (
             self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device),
             self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
         )
 
-def rotate_half(x):
+
+def rotate_half(x: torch.Tensor):
     x1 = x[..., : x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2 :]
     return torch.cat((-x2, x1), dim=-1)
 
+
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
@@ -82,6 +90,11 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
 
+
+# ==========================================
+# 2. MLP
+# ==========================================
+
 class AlinlightMLP(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -91,9 +104,20 @@ class AlinlightMLP(nn.Module):
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = nn.SiLU()
+        self.pre_down_norm = AlinlightRMSNorm(self.intermediate_size, eps=config.rms_norm_eps)
+        
+        # Tag for specialized initialization
+        self.down_proj._is_residual_projection = True
 
     def forward(self, x):
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+        intermediate = self.pre_down_norm(intermediate)
+        return self.down_proj(intermediate)
+
+
+# ==========================================
+# 3. ATTENTION
+# ==========================================
 
 class AlinlightAttention(nn.Module):
     def __init__(self, config, layer_idx: Optional[int] = None):
@@ -106,71 +130,120 @@ class AlinlightAttention(nn.Module):
         self.num_key_value_heads = config.num_key_value_heads
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.sliding_window = config.sliding_window
-        
+        self.attention_dropout = config.attention_dropout
+
         self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
         self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        
+        # Tag for specialized initialization
+        self.o_proj._is_residual_projection = True
+
+        self.use_qk_norm = getattr(config, "use_qk_norm", True)
+        if self.use_qk_norm:
+            self.q_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        
+        self.attn_logit_softcapping = getattr(config, 'attn_logit_softcapping', None)
 
     def forward(
-        self, 
-        hidden_states, 
-        attention_mask=None, 
-        position_ids=None, 
-        past_key_value=None, 
-        output_attentions=False, 
-        use_cache=False, 
-        cos_sin=None
-    ):
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cos_sin: Optional[Tuple[torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+
         bsz, q_len, _ = hidden_states.size()
-        
+
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
-        
+
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 
+        if self.use_qk_norm:
+            query_states = self.q_norm(query_states)
+            key_states = self.k_norm(key_states)
+
+        # 1. RoPE (Applied before caching)
         if cos_sin is not None:
             cos, sin = cos_sin
             query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
+        # 2. KV Cache
         if past_key_value is not None:
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
+        kv_seq_len = key_states.shape[2]
 
-        if self.sliding_window is not None and key_states.shape[2] > self.sliding_window:
-            key_states = key_states[:, :, -self.sliding_window:, :]
-            value_states = value_states[:, :, -self.sliding_window:, :]
+        # 3. Sliding Window (Slicing)
+        if self.sliding_window is not None and kv_seq_len > self.sliding_window:
+            slicing_tokens = kv_seq_len - self.sliding_window
+            key_states = key_states[:, :, slicing_tokens:, :]
+            value_states = value_states[:, :, slicing_tokens:, :]
+            
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, :, :, slicing_tokens:]
 
         past_key_value = (key_states, value_states) if use_cache else None
-        
- 
+
+        # 4. GQA Repeat
         if self.num_key_value_groups > 1:
-            key_states = key_states[:, :, None, :, :].expand(
-                bsz, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_dim
-            ).reshape(bsz, self.num_heads, key_states.shape[-2], self.head_dim)
-            
-            value_states = value_states[:, :, None, :, :].expand(
-                bsz, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_dim
-            ).reshape(bsz, self.num_heads, value_states.shape[-2], self.head_dim)
+            key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+            value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
 
+        # 5. Attention Mechanism
+        attn_weights = None
+        
+        # We must use manual implementation if:
+        # a) Output weights are requested
+        # b) Soft-capping is enabled (SDPA doesn't support intermediate logit transforms)
+        if output_attentions or self.attn_logit_softcapping is not None:
+            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+            
+            if self.attn_logit_softcapping is not None:
+                attn_weights = self.attn_logit_softcapping * torch.tanh(attn_weights / self.attn_logit_softcapping)
+            
+            if attention_mask is not None:
+                attn_weights = attn_weights + attention_mask
+            
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+            
+            if not output_attentions:
+                # If we only calculated weights for soft-capping but user didn't ask for them, drop reference
+                attn_weights_for_output = None
+            else:
+                attn_weights_for_output = attn_weights
 
-        is_causal = q_len > 1
+            attn_weights_dropped = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+            attn_output = torch.matmul(attn_weights_dropped, value_states)
+        else:
+            # Fast Path (SDPA)
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                is_causal=False 
+            )
+            attn_weights_for_output = None
 
-        attn_output = F.scaled_dot_product_attention(
-            query_states, 
-            key_states, 
-            value_states, 
-            attn_mask=None, 
-            dropout_p=0.0, 
-            is_causal=is_causal 
-        )
-        
         attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
-        return self.o_proj(attn_output), None, past_key_value
+        return self.o_proj(attn_output), attn_weights_for_output, past_key_value
+
+
+# ==========================================
+# 4. DECODER LAYER & MODEL
+# ==========================================
 
 class AlinlightDecoderLayer(nn.Module):
     def __init__(self, config, layer_idx: int):
@@ -179,99 +252,218 @@ class AlinlightDecoderLayer(nn.Module):
         self.mlp = AlinlightMLP(config)
         self.input_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        self.resid_pdrop = getattr(config, 'resid_pdrop', 0.0)
+        self.resid_dropout = nn.Dropout(self.resid_pdrop) if self.resid_pdrop > 0 else nn.Identity()
 
-    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False, cos_sin=None):
+    def forward(
+        self, 
+        hidden_states, 
+        attention_mask=None, 
+        position_ids=None, 
+        past_key_value=None, 
+        output_attentions=False, 
+        use_cache=False, 
+        cos_sin=None
+    ):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states, _, present_key_value = self.self_attn(
-            hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cos_sin
+
+        hidden_states, attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cos_sin=cos_sin
         )
-        hidden_states = residual + hidden_states
-        
+        hidden_states = residual + self.resid_dropout(hidden_states)
+
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        return hidden_states, None, present_key_value
+        hidden_states = residual + self.resid_dropout(hidden_states)
+
+        return hidden_states, attn_weights, present_key_value
+
 
 class AlinlightModel(PreTrainedModel):
     config_class = AlinlightConfig
+    
     def __init__(self, config: AlinlightConfig):
         super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        
+        self.embed_scale = math.sqrt(config.hidden_size) if getattr(config, 'embed_scale', False) else 1.0
+        self.embed_dropout = nn.Dropout(config.embed_pdrop) if config.embed_pdrop > 0 else nn.Identity()
+
         self.layers = nn.ModuleList([AlinlightDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
         self.norm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        
+
         scaling_factor = 1.0
         if config.rope_scaling and config.rope_scaling.get("type") == "linear":
              scaling_factor = config.rope_scaling.get("factor", 1.0)
-             
+
         self.rotary_emb = AlinlightRotaryEmbedding(
-            config.hidden_size // config.num_attention_heads, 
-            max_position_embeddings=config.max_position_embeddings, 
-            base=config.rope_theta, 
+            config.hidden_size // config.num_attention_heads,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
             scaling_factor=scaling_factor
         )
+        self.gradient_checkpointing = False
         self.post_init()
 
     def get_input_embeddings(self): return self.embed_tokens
     def set_input_embeddings(self, value): self.embed_tokens = value
 
-    def forward(self, input_ids=None, past_key_values=None, use_cache=None, **kwargs):
-        if input_ids is not None:
-            inputs_embeds = self.embed_tokens(input_ids)
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        bsz, seq_len = input_shape
+        dtype = inputs_embeds.dtype
+        device = inputs_embeds.device
+
+        if attention_mask is not None:
+            current_mask = attention_mask[:, None, None, :].to(dtype=dtype)
+        else:
+            current_mask = torch.ones((bsz, 1, 1, seq_len), dtype=dtype, device=device)
+
+        if past_key_values_length > 0:
+            past_mask = torch.ones((bsz, 1, 1, past_key_values_length), dtype=dtype, device=device)
+            combined_mask = torch.cat([past_mask, current_mask], dim=-1)
         else:
-             inputs_embeds = kwargs.get("inputs_embeds")
+            combined_mask = current_mask
+
+        inverted_mask = (1.0 - combined_mask) * torch.finfo(dtype).min
 
+        if seq_len > 1:
+            causal_mask = torch.triu(
+                torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype),
+                diagonal=1
+            )
+            if past_key_values_length > 0:
+                past_causal = torch.zeros((seq_len, past_key_values_length), dtype=dtype, device=device)
+                causal_mask = torch.cat([past_causal, causal_mask], dim=-1)
+
+            causal_mask = causal_mask[None, None, :, :]
+            inverted_mask = inverted_mask + causal_mask
+
+        return inverted_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
         
-        seq_len = inputs_embeds.shape[1]
+        inputs_embeds = inputs_embeds * self.embed_scale
+        inputs_embeds = self.embed_dropout(inputs_embeds)
+
+        batch_size, seq_length = inputs_embeds.shape[:2]
+        past_key_values_length = 0
         if past_key_values is not None:
-             seq_len += past_key_values[0][0].shape[2]
-             
-        
-        cos, sin = self.rotary_emb(inputs_embeds, seq_len=seq_len)
-        
-        position_ids = kwargs.get("position_ids")
+             past_key_values_length = past_key_values[0][0].shape[2]
+
+        total_seq_len = seq_length + past_key_values_length
+        cos, sin = self.rotary_emb(inputs_embeds, seq_len=total_seq_len)
+
         if position_ids is None:
-             
-             position_ids = torch.arange(seq_len - inputs_embeds.shape[1], seq_len, dtype=torch.long, device=inputs_embeds.device)
-             position_ids = position_ids.unsqueeze(0).expand(inputs_embeds.shape[0], -1)
+             position_ids = torch.arange(
+                 past_key_values_length, total_seq_len, dtype=torch.long, device=inputs_embeds.device
+             ).unsqueeze(0).expand(batch_size, -1)
+
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+        )
 
         hidden_states = inputs_embeds
         next_decoder_cache = () if use_cache else None
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
 
         for idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
             past_key_value = past_key_values[idx] if past_key_values is not None else None
-            layer_outputs = layer(
-                hidden_states, 
-                position_ids=position_ids, 
-                past_key_value=past_key_value, 
-                use_cache=use_cache, 
-                cos_sin=(cos, sin)
-            )
+
+            if self.gradient_checkpointing and self.training:
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions=output_attentions, use_cache=False, cos_sin=(cos, sin))
+                    return custom_forward
+                layer_outputs = checkpoint(
+                    create_custom_forward(layer), hidden_states, attention_mask, position_ids, past_key_value, use_reentrant=True
+                )
+            else:
+                layer_outputs = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cos_sin=(cos, sin)
+                )
+
             hidden_states = layer_outputs[0]
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
             if use_cache:
                 next_decoder_cache += (layer_outputs[2],)
 
         hidden_states = self.norm(hidden_states)
-        
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_decoder_cache, all_hidden_states, all_self_attns] if v is not None)
+
         return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
-            past_key_values=next_decoder_cache
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
         )
 
+
+# ==========================================
+# 5. CAUSAL LM HEAD
+# ==========================================
+
 class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = AlinlightConfig
     _keys_to_ignore_on_load_missing = ["model.rotary_emb.inv_freq"]
+    _supports_gradient_checkpointing = True
 
     def __init__(self, config):
         super().__init__(config)
         self.model = AlinlightModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         
+        self.final_logit_softcapping = getattr(config, 'final_logit_softcapping', None)
+        self.z_loss_weight = getattr(config, 'z_loss_weight', 0.0)
+
         if config.tie_word_embeddings:
             self.lm_head.weight = self.model.embed_tokens.weight
-            
+
         self.post_init()
 
     def get_input_embeddings(self): return self.model.embed_tokens
@@ -279,37 +471,103 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     def get_output_embeddings(self): return self.lm_head
     def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
 
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-
-        if past_key_values:
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            # Scale down residual projections to improve training stability at depth
+            if getattr(module, '_is_residual_projection', False):
+                module.weight.data.normal_(mean=0.0, std=std / math.sqrt(2 * self.config.num_hidden_layers))
+            else:
+                module.weight.data.normal_(mean=0.0, std=std)
+            
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.model.gradient_checkpointing = True
+        self.config.use_cache = False
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        if past_key_values is not None:
             input_ids = input_ids[:, -1:]
-        
+
         position_ids = kwargs.get("position_ids", None)
         if position_ids is None:
             if past_key_values:
-
-                 past_length = past_key_values[0][0].shape[2]
-                 position_ids = torch.tensor([[past_length]], dtype=torch.long, device=input_ids.device)
+                position_ids = (attention_mask.long().sum(dim=-1) - 1).unsqueeze(-1)
             else:
-
-                 position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
+                position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
 
         return {
             "input_ids": input_ids,
             "past_key_values": past_key_values,
             "use_cache": True,
-            "position_ids": position_ids
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
         }
 
-    def forward(self, input_ids=None, past_key_values=None, labels=None, **kwargs):
-        outputs = self.model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
-        hidden_states = outputs.last_hidden_state
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs
+        )
+
+        hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-        
+
+        # Final Logit Soft-Capping
+        if self.final_logit_softcapping is not None:
+            logits = self.final_logit_softcapping * torch.tanh(logits / self.final_logit_softcapping)
+
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
-            loss = F.cross_entropy(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
-
-        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values)
+            
+            loss_fct = nn.CrossEntropyLoss()
+            ce_loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+            
+            # Z-Loss Regularization
+            if self.z_loss_weight > 0 and self.training:
+                # log(sum(exp(x)))^2
+                z_loss = torch.logsumexp(shift_logits, dim=-1).pow(2).mean()
+                loss = ce_loss + self.z_loss_weight * z_loss
+            else:
+                loss = ce_loss
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )