Fix freqs_cis being left uninitialized when loaded via meta-device init

amplify.py

@@ -206,14 +206,24 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
 
         self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
 
-        freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_protein_length * 2)
-
-        # Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
-        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+        # freqs_cis is derived from the config, not learned, so it is not saved in the state_dict. It is built lazily on
+        # first use rather than here: `from_pretrained` may construct the model on the meta device and then materialize
+        # it with uninitialized memory, filling in only the tensors present in the checkpoint. A buffer populated in
+        # __init__ would be left as garbage by that path, so it is registered as None and computed in forward().
+        self.register_buffer("freqs_cis", None, persistent=False)
 
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _ensure_freqs_cis(self, device: torch.device) -> torch.Tensor:
+        """Build the rotary embedding table on first use, then reuse it."""
+        if self.freqs_cis is None or self.freqs_cis.device != device:
+            self.freqs_cis = precompute_freqs_cis(
+                self.config.hidden_size // self.config.num_attention_heads,
+                self.config.max_protein_length * 2,
+            ).to(device)
+        return self.freqs_cis
+
     @classmethod
     def load(cls, checkpoint_path: str, config_path: str, vocab_path: str = None, tag: str = None):
 
@@ -277,10 +287,11 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
             assert input_ids.is_cuda, "Packing uses an implementation of flash-attention and is only supported on GPU."
 
         # RoPE
+        freqs_cis_table = self._ensure_freqs_cis(input_ids.device)
         if position_ids is not None:
-            freqs_cis = self.freqs_cis[position_ids]
+            freqs_cis = freqs_cis_table[position_ids]
         else:
-            freqs_cis = self.freqs_cis[: input_ids.shape[1]].unsqueeze(0).repeat(input_ids.shape[0], 1, 1)
+            freqs_cis = freqs_cis_table[: input_ids.shape[1]].unsqueeze(0).repeat(input_ids.shape[0], 1, 1)
 
         # Embedding
         x = self.encoder(input_ids)