ReturnnTrainingJob, support JAX checkpoints

[albertz]

No description provided.

[michelwi]

should we create a JaxCheckpoint helper class similar to Checkpion/PtCheckpoint ?

(I also think those checkpoints should have a common base class Checkpoint -> {Tf,Pt,Jax}Checkpoint, but that is a different efactor)

[albertz]

should we create a JaxCheckpoint helper class similar to Checkpion/PtCheckpoint ?

I didn't needed that indirection. I'm also not sure how e.g. PtCheckpoint is really helpful. We don't have such wrapper objects for any other jobs.

[michelwi]

I think it was introduced for TF for the purpose of loading a model in returnn. A TF checkpoint consists of multiple files (Graph, weights, index I waguely remember) and returnn expects/expected the common filename without extension to load it. But this filename prefix does not correspond to a physical file on disk, so defining a tk.Path object for it would lead to "Job is not runnable" errors. Hence The Checkpoint class.

PtCheckpoint really wasn't necessary and was introduced for symmetry I guess.

Nowadays we use it to distinguish the checkpoints.. e.g.

def postprocess_checkpopint(checkpoint: Checkpoint | PtCheckpoint):
  if isinstance(checkpoint, Checkpoint):
    # TF specific postprocessing
  elif isinstance(checkpoint, PtCheckpoint):
    # PT specific postprocesing

But of course there are other ways to distinguish those.

For Jax, I guess a tk.Path pointing to a folder that exists would be valid and work.

[albertz]

I think it was introduced for TF for the purpose of loading a model in returnn. A TF checkpoint consists of multiple files (Graph, weights, index I waguely remember) and returnn expects/expected the common filename without extension to load it. But this filename prefix does not correspond to a physical file on disk, so defining a tk.Path object for it would lead to "Job is not runnable" errors.

I don't exactly remember the history, but this is not true since a long time. The TF checkpoint points to self.output_path("models/epoch.%.3d.index" % k). This is an actual physical file. The (TF) Checkpoint class then has this:

    @property
    def ckpt_path(self):
        return self.index_path.get_path()[: -len(".index")]

    def __str__(self):
        return self.ckpt_path

    def __repr__(self):
        return "'%s'" % self.ckpt_path

And the weird/ugly earlier ReturnnConfig serialization the used the Checkpoint.__repr__ here when you put it as ..., "load": chkpt, ... into the config.

Which was always weird to me, because you simply could just directly have "load": index_path in the config as well. There was no need to remove the .index extension. Maybe people just did this because it looked nicer? Or people assumed it must be this way?

So, to summarize, I think there never really was any valid reason to have this Checkpoint class in the first place. Probably just some confusion.

[michelwi]

Or people assumed it must be this way?

Probably.. I did assume it must be this way until now xD