Motivation

According to our simple test results with Tunix, we find out the performance like output_throughput varies from McJAX and Pathways. The former outperforms the latter a lot.

Solution

1. Do profile,
   1.1. compare Pathways and McJAX profile, and optimize, like tune attention kernel blocks on Pathways
   1.2 ensure the overlap works

Ideal solution: Refactor the control place to reduce the device_get impact in resolve_last_batch_result and no operation overhead.

Reproduction

• branch: linchai_deepscaler, latest commit. Note: modify MODEL_PATH_PREFIX

• hardware version: 2 tpu-v5p-8 hosts

• software version

  pip list | egrep 'jax|flax|libtpu'
  flax                         0.12.3
  jax                          0.9.0.1 # use 0.8.1
  jax-smi                      1.0.4
  jaxlib                       0.9.0.1 # use 0.8.1
  jaxtyping                    0.3.7 # 0.3.6
  libtpu                       0.0.34 # 0.0.30
  sglang-jax                   main-6a61469da565d42cf38eb6105a040853a3afffdf
  

• mode: Pathways

Launch Instruction

step1: build your docker image: ./build_docker.sh

step2: xpk workload create-pathways --cluster=$PW_CLUSTER_NAME --workload=$WORKLOAD_NAME --command="TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE='1' HF_TOKEN=$HF_TOKEN python3 examples/deepscaler/train_deepscaler_nb.py" --num-slices=1 --tpu-type="v5p-16" --base-docker-image docker.io/library/tunix_base_image   --priority=high