# Copyright (c) 2023 Graphcore Ltd. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numpy as np from poptransformer import ops from poptransformer.layers import BaseLayer from poptransformer.layers import Linear from poptransformer.utils.param_handler.tensor_parallel_strategy import shard class BaseRWKVAttention(BaseLayer): def __init__(self, context, name, hidden_size, attention_hidden_size, layer_id): super().__init__(context, name) self.hidden_size = hidden_size self.attention_hidden_size = attention_hidden_size self.layer_id = layer_id self.collect_bind_layer_weights() def collect_bind_layer_weights(self): self.key_linear = Linear( self.context, 'key', self.hidden_size, self.attention_hidden_size, use_bias=False) self.receptance_linear = Linear( self.context, 'receptance', self.hidden_size, self.attention_hidden_size, use_bias=False) self.value_linear = Linear( self.context, 'value', self.hidden_size, self.attention_hidden_size, use_bias=False) self.output_linear = Linear( self.context, 'output', self.attention_hidden_size, self.hidden_size, use_bias=False) time_decay_key = '.'.join([self.context, 'time_decay']) time_decay_np = self.get_param_from_state_dict(time_decay_key, [self.hidden_size]) self.time_decay = self.add_initialized_input_tensor(time_decay_np, time_decay_key) time_first_key = '.'.join([self.context, 'time_first']) time_first_np = self.get_param_from_state_dict(time_first_key, [self.hidden_size]) self.time_first = self.add_initialized_input_tensor(time_first_np, time_first_key) time_mix_key_name = '.'.join([self.context, 'time_mix_key']) time_mix_key_np = self.get_param_from_state_dict(time_mix_key_name, [1, 1, self.hidden_size]) self.time_mix_key = self.add_initialized_input_tensor(time_mix_key_np, time_mix_key_name) time_mix_value_name = '.'.join([self.context, 'time_mix_value']) time_mix_value_np = self.get_param_from_state_dict(time_mix_value_name, [1, 1, self.hidden_size]) self.time_mix_value = self.add_initialized_input_tensor(time_mix_value_np, time_mix_value_name) time_mix_receptance_name = '.'.join([self.context, 'time_mix_receptance']) time_mix_receptance_np = self.get_param_from_state_dict(time_mix_receptance_name, [1, 1, self.hidden_size]) self.time_mix_receptance = self.add_initialized_input_tensor(time_mix_receptance_np, time_mix_receptance_name) def gate(self, graph, a, w, b): y1 = ops.mul(graph, a, w) y2 = ops.sub(graph, ops.constant(graph, np.array(1).astype(self.np_float_type)), w) y2 = ops.mul(graph, y2, b) y = ops.add(graph, y1, y2) return y def __call__(self, graph, hidden, layer_state): with graph.nameScope('extract_key_value'): num_state, den_state, max_state = layer_state[2:] key = self.gate(graph, hidden, self.time_mix_key, layer_state[1]) value = self.gate(graph, hidden, self.time_mix_value, layer_state[1]) receptance = self.gate(graph, hidden, self.time_mix_receptance, layer_state[1]) layer_state[1] = hidden key = self.key_linear(graph, key) value = self.value_linear(graph, value) receptance = self.receptance_linear(graph, receptance) receptance = ops.sigmoid(graph, receptance) if self.precision == 'fp16': time_decay = ops.cast(graph, self.time_decay, 'FLOAT') key = ops.cast(graph, key, 'FLOAT') value = ops.cast(graph, value, 'FLOAT') time_first = ops.cast(graph, self.time_first, 'FLOAT') with graph.nameScope('rwkv_linear_attention'): temp1 = ops.add(graph, key, time_first) max_for_output = ops.maximum(graph, max_state, temp1) e1 = ops.exp(graph, ops.sub(graph, max_state, max_for_output)) e2 = ops.exp(graph, ops.sub(graph, temp1, max_for_output)) numerator = ops.add(graph, ops.mul(graph, e1, num_state), ops.mul(graph, e2, value)) denominator = ops.add(graph, ops.mul(graph, e1, den_state), e2) output = ops.div(graph, numerator, denominator) time_decay = ops.exp(graph, time_decay) temp2 = ops.sub(graph, max_state, time_decay) max_for_state = ops.maximum(graph, temp2, key) e1 = ops.exp(graph, ops.sub(graph, temp2, max_for_state)) e2 = ops.exp(graph, ops.sub(graph, key, max_for_state)) layer_state[2] = ops.add(graph, ops.mul(graph, e1, num_state), ops.mul(graph, e2, value)) layer_state[3] = ops.add(graph, ops.mul(graph, e1, den_state), e2) layer_state[4] = max_for_state if self.precision == 'fp16': output = ops.cast(graph, output, self.popart_float_type) output = ops.mul(graph, receptance, output) output = self.output_linear(graph, output) return output, layer_state class TPRWKVAttention(BaseRWKVAttention): def collect_bind_layer_weights(self): key_tp_setting = { 'strategy_name': 'start', } output_tp_setting = { 'strategy_name': 'end', } self.key_linear = Linear( self.context, 'key', self.hidden_size, self.attention_hidden_size, use_bias=False, **key_tp_setting) self.receptance_linear = Linear( self.context, 'receptance', self.hidden_size, self.attention_hidden_size, use_bias=False, **key_tp_setting) self.value_linear = Linear( self.context, 'value', self.hidden_size, self.attention_hidden_size, use_bias=False, **key_tp_setting) self.output_linear = Linear( self.context, 'output', self.attention_hidden_size, self.hidden_size, use_bias=False, **output_tp_setting) vs_setting = {'vs_type': 'consecutive', 'group_size': 1} time_decay_key = '.'.join([self.context, 'time_decay']) time_decay_np = self.get_param_from_state_dict(time_decay_key, [self.hidden_size]) time_decay_np = shard(time_decay_np, self.num_replicas, -1) self.time_decay = self.add_initialized_input_tensor(time_decay_np, time_decay_key, **vs_setting) time_first_key = '.'.join([self.context, 'time_first']) time_first_np = self.get_param_from_state_dict(time_first_key, [self.hidden_size]) time_first_np = shard(time_first_np, self.num_replicas, -1) self.time_first = self.add_initialized_input_tensor(time_first_np, time_first_key, **vs_setting) time_mix_key_name = '.'.join([self.context, 'time_mix_key']) time_mix_key_np = self.get_param_from_state_dict(time_mix_key_name, [1, 1, self.hidden_size]) self.time_mix_key = self.add_initialized_input_tensor(time_mix_key_np, time_mix_key_name) time_mix_value_name = '.'.join([self.context, 'time_mix_value']) time_mix_value_np = self.get_param_from_state_dict(time_mix_value_name, [1, 1, self.hidden_size]) self.time_mix_value = self.add_initialized_input_tensor(time_mix_value_np, time_mix_value_name) time_mix_receptance_name = '.'.join([self.context, 'time_mix_receptance']) time_mix_receptance_np = self.get_param_from_state_dict(time_mix_receptance_name, [1, 1, self.hidden_size]) self.time_mix_receptance = self.add_initialized_input_tensor(time_mix_receptance_np, time_mix_receptance_name) class RWKVAttention(TPRWKVAttention, BaseRWKVAttention): layer_class_map = { 'tp': TPRWKVAttention, 'shard': BaseRWKVAttention} def __init__(self, context, name, hidden_size, attention_hidden_size, layer_id): model_type = self.model_type self.layer_class = self.layer_class_map.get(model_type, None) if not self.layer_class: raise ValueError(f"Invalid model_type {model_type}, options: {self.layer_class_map.keys()}") self.logger.debug(f'initializing model type: {self.layer_class.__name__}') super().__init__(context, name, hidden_size, attention_hidden_size, layer_id) def __call__(self, graph, hidden, layer_state): return self.layer_class.__call__(self, graph, hidden, layer_state) def collect_bind_layer_weights(self): return self.layer_class.collect_bind_layer_weights(self)