Deprecate JAX_RANDOM_WEIGHTS

Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>

Author: kyuyeunk

.buildkite/pipeline_jax.yml

@@ -190,13 +190,12 @@ steps:
        USE_V6E8_QUEUE: "True"
        SKIP_ACCURACY_TESTS: "True"
        VLLM_MLA_DISABLE: "1"
-       JAX_RANDOM_WEIGHTS: "True"
      agents:
        queue: tpu_v6e_8_queue
      commands:
        - |
          if [[ "$$NIGHTLY" == "1" ]]; then
-           .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh -m deepseek-ai/DeepSeek-R1-0528
+           .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh -m deepseek-ai/DeepSeek-R1-0528 --dummy
          else
            echo "Skipping: NIGHTLY environment variable not set"
            exit 0

.buildkite/scripts/run_in_docker.sh

@@ -108,7 +108,6 @@ exec docker run \
   ${QUANTIZATION:+-e QUANTIZATION="$QUANTIZATION"} \
   ${NEW_MODEL_DESIGN:+-e NEW_MODEL_DESIGN="$NEW_MODEL_DESIGN"} \
   ${USE_V6E8_QUEUE:+-e USE_V6E8_QUEUE="$USE_V6E8_QUEUE"} \
-  ${JAX_RANDOM_WEIGHTS:+-e JAX_RANDOM_WEIGHTS="$JAX_RANDOM_WEIGHTS"} \
   ${SKIP_ACCURACY_TESTS:+-e SKIP_ACCURACY_TESTS="$SKIP_ACCURACY_TESTS"} \
   ${VLLM_MLA_DISABLE:+-e VLLM_MLA_DISABLE="$VLLM_MLA_DISABLE"} \
   "${IMAGE_NAME}:${BUILDKITE_COMMIT}" \

tests/e2e/benchmarking/mlperf.sh

@@ -47,6 +47,7 @@ dataset_name=mlperf
 dataset_path=""
 num_prompts=1000
 exit_code=0
+dummy=false
 
 helpFunction()
 {
@@ -57,6 +58,7 @@ helpFunction()
    echo -e "\t-p The path to the processed MLPerf dataset (default: None, which will download the dataset)"
    echo -e "\t-m A space-separated list of HuggingFace model ids to use (default: Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-0.5B-Instruct, meta-llama/Llama-3.1-8B-Instruct and meta-llama/Llama-4-Scout-17B-16E-Instruct)"
    echo -e "\t-n Number of prompts to use for the benchmark (default: 10)"
+   echo -e "\t--dummy Use dummy random weight (default: false)"
    exit 1
 }
 
@@ -87,6 +89,11 @@ while [[ "$#" -gt 0 ]]; do
             shift
             shift
             ;;
+        --dummy)
+            dummy=true
+            shift
+            shift
+            ;;
         -h|--help)
             helpFunction
             ;;
@@ -121,6 +128,10 @@ if [ -z "$dataset_path" ]; then
     fi
 fi
 
+if [ "$dummy" = true ]; then
+    extra_serve_args+=("--load-format=dummy")
+fi
+
 echo "Using the dataset at $dataset_path"
 
 cd "$root_dir"/vllm || exit

tests/models/common/test_model_loader.py

@@ -254,17 +254,13 @@ def test_get_vllm_model(mesh):
     assert callable(compute_logits_fn)
 
 
-@pytest.mark.parametrize("set_in_config", [True, False])
-def test_get_vllm_model_random_weights(mesh, set_in_config):
+def test_get_vllm_model_random_weights(mesh):
     rng = jax.random.PRNGKey(42)
 
     engine_args = EngineArgs(model="Qwen/Qwen3-0.6B")
     vllm_config = engine_args.create_engine_config()
     vllm_config.model_config.dtype = torch.bfloat16
-    if set_in_config:
-        vllm_config.load_config.load_format = "dummy"
-    else:
-        os.environ["JAX_RANDOM_WEIGHTS"] = "True"
+    vllm_config.load_config.load_format = "dummy"
 
     with set_current_vllm_config(vllm_config):
         temp_file = tempfile.mkstemp()[1]

tpu_inference/models/common/model_loader.py

@@ -103,7 +103,7 @@ def create_jit_model(
                                             apply_to_abstract_model=False)
         return model
 
-    if os.getenv("JAX_RANDOM_WEIGHTS", False):
+    if vllm_config.load_config.load_format == "dummy":
         # Create a sharded model with random inited weights.
         # TODO: currently Qwen2ForCausalLM is using legacy model implementation
         # will merge the random init logic when all model are migrated to new model implementation

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -1,6 +1,5 @@
 import copy
 import functools
-import os
 from collections.abc import Sequence
 from contextlib import nullcontext
 from typing import Any, List, Optional, Tuple
@@ -91,12 +90,8 @@ def load_weights(self):
         # may casue errors. Therefore, we disable it during weight loading.
         vllm_config_for_load.parallel_config.enable_expert_parallel = False
 
-        if os.getenv("JAX_RANDOM_WEIGHTS", False):
-            vllm_config_for_load.load_config.load_format = "dummy"
-            use_random_weights = True
-        else:
-            use_random_weights = (
-                vllm_config_for_load.load_config.load_format == "dummy")
+        use_random_weights = (
+            vllm_config_for_load.load_config.load_format == "dummy")
         if use_random_weights:
             logger.info(
                 "Initializing vLLM model with random weights, weight loading skipped."

tpu_inference/platforms/tpu_jax.py

@@ -48,8 +48,7 @@ class TpuPlatform(Platform):
     ]
 
     additional_env_vars: list[str] = [
-        "JAX_RANDOM_WEIGHTS", "PHASED_PROFILING_DIR",
-        "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS",
+        "PHASED_PROFILING_DIR", "TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS",
         "TPU_MULTIHOST_BACKEND", "VLLM_MLA_DISABLE", "TPU_BACKEND_TYPE"
     ]