main.py

import os
import pandas as pd
from modal import Secret, gpu
from sae_lens import SAE
from tqdm import tqdm
from src import auto_interp
from src.auto_interp import auto_interp_pipeline
from src.utils import chunkify, get_device
from src.auto_interp.auto_interp_pipeline import AutoInterpPipeline, batch_interpret_models
from src.dashboard.feature_dataset import FeatureCacheManager, filter_by_quantiles
from src.experiment_utils import reorder_sae
from src.extract_data import batch_create_datasets, build_feature_vis, load_sae_model
from src.models.TopKSaes import MatryoshkaTopKSaeConfig
from src.models.abstractSae import SaeBase
from src.training.train import SaeTrainConfig
from src.training.build_dataset import PILE_PATH, create_acts_dataset_wrapper, create_matched_dataloaders, load_tl_model, load_tokens_from_folder, tokenize_dataset_wrapper
from src.training.dataloader import DataLoaderLite
from src.models.TopKSaes import MatryoshkaTopKSae, TopKSaeConfig, TopKSae
from src.training.train import SaeTrainer
from src.training.common import vol, image, app, dataset_vol
from src.models.MatryoshkaSaeBase import MatryoshkaSaeBaseConfig,MatryoshkaSaeBase
from src.dashboard.feature_dataset import ModifiedFeatureDataset, filter_by_quantiles
from typing import List, Literal, Optional, Union
from src.training.build_dataset import load_activations_from_folder
from torch.utils.data import DataLoader
import torch
from saved_models import GEMMA_SCOPE_VERSIONS, GEMMA_SCOPE_RELEASE_ID, TOPK_SAE_PATHS, MATRYOSHKA_TOPK_SAE_PATHS, GEMMA_MODELS
PATH = "/autoencoder"

def _main(
    cfg : Union[MatryoshkaTopKSaeConfig, TopKSaeConfig],
    possible_batch_size : int,
    use_wandb : bool = True,
    N_tokens : int = 10_000_000,
    save_every_n_tokens : int = 10_000_000,
    eval_every_n_tokens : int = 10_000_000,
    feature_dead_threshold : int = 10_000_000,
    test_granularities : Optional[List[int]] = None,
    debug : bool = False,
    build_feature_vis : bool = False
):
    file_dir = f"{PILE_PATH}/activations"
    
    val_file_dir = '/root/pile_uncopyrighted/activations/token_mapped'
    all_activations = load_activations_from_folder(val_file_dir)
    all_activations = all_activations.view(-1, all_activations.shape[-1]) #[batch_size, d_model]
    val_dataloader = DataLoader(dataset=all_activations, batch_size=possible_batch_size, shuffle=False)
    
    if debug:
        train_dataloader = val_dataloader
    else:
        train_dataloader = DataLoaderLite(
            B=possible_batch_size,
            data_root=file_dir,
            split="train",
        )
       
    
    train_cfg = SaeTrainConfig(
        run_name="first-run",
        sae_cfg=cfg,
        target_batch_size=possible_batch_size,
        max_possible_batch_size=possible_batch_size,
        learning_rate=1e-4,
        wandb=use_wandb,
        n_tokens=N_tokens,
        feature_dead_threshold=feature_dead_threshold,
        val_n_tokens=100_000,
        log_feature_histogram_every_n_tokens=feature_dead_threshold,
        save_every_n_tokens=save_every_n_tokens,
        eval_every_n_tokens=eval_every_n_tokens,
        on_modal=True,
        test_granularities=test_granularities,
        on_save_build_feature_vis=build_feature_vis,
    )
        
    trainer = SaeTrainer(train_cfg, train_dataloader, val_dataloader)
    trainer.fit()
    
    vol.commit()


@app.function(
    gpu=gpu.A100(size="40GB"),
    concurrency_limit=4,
    image=image,
    volumes={
        '/root/pile_uncopyrighted': vol,
    },
    secrets=[
        Secret.from_name("wandb"), 
        Secret.from_name("HF_SECRET")
    ],
    timeout=60*60*14,
)
def wrapper(
    cfg : Union[MatryoshkaSaeBaseConfig], 
    possible_batch_size : int,
    use_wandb : bool = True,
    N_tokens : int = 10_000_000,
    save_every_n_tokens : int = 10_000_000,
    feature_dead_threshold : int = 10_000_000,
    eval_every_n_tokens : int = 10_000_000,
    test_granularities : Optional[List[int]] = None,
    debug : bool = False,
    build_feature_vis : bool = False
):
    _main(
        cfg, 
        possible_batch_size, 
        use_wandb=use_wandb, 
        N_tokens=N_tokens, 
        save_every_n_tokens=save_every_n_tokens, 
        eval_every_n_tokens=eval_every_n_tokens,
        feature_dead_threshold=feature_dead_threshold,
        test_granularities=test_granularities,
        debug=debug,
        build_feature_vis=build_feature_vis
    )
    vol.commit()

KILL_AFTER = 60*60*14
@app.function(
    gpu=gpu.A10G(),
    #gpu=gpu.T4(),
    #gpu=gpu.A100(size="40GB"),
    image=image, 
    timeout=KILL_AFTER,
    volumes={'/root/pile_uncopyrighted': vol},
    secrets=[Secret.from_name("wandb"), Secret.from_name("HF_SECRET")],
)
def ssh_function():
    import subprocess
    import time
    import modal
    import signal
    import torch
    import gc
    import os
    import atexit

    def cleanup_gpu(signum=None, frame=None):
        try:
            # Clear CUDA memory
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                gc.collect()
            
            # More graceful process termination
            os.system('pkill -15 -f python')  # Send SIGTERM first
            time.sleep(2)  # Give processes time to cleanup
            current_pid = os.getpid()
            os.system(f'pkill -9 -f python && kill -9 $(pgrep -f python | grep -v {current_pid})')
            
            if torch.cuda.is_available():
                os.system('nvidia-smi --gpu-reset')
        except Exception as e:
            print(f"Cleanup error: {e}")

    # Register cleanup for various signals and exit
    signal.signal(signal.SIGHUP, cleanup_gpu)
    signal.signal(signal.SIGTERM, cleanup_gpu)
    signal.signal(signal.SIGINT, cleanup_gpu)
    atexit.register(cleanup_gpu)

    # Configure sshd with custom settings
    sshd_config = """
    PrintMotd no
    PrintLastLog no
    UsePAM no
    """
    with open("/etc/ssh/sshd_config.d/custom.conf", "w") as f:
        f.write(sshd_config)

    try:
        subprocess.run(["service", "ssh", "restart"], check=True)
        with modal.forward(port=22, unencrypted=True) as tunnel:
            hostname, port = tunnel.tcp_socket
            connection_cmd = f'ssh -p {port} root@{hostname}'
            print(f"ssh into container using: {connection_cmd}")
            
            while True:
                time.sleep(60)  # Check every minute
                # Verify SSH daemon is still running
                try:
                    subprocess.run(["pgrep", "sshd"], check=True)
                except subprocess.CalledProcessError:
                    print("SSH daemon died, restarting...")
                    subprocess.run(["service", "ssh", "restart"], check=True)
    except Exception as e:
        print(f"SSH server error: {e}")
    finally:
        cleanup_gpu()
        vol.commit()
        
        

@app.function(
    #gpu=gpu.A100(size="40GB"),
    #gpu=gpu.A10G(),
    image=image,
    volumes={
        '/root/pile_uncopyrighted': vol,
    },
    secrets=[
        Secret.from_name("wandb"), 
        Secret.from_name("HF_SECRET")
    ],
    timeout=60*60*2,
)
def run_wrapper():

    from saved_models import MATRYOSHKA_SAMPLED_PATHS


    import os
    os.environ['FIREWORKS_AI_API_KEY'] = "7SZy6719Es1HdhKPoRq7M3t2eRDHtElrt5yV7chGb0DJsjb0"
    path1 = '/root/pile_uncopyrighted/datasets/auto_interp'
    import os
    
    os.makedirs(path1, exist_ok=True)
    from saved_models import MATRYOSHKA_SAMPLED_PATHS
    
    async def main():
        await batch_interpret_models(
            sae_entries=MATRYOSHKA_SAMPLED_PATHS, 
            copy_to_path=path1, 
            device=torch.device('cpu'),
            type='simulator',
            skip_explanations=False,
            n_quantiles=10,
            elems_per_quantile=30,
        )

    import asyncio
    asyncio.run(main())
    vol.commit()
    
    
def upload():
    with vol.batch_upload(force=True) as batch:
        batch.put_directory("datasets/auto_interp", '/root/pile_uncopyrighted/datasets/auto_interp')
            
@app.local_entrypoint()
def main():
    #upload()
    
    #create_feature_cache.remote()
    run_wrapper.remote()
    #interpret_models.remote()
    
# Example usage:
if __name__ == "__main__":
    
    os.environ["WANDB_API_KEY"] = "a3469eb2df23f67e4d6907ebacf50ffb4ee664f7"
    os.environ["HF_TOKEN"] = "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"