run linting

examples/save_state_dict.py

@@ -45,7 +45,9 @@ def main(args):
     Path(args.output).mkdir(exist_ok=True)
     # Dump worker states to output directory
     model_executor = llm.llm_engine.model_executor
-    model_executor._run_workers("save_model", path=args.output, max_size=5 * 1024 ** 3)
+    model_executor._run_workers("save_model",
+                                path=args.output,
+                                max_size=5 * 1024**3)
     # Copy metadata files to output directory
     for file in os.listdir(model_path):
         if not any(

vllm/model_executor/model_loader/loader.py

@@ -362,6 +362,7 @@ def load_model(self, *, model_config: ModelConfig,
                    parallel_config: ParallelConfig,
                    scheduler_config: SchedulerConfig) -> nn.Module:
         from safetensors.torch import load_file
+
         from vllm.distributed import get_tensor_model_parallel_rank
         with set_default_torch_dtype(model_config.dtype):
             with torch.device(device_config.device):

vllm/worker/model_runner.py

@@ -212,13 +212,14 @@ def load_model(self) -> None:
                            "but the KV cache data type is not FP8. "
                            "KV cache scaling factors will not be used.")
 
-    def save_model(self, path: str, max_size: int=None) -> None:
+    def save_model(self, path: str, max_size: int) -> None:
         from safetensors.torch import save_file
+
         from vllm.distributed import get_tensor_model_parallel_rank
         rank = get_tensor_model_parallel_rank()
         idx = 0
         size = 0
-        params = {}
+        params: Dict[str, torch.Tensor] = {}
         for name, param in self.model.named_parameters():
             param_size = param.nelement() * param.element_size()
             if max_size and size + param_size > max_size:

vllm/worker/worker.py

@@ -117,7 +117,7 @@ def init_device(self) -> None:
     def load_model(self):
         self.model_runner.load_model()
 
-    def save_model(self, path: str, max_size: int=None) -> None:
+    def save_model(self, path: str, max_size: int) -> None:
         self.model_runner.save_model(path, max_size=max_size)
 
     @torch.inference_mode()