Causal LM

unsloth/kernels/__init__.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .cross_entropy_loss import fast_cross_entropy_loss
+from .cross_entropy_loss import (
+    fast_cross_entropy_loss,
+    patch_llama_for_causal_lm,
+    unpatch_llama_for_causal_lm,
+)
 from .rms_layernorm import (
     fast_rms_layernorm,
     patch_rms_layernorm,

unsloth/kernels/cross_entropy_loss.py

@@ -375,3 +375,82 @@ def fast_cross_entropy_loss(
     n_items = torch.count_nonzero(labels != -100)
     return loss.sum() / n_items
 pass
+
+
+from transformers.models.llama.modeling_llama import LlamaForCausalLM
+def patch_llama_for_causal_lm():
+    import transformers.models.llama.modeling_llama
+    from transformers.models.llama.modeling_llama import (
+        CausalLMOutputWithPast,
+        Optional,
+        Union,
+        Cache,
+        List,
+        Tuple,
+    )
+    import inspect, re
+    function = inspect.getsource(transformers.models.llama.modeling_llama.LlamaForCausalLM.forward)
+    function = function.split("\n")
+    i = re.match(r"[ ]{1,}", function[0]).span(0)[1]
+    function = [x[i:] for x in function]
+    function = "\n".join(function)
+    function = function[function.find("def forward"):]
+    replacement = """    loss = None
+    logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
+    logit_scaling     = getattr(self.config, "logit_scale", 0)
+    if labels is not None:
+        shift_logits = logits
+        if not hasattr(self, "extra_ignored_labels"):
+            # Fixes https://github.com/unslothai/unsloth/issues/10
+            self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
+        pass
+        
+        shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
+        loss = fast_cross_entropy_loss(
+            logits = shift_logits,
+            labels = shift_labels,
+            logit_softcapping = logit_softcapping,
+            logit_scaling     = logit_scaling,
+        )
+    else:
+        if logit_scaling != 0:
+            if logits.requires_grad:
+                logits = logit_scaling * logits
+            else:
+                logits *= logit_scaling
+            pass
+        pass
+        if logit_softcapping != 0:
+            if logits.requires_grad:
+                logits = (1.0 / logit_softcapping) * logits
+                logits = torch.tanh(logits)
+                logits = logit_softcapping * logits
+            else:
+                logits *= (1.0 / logit_softcapping)
+                torch.tanh(logits, out = logits)
+                logits *= logit_softcapping
+            pass
+        pass
+    pass
+    """
+
+    function = \
+        function[:function.find("    loss = None")] + \
+        replacement + \
+        function[ function.find("    if not return_dict"):]
+    function = function.replace("logits = logits.float()", "\n")
+
+    patched_function = f"class Unsloth_LlamaForCausalLM(LlamaForCausalLM):\n"\
+    f"    {function}\n"
+
+    exec(patched_function)
+    transformers.models.llama.modeling_llama.LlamaForCausalLM = Unsloth_LlamaForCausalLM
+    return
+pass
+
+
+def unpatch_llama_for_causal_lm():
+    import transformers.models.llama.modeling_llama
+    transformers.models.llama.modeling_llama.LlamaForCausalLM = LlamaForCausalLM
+    return
+pass