mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-08 09:03:20 +08:00
Chenjie Luo
e20d218b43
## What does this PR do? **Type of change:** ? New model support **Overview:** ? ## Usage Please see examples/deepseek/README.md <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * New Features * Support for DeepSeek V3.2 quantization and automatic detection of available DeepSeek versions. * Triton-backed weight dequantization utility and MoE-aware calibration mode to improve calibration fidelity. * Documentation * DeepSeek examples README expanded with setup, conversion, calibration, and FP8→FP4 quantization workflows for R1, V3, and V3.2. * Bug Fixes * More robust, failure-tolerant copying of auxiliary files/assets during quantization. * Chores * Updated changelog and lint/ignore rules for example artifacts. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Signed-off-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com>
111 lines
4.7 KiB
Python
111 lines
4.7 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# MIT License
|
|
|
|
# Copyright (c) 2023 DeepSeek
|
|
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
"""Reference: https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py"""
|
|
|
|
|
|
@triton.jit
|
|
def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
|
|
"""
|
|
Dequantizes weights using the provided scaling factors and stores the result.
|
|
|
|
Args:
|
|
x_ptr (tl.pointer): Pointer to the quantized weights.
|
|
s_ptr (tl.pointer): Pointer to the scaling factors.
|
|
y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights.
|
|
M (int): Number of rows in the weight matrix.
|
|
N (int): Number of columns in the weight matrix.
|
|
BLOCK_SIZE (tl.constexpr): Size of the block for tiling.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
pid_m = tl.program_id(axis=0)
|
|
pid_n = tl.program_id(axis=1)
|
|
n = tl.cdiv(N, BLOCK_SIZE)
|
|
offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
|
offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
|
offs = offs_m[:, None] * N + offs_n[None, :]
|
|
mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
|
|
x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
|
|
s = tl.load(s_ptr + pid_m * n + pid_n)
|
|
y = x * s
|
|
tl.store(y_ptr + offs, y, mask=mask)
|
|
|
|
|
|
def weight_dequant(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor:
|
|
"""
|
|
Dequantizes the given weight tensor using the provided scale tensor.
|
|
|
|
Args:
|
|
x (torch.Tensor): The quantized weight tensor of shape (M, N).
|
|
s (torch.Tensor): The scale tensor of shape (M//block_size, N//block_size).
|
|
block_size (int, optional): The block size to use for dequantization. Defaults to 128.
|
|
|
|
Returns:
|
|
torch.Tensor: The dequantized weight tensor of the same shape as `x`.
|
|
|
|
Raises:
|
|
AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2.
|
|
"""
|
|
assert x.is_contiguous() and s.is_contiguous(), "Input tensors must be contiguous"
|
|
assert x.dim() == 2 and s.dim() == 2, "Input tensors must have 2 dimensions"
|
|
M, N = x.size()
|
|
y = torch.empty_like(x, dtype=torch.get_default_dtype())
|
|
grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]))
|
|
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size)
|
|
return y
|