PR #3189: Migrate MaxText/train.py to maxtext/trainers/pre_train/train.py

[bvandermoon]

PR #3189: Migrate MaxText/train.py to maxtext/trainers/pre_train/train.py

Imported from GitHub PR #3189

Description

• Move MaxText/train.py to maxtext/trainers/pre_train/train.py
• Create shim in MaxText/train.py to support old command. Include a deprecation warning
  • TODO: Add deprecation dates to this and other shims we are adding
• Update all old path references in repo
  • Left off links to old file that include a specific commit. Can be found by searching this regex: https://github\.com/AI-Hypercomputer/maxtext/.*trainers/pre_train/train. We can choose to update these links later once the commit id of the actually merged PR is available

Tests

Both new and old commands are working as expected:

New command:

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    run_name=<run_name> \
    base_output_directory=gs://<gcs_bucket> \
    dataset_type=synthetic \
    steps=10 \
    enable_checkpointing=false

Old command:

python3 -m MaxText.train src/maxtext/configs/base.yml \
    run_name=<run_name> \
    base_output_directory=gs://<gcs_bucket> \
    dataset_type=synthetic \
    steps=10 \
    enable_checkpointing=false

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.

Copybara import of the project:

--
3532d64 by Branden Vandermoon bvandermoon@google.com:

Migrate MaxText/train.py to maxtext/trainers/pre_train/train.py

Merging this change closes #3189

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3189 from AI-Hypercomputer:bvandermoon-restructure 3532d64

[codecov]

Codecov Report

❌ Patch coverage is 59.58188% with 116 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	62.73%	78 Missing and 23 partials ⚠️
src/MaxText/train.py	0.00%	14 Missing ⚠️
src/MaxText/estimator.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!