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

```bash
# 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.