add model id

server/text_generation_server/models/bloom.py

@@ -46,6 +46,8 @@ class BLOOMSharded(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/causal_lm.py

@@ -486,6 +486,7 @@ class CausalLM(Model):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
         if use_medusa:
             raise RuntimeError("Medusa decoding is not enabled for AutoModel")
 

server/text_generation_server/models/flash_causal_lm.py

@@ -834,15 +834,21 @@ class FlashCausalLM(Model):
             if os.environ.get("PYTORCH_TUNABLEOP_TUNING", "1"):
                 torch.cuda.tunable.tuning_enable(True)
 
-            tuning_sequences = range(1, 8)
-            tunableop_filename = f"tunableop_tp{self.world_size}_rank{self.rank}.csv"
+            tuning_sequences = list(range(1, 3))
+            tunableop_filepath = os.path.join("/data", f"tunableop_{self.model_id.replace('/', '-')}_tp{self.world_size}_rank{self.rank}.csv")
 
-            logger.info(f"PyTorch TunableOp (https://github.com/fxmarty/pytorch/tree/2.3-patched/aten/src/ATen/cuda/tunable) is enabled. The warmup may take several minutes, picking the ROCm optimal matrix multiplication kernel for the target lengths {', '.join(tuning_sequences)}.")
-            torch.cuda.tunable.read_file(tunableop_filename)
+            logger.info(f"PyTorch TunableOp (https://github.com/fxmarty/pytorch/tree/2.3-patched/aten/src/ATen/cuda/tunable) is enabled. The warmup may take several minutes, picking the ROCm optimal matrix multiplication kernel for the target lengths {', '.join([str(seqlen) for seqlen in tuning_sequences])} (typical decoding lengths). The picked GEMMs are saved in the file {tunableop_filepath}.")
 
-            for seqlen in range(1, 8):
+            if os.path.isfile(tunableop_filepath):
+                logger.info(f"The file {tunableop_filepath} already exists and will be reused.")
+                torch.cuda.tunable.read_file(tunableop_filepath)
+            
+            os.makedirs("/data", exist_ok=True)
+            
+            for seqlen in tuning_sequences:
+                logger.info(f"Warming up TunableOp for seqlen={seqlen}")
                 self.tunableop_warmup(seqlen)
-                torch.cuda.tunable.write_file(tunableop_filename)
+                torch.cuda.tunable.write_file(tunableop_filepath)
             torch.cuda.tunable.tuning_enable(False)
         
         if CUDA_GRAPHS:

server/text_generation_server/models/flash_cohere.py

@@ -28,6 +28,8 @@ class FlashCohere(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_dbrx.py

@@ -30,6 +30,8 @@ class FlashDbrx(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_gemma.py

@@ -29,6 +29,8 @@ class FlashGemma(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_llama.py

@@ -31,6 +31,8 @@ class FlashLlama(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_mistral.py

@@ -318,6 +318,8 @@ class BaseFlashMistral(FlashCausalLM):
         trust_remote_code: bool = False,
         tokenizer_class=AutoTokenizer,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_neox.py

@@ -29,6 +29,8 @@ class FlashNeoXSharded(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_phi.py

@@ -29,6 +29,8 @@ class FlashPhi(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_qwen2.py

@@ -34,6 +34,8 @@ class FlashQwen2(BaseFlashMistral):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_rw.py

@@ -30,6 +30,8 @@ class FlashRWSharded(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_santacoder.py

@@ -33,6 +33,8 @@ class FlashSantacoderSharded(FlashCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/flash_starcoder2.py

@@ -33,6 +33,8 @@ class FlashStarcoder2(BaseFlashMistral):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/galactica.py

@@ -171,6 +171,8 @@ class GalacticaSharded(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/gpt_neox.py

@@ -28,6 +28,8 @@ class GPTNeoxSharded(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/idefics.py

@@ -35,6 +35,8 @@ class IDEFICSSharded(IdeficsCausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/mamba.py

@@ -412,6 +412,8 @@ class Mamba(Model):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, _rank, world_size = initialize_torch_distributed()
         if world_size > 1:
             raise RuntimeError("Mamba does not support Tensor Parallelism (TP)")

server/text_generation_server/models/mpt.py

@@ -47,6 +47,8 @@ class MPTSharded(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/opt.py

@@ -26,6 +26,8 @@ class OPTSharded(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/models/phi.py

@@ -26,6 +26,8 @@ class Phi(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, _rank, _world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device("cuda")

server/text_generation_server/models/santacoder.py

@@ -23,6 +23,8 @@ class SantaCoder(CausalLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         if torch.cuda.is_available():
             device = torch.device("cuda")
             dtype = torch.float16 if dtype is None else dtype

server/text_generation_server/models/seq2seq_lm.py

@@ -536,6 +536,8 @@ class Seq2SeqLM(Model):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         if use_medusa:
             raise RuntimeError("Medusa decoding is not enabled for AutoModel")
 

server/text_generation_server/models/t5.py

@@ -29,6 +29,8 @@ class T5Sharded(Seq2SeqLM):
         dtype: Optional[torch.dtype] = None,
         trust_remote_code: bool = False,
     ):
+        self.model_id = model_id
+        
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")

server/text_generation_server/utils/layers.py

@@ -1018,7 +1018,7 @@ try:
         from flash_attn.layers.rotary import RotaryEmbedding
         import rotary_emb
     elif IS_ROCM_SYSTEM:
-        from vllm import pos_encoding_ops
+        from vllm._C import ops
 
     def _create_inv_freq(dim, base, device):
         inv_freq = 1.0 / (
@@ -1339,6 +1339,5 @@ try:
                 freqs = torch.outer(t, self.inv_freq.to(device=t.device))
                 self._cos_cached = (torch.cos(freqs) * self.mscale).to(dtype)
                 self._sin_cached = (torch.sin(freqs) * self.mscale).to(dtype)
-
-except ImportError:
-    pass
+except ImportError as e:
+    logger.warning(f"ImportError in layers.py, beware that this may cause issues later on: {e}")