fix: shift logits/labels for proper causal LM loss (was predicting current token instead of next token)

open_mythos_hf/modeling.py

@@ -671,12 +671,14 @@ class OpenMythosForCausalLM(OpenMythosPreTrainedModel, GenerationMixin):
         # 6. LM Head
         logits = self.lm_head(x).float()
 
-        # 7. Loss
+        # 7. Loss (shifted for causal LM: logits[i] predicts token[i+1])
         loss = None
         if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
             loss = F.cross_entropy(
-                logits.view(-1, logits.shape[-1]),
-                labels.view(-1),
+                shift_logits.view(-1, shift_logits.shape[-1]),
+                shift_labels.view(-1),
                 ignore_index=-100,
             )