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

Merge remote-tracking branch 'origin/main' into feat/rw

Browse Source
This commit is contained in:
OlivierDehaene 2023-05-30 15:24:44 +02:00
commit 3e517bfc9d
5 changed files with 33 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
launcher/src/main.rs
View File

@ -455,11 +455,12 @@ fn shutdown_shards(shutdown: Arc<Mutex<bool>>, shutdown_receiver: &mpsc::Receive
}
fn num_cuda_devices() -> Option<usize> {
if let Ok(cuda_visible_devices) = env::var("CUDA_VISIBLE_DEVICES") {
let n_devices = cuda_visible_devices.split(',').count();
return Some(n_devices);
}
None
let devices = match env::var("CUDA_VISIBLE_DEVICES") {
Ok(devices) => devices,
Err(_) => env::var("NVIDIA_VISIBLE_DEVICES").ok()?,
};
let n_devices = devices.split(',').count();
Some(n_devices)
}
#[derive(Deserialize)]
@ -509,9 +510,9 @@ fn find_num_shards(sharded: Option<bool>, num_shard: Option<usize>) -> usize {
let num_shard = match (sharded, num_shard) {
(Some(true), None) => {
// try to default to the number of available GPUs
tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES");
let n_devices =
num_cuda_devices().expect("--num-shard and CUDA_VISIBLE_DEVICES are not set");
tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES");
let n_devices = num_cuda_devices()
.expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES are not set");
if n_devices <= 1 {
panic!("`sharded` is true but only found {n_devices} CUDA devices");
}

4
server/text_generation_server/models/bloom.py
View File

@ -246,9 +246,7 @@ class BLOOMSharded(BLOOM):
module.linear = replace_linear(state)
elif quantize == "gptq":
raise NotImplementedError(
"`gptq` is not implemented for now"
)
raise NotImplementedError("`gptq` is not implemented for now")
elif quantize is None:
tensor = tensor.to(device)
else:

4
server/text_generation_server/models/galactica.py
View File

@ -365,9 +365,7 @@ class GalacticaSharded(Galactica):
module.linear = replace_linear(state)
elif quantize == "gptq":
raise NotImplementedError(
"`gptq` is not implemented for now"
)
raise NotImplementedError("`gptq` is not implemented for now")
elif quantize is None:
tensor = tensor.to(device)
else:

4
server/text_generation_server/models/gpt_neox.py
View File

@ -211,9 +211,7 @@ class GPTNeoxSharded(CausalLM):
module.linear = replace_linear(state)
elif quantize == "gptq":
raise NotImplementedError(
"`gptq` is not implemented for now"
)
raise NotImplementedError("`gptq` is not implemented for now")
elif quantize is None:
tensor = tensor.to(device)
else:

6
server/text_generation_server/models/t5.py
View File

@ -224,10 +224,8 @@ class T5Sharded(Seq2SeqLM):
module.linear = replace_linear(state)
elif quantize == "gptq" and not module_name.endswith("wo"):
raise NotImplementedError(
"`gptq` is not implemented for now"
)
elif quantize is None:
raise NotImplementedError("`gptq` is not implemented for now")
elif quantize is None or module_name.endswith("wo"):
tensor = tensor.to(device)
else:
raise ValueError(f"Unexpected quantize `{quantize}`")