test warmup effect on 50M model 6B

[dianaonutu]

In our Scaling Laws experiments, we set a fixed warmup duration (1000 iterations) across all runs, regardless of global batch size and token budget. Fixed warmup was chosen to be able to cool down at various token budgets and thus re-use previous runs and save compute. However, if global batch size and token budget differ, this means that a fixed warmup duration will represent a higher or lower fraction of training iterations. As we've calculated, the warmup duration will represent maximum 35% and minimum 0.25% of the total training iterations.

Initial results with a long warmup fraction leads to higher loss. Here, we will sweep for the highest global batch size and the warmup fraction to verify if we obtain much lower loss or not.

Model 50M, token budget 6B, global batch size 512, which results in 2861 iterations. Sweep over warmup fraction [1, 5, 10, 20] and learning rate [5.e-4, 1.e-3, 2.e-3, 4.e-3, 8.e-3]

[dianaonutu]

Experiments are run on LUMI using the OpenEuroLLM/ROCm-Megatron-LM commit, using the autoexperiment tool, commit b88378e.
To start experiments, install the tool and run:

1. autoexperiment build ./autoexperiment/scripts/dense/warmup_sweep/config.yaml
2. autoexperiment run ./autoexperiment/scripts/dense/warmup_sweep/config.yaml

[dianaonutu]

I ran additional experiments with a 50M model trained on 6B tokens at a global batch size of 512, varying the learning rate and warmup duration across five values. For all configurations, a 35% warmup achieved the highest validation loss, while shorter warmups (5%) seemed better. This suggests that warmup duration may have an effect on the performance. But the results are not conclusive as we need to compare them with the old experiments, which are being re-run on LUMI due to the loss discrepancy, and we don't know if the same would hold for higher token budgets.

The loss spikes appear the most at the highest LR, which makes sense. High batch size and high learning rate lead to abrupt updates. I don't think they are a concern, as there are hyperparameter configurations that don't have them (so much).

warmup_experiments.pptx