Remove exception fall back on checkpoint loading

.ci/scripts/test_model.sh

@@ -96,15 +96,15 @@ test_model() {
       bash examples/models/llama/install_requirements.sh
       # Test export_llama script: python3 -m examples.models.llama.export_llama.
       # Use Llama random checkpoint with Qwen 2.5 1.5b model configuration.
-      "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama --model "${MODEL_NAME}" -c examples/models/llama/params/demo_rand_params.pth -p examples/models/qwen2_5/1_5b_config.json
+      "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama --model "${MODEL_NAME}" -p examples/models/qwen2_5/1_5b_config.json
       rm "./${MODEL_NAME}.pte"
       return  # Skip running with portable executor runnner since portable doesn't support Qwen's biased linears.
   fi
   if [[ "${MODEL_NAME}" == "phi_4_mini" ]]; then
       # Install requirements for export_llama
       bash examples/models/llama/install_requirements.sh
       # Test export_llama script: python3 -m examples.models.llama.export_llama.
-      "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama --model "${MODEL_NAME}" -c examples/models/llama/params/demo_rand_params.pth -p examples/models/phi_4_mini/config.json
+      "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama --model "${MODEL_NAME}" -p examples/models/phi_4_mini/config.json
       run_portable_executor_runner
       rm "./${MODEL_NAME}.pte"
       return

examples/models/llama/model.py

@@ -244,33 +244,24 @@ def __init__(self, **kwargs):
             )
 
         missing, unexpected = None, None
-        try:
-            # assign=True: load params/buffers by assignment instead of performing an in-place copy.
-            # Because we are using device="meta", tensors do not have memory associated with them
-            # and an in-place copy is a no-op. Use assign=True in load_state_dict for this scenario.
-
-            # Also, the checkpoint is loaded and dtype promoted to the transformer's dtype, which is
-            # by default initialized to fp32. This is fine because every other supported type
-            # losslessly converts to fp32, so we don't lose precision here.
-            if checkpoint:
-                missing, unexpected = self.model_.load_state_dict(
-                    checkpoint,
-                    strict=False,
-                    assign=True,
-                )  # self.model_ = Transformer(gptconf)
-            else:
-                print("Checkpoint not provided, defaulting weights to zeros.")
-                self.model_.to_empty(device="cpu")
-                for p in self.model_.parameters():
-                    p.data.fill_(0)
-                for b in self.model_.buffers():
-                    b.data.fill_(0)
-        except RuntimeError as e:
-            print(
-                f"Could not load checkpoint into mode and will defaulting weights to zeros due to error: {e}."
-            )
-            # Need to provide concrete (empty) values for meta-initialized tensors for quantization.
+        # assign=True: load params/buffers by assignment instead of performing an in-place copy.
+        # Because we are using device="meta", tensors do not have memory associated with them
+        # and an in-place copy is a no-op. Use assign=True in load_state_dict for this scenario.
+
+        # Also, the checkpoint is loaded and dtype promoted to the transformer's dtype, which is
+        # by default initialized to fp32. This is fine because every other supported type
+        # losslessly converts to fp32, so we don't lose precision here.
+        if checkpoint:
+            missing, unexpected = self.model_.load_state_dict(
+                checkpoint,
+                strict=False,
+                assign=True,
+            )  # self.model_ = Transformer(gptconf)
+        else:
+            print("Checkpoint not provided, defaulting weights to zeros.")
             self.model_.to_empty(device="cpu")
+            # Need to provide concrete values for meta-initialized tensors for quantization.
+            # otherwise it is just filled with nan's.
             for p in self.model_.parameters():
                 p.data.fill_(0)
             for b in self.model_.buffers():