From 5f5ae7621354d4f4bf61490e3d7e82a4a873ca59 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Thu, 17 Jul 2025 15:32:38 -0400
Subject: [PATCH] add validation around cce + chunked_ce (#2932) [skip ci]

* add validation around cce + chunked_ce

* return on end of validation method
---
 src/axolotl/integrations/cut_cross_entropy/args.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/axolotl/integrations/cut_cross_entropy/args.py b/src/axolotl/integrations/cut_cross_entropy/args.py
index 2729ebe2e..22852479a 100644
--- a/src/axolotl/integrations/cut_cross_entropy/args.py
+++ b/src/axolotl/integrations/cut_cross_entropy/args.py
@@ -41,3 +41,13 @@ class CutCrossEntropyArgs(BaseModel):
             )
 
         return data
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_chunked_cross_entropy_not_set(cls, data):
+        if data.get("chunked_cross_entropy"):
+            raise ValueError(
+                "Cut Cross Entropy does not support chunked cross entropy. "
+                "Please set `chunked_cross_entropy` to `False` or disable Cut Cross Entropy."
+            )
+        return data