social-data-attribution-w2/notes/preconditioner_performance.md

Preconditioner Build Performance

Observed Timings (SOC-152, 100K docs, OLMo-3-7B)

Step	8x H100	1x H100
Value preconditioner (100K docs, 251M tokens)	6m 53s	31m
Query preconditioner (17K queries)	1m 31s	~7m
Mix	0.3s	0.3s
Total	~9 min	~38 min

Throughput: ~243 docs/sec (8-GPU), ~31 docs/sec (1-GPU).

Why It's Fast

With projection_dim=16, the preconditioner matrices are 16x16 per module. The bottleneck is forward/backward passes through the model (~85% of time), not the preconditioner accumulation or eigendecomposition.

Pipeline Architecture

The preconditioner build has 3 steps with these dependencies:

Value preconditioner (Step 1)  ─┐
                                 ├─> Mix (Step 3) ─> done
Query preconditioner (Step 2)  ─┘

Steps 1 and 2 are independent and can run in parallel on separate nodes.

Optimization Levers

Walltime (scheduling)

The 12h SBATCH allocation from SOC-152 was over-allocated by ~80x. Reducing to 30 min lets SLURM backfill the job more aggressively.

Parallel steps 1+2

Value and query preconditioners are independent. Running them on separate nodes saves ~1.5 min wall time (query time) and reduces per-job GPU-hours.

token_batch_size

Higher batch sizes improve GPU utilization. SOC-152 used 4096 (down from 8192 which caused OOM). For H200 with 80GB VRAM, 8192 may be viable.

When to use multi-node chunking

Single 8-GPU node is fastest for 7B models (~9 min). Multi-node chunking adds overhead from model loading + merge and is slower for models that fit on one node. Use chunking only when:

• No 8-GPU nodes are available (only 4-GPU)
• Model is too large for a single node (32B+)
• Data sample is much larger than 100K docs

Pre-cached input data

HF download + consolidation takes 20-30 min. Run precond_cache_data.sbatch once to stage data in shared storage: /storage/ice-shared/cs7634/staff/TDA/preconditioner_data/

The launcher auto-discovers cached data there before falling back to scratch.

Usage

# Single-node (fastest, 8 GPUs):
bash scripts/slurm/attribution/launch_preconditioner.sh \
    --model allenai/Olmo-3-1025-7B

# Chunked multi-node (4 GPUs per node):
bash scripts/slurm/attribution/launch_preconditioner.sh \
    --model allenai/Olmo-3-7B-Instruct \
    --mode chunked --chunks 6 --gpus-per-node 4

# First run (downloads data from HF):
bash scripts/slurm/attribution/launch_preconditioner.sh \
    --model allenai/Olmo-3-1025-7B --prep

# Dry run (print commands only):
bash scripts/slurm/attribution/launch_preconditioner.sh \
    --model allenai/Olmo-3-1025-7B --dry-run

Job Dependency Graph

Single-node mode

Value (GPU) ──┐
              ├─> Mix (CPU) ─> done
Query (GPU) ──┘

Chunked mode

Chunk 0 (GPU) ──┐
Chunk 1 (GPU) ──┤
Chunk 2 (GPU) ──┼─> Merge+Mix (CPU) ─> done
  ...           │
Chunk N (GPU) ──┤
Query   (GPU) ──┘

In both modes, value and query builds run in parallel.