Skip to content

[Core] Support tensor parallelism for GGUF quantization #7520

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mgoin merged 9 commits into from
Aug 19, 2024
Merged

[Core] Support tensor parallelism for GGUF quantization #7520

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
528b358
remove gguf tp raise
Isotr0py Aug 9, 2024
462783e
save draft
Isotr0py Aug 14, 2024
c8ffa77
fix tp=2 inference
Isotr0py Aug 14, 2024
ebf5378
code format
Isotr0py Aug 15, 2024
ac9dbe3
add gguf tp test
Isotr0py Aug 16, 2024
c2f8101
Merge branch 'vllm-project:main' into gguf-tp
Isotr0py Aug 16, 2024
3e3dc91
fix gguf test prompt format
Isotr0py Aug 18, 2024
bc61e56
code format
Isotr0py Aug 18, 2024
7292ba1
Merge branch 'vllm-project:main' into gguf-tp
Isotr0py Aug 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions vllm/model_executor/layers/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,11 +503,16 @@ def weight_loader(self,
loaded_shard_id

if is_gguf_weight:
shard_size = loaded_weight.shape[output_dim]
shard_offset = loaded_weight.shape[output_dim] * \
loaded_shard_id
tp_size = get_tensor_model_parallel_world_size()
output_dim = getattr(param, "output_dim", None)
shard_shape = list(loaded_weight.shape)
shard_shape[output_dim] = shard_shape[output_dim] // tp_size
param.shard_id.append(loaded_shard_id)
param.shard_size[loaded_shard_id] = loaded_weight.shape
param.shard_size[loaded_shard_id] = shard_shape

input_dim = getattr(param, "input_dim", None)
input_size = loaded_weight.shape[input_dim]
param_data = param_data.narrow(input_dim, 0, input_size)

param_data = param_data.narrow(output_dim, shard_offset,
shard_size)
Expand Down Expand Up @@ -855,8 +860,13 @@ def weight_loader(self,
param, orig_qkv_offsets, loaded_shard_id)

if is_gguf_weight:
tp_size = get_tensor_model_parallel_world_size()
output_dim = getattr(param, "output_dim", None)
shard_shape = list(loaded_weight.shape)
shard_shape[output_dim] = shard_shape[output_dim] // tp_size
param.shard_id.append(loaded_shard_id)
param.shard_size[loaded_shard_id] = loaded_weight.shape
param.shard_size[loaded_shard_id] = shard_shape

input_dim = getattr(param, "input_dim", None)
input_size = loaded_weight.shape[input_dim]
param_data = param_data.narrow(input_dim, 0, input_size)
Expand Down Expand Up @@ -968,6 +978,7 @@ def __init__(self,

def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
input_dim = getattr(param, "input_dim", None)

# Special case for GGUF
Expand All @@ -978,7 +989,10 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):

# Materialize GGUF UninitializedParameter
if is_gguf_weight and isinstance(param, UninitializedParameter):
param.materialize(loaded_weight.shape, dtype=loaded_weight.dtype)
weight_shape = list(loaded_weight.shape)
if input_dim:
weight_shape[input_dim] = weight_shape[input_dim] // tp_size
param.materialize(tuple(weight_shape), dtype=loaded_weight.dtype)

param_data = param.data
if input_dim is not None:
Expand Down
4 changes: 0 additions & 4 deletions vllm/model_executor/layers/quantization/gguf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from torch.nn.parameter import Parameter, UninitializedParameter

from vllm import _custom_ops as ops
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig, QuantizeMethodBase)
Expand Down Expand Up @@ -39,9 +38,6 @@ def get_config_filenames(cls) -> List[str]:

@classmethod
def from_config(cls, config: Dict[str, Any]) -> "GGUFConfig":
if get_tensor_model_parallel_world_size() > 1:
raise ValueError(
"GGUF quantization hasn't supported tensor parallelism yet.")
return cls()

def get_quant_method(self, layer: torch.nn.Module,
Expand Down