maxtext icon indicating copy to clipboard operation
maxtext copied to clipboard

Adding ml_goodput_measurement package to jax stable stack requirements

Open Perseus14 opened this issue 8 months ago • 0 comments

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.

Perseus14 avatar May 13 '25 17:05 Perseus14