Skip to content

[Perf] Enable SonicMoE FP8 storage release #2

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

Open
A-nnonymous wants to merge 1 commit into
base: sonicmoe_fp8_dispatch
Choose a base branch
from
Open

[Perf] Enable SonicMoE FP8 storage release #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
22 changes: 18 additions & 4 deletions packages/paddlefleet_ops/src/paddlefleet_ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,13 @@ def _sonic_moe_requirement(
_FLASH_MASK_AVAILABLE = False
_CUDNN_FRONTEND_AVAILABLE = False


def _ecosystem_module_exists(lib_name: str) -> bool:
return (ops_dir / lib_name).exists() or (
ops_dir / f"{lib_name}.py"
).exists()


if paddle.is_compiled_with_cuda():
if paddle.cuda.get_device_capability()[0] >= 9:
_DEEP_GEMM_AVAILABLE = True
Expand All @@ -185,7 +192,7 @@ def _sonic_moe_requirement(
and _cuda_version >= (12, 9)
):
_SONIC_MOE_AVAILABLE = True
if sys.version_info >= (3, 12):
if sys.version_info >= (3, 12) and _ecosystem_module_exists("cudnn"):
_CUDNN_FRONTEND_AVAILABLE = True

if paddle.is_compiled_with_xpu():
Expand Down Expand Up @@ -347,9 +354,16 @@ def _safe_load_ecosystem_lib(
paddle.enable_compat(scope={"cudnn"}, silent=True)
_safe_load_ecosystem_lib("cudnn", ops_dir, globals())
else:
warning, error = _cutedsl_requirement(
"paddlefleet_ops.cudnn", hint=CUDNN_FRONTEND_HINT
)
if sys.version_info < (3, 12):
warning, error = _cutedsl_requirement(
"paddlefleet_ops.cudnn", hint=CUDNN_FRONTEND_HINT
)
else:
warning, error = _build_notice(
"paddlefleet_ops.cudnn",
"cudnn frontend Python module is not installed.",
hint_for_error=CUDNN_FRONTEND_HINT,
)
logger.warning(warning)
blocked_import_messages["paddlefleet_ops.cudnn"] = error

Expand Down
2 changes: 1 addition & 1 deletion packages/paddlefleet_ops/third_party/sonic-moe
Open in desktop
Submodule sonic-moe updated 31 files
+4 −4 sonicmoe/__init__.py
+29 −10 sonicmoe/config.py
+1 −2 sonicmoe/enums.py
+165 −8 sonicmoe/ernie_compat/deepep_metadata.py
+2 −2 sonicmoe/ernie_compat/deepep_metadata_cuda/__init__.py
+4 −2 sonicmoe/ernie_compat/deepep_topk_metadata_cuda/__init__.py
+103 −28 sonicmoe/ernie_compat/deepep_topk_metadata_cuda/kernel.cu
+218 −69 sonicmoe/ernie_compat/mlp_node_v2.py
+1,075 −316 sonicmoe/functional/__init__.py
+1 −1 sonicmoe/functional/backward.py
+1 −1 sonicmoe/functional/forward.py
+2 −2 sonicmoe/functional/fp8_config.py
+1 −1 sonicmoe/functional/topk_softmax.py
+35 −22 sonicmoe/jit.py
+36 −22 sonicmoe/moe.py
+7 −0 sonicmoe/quack_utils/__init__.py
+1,650 −41 sonicmoe/quack_utils/_gated_epilogues.py
+889 −119 sonicmoe/quack_utils/blockscaled_fp8_gemm.py
+14 −13 sonicmoe/quack_utils/fused_quant_kernels.py
+145 −7 sonicmoe/quack_utils/gemm_dgated.py
+88 −6 sonicmoe/quack_utils/gemm_gated.py
+57 −11 sonicmoe/quack_utils/gemm_interface.py
+625 −5 sonicmoe/quack_utils/gemm_sm100_fp8_zeromat.py
+2 −1 sonicmoe/quack_utils/swiglu_triton.py
+1 −0 tests/ops/INDEX.md
+23 −2 tests/ops/test_colwise_quant.py
+179 −0 tests/ops/test_dgated_fp8d_iso32_correctness.py
+199 −0 tests/ops/test_dgated_true2cta_stress.py
+56 −4 tests/ops/test_gemm_dgated.py
+3 −1 tests/ops/test_gemm_dgated_fp8cload.py
+54 −1 tests/ops/test_gemm_gated.py
27 changes: 20 additions & 7 deletions src/paddlefleet/transformer/moe/moe_expert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import paddle
import paddle.nn.functional as F
import paddlefleet_ops
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
build_sharded_state_dict,
shard_weight,
Expand All @@ -38,11 +39,15 @@
k_grouped_bf16_gemm_tn_contiguous_aligned,
)

try:
if paddlefleet_ops.is_deep_gemm_available():
from paddlefleet_ops import deep_gemm as paddlefleet_deep_gemm
from paddlefleet_ops.sonicmoe.quack_utils import quantize_native_fp8_weights
except (ImportError, RuntimeError):
pass

if paddlefleet_ops.is_sonic_moe_available():
from paddlefleet_ops.sonicmoe.functional import clear_all_fp8_weight_caches
from paddlefleet_ops.sonicmoe.quack_utils import (
install_native_fp8_weight_cache,
quantize_native_fp8_weights,
)


class BMMFunction(paddle.autograd.PyLayer):
Expand Down Expand Up @@ -327,6 +332,10 @@ class SonicMoEExpert(GroupedMLPExpert):
_GROUPED_LAYOUT = "grouped"
_SONIC_LAYOUT = "sonic"

@staticmethod
def clear_fp8_weight_cache():
clear_all_fp8_weight_caches()

@staticmethod
def _grouped_w1_to_sonic(weight):
gate, up = paddle.chunk(weight, 2, axis=-1)
Expand Down Expand Up @@ -423,10 +432,14 @@ def step(self):
def quant_weight(self):
self.convert_weights_to_sonic_layout()

payload = quantize_native_fp8_weights(
self.weight1.permute([1, 2, 0]),
self.weight2.permute([1, 2, 0]),
w1_sonic = self.weight1.permute([1, 2, 0])
w2_sonic = self.weight2.permute([1, 2, 0])
payload = quantize_native_fp8_weights(w1_sonic, w2_sonic, iso32=False)
install_native_fp8_weight_cache(
w1_sonic, w2_sonic, payload, iso32=False
)
self._native_fp8_weight_payload = payload

assert payload["format"] == "1x32", (
f"quant strategy {payload.get('format')} is not supported."
)
Expand Down
Loading