Method save_checkpoint

train/model/strategy.py:252–290

Save a checkpoint to the `run_dir` only containing the state_dicts for trainable parameters by default.

(
            self,
            run_dir: Path,
            global_step: int,
            epoch: int,
            train_loss: Optional[float] = None,
            only_trainable: bool = True,
        )

                dist.barrier()

    def save_checkpoint(
            self,
            run_dir: Path,
            global_step: int,
            epoch: int,
            train_loss: Optional[float] = None,
            only_trainable: bool = True,
        ) -> None:
            """Save a checkpoint to the `run_dir` only containing the state_dicts for trainable parameters by default."""
            assert isinstance(self.vlm, FSDP), "FSDPStrategy.save_checkpoint assumes VLM is already wrapped in FSDP!"

            # Summon Full State Dictionary =>> Reconstitute from Shards
            with FSDP.state_dict_type(self.vlm, self.fsdp_state_dict_type, self.fsdp_save_policy):
                full_vlm_state_dict = self.vlm.state_dict()
                model_state_dicts = {
                    mkey: OrderedDict() for mkey in (self.trainable_module_keys if only_trainable else self.all_module_keys)
                }

                # Iterate through `full_vlm_state_dict` and split `mkey.{full_dotted_path}` -> `mkey: {full_dotted_path}`
                for key, param in full_vlm_state_dict.items():
                    for mkey in model_state_dicts:
                        if key.startswith(mprefix := f"{mkey}."):
                            model_state_dicts[mkey][key.removeprefix(mprefix)] = param

                # Save on rank zero *only*
                if overwatch.is_rank_zero():
                    checkpoint_dir = run_dir / "checkpoints"
                    if train_loss is None:
                        checkpoint_path = checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss=inf.pt"
                    else:
                        checkpoint_path = (
                            checkpoint_dir / f"step-{global_step:06d}-epoch-{epoch:02d}-loss={train_loss:.4f}.pt"
                        )

                    # Save Checkpoint & Copy Latest to `latest-checkpoint.pt`
                    torch.save({"model": model_state_dicts}, checkpoint_path)

                    # TODO (siddk) :: This breaks w/ Sagemaker default permissions (root vs. <user>)... skip?
                    # shutil.copy(checkpoint_path, checkpoint_dir / "latest-checkpoint.pt")

    def run_setup(self, run_dir: Path, n_train_examples: int) -> None:
        # Iteratively Assemble FSDP Wrapping Policy by fetching the wrapping policies for each backbone/constituent