Removing Parallelization Causing Inconsistent Shape of Model Parameters
Author: kumulaor
Pull request: no

### System Info

I use the run_flax_glue.py file in transformers/examples/flax/text-classification to test and run the bert model. Now I want to remove the data parallelism used in the file. My modification to the code is to remove the use of **jax.pmap** and In the training loop, the use of **p_train_step** and **p_eval_step** was removed.

Then running my modified code an error occurred " **flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "embedding" in "/bert/embeddings/word_embeddings": (1, 30522, 768), (30522, 768 ).**“

Any help is appreciated. Thanks!

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Remove the use of jax.pmap and replace p_train_step and p_eval_step in the loop with train_step and eval_step
2. Run the modified run_flax_glue.py file
3. See the error

### Expected behavior

I hope it can run normally after deleting the data parallelism.

Comments:
I also removed the use of jax.max.pmean in train_step

The complete error is as follows:
```
Traceback (most recent call last):
  File "run_flax_glue.py", line 660, in <module>
    main()
  File "run_flax_glue.py", line 600, in main
    state, train_metric, dropout_rng = train_step(state, batch, dropout_rng)
  File "run_flax_glue.py", line 559, in train_step
    loss, grad = grad_fn(state.params)
  File "run_flax_glue.py", line 554, in loss_fn
    logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
  File "/public/home/ghfund3_a45/miniconda3/envs/qiuyang/lib/python3.8/site-packages/transformers/models/bert/modeling_flax_bert.py", line 937, in __call__
    outputs = self.module.apply(
  File "/public/home/ghfund3_a45/miniconda3/envs/qiuyang/lib/python3.8/site-packages/transformers/models/bert/modeling_flax_bert.py", line 1329, in __call__
    outputs = self.bert(
  File "/public/home/ghfund3_a45/miniconda3/envs/qiuyang/lib/python3.8/site-packages/transformers/models/bert/modeling_flax_bert.py", line 992, in __call__
    hidden_states = self.embeddings(
  File "/public/home/ghfund3_a45/miniconda3/envs/qiuyang/lib/python3.8/site-packages/transformers/models/bert/modeling_flax_bert.py", line 207, in __call__
    inputs_embeds = self.word_embeddings(input_ids.astype("i4"))
  File "/public/home/ghfund3_a45/miniconda3/envs/qiuyang/lib/python3.8/site-packages/flax/linen/linear.py", line 645, in setup
    self.embedding = self.param('embedding',
flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "embedding" in "/bert/embeddings/word_embeddings": (1, 30522, 768), (30522, 768). (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamShapeError)
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I met this problem too