address issues that prevent using composition for layers like LoRA #177
Conversation
davidkoski
commented
Dec 17, 2024
davidkoski
commented
- see exploration of LoRA using composition mlx-swift-examples#167
- also fixes issue where quantize() could quantize a quantized layer!
- see ml-explore/mlx-swift-examples#167 - also fixes issue where quantize() could quantize a quantized layer!
| @@ -73,7 +73,7 @@ open class Linear: Module, UnaryLayer, Quantizable { | |||
| public let weight: MLXArray | |||
| public let bias: MLXArray? | |||
|
|
|||
| public var shape: (Int, Int) { | |||
| open var shape: (Int, Int) { | |||
There was a problem hiding this comment.
A lot of changes here are public -> open to allow subclasses to override.
| /// ``unfreeze(recursive:keys:strict:)``. | ||
| public private(set) var noGrad = Set<String>() | ||
| /// See ``noGrad()`` | ||
| private var _noGrad = Set<String>() |
There was a problem hiding this comment.
This is a breaking change but unlikely to be used directly (there are methods to manipulate). Subclasses can't override stored properties (they can do so for computed properties). This is replaced with methods that can be overridden.
| /// - Parameters: | ||
| /// - key: module key, see ``ModuleInfo`` | ||
| /// - value: the replacement module | ||
| open func updateModule(key: String, _ value: Any) throws { |
There was a problem hiding this comment.
Primarily exposed for subclasses to override (see test case)
| @@ -922,7 +969,7 @@ extension Module: Updatable, Evaluatable { | |||
| /// ### See Also | |||
| /// - <doc:layers> | |||
| /// - ``Sequential`` | |||
| public protocol UnaryLayer { | |||
| public protocol UnaryLayer: Module { | |||
There was a problem hiding this comment.
Per observation in ml-explore/mlx-swift-examples#167
| } | ||
|
|
||
| public func quantizeSingle(layer: Module, groupSize: Int = 64, bits: Int = 4) -> Quantized? { | ||
| if layer is Quantized { |
There was a problem hiding this comment.
Observed in the past and in the same area -- do not quantize already quantized layers.
| @ModuleInfo(key: "query_proj") public var queryProjection: UnaryLayer | ||
| @ModuleInfo(key: "key_proj") public var keyProjection: UnaryLayer | ||
| @ModuleInfo(key: "value_proj") public var valueProjection: UnaryLayer | ||
| @ModuleInfo(key: "out_proj") public var outProjection: UnaryLayer |
There was a problem hiding this comment.
This is how we would change e.g. attention layers if we wanted to use composition for LoRA.
| open override var shape: (Int, Int) { | ||
| let shape = weight.shape2 | ||
| return (shape.0, shape.1 * 32 / bits) | ||
| } |
There was a problem hiding this comment.
This is also a breaking change but IMHO it was broken before (returning the size of the quantized arrays which was useless)
| // like Linear/QuantizedLinear) but this verifies that it can | ||
| // be written this way. | ||
|
|
||
| class LoRA: Module, UnaryLayer { |
There was a problem hiding this comment.
A simple implementation of LoRA using composition, see ml-explore/mlx-swift-examples#167. This isn't meant as a reference implementation (though it may grow into that), but rather a test of the subclassing to make sure we have it all covered.
