formatting

vllm/engine/llm_engine.py

@@ -101,14 +101,14 @@ def __init__(
         self._init_tokenizer()
         self.seq_counter = Counter()
 
-
         self.model_executor = executor_class(model_config, cache_config,
                                              parallel_config, scheduler_config,
                                              device_config, lora_config)
-        
+
         #If usage stat is enabled, collect relevant info.
         if is_usage_stats_enabled():
-            usage_message.report_usage(self.model_executor.architecture, usage_context)
+            usage_message.report_usage(self.model_executor.architecture,
+                                       usage_context)
 
         # Create the scheduler.
         # NOTE: the cache_config here have been updated with the numbers of

vllm/model_executor/model_loader.py

@@ -20,7 +20,8 @@ def _set_default_torch_dtype(dtype: torch.dtype):
     torch.set_default_dtype(old_dtype)
 
 
-def _get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],str]:
+def _get_model_architecture(
+        model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:
     architectures = getattr(model_config.hf_config, "architectures", [])
     # Special handling for quantized Mixtral.
     # FIXME(woosuk): This is a temporary hack.
@@ -36,9 +37,10 @@ def _get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],
         f"Model architectures {architectures} are not supported for now. "
         f"Supported architectures: {ModelRegistry.get_supported_archs()}")
 
+
 def get_architecture(model_config: ModelConfig) -> str:
     return _get_model_architecture(model_config)[1]
-    
+
 
 def get_model(model_config: ModelConfig, device_config: DeviceConfig,
               **kwargs) -> nn.Module:

vllm/model_executor/neuron_model_loader.py

@@ -21,7 +21,8 @@
 }
 
 
-def _get_model_architecture(config: PretrainedConfig) -> Tuple[Type[nn.Module], str]:
+def _get_model_architecture(
+        config: PretrainedConfig) -> Tuple[Type[nn.Module], str]:
     architectures = getattr(config, "architectures", [])
     for arch in architectures:
         model_cls = ModelRegistry.load_model_cls(arch)
@@ -31,9 +32,11 @@ def _get_model_architecture(config: PretrainedConfig) -> Tuple[Type[nn.Module],
         f"Model architectures {architectures} are not supported for now. "
         f"Supported architectures: {ModelRegistry.get_supported_archs()}")
 
+
 def get_architecture(model_config: ModelConfig) -> str:
     return _get_model_architecture(model_config.hf_config)[1]
 
+
 def get_model(model_config: ModelConfig, device_config: DeviceConfig,
               **kwargs) -> nn.Module:
     from transformers_neuronx.config import (NeuronConfig,

vllm/model_executor/utils.py

@@ -51,9 +51,11 @@ def get_model(model_config: ModelConfig, device_config: DeviceConfig,
     get_model_fn = imported_model_loader.get_model
     return get_model_fn(model_config, device_config, **kwargs)
 
-def get_architecture(model_config: ModelConfig, device_config: DeviceConfig) -> str:
+
+def get_architecture(model_config: ModelConfig,
+                     device_config: DeviceConfig) -> str:
     model_loader_module = DEVICE_TO_MODEL_LOADER_MAP[device_config.device_type]
     imported_model_loader = importlib.import_module(
         f"vllm.model_executor.{model_loader_module}")
     get_architecture_fn = imported_model_loader.get_architecture
-    return get_architecture_fn(model_config)
+    return get_architecture_fn(model_config)

vllm/worker/model_runner.py

@@ -91,7 +91,8 @@ def load_model(self) -> None:
                                    lora_config=self.lora_config,
                                    parallel_config=self.parallel_config,
                                    scheduler_config=self.scheduler_config)
-        self.architecture = get_architecture(self.model_config, self.device_config)
+        self.architecture = get_architecture(self.model_config,
+                                             self.device_config)
         self.model_memory_usage = m.consumed_memory
         logger.info(f"Loading model weights took "
                     f"{self.model_memory_usage / float(2**30):.4f} GB")