From 5ab3848d6ffe71c818ce0d77e264bae0283b5ad1 Mon Sep 17 00:00:00 2001 From: huzhikun Date: Fri, 3 Jan 2025 10:31:14 +0800 Subject: [PATCH] =?UTF-8?q?=E3=80=90dev=E3=80=91page=20attention=E4=B8=ADe?= =?UTF-8?q?mpty=E7=AE=97=E5=AD=90=E5=88=9D=E5=A7=8B=E5=8C=96kv=20cache?= =?UTF-8?q?=E7=9A=84parameter=E6=96=B9=E6=A1=88?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../experimental/infer/core/transformer.py | 14 ++++++++++- mindformers/experimental/infer/core/utils.py | 23 ++++++++++++++++++- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/mindformers/experimental/infer/core/transformer.py b/mindformers/experimental/infer/core/transformer.py index 5694da2c2..3bce4df03 100644 --- a/mindformers/experimental/infer/core/transformer.py +++ b/mindformers/experimental/infer/core/transformer.py @@ -25,7 +25,7 @@ from mindformers.experimental.parallel_core.pynative.transformer.scale_mask_soft from mindformers.experimental.parallel_core.pynative.utils import divide from mindformers.experimental.infer.core import get_act_func, get_attn_mask_func, get_norm from mindformers.experimental.infer.core.layers import ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding -from mindformers.experimental.infer.core.utils import get_tp_world_size +from mindformers.experimental.infer.core.utils import get_tp_world_size, create_empty_parameter from mindformers.modules.flash_attention import FlashAttention from mindformers.modules.infer_attention import InferRotaryEmbedding from mindformers.modules.layers import FreqsMgr, RotaryEmbedding @@ -394,6 +394,18 @@ class ParallelAttention(nn.Cell): kv_shape, config.seq_length, compute_dtype=self.compute_dtype) + self.paged_attention_mgr.key_cache = create_empty_parameter( + shape=self.paged_attention_mgr.key_cache.shape, + dtype=self.paged_attention_mgr.key_cache.dtype, + name=self.paged_attention_mgr.key_cache.name, + requires_grad=self.paged_attention_mgr.key_cache.requires_grad, + ) + self.paged_attention_mgr.value_cache = create_empty_parameter( + shape=self.paged_attention_mgr.value_cache.shape, + dtype=self.paged_attention_mgr.value_cache.dtype, + name=self.paged_attention_mgr.value_cache.name, + requires_grad=self.paged_attention_mgr.value_cache.requires_grad, + ) self.rotary_embedding = InferRotaryEmbedding(rotary_cos_format=2) else: self.apply_rotary_emb = RotaryEmbedding(self.head_dim, config.rotary_dtype) diff --git a/mindformers/experimental/infer/core/utils.py b/mindformers/experimental/infer/core/utils.py index ed46beaea..7304a13d9 100644 --- a/mindformers/experimental/infer/core/utils.py +++ b/mindformers/experimental/infer/core/utils.py @@ -14,8 +14,9 @@ # ============================================================================ """ utils """ -from mindspore import Tensor, ops +from contextlib import contextmanager +from mindspore import Tensor, ops, Parameter, mint from mindformers.experimental.parallel_core.pynative.parallel_state import get_group_size, get_tensor_model_parallel_world_size __all__ = ["get_attn_mask_func", "generate_state_dict"] @@ -93,3 +94,23 @@ def generate_state_dict(network): def get_tp_world_size(): tp_size = get_tensor_model_parallel_world_size() return tp_size if tp_size else 1 + + +def create_empty_parameter(shape, *, dtype=None, device=None, **kwargs): + """Create an empty parameter.""" + def get_param(*args): + return [Tensor, args[0]] + + @contextmanager + def replace_class_method(cls, name, new_method): + old_method = getattr(cls, name) + + setattr(cls, name, new_method) + yield + setattr(cls, name, old_method) + + data = mint.empty(shape, dtype=dtype, device=device) + + with replace_class_method(Parameter, "_get_parameter_new_args", get_param): + param = Parameter(data, **kwargs) + return param -- Gitee