Update generated requirements for post-training with JAX 0.9.2

Author: SurbhiJainUSC

docs/development/update_dependencies.md

@@ -36,7 +36,7 @@ To update dependencies, you will follow these general steps:
 1. **Modify base requirements**: Update the desired dependencies in
    `src/dependencies/requirements/base_requirements/requirements.txt` or the hardware-specific pre-training files
    (`src/dependencies/requirements/base_requirements/tpu-requirements.txt`,
-   `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`).
+   `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`) or the post-training files (`src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`).
 2. **Find the JAX build commit hash**: The dependency generation process is
    pinned to a specific nightly build of JAX. You need to find the commit hash
    for the desired JAX build.
@@ -66,7 +66,7 @@ if you want to build `seed-env` from source.
 
 ## Step 1: Modify base requirements
 
-Update the desired dependencies in `src/dependencies/requirements/base_requirements/requirements.txt` or the hardware-specific pre-training files (`src/dependencies/requirements/base_requirements/tpu-requirements.txt`, `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`).
+Update the desired dependencies in `src/dependencies/requirements/base_requirements/requirements.txt` or the hardware-specific pre-training files (`src/dependencies/requirements/base_requirements/tpu-requirements.txt`, `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`) or the post-training files (`src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`).
 
 ## Step 2: Find the JAX build commit hash
 
@@ -98,6 +98,24 @@ mv generated_tpu_artifacts/tpu-requirements.txt \
   src/dependencies/requirements/generated_requirements/tpu-requirements.txt
 ```
 
+### TPU Post-Training
+
+If you have made changes to the post-training dependencies in `src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt`, you need to regenerate the pinned post-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:
+
+```bash
+seed-env \
+  --local-requirements=src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt \
+  --host-name=MaxText \
+  --seed-commit=<jax-build-commit-hash> \
+  --python-version=3.12 \
+  --requirements-txt=tpu-post-train-requirements.txt \
+  --output-dir=generated_tpu_post_train_artifacts
+
+# Copy generated requirements to src/dependencies/requirements/generated_requirements
+mv generated_tpu_post_train_artifacts/tpu-post-train-requirements.txt \
+  src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt
+```
+
 ### GPU Pre-Training
 
 If you have made changes to the GPU pre-training dependencies in `src/dependencies/requirements/base_requirements/cuda12-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step:

src/dependencies/extra_deps/post_train_github_deps.txt

@@ -1,5 +1,3 @@
 -r post_train_base_deps.txt
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
-mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
 tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/40876e81f04226f9b7b1e4bbdc9051d6b1364b9d.zip
 vllm @ git+https://github.com/vllm-project/vllm@595562651a5a4539ffa910d8570c08fb5169bdc9

src/dependencies/extra_deps/post_train_overrides.txt

@@ -0,0 +1,2 @@
+google-metrax>=0.2.3
+libtpu>=0.0.39

src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt

@@ -0,0 +1,36 @@
+-r requirements.txt
+gepa
+google-cloud-storage
+google-tunix
+hypothesis
+ijson
+mistral_common
+prometheus-fastapi-instrumentator
+pytest-mock
+runai-model-streamer[s3,gcs]
+sortedcontainers
+torchax
+gguf
+pyzmq
+pybase64
+cachetools
+openai
+openai-harmony
+torchvision
+py-cpuinfo
+llguidance
+tpu-info
+xgrammar
+xprof
+yapf
+anthropic
+openai
+ray
+gspread
+lxml
+loguru
+astor
+python-json-logger
+perfetto
+cbor2
+compressed-tensors