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

# OR standalone script where you want to invoke 

    MAXTEXT_CONFIGS_DIR = "src/maxtext/configs"
    import maxtext as mt

    config= mt.pyconfig(model_name=MAXTEXT_MODEL_VERSION, 
    hf_access_token=args.hf_token,
    base_output_directory="gs://path/to/save/artifacts",
    base_config=f"{MAXTEXT_CONFIGS_DIR}/post_train/rl.yml",
    )
    qwen2_actor, _= mt.from_pretrained(config) # ref_mesh and train_mesh are the same for us
  else:
    qwen2_actor = params_lib.create_model_from_safe_tensors(
        MODEL_PATH, config, trainer_mesh, dtype=MODEL_DTYPE
    )

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 47.41379% with 61 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/model_creation_utils.py	48.11%	46 Missing and 9 partials ⚠️
src/maxtext/trainers/post_train/rl/train_rl.py	33.33%	6 Missing ⚠️

📢 Thoughts on this report? Let us know!

[wang2yn84]

Thank you very much!

[richjames0]

lgtm