Skip to content

Applying weight padding to deepseek #421

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
merged 2 commits into from
Feb 13, 2025
Merged

Applying weight padding to deepseek #421

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9c0ce9c
Applying weight padding to deepseek
gshtras Feb 12, 2025
c43f07b
Merge branch 'main' into deepseek_fp8_padding
gshtras Feb 13, 2025
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
24 changes: 14 additions & 10 deletions vllm/model_executor/layers/quantization/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,8 +188,6 @@ def create_weights(
weight_loader = extra_weight_attrs.get("weight_loader")

if self.block_quant:
assert not envs.VLLM_FP8_PADDING, (
"FP8 weight padding is not supported in block quantization.")
tp_size = get_tensor_model_parallel_world_size()
assert self.quant_config.weight_block_size is not None
block_n, block_k = (
Expand Down Expand Up @@ -273,6 +271,17 @@ def create_weights(
else:
layer.register_parameter("input_scale", None)

def add_padding_to_weight(self, weight: torch.Tensor) -> torch.Tensor:
# Pad the weight tensor. This is an optimization on ROCm platform, which
# can benefit from tensors located far enough from one another in memory
if (current_platform.is_rocm() and envs.VLLM_FP8_PADDING
and weight.stride(-1) == 1
and (weight.stride(-2) * weight.element_size()) % 512 == 0):
num_pad = 256 // weight.element_size()
weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
torch.cuda.empty_cache()
return weight

def process_weights_after_loading(self, layer: Module) -> None:
# TODO(rob): refactor block quant into separate class.
if self.block_quant:
Expand All @@ -286,6 +295,8 @@ def process_weights_after_loading(self, layer: Module) -> None:
weight = layer.weight.data
weight_scale_inv = layer.weight_scale_inv.data

weight = self.add_padding_to_weight(weight)

# Torch.compile cannot use Parameter subclasses.
layer.weight = Parameter(weight, requires_grad=False)
layer.weight_scale_inv = Parameter(weight_scale_inv,
Expand Down Expand Up @@ -353,14 +364,7 @@ def process_weights_after_loading(self, layer: Module) -> None:
logical_widths=layer.logical_widths,
)

# Pad the weight
if envs.VLLM_FP8_PADDING and weight.stride(-1) == 1 \
and (weight.stride(-2) * weight.element_size()) % 512 == 0:
num_pad = 256 // weight.element_size()
weight = F.pad(weight, (0, num_pad), "constant",
0)[..., :-num_pad]
torch.cuda.empty_cache()

weight = self.add_padding_to_weight(weight)
# Update layer with new values.
layer.weight = Parameter(weight.t(), requires_grad=False)
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
Expand Down
2 changes: 1 addition & 1 deletion vllm/model_executor/layers/quantization/utils/fp8_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,7 +478,7 @@ def w8a8_block_fp8_matmul(
assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
M = A.numel() // A.shape[-1]

assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert B.ndim == 2 and Bs.ndim == 2
N, K = B.shape
assert triton.cdiv(N, block_n) == Bs.shape[0]
assert triton.cdiv(K, block_k) == Bs.shape[1]
Expand Down
Loading