huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 11:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
OlivierDehaene 2023-06-01 10:55:26 +02:00
parent db2ebe3947
commit f6438ac352
2 changed files with 76 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
launcher/src/main.rs
View File

@ -546,11 +546,7 @@ enum LauncherError {
WebserverCannotStart, WebserverCannotStart,
} }
fn download_convert_model( fn download_convert_model(args: &Args, running: Arc<AtomicBool>) -> Result<(), LauncherError> {
args: &Args,
auto_convert: bool,
running: Arc<AtomicBool>,
) -> Result<(), LauncherError> {
let mut download_argv = vec![ let mut download_argv = vec![
"text-generation-server".to_string(), "text-generation-server".to_string(),
"download-weights".to_string(), "download-weights".to_string(),
@ -562,11 +558,6 @@ fn download_convert_model(
"--json-output".to_string(), "--json-output".to_string(),
]; ];
// Auto convert weights to safetensors
if auto_convert {
download_argv.push("--auto-convert".to_string());
}
// Model optional revision // Model optional revision
if let Some(revision) = &args.revision { if let Some(revision) = &args.revision {
download_argv.push("--revision".to_string()); download_argv.push("--revision".to_string());
@ -932,11 +923,8 @@ fn main() -> Result<(), LauncherError> {
}) })
.expect("Error setting Ctrl-C handler"); .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 and convert model weights
download_convert_model(&args, auto_convert, running.clone())?; download_convert_model(&args, running.clone())?;
// Shared shutdown bool // Shared shutdown bool
let shutdown = Arc::new(Mutex::new(false)); let shutdown = Arc::new(Mutex::new(false));

14
server/text_generation_server/models/flash_santacoder.py
View File

@ -54,12 +54,7 @@ class FlashSantacoder(FlashCausalLM):
) )
# We do not use from_pretrained as we modified the model internal module layout # We do not use from_pretrained as we modified the model internal module layout
try: filenames = weight_files(model_id, revision, ".safetensors")
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)
with init_empty_weights(): with init_empty_weights():
model = FlashSantacoderForCausalLM(config) model = FlashSantacoderForCausalLM(config)
@ -91,8 +86,11 @@ class FlashSantacoder(FlashCausalLM):
transpose: bool, transpose: bool,
): ):
for filename in filenames: for filename in filenames:
state_dict = torch.load(filename, map_location="cpu") with safe_open(
for key, value in state_dict.items(): 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) value = value.to(device if quantize is None else "cpu").to(dtype)
layer_name = ".".join(key.split(".")[:4]) layer_name = ".".join(key.split(".")[:4])