Merge pull request #3684 from AI-Hypercomputer:seed-post-train

PiperOrigin-RevId: 907099588

Author: Google-ML-Automation

docs/development/update_dependencies.md

@@ -36,12 +36,12 @@ 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.
-3. **Generate the requirement files**: Run the `seed-env` CLI tool to generate
-   new, fully-pinned requirements files based on your changes.
+3. **Generate the requirement files**: Run `src/dependencies/scripts/generate_requirements.sh`,
+   which internally invokes `seed-env` to produce fully-pinned requirements files.
 4. **Verify the new dependencies**: Test the new dependencies to ensure the
    project installs and runs correctly.
 
@@ -66,17 +66,17 @@ 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
 
 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 from [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build) and copy its full commit hash.
 
 ## Step 3: Generate the requirements files
 
-Next, run the `seed-env` CLI to generate the new requirements files. You will
-need to do this separately for the TPU and GPU environments. The generated files
-will be placed in a directory specified by `--output-dir`.
+Next, run `generate_requirements.sh` to generate the new requirements files. This
+script wraps the `seed-env` CLI and handles exporting the lock, and applying any
+overrides. You will need to do this separately for the TPU and GPU environments.
 
 > **Note:** The current `src/dependencies/requirements/generated_requirements/` in the repository were generated using JAX build commit hash: [e0d2967b50abbefd651d563dbcd7afbcb963d08c](https://github.com/jax-ml/jax/commit/e0d2967b50abbefd651d563dbcd7afbcb963d08c).
 
@@ -85,35 +85,46 @@ will be placed in a directory specified by `--output-dir`.
 If you have made changes to TPU pre-training dependencies in `src/dependencies/requirements/base_requirements/tpu-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:
 
 ```bash
-seed-env \
-  --local-requirements=src/dependencies/requirements/base_requirements/tpu-base-requirements.txt \
-  --host-name=MaxText \
-  --seed-commit=<jax-build-commit-hash> \
-  --python-version=3.12 \
-  --requirements-txt=tpu-requirements.txt \
-  --output-dir=generated_tpu_artifacts
+bash src/dependencies/scripts/generate_requirements.sh \
+--base-requirements src/dependencies/requirements/base_requirements/tpu-requirements.txt \
+--generated-requirements tpu-requirements.txt \
+--seed-commit <jax-build-commit-hash>
 
 # Copy generated requirements to src/dependencies/requirements/generated_requirements
-mv generated_tpu_artifacts/tpu-requirements.txt \
+mv generated_artifacts/python3_12/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
+
+bash src/dependencies/scripts/generate_requirements.sh \
+--base-requirements src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt \
+--generated-requirements tpu-post-train-requirements.txt \
+--override-requirements src/dependencies/extra_deps/post_train_overrides.txt \
+--seed-commit <jax-build-commit-hash>
+
+# Copy generated requirements to src/dependencies/requirements/generated_requirements
+mv generated_artifacts/python3_12/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:
 
 ```bash
-seed-env \
-  --local-requirements=src/dependencies/requirements/base_requirements/cuda12-requirements.txt \
-  --host-name=MaxText \
-  --seed-commit=<jax-build-commit-hash> \
-  --python-version=3.12 \
-  --requirements-txt=cuda12-requirements.txt \
-  --hardware=cuda12 \
-  --output-dir=generated_gpu_artifacts
+bash src/dependencies/scripts/generate_requirements.sh \
+--base-requirements src/dependencies/requirements/base_requirements/cuda12-requirements.txt \
+--generated-requirements cuda12-requirements.txt \
+--seed-commit <jax-build-commit-hash> \
+--hardware cuda12
 
 # Copy generated requirements to src/dependencies/requirements/generated_requirements
-mv generated_gpu_artifacts/cuda12-requirements.txt \
+mv generated_artifacts/python3_12/cuda12-requirements.txt \
   src/dependencies/requirements/generated_requirements/cuda12-requirements.txt
 ```
 

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,4 @@
+flax==0.12.4
+google-metrax>=0.2.3
+libtpu>=0.0.39
+optax==0.2.6

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

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

src/dependencies/requirements/generated_requirements/cuda12-requirements.txt

@@ -1,5 +1,5 @@
-# Generated by seed-env. Do not edit manually.
-# If you need to modify dependencies, please do so in the host requirements file and run seed-env again.
+# Generated by generate_requirements.sh using seed-env tool. Do not edit manually.
+#    See https://maxtext.readthedocs.io/en/latest/development/update_dependencies.html for details.
 
 absl-py>=2.4.0
 aiofiles>=25.1.0
@@ -30,10 +30,10 @@ cloudpickle>=3.1.2
 clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
-cryptography>=46.0.7
+cryptography>=47.0.0
 cycler>=0.12.1
 dataclasses-json>=0.6.7
-datasets>=4.8.4
+datasets>=4.8.5
 decorator>=5.2.1
 deprecated>=1.3.1
 dill>=0.4.1
@@ -58,9 +58,9 @@ gast>=0.7.0
 gcsfs>=2026.2.0
 google-api-core>=2.30.3
 google-api-python-client>=2.194.0
+google-auth>=2.49.2
 google-auth-httplib2>=0.3.1
 google-auth-oauthlib>=1.3.1
-google-auth>=2.49.2
 google-cloud-aiplatform>=1.148.1
 google-cloud-appengine-logging>=1.9.0
 google-cloud-audit-log>=0.5.0
@@ -70,25 +70,25 @@ google-cloud-logging>=3.15.0
 google-cloud-mldiagnostics>=1.0.2
 google-cloud-monitoring>=2.30.0
 google-cloud-resource-manager>=1.17.0
-google-cloud-storage-control>=1.11.0
 google-cloud-storage>=3.10.1
+google-cloud-storage-control>=1.11.0
 google-crc32c>=1.8.0
 google-genai>=1.73.1
 google-pasta>=0.2.0
 google-resumable-media>=2.8.2
 googleapis-common-protos>=1.74.0
 grain>=0.2.16
 grpc-google-iam-v1>=0.14.4
-grpcio-status>=1.78.0
 grpcio>=1.78.0
+grpcio-status>=1.78.0
 gviz-api>=1.10.0
 h11>=0.16.0
 h5py>=3.14.0
 hf-xet>=1.4.3 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
 httpcore>=1.0.9
 httplib2>=0.31.2
 httpx>=0.28.1
-huggingface-hub>=1.11.0
+huggingface-hub>=1.12.0
 humanize>=4.15.0
 hypothesis>=6.142.1
 identify>=2.6.19
@@ -98,9 +98,9 @@ importlab>=0.8.1
 importlib-metadata>=9.0.0
 iniconfig>=2.3.0
 isort>=8.0.1
+jax>=0.9.2
 jax-cuda12-pjrt>=0.9.2 ; sys_platform == 'linux'
 jax-cuda12-plugin>=0.9.2 ; sys_platform == 'linux'
-jax>=0.9.2
 jaxlib>=0.9.2
 jaxtyping>=0.3.9
 jinja2>=3.1.6
@@ -110,8 +110,8 @@ kiwisolver>=1.5.0
 latex2sympy2-extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.8.6
-markdown-it-py>=4.0.0
 markdown>=3.10.2
+markdown-it-py>=4.0.0
 markupsafe>=3.0.3
 marshmallow>=3.26.2
 math-verify>=0.9.0
@@ -132,11 +132,11 @@ nest-asyncio>=1.6.0 ; sys_platform == 'win32'
 networkx>=3.6.1
 ninja>=1.13.0
 nodeenv>=1.10.0
-numpy-typing-compat>=20251206.2.0
 numpy>=2.0.2
+numpy-typing-compat>=20251206.2.0
 nvidia-cublas-cu12>=12.9.1.4 ; sys_platform == 'linux'
-nvidia-cuda-cccl-cu12>=12.9.27
 nvidia-cuda-cccl>=13.2.27
+nvidia-cuda-cccl-cu12>=12.9.27
 nvidia-cuda-cupti-cu12>=12.9.79 ; sys_platform == 'linux'
 nvidia-cuda-nvcc-cu12>=12.9.86 ; sys_platform == 'linux'
 nvidia-cuda-nvrtc-cu12>=12.9.86 ; sys_platform == 'linux'
@@ -160,8 +160,8 @@ orbax-export>=0.0.8
 packaging>=26.0
 pandas>=3.0.2
 parameterized>=0.9.0
-pathspec>=1.1.0
-pathwaysutils>=0.1.7
+pathspec>=1.1.1
+pathwaysutils>=0.1.8
 pillow>=12.1.1
 platformdirs>=4.9.6
 pluggy>=1.6.0
@@ -173,12 +173,12 @@ proto-plus>=1.27.2
 protobuf>=6.33.6
 psutil>=7.2.2
 pyarrow>=24.0.0
-pyasn1-modules>=0.4.2
 pyasn1>=0.6.3
+pyasn1-modules>=0.4.2
 pycnite>=2024.7.31
 pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy'
-pydantic-core>=2.46.3
 pydantic>=2.13.3
+pydantic-core>=2.46.3
 pydot>=4.0.1
 pyelftools>=0.32
 pyglove>=0.4.5
@@ -187,41 +187,41 @@ pyink>=25.12.0
 pylint>=4.0.5
 pyparsing>=3.3.2
 pyproject-hooks>=1.2.0
-pytest-xdist>=3.8.0
 pytest>=8.4.2
+pytest-xdist>=3.8.0
 python-dateutil>=2.9.0.post0
 pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.6
 regex>=2026.4.4
-requests-oauthlib>=2.0.0
 requests>=2.32.5
+requests-oauthlib>=2.0.0
 rich>=14.3.3
 safetensors>=0.7.0
-scipy-stubs>=1.17.1.2
 scipy>=1.17.1
+scipy-stubs>=1.17.1.2
 sentencepiece>=0.2.1
 seqio>=0.0.20
 setuptools>=82.0.1
 shellingham>=1.5.4
 simple-parsing>=0.1.8
-simplejson>=4.1.0
+simplejson>=4.1.1
 six>=1.17.0
 sniffio>=1.3.1
 sortedcontainers>=2.4.0
 starlette>=1.0.0
 sympy>=1.14.0
 tabulate>=0.10.0
 tenacity>=9.1.4
+tensorboard>=2.20.0
 tensorboard-data-server>=0.7.2
 tensorboard-plugin-profile>=2.13.0
-tensorboard>=2.20.0
 tensorboardx>=2.6.5
+tensorflow>=2.20.0
 tensorflow-datasets>=4.9.9
 tensorflow-metadata>=1.17.3
 tensorflow-text>=2.20.1
-tensorflow>=2.20.0
 tensorstore>=0.1.82
 termcolor>=3.3.0
 tiktoken>=0.12.0
@@ -231,17 +231,17 @@ toml>=0.10.2
 tomlkit>=0.14.0
 toolz>=1.1.0
 tqdm>=4.67.3
-transformer-engine-cu12>=2.13.0
-transformer-engine-jax>=2.13.0
-transformer-engine>=2.13.0
-transformers>=5.6.1
+transformer-engine>=2.14.0
+transformer-engine-cu12>=2.14.0
+transformer-engine-jax>=2.14.0
+transformers>=5.6.2
 treescope>=0.1.10
 typeguard>=2.13.3
-typer>=0.24.2
+typer>=0.25.0
 typing-extensions>=4.15.0
 typing-inspect>=0.9.0
 typing-inspection>=0.4.2
-tzdata>=2026.1 ; sys_platform == 'emscripten' or sys_platform == 'win32'
+tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32'
 uritemplate>=4.2.0
 urllib3>=2.6.3
 uvicorn>=0.46.0
@@ -252,7 +252,7 @@ websockets>=16.0
 werkzeug>=3.1.8
 wheel>=0.46.3
 wrapt>=2.1.2
-xxhash>=3.6.0
+xxhash>=3.7.0
 yarl>=1.23.0
 zipp>=3.23.0
-zstandard>=0.25.0
+zstandard>=0.25.0