Merge branch 'qanthony/mup' of https://github.com/EleutherAI/nanoGPT-mup

 into qanthony/mup

.gitignore

@@ -10,5 +10,4 @@ __pycache__/
 input.txt
 env/
 venv/
-coord_check/*/out/*
-mutransfer_lr/*/out/*
+mup_examples/*/*/out/*

mup_examples/coord_check_shakespeare_char/mup/run.sh

@@ -6,7 +6,7 @@ do
     n_heads=$((width / head_size))
     mup_base_width=256
     mup_width_multiplier=$(echo "scale=8; $width/$mup_base_width" | bc -l)
-    out_dir="coord_check/mup/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/mup/out/width${width}_depth2_seed${seed}"
     python train.py \
         --out_dir=$out_dir \
         --eval_interval=1 \

mup_examples/coord_check_shakespeare_char/sp/run.sh

@@ -4,7 +4,7 @@ do
     do
     head_size=64
     n_heads=$((width / head_size))
-    out_dir="coord_check/sp/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/sp/out/width${width}_depth2_seed${seed}"
     python train.py \
         --out_dir=$out_dir \
         --eval_interval=1 \

mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init/run.sh

@@ -6,7 +6,7 @@ do
     n_heads=$((width / head_size))
     mup_base_width=256
     mup_width_multiplier=$(echo "scale=8; $width/$mup_base_width" | bc -l)
-    out_dir="coord_check/sp_with_mup_hidden_init/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init/out/width${width}_depth2_seed${seed}"
     python train.py \
         --out_dir=$out_dir \
         --eval_interval=1 \

mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr/run.sh

@@ -6,7 +6,7 @@ do
     n_heads=$((width / head_size))
     mup_base_width=256
     mup_width_multiplier=$(echo "scale=8; $width/$mup_base_width" | bc -l)
-    out_dir="coord_check/sp_with_mup_hidden_init_and_lr/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr/out/width${width}_depth2_seed${seed}"
     python train.py \
         --out_dir=$out_dir \
         --eval_interval=1 \

mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr_output_logits/run.sh

@@ -6,7 +6,7 @@ do
     n_heads=$((width / head_size))
     mup_base_width=256
     mup_width_multiplier=$(echo "scale=8; $width/$mup_base_width" | bc -l)
-    out_dir="coord_check/sp_with_mup_hidden_init_and_lr_output_logits/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr_output_logits/out/width${width}_depth2_seed${seed}"
     python train.py \
         --out_dir=$out_dir \
         --eval_interval=1 \

mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr_partial_output_logits/run.sh

@@ -6,7 +6,7 @@ do
     n_heads=$((width / head_size))
     mup_base_width=256
     mup_width_multiplier=$(echo "scale=8; $width/$mup_base_width" | bc -l)
-    out_dir="coord_check/sp_with_mup_hidden_init_and_lr_partial_output_logits/out/width${width}_depth2_seed${seed}"
+    out_dir="mup_examples/coord_check_shakespeare_char/sp_with_mup_hidden_init_and_lr_partial_output_logits/out/width${width}_depth2_seed${seed}"
     mup_output_alpha=$(echo "scale=8; sqrt($mup_width_multiplier)" | bc -l)
     python train.py \
         --out_dir=$out_dir \

mup_examples/mutransfer_lr_shakespeare_char/mup/run.sh

@@ -9,7 +9,7 @@ LAYERS=2
 
 for width in 256 512 1024 2048
 do
-    for lr in 0.0009765625 0.00048828125 0.000244140625 0.0001220703125 0.00006103515625
+    for lr in 0.125 0.0625 0.03125 0.015625 0.0078125 0.00390625 0.001953125 0.0009765625 0.00048828125 0.000244140625 0.0001220703125 0.00006103515625
     do
         for seed in 1 2 3
         do

mup_examples/mutransfer_lr_shakespeare_char/sp/run.sh

@@ -7,9 +7,9 @@ LAUNCHER=python
 
 LAYERS=2
 
-for width in 1024 2048
+for width in 256 512 1024 2048
 do
-    for lr in 0.00390625 0.001953125 0.0009765625 0.00048828125 0.000244140625 0.0001220703125 0.00006103515625 0.00003051757812
+    for lr in 0.00390625 0.001953125 0.0009765625 0.00048828125 0.000244140625 0.0001220703125 0.00006103515625 0.00003051757812 0.00048828125 0.000244140625 0.0001220703125 0.00006103515625 0.00003051757812 0.00001525878906 0.000007629394531 0.000003814697266
     do
         for seed in 1 2 3
         do