keepsloading/icml_repro_scratch
The keepsloading/icml_repro_scratch model is a reproducibility package for training and evaluating sparse-attention variants of the Qwen3-1.7B architecture. Developed by keepsloading, this 1.7 billion parameter model focuses on comparing AHA and L2A-style shared-gate sparse attention mechanisms. It provides a self-contained environment to reproduce experiments on router granularity, offering insights into efficient attention implementations for large language models.
Loading preview...
Overview
This repository provides a self-contained reproducibility package for training and evaluating two distinct sparse-attention variants derived from a tuned Qwen3-1.7B checkpoint. The primary goal is to compare the impact of router granularity on model performance and sparsity.
Key Variants & Experimentation
- AHA (Attention Head Activation): Uses one gate per token, KV head, and layer (
token_kv_head). - L2A-style shared-gate: Employs one gate per token and layer, shared across all heads (
token).
Both variants utilize the same underlying AHA code, local attention, loss functions, data, optimizer settings, and a two-stage training protocol. The key experimental variable is the router granularity.
Training & Evaluation Protocol
- Starting Point: All experiments begin from a tuned-vanilla Qwen3-1.7B checkpoint included in the repository.
- Hardware: Requires at least one NVIDIA GPU with 80 GB memory (e.g., A100 80 GB) and 120 GB free disk space.
- Training: Each variant is trained as a single-GPU run. The process involves a hot start, Stage 1 (Gate distillation), and Stage 2 (Grouped-LR recovery).
- Evaluation: Strict evaluation is performed on benchmarks including local RULER, BabiLong, HELMET-ICL, and MRCR. It scans various inference thresholds to identify the highest effective sparsity that maintains at least 95% of the tuned vanilla model's macro score across all suites.
Integrity and Environment
The package emphasizes reproducibility through:
- Pinned Docker environment and Python dependencies.
- SHA-256 verification of starting weights, datasets, and source code.
- Comprehensive integrity checks and implementation tests before training.
Use Cases
This package is ideal for researchers and developers interested in:
- Reproducing sparse attention experiments: Specifically, comparing AHA and L2A-style router granularities.
- Understanding sparse attention mechanisms: Gaining insights into how different gating strategies affect model performance and efficiency.
- Developing efficient LLMs: Exploring methods to reduce computational overhead while maintaining model quality.