fix: load checkpoints with torch.load(weights_only=True)

litgpt/api.py

@@ -115,7 +115,7 @@ def trainer_setup(self, trainer_ckpt: Path | None = None) -> None:
             self.load_state_dict(state_dict, strict=True)
 
         elif self.checkpoint_dir is not None:
-            state_dict = torch.load(self.checkpoint_dir / "lit_model.pth", weights_only=False)
+            state_dict = torch.load(self.checkpoint_dir / "lit_model.pth", weights_only=True)
             self.load_state_dict(state_dict, strict=False)
 
         else:
@@ -395,7 +395,7 @@ def distribute(
 
             if generate_strategy == "sequential":
                 state_dict = torch.load(
-                    str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu", weights_only=False
+                    str(self.checkpoint_dir / "lit_model.pth"), mmap=True, map_location="cpu", weights_only=True
                 )
                 model.load_state_dict(state_dict, assign=True)
                 model = fabric.setup_module(model, move_to_device=False)
@@ -419,7 +419,7 @@ def distribute(
                             str(self.checkpoint_dir / "lit_model.pth"),
                             mmap=True,
                             map_location="cpu",
-                            weights_only=False,
+                            weights_only=True,
                         )
                         model.load_state_dict(state_dict, assign=True)
 

litgpt/eval/evaluate.py

@@ -98,7 +98,7 @@ def convert_and_evaluate(
         # Hack: LitGPT's conversion doesn't save a pickle file that is compatible to be loaded with
         # `torch.load(..., weights_only=True)`, which is a requirement in HFLM.
         # So we're `torch.load`-ing and `torch.save`-ing it again to work around this.
-        state_dict = torch.load(out_dir / "model.pth")
+        state_dict = torch.load(out_dir / "model.pth", weights_only=True)
         torch.save(state_dict, model_path)
         os.remove(out_dir / "model.pth")
 

litgpt/generate/sequentially.py

@@ -239,7 +239,7 @@ def main(
     print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
 
     t0 = time.perf_counter()
-    state_dict = torch.load(str(checkpoint_path), mmap=True, map_location="cpu")
+    state_dict = torch.load(str(checkpoint_path), mmap=True, map_location="cpu", weights_only=True)
     # TODO: this assumes that the model fits on CPU. Use lazy_load and make the materialization checkpoint aware
     model.load_state_dict(state_dict, assign=True)
     print(f"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)

litgpt/generate/tp.py

@@ -204,7 +204,7 @@ def main(
     for rank in range(fabric.world_size):
         if fabric.global_rank == rank:
             t0 = time.perf_counter()
-            state_dict = torch.load(str(checkpoint_path), mmap=True, map_location="cpu")
+            state_dict = torch.load(str(checkpoint_path), mmap=True, map_location="cpu", weights_only=True)
             model.load_state_dict(state_dict, assign=True)
             print(f"[{rank}] Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
 

litgpt/scripts/convert_pretrained_checkpoint.py

@@ -37,7 +37,7 @@ def convert_pretrained_checkpoint(checkpoint_dir: Path, output_dir: Path) -> Non
     # Extract the model state dict and save to output folder
     with incremental_save(output_checkpoint_file) as saver:
         print("Processing", checkpoint_file)
-        full_checkpoint = torch.load(str(checkpoint_file), mmap=True)
+        full_checkpoint = torch.load(str(checkpoint_file), mmap=True, weights_only=True)
         loaded_state_dict = full_checkpoint["model"]
         converted_state_dict = {}
         for param_name, param in loaded_state_dict.items():

litgpt/scripts/merge_lora.py

@@ -63,8 +63,8 @@ def merge_lora(
         model.sin = None
 
     lora_path = checkpoint_dir / "lit_model.pth.lora"
-    pretrained_checkpoint = torch.load(str(pretrained_checkpoint_dir / "lit_model.pth"), mmap=True)
-    lora_checkpoint = torch.load(str(lora_path), mmap=True)
+    pretrained_checkpoint = torch.load(str(pretrained_checkpoint_dir / "lit_model.pth"), mmap=True, weights_only=True)
+    lora_checkpoint = torch.load(str(lora_path), mmap=True, weights_only=True)
     lora_checkpoint = lora_checkpoint.get("model", lora_checkpoint)
 
     # Merge LoRA weights into the base model

litgpt/utils.py

@@ -384,7 +384,7 @@ def load_checkpoint(fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, s
     if isinstance(fabric.strategy, FSDPStrategy):
         fabric.load_raw(checkpoint_path, model, strict=strict)
     elif isinstance(fabric.strategy, ModelParallelStrategy):
-        state_dict = torch.load(checkpoint_path, mmap=True)
+        state_dict = torch.load(checkpoint_path, mmap=True, weights_only=True)
         load_from_full_model_state_dict(
             model=model,
             full_sd=state_dict,