Repository navigation
Add native LoRA and NF4 QLoRA layers - #466
Conversation
|
Hi @surya-mp Thank you for the contribution! I've been meaning to do this for a while and never had the time. This is very welcome! Unfortunately, I'm in a couple of weeks trip without a laptop ... But I'll review/merge it as soon as I'm back, 1st week of October. Apologies. Cheers PS: for a later time: this is a useful use case, it could be interesting to create a page in gomlx.github.io, if it's of interest to you. I'd also love to add/improve fused ops to optimize the execution of Lora models with the Go backend. |
There was a problem hiding this comment.
First, thanks again for the contribution, PEFT is something that has been lacking 😃
Even though I understand the math around LoRA/QLoRA, and after reading the peft.md, I'm still not following how should I use ml/layers/peft :) Probably, it's me being slow ... But in any case, we would like to make it simple to normal users.
A few questions: I assume this works by having the user rewrite the model replacing calls to layers.fnn.New (or ml/nn/Dense) with calls to peft.NewLinear or peft.NewNF4Linear ? I don't follow when one would call peft.Inject -- if within the model building function or outside of it. Also whether it applies to all layers, or only to one layer.
In GoMLX usually all the state of a model, when possible, is store in the model.Store. Rarely there is an object that is used both, outside the model building function and inside the building function. Not saying it can't exist, but it feels to me the LoRA adapter could be completely exist inside the model building function -- and largely configured by hyperparameters stored in model.Store (accessed through the scope). You can always add extra configuration methods, but it would be nice if the defaults are read from the scope. See the package ml/layers/fnn for an example.
A few meta comments:
- During model building in GoMLX we "throw" exceptions (errors would make it unreadable, see discussion in docs), I suspect some of the methods/functions are used during model building and should not return an error, but panic the error (or use
exceptions.Panicf. - The conventions is that errors must always have a stack-trace: this greatly facilitate pin-pointing building errors. Everywhere you return the pre-defined errors, use instead a
return errors.Wrap(ErrMyErrorType)instead (orpanic(errors.Wrap(ErrMyErrorType))if inside a model building function. They can still be checked later for the different types. - If using AI (I use it often for straight-forward implementations), consider including the
.agents/AGENT.mdof the GoMLX repo, as well as the included skill in.agents/skills/golang-gomlx/: it already explains to the AI about the errors, and other conventions. - The
rngparameter: the convention is always to use the random methods from themodel.Store(you can get it withscope.Store()). You can have an optional.WithRNG()method if someone wants something different.
Apologies, if that is me not understanding your abstractions. If this doesn't make sense, maybe we could schedule a quick VC to chat ? (sometimes it is easier)
|
Thank you. You were right that the original injection/registry API did not feel native to GoMLX. I refactored the package so that LoRA and QLoRA are graph-building layers configured from I removed the architecture-level injection and named-adapter machinery, moved validation failures in graph construction to GoMLX exceptions with stack traces, and used the model store RNG for adapter initialization. The documentation and tests now reflect this model. |
There was a problem hiding this comment.
Thanks @surya-mp, this is looking much better 👍🏻
There are still a few inconsistencies on how GoMLX is organized, I'll describe below. But I don't want to impose too much on you, if you don't want to do the change, let me know, and I will merge your changes as is (with the core of LoRA/QLoRA) and then change the API to match GoMLX, your choice 😃
My conceptual understanding is that (Q)LoRA should be a "meta" layer, that changes the way a dense layer (or other layers like attention) work. Ideally, it would have the exact same API, just add LoRA adapters to it, if configured (and if not configured, simply fall back and work exactly like the original).
With that in mind, I have the following comments:
- The
peft.Linearobject you created is something in between ann.DenseAPI and aml/layers/fnnAPI.- The
nnpackage doesn't handle variables (ormodel.Storeormodel.Scope), it handles the routing to fused-ops when available, and otherwise implements the fallback computation math. Usually, end users building models will not call these directly, and instead will use the variousml/layers(like thefnn). - The
layers(likeml/layers/fnn) packages are the one most folks will use, and handle the creation the variables and such.
- The
And then I have the following is a more radical suggestion/proposal:
- Add a
ParamAdapter = "peft_adapter"parameter, that defines the name of the (LoRA) adapter to use. If empty, no adapter is used (and it falls back to the usualfnn). The nice thing is that the same model can support various adapters, by simply changing this hyper-parameter -- meaning one can have a checkpoint with multiple different lora adapters -- or multiple checkpoints one per LoRA adapter, either way works. - We change the
fnnpackage to recognize thepeft.ParamAdapter("peft_adapter") parameter (maybe create a functionpeft.HasAdapter(scope) bool. If it returns true, thefnnpackage would freeze the normal weights (weights.SetTrainable(false)) calls apeft.AddLinear(layerScope, x)to add the LoRA projection (or similar). This way everyone gets LoRA adapter for free by simply configuring the "peft_adapter" hyperparameter, without having to do anything. - When creating the LoRA variables (A and B) create it under the
<FNNs scope> / <adapterName> / {A|B}. This keeps things cleanly organized. - One can eventually (future PR) update the variable checking CLI tool in
cmd/gomlx_checkpointswith flags to trim or copy (merge) LoRA adapter weights across checkpoints if one wants. It can also support the functionality to merge some LoRA weights into the main weights (or we can create another tool for that). - Another future suggested changes (for future PRs): add LoRA adapters to the
attentionlayer (and other layers) using the same format. - Yet another suggestion: I can work on updating
go-huggingfaceto be able to load models with LoRA adapters such that it would simply work for inference (including handling multi-tenancy, with multiple different LoRA adapters).
This way, the user doesn't need to do anything in their code. They simple set the hyperparameter "peft_adapter" to the name of the LoRA adapter they want to train or use (inference) and things just work.
What do you think ?
ps.: The fnn package doesn't handle yet PF4 quantized. I suggest we postpone QLoRA for a next PR, and I can better formalize the PF4 use case -- it also could benefit from some more optimized fused ops in the "go" backend, which I'm happy to do next.
| non-trainable. | ||
|
|
||
| ```go | ||
| modelScope := store.RootScope().In("model") |
There was a problem hiding this comment.
You can set the parameters in the store object directly (if no scope is given, it will set on the root scope), so here simply do:
store.SetParams(map[string]any{
peft.ParamRank: 16,
peft.ParamAlpha: float32(32),
peft.ParamDropout: float32(0.05),
peft.ParamBias: peft.BiasNone,
})|
hi @surya-mp , I was thinking: I want to implement the If you are ok, I could merge your changes as is, and then make the changes in the |
|
Hi @janpfeifer, you can go ahead and merge the PR. Regarding the go-huggingface package you mentioned for loading models with adapters just for inference, I've built something similar for my own use case. I plan to update that later once this is done, but let me know if you've already started working on it or if I can help. |
|
Thanks! Cool that you are already using it btw!! I'll update the other layers to automatically use |
|
See follow up in #469, where I make it such one just need to import The #469 was also merged, and the documentation also updated in https://gomlx.github.io/docs/layers/#parameter-efficient-fine-tuning-peft--lora |
Summary
Adds native parameter-efficient fine-tuning support to GoMLX through a new
ml/layers/peftpackage.The implementation provides:
docs/peft.md.Design
The package keeps GoMLX-specific PEFT behavior close to the graph and variable APIs:
Linearprovides LoRA over standard GoMLX dense weights.NF4Linearprovides QLoRA over packed NF4 base weights.InjectandInjectNF4let model architectures expose named modules and replace only configuredtargets.
matrices and configured biases.
Validation
go test ./ml/layers/peftThe focused test suite covers LoRA and QLoRA forward paths, frozen base weights, adapter-only gradients, named adapter selection, target injection, validation, and rollback behavior.