Add a Ray trainer for MaxText

[richardsliu]

Description

Create a Ray training script that can train MaxText on a multislice TPU. Uses ray-tpu (https://github.com/AI-Hypercomputer/ray-tpu) to initialze runtime environment on multislice TPU deployments.

Tests

Manually tested with:

ray job submit \
    --working-dir . \
    --runtime-env-json='{"excludes": [".git", "MaxText/test_assets"]}' \
    -- python MaxText/ray_trainer.py MaxText/configs/base.yml \
        base_output_directory=/tmp/maxtext \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        max_target_length=8192 \
        model_name=default \
        steps=100

Checklist

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

• [ X ] I have performed a self-review of my code.
• [ X ] I have necessary comments in my code, particularly in hard-to-understand areas.
• [ X ] I have run end-to-end tests tests and provided workload links above if applicable.
• [ X ] I have made or will make corresponding changes to the doc if needed.