Unverified Commit 2ee60b75 authored by Dongkeun Yoon's avatar Dongkeun Yoon Committed by GitHub
Browse files

fix warning trigger for embed_positions when loading xglm (#25798)



* fix warning triggering for xglm.embed_positions

* Make TF variable a tf.constant to match (and fix some spelling)

---------
Co-authored-by: default avatarMatt <rocketknight1@gmail.com>
parent 5b5ee235
......@@ -63,7 +63,7 @@ TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST = [
LARGE_NEGATIVE = -1e8
def create_sinusiodal_positions(num_positions: int, embedding_dim: int, padding_idx: Optional[int]) -> tf.Tensor:
def create_sinusoidal_positions(num_positions: int, embedding_dim: int, padding_idx: Optional[int]) -> tf.Tensor:
half_dim = embedding_dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = tf.exp(tf.range(half_dim, dtype=tf.float32) * -emb)
......@@ -83,7 +83,7 @@ def create_sinusiodal_positions(num_positions: int, embedding_dim: int, padding_
)
emb *= _padding_mask
return tf.Variable(emb, trainable=False, name="model.embed_positions.weights")
return tf.constant(emb, name="embed_positions")
def _create_position_ids_from_input_ids(
......@@ -438,7 +438,7 @@ class TFXGLMMainLayer(tf.keras.layers.Layer):
)
self.offset = 2
self._embed_positions_weights = create_sinusiodal_positions(
self._embed_positions_weights = create_sinusoidal_positions(
num_positions=config.max_position_embeddings + self.offset,
embedding_dim=config.d_model,
padding_idx=config.pad_token_id,
......
......@@ -176,7 +176,7 @@ class XGLMSinusoidalPositionalEmbedding(nn.Module):
# in forward put the weights on the correct dtype and device of the param
emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
self.register_buffer("weights", emb_weights)
self.register_buffer("weights", emb_weights, persistent=False)
@staticmethod
def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment