#!/bin/bash
s_token=2200000000
m_token=6700000000
l_token=14580000000
xl_token=25500000000
s_lr=0.0006
m_lr=0.0003
l_lr=0.00025
xl_lr=0.0002
train_batch=8
test_batch=8
s_lr_ratio=0.30
m_lr_ratio=0.38
l_lr_ratio=0.41
xl_lr_ratio=0.43
# The parameters for BlockDense with about 32\% parameters are not matched with the LowRank and BlockShuffle exactly. Thus, we provide the max_step_ratio of self-guided training for BlockDense separately to exactly match the training FLOPs.

s_bld_ratio=0.30
m_bld_ratio=0.40
l_bld_ratio=0.41
xl_bld_ratio=0.45

# apply self-guided training for the first half of training
./run_gpt.sh "gpt2s-lr-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=lowrank method.kwargs.rank=192 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"
./run_gpt.sh "gpt2m-lr-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=lowrank method.kwargs.rank=256 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"
./run_gpt.sh "gpt2s-bld-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=blockdense method.kwargs.nblocks=2 method.kwargs.rank=256 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"
./run_gpt.sh "gpt2m-bld-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=blockdense method.kwargs.nblocks=4 method.kwargs.rank=384 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"
./run_gpt.sh "gpt2s-bls-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"
./run_gpt.sh "gpt2m-bls-0x33-sgd-fixedstep" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedstep optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=0.5"

# to match the flops
# LowRank
./run_gpt.sh "gpt2s-lr-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=lowrank method.kwargs.rank=192 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${s_lr_ratio}"
./run_gpt.sh "gpt2m-lr-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=lowrank method.kwargs.rank=256 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${m_lr_ratio}"
./run_gpt.sh "gpt2l-lr-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2l method=lowrank method.kwargs.rank=384  optimization.optimizer.kwargs.lr=${l_lr} optimization.max_tokens=${l_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${l_lr_ratio}"
./run_gpt.sh "gpt2xl-lr-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2xl method=lowrank method.kwargs.rank=512 optimization.optimizer.kwargs.lr=${xl_lr} optimization.max_tokens=${xl_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${xl_lr_ratio}"

# BlockDense
./run_gpt.sh "gpt2s-bld-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=blockdense method.kwargs.nblocks=2 method.kwargs.rank=256 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${s_bld_ratio}"
./run_gpt.sh "gpt2m-bld-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=blockdense method.kwargs.nblocks=4 method.kwargs.rank=384 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${m_bld_ratio}"
./run_gpt.sh "gpt2l-bld-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2l method=blockdense method.kwargs.nblocks=2 method.kwargs.rank=512  optimization.optimizer.kwargs.lr=${l_lr} optimization.max_tokens=${l_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${l_bld_ratio}"
./run_gpt.sh "gpt2xl-bld-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2xl method=blockdense method.kwargs.nblocks=4 method.kwargs.rank=768 optimization.optimizer.kwargs.lr=${xl_lr} optimization.max_tokens=${xl_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${xl_bld_ratio}"

# BlockShuffle
./run_gpt.sh "gpt2s-bls-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2 method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${s_lr} optimization.max_tokens=${s_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${s_lr_ratio}"
./run_gpt.sh "gpt2m-bls-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2m method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${m_lr} optimization.max_tokens=${m_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${m_lr_ratio}"
./run_gpt.sh "gpt2l-bls-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2l method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${l_lr} optimization.max_tokens=${l_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${l_lr_ratio}"
./run_gpt.sh "gpt2xl-bls-0x33-sgd" 1 "torchrun --nnodes=1 --nproc_per_node=1 refinedweb_experiment.py model=gpt2xl method=blockshuffle method.kwargs.nblocks=4 optimization.optimizer.kwargs.lr=${xl_lr} optimization.max_tokens=${xl_token} data.train.train_batch=${train_batch} data.test.test_batch=${test_batch} optimization.log_interval=100 optimization/training=self_guided_training optimization.training.kwargs.mode=fixedflop optimization.training.kwargs.reduce_flop=true optimization.training.kwargs.max_step_ratio=${xl_lr_ratio}"


