Testing muP.

[cloneofsimo]

Initial Tests on muP running from 110M ~ 4.4B models.

Our goal is to see if results from TP-V5 and muScaling results are reproducible, and test to see if muP provides stable scaling-law predictions.

export WORLD_SIZE=$(nvidia-smi -L | wc -l)
#!/bin/bash

lrs=(1e-4 2e-4 5e-4 1e-3)
widths=(16 32 64 128 256)

for lr in "${lrs[@]}"; do
    for width in "${widths[@]}"; do
        run_name="lr_${lr}_width_${width}"
        deepspeed --num_gpus $WORLD_SIZE run_trainer.py --lr $lr --width $width --run_name $run_name
    done
done

---

Current Status:

[cloneofsimo]

Looks correct to me!

[xszheng2020]

Could you please share the loss curves? with respect to training steps

[cloneofsimo]

Yeah sure, here plots of some more recent runs