Add Parallel Linear in ChannelWise Quant (#1739)

4389a804 · Chang Xu · GitHub · 8ff2de56 · 4389a804
隐藏空白更改
内联并排

Showing with 6 addition and 1 deletion

paddleslim/quant/observers/channel_wise.py paddleslim/quant/observers/channel_wise.py +6 -1

未找到文件。
--- a/paddleslim/quant/observers/channel_wise.py
+++ b/paddleslim/quant/observers/channel_wise.py
@@ -19,7 +19,12 @@ from .mse import MSEObserverLayer
 from .uniform import UniformObserver
 from paddle.quantization.factory import ObserverFactory
-CHANNEL_AXIS: Dict[type, int] = {paddle.nn.Conv2D: 0, paddle.nn.Linear: 1}
+CHANNEL_AXIS: Dict[type, int] = {
+    paddle.nn.Conv2D: 0,
+    paddle.nn.Linear: 1,
+    paddle.distributed.fleet.meta_parallel.ColumnParallelLinear: 1,
+    paddle.distributed.fleet.meta_parallel.RowParallelLinear: 1
+}
 class ChannelWiseObserver(UniformObserver):