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

Fix #366

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_tensor(key)
                     value = value.to(device if quantize is None else "cpu").to(dtype)
 
                     layer_name = ".".join(key.split(".")[:4])
@@ -167,9 +165,21 @@ class FlashSantacoder(FlashCausalLM):
 
                     del value
 
+        if model.lm_head.weight.device == torch.device("meta"):
+            model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
+
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)
 
+        uninitialized_parameters = []
+        for n, p in model.named_parameters():
+            if p.data.device == torch.device("meta"):
+                uninitialized_parameters.append(n)
+        if uninitialized_parameters:
+            raise RuntimeError(
+                f"found uninitialized parameters in model : {uninitialized_parameters}"
+            )
+
     def decode(self, generated_ids: List[int]) -> str:
         # Do not skip special tokens as they are used for custom parsing rules of the generated text
         return self.tokenizer.decode(
@@ -389,6 +399,8 @@ class FlashSantacoderSharded(FlashSantacoder):
                     else:
                         module._buffers[param_name] = tensor
 
+        if model.lm_head.weight.device == torch.device("meta"):
             model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight)
+
         torch.cuda.empty_cache()
         model.post_load_weights(quantize)