Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

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)]
weights_cache_override: Option<String>,
#[clap(long, env)]
disable_custom_kernels: bool,
#[clap(long, env)]
json_output: bool,
#[clap(long, env)]
otlp_endpoint: Option<String>,
@ -79,6 +81,7 @@ fn main() -> ExitCode {
master_port,
huggingface_hub_cache,
weights_cache_override,
disable_custom_kernels,
json_output,
otlp_endpoint,
} = args;
@ -242,6 +245,7 @@ fn main() -> ExitCode {
master_port,
huggingface_hub_cache,
weights_cache_override,
disable_custom_kernels,
otlp_endpoint,
status_sender,
shutdown,
@ -406,6 +410,7 @@ fn shard_manager(
master_port: usize,
huggingface_hub_cache: Option<String>,
weights_cache_override: Option<String>,
disable_custom_kernels: bool,
otlp_endpoint: Option<String>,
status_sender: mpsc::Sender<ShardStatus>,
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
// 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") {