Repository navigation
GLM: an optional residual-direction ablation (carved out of #642) - #901
Open
ClintMoody wants to merge 3 commits into
Open
ClintMoody wants to merge 3 commits into
ClintMoody wants to merge 3 commits into
Conversation
… work STRATA_GLM_ABLATE=<file> projects directions out of the GLM fast path's residual streams (research; off unless the variable is set). Outputs mode removes K orthonormal vectors from the token embedding and every block output; per-layer mode (STRATA_GLM_ABLATE_LAYERS) gives each layer its own vectors, the same operation as exllamav3's EXL3_CVEC hook. This is the harness alone, kept off the paged-attention branch per review. It is not wired into a translation unit yet; the follow-up wires it.
…cument them The harness in glm_ablate.cuh was carved out of the paged-attention work, where it had no business: it is a setting of its own. This adds the wiring back, on top of the carved header: - the include, and the direction's upload on the device before the expert pool is sized (so the pool sees the bytes it takes); - the mixer's and the FFN's writes, and the token embedding - x - s (x . v) v before each mHC write, in outputs mode; - both operands of every layer's FFN write, in per-layer mode; - the docs/SETTINGS.md section and the three MAYA_GLM_ABLATE* rows. Off unless MAYA_GLM_ABLATE is set: nothing runs, and the answers do not change.
glm_ablate_test runs the rows kernel on a device and compares it with a double-precision reference: outputs mode, the scale at 0.5, and per-layer mode (layer 0 steered, layer 1 all zero and untouched). A device that cannot be opened is a SKIP, as the other GPU tests do. Adds the target and three ctest entries. The engine does not link it.
Owner
|
Thanks for splitting this out. It is off unless the env var names a file, which is good. Which GPUs/model/quant did you test on, and can you confirm with it unset the outputs and speed are identical to main (greedy output match, decode and prefill)? We will look at whether this belongs in Maya. |
This file contains hidden or 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
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.
The paged-attention work (#642) grew a setting of its own: a harness that removes chosen directions from the GLM path's residual streams. @catonooka asked for one topic per PR, so it moved here. #642 carries no
glm_ablatereference now.What it does
x <- x - s (x . v) v, applied to the way into every write into the residual streams.src/core/glm_ablate.cuhholds the projection and its three settings. Off unlessMAYA_GLM_ABLATEnames a file of raw float32 directions.K x n_embdvalues, K orthonormal directions. Applied to the token embedding and every block output, the mixer's and the FFN's, before the mHC write.MAYA_GLM_ABLATE_LAYERS=n_layers):n_layers x K x n_embdvalues. Layer l's own vectors are applied to the two operands of that layer's FFN write, so the fused read needs no change. A layer whose vectors are all zero is not steered.MAYA_GLM_ABLATE_SCALE(default 1) sets s: 1 removes the component along each direction, above 1 over-projects.The same operation as exllamav3's
EXL3_CVEChook after each block, and the knob DeepSeek's--ablategives: find a direction that carries a behaviour, remove it, watch the model.What it costs
Nothing when unset: each call site checks one bool and returns. The NextN draft block is not touched (a draft only has to be a good guess).
Wiring
src/core/glm_fast_path.cu: the include; the direction's upload on the device before the expert pool is sized; the mixer's and the FFN's writes; the last layer's write; the token embedding; and both operands of every layer's FFN write in per-layer mode. Nine call sites in all.Tests
tests/glm_ablate_test.curuns the kernel on a device and compares every value with a double-precision reference:x - (x . v) vfor every direction, on every row, and no component along either direction survives;Run here on a 3090 Ti against this tree: 4, 4 and 5 checks, 0 failures. Three ctest entries (
glm_ablate_test,glm_ablate_scale_test,glm_ablate_layers_test). A device that cannot be opened is a SKIP, as the other GPU tests do, so a CPU-only build still runs the suite.Notes
MAYA_*is the alias:include/strata/platform/env_alias.hppmaps everyMAYA_XtoSTRATA_X, andMAYA_Xwins when both are set. Each program callsenv_aliases()first thing inmain().glm ablate:line and leaves the stock model running. No GPU room turns the setting off.docs/SETTINGS.mdhas the section, in the shape of the other rows there.This branch changes no default path, so the
CONTRIBUTING.mdbefore/aftermaya.py --benchpair does not apply to it: withMAYA_GLM_ABLATEunset the work is the same, plus one bool check per call site.