It's kind of working!

Still unclear how to set input ids

decoder_only_t5/modeling.py

@@ -1,4 +1,5 @@
 import copy
+import math
 from typing import Optional, Tuple, Union
 
 import torch
@@ -19,6 +20,39 @@ logger = logging.get_logger(__name__)
 _CONFIG_FOR_DOC = "DecoderOnlyT5Config"
 
 
+class DecoderOnlyT5LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6, use_scale=True, center_scale_at_zero=False):
+        """
+        Construct a layernorm module in the T5 style No bias and no subtraction of mean.
+        """
+        super().__init__()
+        if use_scale:
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+        else:
+            assert not center_scale_at_zero
+            self.weight = None
+        self.center_scale_at_zero = center_scale_at_zero
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/components/layer_norm.py#L30
+
+        # layer norm should always be calculated in float32
+        mean2 = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(mean2 + self.variance_epsilon)
+
+        # convert into float16 if necessary
+        if self.weight is None:
+            return hidden_states        
+        if self.weight.dtype == torch.float16:
+            hidden_states = hidden_states.to(torch.float16)
+        if self.center_scale_at_zero:
+            return (self.weight + 1.0) * hidden_states
+        else:
+            return self.weight * hidden_states
+            
+
+
 class DecoderOnlyT5LayerFF(modeling_t5.T5LayerFF):
     def __init__(self, config: DecoderOnlyT5Config):
         super(modeling_t5.T5LayerFF, self).__init__()
@@ -28,7 +62,7 @@ class DecoderOnlyT5LayerFF(modeling_t5.T5LayerFF):
             self.DenseReluDense = modeling_t5.T5DenseActDense(config)
 
         if not config.parallel_layers:
-            self.layer_norm = modeling_t5.T5LayerNorm(
+            self.layer_norm = modeling_t5.DecoderOnlyT5LayerNorm(
                 config.d_model, eps=config.layer_norm_epsilon
             )
         else:
@@ -37,7 +71,7 @@ class DecoderOnlyT5LayerFF(modeling_t5.T5LayerFF):
 
 
 # LlamaRotaryEmbedding
-class T5DecoderOnlyRotaryEmbedding(nn.Module):
+class DecoderOnlyT5RotaryEmbedding(nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
         super().__init__()
 
@@ -139,25 +173,21 @@ class DecoderOnlyT5Attention(modeling_t5.T5Attention):
     def __init__(self, config: DecoderOnlyT5Config, has_relative_attention_bias=False):
         super(modeling_t5.T5Attention, self).__init__()
         self.is_decoder = config.is_decoder
-        self.has_relative_attention_bias = has_relative_attention_bias
-        self.relative_attention_num_buckets = config.relative_attention_num_buckets
-        self.relative_attention_max_distance = config.relative_attention_max_distance
+        assert not has_relative_attention_bias
+        assert config.use_rotary_embedding
         self.d_model = config.d_model
-        self.key_value_proj_dim = config.d_kv
-        self.n_heads = config.num_heads
-        self.n_kv_heads = 1 if config.multi_query_attention else self.n_heads
-        self.n_kv_groups = self.n_heads // self.n_kv_heads
-        self.dropout = config.dropout_rate
-        self.inner_dim = self.n_heads * self.key_value_proj_dim
-        self.kv_inner_dim = self.n_kv_heads * self.key_value_proj_dim
-        if config.use_rotary_embedding:
-            self.rotary_embedding = T5DecoderOnlyRotaryEmbedding(
-                self.key_value_proj_dim,
-                max_position_embeddings=config.relative_attention_max_distance,
-                base=config.rotary_embedding_max_timescale,
-            )
-        else:
-            self.rotary_embedding = None
+        self.head_dim = config.d_kv
+        self.num_heads = config.num_heads
+        self.num_key_value_heads = 1 if config.multi_query_attention else self.n_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.attention_dropout = config.dropout_rate
+        self.inner_dim = self.num_heads * self.head_dim
+        self.kv_inner_dim = self.num_key_value_heads * self.head_dim    
+        self.rotary_emb = DecoderOnlyT5RotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=config.relative_attention_max_distance,
+            base=config.rotary_embedding_max_timescale,
+        )
 
         # Mesh TensorFlow initialization to avoid scaling before softmax
         self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)
@@ -165,179 +195,79 @@ class DecoderOnlyT5Attention(modeling_t5.T5Attention):
         self.v = nn.Linear(self.d_model, self.kv_inner_dim, bias=False)
         self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)
 
-        if self.has_relative_attention_bias:
-            self.relative_attention_bias = nn.Embedding(
-                self.relative_attention_num_buckets, self.n_heads
-            )
         self.pruned_heads = set()
         self.gradient_checkpointing = False
 
     def forward(
         self,
-        hidden_states,
-        mask=None,
+        hidden_states: torch.Tensor,
         key_value_states=None,
         position_bias=None,
-        position_ids=None,
-        past_key_value=None,
+        mask: Optional[torch.Tensor] = None,
         layer_head_mask=None,
-        query_length=None,
-        use_cache=False,
-        output_attentions=False,
-    ):
-        """
-        Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
-        """
-        # Input is (batch_size, seq_length, dim)
-        # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)
-        # past_key_value[0] is (batch_size, n_kv_heads, q_len - 1, dim_per_head)
-        batch_size, seq_length = hidden_states.shape[:2]
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        assert key_value_states is None
+        assert position_bias is None
+        assert layer_head_mask is None
+
+        bsz, q_len, _ = hidden_states.size()
 
-        real_seq_length = seq_length
+        query_states = self.q(hidden_states)
+        key_states = self.k(hidden_states)
+        value_states = self.v(hidden_states)
 
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
-            if len(past_key_value) != 2:
-                raise ValueError(
-                    f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
-                )
-            real_seq_length += (
-                past_key_value[0].shape[2] if query_length is None else query_length
-            )
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
-        key_length = (
-            real_seq_length if key_value_states is None else key_value_states.shape[1]
-        )
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        def shape(states, n_heads):
-            """projection"""
-            return states.view(
-                batch_size, -1, n_heads, self.key_value_proj_dim
-            ).transpose(1, 2)
+        past_key_value = (key_states, value_states) if use_cache else None
 
-        def unshape(states):
-            """reshape"""
-            return (
-                states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)
-            )
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
 
-        def project(hidden_states, proj_layer, key_value_states, past_key_value):
-            """projects hidden states correctly to key/query states"""
-            if key_value_states is None:
-                # self-attn
-                # (batch_size, n_kv_heads, seq_length, dim_per_head)
-                hidden_states = shape(proj_layer(hidden_states), self.n_kv_heads)
-            elif past_key_value is None:
-                # cross-attn
-                # (batch_size, n_kv_heads, seq_length, dim_per_head)
-                hidden_states = shape(proj_layer(key_value_states), self.n_kv_heads)
-            return hidden_states
-
-        def concat_past_key_value(hidden_states, past_key_value, key_value_states):
-            if key_value_states is None:
-                # self-attn
-                # (batch_size, n_kv_heads, key_length, dim_per_head)
-                hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
-            elif past_key_value.shape[2] != key_value_states.shape[1]:
-                # checking that the `sequence_length` of the `past_key_value` is the same as
-                # the provided `key_value_states` to support prefix tuning
-                # cross-attn
-                # (batch_size, n_kv_heads, seq_length, dim_per_head)
-                raise NotImplementedError(
-                    "cross attention with RoPE and past KV is not implemented"
-                )
-                # hidden_states = shape(proj_layer(key_value_states), self.n_kv_heads)
-            else:
-                # cross-attn
-                hidden_states = past_key_value
-            return hidden_states
-
-        # get query states
-        query_states = shape(
-            self.q(hidden_states), self.n_heads
-        )  # (batch_size, n_heads, seq_length, dim_per_head)
-
-        # get key/value states
-        key_states = project(hidden_states, self.k, key_value_states, past_key_value)
-        value_states = project(hidden_states, self.v, key_value_states, past_key_value)
-
-        # RoPE
-        if self.rotary_embedding is not None:
-            kv_seq_len = key_states.shape[-2]
-            if past_key_value:
-                kv_seq_len += past_key_value[0].shape[-2]
-            cos, sin = self.rotary_embedding(query_states, seq_len=kv_seq_len)
-            query_states, key_states = apply_rotary_pos_emb(
-                query_states, key_states, cos, sin, position_ids
-            )
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
 
-        # concat past
-        if past_key_value is not None:
-            key_states = concat_past_key_value(
-                key_states,
-                past_key_value[0],
-                key_value_states,
-            )
-            value_states = concat_past_key_value(
-                value_states,
-                past_key_value[1],
-                key_value_states,
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
             )
 
-        # MultiQueryDotProductAttention
-        key_states = repeat_kv(key_states, self.n_kv_groups)
-        value_states = repeat_kv(value_states, self.n_kv_groups)
-
-        # compute scores
-        scores = torch.matmul(
-            query_states, key_states.transpose(3, 2)
-        )  # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
-
-        if position_bias is None:
-            if not self.has_relative_attention_bias:
-                position_bias = torch.zeros(
-                    (1, self.n_heads, real_seq_length, key_length),
-                    device=scores.device,
-                    dtype=scores.dtype,
-                )
-                if self.gradient_checkpointing and self.training:
-                    position_bias.requires_grad = True
-            else:
-                position_bias = self.compute_bias(
-                    real_seq_length, key_length, device=scores.device
+        if mask is not None:
+            if mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {mask.size()}"
                 )
+            attn_weights = attn_weights + mask
 
-            # if key and values are already calculated
-            # we want only the last query position bias
-            if past_key_value is not None:
-                position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout)
+        attn_output = torch.matmul(attn_weights, value_states)
 
-            if mask is not None:
-                position_bias = (
-                    position_bias + mask
-                )  # (batch_size, n_heads, seq_length, key_length)
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
 
-        if self.pruned_heads:
-            mask = torch.ones(position_bias.shape[1])
-            mask[list(self.pruned_heads)] = 0
-            position_bias_masked = position_bias[:, mask.bool()]
-        else:
-            position_bias_masked = position_bias
-
-        scores += position_bias_masked
-        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
-            scores
-        )  # (batch_size, n_heads, seq_length, key_length)
-        attn_weights = nn.functional.dropout(
-            attn_weights, p=self.dropout, training=self.training
-        )  # (batch_size, n_heads, seq_length, key_length)
-
-        # Mask heads if we want to
-        if layer_head_mask is not None:
-            attn_weights = attn_weights * layer_head_mask
-
-        attn_output = unshape(
-            torch.matmul(attn_weights, value_states)
-        )  # (batch_size, seq_length, dim)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.inner_dim)
         attn_output = self.o(attn_output)
 
         present_key_value_state = (
@@ -356,8 +286,11 @@ class DecoderOnlyT5LayerSelfAttention(modeling_t5.T5LayerSelfAttention):
         self.SelfAttention = DecoderOnlyT5Attention(
             config, has_relative_attention_bias=has_relative_attention_bias
         )
-        self.layer_norm = modeling_t5.T5LayerNorm(
-            config.d_model, eps=config.layer_norm_epsilon
+        self.layer_norm = DecoderOnlyT5LayerNorm(
+            config.d_model,
+            eps=config.layer_norm_epsilon,
+            use_scale=True,
+            center_scale_at_zero=True,
         )
         self.dropout = nn.Dropout(config.dropout_rate)
         self.parallel_layers = config.parallel_layers
@@ -425,20 +358,19 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
         position_bias=None,
         position_ids=None,
         encoder_hidden_states=None,
-        encoder_attention_mask=None,
-        encoder_decoder_position_bias=None,
         layer_head_mask=None,
-        cross_attn_layer_head_mask=None,
         past_key_value=None,
         use_cache=False,
         output_attentions=False,
+        encoder_attention_mask=None,
+        encoder_decoder_position_bias=None,
+        cross_attn_layer_head_mask=None,
         return_dict=True,
     ):
+        assert encoder_attention_mask is None
+        assert encoder_decoder_position_bias is None
+        assert cross_attn_layer_head_mask is None
         if past_key_value is not None:
-            if not self.is_decoder:
-                logger.warning(
-                    "`past_key_values` is passed to the encoder. Please make sure this is intended."
-                )
             expected_num_past_key_values = 2 if encoder_hidden_states is None else 4
 
             if len(past_key_value) != expected_num_past_key_values:
@@ -447,11 +379,9 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
                     f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}"
                     f"Got {len(past_key_value)} past key / value states"
                 )
-
             self_attn_past_key_value = past_key_value[:2]
-            cross_attn_past_key_value = past_key_value[2:]
         else:
-            self_attn_past_key_value, cross_attn_past_key_value = None, None
+            self_attn_past_key_value = None
 
         ff_layer = self.layer[-1]
         if self.parallel_layers:
@@ -490,45 +420,7 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
             and not self.is_decoder_only
             and encoder_hidden_states is not None
         )
-        if do_cross_attention:
-            # the actual query length is unknown for cross attention
-            # if using past key value states. Need to inject it here
-            if present_key_value_state is not None:
-                query_length = present_key_value_state[0].shape[2]
-            else:
-                query_length = None
-
-            cross_attention_outputs = self.layer[1](
-                x,
-                key_value_states=encoder_hidden_states,
-                attention_mask=encoder_attention_mask,
-                position_bias=encoder_decoder_position_bias,
-                # position_ids ?
-                layer_head_mask=cross_attn_layer_head_mask,
-                past_key_value=cross_attn_past_key_value,
-                query_length=query_length,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-            )
-            x = cross_attention_outputs[0]
-
-            # clamp inf values to enable fp16 training
-            if x.dtype == torch.float16:
-                clamp_value = torch.where(
-                    torch.isinf(x).any(),
-                    torch.finfo(x.dtype).max - 1000,
-                    torch.finfo(x.dtype).max,
-                )
-                x = torch.clamp(x, min=-clamp_value, max=clamp_value)
-
-            # Combine self attn and cross attn key value states
-            if present_key_value_state is not None:
-                present_key_value_state = (
-                    present_key_value_state + cross_attention_outputs[1]
-                )
-
-            # Keep cross-attention outputs and relative position weights
-            attention_outputs = attention_outputs + cross_attention_outputs[2:]
+        assert not do_cross_attention
 
         if self.parallel_layers:
             # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/architectures/t5/t5_architecture.py#L534-L578
@@ -577,12 +469,12 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
                 for i in range(config.num_layers)
             ]
         )
-        if not config.parallel_layers:
-            self.final_layer_norm = modeling_t5.T5LayerNorm(
-                config.d_model, eps=config.layer_norm_epsilon
-            )
-        else:
-            self.final_layer_norm = nn.Identity()
+        self.final_layer_norm = DecoderOnlyT5LayerNorm(
+            config.d_model,
+            eps=config.layer_norm_epsilon,
+            use_scale=False,
+            center_scale_at_zero=False,
+        )
         self.dropout = nn.Dropout(config.dropout_rate)
 
         # Initialize weights and apply final processing
@@ -654,8 +546,7 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
                 seq_length + past_key_values_length,
                 dtype=torch.long,
                 device=device,
-            )
-            position_ids = position_ids.unsqueeze(0)
+            ).unsqueeze(0)
 
         if inputs_embeds is None:
             if self.embed_tokens is None:
@@ -683,18 +574,6 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
             attention_mask = torch.ones(
                 batch_size, mask_seq_length, device=inputs_embeds.device
             )
-        if (
-            self.is_decoder
-            and encoder_attention_mask is None
-            and encoder_hidden_states is not None
-        ):
-            encoder_seq_length = encoder_hidden_states.shape[1]
-            encoder_attention_mask = torch.ones(
-                batch_size,
-                encoder_seq_length,
-                device=inputs_embeds.device,
-                dtype=torch.long,
-            )
 
         # initialize past_key_values with `None` if past does not exist
         if past_key_values is None:
@@ -706,25 +585,6 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
             attention_mask, input_shape
         )
 
-        # If a 2D or 3D attention mask is provided for the cross-attention
-        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
-        if self.is_decoder and encoder_hidden_states is not None:
-            (
-                encoder_batch_size,
-                encoder_sequence_length,
-                _,
-            ) = encoder_hidden_states.size()
-            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
-            if encoder_attention_mask is None:
-                encoder_attention_mask = torch.ones(
-                    encoder_hidden_shape, device=inputs_embeds.device
-                )
-            encoder_extended_attention_mask = self.invert_attention_mask(
-                encoder_attention_mask
-            )
-        else:
-            encoder_extended_attention_mask = None
-
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(
@@ -742,7 +602,6 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
         all_attentions = () if output_attentions else None
         all_cross_attentions = () if (output_attentions and self.is_decoder) else None
         position_bias = None
-        encoder_decoder_position_bias = None
 
         hidden_states = self.dropout(inputs_embeds)
 
@@ -758,25 +617,10 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
                 if attention_mask is not None:
                     attention_mask = attention_mask.to(hidden_states.device)
                 if position_bias is not None:
-                    position_bias = position_bias.to(hidden_states.device)
-                if encoder_hidden_states is not None:
-                    encoder_hidden_states = encoder_hidden_states.to(
-                        hidden_states.device
-                    )
-                if encoder_extended_attention_mask is not None:
-                    encoder_extended_attention_mask = (
-                        encoder_extended_attention_mask.to(hidden_states.device)
-                    )
-                if encoder_decoder_position_bias is not None:
-                    encoder_decoder_position_bias = encoder_decoder_position_bias.to(
-                        hidden_states.device
-                    )
+                    position_bias = position_bias.to(hidden_states.device)                
                 if layer_head_mask is not None:
                     layer_head_mask = layer_head_mask.to(hidden_states.device)
-                if cross_attn_layer_head_mask is not None:
-                    cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(
-                        hidden_states.device
-                    )
+
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
@@ -786,9 +630,9 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
                     hidden_states,
                     extended_attention_mask,
                     position_bias,
-                    encoder_hidden_states,
-                    encoder_extended_attention_mask,
-                    encoder_decoder_position_bias,
+                    None,
+                    None,
+                    None,
                     layer_head_mask,
                     cross_attn_layer_head_mask,
                     None,  # past_key_value is always None with gradient checkpointing
@@ -801,9 +645,9 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
                     attention_mask=extended_attention_mask,
                     position_bias=position_bias,
                     position_ids=position_ids,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_attention_mask=encoder_extended_attention_mask,
-                    encoder_decoder_position_bias=encoder_decoder_position_bias,
+                    encoder_hidden_states=None,
+                    encoder_attention_mask=None,
+                    encoder_decoder_position_bias=None,
                     layer_head_mask=layer_head_mask,
                     cross_attn_layer_head_mask=cross_attn_layer_head_mask,
                     past_key_value=past_key_value,
@@ -822,10 +666,6 @@ class DecoderOnlyT5Stack(modeling_t5.T5Stack):
             # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights),
             # (cross-attention position bias), (cross-attention weights)
             position_bias = layer_outputs[2]
-            if self.is_decoder and encoder_hidden_states is not None:
-                encoder_decoder_position_bias = layer_outputs[
-                    4 if output_attentions else 3
-                ]
             # append next layer key value states
             if use_cache:
                 present_key_value_states = present_key_value_states + (
@@ -900,8 +740,6 @@ class DecoderOnlyT5Model(modeling_t5.T5ForConditionalGeneration):
     def _tie_weights(self):
         if not self.config.tie_word_embeddings:
             return
-        if self.encoder:
-            self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared)
         if self.decoder:
             self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared)