Hybrid Estimator

[ZidongLiu]

TODO

• add another flag to by iteration training/eval for zo_rge_main
• Reuse same lr scheduling logic as decomfl in zo_rge_main
• different learning rate for 2 parameter set

Questions

• learning rate should differ?
• how to split
• Is implementation correct.
• can adjust split to test if it's about the same as decomfl and hiso

Summary

• Implement Batch version
• Implement Paramwise version
• Integrate with decomfl_main
• Integrate with zo_rge_main

Result

• Paramwise can train to 80%, speed 1.6s / iteration
  uv run decomfl_main.py --large-model=opt-125m --dataset=sst2 --iterations=1000 --train-batch-size=32 --test-batch-size=200 --eval-iterations=25 --num-clients=3 --num-sample-clients=2 --local-update-steps=1 --num-pert=5 --lr=1e-5 --mu=1e-3 --grad-estimate-method=rge-forward --no-optim --estimator-type=hybrid

• Batch can train to 80%, speed 1s / iteration
  uv run decomfl_main.py --large-model=opt-125m --dataset=sst2 --iterations=1000 --train-batch-size=32 --test-batch-size=200 --eval-iterations=25 --num-clients=3 --num-sample-clients=2 --local-update-steps=1 --num-pert=5 --lr=1e-5 --mu=1e-3 --grad-estimate-method=rge-forward --estimator-type=hybrid

• zo_rge_main, train by iteration

• zo_rge_main, train_by_epoch

• 2 lrs for 2 set of hybrid parameters, uv run python zo_rge_main.py --dataset=sst2 --eval-iterations=25 --large-model=opt-125m --model-dtype=float32 --seed=365 --iterations=2000 --train-batch-size=32 --test-batch-size=32 --num-pert=5 --lr=0.000005 --momentum=0 --grad-estimate-method=rge-central --mu=0.001 --estimator-type=hybrid --no-optim --lr2=1e-4