Should gpt-j-6B model's embedding layer have bias?

xiaoda99 · July 20, 2022, 3:29am

In the original mesh-transformer-jax code , the embedding layer is implemented as haiku Linear layer, which has a bias parameter by default:

github.com

kingoflolz/mesh-transformer-jax/blob/master/mesh_transformer/layers.py#L185


      
              self.out_dim = out_dim
              self.in_dim_per_shard = in_dim // shards
              self.out_dim_per_shard = out_dim // shards
          
          
    if config["pe"] == "fixed":
                  embed_init = hk.initializers.TruncatedNormal(stddev=0.02)
                  self.positional_embeddings = hk.get_parameter('pos_embs', [config["seq"], self.out_dim_per_shard], init=embed_init)
              else:
                  self.positional_embeddings = None
          
          
    self.proj = hk.Linear(self.out_dim, w_init=hk.initializers.TruncatedNormal(stddev=1 / np.sqrt(in_dim)))
          
          
def __call__(self, x, dtype=jnp.bfloat16):
              shard_start_index = jax.lax.axis_index('shard') * self.in_dim_per_shard
          
          
    input_onehot = jax.nn.one_hot(x - shard_start_index, self.in_dim_per_shard)
              proj_out = self.proj(input_onehot)
          
          
    proj_out = g_psum(proj_out)
          
          
    if self.positional_embeddings is not None:

However, in HF transformers code, the embedding layer is implemented as plain nn.Embedding, which has no bias:

github.com

huggingface/transformers/blob/main/src/transformers/models/gptj/modeling_gptj.py#L487


      
          @add_start_docstrings(
              "The bare GPT-J Model transformer outputting raw hidden-states without any specific head on top.",
              GPTJ_START_DOCSTRING,
          )
          class GPTJModel(GPTJPreTrainedModel):
              def __init__(self, config):
                  super().__init__(config)
          
          
        self.embed_dim = config.n_embd
                  self.vocab_size = config.vocab_size
                  self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
                  self.drop = nn.Dropout(config.embd_pdrop)
                  self.h = nn.ModuleList([GPTJBlock(config) for _ in range(config.n_layer)])
                  self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
          
          
        # Model parallel
                  self.model_parallel = False
                  self.device_map = None
                  self.gradient_checkpointing = False
          
          
        # Initialize weights and apply final processing

Is this a bug when porting the mesh-transformer-jax gpt-j-6B model to HF?

Da Xiao

Topic		Replies	Views
How to understand the bias term in language model head (when we tie the word embeddings)? 🤗Transformers	0	907	September 12, 2022
How can i skip GPT2LMHeadModel embedding layers? 🤗Transformers	4	1004	February 17, 2023
Positional embedding in GPT-J when using `past_layer` Models	0	402	January 13, 2023
Positional Embeddings in Transformer Implementations 🤗Transformers	1	1717	September 3, 2024
Pretrain GPT-J-6B from scratch on Arabic Flax/JAX Projects	8	2390	August 15, 2021

Should gpt-j-6B model's embedding layer have bias?

Related topics