maxtext
maxtext copied to clipboard
Adding ml_goodput_measurement package to jax stable stack requirements
Description
Add ml_goodput_measurement library to jax stable stack requirements file
If the change fixes a bug or a Github issue, please include a link, e.g.,: FIXES: b/416288633 (partially)
Tests
Please describe how you tested this change, and include any instructions and/or commands to reproduce.
- SSH into a TPU-v4p VM
- Run the following commands
BASE_IMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.5.2-rev1
bash docker_build_dependency_image.sh DEVICE=tpu MODE=stable_stack BASEIMAGE=${BASE_IMAGE}
docker run -v $PWD:/deps --rm -it --privileged --net host --entrypoint bash maxtext_base_image
python3 -m MaxText.train MaxText/configs/base.yml run_name=$RUN_NAME base_output_directory=$DIR dataset_path=gs://maxtext-dataset steps=10000
Goodput Logs: https://cloudlogging.app.goo.gl/Ewy48qU5oSNURfWq8
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.