save

Author: luciaquirke

bergson/trackstar.py

@@ -6,7 +6,7 @@
     IndexConfig,
     PreprocessConfig,
     ScoreConfig,
-    TrackstarConfig,
+    TrackStarConfig,
 )
 from .process_grads import mix_preconditioners
 from .score.score import score_dataset
@@ -48,7 +48,7 @@ def trackstar(
     index_cfg: IndexConfig,
     score_cfg: ScoreConfig,
     preprocess_cfg: PreprocessConfig,
-    trackstar_cfg: TrackstarConfig,
+    trackstar_cfg: TrackStarConfig,
 ):
     """Run the full trackstar pipeline: preconditioners -> mix -> build -> score."""
     run_path = index_cfg.run_path
@@ -57,7 +57,6 @@ def trackstar(
     mixed_preconditioner_path = f"{run_path}/mixed_preconditioner"
     query_path = f"{run_path}/query"
     scores_path = f"{run_path}/scores"
-    resume = trackstar_cfg.resume
 
     # Steps 1-2 only compute preconditioners, so don't preprocess grads.
     precond_preprocess_cfg = PreprocessConfig()

examples/filter_data.py

@@ -115,6 +115,13 @@ class FilterConfig:
     projection_dim: int = 16
     """Projection dimension for gradient index."""
 
+    test_size: float = 0.05
+
+    tag: str = ""
+
+    pdbs: int = 8
+    "Per-device batch size"
+
 
 def run_sft(
     cfg: FilterConfig,
@@ -145,10 +152,13 @@ def run_sft(
             bias="none",
             task_type="CAUSAL_LM",
         )
-        model = get_peft_model(model, lora_config)
-        model.print_trainable_parameters()
+        model = get_peft_model(model, lora_config)  # type: ignore
+        model.print_trainable_parameters()  # type: ignore
 
-    num_train_steps = (len(train) // 32) * cfg.num_epochs
+    effective_batch_size = 32
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    grad_acc_steps = effective_batch_size / world_size / cfg.pdbs
+    num_train_steps = (len(train) // effective_batch_size) * cfg.num_epochs
     eval_steps = max(1, num_train_steps // 10)
 
     trainer = SFTTrainer(
@@ -158,9 +168,9 @@ def run_sft(
         args=SFTConfig(
             max_length=2048,
             output_dir=output_dir,
-            per_device_train_batch_size=1,
-            per_device_eval_batch_size=1,
-            gradient_accumulation_steps=32,
+            per_device_train_batch_size=cfg.pdbs,
+            per_device_eval_batch_size=cfg.pdbs,
+            gradient_accumulation_steps=grad_acc_steps,
             gradient_checkpointing=True,
             learning_rate=3e-4,
             num_train_epochs=cfg.num_epochs,
@@ -291,6 +301,8 @@ def run_trackstar(
         "--nproc_per_node",
         str(num_gpus),
         "--overwrite",
+        "--index_cfg.precision",
+        "bf16",
     ]
     # PEFT models need explicit tokenizer since adapter dir has no tokenizer config
     if args.use_lora:
@@ -332,7 +344,7 @@ def sft_full(args: FilterConfig, output_dir: str) -> str:
     if args.max_samples:
         dataset = dataset.select(range(min(args.max_samples, len(dataset))))
 
-    split = dataset.train_test_split(test_size=0.05, seed=args.seed)
+    split = dataset.train_test_split(test_size=args.test_size, seed=args.seed)
     train_ds, eval_ds = split["train"], split["test"]
 
     tokenizer = AutoTokenizer.from_pretrained(args.model, max_length=8192)
@@ -506,18 +518,16 @@ def main(
 
     # Always load the original text dataset for training.
     # Don't shuffle here — order must match the gradient index built by bergson.
-    orig_dataset = assert_type(Dataset, load_dataset(args.dataset, split=args.split))
+    ds = assert_type(Dataset, load_dataset(args.dataset, split=args.split))
     if args.max_samples:
-        orig_dataset = orig_dataset.select(
-            range(min(args.max_samples, len(orig_dataset)))
-        )
+        ds = ds.select(range(min(args.max_samples, len(ds))))
 
     # Add original index column so we can map back after train_test_split shuffles
-    orig_dataset = orig_dataset.add_column("_orig_idx", list(range(len(orig_dataset))))
+    ds = ds.add_column("_orig_idx", list(range(len(ds))))
 
     # Split original dataset (same seed ensures consistent eval set)
     print("Splitting...")
-    orig_split = orig_dataset.train_test_split(test_size=0.05, seed=args.seed)
+    orig_split = ds.train_test_split(test_size=args.test_size, seed=args.seed)
     orig_train, orig_eval = orig_split["train"], orig_split["test"]
 
     model_name = args.model.split("/")[-1]
@@ -526,12 +536,13 @@ def main(
     lora_suffix = "_lora" if args.use_lora else ""
     proj_suffix = f"_p{args.projection_dim}" if args.projection_dim != 16 else ""
 
-    if args.filter in ("attribution", "loss"):
-        # Step 1: SFT on the full dataset so gradients are meaningful
+    if args.filter in ("attribution", "loss", "trackstar"):
+        # SFT on the full dataset so training statistics can be collected
         sft_dir = f"examples/runs/{model_name}_{dataset_name}_sft{lora_suffix}"
         sft_model_path = sft_full(args, sft_dir)
 
-        # Step 2: Build gradient index using the finetuned checkpoint
+    if args.filter in ("attribution", "loss"):
+        # Collect gradients and losses using the fine-tuned model
         if not args.index_dataset:
             args.index_dataset = (
                 f"examples/runs/{model_name}_{dataset_name}"
@@ -541,20 +552,17 @@ def main(
         build_index(args, args.index_dataset, model=sft_model_path)
         grad_dataset = load_gradient_dataset(Path(args.index_dataset), structured=False)
 
-        # Split gradient dataset the same way
-        grad_split = grad_dataset.train_test_split(test_size=0.05, seed=args.seed)
+        # Split resulting data to match the original train/test split
+        grad_split = grad_dataset.train_test_split(
+            test_size=args.test_size, seed=args.seed
+        )
         grad_train = grad_split["train"]
         grad_train.set_format("torch")
-
     elif args.filter == "trackstar":
-        # Step 1: SFT on the full dataset so gradients are meaningful
-        sft_dir = f"examples/runs/{model_name}_{dataset_name}_sft{lora_suffix}"
-        sft_model_path = sft_full(args, sft_dir)
-
         # Step 2: Run trackstar pipeline for scoring
         trackstar_path = (
             f"examples/runs/{model_name}_{dataset_name}"
-            f"_trackstar{lora_suffix}{proj_suffix}"
+            f"_trackstar{lora_suffix}{proj_suffix}{args.tag}"
         )
         run_trackstar(
             args,