Store last.ckpt as symlink when appropriate to save space

[ZhaofengWu]

🚀 Feature

Currently, if I'm understanding correctly, last.ckpt is always stored as a separate model file (unless it's disabled, of course). But sometimes the same checkpoint already exists, for example as the best checkpoint, or as one of the top-k checkpoints. When this happens, it'd be very helpful that only a symlink is stored to reduce space usage.

Motivation

Modern models could take up a lot of disk space. In the typical use case where only the best and last checkpoints are stored, this could reduce the space usage by ~half.

cc @Borda @awaelchli

[carmocca]

This has been already proposed in #4335, but this cannot be done because the fsspec API does not support it. It wouldn't work for remote filesystems

[ZhaofengWu]

Thanks for the context! I'm not very familiar with fsspec and how lightning uses it, but why wouldn't something like os.symlink work? Because some people might work with two sets of file systems, and the path doesn't distinguish between them? If so, I believe it's a valid concern, but maybe not the most common use case? Why not just implement the proposed flag symlink_to_last in that issue and let the user decide when to enable it, and users who work with multiple file systems can choose to disable it?

[carmocca]

Yes, it can be implemented. It's open to community contributions.

[ZhaofengWu]

I'm happy to give this a try. But to verify the best way to do it: current I'm considering changing _save_topk_checkpoint to return the ckpt path iff it saved the current state. And we pass this flag into _save_last_checkpoint (most of the calls to these methods are preceded by _save_topk_checkpoint), which uses the flag to determine if it should symlink or not. Does that make sense?

[ZhaofengWu]

And of course, if we don't want this behavior to be the default, we can introduce a flag like symlink_last_ckpt.

[stas00]

could this be revived? This is not only a problem with disc space, the way it's done now is also taking 2x longer to save the checkpoint which for 100B model could take a really long time.

Additionally as I shared in #18670 there is a race condition and one could end up with 2 last files.

I proposed an alternatively solution used by other frameworks which uses an actual file latest which is updated to point to the last checkpoint's filename - no race condition in this situation since you only ever have one latest file.

But my main need is to remove the overhead of saving the file 2 times, which is a significant problem, especially on a slow filesystem.

Thank you!