add quant cache (#10969)

Author: liuruyan

paddlenlp/trainer/trainer_callback.py

@@ -27,6 +27,7 @@
 import numpy as np
 from tqdm.auto import tqdm
 
+from paddlenlp.transformers.moe_utils import offload, reload
 from paddlenlp.utils.log import logger
 
 from .trainer_utils import IntervalStrategy, has_length
@@ -646,5 +647,23 @@ def on_step_begin(self, args, state, control, **kwargs):
         if not g_shard_bypass_dygraph_optimizer or skip_count == 0:
             model.fp8_quant_weight(True)
             optimizer.clear_param_storage("moe_expert")
+            optimizer.clear_param_storage("rms_linear")
+            optimizer.clear_param_storage("memory_attn")
+            optimizer.clear_param_storage("attn_out_project")
+            optimizer.clear_param_storage("shared_expert")
+
+            self.moe_weights_name = []
+            for param in optimizer._inner_opt._parameter_list:
+                color = getattr(param, "color", -1)
+                if isinstance(color, dict) and color["color"] == "moe_expert":
+                    self.moe_weights_name.append(param.name)
+
+            for name in self.moe_weights_name:
+                offload(optimizer._master_weights[name])
 
         skip_count += 1
+
+    def on_optimizer_begin(self, args, state, control, **kwargs):
+        optimizer = kwargs["optimizer"]
+        for name in self.moe_weights_name:
+            reload(optimizer._master_weights[name])

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -88,6 +88,7 @@
     FP8LinearFunctionBase,
     FP8Mlp,
     cache_fp8_weight,
+    set_parameter_color,
 )
 from .fp8_linear import Linear
 
@@ -106,6 +107,7 @@ def swiglu(x, y=None):
             x, y = paddle.chunk(x, chunks=2, axis=-1)
         return F.silu(x) * y
 
+
 try:
     from paddle.incubate.nn.functional import fused_partial_rope
 except ImportError:
@@ -752,6 +754,7 @@ def forward(self, x):
 
 class FusedNormGateFunc(paddle.autograd.PyLayer):
     """recompute of postnorm and gate"""
+
     _current_norm_output = None
     _current_invar = None
 
@@ -799,6 +802,7 @@ def backward(ctx, d_gate_logits, d_norm_output):
 
         return dx, d_rms_norm_weight, d_moe_gate_weight
 
+
 class TemporaryVarContext:
     def __init__(self, norm_output, invar):
         self.norm_output = norm_output
@@ -810,6 +814,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         FusedNormGateFunc.clear_temporary_vars()
 
+
 def balance_expert_assignment(n, m, k):
     assert k * n % m == 0
     matrix = paddle.zeros((n, m), dtype=paddle.int32)
@@ -999,7 +1004,11 @@ def __init__(self, config: DeepseekV2Config, norm_weight=None, norm_eps=None):
 
         if config.offline_quant_expert_weight and config.clear_origin_weight_when_offline_quant:
             moe_grad_group = fleet.get_hybrid_communicate_group().expert_grad_comm_group
-            for p in self.experts.parameters():
+            expert_w1_list = [expert.w1 for expert in self.experts if expert is not None]
+            expert_w2_list = [expert.w2 for expert in self.experts if expert is not None]
+            for p in expert_w1_list:
+                setattr(p, "color", {"color": "moe_expert", "group": moe_grad_group})
+            for p in expert_w2_list:
                 setattr(p, "color", {"color": "moe_expert", "group": moe_grad_group})
 
         self.alpha = config.aux_loss_alpha
@@ -1019,6 +1028,7 @@ def __init__(self, config: DeepseekV2Config, norm_weight=None, norm_eps=None):
                 self.shared_experts = DeepseekV2MLPClass(
                     config=config, intermediate_size=intermediate_size, is_moe=False
                 )
+            set_parameter_color([self.shared_experts.w1, self.shared_experts.w2], "shared_expert")
 
     def fp8_quant_weight(self, batch_mode=False):
         """Quantize weights in FP8 format.
@@ -1171,7 +1181,16 @@ def qkv_pre_process(
 ):
     if (fused_partial_rope is None) or (position_ids is not None):
         return qkv_pre_process_no_fuse(
-            q, kv, k_pe, rotary_emb, num_heads, q_head_dim, qk_nope_head_dim, v_head_dim, qk_rope_head_dim, position_ids
+            q,
+            kv,
+            k_pe,
+            rotary_emb,
+            num_heads,
+            q_head_dim,
+            qk_nope_head_dim,
+            v_head_dim,
+            qk_rope_head_dim,
+            position_ids,
         )
 
     bsz, q_len, _ = q.shape
@@ -1712,6 +1731,7 @@ def __init__(
             kv_lora_rank,
             softmax_scale,
         )
+        set_parameter_color([self.q_up_weight, self.kv_up_weight], "memory_attn")
 
     def fp8_quant_weight(self):
         cache_fp8_weight(self.q_up_weight)
@@ -1839,6 +1859,7 @@ def __init__(self, hidden_size, q_out_dim, kv_outdim, eps=1e-6) -> None:
             is_bias=False,
         )
         self.eps = eps
+        set_parameter_color([self.q_down_weight], "rms_linear")
 
     def fp8_quant_weight(self):
         cache_fp8_weight(self.q_down_weight)
@@ -2237,6 +2258,8 @@ def fp8_quant_weight(self, batch_mode=False):
             # logger.info(f"fp8 quant weight for mlp {type(self.mlp)}")
             self.mlp.fp8_quant_weight(batch_mode)
             self.self_attn.fp8_quant_weight()
+        elif isinstance(self.mlp, FP8Mlp):
+            self.self_attn.fp8_quant_weight()
 
     def forward(
         self,

paddlenlp/transformers/fp8_utils.py

@@ -50,6 +50,22 @@ def swiglu(x, y=None):
 ]
 
 
+def set_parameter_color(
+    parameters, color, group=None, offline_quant_expert_weight=True, clear_origin_weight_when_offline_quant=True
+):
+    if offline_quant_expert_weight and clear_origin_weight_when_offline_quant:
+        if group is None:
+            for p in parameters:
+                if hasattr(p, "color") and p.color is not None:
+                    continue
+                setattr(p, "color", {"color": color})
+        else:
+            for p in parameters:
+                if hasattr(p, "color") and p.color is not None:
+                    continue
+                setattr(p, "color", {"color": color, "group": group})
+
+
 def extract_first_if_tuple(x):
     return x[0] if isinstance(x, tuple) else x
 
@@ -601,6 +617,7 @@ def __init__(self, in_features: int, out_features: int, bias_attr: bool = False)
             dtype="bfloat16",
             is_bias=False,
         )
+        set_parameter_color([self.weight], "attn_out_project")
 
     def fp8_quant_weight(self):
         cache_fp8_weight(self.weight)