jlens_glm52: GLM-5.2 serving + Jacobian-lens steering demo (examples)#261
Draft
maxsloef-goodfire wants to merge 29 commits into
Draft
jlens_glm52: GLM-5.2 serving + Jacobian-lens steering demo (examples)#261maxsloef-goodfire wants to merge 29 commits into
maxsloef-goodfire wants to merge 29 commits into
Conversation
…0/1)
Serve script: single B200 node, TP8, on-the-fly fp8 (bf16 ckpt is 1.4TB),
--enable-steering + filesystem capture consumer + graphsafe keys for the
lens band. build-env.sh installs the fork editable with precompiled
kernels (steering/capture are torch-level). pack_steering.py converts
MegaFire jlens steering exports ({layer: [hidden] fp32}) to
SteeringVectorSpecPacked; smoke_client.py exercises chat + steering
surface + steered-vs-baseline.
Lens artifacts: /mnt/data/artifacts/jlens/ (lens_glm52_1k.pt, unembed,
norms, steering_dog_demo.pt). Fit side: goodfire-ai/MegaFire PR RhizoNymph#18.
setup.py resolves the precompiled-wheel commit via git merge-base with upstream main; a shallow clone silently yields an install without vllm._C. Fail fast on shallow checkouts and treat a venv whose vllm._C does not import as broken (rebuild).
…potent)
Dose-response on 'Tell me about your weekend plans', greedy: baseline
(0 dog-words) -> L40 x0.65 ('process data, fetch balls, and maybe take a
nap') -> L40 x0.8 (model critiques its own draft as 'too dog-like') ->
3-layer x0.35-0.5 (coherent dog-persona weekend planning) -> x>=1.0
(token takeover). fp8 TP8 single node, decode-only steering.
Sends the reference prompts through /v1/completions with filesystem capture at the lens band and compares per-position cosine against Megatron-harvested residuals. Results (fp8 TP8 serving): mean cosine 0.98-0.99 at L30/40/50 across all 8 prompts (worst single position 0.85) — hook semantics (post_block == block output), tokenization and row order all agree; the gap is fp8 weight-quantization drift. Lens readout on the served residuals reproduces the reference top-k (Italy/Italian/currency at L50 on the boot-country prompt). Gotcha encoded in the script: bf16 captures decode as raw uint16 bit patterns (numpy has no bfloat16); a naive float cast anti-correlates (cos ~ -0.35) with huge norms. _to_float() reinterprets first.
Direct CaptureSink (per-step streaming, no request-end buffering): on
each captured chunk computes topk(unembed(J_l @ h)) for the configured
band and appends to {out}/{request}.jsonl -- a live 'what is the model
disposed to say at layer l' ticker during decode. Lazy weight load so
only the TP-rank-0 worker that receives captures pays the memory
(J band fp16 + lm_head fp16, ~2.2GB on the configured device).
live_demo.py renders generated-token vs band-readout tables and
composes with per-request steering (phase 4).
…3/4) Currency prompt: while generating ' shaped'/' like' the L50 readout shows ' coastline'/' peninsula'; the answer tokens read Italian/Italy — the multihop bridge entity live in the band during decode. Steered run (dog direction at L40 x0.8): L50 readout is ' dog' at every generated position while L30 (upstream of the injection) stays clean — inject-at-one-layer, detect-downstream, on one request.
Author
|
All four phases are now on this branch, validated end-to-end against a live fp8 TP8 single-node GLM-5.2 server:
Still draft pending a perf pass (decode overhead with the consumer armed) and placement feedback on the plugin location. 🤖 Generated with Claude Code |
FastAPI sidecar + single-page console: prompt box, steering panel (any word -> lens-derived direction packed on the fly; strength slider with takeover warning; band-layer selection), side-by-side baseline/steered completions, and a live per-token readout ticker (top-1 per band layer, top-5 on hover; amber = injected concept, teal = readout token that later appears in the generated text). One SSE stream carries token + readout events; the sidecar must run on the vllm node so the consumer's JSONL is visible without NFS attribute-cache lag (serve script now launches it alongside the server). Interaction patterns borrowed from Neuronpedia's steering UI (MIT).
- Multi-user safe: client-supplied request_id embeds in the consumer's readout filename; the tailer matches by prefix instead of racing on mtimes. - GLM-5.2's template pre-opens a <think> block; the sidecar segments the stream at </think> (tags may split across deltas) and the UI renders thinking as a labeled collapsible block, with a toggle that disables it via chat_template_kwargs. - max_tokens cap 256 -> 2048 (a 1000-token request was silently truncated mid-sentence).
The server's own context budget (--max-model-len minus prompt) is the real limit and vllm errors informatively past it; relay that error into the output pane instead of second-guessing with an arbitrary cap.
- Generate morphs into Stop while streaming; abort propagates through the sidecar (httpx stream close) so vllm cancels the request. - Ticker was O(n^2): per-event full-table rewrites, re-sorts, full generated-text substring scans, and unconditional reflow-forcing autoscroll made 4k-token runs crawl. Now: append-or-binary-insert rows, one cell write per event, Set-membership highlights, rAF-gated autoscroll that yields when the user scrolls up, and single merged text nodes per output region. - Defaults: inject at L40+L50, strength 0.30.
- lens_api.py: POST /api/lens/prompt speaking the Neuronpedia jlens UI NDJSON contract (meta -> prompt -> token* -> done): chat templating with prefill + enable_thinking, generation via /v1/completions with token-id prompts and return_tokens_as_token_ids, per-position readout assembly from the consumer stream, steer fields mapped onto packed steering vectors (additive only; ablate/swap rejected in-stream; reuse_len always 0; filterNonWordTokens ignored). - consumer: emits both JACOBIAN_LENS and LOGIT_LENS top-k with probabilities per (position, layer); final layer = identity transport. - serve script: stride-4 depth sweep (0..76 step 4, +77) as the capture band with matching graphsafe keys, jlens consumer baked in, MAX_MODEL_LEN default 32768. Deploys on next server rotation.
…ged hot-swap support)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CtV8yi9oZnBewk95RnUJtY
Pane events (tokens + readouts) cached keyed on everything that affects that pane — the baseline key excludes steering, so strength iteration replays the unsteered pass instantly (and identical multi-user queries hit too). Memory LRU + disk persistence, version-salted; partial/aborted panes never cached; cached panes get a 'cached' tag in the UI.
…ing parser) Sidecar prefers server-parsed reasoning_content deltas when the parser is active; the </think> tag-splitter remains as fallback for servers without it.
Third pane steers with the raw unembed direction w_t (no J transport) at the same layers/strength — the interventional analogue of the paper's readout methods-comparison. Cached independently; ticker stays on the jacobian pane.
… NP UI with /np/ base The stride-4 capture grid does not contain L30/L50 — console requests were rejected at admission. Sidecar now keeps every fitted J resident so steering can target any depth. NP UI assets 404'd at the /np/ mount (root-absolute Vite base) -> rebuilt with --base=/np/.
Two bugs from adopting --reasoning-parser glm45: (1) thinking arrives as reasoning_content deltas with no <think> tags in content, but the segmenter defaulted content to the think channel and waited for a </think> that never came -> whole answer landed in thinking, answer pane blank. Now route by field (reasoning_content -> think, content -> answer; literal-tag splitter kept only as a no-parser fallback). (2) GLM-5.2's template defaults thinking OFF, so omitting the kwarg meant the 'thinking' checkbox never enabled it — set enable_thinking explicitly both ways.
The glm45 parser buffers GLM-5.2's thinking into an unclosed reasoning block and emits nothing (finish=length, empty content+reasoning_content) when thinking is on; thinking-off is unaffected. Default the flag off; the sidecar's literal-tag splitter handles thinking when the parser is absent. Console defaults thinking-off for immediate clean answers.
…tag) GLM-5.2's chat template pre-opens the think block in the PROMPT, so the generation stream is reasoning terminated by a bare </think> then the answer — no opening <think> in the output. Start in the think channel when thinking is enabled and flip to answer at </think> (the rewrite had wrongly defaulted to answer). Verified: 285 think tokens then 46 answer tokens, cleanly split, no leaked tags. Re-enabled thinking by default.
…wser) fp8 residuals occasionally carry inf/nan, so a readout row's softmax produced NaN probs; Python json.dumps writes a bare NaN literal that the browser's JSON.parse rejects (SyntaxError: Unexpected token N). Fix at both ends: consumer nan_to_num's residuals+logits before softmax and zeroes non-finite probs; sidecar/lens_api recursively sanitize any non-finite float before re-serializing (deploys via sidecar restart, no server rotation needed).
GLM-5.2 fell through to the ChatML fallback (markers never match), so the token stream rendered ungrouped and the pre-assistant boundary position — the token the workspace paper cares most about — was not cleanly inspectable. Adds a glm JlensChatFormat with role-marker-driven groupTokens ([gMASK]<sop> prefix folds into the first header; <think> tags stay inline for the chat renderer's think styling) and registers it ahead of the qwen fallback. vite base pinned so npm run build cannot regress the /np mount.
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
1 participant
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.
What
Examples-only branch: serve GLM-5.2 (
GlmMoeDsaForCausalLM, 744B MoE) on this fork and drive it with steering vectors derived from a fitted Jacobian lens — the first live use of the steering subsystem with lens-transport directions on a ~750B model. No engine changes.New under
examples/jlens_glm52/:build-env.sh— editable install withVLLM_USE_PRECOMPILED(+ guards: shallow clones silently break the merge-base wheel resolution and yield an install withoutvllm._C; the script fails fast and treats avllm._C-less venv as broken)sbatch-serve-glm52.sh— single-node 8-GPU serve, TP8, on-the-fly fp8 (the bf16 checkpoint does not fit one node),--enable-steering, filesystem capture consumer, graphsafe keys for the lens bandpack_steering.py— converts lens steering exports ({layer: [hidden] fp32}.pt) toSteeringVectorSpecPackedsmoke_client.py— baseline chat + steering-surface check + steered comparisonsteering_demo_transcripts.md— dose-response transcripts (greedy, seed 0, decode-only steering of a " dog" direction atpost_block): baseline → "process data, fetch balls, and maybe take a nap" (L40 ×0.65) → the model critiquing its own draft as "too dog-like" (×0.8) → coherent dog-persona (3 layers ×0.35–0.5) → token takeover (≥×1.0)Notes
phase2_parity.py).jlenscapture consumer (CaptureSink) streaming per-token lens readoutsunembed(J_l @ h)during decode.🤖 Generated with Claude Code
https://claude.ai/code/session_01CtV8yi9oZnBewk95RnUJtY