Skip to content

Add build → score pipeline example#253

Open
jammastergirish wants to merge 4 commits intoadd_yaml-mediated_search_pipelinefrom
add_build_score_pipeline_example_yaml
Open

Add build → score pipeline example#253
jammastergirish wants to merge 4 commits intoadd_yaml-mediated_search_pipelinefrom
add_build_score_pipeline_example_yaml

Conversation

@jammastergirish
Copy link
Copy Markdown
Collaborator

@jammastergirish jammastergirish commented Apr 29, 2026

Summary

Adds a runnable two-step pipeline example demonstrating the no-compression scoring path:

bergson pipeline examples/pipelines/build_then_score.yaml

Step 1 (build) creates a one-gradient on-disk query index with projection_dim: 0 and aggregation: mean. Step 2 (score) loads that query into memory and dot-products it against each item in a small slice of NeelNanda/pile-10k, producing per-item influence scores. Until now, this required two CLI invocations — this single YAML now drives the whole flow through bergson pipeline.

Why

  1. Showcase the YAML pipeline functionality introduced in ELE-11: Add YAML-mediated search pipeline to Bergson #246 with a real, useful example.
  2. Demonstrate a feature that previously required multiple commands (bergson build + bergson score) — now reducible to a single bergson pipeline invocation.
  3. Lay groundwork for EK-FAC support in score — having an end-to-end no-compression scoring pipeline in YAML form is a prerequisite for then layering Hessian application onto the uncompressed gradients (Lewis's worker has the existing reference for that), without entangling Hessian work with the rest of the codebase.

Files

  • examples/pipelines/build_then_score.yaml — the example pipeline (small model: gpt2, small dataset: NeelNanda/pile-10k with train[:20] for the query and train[:100] for scoring, chunk_length: 1024).
  • README.md — extends the "Run a Multi-Step Pipeline" section introduced in ELE-11: Add YAML-mediated search pipeline to Bergson #246 with a one-sentence reference to build_then_score.yaml as a second example covering the no-compression scoring path. (The section itself, the description of bergson pipeline, and the link to hessian_then_build.yaml all live in ELE-11: Add YAML-mediated search pipeline to Bergson #246.)
  • tests/test_yaml_pipeline.pytest_build_then_score_example_parses asserts the shipped YAML hydrates into the right Build/Score commands with the expected configs (parse-only, no GPU needed).

Configuration notes

  • skip_preconditioners: true is set on both build and score steps. With projection_dim: 0 the per-token gradient P retains the full flattened module-gradient dimension; the default per-module Adafactor preconditioner update P.mT @ P would otherwise materialize a (full_grad_dim × full_grad_dim) matrix — multi-TB even for GPT-2's c_attn — and OOM immediately. The no-compression scoring path uses raw dot products of full gradients, so preconditioners aren't needed here. See issue Default skip_preconditioners=True for score (or have projection_dim=0 imply it) #255 for a suggested follow-up to make this implicit.
  • unit_normalize and --preconditioner_path from the CLI command in the Linear task aren't replicated here. Those belong to the compressed/preconditioned path. This example is deliberately the no-compression demo.

Dependencies

jammastergirish and others added 4 commits April 29, 2026 11:12
With projection_dim=0 the per-token gradient P retains the full
flattened module-gradient dimension, so the default per-module
Adafactor preconditioner update `P.mT @ P` materializes a
(full_grad_dim x full_grad_dim) matrix — multiple TB even for
GPT-2's c_attn — and OOMs immediately.

The no-compression scoring path uses raw dot products of full
gradients, not preconditioned ones, so the example does not need
preconditioners. Set skip_preconditioners: true on the build step
to bypass the OOM-prone allocation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The score step also runs the gradient collector with
projection_dim=0, so it hits the same multi-TB `P.mT @ P`
allocation on the first backward pass. Mirror the build step
and set skip_preconditioners: true.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@jammastergirish jammastergirish force-pushed the add_build_score_pipeline_example_yaml branch from 52d766c to 0a05763 Compare April 29, 2026 18:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant