Update modeling_smartcoder_moe.py

modeling_smartcoder_moe.py

@@ -1,5 +1,4 @@
-
-# modeling_smartcoder_moe.py
+#modeling_smartcoder_moe.py
 
 #Architecture (from tensor inspection):
 #- vocab_size: 65536, hidden: 2048, layers: 40
@@ -17,6 +16,7 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils.checkpoint
 from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
@@ -223,11 +223,22 @@ class SmartCoderMoEModel(nn.Module):
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList([SmartCoderDecoderLayer(config) for _ in range(config.num_hidden_layers)])
         self.norm   = LayerNormWithBias(config.hidden_size, config.rms_norm_eps)
+        # Required for transformers' _set_gradient_checkpointing() to have
+        # something to toggle. Without this attribute + the checkpoint()
+        # call in forward(), declaring supports_gradient_checkpointing=True
+        # at the PreTrainedModel level is a lie transformers will catch and
+        # raise on -- which is exactly the ValueError this fixes.
+        self.gradient_checkpointing = True
 
     def forward(self, input_ids, attention_mask=None, **kwargs):
         hidden_states = self.embed_tokens(input_ids)
         for layer in self.layers:
-            hidden_states = layer(hidden_states, attention_mask=attention_mask)
+            if self.gradient_checkpointing and self.training:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    layer, hidden_states, attention_mask, use_reentrant=True
+                )
+            else:
+                hidden_states = layer(hidden_states, attention_mask=attention_mask)
         return self.norm(hidden_states)
 
 
@@ -254,6 +265,18 @@ class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
     def get_input_embeddings(self): return self.model.embed_tokens
     def get_output_embeddings(self): return self.lm_head
 
+    # transformers' _set_gradient_checkpointing (called by Unsloth/Trainer)
+    # looks for this attribute on the *PreTrainedModel* root, finds the
+    # submodule that has it, and toggles it. Exposing it here as a property
+    # delegating to self.model keeps both objects in sync.
+    @property
+    def gradient_checkpointing(self):
+        return self.model.gradient_checkpointing
+
+    @gradient_checkpointing.setter
+    def gradient_checkpointing(self, value):
+        self.model.gradient_checkpointing = value
+
     def forward(
         self,
         input_ids=None,