[Flax tests/FlaxBert] make from_pretrained test faster
# What does this PR do?

Use flax checkpoint for the `from_pretrained` test and only test this for the base model to speed up the tests.
_The documentation is not available anymore as the PR was closed or merged._<|||||>Thanks!