[examples/flax] add adafactor optimizer
# What does this PR do?

This PR adds the Adafactor optimizer in all language modeling scripts. This enables fine-tuning large models like 1.3B GPTNeo. Without adafactor the current CLM script can't even load the optimizer states. With Adafactor and bf16 it should now be possible to fit 1.3B GPTNeo on a single v3-8.

With 1.3B GPT Neo, one should be able to fit
- 8 per device batch size with max length 512
- 2 per device batch size with max_length 1024

```
python examples/flax/language-modeling/run_clm_flax.py \
    --model_name_or_path EleutherAI/gpt-neo-1.3B \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --do_train --do_eval \
    --block_size 512 \
    --num_train_epochs 1 \
    --learning_rate 1e-5 \
    --per_device_train_batch_size 8  --per_device_eval_batch_size 8 \
    --dtype bfloat16 \
    --adafactor \
    --overwrite_output_dir \
    --output_dir ~/tmp \