Add simplified APIs for model obtaining maxtext models

[A9isha]

Description

Adding a simplified API of from_pretrained() which can work minimally as

For creating all the necessary configs and mesh for RL, we can use the APIs.

import maxtext as mt

config= mt.pyconfig(model_name="llama3.1-8b-Instruct")
model,mesh = mt.from_pretrained(config)

#OR
config= mt.pyconfig(model_name="llama3.1-8b-Instruct")
## your own mesh
model = mt.from_pretrained(config, mesh)


# OR for train_rl.py
trainer_config, sampler_config, trainer_devices, sampler_devices = mt.setup_configs_and_devices(model_name="llama3.1-8b-Instruct")

reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = mt.create_models_and_meshes(
      trainer_config, sampler_config, trainer_devices, sampler_devices
  )

# OR regular invocation of train_rl.py

run_name=maz-8b-$RANDOM python3 -m src.maxtext.trainers.post_train.rl.train_rl \
 model_name=llama3.1-8b-Instruct run_name=$run_name\
    steps=4 rollout_tensor_parallelism=-1\
 rollout_data_parallelism=1 rollout_expert_parallelism=1\
 test_batch_start_index=10  num_test_batches=15 
base_output_directory=/path/to/say/gcs/bucket

FIXES: b/492376313

Tests

Ran locally

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 9.48276% with 105 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/model_creation_utils.py	7.47%	99 Missing ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	37.50%	5 Missing ⚠️
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!