Request to implement SCALE optimizer

[jim-plus]

Although the proposed SCALE optimizer was only validated for pretraining, the claimed memory savings make it appealing for a fine-tuning optimizer as well.

https://arxiv.org/abs/2506.16659

[krammnic]

Hey, thank you for the issue, would you like to try to implement it by yourself?

[jim-plus]

Optimizers aren't my forte, but I could give it a look and at least share what I've found.

This appears to be a reference implementation of SCALE:
https://github.com/OptimAI-Lab/Minimalist_LLM_Pretraining

This appears to be the optimizer code itself:
https://github.com/OptimAI-Lab/Minimalist_LLM_Pretraining/blob/main/mem_eff_pt/pt_scale/scale_optimizer.py

They provide example hyperparameters:

torchrun --standalone --nproc_per_node 4 torchrun_main_DDP.py
--model_name 350m_scale
--model_config configs/llama_350m.json
--optimizer scale
--lr 1e-3
--momentum 0.9
--weight_decay 0.0
--batch_size 128
--total_batch_size 512
--num_training_steps 60000
--warmup_steps 6000
--dtype bfloat16
--eval_every 1000
--save_every 99999
--seed 42
--scheduler cosine
--dataset_path /path/to/c4/en \

[jim-plus]

The optimizer seems mostly straightforward to port except for the problem of how to identify if the layer being operated on is the last layer.