On-device training of small LLM LoRA weights

[martinkorelic]

Bit of a beginner with the onnx runtime framework, but I noticed that it is one of the rare frameworks that allow on-device retraining of models. So here’s my question:

Does a framework like ONNX runtime offer such a thing for LLM models, by converting from existing transformer models and then finetuning them on device? Instead of re-training the whole model, we would only finetune the LoRA weights.

I understand the hardware restrictions of edge devices, but I would still like to know if such a thing is still feasible with this framework.

Also is it possible that on-device learning can be used with the underlying device GPU (like for Android…)?

Open to any answers, tips or discussion about this. Thank you!

[carzh]

On-device training does provide regular finetuning abilities by enabling users to freeze most layers and unfreeze some layers. AFAIK, we haven't yet tested on-device training with finetuning the LoRA weights, but it should be possible with the following:

• Apply LoRA changes and config to PyTorch / transformers model
• Export the model using an exporter.
• Generate the training artifacts. Mark all the lora linear weights as requires_grad, everything else as frozen
• Continue the rest of the training as per usual.

If you give it a try, I'd be curious to hear how it goes.

For your second question, on-device training is currently enabled on CUDA, but not other device GPUs.

[martinkorelic]

Is GPU training limited by the framework or is it possible to implement such a thing by using ONNX runtime execution provider? Or are they only used for inference at this point?

[carzh]

We're limited by op support for a given execution provider. Gradient ops would have to be implemented to train on the device GPU. Otherwise, it'll fall back to the CPU gradient op implementation.