242 lines
8.1 KiB
Python
242 lines
8.1 KiB
Python
import torch
|
|
import fastvideo.envs as envs
|
|
from fastvideo.logger import init_logger
|
|
import argparse
|
|
import sys
|
|
from typing import List, Dict, Union
|
|
import yaml
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
def find_nccl_library() -> str:
|
|
"""
|
|
We either use the library file specified by the `VLLM_NCCL_SO_PATH`
|
|
environment variable, or we find the library file brought by PyTorch.
|
|
After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
|
|
found by `ctypes` automatically.
|
|
"""
|
|
so_file = envs.FASTVIDEO_NCCL_SO_PATH
|
|
|
|
# manually load the nccl library
|
|
if so_file:
|
|
logger.info(
|
|
"Found nccl from environment variable FASTVIDEO_NCCL_SO_PATH=%s",
|
|
so_file)
|
|
else:
|
|
if torch.version.cuda is not None:
|
|
so_file = "libnccl.so.2"
|
|
elif torch.version.hip is not None:
|
|
so_file = "librccl.so.1"
|
|
else:
|
|
raise ValueError("NCCL only supports CUDA and ROCm backends.")
|
|
logger.info("Found nccl from library %s", so_file)
|
|
return so_file
|
|
|
|
prev_set_stream = torch.cuda.set_stream
|
|
|
|
_current_stream = None
|
|
|
|
|
|
def _patched_set_stream(stream: torch.cuda.Stream) -> None:
|
|
global _current_stream
|
|
_current_stream = stream
|
|
prev_set_stream(stream)
|
|
|
|
|
|
torch.cuda.set_stream = _patched_set_stream
|
|
|
|
|
|
def current_stream() -> torch.cuda.Stream:
|
|
"""
|
|
replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
|
|
it turns out that `torch.cuda.current_stream()` is quite expensive,
|
|
as it will construct a new stream object at each call.
|
|
here we patch `torch.cuda.set_stream` to keep track of the current stream
|
|
directly, so that we can avoid calling `torch.cuda.current_stream()`.
|
|
|
|
the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
|
|
from C/C++ code.
|
|
"""
|
|
from fastvideo.platforms import current_platform
|
|
global _current_stream
|
|
if _current_stream is None:
|
|
# when this function is called before any stream is set,
|
|
# we return the default stream.
|
|
# On ROCm using the default 0 stream in combination with RCCL
|
|
# is hurting performance. Therefore creating a dedicated stream
|
|
# per process
|
|
_current_stream = torch.cuda.Stream() if current_platform.is_rocm(
|
|
) else torch.cuda.current_stream()
|
|
return _current_stream
|
|
|
|
|
|
class StoreBoolean(argparse.Action):
|
|
|
|
def __call__(self, parser, namespace, values, option_string=None):
|
|
if values.lower() == "true":
|
|
setattr(namespace, self.dest, True)
|
|
elif values.lower() == "false":
|
|
setattr(namespace, self.dest, False)
|
|
else:
|
|
raise ValueError(f"Invalid boolean value: {values}. "
|
|
"Expected 'true' or 'false'.")
|
|
|
|
|
|
class SortedHelpFormatter(argparse.HelpFormatter):
|
|
"""SortedHelpFormatter that sorts arguments by their option strings."""
|
|
|
|
def add_arguments(self, actions):
|
|
actions = sorted(actions, key=lambda x: x.option_strings)
|
|
super().add_arguments(actions)
|
|
|
|
|
|
class FlexibleArgumentParser(argparse.ArgumentParser):
|
|
"""ArgumentParser that allows both underscore and dash in names."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
# Set the default 'formatter_class' to SortedHelpFormatter
|
|
if 'formatter_class' not in kwargs:
|
|
kwargs['formatter_class'] = SortedHelpFormatter
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def parse_args(self, args=None, namespace=None):
|
|
if args is None:
|
|
args = sys.argv[1:]
|
|
|
|
if '--config' in args:
|
|
args = self._pull_args_from_config(args)
|
|
|
|
# Convert underscores to dashes and vice versa in argument names
|
|
processed_args = []
|
|
for arg in args:
|
|
if arg.startswith('--'):
|
|
if '=' in arg:
|
|
key, value = arg.split('=', 1)
|
|
key = '--' + key[len('--'):].replace('_', '-')
|
|
processed_args.append(f'{key}={value}')
|
|
else:
|
|
processed_args.append('--' +
|
|
arg[len('--'):].replace('_', '-'))
|
|
elif arg.startswith('-O') and arg != '-O' and len(arg) == 2:
|
|
# allow -O flag to be used without space, e.g. -O3
|
|
processed_args.append('-O')
|
|
processed_args.append(arg[2:])
|
|
else:
|
|
processed_args.append(arg)
|
|
|
|
return super().parse_args(processed_args, namespace)
|
|
|
|
def _pull_args_from_config(self, args: List[str]) -> List[str]:
|
|
"""Method to pull arguments specified in the config file
|
|
into the command-line args variable.
|
|
|
|
The arguments in config file will be inserted between
|
|
the argument list.
|
|
|
|
example:
|
|
```yaml
|
|
port: 12323
|
|
tensor-parallel-size: 4
|
|
```
|
|
```python
|
|
$: vllm {serve,chat,complete} "facebook/opt-12B" \
|
|
--config config.yaml -tp 2
|
|
$: args = [
|
|
"serve,chat,complete",
|
|
"facebook/opt-12B",
|
|
'--config', 'config.yaml',
|
|
'-tp', '2'
|
|
]
|
|
$: args = [
|
|
"serve,chat,complete",
|
|
"facebook/opt-12B",
|
|
'--port', '12323',
|
|
'--tensor-parallel-size', '4',
|
|
'-tp', '2'
|
|
]
|
|
```
|
|
|
|
Please note how the config args are inserted after the sub command.
|
|
this way the order of priorities is maintained when these are args
|
|
parsed by super().
|
|
"""
|
|
assert args.count(
|
|
'--config') <= 1, "More than one config file specified!"
|
|
|
|
index = args.index('--config')
|
|
if index == len(args) - 1:
|
|
raise ValueError("No config file specified! \
|
|
Please check your command-line arguments.")
|
|
|
|
file_path = args[index + 1]
|
|
|
|
config_args = self._load_config_file(file_path)
|
|
|
|
# 0th index is for {serve,chat,complete}
|
|
# followed by model_tag (only for serve)
|
|
# followed by config args
|
|
# followed by rest of cli args.
|
|
# maintaining this order will enforce the precedence
|
|
# of cli > config > defaults
|
|
if args[0] == "serve":
|
|
if index == 1:
|
|
raise ValueError(
|
|
"No model_tag specified! Please check your command-line"
|
|
" arguments.")
|
|
args = [args[0]] + [
|
|
args[1]
|
|
] + config_args + args[2:index] + args[index + 2:]
|
|
else:
|
|
args = [args[0]] + config_args + args[1:index] + args[index + 2:]
|
|
|
|
return args
|
|
|
|
def _load_config_file(self, file_path: str) -> List[str]:
|
|
"""Loads a yaml file and returns the key value pairs as a
|
|
flattened list with argparse like pattern
|
|
```yaml
|
|
port: 12323
|
|
tensor-parallel-size: 4
|
|
```
|
|
returns:
|
|
processed_args: list[str] = [
|
|
'--port': '12323',
|
|
'--tensor-parallel-size': '4'
|
|
]
|
|
|
|
"""
|
|
|
|
extension: str = file_path.split('.')[-1]
|
|
if extension not in ('yaml', 'yml'):
|
|
raise ValueError(
|
|
"Config file must be of a yaml/yml type.\
|
|
%s supplied", extension)
|
|
|
|
# only expecting a flat dictionary of atomic types
|
|
processed_args: List[str] = []
|
|
|
|
config: Dict[str, Union[int, str]] = {}
|
|
try:
|
|
with open(file_path) as config_file:
|
|
config = yaml.safe_load(config_file)
|
|
except Exception as ex:
|
|
logger.error(
|
|
"Unable to read the config file at %s. \
|
|
Make sure path is correct", file_path)
|
|
raise ex
|
|
|
|
store_boolean_arguments = [
|
|
action.dest for action in self._actions
|
|
if isinstance(action, StoreBoolean)
|
|
]
|
|
|
|
for key, value in config.items():
|
|
if isinstance(value, bool) and key not in store_boolean_arguments:
|
|
if value:
|
|
processed_args.append('--' + key)
|
|
else:
|
|
processed_args.append('--' + key)
|
|
processed_args.append(str(value))
|
|
|
|
return processed_args
|