feat(server): load santacoder/starcoder models with safetensors

launcher/src/main.rs

@@ -546,11 +546,7 @@ enum LauncherError {
     WebserverCannotStart,
 }
 
-fn download_convert_model(
-    args: &Args,
-    auto_convert: bool,
-    running: Arc<AtomicBool>,
-) -> Result<(), LauncherError> {
+fn download_convert_model(args: &Args, running: Arc<AtomicBool>) -> Result<(), LauncherError> {
     let mut download_argv = vec![
         "text-generation-server".to_string(),
         "download-weights".to_string(),
@@ -562,11 +558,6 @@ fn download_convert_model(
         "--json-output".to_string(),
     ];
 
-    // Auto convert weights to safetensors
-    if auto_convert {
-        download_argv.push("--auto-convert".to_string());
-    }
-
     // Model optional revision
     if let Some(revision) = &args.revision {
         download_argv.push("--revision".to_string());
@@ -932,11 +923,8 @@ fn main() -> Result<(), LauncherError> {
     })
     .expect("Error setting Ctrl-C handler");
 
-    // auto_convert is only needed for sharded models as we do not require safetensors in
-    // single shard mode
-    let auto_convert = num_shard > 1;
     // Download and convert model weights
-    download_convert_model(&args, auto_convert, running.clone())?;
+    download_convert_model(&args, running.clone())?;
 
     // Shared shutdown bool
     let shutdown = Arc::new(Mutex::new(false));

server/text_generation_server/models/flash_santacoder.py

@@ -54,12 +54,7 @@ class FlashSantacoder(FlashCausalLM):
         )
 
         # We do not use from_pretrained as we modified the model internal module layout
-        try:
-            filenames = weight_files(model_id, revision, ".bin")
-        # Local files not found
-        except LocalEntryNotFoundError:
-            hub_files = weight_hub_files(model_id, revision, ".bin")
-            filenames = download_weights(hub_files, model_id, revision)
+        filenames = weight_files(model_id, revision, ".safetensors")
 
         with init_empty_weights():
             model = FlashSantacoderForCausalLM(config)
@@ -91,8 +86,11 @@ class FlashSantacoder(FlashCausalLM):
         transpose: bool,
     ):
         for filename in filenames:
-            state_dict = torch.load(filename, map_location="cpu")
-            for key, value in state_dict.items():
+            with safe_open(
+                    filename, framework="pt", device=str(device) if quantize is None else "cpu"
+            ) as f:
+                for key in f.keys():
+                    value = f.get_slice(key)
                     value = value.to(device if quantize is None else "cpu").to(dtype)
 
                     layer_name = ".".join(key.split(".")[:4])