Update the network setting to match internal dev settings.

mend

MaxText/configs/README.md

@@ -19,7 +19,7 @@ This directory contains high performance model configurations for different gene
 
 These configurations do 3 things:
 * Sets various XLA compiler flags as `LIBTPU_INIT_ARGS` to optimize runtime performance.
-* Runs [rto_setup.sh](https://github.com/google/maxtext/blob/main/rto_setup.sh) or [gke_rto_setup.sh](https://github.com/google/maxtext/blob/main/gke_rto_setup.sh) to optimize communication protocols for network performance. 
+* Runs [rto_setup.sh](https://github.com/google/maxtext/blob/main/rto_setup.sh) to optimize communication protocols for network performance. 
 (This only needs to be run once on each worker)
 * Runs [train.py](https://github.com/google/maxtext/blob/main/MaxText/train.py) with specific hyper-parameters (batch size, etc.)
 

MaxText/configs/experimental/1024b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

MaxText/configs/experimental/128b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

MaxText/configs/experimental/256b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

MaxText/configs/experimental/32b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

MaxText/configs/experimental/512b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

MaxText/configs/experimental/64b.sh

@@ -14,19 +14,8 @@ for ARGUMENT in "$@"; do
  export "$KEY"="$VALUE"
 done
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- bash rto_setup.sh
- # For DNS lookup when running on large number of VMs on gce
- echo '142.250.123.95 www.googleapis.com' | sudo tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | sudo tee -a /etc/hosts
-else
- # Set up network for running on gke
- bash gke_rto_setup.sh
- # For DNS lookup when running on large number of VMs on gke
- echo '142.250.123.95 www.googleapis.com' | tee -a /etc/hosts
- echo '142.251.4.128 storage.googleapis.com' | tee -a /etc/hosts
-fi
+# Use preflight.sh to set up env based on platform
+bash preflight.sh PLATFORM=$PLATFORM
 
 # Train
 export LIBTPU_INIT_ARGS="--xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true"

preflight.sh

@@ -20,12 +20,13 @@ if [ -z $PLATFORM ]; then
  exit 1
 fi
 
-if [[ $PLATFORM == "gce" ]]; then
- # Set up network for running on gce
- echo "Setting up network for GCE"
- bash rto_setup.sh
+# Check if sudo is available
+if command -v sudo >/dev/null 2>&1; then
+ # sudo is available, use it
+ echo "running rto_setup.sh with sudo"
+ sudo bash rto_setup.sh
 else
- # Set up network for running on gke
- echo "Setting up network for GKE"
- bash gke_rto_setup.sh
+ # sudo is not available, run the script without sudo
+ echo "running rto_setup.sh without sudo"
+ bash rto_setup.sh
 fi

rto_setup.sh

@@ -1,15 +1,61 @@
-echo "Running rto_setup.sh for GCE..."
+echo "Running rto_setup.sh"
 
 # Stop execution if any command exits with error
 set -e
 
-echo "Adjust RTO and apply non cache copy"
-first_line_res=$(ip route show | head -n 1)
-if [[ "$(echo "$first_line_res" | grep "rto_min lock 5ms" | wc -l)" -eq 0 ]]; then
- sudo ip route change "${first_line_res}" rto_min 5ms
+echo "Adjust Network settings and apply non cache copy"
+
+# Install ip.
+apt-get update
+yes | apt-get install net-tools
+yes | apt-get install iproute2
+yes | apt-get install procps
+yes | apt-get install ethtool
+
+# Disable slow start after idle
+sysctl net.ipv4.tcp_slow_start_after_idle=0
+
+# Disable metrics cache
+sysctl net.ipv4.tcp_no_metrics_save=1
+
+# Address rto_min issue with two default routing entries: screen/7RGgkiXkGXSeYF2
+route=$(ip route show | sed -n 1p)
+second_route=$(ip route show | sed -n 2p)
+if [[ "${second_route}" =~ ^default.* ]]; then
+ modified_route=${route//" lock"/}
+ ip route delete ${modified_route}
+fi
+route=$(ip route show | sed -n 1p)
+echo "route rto before change: $route"
+if [[ "${route}" =~ .*lock.*5ms.* ]]; then
+ echo "${route}"
+else
+ # shellcheck disable=SC2086
+ ip route change $route rto_min 5ms
 fi
+route=$(ip route show | sed -n 1p)
+echo "route rto after change: $route"
+
+# Disable Cubic Hystart Ack-Train
+echo 2 > /sys/module/tcp_cubic/parameters/hystart_detect
+
+# Improve handling SYN burst
+echo 4096 > /proc/sys/net/core/somaxconn
+echo 4096 > /proc/sys/net/ipv4/tcp_max_syn_backlog
+
+# Disable MTU Discovery
+echo 0 > /proc/sys/net/ipv4/tcp_mtu_probing
+
+# Increase TCP Zerocopy control memory
+sysctl -w net.core.optmem_max=131072
+
+# Printing output of `ip route show`
+echo -e "\nPrinting output of 'ip route show':"
+ip route show
+
+first_line_res=$(ip route show | head -n 1)
 dev_name=$(echo "$first_line_res" | awk -F'[[:space:]]' '{ print $5 }')
 echo "dev_name=${dev_name}"
-sudo ethtool -K "${dev_name}" tx-nocache-copy on
+ethtool -K "${dev_name}" tx-nocache-copy on
 
 echo "rto_setup.sh finished"