[llm] support tensorwise fp8/int8 training (#10612)

Author: lugimzzz

llm/run_finetune.py

@@ -166,6 +166,13 @@ def main():
         qlora_weight_blocksize=model_args.qlora_weight_blocksize,
         qlora_weight_double_quant=model_args.qlora_weight_double_quant,
         qlora_weight_double_quant_block_size=model_args.qlora_weight_double_quant_block_size,
+        apply_hadamard=model_args.apply_hadamard,
+        hadamard_block_size=model_args.hadamard_block_size,
+        quant_input_grad=model_args.quant_input_grad,
+        quant_weight_grad=model_args.quant_weight_grad,
+        apply_online_actscale_step=model_args.apply_online_actscale_step,
+        actscale_moving_rate=model_args.actscale_moving_rate,
+        fp8_format_type=model_args.fp8_format_type,
     )
 
     model_config = AutoConfig.from_pretrained(
@@ -447,7 +454,9 @@ def compute_metrics_do_generation(eval_preds):
         gen_args=gen_args,
         data_args=data_args,
     )
-    trainable_parameters = [p for p in model.parameters() if not p.stop_gradient]
+    trainable_parameters = [
+        p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name)
+    ]
     trainer.set_optimizer_grouped_parameters(trainable_parameters)
     if model_args.lorapro:
         optimizer = AdamWLoRAPro(

paddlenlp/quantization/hadamard_utils.py

@@ -14,6 +14,8 @@
 
 import paddle
 
+from paddlenlp.utils import infohub
+
 
 def matmul_hadU(X):
 
@@ -31,22 +33,37 @@ def matmul_hadU(X):
     return input.reshape(X.shape)
 
 
-def random_hadamard_matrix(size, dtype, is_block=False):
-    if not is_block:
-        A = paddle.randint(low=0, high=2, shape=[size, size]).astype("float32") * 2 - 1
-        Q, _ = paddle.linalg.qr(A)
-        return Q.astype(dtype), 1
+def create_hadamard_matrix(block_size, dtype):
+    Q = paddle.diag(paddle.ones((block_size), dtype=dtype))
+    block = matmul_hadU(Q)
+    return block
+
+
+def hadamard_matmul(input, side, hadamard_matrix, block_size):
+    # left -> H.T@input right -> input@H
+    origin_shape = input.shape
+    input = input.reshape([-1, origin_shape[-1]])
+    if side == "left":
+        # H.T@input -> (input.T@H).T
+        input = input.transpose([1, 0])
+    block_num = input.shape[-1] // block_size
+    output = input.reshape([-1, block_num, block_size]) @ hadamard_matrix
+    output = output.reshape([-1, block_num * block_size])
+    if side == "left":
+        output = output.transpose([1, 0])
+    output = output.reshape(origin_shape)
+
+    return output
+
+
+def apply_hadamard_matmul(x, side, block_size):
+    if getattr(infohub, "hadamard") is None:
+        setattr(infohub, "hadamard", {})
+
+    if block_size in infohub.hadamard:
+        hadamard_matrix = infohub.hadamard[block_size]
     else:
-        num_blocks = size
-        while not (num_blocks % 2):
-            num_blocks = num_blocks // 2
-        block_size = size // num_blocks
-        Q = paddle.diag(paddle.ones((block_size,), dtype="float32"))
-        block = matmul_hadU(Q)
-        large_matrix = paddle.zeros([size, size])
-
-        for i in range(num_blocks):
-            start_row = i * block_size
-            start_col = i * block_size
-            large_matrix[start_row : start_row + block_size, start_col : start_col + block_size] = block
-    return large_matrix.cast(dtype), block_size
+        hadamard_matrix = create_hadamard_matrix(block_size, x.dtype)
+        infohub.hadamard[block_size] = hadamard_matrix
+    target_x = hadamard_matmul(x, side, hadamard_matrix, block_size)
+    return target_x