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

feat(launcher): add disable_custom_kernels arg (#67)

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-15 16:23:45 +01:00 committed by GitHub
parent c5a4a1faf3
commit 68455353f5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
launcher/src/main.rs
View File

@ -47,6 +47,8 @@ struct Args {
#[clap(long, env)] #[clap(long, env)]
weights_cache_override: Option<String>, weights_cache_override: Option<String>,
#[clap(long, env)] #[clap(long, env)]
disable_custom_kernels: bool,
#[clap(long, env)]
json_output: bool, json_output: bool,
#[clap(long, env)] #[clap(long, env)]
otlp_endpoint: Option<String>, otlp_endpoint: Option<String>,
@ -79,6 +81,7 @@ fn main() -> ExitCode {
master_port, master_port,
huggingface_hub_cache, huggingface_hub_cache,
weights_cache_override, weights_cache_override,
disable_custom_kernels,
json_output, json_output,
otlp_endpoint, otlp_endpoint,
} = args; } = args;
@ -242,6 +245,7 @@ fn main() -> ExitCode {
master_port, master_port,
huggingface_hub_cache, huggingface_hub_cache,
weights_cache_override, weights_cache_override,
disable_custom_kernels,
otlp_endpoint, otlp_endpoint,
status_sender, status_sender,
shutdown, shutdown,
@ -406,6 +410,7 @@ fn shard_manager(
master_port: usize, master_port: usize,
huggingface_hub_cache: Option<String>, huggingface_hub_cache: Option<String>,
weights_cache_override: Option<String>, weights_cache_override: Option<String>,
disable_custom_kernels: bool,
otlp_endpoint: Option<String>, otlp_endpoint: Option<String>,
status_sender: mpsc::Sender<ShardStatus>, status_sender: mpsc::Sender<ShardStatus>,
shutdown: Arc<Mutex<bool>>, shutdown: Arc<Mutex<bool>>,
@ -474,6 +479,11 @@ fn shard_manager(
)); ));
}; };
// If disable_custom_kernels is true, pass it to the shard as an env var
if disable_custom_kernels {
env.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
}
// If the NCCL_SHM_DISABLE env var is set, pass it to the shard // If the NCCL_SHM_DISABLE env var is set, pass it to the shard
// needed when running NCCL inside a docker container and when you can't increase shm size // needed when running NCCL inside a docker container and when you can't increase shm size
if let Ok(nccl_shm_disalbe) = env::var("NCCL_SHM_DISABLE") { if let Ok(nccl_shm_disalbe) = env::var("NCCL_SHM_DISABLE") {