Training Fast Weight Programmers by backpropagating through the Delta Rule #9
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ptxas info : 19 bytes gmem
ptxas info : Compiling entry function '_Z20causal_attend_kernelIN3c108BFloat16ELi4ELi2ELi4EEviPKT_S4_S4_PKfPS2_' for 'sm_86'
ptxas info : Function properties for _Z20causal_attend_kernelIN3c108BFloat16ELi4ELi2ELi4EEviPKT_S4_S4_PKfPS2_
72 bytes stack frame, 64 bytes spill stores, 116 bytes spill loads
ptxas info : Used 255 registers, 400 bytes cmem[0]
…ort sequences for now
proger
changed the title
proger
changed the title
proger
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Introducing a kernel for training a fast weight programmer by backpropagating through the delta rule (online linear regression) with @ischlag.
Improving on top of first order recurrence with scalar hidden state, this kernel uses vector-valued updates like the transformer, allowing use of matrix multiplication hardware, and avoiding saturation of capacity of the fast weights network, thanks to the delta rule.
The implementation uses chunking provided by @sustcsonglin's equation 6 and currently excels at the head dimension of 32, perfectly fitting into the registers of a 3090 warp. The code for tensor cores uses ThunderKittens which enable effortless WMMA.