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.

[devangpratap]

Hi, I'd like to take a crack at fixing this. The bug is clear — initial and final validation in both lora.py and full.py skip the fabric.all_reduce that periodic validation correctly uses. I'll open a PR shortly.

[discobot]

The same missing reduction also exists in adapter.py and adapter_v2.py, and the stale val_loss in the progress line is because the periodic eval reduces into a temporary val_loss_tensor that is never assigned back. Fix in #2266: all three validation paths in the four finetune scripts now report the cross-device mean, with CPU regression tests that stub Fabric.all_reduce.