Partial write is actually O(N^2) for non copy-on-write FS

[qGentry]

Seems like that due to atomicity we create a recursive copy of the checkpoint on every new entry.
https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/path/snapshot/snapshot.py#L72

For the regular NFS/S3 that doesn't support copy-on-write, this actually leads to O(N^2) traffic usage, because we make a full copy of existing checkpoint on each update. Is there a way to avoid that, even at the cost of disabling the atomicity? I don't really need it in my case by still want to use partial writes.

[lyglst]

We appreciate you reporting this issue - overhauling the existing design to break the atomicity is not an option unfortunately, as it would break our core API contract. As a workaround, I recommend you to partial write into a local cache (like NVME) before finalizing the checkpoint into an object storage system.

In the meantime, I definitely recognize this is a performance bottleneck. I will discuss with my teammates to figure out a way to mitigate this issue.

[BlaziusMaximus]

Hi Filipp, thanks for bringing this to our attention! The DefaultSnapshot mechanism relying on fully copying the checkpoint is definitely way too expensive in the partial saving case. But like Shutong said, breaking atomicity isn't an option since it could result in really nasty, hard to debug checkpoints on disk.

I think the best approach here would be to write each partial save to a new directory:

.partial_save/
  pytree/              # committed data from save 1
  pytree.pending_1/    # new leaves from save 2 (brand new dir, nothing to copy)
  pytree.pending_2/    # new leaves from save 3 (brand new dir, nothing to copy)

Then we'd perform a cheap merge during ocp.partial.finalize. This would reduce the io complexity from O(N^2) to O(N), as it should be, and preserve atomicity to boot. We'll work this into our priorities, and I'll make sure to update this issue with any progress.

Also, as the main author of partial saving, I'm glad to see it being adopted in OSS! I'd love to know more about your use case if you're willing to share. Thanks!

[qGentry]

Hey guys, thanks for your immediate attention for this matter!

To clarify, my main use-case for the partial checkpoints is to avoid materializing entire model state in-memory when doing offline checkpoint import from pytorch (i.e. HF safetensors) for large models (Kimi-K2, DeepSeekV3) to my internal JAX-based training framework.

I recommend you to partial write into a local cache (like NVME) before finalizing the checkpoint into an object storage system.

I wish I could but unfortunately this is not an option for me.

After some quick glance on the existing code-base, I've manager to come up with a quick hack for the partial writes that gives up on atomicity (and avoids expensive copies) which seems to work in my case, though it includes a lot of hacks (creating dirs manually, calling the finalize, etc.).

Would love to see proper solution in the upstream at some point!

    output_location = _to_epath(dst.location)
    if output_location.exists():
        raise OSError(f"Checkpoint directory {output_location} already exists.")

    step = dst.step if dst.step is not None else 0
    step_dir = output_location / str(step)
    metadata_dir = step_dir / "metadata"
    state_dir = step_dir / "state"
    state_work_root = step_dir / "state.partial_save"
    state_work_dir = state_work_root / "state"
    step_dir.mkdir(parents=True, exist_ok=False)
    metadata_dir.mkdir(parents=False, exist_ok=False)
    state_work_root.mkdir(parents=False, exist_ok=False)

    try:
        # loads lazily, tensors are loaded in RAM on-demand
        src_weights, _ = load_hf_safetensors_weight_source(
            src,
            cache_dir=cache_dir,
        )
        logger.info("Loaded source weights")

        metadata_handler = ocp.JsonCheckpointHandler()
        metadata_handler.save(
            epath.Path(metadata_dir), args=ocp.args.JsonSave(metadata)
        )
        logger.info("Saved checkpoint metadata")
        state_handler = ocp.StandardCheckpointHandler()
        logger.info(
            "Using StandardCheckpointHandler on %s to reuse Orbax additions-only "
            "partial-save merge logic without the v1 snapshot wrapper.",
            state_work_dir,
        )

        total_rules = len(flat_rules)
        for idx, (dst_key, rule_config) in enumerate(flat_rules.items(), start=1):
            rule = create_transform_rule(rule_config)
            logger.info(
                "Transforming %s (%d/%d)",
                "/".join(dst_key),
                idx,
                total_rules,
            )
            result = rule.transform(
                src_weights,
                dst_key=dst_key,
                dst_index=(),
            )
            nested_result = nnx.traversals.unflatten_mapping({dst_key: result})
            state_handler.save(
                epath.Path(state_work_dir),
                item={state_params_key: nested_result},
            )
            del result
            del nested_result
            gc.collect()

        state_work_dir.rename(state_dir)
        state_work_root.rmdir()
        BasePyTreeCheckpointHandler().finalize(state_dir)
        fix_missing_commit_success_txt(str(output_location), str(step))
        logger.info("Saved streamed checkpoint to %s", output_location)
    except Exception:
        if output_location.exists():
            logger.warning(
                "Import failed, deleting incomplete output dir %s", output_location
            )
            output_location.rmtree()
        raise