56

0.3.0

@@ -0,0 +1,13 @@
+Defaulting to user installation because normal site-packages is not writeable
+Looking in links: https://storage.googleapis.com/jax-releases/libtpu_releases.html
+Requirement already satisfied: jax[tpu] in /home/perk/.local/lib/python3.8/site-packages (0.4.1)
+Requirement already satisfied: scipy>=1.5 in /home/perk/.local/lib/python3.8/site-packages (from jax[tpu]) (1.10.0)
+Requirement already satisfied: opt-einsum in /home/perk/.local/lib/python3.8/site-packages (from jax[tpu]) (3.3.0)
+Requirement already satisfied: numpy>=1.20 in /home/perk/.local/lib/python3.8/site-packages (from jax[tpu]) (1.24.1)
+Requirement already satisfied: libtpu-nightly==0.1.dev20221212 in /home/perk/.local/lib/python3.8/site-packages (from jax[tpu]) (0.1.dev20221212)
+Requirement already satisfied: jaxlib==0.4.1 in /home/perk/.local/lib/python3.8/site-packages (from jax[tpu]) (0.4.1)
+Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from jax[tpu]) (2.27.1)
+Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->jax[tpu]) (2019.11.28)
+Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/perk/.local/lib/python3.8/site-packages (from requests->jax[tpu]) (1.26.13)
+Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->jax[tpu]) (2.8)
+Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.8/dist-packages (from requests->jax[tpu]) (2.0.12)

run.sh

@@ -5,7 +5,7 @@ python run_mlm_flax_stream.py \
     --dataset_name="NbAiLab/nbailab_extended" \
     --max_seq_length="512" \
     --weight_decay="0.01" \
-    --per_device_train_batch_size="62" \
+    --per_device_train_batch_size="56" \
     --per_device_eval_batch_size="16" \
     --learning_rate="3e-4" \
     --warmup_steps="25000" \