Update modeling_neollm.py

modeling_neollm.py

@@ -6095,14 +6095,26 @@ _EXTENDED_CCE_INSTALL = (
     "cut-cross-entropy @ "
     "git+https://github.com/Kitsunp/ml-cross-entropy.git@main"
 )
+_EXTENDED_CCE_PARAMETERS = (
+    frozenset(inspect.signature(linear_cross_entropy).parameters)
+    if linear_cross_entropy is not None
+    else frozenset()
+)
+
+# The compiler-safe CCE operator returns compact saved tensors whose first
+# dimension is the number of non-padding labels.  Let Dynamo represent that
+# value as an unbacked SymInt; otherwise it splits the model immediately after
+# the otherwise opaque custom op.  The data-dependent compaction itself remains
+# inside CCE and does not become an Inductor graph.
+if linear_cross_entropy is not None:
+    torch._dynamo.config.capture_dynamic_output_shape_ops = True
 
 
 def _require_extended_cce_options(*option_names: str) -> None:
     """Fail clearly only when an explicitly enabled CCE extension is unavailable."""
     if linear_cross_entropy is None:
         return
-    parameters = inspect.signature(linear_cross_entropy).parameters
-    missing = [name for name in option_names if name not in parameters]
+    missing = [name for name in option_names if name not in _EXTENDED_CCE_PARAMETERS]
     if missing:
         raise ImportError(
             "The installed cut-cross-entropy package does not provide the enabled "
@@ -6112,7 +6124,6 @@ def _require_extended_cce_options(*option_names: str) -> None:
         )
 
 
-@torch.compiler.disable
 def compute_cce_loss(
     hidden_states,
     labels,
@@ -6126,7 +6137,13 @@ def compute_cce_loss(
     mu_loss_lambda=1e-4,
     return_loss_metrics=False,
 ):
-    """CCE loss excluded from torch.compile, preserving the configured CCE impl."""
+    """CCE loss with a compiler-visible custom-op boundary in the extended backend.
+
+    The backend keeps data-dependent label compaction and Triton autotuning
+    opaque while exposing one forward/backward operator to Dynamo.  Older CCE
+    installations retain their normal eager fallback and fail clearly only
+    when an explicitly requested extension is unavailable.
+    """
     if linear_cross_entropy is None:
         raise ImportError(
             "NeoLLM was configured with `ntp_loss_backend='cce'`, but "
@@ -7928,7 +7945,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
-        meap_seed: Optional[int] = None,
+        meap_seed: Optional[Union[int, torch.Tensor]] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
@@ -7957,11 +7974,18 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             mask_token_id = getattr(self.config, "meap_mask_token_id", None)
             if mask_token_id is None:
                 raise ValueError("`meap_mask_token_id` is required when MEAP is enabled.")
-            effective_meap_seed = int(
-                getattr(self.config, "meap_seed", 0)
-                if meap_seed is None
-                else meap_seed
-            )
+            if meap_seed is None:
+                effective_meap_seed = input_ids.new_tensor(
+                    int(getattr(self.config, "meap_seed", 0)), dtype=torch.long
+                )
+            elif isinstance(meap_seed, torch.Tensor):
+                effective_meap_seed = meap_seed.to(
+                    device=input_ids.device, dtype=torch.long
+                )
+            else:
+                effective_meap_seed = input_ids.new_tensor(
+                    int(meap_seed), dtype=torch.long
+                )
             eligible_mask = (
                 attention_mask.to(device=input_ids.device, dtype=torch.bool)
                 if attention_mask is not None
@@ -7989,9 +8013,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 masked_count.float() / eligible_count.clamp_min(1).float()
             ).detach()
             self._last_meap_masked_tokens = masked_count.detach()
-            self._last_meap_seed = input_ids.new_tensor(
-                effective_meap_seed, dtype=torch.long
-            )
+            self._last_meap_seed = effective_meap_seed.detach()
 
         tweo_enabled = (
             bool(getattr(self.config, "use_tweo", False))