Skip to content

Add a usage guide for the checkpoint benchmark framework#3333

Open
mridul-sahu wants to merge 1 commit into
google:mainfrom
mridul-sahu:feature/benchmark-guide
Open

Add a usage guide for the checkpoint benchmark framework#3333
mridul-sahu wants to merge 1 commit into
google:mainfrom
mridul-sahu:feature/benchmark-guide

Conversation

@mridul-sahu

Copy link
Copy Markdown
Collaborator

Description

Adds GUIDE.md under checkpoint/orbax/checkpoint/_src/testing/benchmarks/: a
structured usage guide for the checkpoint benchmark framework.

It documents:

  • the two-part model — a benchmark generator class (its test_fn runs the
    measured code) plus a YAML config that selects and parametrizes it;
  • how to write a benchmark and capture metrics with the one-line measure()
    block (including nested per-step measurement), and how to reuse the built-in
    generators;
  • the full config reference (checkpoint_config, mesh_config, benchmarks,
    baselines) and the metrics reference (collectors + the TensorBoard tag
    namespaces, with units);
  • the feature set — parameter sweeps, repeats + cross-host aggregation, baseline
    capture/compare, profiler traces / HLO dumps, TensorBoard output;
  • running on Pathways and cloud — automatic backend init, the colocated-Python
    load dispatcher (load_with_colocated_python), and the XPK launcher;
  • worked recipes, a cheat sheet, and a source map linking each component on
    GitHub.

Docs only — no code changes.

Type of change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (alters existing behavior or a public API)
  • Docs / tests / internal tooling (no user-facing behavior change)

Checklist

  • I have read the contribution guidelines.
  • Tests cover this change and pass locally.
  • Public APIs and non-trivial functions are documented.
  • If this change is user-facing, I have updated CHANGELOG.md

@github-actions github-actions Bot added the pull ready Ready to be pulled from GitHub into Google label Jun 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pull ready Ready to be pulled from GitHub into Google

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant