Restore derived LNS scale buffers after checkpoint loading

modeling_neollm.py

@@ -6752,7 +6752,7 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        """Initialize only a legacy checkpoint that lacks ``knot_grid``.
+        """Restore derived buffers and support legacy ``knot_grid`` loading.
 
         Current checkpoints persist this buffer and must retain its exact saved
         values.  The previous compatibility path regenerated the grid after
@@ -6767,6 +6767,27 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             **kwargs,
         )
         missing_keys = set(loading_info.get("missing_keys", ()))
+        # LNS.scale is derived, non-persistent state. Low-memory loaders can
+        # materialize its meta storage without running the original constructor
+        # initializer. Rebuild the constant; never infer validity from finiteness.
+        for module in model.modules():
+            if isinstance(module, LNS):
+                scale_device = module.scale.device
+                if scale_device.type == "meta":
+                    scale_device = model.device
+                    if scale_device.type == "meta":
+                        scale_device = torch.device("cpu")
+                scale_dtype = (
+                    requested_dtype
+                    if isinstance(requested_dtype, torch.dtype)
+                    and requested_dtype.is_floating_point
+                    else module.scale.dtype
+                )
+                module._buffers["scale"] = torch.tensor(
+                    1.0 / math.sqrt(module.layer_idx),
+                    device=scale_device,
+                    dtype=torch.float32,
+                ).to(dtype=scale_dtype)
         for module_name, module in model.named_modules():
             if not isinstance(module, LeviathanGenerator):
                 continue