Materialize meta buffers during model load

src/evaluator.py

@@ -71,6 +71,28 @@ def _load_model_state(model, state_dict, strict=True):
         return model.load_state_dict(state_dict, strict=strict)
 
 
+def _set_module_attr(root, dotted_name, value):
+    module = root
+    parts = dotted_name.split(".")
+    for part in parts[:-1]:
+        module = getattr(module, part)
+    setattr(module, parts[-1], value)
+
+
+def _materialize_meta_buffers(model, device):
+    """Create real tensors for buffers that are not stored in older checkpoints."""
+    for name, buf in list(model.named_buffers()):
+        if not getattr(buf, "is_meta", False):
+            continue
+        shape = tuple(buf.shape)
+        dtype = buf.dtype
+        if name.endswith("position_ids") and len(shape) == 2:
+            value = torch.arange(shape[1], dtype=dtype, device=device).expand(shape[0], -1)
+        else:
+            value = torch.zeros(shape, dtype=dtype, device=device)
+        _set_module_attr(model, name, value)
+
+
 
 
 def _refresh_runtime_buffers(model, device):
@@ -87,12 +109,12 @@ def _refresh_runtime_buffers(model, device):
             prior = torch.tensor(prior_rows, dtype=torch.float, device=device)
             prior = prior - prior.mean(dim=-1, keepdim=True)
             module.source_prior_bias = prior
+    _materialize_meta_buffers(model, device)
     meta_params = [name for name, param in model.named_parameters() if getattr(param, "is_meta", False)]
-    meta_buffers = [name for name, buf in model.named_buffers() if getattr(buf, "is_meta", False)]
-    if meta_params or meta_buffers:
+    if meta_params:
         raise RuntimeError(
-            "Checkpoint did not materialize model tensors. First meta tensors: "
-            + ", ".join((meta_params + meta_buffers)[:20])
+            "Checkpoint did not materialize model parameters. First meta parameters: "
+            + ", ".join(meta_params[:20])
         )
     return model