Initial and final evaluation in `finetune` scripts do not accumulate over devices

[mseeger]

Bug description

I am looking at litgpt/finetune/lora.py and litgpt/finetune/full.py. In the LoRA code, the periodic evaluation L401-418 runs validate on each device, then accumulates parts by all_reduce.

But this does not happen for initial evaluation L310 and final evaluation L260. This seems pretty wrong to me, the loss values would just be the one on the rank 0 device.

Another issue is the val_loss value which is printed in L395, but which seems never updated in L401-418.

I'd be happy to submit a PR fixing all this, but first wanted to check whether I understand something wrong here?

Reproduced in studio

No response

What operating system are you using?

macOS

LitGPT Version

0.5.9

[mseeger]

I actually found further issues with these finetune scripts:

• Optimizer params cannot be set as args, not even AdamW.lr. This is maybe the most important parameter to adjust
• LR scheduler sets LR for step 0 to 0. Fine-tuning often has only a small number of steps, especially with parallel devices. No gradient should go wasted
• TrainArgs.lr_warmup_fraction can be set, but is not used. Only TrainArgs.lr_warmup_steps is used. The default is 100, which is used even if the max number of steps is smaller than that

[mseeger]

I could submit a PR to change some of this, but I am moving away from using fabric, because it does not work for my non-standard use case.

You guys are super busy I know. But it would be awesome if the fine-tuning scripts are really good. I suppose most folks want to do that, not pretraining or other things. Ideally, the scripts would not use fabric, because it is not fully general. I found fabric to fail if the gradient computation is not just done by loss.backward() on some model.