OpenMOSS
diff --git a/‎.github/workflows/checks.yml‎
Lines changed: 2 additions & 1 deletion b/‎.github/workflows/checks.yml‎
Lines changed: 2 additions & 1 deletion
diff --git a/‎.github/workflows/ui-ssr-checks.yml‎
Lines changed: 42 additions & 0 deletions b/‎.github/workflows/ui-ssr-checks.yml‎
Lines changed: 42 additions & 0 deletions
diff --git a/‎.pre-commit-config.yaml‎
Lines changed: 11 additions & 0 deletions b/‎.pre-commit-config.yaml‎
Lines changed: 11 additions & 0 deletions
diff --git a/‎README.md‎
Lines changed: 20 additions & 6 deletions b/‎README.md‎
Lines changed: 20 additions & 6 deletions
diff --git a/‎TransformerLens‎ b/‎TransformerLens‎
diff --git a/‎docs/index.md‎
Lines changed: 36 additions & 15 deletions b/‎docs/index.md‎
Lines changed: 36 additions & 15 deletions
diff --git a/‎docs/workthrough.md‎
Lines changed: 11 additions & 0 deletions b/‎docs/workthrough.md‎
Lines changed: 11 additions & 0 deletions
diff --git a/‎examples/README.md‎
Lines changed: 17 additions & 0 deletions b/‎examples/README.md‎
Lines changed: 17 additions & 0 deletions
diff --git a/‎examples/analyze_pythia_sae.py‎
Lines changed: 71 additions & 0 deletions b/‎examples/analyze_pythia_sae.py‎
Lines changed: 71 additions & 0 deletions
diff --git a/‎examples/analyze_pythia_sae_with_pre_generated_activations.py‎
Lines changed: 62 additions & 0 deletions b/‎examples/analyze_pythia_sae_with_pre_generated_activations.py‎
Lines changed: 62 additions & 0 deletions
@@ -46,14 +46,15 @@ jobs:
         uses: actions/checkout@v4
         with:
           submodules: "true"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v5
         with:
           enable-cache: true
           cache-dependency-glob: "uv.lock"
 
       - name: Install the project
-        run: uv sync --extra default --dev
+        run: uv sync --all-extras --dev
 
       - name: Type check
         run: uv run basedpyright .
 
@@ -0,0 +1,42 @@
+name: UI SSR Checks
+
+on:
+  push:
+    branches:
+      - main
+      - dev
+    paths:
+      - "ui-ssr/**"
+      - ".github/workflows/ui-ssr-checks.yml"
+  pull_request:
+    branches:
+      - main
+      - dev
+    paths:
+      - "ui-ssr/**"
+      - ".github/workflows/ui-ssr-checks.yml"
+
+permissions:
+  contents: read
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ui-ssr
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Setup Bun
+        uses: oven-sh/setup-bun@v1
+        with:
+          bun-version: latest
+
+      - name: Install dependencies
+        run: bun install
+
+      - name: Run CI checks
+        run: bun run ci:check
@@ -8,3 +8,14 @@ repos:
         args: [--fix]
       # Run the formatter.
       - id: ruff-format
+
+  - repo: local
+    hooks:
+      - id: ui-ssr-lint-staged
+        name: ui-ssr lint-staged
+        # Enter the directory and run lint-staged.
+        # lint-staged will handle the staged files selection and running prettier/eslint.
+        entry: bash -c 'cd ui-ssr && bun run check'
+        language: system
+        pass_filenames: false
+        files: ^ui-ssr/
@@ -30,10 +30,26 @@
 
 ## Installation
 
+Use [pip](https://pypi.org/project/pip/) to install Language-Model-SAEs:
+
+```bash
+pip install lm-saes==2.0.0b4
+```
+
+We also highly recommend using [uv](https://docs.astral.sh/uv/) to manage your own project dependencies. You can use
+
+```bash
+uv add lm-saes==2.0.0b4
+```
+
+to add Language-Model-SAEs as your project dependency.
+
+## Development
+
 We use [uv](https://docs.astral.sh/uv/) to manage the dependencies, which is an alternative to [poetry](https://python-poetry.org/) or [pdm](https://pdm-project.org/). To install the required packages, just install [uv](https://docs.astral.sh/uv/getting-started/installation/), and run the following command:
 
 ```bash
-uv sync --extra default
+uv sync
 ```
 
 This will install all the required packages for the codebase in `.venv` directory. For Ascend NPU support, run
@@ -47,15 +63,13 @@ A forked version of `TransformerLens` is also included in the dependencies to pr
 If you want to use the visualization tools, you also need to install the required packages for the frontend, which uses [bun](https://bun.sh/) for dependency management. Follow the instructions on the website to install it, and then run the following command:
 
 ```bash
-cd ui
+cd ui-ssr
 bun install
 ```
 
-`bun` is not well-supported on Windows, so you may need to use WSL or other Linux-based solutions to run the frontend, or consider using a different package manager, such as `pnpm` or `yarn`.
-
 ## Launch an Experiment
 
-The guidelines and examples for launching experiments are generally outdated. At this moment, you may explore `src/lm_saes/runners` folder for the interface for generating activations and training & analyzing SAE variants. For analyzing SAEs, a MongoDB instance is required. More instructions will be provided in near future.
+Explore the `examples` to check the basic usage of training/analyzing SAEs in different configurations. Note a MongoDB is recommended for recording the model/dataset/SAE configurations and required for storing analyses. For more advanced usage, you may explore `src/lm_saes/runners` folder for the interface for generating activations and training & analyzing SAE variants, and directly write your own variant of training/analyzing script at the runner level.
 
 ## Visualizing the Learned Dictionary
 
@@ -65,7 +79,7 @@ The analysis results will be saved using MongoDB, and you can use the provided v
 uvicorn server.app:app --port 24577 --env-file server/.env
 ```
 
-Then, copy the `ui/.env.example` file to `ui/.env` and modify the `VITE_BACKEND_URL` to fit your server settings (by default, it's `http://localhost:24577`), and start the frontend by running the following command:
+Then, copy the `ui/.env.example` file to `ui/.env` and modify the `BACKEND_URL` to fit your server settings (by default, it's `http://localhost:24577`), and start the frontend by running the following command:
 
 ```bash
 cd ui
 
@@ -60,40 +60,61 @@ To train a simple Sparse Autoencoder on `blocks.5.hook_resid_post` of a Pythia-1
 ```python
 settings = TrainSAESettings(
     sae=SAEConfig(
-        hook_point_in=f"blocks.5.hook_resid_post",
+        hook_point_in="blocks.6.hook_resid_post",
+        hook_point_out="blocks.6.hook_resid_post",
         d_model=768,
         expansion_factor=8,
-        act_fn="jumprelu",
+        act_fn="topk",
+        top_k=50,
+        dtype=torch.float32,
+        device="cuda",
     ),
     initializer=InitializerConfig(
         grid_search_init_norm=True,
     ),
     trainer=TrainerConfig(
-        lr=5e-5,
-        l1_coefficient=0.3,
+        lr=1e-4,
+        initial_k=50,
+        k_warmup_steps=0.1,
+        k_schedule_type="linear",
         total_training_tokens=800_000_000,
-        sparsity_loss_type="tanh-quad",
-        jumprelu_lr_factor=0.1,
+        log_frequency=1000,
+        eval_frequency=1000000,
+        n_checkpoints=0,
+        check_point_save_mode="linear",
+        exp_result_path="results",
     ),
+    model=LanguageModelConfig(
+        model_name="EleutherAI/pythia-160m",
+        device="cuda",
+        dtype="torch.float16",
+    ),
+    model_name="pythia-160m",
+    datasets={
+        "SlimPajama-3B": DatasetConfig(
+            dataset_name_or_path="Hzfinfdu/SlimPajama-3B",
+        )
+    },
     wandb=WandbConfig(
         wandb_project="lm-saes",
-        exp_name=name,
+        exp_name="pythia-160m-sae",
     ),
     activation_factory=ActivationFactoryConfig(
         sources=[
-            ActivationFactoryActivationsSource(
-                path=Path(args.activation_path).expanduser(),
-                name=f"pythia-160m-1d",
-                device="cuda",
-                dtype=torch.float32,
+            ActivationFactoryDatasetSource(
+                name="SlimPajama-3B",
             )
         ],
         target=ActivationFactoryTarget.ACTIVATIONS_1D,
-        hook_points=["blocks.5.hook_resid_post"],
+        hook_points=["blocks.6.hook_resid_post"],
         batch_size=4096,
-        buffer_size=None,
+        buffer_size=4096 * 4,
+        buffer_shuffle=BufferShuffleConfig(
+            perm_seed=42,
+            generator_device="cuda",
+        ),
     ),
-    sae_name="L5R",
+    sae_name="pythia-160m-sae",
     sae_series="pythia-sae",
 )
 train_sae(settings)
 
@@ -0,0 +1,11 @@
+# Workthrough
+
+`Language-Model-SAEs` provides a general way to train, analyze and visualize Sparse Autoencoders and their variants. To help you get started quickly, we've included [example scripts]() that guide you through each stage of working with SAEs. This guide begins with a foundational example and progressively introduces the core features and capabilities of the library.
+
+## Training Basic Sparse Autoencoders
+
+A [Sparse Autoencoder]() is trained to reconstruct model activations at specific position. We depend on [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) to take activations out of model forward pass, specified by hook points. To train a vanilla SAE on Pythia 160M Layer 6 output, you can create the following `TrainSAESetting`:
+
+```python
+
+```
@@ -0,0 +1,17 @@
+# Example setups of Language-Model-SAEs
+
+The standard SAE-based pipeline of mechanistically interpreting internal representations of language models contains the following steps: Generating activations (optional) -> Training SAEs -> Analyzing SAEs -> Visualizing analyses.
+
+Here present example setups of generating, training and analyzing, with variants of SAE architectures, activation functions and whether to use pre-generated activations.
+
+## Use on-the-fly model activations
+
+SAE training requires stream of model activations at certain hook points (i.e. specefic location of model internal representation). Model activations can either be cached ahead-of-time on the disk, or produced on the fly. 
+
+For on-the-fly model activation usage, the _Generating activations_ step can be skipped, and thus the overall pipeline is simplified. You can refer to [train_pythia_sae_topk](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/train_pythia_sae_topk.py) and [analyze_pythia_sae](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/analyze_pythia_sae.py) and other scripts without a `with_pre_generated_activations` suffix to launch the experiments on Pythia. Note the analyzing requires a MongoDB instance (default to `mongodb://localhost:27017`) running to save the analyzing results.
+
+## Use cached activations
+
+Cached activations are more common usage in practical SAE training and analyzing. It enables effective hyperparameter sweeping with reuse of generated activations, and also enables parallelled training and analyzing (DP/TP). However, it requires a non-trivial amount of disk space, e.g., caching 800M tokens of one layer activation of Pythia 160M requires about 6TB space.
+
+To launch experiments with cached activations, you should first generate activations with 1d shape (`(batch, d_model)`, for training use), and 2d shape (`(batch, n_context, d_model)`, for analyzing use), by running [generate_pythia_activation_1d](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/generate_pythia_activation_1d.py) and [generate_pythia_activation_2d](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/generate_pythia_activation_2d.py). Then, you can use [train_pythia_sae_with_pre_generated_activations](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/train_pythia_sae_with_pre_generated_activations.py) and [analyze_pythia_sae_with_pre_generated_activations](https://github.com/OpenMOSS/Language-Model-SAEs/blob/main/examples/analyze_pythia_sae_with_pre_generated_activations.py) to run training and analyzing respectively, with a pre-generated activation path specified. Note the analyzing still requires a MongoDB instance running.
@@ -0,0 +1,71 @@
+import argparse
+import os
+
+import torch
+
+from lm_saes import (
+    ActivationFactoryConfig,
+    ActivationFactoryDatasetSource,
+    ActivationFactoryTarget,
+    AnalyzeSAESettings,
+    DatasetConfig,
+    FeatureAnalyzerConfig,
+    LanguageModelConfig,
+    MongoDBConfig,
+    SAEConfig,
+    analyze_sae,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sae_path", type=str, required=True)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+    args = parse_args()
+
+    sae_cfg = SAEConfig.from_pretrained(os.path.expanduser(args.sae_path), device="cuda", dtype=torch.float16)
+    analyze_sae(
+        AnalyzeSAESettings(
+            sae=sae_cfg,
+            sae_name="pythia-160m-sae",
+            sae_series="pythia-sae",
+            activation_factory=ActivationFactoryConfig(
+                sources=[
+                    ActivationFactoryDatasetSource(
+                        name="SlimPajama-3B",
+                    )
+                ],
+                target=ActivationFactoryTarget.ACTIVATIONS_2D,
+                hook_points=["blocks.6.hook_resid_post"],
+                batch_size=16,
+                context_size=2048,
+            ),
+            model=LanguageModelConfig(
+                model_name="EleutherAI/pythia-160m",
+                device="cuda",
+                dtype="torch.float16",
+            ),
+            model_name="pythia-160m",
+            datasets={
+                "SlimPajama-3B": DatasetConfig(
+                    dataset_name_or_path="Hzfinfdu/SlimPajama-3B",
+                )
+            },
+            analyzer=FeatureAnalyzerConfig(
+                total_analyzing_tokens=100_000_000,
+                subsamples={
+                    "top_activations": {"proportion": 1.0, "n_samples": 20},
+                    "subsampling_80%": {"proportion": 0.8, "n_samples": 10},
+                    "subsampling_60%": {"proportion": 0.6, "n_samples": 10},
+                    "subsampling_40%": {"proportion": 0.4, "n_samples": 10},
+                    "non_activating": {"proportion": 0.3, "n_samples": 20, "max_length": 50},
+                },
+            ),
+            mongo=MongoDBConfig(),
+            device_type="cuda",
+        )
+    )
@@ -0,0 +1,62 @@
+import argparse
+import os
+
+import torch
+
+from lm_saes import (
+    ActivationFactoryActivationsSource,
+    ActivationFactoryConfig,
+    ActivationFactoryTarget,
+    AnalyzeSAESettings,
+    FeatureAnalyzerConfig,
+    MongoDBConfig,
+    SAEConfig,
+    analyze_sae,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sae_path", type=str, required=True)
+    parser.add_argument("--activation_path", type=str, required=True)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
+    args = parse_args()
+
+    sae_cfg = SAEConfig.from_pretrained(os.path.expanduser(args.sae_path), device="cuda", dtype=torch.float16)
+    analyze_sae(
+        AnalyzeSAESettings(
+            sae=sae_cfg,
+            sae_name="pythia-160m-sae",
+            sae_series="pythia-sae",
+            activation_factory=ActivationFactoryConfig(
+                sources=[
+                    ActivationFactoryActivationsSource(
+                        path=str(args.activation_path),
+                        name="pythia-160m-2d",
+                        device="cuda",
+                        dtype=torch.float16,
+                    )
+                ],
+                target=ActivationFactoryTarget.ACTIVATIONS_2D,
+                hook_points=["blocks.6.hook_resid_post"],
+                batch_size=16,
+                context_size=2048,
+            ),
+            analyzer=FeatureAnalyzerConfig(
+                total_analyzing_tokens=100_000_000,
+                subsamples={
+                    "top_activations": {"proportion": 1.0, "n_samples": 20},
+                    "subsampling_80%": {"proportion": 0.8, "n_samples": 10},
+                    "subsampling_60%": {"proportion": 0.6, "n_samples": 10},
+                    "subsampling_40%": {"proportion": 0.4, "n_samples": 10},
+                    "non_activating": {"proportion": 0.3, "n_samples": 20, "max_length": 50},
+                },
+            ),
+            mongo=MongoDBConfig(),
+            device_type="cuda",
+        )
+    )