Improve Activation Scoring of Puzzletron

Feature Request: Improve Activation Scoring of Puzzletron

Description

The current activation scoring implementation in Puzzletron has architectural limitations that hinder scalability, model support, and training efficiency. This request addresses three critical gaps: the lack of multi-hook execution, limited architectural support compared to Minitron, and missing distributed parallelization capabilities.

1. Multi-Hook Activation Scoring Runs

Currently, the system executes separate, suboptimal scoring rounds for individual hooks (e.g., one for FFN channel importance and another for query heads). This approach becomes entirely unmanageable with complex architectures like nemotron3-super, which require concurrent hooks for attention, Mamba, shared experts, individual expert channels, expert removal, and latent MoE. The scoring pipeline must be refactored to handle unified, multi-hook activation scoring within a single run to reduce redundant overhead.

2. Expanded Architecture Support

Puzzletron is currently limited to FFN channel scoring. In contrast, Minitron already includes mature support for Mamba layers and baseline capabilities for Mixture of Experts (MoE) layers. These existing layer-scoring features need to be ported directly from Minitron to Puzzletron to achieve feature parity and enable broader model profiling.

3. Comprehensive Distributed Parallelization

As model sizes scale and initial checkpoints utilize increasingly massive sequence lengths, Puzzletron requires robust distributed training infrastructure. We need to add full support for Expert Parallelism (EP), Tensor Parallelism (TP), Context Parallelism (CP), and Fully Sharded Data Parallel (FSDP). This should be achieved by introducing comprehensive AutoModel support that functions seamlessly under both configurations:

• force_hf = True
• force_hf = False

Business & Engineering Impact

Resolving these bottlenecks will significantly reduce the computational overhead of profiling next-generation architectures, unblock pruning workflows for MoE/Mamba models, and ensure our optimization tools scale to enterprise-level foundation checkpoints.