How does Orbax merge per-process checkpoint files into one checkpoint?

merge_ocdbt_per_process_files globs the per-process subdirectories (each named ocdbt.process_ ), opens every one of them as a child TensorStore KvStore, and opens the parent (global) KvStore with OCDBT write options. It then creates a single transaction with ts.Transaction(atomic=True), stages each child.experimental_copy_range_to(parent.with_transaction(txn)) into it, awaits them all, optionally validates the merged view, and finally calls txn.commit_async() once. The global manifest gains references to every per-process store in one atomic step.

Why does Orbax write checkpoints to per-process subdirectories first instead of one shared file?

Coordination and contention. With TensorStore's OCDBT format, arrays are initially written to per-process subdirectories depending on which host is doing the writing, so hundreds or thousands of hosts each write their own B-tree independently — no locks, no cross-host RPC, no single writer bottlenecking the save. The merge is a separate, cheap step that stitches those independent B-trees into one global key-value store afterward, so the expensive parallel write stays fully uncoordinated.

What makes the OCDBT checkpoint merge atomic?

The single ts.Transaction(atomic=True). Every child copy is staged into that one transaction rather than written directly, and nothing lands in the global manifest until the final txn.commit_async(). So a crash at any point before commit leaves the manifest exactly as it was — there is no intermediate state where a reader could see a manifest referencing some per-process stores but not others. It is all-or-nothing: either every per-process store is promoted into the global manifest, or none is.

Are the per-process subdirectories deleted after the merge?

No, and deleting them would corrupt the checkpoint. The file states it directly: the original per-process subdirectories are not and should not be deleted — the global kvstore continues to reference them. The merge writes manifest entries into the parent that point back at the per-process data, so the merged checkpoint depends on those subdirectories still existing. Removing them leaves a global manifest with dangling references.

What does Orbax validate after merging a checkpoint, and why?

_validate_params lists every key in the merged KvStore and classifies each as a data chunk or a .zarray metadata entry, grouping them by parameter name. It then checks that no parameter is missing its data chunks and no parameter is missing its .zarray metadata — raising a ValueError naming exactly which params are incomplete. Crucially, this runs on the transaction's uncommitted view before commit_async(), so a partial or truncated write is caught and the transaction is abandoned rather than committing a checkpoint that would load with silently missing parameters.

The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest

Every host in a training run writes its own OCDBT checkpoint files into a per-process subdirectory — no coordination, no contention. Turning those scattered per-process B-trees into one global manifest a reader can trust is a single atomic TensorStore transaction. We walk the real Orbax file that does it: glob the subdirs, stage every copy_range under ts.Transaction(atomic=True), validate the would-be-merged store, then commit once — all of it, or none of it.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest — built to be understood, not skimmed.

Questions

How does Orbax merge per-process checkpoint files into one checkpoint?: merge_ocdbt_per_process_files globs the per-process subdirectories (each named ocdbt.process_<n>), opens every one of them as a child TensorStore KvStore, and opens the parent (global) KvStore with OCDBT write options. It then creates a single transaction with ts.Transaction(atomic=True), stages each child.experimental_copy_range_to(parent.with_transaction(txn)) into it, awaits them all, optionally validates the merged view, and finally calls txn.commit_async() once. The global manifest gains references to every per-process store in one atomic step.
Why does Orbax write checkpoints to per-process subdirectories first instead of one shared file?: Coordination and contention. With TensorStore's OCDBT format, arrays are initially written to per-process subdirectories depending on which host is doing the writing, so hundreds or thousands of hosts each write their own B-tree independently — no locks, no cross-host RPC, no single writer bottlenecking the save. The merge is a separate, cheap step that stitches those independent B-trees into one global key-value store afterward, so the expensive parallel write stays fully uncoordinated.
What makes the OCDBT checkpoint merge atomic?: The single ts.Transaction(atomic=True). Every child copy is staged into that one transaction rather than written directly, and nothing lands in the global manifest until the final txn.commit_async(). So a crash at any point before commit leaves the manifest exactly as it was — there is no intermediate state where a reader could see a manifest referencing some per-process stores but not others. It is all-or-nothing: either every per-process store is promoted into the global manifest, or none is.
Are the per-process subdirectories deleted after the merge?: No, and deleting them would corrupt the checkpoint. The file states it directly: the original per-process subdirectories are not and should not be deleted — the global kvstore continues to reference them. The merge writes manifest entries into the parent that point back at the per-process data, so the merged checkpoint depends on those subdirectories still existing. Removing them leaves a global manifest with dangling references.
What does Orbax validate after merging a checkpoint, and why?: _validate_params lists every key in the merged KvStore and classifies each as a data chunk or a .zarray metadata entry, grouping them by parameter name. It then checks that no parameter is missing its data chunks and no parameter is missing its .zarray metadata — raising a ValueError naming exactly which params are incomplete. Crucially, this runs on the transaction's uncommitted view before commit_async(), so a partial or truncated write is caught and the transaction is abandoned rather than committing a checkpoint that would load with silently missing parameters.

The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest

Questions

Related explainers