Add eager flag and set compile to be always true (#734)

launcher/src/main.rs

@@ -342,9 +342,18 @@ struct Args {
 
     /// Whether you want to compile the model into a CUDA graph.
     /// This will speed up decoding but increase GPU memory usage.
+    /// Only use either `--compile` or `--eager`. Using both at the same time will
+    /// result in an error.
     #[clap(long, env, value_enum)]
     compile: bool,
 
+    /// Whether you want to run the model in eager mode, without
+    /// CUDA mode compilation, or run it with compilation.
+    /// Only use either `--compile` or `--eager`. Using both at the same time will
+    /// result in an error.
+    #[clap(long, env, value_enum)]
+    eager: bool,
+
     // The maximum batch size past which CUDA graphs are disabled.
     #[clap(default_value = "128", long, env)]
     compile_max_batch_size: usize,
@@ -656,6 +665,7 @@ fn shard_manager(
     adapter_source: String,
     quantize: Option<Quantization>,
     compile: bool,
+    eager: bool,
     compile_max_batch_size: usize,
     compile_max_rank: usize,
     compile_batch_size: usize,
@@ -738,10 +748,14 @@ fn shard_manager(
     }
 
     // CUDA graph compilation
-    if compile {
+    if !eager {
         shard_args.push("--compile".to_string());
     }
 
+    if compile && eager {
+        panic!("Cannot use both --compile and --eager at the same time.");
+    }
+
     // Speculative decoding
     if let Some(speculative_tokens) = speculative_tokens {
         shard_args.push("--speculative-tokens".to_string());
@@ -1310,6 +1324,7 @@ fn spawn_shards(
         let otlp_endpoint = args.otlp_endpoint.clone();
         let quantize = args.quantize;
         let compile = args.compile;
+        let eager = args.eager;
         let compile_max_batch_size = args.compile_max_batch_size;
         let compile_max_rank = args.compile_max_rank;
         let compile_batch_size = args.compile_batch_size;
@@ -1342,6 +1357,7 @@ fn spawn_shards(
                 adapter_source,
                 quantize,
                 compile,
+                eager,
                 compile_max_batch_size,
                 compile_max_rank,
                 compile_batch_size,

server/lorax_server/models/bloom.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -66,7 +67,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with Bloom")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():

server/lorax_server/models/causal_lm.py

@@ -2,6 +2,7 @@
 from typing import Dict, List, Optional, Tuple, Type
 
 import torch
+from loguru import logger
 from opentelemetry import trace
 from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase
 
@@ -499,7 +500,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with CausalLM")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         if torch.cuda.is_available():
             device = torch.device("cuda")

server/lorax_server/models/flash_causal_lm.py

@@ -1182,6 +1182,7 @@ def __init__(
             SLIDING_WINDOW_BLOCKS = math.ceil(sliding_window / BLOCK_SIZE)
 
         self.compile = compile
+
         self.model_graph_wrapper: GraphCache = None
         self.kv_cache = []
 

server/lorax_server/models/galactica.py

@@ -3,6 +3,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -161,7 +162,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with GalacticaSharded")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():

server/lorax_server/models/gpt_neox.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -29,7 +30,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with GPT-NeoX")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():

server/lorax_server/models/model.py

@@ -228,6 +228,10 @@ def check_initialized(self):
                 f"found uninitialized parameters in model {self.__class__.__name__}: {uninitialized_parameters}"
             )
 
+    @property
+    def supports_cuda_graph_compilation(self) -> bool:
+        return True
+
     @property
     def supports_adapter_loading(self) -> bool:
         return False

server/lorax_server/models/mpt.py

@@ -5,6 +5,7 @@
 import torch
 import torch.distributed
 from huggingface_hub import hf_hub_download
+from loguru import logger
 from opentelemetry import trace
 from transformers import AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase
 
@@ -58,7 +59,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with MPT")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():

server/lorax_server/models/opt.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -27,7 +28,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with OPT")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():

server/lorax_server/models/rw.py

@@ -1,6 +1,7 @@
 from typing import List, Optional, Tuple
 
 import torch
+from loguru import logger
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from lorax_server.models.causal_lm import CausalLM
@@ -17,7 +18,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with RW")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         if torch.cuda.is_available():
             device = torch.device("cuda")

server/lorax_server/models/santacoder.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from lorax_server.models.causal_lm import CausalLM
@@ -24,7 +25,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with SantaCoder")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         if torch.cuda.is_available():
             device = torch.device("cuda")

server/lorax_server/models/seq2seq_lm.py

@@ -2,6 +2,7 @@
 from typing import Dict, List, Optional, Tuple, Type
 
 import torch
+from loguru import logger
 from opentelemetry import trace
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, PreTrainedTokenizerBase
 
@@ -488,7 +489,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with Seq2SeqLM")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         if torch.cuda.is_available():
             device = torch.device("cuda")

server/lorax_server/models/t5.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.distributed
+from loguru import logger
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -29,7 +30,7 @@ def __init__(
         trust_remote_code: bool = False,
     ):
         if compile:
-            raise ValueError("`--compile` is not supported with T5")
+            logger.info(f"Model {model_id} does not support CUDA graph compilation. Skipping compilation.")
 
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():