JAX-vLLM Offloading k8s (GKE)

.github/actions/gke-xpk/action.yml

@@ -28,7 +28,7 @@ inputs:
     required: false
     type: string
   MAIN_CONTAINER_NAME: 
-    description: 'Name of the main contianer in an XPK JobSet (fixed)'
+    description: 'Name of the main contianer in an XPK JobSet (fixed in xpk)'
     default: gpu-image
     required: false
     type: string
@@ -57,6 +57,11 @@ inputs:
     required: false
     default: 'nvidia-smi; free -h;'
     type: string
+  ENVS:
+    description: 'Environment variables to pass to xpk for setting in JobSet (delimited by ;)'
+    required: false
+    default: ''
+    type: string
   EXIT_COMMAND:
     description: 'Command to set exit code'
     required: false
@@ -124,6 +129,7 @@ runs:
   - name: Set workload commands
     shell: bash -x -e -u {0}
     run: |
+      # install dependencies to enable export artifacts from container to gcs bucket
       PRELUDE="
           apt install -y ripgrep > /dev/null;
           curl -LO https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-linux-x86_64.tar.gz;
@@ -133,6 +139,7 @@ runs:
 
           mkdir -p /usr/share/workload;
           mkdir -p ${{ inputs.CONTAINER_OUTPUT_PATH }};"
+
       # Work around GCP's deployment model that munges together three
       # mostly unrelated things: (1) the host machine's CUDA driver/libs,
       # (2) the version of NCCL installed on the host machine, and (3)
@@ -151,6 +158,7 @@ runs:
           env;
       "
 
+      # gsutil command to export logs from container's /opt/output to bucket
       POSTLUDE="
           ./google-cloud-sdk/bin/gsutil cp -r ${{ inputs.CONTAINER_OUTPUT_PATH }}/ ${GCS_ARTIFACT_PATH}/node-0\$NODE_RANK;
           ${{ inputs.EXIT_COMMAND }}
@@ -163,9 +171,7 @@ runs:
       POSTLUDE=$(echo ${POSTLUDE} | sed 's/\n/\ /g')
       CMD=$(echo ${CMD} | sed 's/\n/\ /g')
 
-      echo "PRELUDE=${PRELUDE}" >> ${GITHUB_ENV}
-      echo "CMD=${CMD}" >> ${GITHUB_ENV}
-      echo "POSTLUDE=${POSTLUDE}" >> ${GITHUB_ENV}
+      echo "CMD=${PRELUDE} ${CMD} ${POSTLUDE}" >> ${GITHUB_ENV}
 
   - name: Create workload on cluster with XPK
     shell: bash -x -e -u {0}
@@ -194,16 +200,21 @@ runs:
       }
 
       if version_greater "${{ inputs.XPK_VERSION }}" "v0.10.0"; then
-          args+=(
-              --docker-image-pull-secret=${{ inputs.IMAGE_PULL_SECRET_NAME }}
-              --env="JAX_COORDINATOR_PORT=3389"
-              --env="JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):3389"
-          )
-      fi
+        args+=(
+            --docker-image-pull-secret=${{ inputs.IMAGE_PULL_SECRET_NAME }}
+        )
 
+        envs_flat=$(echo "${{ inputs.ENVS }}" | tr '\n' ' ')
+        IFS=';' read -ra env_vars <<< "${envs_flat}"
+        for env in "${env_vars[@]}"; do
+            env=$(echo "${env}" | xargs)
+            [[ -n "${env}" ]] && args+=(--env="${env}")
+        done
+      fi
+
       python xpk.py workload create \
-        ${args[@]} \
-        --command="${PRELUDE} ${CMD} ${POSTLUDE}"
+        "${args[@]}" \
+        --command="${CMD}"
 
   - name: Wait for JobSet to unsuspend on cluster
     shell: bash -u {0}

.github/gke-workflow/jax-vllm-offloading/huggingface-secret.yml

@@ -0,0 +1,8 @@
+apiVersion: v1
+kind: Secret
+metadata:
+  name: hf-token-secret
+  namespace: default
+type: Opaque
+stringData:
+  token: {{ HF_TOKEN}}

.github/gke-workflow/xpk/v0.13.0/tcpxo_decorator.patch

@@ -1,5 +1,5 @@
 diff --git a/src/xpk/core/workload_decorators/tcpxo_decorator.py b/src/xpk/core/workload_decorators/tcpxo_decorator.py
-index 3734f87..dc3b24a 100644
+index 3734f87..4a35459 100644
 --- a/src/xpk/core/workload_decorators/tcpxo_decorator.py
 +++ b/src/xpk/core/workload_decorators/tcpxo_decorator.py
 @@ -181,7 +181,9 @@ def update_gpu_containers(job_manifest):
@@ -13,3 +13,10 @@ index 3734f87..dc3b24a 100644
        )
        container['env'].append({
            'name': 'NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY',
+@@ -197,3 +199,6 @@ def update_gpu_containers(job_manifest):
+       container['volumeMounts'].append(
+           {'name': 'dshm', 'mountPath': '/dev/shm'}
+       )
++      container['env'].append(
++          {'name': 'HF_TOKEN', 'valueFrom': {'secretKeyRef': {'name': 'hf-token-secret', 'key': 'token'}}}
++      )

.github/workflows/_test_maxtext_gke_xpk.yaml

@@ -50,8 +50,11 @@ jobs:
         IMAGE: ${{ env.MAXTEXT_IMAGE }}
         IMAGE_PULL_SECRET_NAME: ${{ steps.store-token.outputs.token-name }}
         WORKLOAD_NAME_PREFIX: ${{ env.WORKLOAD_NAME_PREFIX }}
-        COMMAND: |
+        ENVS: |
+          JAX_COORDINATOR_PORT=3389;
+          JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):\$(JAX_COORDINATOR_PORT);
           console=/dev/stdout;
+        COMMAND: |
           nsys-jax --capture-range=cudaProfilerApi
                    --capture-range-end=stop
                    -o /opt/output/profile.zip

.github/workflows/_test_nccl_gke.yaml

@@ -96,16 +96,17 @@ jobs:
         IMAGE: ${{ env.BASE_IMAGE }}
         IMAGE_PULL_SECRET_NAME: ${{ steps.store-token.outputs.token-name }}
         WORKLOAD_NAME_PREFIX: ${{ steps.workload-name.outputs.WORKLOAD_PREFIX }}
+        ENVS: |
+          JAX_COORDINATOR_PORT=3389;
+          JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):\$(JAX_COORDINATOR_PORT);
+          NHOSTS=${{ env.NHOSTS }};
+          NCCL_LIB_DIR=/opt/nvida/nccl/lib;
+          SCRIPT_DIR=/scripts;
+          NCCL_MINBYTES=${{ env.NCCL_MINBYTES }};
+          NCCL_MAXBYTES=${{ env.NCCL_MAXBYTES }};
+          NCCL_STEPFACTOR=${{ env.NCCL_STEPFACTOR }};
+          NCCL_ITERS=${{ env.NCCL_ITERS }};
         COMMAND: |
-          export NHOSTS=${{ env.NHOSTS }};
-          export NCCL_LIB_DIR=/opt/nvida/nccl/lib;
-          export SCRIPT_DIR=/scripts;
-
-          export NCCL_MINBYTES=${{ env.NCCL_MINBYTES }};
-          export NCCL_MAXBYTES=${{ env.NCCL_MAXBYTES }};
-          export NCCL_STEPFACTOR=${{ env.NCCL_STEPFACTOR }};
-          export NCCL_ITERS=${{ env.NCCL_ITERS }};
-
           service ssh restart;
           console=/dev/stdout;
           declare -a hosts=('nccl-test-host-1' 'nccl-test-host-2');

.github/workflows/jax-vllm-offloading-gke-grpo.yml

@@ -0,0 +1,102 @@
+name: JAX-vLLM offloading GRPO (GKE, XPK)
+
+on:
+  workflow_call:
+    inputs:
+      JAX_VLLM_OFFLOADING_IMAGE:
+        type: string
+        description: MaxText image from ghcr.io/nvidia
+        default: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
+        required: false
+
+jobs:
+  jax-vllm-offloading-grpo-gke-xpk:
+    runs-on: gke-a3mega
+    strategy:
+      matrix:
+        model: ["meta-llama/Llama-3.1-8B-Instruct"]
+    env:
+      WORKLOAD_NAME_PREPREFIX: vllm-grpo
+      JAX_VLLM_OFFLOADING_IMAGE: ${{ inputs.JAX_VLLM_OFFLOADING_IMAGE }}
+
+      NUM_NODES: 2
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Login to GitHub Container Registry
+      uses: docker/login-action@v3
+      with:
+        registry: ghcr.io
+        username: ${{ github.repository_owner }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
+    - name: K8s GHCR store and delete token
+      id: store-token
+      uses: ./.github/actions/store-delete-k8s-ghcr
+
+    - name: Format workload name
+      id: workload-name
+      run: |
+        WORKLOAD_NAME_PREFIX="${WORKLOAD_NAME_PREPREFIX}-$(echo ${{ matrix.model }} | sed 's|.*/\(.*\)-[^-]*|\1|')"
+        WORKLOAD_NAME_PREFIX=$(echo ${WORKLOAD_NAME_PREFIX} | tr '.' '-')
+        echo "WORKLOAD_NAME_PREFIX=${WORKLOAD_NAME_PREFIX,,}" >> ${GITHUB_OUTPUT}
+
+    - name: Run XPK workload on cluster
+      uses: ./.github/actions/gke-xpk
+      with:
+        IMAGE: ${{ env.JAX_VLLM_OFFLOADING_IMAGE }}
+        IMAGE_PULL_SECRET_NAME: ${{ steps.store-token.outputs.token-name }}
+        WORKLOAD_NAME_PREFIX: ${{ steps.workload-name.outputs.WORKLOAD_NAME_PREFIX }}
+        ENVS: |
+          CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7;
+          CUDA_DEVICE_ORDER=PCI_BUS_ID;
+          CUDA_DEVICE_MAX_CONNECTIONS=16;
+          VLLM_ENFORCE_EAGER=1;
+          VLLM_GPU_MEMORY_UTILIZATION=0.7;
+          VLLM_TENSOR_PARALLEL_SIZE=8;
+          VLLM_DISTRIBUTED_BACKEND=mp;
+          VLLM_ATTENTION_BACKEND=TRITON_ATTN;
+          VLLM_LOAD_FORMAT=dummy;
+          NCCL_NET_PLUGIN=/opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so;
+          NCCL_TUNER_PLUGIN=none;
+          MODEL_NAME=${{ matrix.model }};
+          NCCL_CUMEM_ENABLE=0;
+          NCCL_BUFFSIZE=16777216;
+          XLA_FLAGS=--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL --xla_gpu_collective_permute_combine_threshold_bytes=8589934592 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_all_reduce_combine_threshold_bytes=8589934592;
+          TRANSFER_MODE=grouped;
+          USE_POLYMORPHIC_MESH=0;
+          JAX_COORDINATOR_PORT=3389;
+          JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):\$(JAX_COORDINATOR_PORT);
+          GATEWAY_PORT=50051;
+          GATEWAY_URL=\$(JOBSET_NAME):\$(GATEWAY_PORT);
+          OUTPUT_DIR=/opt/output;
+          LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/cuda-12.9/compat/lib.real:/usr/local/nvidia/lib64;
+
+        COMMAND: |
+          set -x;
+
+          pip install jax[k8s];
+          python -c 'import jax; jax.distributed.initialize(); print(jax.devices()); print(jax.local_devices()); assert jax.process_count() > 1; assert len(jax.devices()) > len(jax.local_devices());';
+
+          PIDS=();
+          if [ \${NODE_RANK} = 0 ]; then
+            echo Starting gateway;
+            cd /opt/jtbx/jax-inference-offloading;
+            python jax_inference_offloading/controller/gateway.py 2>&1 | tee -a gateway.log &
+            PIDS+=(\$!);
+
+            echo Starting rollout;
+            cd /opt/jtbx/jax-inference-offloading/examples;
+            python rollout.py 2>&1 | tee -a rollout.log &
+            PIDS+=(\$!);
+          else
+            export MODEL_PATH=\$(python download_model.py --hub=hf --model=\${MODEL_NAME} --ignore='*.pth');
+
+            echo Starting GRPO trainer;
+            python trainer_grpo.py 2>&1 | tee -a trainer_grpo.log &
+            PIDS+=(\$!);
+          fi;
+
+          wait \${PIDS[@]};
+          EXIT_CODE=\$PIPESTATUS;

.github/workflows/jax-vllm-offloading-gke-transfer.yml

@@ -0,0 +1,99 @@
+name: JAX-vLLM offloading transfer (GKE, XPK)
+
+on:
+  workflow_call:
+    inputs:
+      JAX_VLLM_OFFLOADING_IMAGE:
+        type: string
+        description: MaxText image from ghcr.io/nvidia
+        default: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
+        required: false
+
+jobs:
+  jax-vllm-offloading-transfer-gke-xpk:
+    runs-on: gke-a3mega
+    strategy:
+      matrix:
+        model: ["meta-llama/Llama-3.1-8B-Instruct", "meta-llama/Llama-3.1-70B-Instruct"]
+    env:
+      WORKLOAD_NAME_PREPREFIX: vllm-transf # due to 40 character workload name limit
+      JAX_VLLM_OFFLOADING_IMAGE: ${{ inputs.JAX_VLLM_OFFLOADING_IMAGE }}
+
+      NUM_NODES: 2
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Login to GitHub Container Registry
+      uses: docker/login-action@v3
+      with:
+        registry: ghcr.io
+        username: ${{ github.repository_owner }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
+    - name: K8s GHCR store and delete token
+      id: store-token
+      uses: ./.github/actions/store-delete-k8s-ghcr
+
+    - name: Format workload name
+      id: workload-name
+      run: |
+        WORKLOAD_NAME_PREFIX="${WORKLOAD_NAME_PREPREFIX}-$(echo ${{ matrix.model }} | sed 's|.*/\(.*\)-[^-]*|\1|')"
+        WORKLOAD_NAME_PREFIX=$(echo ${WORKLOAD_NAME_PREFIX} | tr '.' '-')
+        echo "WORKLOAD_NAME_PREFIX=${WORKLOAD_NAME_PREFIX,,}" >> ${GITHUB_OUTPUT}
+
+    - name: Run XPK workload on cluster
+      uses: ./.github/actions/gke-xpk
+      with:
+        IMAGE: ${{ env.JAX_VLLM_OFFLOADING_IMAGE }}
+        IMAGE_PULL_SECRET_NAME: ${{ steps.store-token.outputs.token-name }}
+        WORKLOAD_NAME_PREFIX: ${{ steps.workload-name.outputs.WORKLOAD_NAME_PREFIX }}
+        ENVS: |
+          CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7;
+          CUDA_DEVICE_ORDER=PCI_BUS_ID;
+          CUDA_DEVICE_MAX_CONNECTIONS=16;
+          VLLM_ENFORCE_EAGER=1;
+          VLLM_GPU_MEMORY_UTILIZATION=0.7;
+          VLLM_TENSOR_PARALLEL_SIZE=8;
+          VLLM_DISTRIBUTED_BACKEND=mp;
+          VLLM_ATTENTION_BACKEND=TRITON_ATTN;
+          VLLM_LOAD_FORMAT=dummy;
+          MODEL_NAME=${{ matrix.model }};
+          NCCL_NET_PLUGIN=/opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so;
+          NCCL_TUNER_PLUGIN=none;
+          NCCL_CUMEM_ENABLE=0;
+          NCCL_BUFFSIZE=16777216;
+          XLA_FLAGS=--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL --xla_gpu_collective_permute_combine_threshold_bytes=8589934592 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_all_reduce_combine_threshold_bytes=8589934592;
+          TRANSFER_MODE=grouped;
+          USE_POLYMORPHIC_MESH=0;
+          JAX_COORDINATOR_PORT=3389;
+          JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):\$(JAX_COORDINATOR_PORT);
+          GATEWAY_PORT=50051;
+          GATEWAY_URL=\$(JOBSET_NAME):\$(GATEWAY_PORT);
+          OUTPUT_DIR=/opt/output;
+          LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/cuda-12.9/compat/lib.real:/usr/local/nvidia/lib64;
+        COMMAND: |
+          set -x;
+
+          pip install jax[k8s];
+          python -c 'import jax; jax.distributed.initialize(); print(jax.devices()); print(jax.local_devices()); assert jax.process_count() > 1; assert len(jax.devices()) > len(jax.local_devices());';
+
+          PIDS=();
+          if [ \${NODE_RANK} = 0 ]; then
+            echo Starting gateway;
+            cd /opt/jtbx/jax-inference-offloading;
+            python jax_inference_offloading/controller/gateway.py 2>&1 | tee -a gateway.log &
+            PIDS+=(\$!);
+
+            echo Starting rollout;
+            cd /opt/jtbx/jax-inference-offloading/examples;
+            python rollout.py 2>&1 | tee -a rollout.log &
+            PIDS+=(\$!);
+          else
+            echo Starting trainer;
+            python trainer.py 2>&1 | tee -a trainer.log &
+            PIDS+=(\$!);
+          fi;
+
+          wait \${PIDS[@]};
+          EXIT_CODE=\$PIPESTATUS;