huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-19 22:02:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Fix #366
This commit is contained in:
OlivierDehaene 2023-06-01 12:10:35 +02:00 committed by GitHub
parent c0928e6f26
commit 95d3546976
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 91 additions and 91 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,
}
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));

28
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
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)