Fixed exhausted bug between head and workers

README.md

@@ -184,6 +184,7 @@ Note: Get address ip and port information from ray head.
 Here is an example to run the server with ray for llama2 7B model:
 
 ```bash
+export DISABLE_XLA2_PJRT_TEST="true"
 python run_server_with_ray.py --tpu_chips=16 --num_hosts=4 --worker_chips=4 -model_name=$model_name          --size=7b --batch_size=96 --max_cache_length=2048 --quantize_weights=$quantize --quantize_type=$quantize_type --quantize_kv_cache=$quantize --checkpoint_path=$output_ckpt_dir   --tokenizer_path=$tokenizer_path --sharding_config="default_shardings/llama.yaml"
 ```
 

run_interactive_multiple_host.py

@@ -57,7 +57,7 @@ def create_engine():
       sharding_config=FLAGS.sharding_config,
       num_hosts=_NUM_HOSTS.value,
       worker_chips=_WORKER_CHIPS.value,
-      tpu_chips=_TPU_CHIPS,
+      tpu_chips=_TPU_CHIPS.value,
   )
 
   print("Initialize engine", time.perf_counter() - start)