social-data-attribution-w2/docs/TRACKSTAR_PIPELINE.md

# TrackStar Pipeline

Pipeline for scoring training data against evaluation queries using Bergson attribution. Takes evaluation queries (GSM8K, SocialIQA, MMLU Social Sciences, MMLU STEM) and a training corpus, then ranks which training documents are most influential for each query using gradient-based attribution scores.

For Bergson tool reference, see [BERGSON_REFERENCE.md](BERGSON_REFERENCE.md). For step-by-step run instructions, see [ATTRIBUTION_RUNBOOK.md](ATTRIBUTION_RUNBOOK.md).

## Architecture

The pipeline uses Bergson's Mode B (reduce/score) workflow. Three phases run as SLURM jobs with dependency chaining:

**Phase 1: Reduce** runs `bergson reduce` per query benchmark. Computes aggregated query gradients. One GPU job per benchmark (4 benchmarks = 4 jobs).

**Phase 2: Score** runs `bergson score` per (data shard x query). Streams training documents and scores each against the reduced query gradient. N shards x Q queries = N*Q GPU jobs.

**Phase 3: Aggregate** runs on CPU. Loads all shard scores for one variant/run_id, concatenates per query, selects top-k, writes ranked JSONL.

Jobs are chained with `--dependency=afterok` so each phase waits for the previous to succeed. If any reduce job fails, all score jobs are blocked.

## Code layout

### Python library (`src/data_attribution/attribution/trackstar/`)

| Module | Entry point | Purpose |
|--------|-------------|---------|
| `sharding.py` | (library) | JSONL splitting, zst decompression, shard discovery, path helpers |
| `queries.py` | `data-attribution-trackstar-query` | Download query JSONL from HuggingFace per variant |
| `prepare.py` | `data-attribution-trackstar-prepare` | Discover/download data, consolidate into N shards |
| `aggregate.py` | `data-attribution-trackstar-aggregate` | Load shard scores, rank top-k documents |
| `extract.py` | `data-attribution-trackstar-extract` | Extract all scores, join with manifest labels |

### Shell orchestration (`scripts/attribution/`)

| Script | Purpose |
|--------|---------|
| `launch_trackstar.sh` | Main pipeline launcher with arg parsing and dependency chaining |
| `trackstar_common.sh` | Variant configuration, SLURM flag builder, run ID generation |
| `trackstar_phases.sh` | Phase submission functions (reduce, score, aggregate) |
| `orchestrate_trackstar.sh` | Unattended runner: waits for data sync, shards, launches pipeline |
| `sync_sample.py` | Parallel R2 download (64 threads, skip-existing) |

### SLURM scripts (`scripts/slurm/attribution/`)

| Script | Phase | Resources |
|--------|-------|-----------|
| `trackstar_reduce.sbatch` | Reduce | 1 GPU, 12h walltime |
| `trackstar_score.sbatch` | Score | 1 GPU, 4h walltime |
| `trackstar_aggregate.sbatch` | Aggregate | CPU only, 64GB RAM, 1h |

### Tests (`tests/data_attribution/attribution/trackstar/`)

| File | Coverage |
|------|----------|
| `test_sharding.py` | `_open_lines`, `split_jsonl`, `split_multi_jsonl` (14 tests) |
| `test_pipeline.py` | Launch integration, path builders, SLURM script validation (8 tests) |

## Model variants

| Variant | Model | Query columns | Query dataset |
|---------|-------|---------------|---------------|
| base | `allenai/Olmo-3-1025-7B` | `text` | `HCAI-Lab/base-query-data` |
| instruct_base | `allenai/Olmo-3-7B-Instruct` | `prompt`, `completion` | `HCAI-Lab/instruct-query-data` |
| instruct_cot | `allenai/Olmo-3-7B-Instruct` | `prompt`, `completion` | `HCAI-Lab/instruct-cot-query-data` |

## Bergson parameters used

| Parameter | Value | Rationale |
|-----------|-------|-----------|
| precision | fp32 | Numerical stability for attribution |
| projection_dim | 16 | Exploration-grade compression (256-dim per module) |
| token_batch_size | 1024 | Conservative; fits H100 80GB with OLMo-3-7B |
| truncation | true | Dolma docs can be long |
| skip_preconditioners | true | Faster; add for final runs |
| normalizer | none | Raw gradients |
| unit_normalize | true | Cosine-like similarity across modules |
| loss_fn | ce | Cross-entropy |
| loss_reduction | mean | Mean across tokens |
| reduce aggregation | mean | Mean-reduce query gradients before scoring |
| score mode | individual | Emit one score column per document |

## Data flow

```
Query JSONL files (HuggingFace)
  |
  v
bergson reduce (Phase 1, GPU)
  -> reduces/{variant}/{run_id}/queries_{name}/
  |
Training JSONL shards (from prepare.py)
  |
  v
bergson score (Phase 2, GPU)
  -> scores/{variant}/{run_id}/{query}/{source}/{shard}/
     - scores.bin (float32 per doc)
     - data.hf (document metadata)
     - info.json
  |
  v
aggregate (Phase 3, CPU)
  -> results/{variant}/{run_id}/{query}_top{k}.jsonl
  |
  v
extract (optional, CPU)
  -> results/{variant}/{run_id}/all_scores_combined.parquet
     (all docs x all benchmarks, joined with manifest bin labels)
```

## SLURM configuration

All SLURM resource parameters are configurable via environment variables:

```bash
export SLURM_ACCOUNT=gts-schava6-qcf  # charge account (omit for ICE)
export SLURM_PARTITION=coe-gpu,ice-gpu  # partition(s)
export SLURM_QOS=coe-ice               # QOS
export SLURM_CONSTRAINT=HX00           # GPU constraint (H100/H200 for 7B fp32)
export GPU_WALLTIME=12:00:00           # walltime for GPU jobs
export CPUS_PER_TASK=4                 # CPUs per GPU job
export MEM=64G                         # memory per job
export PRECISION=fp32                  # fp32
export RUN_ROOT=runs/trackstar          # base directory for all outputs
export SHARD_DIR=runs/trackstar/shards  # location of sharded training data
export QUERY_DIR=queries/base          # query JSONL directory
```

## GPU requirements

OLMo-3-1025-7B in fp32 requires ~28GB VRAM for gradient computation. Use `--constraint HX00` on PACE ICE to target H100/H200 GPUs (80GB). Smaller GPUs (L40 22GB, V100 16-32GB) will OOM.

## Score output format

Each `bergson score` run writes per-shard:

- `scores.bin`: Structured numpy memmap. Dtype: `{score_0: float32, written_0: bool}`, 8 bytes per doc.
- `info.json`: `num_items`, `num_scores`, full dtype definition.
- `data.hf/`: HuggingFace datasets Arrow format with `length` and `loss` columns.

Scores are signed floats. Positive = training doc pushes model toward query behavior. Negative = pushes away. Magnitude indicates strength.

## Known limitations

- Mode B recomputes all document gradients for each benchmark. For N benchmarks, the training corpus is processed N times. Mode A (build index once) would be more efficient for 3+ benchmarks.
- The reduce step collapses all queries per benchmark into one mean gradient. Per-query attribution detail is lost. Mode A with Attributor preserves individual query scores.
- The `afterok` dependency chains all reduce jobs together. One failed reduce blocks all score jobs, not just the affected query.
- Score jobs target H100/H200 GPUs for OLMo-3-7B in fp32.