Saving weights and logs of step 1000

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "roberta-base",
+  "_name_or_path": "./",
   "architectures": [
     "RobertaForMaskedLM"
   ],
@@ -19,6 +19,7 @@
   "num_hidden_layers": 12,
   "pad_token_id": 1,
   "position_embedding_type": "absolute",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.16.0.dev0",
   "type_vocab_size": 1,
   "use_cache": true,

events.out.tfevents.1642353687.t1v-n-00e295a4-w-0.437639.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64a2c72eddd38a2e2c1f27d2fbc653a4aadf40419494842b53dd7f7559a30129
+size 146996

events.out.tfevents.1642508113.t1v-n-00e295a4-w-0.632677.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ff28af8474a1a3bed992c3c440935d0776de435ca3b3eab3d8a8ed8cec34a2f
+size 40

events.out.tfevents.1642508176.t1v-n-00e295a4-w-0.645011.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79ef23d9737ae60fddcddd281a6127f5325f3d0c8080d8b3b17c8105bf147c82
+size 40

events.out.tfevents.1642508417.t1v-n-00e295a4-w-0.662728.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2cd702f0bd422826437b0d591702c3d773f8a3c8a4b25c8a2a91f848f42b2ad
+size 40

events.out.tfevents.1642509967.t1v-n-00e295a4-w-0.712066.0.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0694b7f11a53d724a4552a8be128d5008f0c34de819b6925531da247835c6afe
+size 40

flax_model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3ef24399f6318fe5bd62ae854e2d6bfd60c2ce162537aafd46d63d7221cfe3e
+size 498796983

run_mlm_flax.py

@@ -640,17 +640,17 @@ def main():
         # - mask for weight decay is not implemented but we don't use it anyway
         optimizer = distributed_shampoo(
             linear_decay_lr_schedule_fn,
-            block_size=1024,  # recommended default for large LM is 1536
+            block_size=1536,  # 1024 by Boris, recommended default for large LM is 1536
             beta1=training_args.adam_beta1,  # 0.9,
             beta2=training_args.adam_beta2,  # 0.999,
             diagonal_epsilon=training_args.adam_epsilon,  # 1e-10,
             matrix_epsilon=1e-8,
             weight_decay=training_args.weight_decay,  # 0.0,
             start_preconditioning_step=1001,
-            preconditioning_compute_steps=10,
+            preconditioning_compute_steps=1,
             statistics_compute_steps=1,
             best_effort_shape_interpretation=True,
-            graft_type=GraftingType.RMSPROP_NORMALIZED,
+            graft_type=GraftingType.SGD,  # GraftingType.RMSPROP_NORMALIZED,
             nesterov=False,
             exponent_override=0,
             batch_axis_name="batch",
@@ -658,7 +658,7 @@ def main():
             moving_average_for_momentum=True,
             skip_preconditioning_dim_size_gt=4096,
             clip_by_scaled_gradient_norm=None,
-            precision=jax.lax.Precision.HIGHEST,
+            precision=jax.lax.Precision.DEFAULT,  # DEFAULT is  bfloat16
         )
     else:
         optimizer = optax.adamw(

train.128.sh

@@ -1,18 +1,18 @@
 python run_mlm_flax.py \
     --output_dir="./" \
     --model_type="roberta" \
-    --config_name="roberta-base" \
+    --config_name="./" \
     --tokenizer_name="NbAiLab/nb-roberta-base" \
     --dataset_name="NbAiLab/NCC" \
     --max_seq_length="128" \
     --weight_decay="0.0" \
-    --per_device_train_batch_size="232" \
-    --per_device_eval_batch_size="232" \
+    --per_device_train_batch_size="210" \
+    --per_device_eval_batch_size="210" \
     --pad_to_max_length \
     --learning_rate="6e-4" \
     --warmup_steps="10000" \
     --overwrite_output_dir \
-    --num_train_epochs="3" \
+    --num_train_epochs="5" \
     --distributed_shampoo  \
     --adam_beta1="0.9" \
     --adam_beta2="0.99" \