Should gpt-j-6B model's embedding layer have bias?
In the original mesh-transformer-jax code , the embedding layer is implemented as haiku Linear layer, which has a bias parameter by default:github.comkingoflolz/mesh-transformer-jax/blob/master/mesh_transformer/layers.py#L185self.out_dim = out_dimself.in_dim_per_shard = in_dim // shardsself.out_dim_per_shard = out_dim // shardsif 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 = Noneself.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_shardinput_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.comhuggingface/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_embdself.vocab_size = config.vocab_sizeself.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 parallelself.model_parallel = Falseself.device_map = Noneself.gradient_checkpointing = False# Initialize weights and apply final processingIs this a bug when porting the mesh-transformer-jax gpt-j-6B model to HF?Da Xiao