fix: only run save full model on main process

[ErwannMillon]

was running into a bug where _save_full_model would break during multi-gpu full finetuning.
The function was being run by all processes, and

        for item in os.listdir(temporary_dir):
            s = os.path.join(temporary_dir, item)
            d = os.path.join(output_dir, item)
            if os.path.isdir(s):
                shutil.copytree(s, d, dirs_exist_ok=True)  # Python 3.8+
            else:
                shutil.copy2(s, d)

        # Remove the temporary directory
        shutil.rmtree(temporary_dir)

would error out because the temporary_dir would get deleted while other processes were trying to copy files that no longer existed.

Fixed by returning early if not main process.

[bghira]

this will break DeepSpeed

[ErwannMillon]

hm, i'm running w/ deepspeed level 1 and still seems to work. The return is only in the save hook, the sharded deepspeed params still get saved on each gpu if i'm not mistaken

[ErwannMillon]

alternatively maybe wrapping only:

        for item in os.listdir(temporary_dir):
            s = os.path.join(temporary_dir, item)
            d = os.path.join(output_dir, item)
            if os.path.isdir(s):
                shutil.copytree(s, d, dirs_exist_ok=True)  # Python 3.8+
            else:
                shutil.copy2(s, d)

        # Remove the temporary directory
        shutil.rmtree(temporary_dir)
        

in the if self.accelerator.is_main_process ?

[bghira]

ah, ok.

[bghira]

oh, but multi-node training will probably want a local copy of the training state on each node, right?

[ErwannMillon]

then maybe is_local_main_process ? just pushed a commit w/ this change