Run "on_stage_end" on all processes and save on only a single process - #2059

This seems better to me. But in the current implementation any registered save hooks have expected to be called on the main process. I don't know if Checkpointer should wrap saving in run_on_main or the individual objects. Maybe Checkpointer should do it?

Can saving still only run on the main process under FSDP, or do the other processes need to explicitly do something?

Also, the intra-epoch checkpoint saving as currently implemented (based on time) is not compatible with this:

Note: Checkpointer should also wrap the checkpoint deletion in run_on_main.

This seems better to me. But in the current implementation any registered save hooks have expected to be called on the main process. I don't know if Checkpointer should wrap saving in run_on_main or the individual objects. Maybe Checkpointer should do it?

Can saving still only run on the main process under FSDP, or do the other processes need to explicitly do something?

You've correctly deduced the reason why I put this in individual classes -- some (including FSDP) may want to run code on all processes. Perhaps a more backwards-compatible way of doing this would be to add some sort of argument (to mark_as_saver?) that would by default run only on main but would allow running on all if passed.... Thinking aloud on this one.

Also, the intra-epoch checkpoint saving as currently implemented (based on time) is not compatible with this:

speechbrain/speechbrain/core.py

Lines 1130 to 1144 in 5d479fc

	if (
	self.checkpointer is not None
	and self.ckpt_interval_minutes > 0
	and time.time() - last_ckpt_time
	>= self.ckpt_interval_minutes * 60.0
	):
	# This should not use run_on_main, because that
	# includes a DDP barrier. That eventually leads to a
	# crash when the processes'
	# time.time() - last_ckpt_time differ and some
	# processes enter this block while others don't,
	# missing the barrier.
	if sb.utils.distributed.if_main_process():
	self._save_intra_epoch_ckpt()
	last_ckpt_time = time.time()

This makes me realise that if we put run_on_main somewhere, then all processes must hit that code, so if the run_on_main goes into classes, we must ensure that all processes run that (and in the same order, i.e. if one process runs into save barrier and another into gradient sync barrier then it ends up in deadlock). Maybe we don't need the barrier inside run_on_main in many cases? There is already the if_main_process() (which really should be called is_main_process()), that might be enough in most cases. The other processes will wait at the next barrier, while main process saves or does whatever errands we need.
Note: Checkpointer should also wrap the checkpoint deletion in run_on_main.

Perhaps by default we should call if_main_process() unless the barrier is really needed. Wondering if there would be a reason for the barrier that I haven't thought of -- a timeout in the gradient sync or something like this.

Continuing on the checkpoints with time intervals: basically, time cannot be relied on for multi-process synchronisation (almost by definition). If FSDP needs Checkpointer code to run on all nodes then I think that this checkpoint with time interval is not possible with FSDP.

With the coming changes in SpeechBrain, I think we can afford to break some backwards compatibility and move to checkpoints every x steps. Some people also requested that anyway.

I think the mark_as_saver decorator could indeed have an optional parameter (something like @mark_as_saver(main_process_only=False)), and checkpointer could then keep track of this.

With the coming changes in SpeechBrain, I think we can afford to break some backwards compatibility and move to checkpoints every x steps. Some people also requested that anyway.

Should that be done as part of this PR or could it be a separate one?

I think the mark_as_saver decorator could indeed have an optional parameter (something like @mark_as_saver(main_process_only=False)), and checkpointer could then keep track of this.

Implementation attempted here -- does this look reasonable?

I think the two separated sets of default save hooks (main proc only vs. all procs) will lead to many errors. When searching for a default hook, the code looks through the class hierarchy, where it will often find e.g. torch.Module, and use that default hook. While we could change the order so that all_procs hooks are prioritized (as they are less common), I fear that this will lead to similar weird and hard to find bugs.

I think we'd want to keep one set of save hooks. Perhaps the checkpoint registering could wrap the saver hook callable in a wrapper that only runs on the main process:

from functools import wraps
def main_process_only(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if sb.utils.distributed.if_main_process():
            func(*args, **kwargs)
    return wrapper

Or we could keep more information about those hooks, e.g. the hooks dict could map types to some Hook dataclass with entries like callable, run_on_all_procs, etc. but then that might require changing all existing default hooks.

With the coming changes in SpeechBrain, I think we can afford to break some backwards compatibility and move to checkpoints every x steps. Some people also requested that anyway.

Should that be done as part of this PR or could it be a separate one?

It's something that breaks from this, so it could be done as part of this change or if we want to be tidy it could be a separate change implemented before this.

I think we'd want to keep one set of save hooks. Perhaps the checkpoint registering could wrap the saver hook callable in a wrapper that only runs on the main process

This is a neat solution, I wish I'd thought of it first!

With the coming changes in SpeechBrain, I think we can afford to break some backwards compatibility and move to checkpoints every x steps. Some people also requested that anyway.

Okay, I've made the change in core.py. Pretty much every recipe uses it, so it'll take some time to change over, but fortunately it's only a one-line change. I looked and the number of minutes varies from recipe to recipe. I'm not sure how to select a sane value for the checkpoint steps... any ideas?

I was thinking of assuming there's 5 iterations per second (a tremendously inaccurate assumption) and converting the minutes to steps (* 60 * 5), but there's probably a better way to do this.

Okay, I've made the change in core.py. Pretty much every recipe uses it, so it'll take some time to change over, but fortunately it's only a one-line change. I looked and the number of minutes varies from recipe to recipe. I'm not sure how to select a sane value for the checkpoint steps... any ideas?

Oof. Well, one idea is to keep the time based checkpointing, it is kind of useful and works in most cases (and now that you added step based checkpoints, keep those as well :D). We could also try to detect FSDP and raise an error, though I guess at the moment we cannot easily detect FSDP. Maybe a hack for now, something like checking if the name "FSDP" is in globals, I don't know.
EDIT: Maybe detect distributed situations (easier), and warn that if it's FSDP, time interval checkpoints don't work (but don't raise an error)

Otherwise, I think this could be a good implementation. You need to remember to wrap torch_recovery with main_process_only, and same for checkpoint deletion. Also, I wonder if many recipes do things like write some WER file in an on_stage_end, expecting this to happen only on main?

You need to remember to wrap torch_recovery with main_process_only, and same for checkpoint deletion.

Ah yes, I've added this for checkpoint deletion. However, for torch_recovery I'm pretty confident we don't want it to happen only on main, but rather on all processes, since we use this to load models on all devices (and there's no contested resource).

Also, I wonder if many recipes do things like write some WER file in an on_stage_end, expecting this to happen only on main?

These should be fixed:

$ grep -r 'with open(self.hparams.wer_file' recipes/ | wc -l
      33

But there are some cases of saving that don't necessarily need to be fixed:

$ grep -r 'torchaudio.save(' recipes/ | wc -l
      63

Looking at these, they mostly happen in places other than on_stage_end so it seems like the recipes are not expected to work on multiple processes.

You need to remember to wrap torch_recovery with main_process_only, and same for checkpoint deletion.

Ah yes, I've added this for checkpoint deletion. However, for torch_recovery I'm pretty confident we don't want it to happen only on main, but rather on all processes, since we use this to load models on all devices (and there's no contested resource).

Oops yea sorry, I meant torch_save

Side note: I was going to nitpick about using sys.exit to crash Brain in case of wrong arguments, but I noticed that this is done for some other options as well and you were using the same convention. I propose a separate PR to change all those to raising an
error like a proper respectful Python class should.

There is an issue with this PR, as reported by @Adel-Moumen:

"I trained a model with 3 GPUs and basically, I got [three] save folders"

and

"they all have:

# yamllint disable
WER: 5.784713797286864
end-of-epoch: true
unixtime: 1691803456.4516888

except one with the right .ckpt. however, the problem is that when loading the new ckpt we are getting an error saying that we are failing to load the .ckpt as there is no ckpt, and indeed, the issue is that the last ckpt (the most recent one) is lacking from this .ckpt"