29 | if (!info.ablate_loss_diff) {
30 | return <> ;
31 | }
32 | const lossDiffsNorm = normalizeToUnitInterval(info.ablate_loss_diff.map((x) => (-x)));
33 | return (
34 |
35 | {info.tokens.map((token, idx) => {
36 | const highlight = idx === info.idx;
37 | const loss_diff = (idx === 0) ? 0: info.ablate_loss_diff[idx-1];
38 | const kl = (idx === 0) ? 0: info.kl[idx-1];
39 | const activation = info.acts[idx];
40 | const top_downvotes = (idx === 0) ? [] : info.top_downvotes_logits[idx-1];
41 | const top_downvote_tokens = (idx === 0) ? [] : info.top_downvote_tokens_logits[idx-1];
42 | const top_upvotes = (idx === 0) ? [] : info.top_upvotes_logits[idx-1];
43 | const top_upvote_tokens = (idx === 0) ? [] : info.top_upvote_tokens_logits[idx-1];
44 | // const top_downvotes_weighted = (idx === 0) ? [] : info.top_downvotes_weighted[idx-1];
45 | // const top_downvote_tokens_weighted = (idx === 0) ? [] : info.top_downvote_tokens_weighted[idx-1];
46 | // const top_upvotes_weighted = (idx === 0) ? [] : info.top_upvotes_weighted[idx-1];
47 | // const top_upvote_tokens_weighted = (idx === 0) ? [] : info.top_upvote_tokens_weighted[idx-1];
48 | const top_downvotes_probs = (idx === 0) ? [] : info.top_downvotes_probs[idx-1];
49 | const top_downvote_tokens_probs = (idx === 0) ? [] : info.top_downvote_tokens_probs[idx-1];
50 | const top_upvotes_probs = (idx === 0) ? [] : info.top_upvotes_probs[idx-1];
51 | const top_upvote_tokens_probs = (idx === 0) ? [] : info.top_upvote_tokens_probs[idx-1];
52 | const color = getInterpolatedColor(colors, [-1, 0, 1], (idx === 0) ? 0 : lossDiffsNorm[idx-1]);
53 | if (!renderNewlines) {
54 | token = token.replace(/\n/g, '↵')
55 | }
56 | return
65 | {token}
66 |
67 | }
68 | tooltip={(idx <= info.idx) ? (prediction prior to ablated token)
:
69 | Loss diff: {loss_diff.toExponential(2)}
70 | KL(clean || ablated): {kl.toExponential(2)}
71 | Logit diffs:
72 |
73 |
74 |
75 | {
76 | ['', /*' (weighted)',*/ ' (probs)'].map((suffix, i) => {
77 | return
78 | Upvoted {suffix} Downvoted {suffix}
81 | })
82 | }
83 |
84 |
85 |
86 | {
87 | top_upvotes.map((upvote, j) => {
88 | const downvote = top_downvotes[j];
89 | return
90 | {upvote.toExponential(1)}
91 | {top_upvote_tokens[j]}
92 | {downvote.toExponential(1)}
93 | {top_downvote_tokens[j]}
94 | {top_upvotes_probs[j].toExponential(1)}
95 | {top_upvote_tokens_probs[j]}
96 | {top_downvotes_probs[j].toExponential(1)}
97 | {top_downvote_tokens_probs[j]}
98 |
99 | })
100 | }
101 |
102 |
103 |
108 | )
109 |
110 | }
111 |
--------------------------------------------------------------------------------
/sae-viewer/src/components/featureSelect.tsx:
--------------------------------------------------------------------------------
1 | import FeatureInfo from "./components/featureInfo"
2 | import React, { useEffect, useState, FormEvent } from "react"
3 | import Welcome from "./welcome"
4 | import { Feature } from "./types"
5 | import { useState } from "react"
6 |
7 | import { pathForFeature, DEFAULT_AUTOENCODER, SUBJECT_MODELS, AUTOENCODER_FAMILIES } from "../autoencoder_registry"
8 |
9 | type FeatureSelectProps = {
10 | init_feature: Feature,
11 | onFeatureChange?: (feature: Feature) => void,
12 | onFeatureSubmit: (feature: Feature) => void,
13 | }
14 |
15 | export default function FeatureSelect({init_feature, onFeatureChange, onFeatureSubmit, show_go}: FeatureSelectProps) {
16 | let [feature, setFeature] = useState(init_feature);
17 | let family = AUTOENCODER_FAMILIES[feature.autoencoder.family];
18 | let [warningAcknowledged, setWarningAcknowledged] = useState
(localStorage.getItem('warningAcknowledged') === 'true');
19 | let changeFeature = (feature: Feature) => {
20 | onFeatureChange && onFeatureChange(feature);
21 | setFeature(feature);
22 | }
23 | // console.log('features', feature.autoencoder.num_features)
24 | const feelingLuckySubmit = () => {
25 | const atom = Math.floor(Math.random() * feature.autoencoder.num_features);
26 | const random_feature: Feature = {atom, autoencoder: feature.autoencoder}
27 | setFeature(random_feature);
28 | onFeatureSubmit(random_feature);
29 | return false
30 | }
31 | const acknowledgeWarning = () => {
32 | localStorage.setItem('warningAcknowledged', 'true');
33 | setWarningAcknowledged(true);
34 | }
35 |
36 | if (!warningAcknowledged) {
37 | return (
41 | Note: by clicking this button you acknowledge that the content of the documents are taken randomly from the internet, and may contain offensive or inappropriate content.
42 | )
43 | }
44 |
45 | return (
46 | <>
47 |
48 |
49 | Subject model {" "}
50 | {
51 | let family = Object.values(AUTOENCODER_FAMILIES).find((family) => (family.subject === e.target.value));
52 | changeFeature({
53 | atom: 0, autoencoder: family.get_ae(family.default_H(feature.autoencoder.H))
54 | })
55 | }}>
56 | {Object.values(SUBJECT_MODELS).map((subject_model) => (
57 | {subject_model}
58 | ))}
59 |
60 |
61 | {" "}
62 | Autoencoder {" "}
65 | family {" "}
66 | {
67 | let family = AUTOENCODER_FAMILIES[e.target.value];
68 | changeFeature({
69 | atom: 0, autoencoder: family.get_ae(family.default_H(feature.autoencoder.H))
70 | })
71 | }}>
72 | {Object.values(AUTOENCODER_FAMILIES).filter((family) => (family.subject === feature.autoencoder.subject)).map((family) => (
73 | {AUTOENCODER_FAMILIES[family.name].label}
74 | )) }
75 |
76 |
77 | {
78 | AUTOENCODER_FAMILIES[feature.autoencoder.family].selectors.map((selector) => (
79 |
80 | {" "}
81 | {selector.label || selector.key}
82 | {
84 | let family = AUTOENCODER_FAMILIES[feature.autoencoder.family];
85 | changeFeature({
86 | atom: 0, autoencoder: family.get_ae({...feature.autoencoder.H, [selector.key]: e.target.value})
87 | })
88 | }}
89 | >
90 | {selector.values.map((value) => (
91 | {value}
92 | ))}
93 |
94 |
95 | ))
96 | }
97 | {
98 | family.warning ? Note: {family.warning} : null
99 | }
100 |
101 | Feature
103 | {
118 | e.preventDefault()
119 | onFeatureSubmit(feature)
120 | return false
121 | }}
122 | className="border border-gray-300 rounded-md p-2"
123 | style={{ width: 200 }}
124 | disabled={!warningAcknowledged}
125 | >
126 | Go to feature {feature.atom}
127 |
128 | }
129 |
135 | I'm feeling lucky
136 |
137 |
138 |
139 |
142 |
143 | )
144 | }
145 |
--------------------------------------------------------------------------------
/sae-viewer/src/components/featureInfo.tsx:
--------------------------------------------------------------------------------
1 | import React, { useEffect, useState, useRef } from "react"
2 | import { normalizeSequences, SequenceInfo, Feature, FeatureInfo } from "../types"
3 | import TokenHeatmap from "./tokenHeatmap";
4 | import TokenAblationmap from "./tokenAblationmap";
5 | import Histogram from "./histogram"
6 | import Tooltip from "./tooltip"
7 |
8 | import {get_feature_info} from "../interpAPI"
9 |
10 | export default ({ feature }: {feature: Feature}) => {
11 | const [data, setData] = useState(null as FeatureInfo | null)
12 | const [showingMore, setShowingMore] = useState({})
13 | const [renderNewlines, setRenderNewlines] = useState(false)
14 | const [isLoading, setIsLoading] = useState(true)
15 | const [got_error, setError] = useState(null)
16 | const currentFeatureRef = useRef(feature);
17 |
18 |
19 | useEffect(() => {
20 | async function fetchData() {
21 | setIsLoading(true)
22 | try {
23 | currentFeatureRef.current = feature; // Update current feature in ref on each effect run
24 | const result = await get_feature_info(feature)
25 | if (currentFeatureRef.current !== feature) {
26 | return;
27 | }
28 | normalizeSequences(result.top, result.random)
29 | result.top.sort((a, b) => b.act - a.act);
30 | setData(result)
31 | setIsLoading(false)
32 | setError(null);
33 | } catch (e) {
34 | setError(e);
35 | }
36 | try {
37 | const result = await get_feature_info(feature, true)
38 | if (currentFeatureRef.current !== feature) {
39 | return;
40 | }
41 | normalizeSequences(result.top, result.random)
42 | result.top.sort((a, b) => b.act - a.act);
43 | setData(result)
44 | setIsLoading(false)
45 | setError(null);
46 | } catch (e) {
47 | setError('Note: ablation effects data not available for this model');
48 | }
49 | }
50 | fetchData()
51 | }, [feature])
52 |
53 | if (isLoading) {
54 | return (
55 |
56 |
57 |
loading top dataset examples
58 | {
59 | got_error ?
Error loading data: {got_error} : null
60 | }
61 |
85 | {
86 | got_error ?
{got_error} : null
87 | }
88 |
89 |
90 |
92 |
93 |
94 |
95 | 0]'}/>
96 | {data.density.toExponential(2)}
98 |
99 |
100 |
101 | {data.mean_act ? data.mean_act.toExponential(2) : 'data not available'}
103 |
104 |
105 | E[a2 ]}/> {data.mean_act_squared ? data.mean_act_squared.toExponential(2): 'data not available'}
107 |
108 |
109 | E[a3 ]/(E[a2 ])1.5 }/> {data.skew ? data.skew.toExponential(2) : 'data not available'}
111 |
112 |
113 | E[a4 ]/(E[a2 ])2 }/> {data.kurtosis ? data.kurtosis.toExponential(2) : 'data not available'}
115 |
116 |
117 |
118 |
119 | {
120 | all_sequences.map(({label, sequences, default_show}, idx) => {
121 | // console.log('sequences', sequences)
122 | const n_show = showingMore[label] ? sequences.length : default_show;
123 | return (
124 |
125 |
126 | {label}
127 | setShowingMore({...showingMore, [label]: !showingMore[label]})}>
129 | {showingMore[label] ? 'show less' : 'show more'}
130 |
131 | setRenderNewlines(!renderNewlines)}>
133 | {renderNewlines ? 'collapse newlines' : 'show newlines'}
134 |
135 |
136 |
137 |
138 |
139 | Doc ID Token Activation Activations
140 | {sequences.length && sequences[0].ablate_loss_diff && Effects }
141 |
142 |
143 |
144 | {sequences.slice(0, n_show).map((sequence, i) => (
145 |
146 | {sequence.doc_id} {sequence.idx} {sequence.act.toFixed(2)}
147 |
148 |
149 |
150 | {
151 | sequence.ablate_loss_diff &&
152 |
153 |
154 |
155 | }
156 |
157 | ))}
158 |
159 |
160 |
161 | )
162 | })
163 | }
164 |
113 |
Welcome! This is a viewer for sparse autoencoders features trained in this paper
114 |
Pick a feature:
115 |
setFeature(f)}
118 | onFeatureSubmit={(f: Feature) => navigate(pathForFeature(f))}
119 | show_go={true}
120 | />
121 |
122 |
123 |
Interesting features:
124 |
125 |
128 | {displayFeatures.map(({ heading, heading_type, feature, label, description }, j) => (
129 | heading ?
130 | {React.createElement(heading_type, {}, heading)}
131 |
:
handleClick(feature)}
134 | style={{ width: 200 }}
135 | className="text-blue-500 hover:text-blue-700"
136 | title={description}
137 | >
138 | {label}
139 |
140 | ))}
141 |
142 |
143 |
"
33 |
34 |
35 | class NtgExplanation(Explanation):
36 | def __init__(self, trie: dict):
37 | self.trie = trie
38 |
39 | def todict(self) -> dict:
40 | return {
41 | "trie": self.trie,
42 | }
43 |
44 | @classmethod
45 | def load(cls, serialized: dict) -> "Explanation":
46 | assert isinstance(serialized, dict)
47 | return cls(serialized["trie"])
48 |
49 | def predict(self, tokens: list[str]) -> list[float]:
50 | predicted_acts = []
51 | # for each token, traverse the trie beginning from that token and proceeding in reverse order until we match
52 | # a pattern or are no longer able to traverse.
53 | for i in range(len(tokens)):
54 | curr = self.trie
55 | for j in range(i, -1, -1):
56 | if tokens[j] not in curr and _ANY_TOKEN not in curr:
57 | predicted_acts.append(0)
58 | break
59 | if tokens[j] in curr:
60 | curr = curr[tokens[j]]
61 | else:
62 | curr = curr[_ANY_TOKEN]
63 | if _SALIENCY_KEY in curr:
64 | predicted_acts.append(curr[_SALIENCY_KEY])
65 | break
66 | # if we"ve reached the end of the sequence and haven't found a saliency value, append 0.
67 | elif j == 0:
68 | if _START_TOKEN in curr:
69 | curr = curr[_START_TOKEN]
70 | assert _SALIENCY_KEY in curr
71 | predicted_acts.append(curr[_SALIENCY_KEY])
72 | break
73 | predicted_acts.append(0)
74 | # We should have appended a value for each token in the sequence.
75 | assert len(predicted_acts) == len(tokens)
76 | return predicted_acts
77 |
78 | # TODO make this more efficient
79 | def predict_many(self, tokens_batch: list[list[str]]) -> list[list[float]]:
80 | return [self.predict(t) for t in tokens_batch]
81 |
82 |
83 | def batched(iterable, bs):
84 | batch = []
85 | it = iter(iterable)
86 | while True:
87 | batch = []
88 | try:
89 | for _ in range(bs):
90 | batch.append(next(it))
91 | yield batch
92 | except StopIteration:
93 | if len(batch) > 0:
94 | yield batch
95 | return
96 |
97 |
98 | def apply_batched(fn, iterable, bs):
99 | for batch in batched(iterable, bs):
100 | ret = fn(batch)
101 | assert len(ret) == len(batch)
102 | yield from ret
103 |
104 |
105 | def batch_parallelize(algos, fn, batch_size):
106 | """
107 | Algorithms are coroutines that yield items to be processed in parallel.
108 | We concurrently run the algorithm on all items in the batch.
109 | """
110 | inputs = []
111 | for i, algo in enumerate(algos):
112 | inputs.append((i, next(algo)))
113 | results = [None] * len(algos)
114 | while len(inputs) > 0:
115 | ret = list(apply_batched(fn, [x[1] for x in inputs], batch_size))
116 | assert len(ret) == len(inputs)
117 | inds = [x[0] for x in inputs]
118 | inputs = []
119 | for i, r in zip(inds, ret):
120 | try:
121 | next_input = algos[i].send(r)
122 | inputs.append((i, next_input))
123 | except StopIteration as e:
124 | results[i] = e.value
125 | return results
126 |
127 |
128 | def create_n2g_explanation(
129 | model_fn: Callable, train_set: list[dict], batch_size: int = 16,
130 | padding_token=4808 # " _" for GPT-2
131 | ) -> NtgExplanation:
132 | truncated = []
133 | # for each one find the index of the selected activation in the doc. truncate the sequences after this point.
134 | for doc in train_set:
135 | # get index of selected activation. for docs stored in 'top', this is the max activation.
136 | # for docs stored in 'random', it is a random positive activation (we sample activations, not docs
137 | # to populate 'random', so docs with more positive activations are more likely to be included).
138 | max_idx = doc["idx"]
139 | truncated.append(
140 | {
141 | "act": doc["act"],
142 | "acts": doc["acts"][: max_idx + 1],
143 | "tokens": doc["tokens"][: max_idx + 1],
144 | "token_ints": doc["token_ints"][: max_idx + 1],
145 | }
146 | )
147 |
148 | def get_minimal_subsequence(doc):
149 | for i in range(len(doc["token_ints"]) - 1, -1, -1):
150 | atom_acts = yield doc["token_ints"][i:]
151 | assert (
152 | len(atom_acts) == len(doc["token_ints"]) - i
153 | ), f"{len(atom_acts)} != {len(doc['token_ints']) - i}"
154 | if atom_acts[-1] / doc["act"] >= 0.5:
155 | return {
156 | "tokens": doc["tokens"][i:],
157 | "token_ints": doc["token_ints"][i:],
158 | "subsequence_act": atom_acts[-1],
159 | "orig_act": doc["act"],
160 | }
161 | print("Warning: no minimal subsequence found")
162 | # raise ValueError("No minimal subsequence found")
163 | return {
164 | "tokens": doc["tokens"],
165 | "token_ints": doc["token_ints"],
166 | "subsequence_act": doc["act"],
167 | "orig_act": doc["act"],
168 | }
169 |
170 | minimal_subsequences = batch_parallelize(
171 | [get_minimal_subsequence(doc) for doc in truncated], model_fn, batch_size
172 | )
173 |
174 | start_padded = apply_batched(
175 | model_fn,
176 | [[padding_token] + doc["token_ints"] for doc in minimal_subsequences],
177 | batch_size,
178 | )
179 | for min_seq, pad_atom_acts in zip(minimal_subsequences, start_padded):
180 | min_seq["can_pad_start"] = pad_atom_acts[-1] / min_seq["orig_act"] >= 0.5
181 |
182 | # for m in minimal_subsequences:
183 | # print("\t" + "".join(m["tokens"]))
184 |
185 | # for each token in a minimal subsequence, replace it with a padding token and compute the saliency value (1 - (orig act / new act))
186 | for doc in minimal_subsequences:
187 | all_seqs = []
188 | for i in range(len(doc["token_ints"])):
189 | tokens = doc["token_ints"][:i] + [padding_token] + doc["token_ints"][i + 1 :]
190 | assert len(tokens) == len(doc["token_ints"])
191 | all_seqs.append(tokens)
192 | saliency_vals = []
193 | all_atom_acts = apply_batched(model_fn, all_seqs, batch_size)
194 | for atom_acts, tokens in zip(all_atom_acts, all_seqs):
195 | assert len(atom_acts) == len(tokens)
196 | saliency_vals.append(1 - (atom_acts[-1] / doc["subsequence_act"]))
197 | doc["saliency_vals"] = saliency_vals
198 |
199 | trie = {}
200 | for doc in minimal_subsequences:
201 | curr = trie
202 | for i, (token, saliency) in enumerate(zip(doc["tokens"][::-1], doc["saliency_vals"][::-1])):
203 | if saliency < 0.5:
204 | token = _ANY_TOKEN
205 | if token not in curr:
206 | curr[token] = {}
207 | curr = curr[token]
208 | if i == len(doc["tokens"]) - 1:
209 | if not doc["can_pad_start"]:
210 | curr[_START_TOKEN] = {}
211 | curr = curr[_START_TOKEN]
212 | curr[_SALIENCY_KEY] = doc["subsequence_act"]
213 |
214 | return NtgExplanation(trie)
215 |
216 |
217 | if __name__ == "__main__":
218 |
219 | def first_position_fn(tokens: list[list[float]]) -> list[list[float]]:
220 | return [[1.0] + [0.0] * (len(toks) - 1) for toks in tokens]
221 |
222 | expl = create_n2g_explanation(
223 | first_position_fn,
224 | [
225 | {
226 | "idx": 0,
227 | "act": 1.0,
228 | "acts": [1.0, 0.0, 0.0],
229 | "tokens": ["a", "b", "c"],
230 | "token_ints": [0, 1, 2],
231 | },
232 | {
233 | "idx": 0,
234 | "act": 1.0,
235 | "acts": [1.0, 0.0, 0.0],
236 | "tokens": ["b", "c", "d"],
237 | "token_ints": [1, 2, 3],
238 | },
239 | ],
240 | )
241 | print(expl.trie)
242 | assert expl.predict(["a", "b", "c"]) == [1.0, 0.0, 0.0]
243 | assert expl.predict(["c", "b", "a"]) == [1.0, 0.0, 0.0]
244 |
245 | def c_fn(tokens: list[list[float]]) -> list[list[float]]:
246 | return [[1.0 if tok == 2 else 0.0 for tok in toks] for toks in tokens]
247 |
248 | expl = create_n2g_explanation(
249 | c_fn,
250 | [
251 | {
252 | "idx": 2,
253 | "act": 1.0,
254 | "acts": [0.0, 0.0, 1.0],
255 | "tokens": ["a", "b", "c"],
256 | "token_ints": [0, 1, 2],
257 | },
258 | {
259 | "idx": 1,
260 | "act": 1.0,
261 | "acts": [0.0, 1.0, 0.0],
262 | "tokens": ["b", "c", "d"],
263 | "token_ints": [1, 2, 3],
264 | },
265 | ],
266 | )
267 |
268 | print(expl.trie)
269 | assert expl.predict(["b", "c"]) == [0.0, 1.0]
270 | assert expl.predict(["a", "a", "a"]) == [0.0, 0.0, 0.0]
271 | assert expl.predict(["c", "b", "c"]) == [1.0, 0.0, 1.0]
272 |
273 | def a_star_c_fn(tokens: list[list[float]]) -> list[list[float]]:
274 | return [
275 | [1.0 if (tok == 2 and 0 in toks[:i]) else 0.0 for i, tok in enumerate(toks)]
276 | for toks in tokens
277 | ]
278 |
279 | expl = create_n2g_explanation(
280 | a_star_c_fn,
281 | [
282 | {
283 | "idx": 2,
284 | "act": 1.0,
285 | "acts": [0.0, 0.0, 1.0],
286 | "tokens": ["a", "b", "c"],
287 | "token_ints": [0, 1, 2],
288 | },
289 | {
290 | "idx": 2,
291 | "act": 1.0,
292 | "acts": [0.0, 0.0, 1.0],
293 | "tokens": ["b", "a", "c"],
294 | "token_ints": [1, 0, 2],
295 | },
296 | {
297 | "idx": 1,
298 | "act": 1.0,
299 | "acts": [0.0, 1.0, 0.0],
300 | "tokens": ["a", "c", "d"],
301 | "token_ints": [0, 2, 3],
302 | },
303 | ],
304 | )
305 |
306 | print(expl.trie)
307 | assert expl.predict(["b", "c"]) == [0.0, 0.0]
308 | assert expl.predict(["a", "c"]) == [0.0, 1.0]
309 | assert expl.predict(["a", "e", "c", "a", "c"]) == [0.0, 0.0, 1.0, 0.0, 1.0]
310 | # NOTE: should be 0, 0, 0, 1 but we're not smart enough to handle this yet
311 | assert expl.predict(["a", "b", "b", "c"]) == [0.0, 0.0, 0.0, 0.0]
312 |
313 | def zero_fn(tokens: list[list[float]]) -> list[list[float]]:
314 | return [[0.0 for tok in toks] for toks in tokens]
315 |
316 | expl = create_n2g_explanation(zero_fn, [])
317 |
318 | print(expl.trie)
319 | assert expl.predict(["b", "c"]) == [0.0, 0.0]
320 | assert expl.predict(["a", "c"]) == [0.0, 0.0]
321 | assert expl.predict(["a", "e", "c", "a", "c"]) == [0.0, 0.0, 0.0, 0.0, 0.0]
322 |
--------------------------------------------------------------------------------
/sparse_autoencoder/kernels.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | import triton
4 | import triton.language as tl
5 |
6 |
7 | ## kernels
8 |
9 |
10 | def triton_sparse_transpose_dense_matmul(
11 | sparse_indices: torch.Tensor,
12 | sparse_values: torch.Tensor,
13 | dense: torch.Tensor,
14 | N: int,
15 | BLOCK_SIZE_AK=128,
16 | ) -> torch.Tensor:
17 | """
18 | calculates sparse.T @ dense (i.e reducing along the collated dimension of sparse)
19 | dense must be contiguous along dim 0 (in other words, dense.T is contiguous)
20 |
21 | sparse_indices is shape (A, k)
22 | sparse_values is shape (A, k)
23 | dense is shape (A, B)
24 |
25 | output is shape (N, B)
26 | """
27 |
28 | assert sparse_indices.shape == sparse_values.shape
29 | assert sparse_indices.is_contiguous()
30 | assert sparse_values.is_contiguous()
31 | assert dense.is_contiguous() # contiguous along B
32 |
33 | K = sparse_indices.shape[1]
34 | A = dense.shape[0]
35 | B = dense.shape[1]
36 | assert sparse_indices.shape[0] == A
37 |
38 | # COO-format and sorted
39 | sorted_indices = sparse_indices.view(-1).sort()
40 | coo_indices = torch.stack(
41 | [
42 | torch.arange(A, device=sparse_indices.device).repeat_interleave(K)[
43 | sorted_indices.indices
44 | ],
45 | sorted_indices.values,
46 | ]
47 | ) # shape (2, A * K)
48 | coo_values = sparse_values.view(-1)[sorted_indices.indices] # shape (A * K,)
49 | return triton_coo_sparse_dense_matmul(coo_indices, coo_values, dense, N, BLOCK_SIZE_AK)
50 |
51 |
52 | def triton_coo_sparse_dense_matmul(
53 | coo_indices: torch.Tensor,
54 | coo_values: torch.Tensor,
55 | dense: torch.Tensor,
56 | N: int,
57 | BLOCK_SIZE_AK=128,
58 | ) -> torch.Tensor:
59 | AK = coo_indices.shape[1]
60 | B = dense.shape[1]
61 |
62 | out = torch.zeros(N, B, device=dense.device, dtype=coo_values.dtype)
63 |
64 | grid = lambda META: (
65 | triton.cdiv(AK, META["BLOCK_SIZE_AK"]),
66 | 1,
67 | )
68 | triton_sparse_transpose_dense_matmul_kernel[grid](
69 | coo_indices,
70 | coo_values,
71 | dense,
72 | out,
73 | stride_da=dense.stride(0),
74 | stride_db=dense.stride(1),
75 | B=B,
76 | N=N,
77 | AK=AK,
78 | BLOCK_SIZE_AK=BLOCK_SIZE_AK,
79 | BLOCK_SIZE_B=triton.next_power_of_2(B),
80 | )
81 | return out
82 |
83 |
84 | @triton.jit
85 | def triton_sparse_transpose_dense_matmul_kernel(
86 | coo_indices_ptr,
87 | coo_values_ptr,
88 | dense_ptr,
89 | out_ptr,
90 | stride_da,
91 | stride_db,
92 | B,
93 | N,
94 | AK,
95 | BLOCK_SIZE_AK: tl.constexpr,
96 | BLOCK_SIZE_B: tl.constexpr,
97 | ):
98 | """
99 | coo_indices is shape (2, AK)
100 | coo_values is shape (AK,)
101 | dense is shape (A, B), contiguous along B
102 | out is shape (N, B)
103 | """
104 |
105 | pid_ak = tl.program_id(0)
106 | pid_b = tl.program_id(1)
107 |
108 | coo_offsets = tl.arange(0, BLOCK_SIZE_AK)
109 | b_offsets = tl.arange(0, BLOCK_SIZE_B)
110 |
111 | A_coords = tl.load(
112 | coo_indices_ptr + pid_ak * BLOCK_SIZE_AK + coo_offsets,
113 | mask=pid_ak * BLOCK_SIZE_AK + coo_offsets < AK,
114 | )
115 | K_coords = tl.load(
116 | coo_indices_ptr + pid_ak * BLOCK_SIZE_AK + coo_offsets + AK,
117 | mask=pid_ak * BLOCK_SIZE_AK + coo_offsets < AK,
118 | )
119 | values = tl.load(
120 | coo_values_ptr + pid_ak * BLOCK_SIZE_AK + coo_offsets,
121 | mask=pid_ak * BLOCK_SIZE_AK + coo_offsets < AK,
122 | )
123 |
124 | last_k = tl.min(K_coords)
125 | accum = tl.zeros((BLOCK_SIZE_B,), dtype=tl.float32)
126 |
127 | for ind in range(BLOCK_SIZE_AK):
128 | if ind + pid_ak * BLOCK_SIZE_AK < AK:
129 | # workaround to do A_coords[ind]
130 | a = tl.sum(
131 | tl.where(
132 | tl.arange(0, BLOCK_SIZE_AK) == ind,
133 | A_coords,
134 | tl.zeros((BLOCK_SIZE_AK,), dtype=tl.int64),
135 | )
136 | )
137 |
138 | k = tl.sum(
139 | tl.where(
140 | tl.arange(0, BLOCK_SIZE_AK) == ind,
141 | K_coords,
142 | tl.zeros((BLOCK_SIZE_AK,), dtype=tl.int64),
143 | )
144 | )
145 |
146 | v = tl.sum(
147 | tl.where(
148 | tl.arange(0, BLOCK_SIZE_AK) == ind,
149 | values,
150 | tl.zeros((BLOCK_SIZE_AK,), dtype=tl.float32),
151 | )
152 | )
153 |
154 | tl.device_assert(k < N)
155 |
156 | if k != last_k:
157 | tl.atomic_add(
158 | out_ptr + last_k * B + BLOCK_SIZE_B * pid_b + b_offsets,
159 | accum,
160 | mask=BLOCK_SIZE_B * pid_b + b_offsets < B,
161 | )
162 | accum *= 0
163 | last_k = k
164 |
165 | if v != 0:
166 | accum += v * tl.load(dense_ptr + a * stride_da + b_offsets, mask=b_offsets < B)
167 |
168 | tl.atomic_add(
169 | out_ptr + last_k * B + BLOCK_SIZE_B * pid_b + b_offsets,
170 | accum,
171 | mask=BLOCK_SIZE_B * pid_b + b_offsets < B,
172 | )
173 |
174 |
175 | def triton_sparse_dense_matmul(
176 | sparse_indices: torch.Tensor,
177 | sparse_values: torch.Tensor,
178 | dense: torch.Tensor,
179 | ) -> torch.Tensor:
180 | """
181 | calculates sparse @ dense (i.e reducing along the uncollated dimension of sparse)
182 | dense must be contiguous along dim 0 (in other words, dense.T is contiguous)
183 |
184 | sparse_indices is shape (A, k)
185 | sparse_values is shape (A, k)
186 | dense is shape (N, B)
187 |
188 | output is shape (A, B)
189 | """
190 | N = dense.shape[0]
191 | assert sparse_indices.shape == sparse_values.shape
192 | assert sparse_indices.is_contiguous()
193 | assert sparse_values.is_contiguous()
194 | assert dense.is_contiguous() # contiguous along B
195 |
196 | A = sparse_indices.shape[0]
197 | K = sparse_indices.shape[1]
198 | B = dense.shape[1]
199 |
200 | out = torch.zeros(A, B, device=dense.device, dtype=sparse_values.dtype)
201 |
202 | triton_sparse_dense_matmul_kernel[(A,)](
203 | sparse_indices,
204 | sparse_values,
205 | dense,
206 | out,
207 | stride_dn=dense.stride(0),
208 | stride_db=dense.stride(1),
209 | A=A,
210 | B=B,
211 | N=N,
212 | K=K,
213 | BLOCK_SIZE_K=triton.next_power_of_2(K),
214 | BLOCK_SIZE_B=triton.next_power_of_2(B),
215 | )
216 | return out
217 |
218 |
219 | @triton.jit
220 | def triton_sparse_dense_matmul_kernel(
221 | sparse_indices_ptr,
222 | sparse_values_ptr,
223 | dense_ptr,
224 | out_ptr,
225 | stride_dn,
226 | stride_db,
227 | A,
228 | B,
229 | N,
230 | K,
231 | BLOCK_SIZE_K: tl.constexpr,
232 | BLOCK_SIZE_B: tl.constexpr,
233 | ):
234 | """
235 | sparse_indices is shape (A, K)
236 | sparse_values is shape (A, K)
237 | dense is shape (N, B), contiguous along B
238 | out is shape (A, B)
239 | """
240 |
241 | pid = tl.program_id(0)
242 |
243 | offsets_k = tl.arange(0, BLOCK_SIZE_K)
244 | sparse_indices = tl.load(
245 | sparse_indices_ptr + pid * K + offsets_k, mask=offsets_k < K
246 | ) # shape (K,)
247 | sparse_values = tl.load(
248 | sparse_values_ptr + pid * K + offsets_k, mask=offsets_k < K
249 | ) # shape (K,)
250 |
251 | accum = tl.zeros((BLOCK_SIZE_B,), dtype=tl.float32)
252 |
253 | offsets_b = tl.arange(0, BLOCK_SIZE_B)
254 |
255 | for k in range(K):
256 | # workaround to do sparse_indices[k]
257 | i = tl.sum(
258 | tl.where(
259 | tl.arange(0, BLOCK_SIZE_K) == k,
260 | sparse_indices,
261 | tl.zeros((BLOCK_SIZE_K,), dtype=tl.int64),
262 | )
263 | )
264 | # workaround to do sparse_values[k]
265 | v = tl.sum(
266 | tl.where(
267 | tl.arange(0, BLOCK_SIZE_K) == k,
268 | sparse_values,
269 | tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32),
270 | )
271 | )
272 |
273 | tl.device_assert(i < N)
274 | if v != 0:
275 | accum += v * tl.load(
276 | dense_ptr + i * stride_dn + offsets_b * stride_db, mask=offsets_b < B
277 | )
278 |
279 | tl.store(out_ptr + pid * B + offsets_b, accum.to(sparse_values.dtype), mask=offsets_b < B)
280 |
281 |
282 | def triton_dense_dense_sparseout_matmul(
283 | dense1: torch.Tensor,
284 | dense2: torch.Tensor,
285 | at_indices: torch.Tensor,
286 | ) -> torch.Tensor:
287 | """
288 | dense1: shape (A, B)
289 | dense2: shape (B, N)
290 | at_indices: shape (A, K)
291 | out values: shape (A, K)
292 | calculates dense1 @ dense2 only for the indices in at_indices
293 |
294 | equivalent to (dense1 @ dense2).gather(1, at_indices)
295 | """
296 | A, B = dense1.shape
297 | N = dense2.shape[1]
298 | assert dense2.shape[0] == B
299 | assert at_indices.shape[0] == A
300 | K = at_indices.shape[1]
301 | assert at_indices.is_contiguous()
302 |
303 | assert dense1.stride(1) == 1, "dense1 must be contiguous along B"
304 | assert dense2.stride(0) == 1, "dense2 must be contiguous along B"
305 |
306 | if K > 512:
307 | # print("WARN - using naive matmul for large K")
308 | # naive is more efficient for large K
309 | return (dense1 @ dense2).gather(1, at_indices)
310 |
311 | out = torch.zeros(A, K, device=dense1.device, dtype=dense1.dtype)
312 |
313 | # grid = lambda META: (triton.cdiv(A, META['BLOCK_SIZE_A']),)
314 |
315 | triton_dense_dense_sparseout_matmul_kernel[(A,)](
316 | dense1,
317 | dense2,
318 | at_indices,
319 | out,
320 | stride_d1a=dense1.stride(0),
321 | stride_d1b=dense1.stride(1),
322 | stride_d2b=dense2.stride(0),
323 | stride_d2n=dense2.stride(1),
324 | A=A,
325 | B=B,
326 | N=N,
327 | K=K,
328 | BLOCK_SIZE_B=triton.next_power_of_2(B),
329 | BLOCK_SIZE_N=triton.next_power_of_2(N),
330 | BLOCK_SIZE_K=triton.next_power_of_2(K),
331 | )
332 |
333 | return out
334 |
335 |
336 | @triton.jit
337 | def triton_dense_dense_sparseout_matmul_kernel(
338 | dense1_ptr,
339 | dense2_ptr,
340 | at_indices_ptr,
341 | out_ptr,
342 | stride_d1a,
343 | stride_d1b,
344 | stride_d2b,
345 | stride_d2n,
346 | A,
347 | B,
348 | N,
349 | K,
350 | BLOCK_SIZE_B: tl.constexpr,
351 | BLOCK_SIZE_N: tl.constexpr,
352 | BLOCK_SIZE_K: tl.constexpr,
353 | ):
354 | """
355 | dense1: shape (A, B)
356 | dense2: shape (B, N)
357 | at_indices: shape (A, K)
358 | out values: shape (A, K)
359 | """
360 |
361 | pid = tl.program_id(0)
362 |
363 | offsets_k = tl.arange(0, BLOCK_SIZE_K)
364 | at_indices = tl.load(at_indices_ptr + pid * K + offsets_k, mask=offsets_k < K) # shape (K,)
365 |
366 | offsets_b = tl.arange(0, BLOCK_SIZE_B)
367 | dense1 = tl.load(
368 | dense1_ptr + pid * stride_d1a + offsets_b * stride_d1b, mask=offsets_b < B
369 | ) # shape (B,)
370 |
371 | accum = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32)
372 |
373 | for k in range(K):
374 | # workaround to do at_indices[b]
375 | i = tl.sum(
376 | tl.where(
377 | tl.arange(0, BLOCK_SIZE_K) == k,
378 | at_indices,
379 | tl.zeros((BLOCK_SIZE_K,), dtype=tl.int64),
380 | )
381 | )
382 | tl.device_assert(i < N)
383 |
384 | dense2col = tl.load(
385 | dense2_ptr + offsets_b * stride_d2b + i * stride_d2n, mask=offsets_b < B
386 | ) # shape (B,)
387 | accum += tl.where(
388 | tl.arange(0, BLOCK_SIZE_K) == k,
389 | tl.sum(dense1 * dense2col),
390 | tl.zeros((BLOCK_SIZE_K,), dtype=tl.int64),
391 | )
392 |
393 | tl.store(out_ptr + pid * K + offsets_k, accum, mask=offsets_k < K)
394 |
395 |
396 | class TritonDecoderAutograd(torch.autograd.Function):
397 | @staticmethod
398 | def forward(ctx, sparse_indices, sparse_values, decoder_weight):
399 | ctx.save_for_backward(sparse_indices, sparse_values, decoder_weight)
400 | return triton_sparse_dense_matmul(sparse_indices, sparse_values, decoder_weight.T)
401 |
402 | @staticmethod
403 | def backward(ctx, grad_output):
404 | sparse_indices, sparse_values, decoder_weight = ctx.saved_tensors
405 |
406 | assert grad_output.is_contiguous(), "grad_output must be contiguous; this is probably because the subsequent op was a .sum() or something like that, which returns a non contiguous gradient"
407 |
408 | decoder_grad = triton_sparse_transpose_dense_matmul(
409 | sparse_indices, sparse_values, grad_output, N=decoder_weight.shape[1]
410 | ).T
411 |
412 | return (
413 | None,
414 | triton_dense_dense_sparseout_matmul(grad_output, decoder_weight, sparse_indices),
415 | # decoder is contiguous when transposed so this is a matching layout
416 | decoder_grad,
417 | None,
418 | )
419 |
420 |
421 | def triton_add_mul_(
422 | x: torch.Tensor,
423 | a: torch.Tensor,
424 | b: torch.Tensor,
425 | c: float,
426 | ):
427 | """
428 | does
429 | x += a * b * c
430 |
431 | x : [m, n]
432 | a : [m, n]
433 | b : [m, n]
434 | c : float
435 | """
436 |
437 | if len(a.shape) == 1:
438 | a = a[None, :].broadcast_to(x.shape)
439 |
440 | if len(b.shape) == 1:
441 | b = b[None, :].broadcast_to(x.shape)
442 |
443 | assert x.shape == a.shape == b.shape
444 |
445 | BLOCK_SIZE_M = 64
446 | BLOCK_SIZE_N = 64
447 | grid = lambda META: (
448 | triton.cdiv(x.shape[0], META["BLOCK_SIZE_M"]),
449 | triton.cdiv(x.shape[1], META["BLOCK_SIZE_N"]),
450 | )
451 | triton_add_mul_kernel[grid](
452 | x,
453 | a,
454 | b,
455 | c,
456 | x.stride(0),
457 | x.stride(1),
458 | a.stride(0),
459 | a.stride(1),
460 | b.stride(0),
461 | b.stride(1),
462 | BLOCK_SIZE_M,
463 | BLOCK_SIZE_N,
464 | x.shape[0],
465 | x.shape[1],
466 | )
467 |
468 |
469 | @triton.jit
470 | def triton_add_mul_kernel(
471 | x_ptr,
472 | a_ptr,
473 | b_ptr,
474 | c,
475 | stride_x0,
476 | stride_x1,
477 | stride_a0,
478 | stride_a1,
479 | stride_b0,
480 | stride_b1,
481 | BLOCK_SIZE_M: tl.constexpr,
482 | BLOCK_SIZE_N: tl.constexpr,
483 | M: tl.constexpr,
484 | N: tl.constexpr,
485 | ):
486 | pid_m = tl.program_id(0)
487 | pid_n = tl.program_id(1)
488 |
489 | offsets_m = tl.arange(0, BLOCK_SIZE_M) + pid_m * BLOCK_SIZE_M
490 | offsets_n = tl.arange(0, BLOCK_SIZE_N) + pid_n * BLOCK_SIZE_N
491 |
492 | x = tl.load(
493 | x_ptr + offsets_m[:, None] * stride_x0 + offsets_n[None, :] * stride_x1,
494 | mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N),
495 | )
496 | a = tl.load(
497 | a_ptr + offsets_m[:, None] * stride_a0 + offsets_n[None, :] * stride_a1,
498 | mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N),
499 | )
500 | b = tl.load(
501 | b_ptr + offsets_m[:, None] * stride_b0 + offsets_n[None, :] * stride_b1,
502 | mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N),
503 | )
504 |
505 | x_dtype = x.dtype
506 | x = (x.to(tl.float32) + a.to(tl.float32) * b.to(tl.float32) * c).to(x_dtype)
507 |
508 | tl.store(
509 | x_ptr + offsets_m[:, None] * stride_x0 + offsets_n[None, :] * stride_x1,
510 | x,
511 | mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N),
512 | )
513 |
514 |
515 |
516 | def triton_sum_dim0_in_fp32(xs):
517 | a, b = xs.shape
518 |
519 | assert xs.is_contiguous()
520 | assert xs.dtype == torch.float16
521 |
522 | BLOCK_SIZE_A = min(triton.next_power_of_2(a), 512)
523 | BLOCK_SIZE_B = 64 # cache line is 128 bytes
524 |
525 | out = torch.zeros(b, dtype=torch.float32, device=xs.device)
526 |
527 | grid = lambda META: (triton.cdiv(b, META["BLOCK_SIZE_B"]),)
528 |
529 | triton_sum_dim0_in_fp32_kernel[grid](
530 | xs,
531 | out,
532 | stride_a=xs.stride(0),
533 | a=a,
534 | b=b,
535 | BLOCK_SIZE_A=BLOCK_SIZE_A,
536 | BLOCK_SIZE_B=BLOCK_SIZE_B,
537 | )
538 |
539 | return out
540 |
541 |
542 | @triton.jit
543 | def triton_sum_dim0_in_fp32_kernel(
544 | xs_ptr,
545 | out_ptr,
546 | stride_a,
547 | a,
548 | b,
549 | BLOCK_SIZE_A: tl.constexpr,
550 | BLOCK_SIZE_B: tl.constexpr,
551 | ):
552 | # each program handles 64 columns of xs
553 | pid = tl.program_id(0)
554 | offsets_b = tl.arange(0, BLOCK_SIZE_B) + pid * BLOCK_SIZE_B
555 |
556 | all_out = tl.zeros((BLOCK_SIZE_B,), dtype=tl.float32)
557 |
558 | for i in range(0, a, BLOCK_SIZE_A):
559 | offsets_a = tl.arange(0, BLOCK_SIZE_A) + i
560 | xs = tl.load(
561 | xs_ptr + offsets_a[:, None] * stride_a + offsets_b[None, :],
562 | mask=(offsets_a < a)[:, None] & (offsets_b < b)[None, :],
563 | other=0,
564 | )
565 | xs = xs.to(tl.float32)
566 | out = tl.sum(xs, axis=0)
567 | all_out += out
568 |
569 | tl.store(out_ptr + offsets_b, all_out, mask=offsets_b < b)
570 |
571 |
572 | def mse(
573 | output,
574 | target,
575 | ): # fusing fp32 cast and MSE to save memory
576 | assert output.shape == target.shape
577 | assert len(output.shape) == 2
578 | assert output.stride(1) == 1
579 | assert target.stride(1) == 1
580 |
581 | a, b = output.shape
582 |
583 | BLOCK_SIZE_B = triton.next_power_of_2(b)
584 |
585 | class _MSE(torch.autograd.Function):
586 | @staticmethod
587 | def forward(ctx, output, target):
588 | ctx.save_for_backward(output, target)
589 | out = torch.zeros(a, dtype=torch.float32, device=output.device)
590 |
591 | triton_mse_loss_fp16_kernel[(a,)](
592 | output,
593 | target,
594 | out,
595 | stride_a_output=output.stride(0),
596 | stride_a_target=target.stride(0),
597 | a=a,
598 | b=b,
599 | BLOCK_SIZE_B=BLOCK_SIZE_B,
600 | )
601 |
602 | return out
603 |
604 | @staticmethod
605 | def backward(ctx, grad_output):
606 | output, target = ctx.saved_tensors
607 | res = (output - target).float()
608 | res *= grad_output[:, None] * 2 / b
609 | return res, None
610 |
611 | return _MSE.apply(output, target).mean()
612 |
613 |
614 | def normalized_mse(recon: torch.Tensor, xs: torch.Tensor) -> torch.Tensor:
615 | # only used for auxk
616 | xs_mu = (
617 | triton_sum_dim0_in_fp32(xs) / xs.shape[0]
618 | if xs.dtype == torch.float16
619 | else xs.mean(dim=0)
620 | )
621 |
622 | loss = mse(recon, xs) / mse(
623 | xs_mu[None, :].broadcast_to(xs.shape), xs
624 | )
625 |
626 | return loss
627 |
628 |
629 | @triton.jit
630 | def triton_mse_loss_fp16_kernel(
631 | output_ptr,
632 | target_ptr,
633 | out_ptr,
634 | stride_a_output,
635 | stride_a_target,
636 | a,
637 | b,
638 | BLOCK_SIZE_B: tl.constexpr,
639 | ):
640 | pid = tl.program_id(0)
641 | offsets_b = tl.arange(0, BLOCK_SIZE_B)
642 |
643 | output = tl.load(
644 | output_ptr + pid * stride_a_output + offsets_b,
645 | mask=offsets_b < b,
646 | )
647 | target = tl.load(
648 | target_ptr + pid * stride_a_target + offsets_b,
649 | mask=offsets_b < b,
650 | )
651 |
652 | output = output.to(tl.float32)
653 | target = target.to(tl.float32)
654 |
655 | mse = tl.sum((output - target) * (output - target)) / b
656 |
657 | tl.store(out_ptr + pid, mse)
658 |
659 |
660 | def triton_add_mul_(
661 | x: torch.Tensor,
662 | a: torch.Tensor,
663 | b: torch.Tensor,
664 | c: float,
665 | ):
666 | """
667 | does
668 | x += a * b * c
669 |
670 | x : [m, n]
671 | a : [m, n]
672 | b : [m, n]
673 | c : float
674 | """
675 |
676 | if len(a.shape) == 1:
677 | a = a[None, :].broadcast_to(x.shape)
678 |
679 | if len(b.shape) == 1:
680 | b = b[None, :].broadcast_to(x.shape)
681 |
682 | assert x.shape == a.shape == b.shape
683 |
684 | BLOCK_SIZE_M = 64
685 | BLOCK_SIZE_N = 64
686 | grid = lambda META: (
687 | triton.cdiv(x.shape[0], META["BLOCK_SIZE_M"]),
688 | triton.cdiv(x.shape[1], META["BLOCK_SIZE_N"]),
689 | )
690 | triton_add_mul_kernel[grid](
691 | x,
692 | a,
693 | b,
694 | c,
695 | x.stride(0),
696 | x.stride(1),
697 | a.stride(0),
698 | a.stride(1),
699 | b.stride(0),
700 | b.stride(1),
701 | BLOCK_SIZE_M,
702 | BLOCK_SIZE_N,
703 | x.shape[0],
704 | x.shape[1],
705 | )
706 |
--------------------------------------------------------------------------------
/sparse_autoencoder/train.py:
--------------------------------------------------------------------------------
1 | # bare bones training script using sparse kernels and sharding/data parallel.
2 | # the main purpose of this code is to provide a reference implementation to compare
3 | # against when implementing our training methodology into other codebases, and to
4 | # demonstrate how sharding/DP can be implemented for autoencoders. some limitations:
5 | # - many basic features (e.g checkpointing, data loading, validation) are not implemented,
6 | # - the codebase is not designed to be extensible or easily hackable.
7 | # - this code is not guaranteed to run efficiently out of the box / in
8 | # combination with other changes, so you should profile it and make changes as needed.
9 | #
10 | # example launch command:
11 | # torchrun --nproc-per-node 8 train.py
12 |
13 |
14 | import os
15 | from dataclasses import dataclass
16 | from typing import Callable, Iterable, Iterator
17 |
18 | import torch
19 | import torch.distributed as dist
20 | import torch.nn as nn
21 | import torch.nn.functional as F
22 | import triton
23 | import triton.language as tl
24 | from sparse_autoencoder.kernels import *
25 | from torch.distributed import ReduceOp
26 |
27 | RANK = int(os.environ.get("RANK", "0"))
28 |
29 |
30 | ## parallelism
31 |
32 |
33 | @dataclass
34 | class Comm:
35 | group: torch.distributed.ProcessGroup
36 |
37 | def all_reduce(self, x, op=ReduceOp.SUM, async_op=False):
38 | return dist.all_reduce(x, op=op, group=self.group, async_op=async_op)
39 |
40 | def all_gather(self, x_list, x, async_op=False):
41 | return dist.all_gather(list(x_list), x, group=self.group, async_op=async_op)
42 |
43 | def broadcast(self, x, src, async_op=False):
44 | return dist.broadcast(x, src, group=self.group, async_op=async_op)
45 |
46 | def barrier(self):
47 | return dist.barrier(group=self.group)
48 |
49 | def size(self):
50 | return self.group.size()
51 |
52 |
53 | @dataclass
54 | class ShardingComms:
55 | n_replicas: int
56 | n_op_shards: int
57 | dp_rank: int
58 | sh_rank: int
59 | dp_comm: Comm | None
60 | sh_comm: Comm | None
61 | _rank: int
62 |
63 | def sh_allreduce_forward(self, x: torch.Tensor) -> torch.Tensor:
64 | if self.sh_comm is None:
65 | return x
66 |
67 | class AllreduceForward(torch.autograd.Function):
68 | @staticmethod
69 | def forward(ctx, input):
70 | assert self.sh_comm is not None
71 | self.sh_comm.all_reduce(input, async_op=True)
72 | return input
73 |
74 | @staticmethod
75 | def backward(ctx, grad_output):
76 | return grad_output
77 |
78 | return AllreduceForward.apply(x) # type: ignore
79 |
80 | def sh_allreduce_backward(self, x: torch.Tensor) -> torch.Tensor:
81 | if self.sh_comm is None:
82 | return x
83 |
84 | class AllreduceBackward(torch.autograd.Function):
85 | @staticmethod
86 | def forward(ctx, input):
87 | return input
88 |
89 | @staticmethod
90 | def backward(ctx, grad_output):
91 | grad_output = grad_output.clone()
92 | assert self.sh_comm is not None
93 | self.sh_comm.all_reduce(grad_output, async_op=True)
94 | return grad_output
95 |
96 | return AllreduceBackward.apply(x) # type: ignore
97 |
98 | def init_broadcast_(self, autoencoder):
99 | if self.dp_comm is not None:
100 | for p in autoencoder.parameters():
101 | self.dp_comm.broadcast(
102 | maybe_transpose(p.data),
103 | replica_shard_to_rank(
104 | replica_idx=0,
105 | shard_idx=self.sh_rank,
106 | n_op_shards=self.n_op_shards,
107 | ),
108 | )
109 |
110 | if self.sh_comm is not None:
111 | # pre_bias is the same across all shards
112 | self.sh_comm.broadcast(
113 | autoencoder.pre_bias.data,
114 | replica_shard_to_rank(
115 | replica_idx=self.dp_rank,
116 | shard_idx=0,
117 | n_op_shards=self.n_op_shards,
118 | ),
119 | )
120 |
121 | def dp_allreduce_(self, autoencoder) -> None:
122 | if self.dp_comm is None:
123 | return
124 |
125 | for param in autoencoder.parameters():
126 | if param.grad is not None:
127 | self.dp_comm.all_reduce(maybe_transpose(param.grad), op=ReduceOp.AVG, async_op=True)
128 |
129 | # make sure statistics for dead neurons are correct
130 | self.dp_comm.all_reduce( # type: ignore
131 | autoencoder.stats_last_nonzero, op=ReduceOp.MIN, async_op=True
132 | )
133 |
134 | def sh_allreduce_scale(self, scaler):
135 | if self.sh_comm is None:
136 | return
137 |
138 | if hasattr(scaler, "_scale") and scaler._scale is not None:
139 | self.sh_comm.all_reduce(scaler._scale, op=ReduceOp.MIN, async_op=True)
140 | self.sh_comm.all_reduce(scaler._growth_tracker, op=ReduceOp.MIN, async_op=True)
141 |
142 | def _sh_comm_op(self, x, op):
143 | if isinstance(x, (float, int)):
144 | x = torch.tensor(x, device="cuda")
145 |
146 | if not x.is_cuda:
147 | x = x.cuda()
148 |
149 | if self.sh_comm is None:
150 | return x
151 |
152 | out = x.clone()
153 | self.sh_comm.all_reduce(x, op=op, async_op=True)
154 | return out
155 |
156 | def sh_sum(self, x: torch.Tensor) -> torch.Tensor:
157 | return self._sh_comm_op(x, ReduceOp.SUM)
158 |
159 | def all_broadcast(self, x: torch.Tensor) -> torch.Tensor:
160 | if self.dp_comm is not None:
161 | self.dp_comm.broadcast(
162 | x,
163 | replica_shard_to_rank(
164 | replica_idx=0,
165 | shard_idx=self.sh_rank,
166 | n_op_shards=self.n_op_shards,
167 | ),
168 | )
169 |
170 | if self.sh_comm is not None:
171 | self.sh_comm.broadcast(
172 | x,
173 | replica_shard_to_rank(
174 | replica_idx=self.dp_rank,
175 | shard_idx=0,
176 | n_op_shards=self.n_op_shards,
177 | ),
178 | )
179 |
180 | return x
181 |
182 |
183 | def make_torch_comms(n_op_shards=4, n_replicas=2):
184 | if "RANK" not in os.environ:
185 | assert n_op_shards == 1
186 | assert n_replicas == 1
187 | return TRIVIAL_COMMS
188 |
189 | rank = int(os.environ.get("RANK"))
190 | world_size = int(os.environ.get("WORLD_SIZE", 1))
191 | os.environ["CUDA_VISIBLE_DEVICES"] = str(rank % 8)
192 |
193 | print(f"{rank=}, {world_size=}")
194 | dist.init_process_group("nccl")
195 |
196 | my_op_shard_idx = rank % n_op_shards
197 | my_replica_idx = rank // n_op_shards
198 |
199 | shard_rank_lists = [list(range(i, i + n_op_shards)) for i in range(0, world_size, n_op_shards)]
200 |
201 | shard_groups = [dist.new_group(shard_rank_list) for shard_rank_list in shard_rank_lists]
202 |
203 | my_shard_group = shard_groups[my_replica_idx]
204 |
205 | replica_rank_lists = [
206 | list(range(i, n_op_shards * n_replicas, n_op_shards)) for i in range(n_op_shards)
207 | ]
208 |
209 | replica_groups = [dist.new_group(replica_rank_list) for replica_rank_list in replica_rank_lists]
210 |
211 | my_replica_group = replica_groups[my_op_shard_idx]
212 |
213 | torch.distributed.all_reduce(torch.ones(1).cuda())
214 | torch.cuda.synchronize()
215 |
216 | dp_comm = Comm(group=my_replica_group)
217 | sh_comm = Comm(group=my_shard_group)
218 |
219 | return ShardingComms(
220 | n_replicas=n_replicas,
221 | n_op_shards=n_op_shards,
222 | dp_comm=dp_comm,
223 | sh_comm=sh_comm,
224 | dp_rank=my_replica_idx,
225 | sh_rank=my_op_shard_idx,
226 | _rank=rank,
227 | )
228 |
229 |
230 | def replica_shard_to_rank(replica_idx, shard_idx, n_op_shards):
231 | return replica_idx * n_op_shards + shard_idx
232 |
233 |
234 | TRIVIAL_COMMS = ShardingComms(
235 | n_replicas=1,
236 | n_op_shards=1,
237 | dp_rank=0,
238 | sh_rank=0,
239 | dp_comm=None,
240 | sh_comm=None,
241 | _rank=0,
242 | )
243 |
244 |
245 | def sharded_topk(x, k, sh_comm, capacity_factor=None):
246 | batch = x.shape[0]
247 |
248 | if capacity_factor is not None:
249 | k_in = min(int(k * capacity_factor // sh_comm.size()), k)
250 | else:
251 | k_in = k
252 |
253 | topk = torch.topk(x, k=k_in, dim=-1)
254 | inds = topk.indices
255 | vals = topk.values
256 |
257 | if sh_comm is None:
258 | return inds, vals
259 |
260 | all_vals = torch.empty(sh_comm.size(), batch, k_in, dtype=vals.dtype, device=vals.device)
261 | sh_comm.all_gather(all_vals, vals, async_op=True)
262 |
263 | all_vals = all_vals.permute(1, 0, 2) # put shard dim next to k
264 | all_vals = all_vals.reshape(batch, -1) # flatten shard into k
265 |
266 | all_topk = torch.topk(all_vals, k=k, dim=-1)
267 | global_topk = all_topk.values
268 |
269 | dummy_vals = torch.zeros_like(vals)
270 | dummy_inds = torch.zeros_like(inds)
271 |
272 | my_inds = torch.where(vals >= global_topk[:, [-1]], inds, dummy_inds)
273 | my_vals = torch.where(vals >= global_topk[:, [-1]], vals, dummy_vals)
274 |
275 | return my_inds, my_vals
276 |
277 |
278 | ## autoencoder
279 |
280 |
281 | class FastAutoencoder(nn.Module):
282 | """
283 | Top-K Autoencoder with sparse kernels. Implements:
284 |
285 | latents = relu(topk(encoder(x - pre_bias) + latent_bias))
286 | recons = decoder(latents) + pre_bias
287 | """
288 |
289 | def __init__(
290 | self,
291 | n_dirs_local: int,
292 | d_model: int,
293 | k: int,
294 | auxk: int | None,
295 | dead_steps_threshold: int,
296 | comms: ShardingComms | None = None,
297 | ):
298 | super().__init__()
299 | self.n_dirs_local = n_dirs_local
300 | self.d_model = d_model
301 | self.k = k
302 | self.auxk = auxk
303 | self.comms = comms if comms is not None else TRIVIAL_COMMS
304 | self.dead_steps_threshold = dead_steps_threshold
305 |
306 | self.encoder = nn.Linear(d_model, n_dirs_local, bias=False)
307 | self.decoder = nn.Linear(n_dirs_local, d_model, bias=False)
308 |
309 | self.pre_bias = nn.Parameter(torch.zeros(d_model))
310 | self.latent_bias = nn.Parameter(torch.zeros(n_dirs_local))
311 |
312 | self.stats_last_nonzero: torch.Tensor
313 | self.register_buffer("stats_last_nonzero", torch.zeros(n_dirs_local, dtype=torch.long))
314 |
315 | def auxk_mask_fn(x):
316 | dead_mask = self.stats_last_nonzero > dead_steps_threshold
317 | x.data *= dead_mask # inplace to save memory
318 | return x
319 |
320 | self.auxk_mask_fn = auxk_mask_fn
321 |
322 | ## initialization
323 |
324 | # "tied" init
325 | self.decoder.weight.data = self.encoder.weight.data.T.clone()
326 |
327 | # store decoder in column major layout for kernel
328 | self.decoder.weight.data = self.decoder.weight.data.T.contiguous().T
329 |
330 | unit_norm_decoder_(self)
331 |
332 | @property
333 | def n_dirs(self):
334 | return self.n_dirs_local * self.comms.n_op_shards
335 |
336 | def forward(self, x):
337 | class EncWrapper(torch.autograd.Function):
338 | @staticmethod
339 | def forward(ctx, x, pre_bias, weight, latent_bias):
340 | x = x - pre_bias
341 | latents_pre_act = F.linear(x, weight, latent_bias)
342 |
343 | inds, vals = sharded_topk(
344 | latents_pre_act,
345 | k=self.k,
346 | sh_comm=self.comms.sh_comm,
347 | capacity_factor=4,
348 | )
349 |
350 | ## set num nonzero stat ##
351 | tmp = torch.zeros_like(self.stats_last_nonzero)
352 | tmp.scatter_add_(
353 | 0,
354 | inds.reshape(-1),
355 | (vals > 1e-3).to(tmp.dtype).reshape(-1),
356 | )
357 | self.stats_last_nonzero *= 1 - tmp.clamp(max=1)
358 | self.stats_last_nonzero += 1
359 | ## end stats ##
360 |
361 | ## auxk
362 | if self.auxk is not None: # for auxk
363 | # IMPORTANT: has to go after stats update!
364 | # WARN: auxk_mask_fn can mutate latents_pre_act!
365 | auxk_inds, auxk_vals = sharded_topk(
366 | self.auxk_mask_fn(latents_pre_act),
367 | k=self.auxk,
368 | sh_comm=self.comms.sh_comm,
369 | capacity_factor=2,
370 | )
371 | ctx.save_for_backward(x, weight, inds, auxk_inds)
372 | else:
373 | ctx.save_for_backward(x, weight, inds)
374 | auxk_inds = None
375 | auxk_vals = None
376 |
377 | ## end auxk
378 |
379 | return (
380 | inds,
381 | vals,
382 | auxk_inds,
383 | auxk_vals,
384 | )
385 |
386 | @staticmethod
387 | def backward(ctx, _, grad_vals, __, grad_auxk_vals):
388 | # encoder backwards
389 | if self.auxk is not None:
390 | x, weight, inds, auxk_inds = ctx.saved_tensors
391 |
392 | all_inds = torch.cat((inds, auxk_inds), dim=-1)
393 | all_grad_vals = torch.cat((grad_vals, grad_auxk_vals), dim=-1)
394 | else:
395 | x, weight, inds = ctx.saved_tensors
396 |
397 | all_inds = inds
398 | all_grad_vals = grad_vals
399 |
400 | grad_sum = torch.zeros(self.n_dirs_local, dtype=torch.float32, device=grad_vals.device)
401 | grad_sum.scatter_add_(
402 | -1, all_inds.flatten(), all_grad_vals.flatten().to(torch.float32)
403 | )
404 |
405 | return (
406 | None,
407 | # pre_bias grad optimization - can reduce before mat-vec multiply
408 | -(grad_sum @ weight),
409 | triton_sparse_transpose_dense_matmul(all_inds, all_grad_vals, x, N=self.n_dirs_local),
410 | grad_sum,
411 | )
412 |
413 | pre_bias = self.comms.sh_allreduce_backward(self.pre_bias)
414 |
415 | # encoder
416 | inds, vals, auxk_inds, auxk_vals = EncWrapper.apply(
417 | x, pre_bias, self.encoder.weight, self.latent_bias
418 | )
419 |
420 | vals = torch.relu(vals)
421 | if auxk_vals is not None:
422 | auxk_vals = torch.relu(auxk_vals)
423 |
424 | recons = self.decode_sparse(inds, vals)
425 |
426 | return recons, {
427 | "auxk_inds": auxk_inds,
428 | "auxk_vals": auxk_vals,
429 | }
430 |
431 | def decode_sparse(self, inds, vals):
432 | recons = TritonDecoderAutograd.apply(inds, vals, self.decoder.weight)
433 | recons = self.comms.sh_allreduce_forward(recons)
434 |
435 | return recons + self.pre_bias
436 |
437 |
438 | def unit_norm_decoder_(autoencoder: FastAutoencoder) -> None:
439 | """
440 | Unit normalize the decoder weights of an autoencoder.
441 | """
442 | autoencoder.decoder.weight.data /= autoencoder.decoder.weight.data.norm(dim=0)
443 |
444 |
445 | def unit_norm_decoder_grad_adjustment_(autoencoder) -> None:
446 | """project out gradient information parallel to the dictionary vectors - assumes that the decoder is already unit normed"""
447 |
448 | assert autoencoder.decoder.weight.grad is not None
449 |
450 | triton_add_mul_(
451 | autoencoder.decoder.weight.grad,
452 | torch.einsum("bn,bn->n", autoencoder.decoder.weight.data, autoencoder.decoder.weight.grad),
453 | autoencoder.decoder.weight.data,
454 | c=-1,
455 | )
456 |
457 |
458 | def maybe_transpose(x):
459 | return x.T if not x.is_contiguous() and x.T.is_contiguous() else x
460 |
461 |
462 | def sharded_grad_norm(autoencoder, comms, exclude=None):
463 | if exclude is None:
464 | exclude = []
465 | total_sq_norm = torch.zeros((), device="cuda", dtype=torch.float32)
466 | exclude = set(exclude)
467 |
468 | total_num_params = 0
469 | for param in autoencoder.parameters():
470 | if param in exclude:
471 | continue
472 | if param.grad is not None:
473 | sq_norm = ((param.grad).float() ** 2).sum()
474 | if param is autoencoder.pre_bias:
475 | total_sq_norm += sq_norm # pre_bias is the same across all shards
476 | else:
477 | total_sq_norm += comms.sh_sum(sq_norm)
478 |
479 | param_shards = comms.n_op_shards if param is autoencoder.pre_bias else 1
480 | total_num_params += param.numel() * param_shards
481 |
482 | return total_sq_norm.sqrt()
483 |
484 |
485 | def batch_tensors(
486 | it: Iterable[torch.Tensor],
487 | batch_size: int,
488 | drop_last=True,
489 | stream=None,
490 | ) -> Iterator[torch.Tensor]:
491 | """
492 | input is iterable of tensors of shape [batch_old, ...]
493 | output is iterable of tensors of shape [batch_size, ...]
494 | batch_old does not need to be divisible by batch_size
495 | """
496 |
497 | tensors = []
498 | batch_so_far = 0
499 |
500 | for t in it:
501 | tensors.append(t)
502 | batch_so_far += t.shape[0]
503 |
504 | if sum(t.shape[0] for t in tensors) < batch_size:
505 | continue
506 |
507 | while batch_so_far >= batch_size:
508 | if len(tensors) == 1:
509 | (concat,) = tensors
510 | else:
511 | with torch.cuda.stream(stream):
512 | concat = torch.cat(tensors, dim=0)
513 |
514 | offset = 0
515 | while offset + batch_size <= concat.shape[0]:
516 | yield concat[offset : offset + batch_size]
517 | batch_so_far -= batch_size
518 | offset += batch_size
519 |
520 | tensors = [concat[offset:]] if offset < concat.shape[0] else []
521 |
522 | if len(tensors) > 0 and not drop_last:
523 | yield torch.cat(tensors, dim=0)
524 |
525 |
526 | def print0(*a, **k):
527 | if RANK == 0:
528 | print(*a, **k)
529 |
530 |
531 | import wandb
532 |
533 |
534 | class Logger:
535 | def __init__(self, **kws):
536 | self.vals = {}
537 | self.enabled = (RANK == 0) and not kws.pop("dummy", False)
538 | if self.enabled:
539 | wandb.init(
540 | **kws
541 | )
542 |
543 | def logkv(self, k, v):
544 | if self.enabled:
545 | self.vals[k] = v.detach() if isinstance(v, torch.Tensor) else v
546 | return v
547 |
548 | def dumpkvs(self):
549 | if self.enabled:
550 | wandb.log(self.vals)
551 | self.vals = {}
552 |
553 |
554 | def training_loop_(
555 | ae, train_acts_iter, loss_fn, lr, comms, eps=6.25e-10, clip_grad=None, ema_multiplier=0.999, logger=None
556 | ):
557 | if logger is None:
558 | logger = Logger(dummy=True)
559 |
560 | scaler = torch.cuda.amp.GradScaler()
561 | autocast_ctx_manager = torch.cuda.amp.autocast()
562 |
563 | opt = torch.optim.Adam(ae.parameters(), lr=lr, eps=eps, fused=True)
564 | if ema_multiplier is not None:
565 | ema = EmaModel(ae, ema_multiplier=ema_multiplier)
566 |
567 | for i, flat_acts_train_batch in enumerate(train_acts_iter):
568 | flat_acts_train_batch = flat_acts_train_batch.cuda()
569 |
570 | with autocast_ctx_manager:
571 | recons, info = ae(flat_acts_train_batch)
572 |
573 | loss = loss_fn(ae, flat_acts_train_batch, recons, info, logger)
574 |
575 | print0(i, loss)
576 |
577 | logger.logkv("loss_scale", scaler.get_scale())
578 |
579 | if RANK == 0:
580 | wandb.log({"train_loss": loss.item()})
581 |
582 | loss = scaler.scale(loss)
583 | loss.backward()
584 |
585 | unit_norm_decoder_(ae)
586 | unit_norm_decoder_grad_adjustment_(ae)
587 |
588 | # allreduce gradients
589 | comms.dp_allreduce_(ae)
590 |
591 | # keep fp16 loss scale synchronized across shards
592 | comms.sh_allreduce_scale(scaler)
593 |
594 | # if you want to do anything with the gradients that depends on the absolute scale (e.g clipping, do it after the unscale_)
595 | scaler.unscale_(opt)
596 |
597 | # gradient clipping
598 | if clip_grad is not None:
599 | grad_norm = sharded_grad_norm(ae, comms)
600 | logger.logkv("grad_norm", grad_norm)
601 | grads = [x.grad for x in ae.parameters() if x.grad is not None]
602 | torch._foreach_mul_(grads, clip_grad / torch.clamp(grad_norm, min=clip_grad))
603 |
604 | if ema_multiplier is not None:
605 | ema.step()
606 |
607 | # take step with optimizer
608 | scaler.step(opt)
609 | scaler.update()
610 |
611 | logger.dumpkvs()
612 |
613 |
614 | def init_from_data_(ae, stats_acts_sample, comms):
615 | from geom_median.torch import compute_geometric_median
616 |
617 | ae.pre_bias.data = (
618 | compute_geometric_median(stats_acts_sample[:32768].float().cpu()).median.cuda().float()
619 | )
620 | comms.all_broadcast(ae.pre_bias.data)
621 |
622 | # encoder initialization (note: in our ablations we couldn't find clear evidence that this is beneficial, this is just to ensure exact match with internal codebase)
623 | d_model = ae.d_model
624 | with torch.no_grad():
625 | x = torch.randn(256, d_model).cuda().to(stats_acts_sample.dtype)
626 | x /= x.norm(dim=-1, keepdim=True)
627 | x += ae.pre_bias.data
628 | comms.all_broadcast(x)
629 | recons, _ = ae(x)
630 | recons_norm = (recons - ae.pre_bias.data).norm(dim=-1).mean()
631 |
632 | ae.encoder.weight.data /= recons_norm.item()
633 | print0("x norm", x.norm(dim=-1).mean().item())
634 | print0("out norm", (ae(x)[0] - ae.pre_bias.data).norm(dim=-1).mean().item())
635 |
636 |
637 | from contextlib import contextmanager
638 |
639 |
640 | @contextmanager
641 | def temporary_weight_swap(model: torch.nn.Module, new_weights: list[torch.Tensor]):
642 | for _p, new_p in zip(model.parameters(), new_weights, strict=True):
643 | assert _p.shape == new_p.shape
644 | _p.data, new_p.data = new_p.data, _p.data
645 |
646 | yield
647 |
648 | for _p, new_p in zip(model.parameters(), new_weights, strict=True):
649 | assert _p.shape == new_p.shape
650 | _p.data, new_p.data = new_p.data, _p.data
651 |
652 |
653 | class EmaModel:
654 | def __init__(self, model, ema_multiplier):
655 | self.model = model
656 | self.ema_multiplier = ema_multiplier
657 | self.ema_weights = [torch.zeros_like(x, requires_grad=False) for x in model.parameters()]
658 | self.ema_steps = 0
659 |
660 | def step(self):
661 | torch._foreach_lerp_(
662 | self.ema_weights,
663 | list(self.model.parameters()),
664 | 1 - self.ema_multiplier,
665 | )
666 | self.ema_steps += 1
667 |
668 | # context manager for setting the autoencoder weights to the EMA weights
669 | @contextmanager
670 | def use_ema_weights(self):
671 | assert self.ema_steps > 0
672 |
673 | # apply bias correction
674 | bias_correction = 1 - self.ema_multiplier**self.ema_steps
675 | ema_weights_bias_corrected = torch._foreach_div(self.ema_weights, bias_correction)
676 |
677 | with torch.no_grad():
678 | with temporary_weight_swap(self.model, ema_weights_bias_corrected):
679 | yield
680 |
681 |
682 | @dataclass
683 | class Config:
684 | n_op_shards: int = 1
685 | n_replicas: int = 8
686 |
687 | n_dirs: int = 32768
688 | bs: int = 131072
689 | d_model: int = 768
690 | k: int = 32
691 | auxk: int = 256
692 |
693 | lr: float = 1e-4
694 | eps: float = 6.25e-10
695 | clip_grad: float | None = None
696 | auxk_coef: float = 1 / 32
697 | dead_toks_threshold: int = 10_000_000
698 | ema_multiplier: float | None = None
699 |
700 | wandb_project: str | None = None
701 | wandb_name: str | None = None
702 |
703 |
704 | def main():
705 | cfg = Config()
706 | comms = make_torch_comms(n_op_shards=cfg.n_op_shards, n_replicas=cfg.n_replicas)
707 |
708 | ## dataloading is left as an exercise for the reader
709 | acts_iter = ...
710 | stats_acts_sample = ...
711 |
712 | n_dirs_local = cfg.n_dirs // cfg.n_op_shards
713 | bs_local = cfg.bs // cfg.n_replicas
714 |
715 | ae = FastAutoencoder(
716 | n_dirs_local=n_dirs_local,
717 | d_model=cfg.d_model,
718 | k=cfg.k,
719 | auxk=cfg.auxk,
720 | dead_steps_threshold=cfg.dead_toks_threshold // cfg.bs,
721 | comms=comms,
722 | )
723 | ae.cuda()
724 | init_from_data_(ae, stats_acts_sample, comms)
725 | # IMPORTANT: make sure all DP ranks have the same params
726 | comms.init_broadcast_(ae)
727 |
728 | mse_scale = (
729 | 1 / ((stats_acts_sample.float().mean(dim=0) - stats_acts_sample.float()) ** 2).mean()
730 | )
731 | comms.all_broadcast(mse_scale)
732 | mse_scale = mse_scale.item()
733 |
734 | logger = Logger(
735 | project=cfg.wandb_project,
736 | name=cfg.wandb_name,
737 | dummy=cfg.wandb_project is None,
738 | )
739 |
740 | training_loop_(
741 | ae,
742 | batch_tensors(
743 | acts_iter,
744 | bs_local,
745 | drop_last=True,
746 | ),
747 | lambda ae, flat_acts_train_batch, recons, info, logger: (
748 | # MSE
749 | logger.logkv("train_recons", mse_scale * mse(recons, flat_acts_train_batch))
750 | # AuxK
751 | + logger.logkv(
752 | "train_maxk_recons",
753 | cfg.auxk_coef
754 | * normalized_mse(
755 | ae.decode_sparse(
756 | info["auxk_inds"],
757 | info["auxk_vals"],
758 | ),
759 | flat_acts_train_batch - recons.detach() + ae.pre_bias.detach(),
760 | ).nan_to_num(0),
761 | )
762 | ),
763 | lr=cfg.lr,
764 | eps=cfg.eps,
765 | clip_grad=cfg.clip_grad,
766 | ema_multiplier=cfg.ema_multiplier,
767 | logger=logger,
768 | comms=comms,
769 | )
770 |
771 |
772 | if __name__ == "__main__":
773 | main()
774 |
--------------------------------------------------------------------------------