8000 Refactor PEFT to use post-linear lifecycle hooks (`peft_adapter`) by janpfeifer · Pull Request #469 · gomlx/gomlx · GitHub
[go: up one dir, main page]

Skip to content
< 7438 /div>

Refactor PEFT to use post-linear lifecycle hooks (peft_adapter) - #469

Merged
janpfeifer merged 4 commits into
mainfrom
peft-api-change
Oct 9, 2026
Merged

janpfeifer merged 4 commits into
mainfrom
peft-api-change

Conversation

@janpfeifer
@janpfeifer janpfeifer commented Oct 9, 2026 •
Copy link
Copy Markdown
Contributor

Description

This PR refactors the PEFT (Parameter-Efficient Fine-Tuning) API to operate as a meta-layer via post-linear lifecycle hooks rather than requiring manual projection replacements.

A follow up to #466

Key Changes

  • Post-Linear Hook Registry (ml/layers/hooks.go):
    • Adds PostLinearHook (func(scope *model.Scope, input, output *Node) *Node).
    • Implements RegisterPostLinearHook(priority, hook) with priority-based sorting and thread-safe registration.
    • Implements ApplyPostLinearHooks and ClearPostLinearHooks.
  • Layer Hook Integration:
    • layers.DenseWithLayout: Dispatches post-linear hooks immediately after nn.Dense.
    • fnn.Done(): Dispatches post-linear hooks after each layer's linear projection, before non-linearities (activation/dropout/normalization/residual).
    • attention.MultiHeadAttention: Delegates query, key, value, and output projections through layers.DenseWithLayout, routing them through the hooks automatically.
  • PEFT Hook Implementation (ml/layers/peft):
    • Automatically registers in init() with HookPriorityDefault.
    • peft_adapter: Activates LoRA with adapter sub-scope name (e.g., scope.In(adapterName)/A,B). If empty (""), PEFT is inactive with zero overhead.
    • Freeze-only mode: Setting peft_adapter to "off", "disabled", "false", or "none" freezes base weights without creating adapters.
    • Selective module targeting: peft_target_modules (comma-separated or []string) filters matching module scopes; unmatched linear layers are frozen without adapters.
    • Multi-adapter coexistence: Multiple adapters can coexist in the same model.Store under their respective sub-scopes.
  • Documentation & Agent Skills:
    • Updated docs/peft.md with usage, CLI -set parameters, multi-adapter patterns, and a Developer Guide.
    • Updated .agents/AGENTS.md and .agents/skills/golang-gomlx/.
  • Tests:
    • Unit tests in ml/layers/hooks_test.go.
    • Comprehensive end-to-end tests in ml/layers/peft/peft_test.go covering FNN, Dense, and MultiHeadAttention with LoRA, target module filtering, freeze-only mode, and multi-adapter switching.

…ions.

Added `ml/layers/hooks.go` defining the `PostLinearHook` for after-linear projections.
Updated documentation.
@janpfeifer
janpfeifer merged commit e01d4a1 into main Oct 9, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

0