Merge branch 'main' into ajkv/offline-distillation-soft

Author: ajkv-google

.github/CODEOWNERS

@@ -1,36 +1,34 @@
 * @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git
 
 # Model bring-up
-src/MaxText/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande
-src/MaxText/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google @NuojCheng
+src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande
+src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
 src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande
-src/MaxText/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google @NuojCheng
+src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
+src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
 
 # Features
 src/maxtext/experimental/rl @A9isha @khatwanimohit @xuefgu @gagika @richjames0 @shralex @NicoGrande
-src/MaxText/input_pipeline @aireenmei @SurbhiJainUSC @richjames0 @shralex @NicoGrande
-src/MaxText/kernels/megablox @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
-src/MaxText/kernels/ragged_attention.py @patemotter @vipannalla @richjames0 @shralex
-src/MaxText/layers/pipeline.py @gobbleturk @richjames0 @shralex @NuojCheng
-src/MaxText/layers/moe.py @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
-src/MaxText/layers/multi_token_prediction.py @parambole @RissyRan @gagika @richjames0 @shralex
-src/MaxText/elastic_train.py @lukebaumann @shauryagup @richjames0 @shralex
-src/MaxText/layers/quantizations.py @khatwanimohit @jshin1394 @liudangyi @richjames0 @shralex
+src/maxtext/input_pipeline @aireenmei @SurbhiJainUSC @richjames0 @shralex @NicoGrande
+src/maxtext/kernels/megablox @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
+src/maxtext/kernels/ragged_attention.py @patemotter @vipannalla @richjames0 @shralex
+src/maxtext/layers/pipeline.py @gobbleturk @richjames0 @shralex @NuojCheng
+src/maxtext/layers/moe.py @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
+src/maxtext/layers/multi_token_prediction.py @parambole @RissyRan @gagika @richjames0 @shralex
+src/maxtext/layers/quantizations.py @khatwanimohit @jshin1394 @liudangyi @richjames0 @shralex
 
 # Inference
-src/maxtext/tests/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
+tests/inference/ @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
 src/maxtext/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
-src/maxtext/inference_mlperf @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @richjames0
 
 # Dockerfiles and dependencies
-*.Dockerfile @bvandermoon @parambole @richjames0 @shralex
-*.txt @bvandermoon @parambole @richjames0 @shralex
+src/dependencies/ @bvandermoon @parambole @richjames0 @shralex
 
 # Docs
-*.md @jacoguzo @bvandermoon @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla
+docs/ @jacoguzo @bvandermoon @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla
 
 # Workflow files
-.github/workflows @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0
+.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0
 
 # Benchmarking/Recipes
-benchmarks @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025
+benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025

.github/workflows/UploadDockerImages.yml

@@ -65,11 +65,11 @@ jobs:
           - device: tpu
             build_mode: stable
             image_name: maxtext_jax_stable
-            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
           - device: tpu
             build_mode: nightly
             image_name: maxtext_jax_nightly
-            dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}
@@ -79,32 +79,18 @@ jobs:
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
 
-  tpu-post-training-stable:
-    name: tpu-post-training-stable
-    needs: setup
-    uses: ./.github/workflows/build_and_push_docker_image.yml
-    with:
-      image_name: maxtext_post_training_stable
-      device: tpu
-      build_mode: stable
-      workflow: post-training
-      dockerfile: ./dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
-      maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
-      image_date: ${{ needs.setup.outputs.image_date }}
-
   tpu-post-training-nightly:
     name: tpu-post-training-nightly
-    needs: [setup, tpu-post-training-stable]
+    needs: [setup]
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: maxtext_post_training_nightly
       device: tpu
       build_mode: nightly
       workflow: post-training
-      dockerfile: ./dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile
+      dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
-      base_image: gcr.io/tpu-prod-env-multipod/maxtext_post_training_stable:${{ needs.setup.outputs.image_date }}
 
   gpu-pre-training:
     name: ${{ matrix.image_name }}
@@ -116,11 +102,11 @@ jobs:
           - device: gpu
             build_mode: stable
             image_name: maxtext_gpu_jax_stable
-            dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
           - device: gpu
             build_mode: nightly
             image_name: maxtext_gpu_jax_nightly
-            dockerfile: ./dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
     uses: ./.github/workflows/build_and_push_docker_image.yml
     with:
       image_name: ${{ matrix.image_name }}

.github/workflows/build_and_push_docker_image.yml

@@ -35,16 +35,16 @@ on:
         required: true
         type: string
       image_date:
-        required: true
-        type: string
-      base_image:
         required: false
         type: string
-        default: ''
       workflow:
         required: false
         type: string
         default: 'pre-training'
+      version_name:
+        required: false
+        type: string
+        default: ''
 
 permissions:
   contents: read
@@ -115,7 +115,7 @@ jobs:
           push: true
           context: .
           file: ${{ inputs.dockerfile }}
-          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:latest
+          tags: gcr.io/tpu-prod-env-multipod/${{ inputs.image_name }}:${{ github.run_id }}
           cache-from: type=gha
           outputs: type=image,compression=zstd,force-compression=true
           build-args: |
@@ -125,34 +125,40 @@ jobs:
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
             INCLUDE_TEST_ASSETS=true
-            ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'
         shell: bash
         run: |
           SOURCE_IMAGE="gcr.io/tpu-prod-env-multipod/${INPUTS_IMAGE_NAME}"
 
-          # Add date tag
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
+          if [[ $INPUTS_VERSION_NAME ]]; then
+            echo "Tagging docker images corresponding to PyPI release..."
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_VERSION_NAME}" --quiet
+          else
+            echo "Tagging docker images corresponding to nightly release..."
 
-          # Convert date to YYYYMMDD format
-          clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
+            # Add date tag
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${INPUTS_IMAGE_DATE}" --quiet
 
-          # Add MaxText tag
-          maxtext_hash=$(git rev-parse --short HEAD)
-          gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
+            # Convert date to YYYYMMDD format
+            clean_date=$(echo "${INPUTS_IMAGE_DATE}" | sed 's/[-:]//g' | cut -c1-8)
 
+            # Add MaxText tag
+            maxtext_hash=$(git rev-parse --short HEAD)
+            gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:maxtext_${maxtext_hash}_${clean_date}" --quiet
 
           # Add post-training dependencies tags
           if [ "${{ inputs.workflow }}" == "post-training" ]; then
             for dir in tunix vllm tpu-inference; do
               if [ -d "./$dir" ]; then
                 dir_hash=$(git -C "$dir" rev-parse --short HEAD)
-                gcloud container images add-tag "$SOURCE_IMAGE:latest" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
-              fi
-            done
+                gcloud container images add-tag "$SOURCE_IMAGE:${{ github.run_id }}" "$SOURCE_IMAGE:${dir}_${dir_hash}_${clean_date}" --quiet
+                fi
+              done
+            fi
           fi
         env:
           INPUTS_IMAGE_NAME: ${{ inputs.image_name }}
           INPUTS_IMAGE_DATE: ${{ inputs.image_date }}
+          INPUTS_VERSION_NAME: ${{ inputs.version_name }}

.github/workflows/build_and_upload_images.sh

@@ -49,7 +49,7 @@ if [[ ! -v CLOUD_IMAGE_NAME ]] || [[ ! -v PROJECT ]] || [[ ! -v MODE ]] || [[ !
 fi
 
 gcloud auth configure-docker us-docker.pkg.dev --quiet
-bash "$MAXTEXT_REPO_ROOT"'/dependencies/scripts/docker_build_dependency_image.sh' LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE="$MODE" DEVICE="$DEVICE"
+bash "$MAXTEXT_REPO_ROOT"'/src/dependencies/scripts/docker_build_dependency_image.sh' LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE="$MODE" DEVICE="$DEVICE"
 image_date=$(date +%Y-%m-%d)
 
 # Upload only dependencies image
@@ -65,7 +65,7 @@ if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-$
 fi
 
 # Build then upload "dependencies + code" image
-docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f "$MAXTEXT_REPO_ROOT"'/dependencies/dockerfiles/maxtext_runner.Dockerfile' -t ${LOCAL_IMAGE_NAME}_runner .
+docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f "$MAXTEXT_REPO_ROOT"'/src/dependencies/dockerfiles/maxtext_runner.Dockerfile' -t ${LOCAL_IMAGE_NAME}_runner .
 docker tag ${LOCAL_IMAGE_NAME}_runner gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest
 docker push gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest
 docker tag ${LOCAL_IMAGE_NAME}_runner gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:${image_date}

.github/workflows/check_docs_build.yml

@@ -27,7 +27,7 @@ jobs:
         run: uv venv --python 3.12 $GITHUB_WORKSPACE/venv
 
       - name: Install dependencies
-        run: . $GITHUB_WORKSPACE/venv/bin/activate && uv pip install -r dependencies/requirements/requirements_docs.txt
+        run: . $GITHUB_WORKSPACE/venv/bin/activate && uv pip install -r src/dependencies/requirements/requirements_docs.txt
 
       - name: Build documentation
         run: |

.github/workflows/pypi_release.yml

@@ -41,12 +41,11 @@ jobs:
     name: Build and Test MaxText Package
     needs: [release_approval]
     uses: ./.github/workflows/build_and_test_maxtext.yml
+    secrets: inherit
 
   publish_maxtext_package_to_pypi:
     name: Publish MaxText package to PyPI
-    # Temporarily only require release_approval for a one-time upload.
-    # Immediately revert this to `needs: [build_and_test_maxtext_package]`.
-    needs: [release_approval]
+    needs: [build_and_test_maxtext_package]
     runs-on: ubuntu-latest
     environment: release
     steps:
@@ -61,3 +60,66 @@ jobs:
       uses: pypa/gh-action-pypi-publish@release/v1
       with:
         packages-dir: dist/
+
+  get_latest_maxtext_pypi_version:
+    name: Get latest MaxText PyPI version
+    needs: [publish_maxtext_package_to_pypi]
+    runs-on: ubuntu-latest
+    outputs:
+      latest_pypi_version: ${{ steps.get_version.outputs.version }}
+    steps:
+      - name: Install jq
+        run: sudo apt-get update && sudo apt-get install -y jq
+      - name: Fetch latest version of maxtext from PyPI
+        id: get_version
+        run: |
+          # Fetch JSON from PyPI for 'maxtext'
+          echo "Fetching latest version from https://pypi.org/pypi/maxtext/json"
+          pypi_json=$(curl -s https://pypi.org/pypi/maxtext/json)
+
+          # Extract the version from the "info" section using jq
+          latest_version=$(echo "$pypi_json" | jq -r ".info.version")
+
+          if [ -z "$latest_version" ] || [ "$latest_version" == "null" ]; then
+            echo "Error: Could not parse latest version from PyPI JSON."
+            exit 1
+          fi
+
+          echo "Successfully fetched latest MaxText version on PyPI: $latest_version"
+          # Set the output variable for other jobs to consume
+          echo "version=$latest_version" >> "$GITHUB_OUTPUT"
+
+  # This job builds and pushes MaxText stable Docker images for both TPU and GPU devices.
+  # It runs only after a new release is published to PyPI.
+  # Creates docker image for MaxText commit corresponding to the release.
+  upload_maxtext_docker_images:
+    name: ${{ matrix.image_name }}
+    needs: [get_latest_maxtext_pypi_version]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_jax_stable
+            workflow: pre-training
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+          - device: gpu
+            build_mode: stable
+            image_name: maxtext_gpu_jax_stable
+            workflow: pre-training
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+          - device: tpu
+            build_mode: stable
+            image_name: maxtext_post_training_stable
+            workflow: post-training
+            dockerfile: ./src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: ${{ matrix.image_name }}
+      device: ${{ matrix.device }}
+      build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
+      dockerfile: ${{ matrix.dockerfile }}
+      maxtext_sha: ${{ github.sha }}
+      version_name: ${{ needs.get_latest_maxtext_pypi_version.outputs.latest_pypi_version }}

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
       - id: codespell
         args:
           - '-w'
-          - '--skip="*.txt,pylintrc,.*,src/maxtext/assets/*"'
+          - '--skip="*.txt,pylintrc,.*,src/maxtext/assets/*,src/maxtext/input_pipeline/protos/*"'
           - '-L ND,nd,sems,TE,ROUGE,rouge,astroid,ags,dout'
           - '.'
         additional_dependencies:
@@ -30,6 +30,7 @@ repos:
         args:
           - '--disable=R0401,R0917,W0201,W0613'
           - "--ignore-patterns='.pytype,.*pyi$'"
+          - '--ignore-paths=src/maxtext/input_pipeline/protos'
           - 'benchmarks'
           - 'src'
           - 'tests'
@@ -47,6 +48,7 @@ repos:
     rev: 24.10.1
     hooks:
       - id: pyink
+        exclude: src/maxtext/input_pipeline/protos/
         args:
           - '--pyink-indentation=2'
           - '--line-length=122'

.readthedocs.yml

@@ -21,4 +21,4 @@ sphinx:
 # See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
 python:
   install:
-    - requirements: dependencies/requirements/requirements_docs.txt
+    - requirements: src/dependencies/requirements/requirements_docs.txt

PREFLIGHT.md

@@ -7,12 +7,12 @@ Before you run ML workload on Multihost with GCE or GKE, simply apply `bash pref
 
 Here is an example for GCE:
 ```
-bash preflight.sh PLATFORM=GCE && python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name=${YOUR_JOB_NAME?}
+bash preflight.sh PLATFORM=GCE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 Here is an example for GKE:
 ```
-bash preflight.sh PLATFORM=GKE && python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name=${YOUR_JOB_NAME?}
+bash preflight.sh PLATFORM=GKE && python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 # Optimization 2: Numa binding (You can only apply this to v4 and v5p)
@@ -22,14 +22,14 @@ For GCE,
 [preflight.sh](https://github.com/google/maxtext/blob/main/preflight.sh) will help you install `numactl` dependency, so you can use it directly, here is an example:
 
 ```
-bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name=${YOUR_JOB_NAME?}
+bash preflight.sh PLATFORM=GCE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 For GKE,
-`numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
+`numactl` should be built into your docker image from [maxtext_tpu_dependencies.Dockerfile](https://github.com/google/maxtext/blob/main/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile), so you can use it directly if you built the maxtext docker image. Here is an example
 
 ```
-bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml run_name=${YOUR_JOB_NAME?}
+bash preflight.sh PLATFORM=GKE && numactl --membind 0 --cpunodebind=0 python3 -m maxtext.trainers.pre_train.train run_name=${YOUR_JOB_NAME?}
 ```
 
 1. `numactl`: This is the command-line tool used for controlling NUMA policy for processes or shared memory. It's particularly useful on multi-socket systems where memory locality can impact performance.

benchmarks/maxtext_xpk_runner.py

@@ -746,6 +746,7 @@ def xpk_benchmark_runner(
     command, name = generate_xpk_workload_cmd(
         cluster_config=cluster_config,
         wl_config=wl_config,
+        workload_name=wl_config.run_name,
         user=user,
         exp_name=exp_name,
     )