The engineering challenges of scaling interpretability

This post details the engineering challenges and solutions for scaling interpretability research, specifically focusing on distributed data shuffling and feature visualization generation. For shuffling, the team evolved from a single-GPU approach to a distributed strategy involving K jobs reading and writing portions of data, and later implemented a multi-pass approach to reduce shuffle size by N*K. For feature visualization, they developed a sharded approach over datasets and features, caching activations and then re-sharding over features to aggregate results, followed by a dataset-sharded pass to efficiently compute feature firing on surrounding tokens.