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. For step-by-step run instructions, see 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:

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.