[BugFix] FSDP2 SLURM restart

Author: RandallBalestriero

examples/imagenet1k_supervised_vit_fsdp2.py

@@ -50,30 +50,53 @@
 
 
 class TopkAccuracy(pl.Callback):
-    """Validation top-1 accuracy on the trained head's logits (``batch["logits"]``)."""
+    """Validation top-1 / top-5 accuracy on the head's logits (``batch["logits"]``).
+
+    Uses ``average="micro"`` (correct / total) — the standard ImageNet top-1
+    every paper reports. ``torchmetrics``'s
+    :class:`~torchmetrics.classification.MulticlassAccuracy` defaults to
+    ``average="macro"`` (mean of per-class recall); on the full balanced
+    ImageNet val set macro and micro nearly coincide, but micro is the
+    unambiguous convention, so we pin it and add top-5 alongside.
+    """
 
     def __init__(self):
         super().__init__()
-        self.acc = torchmetrics.classification.MulticlassAccuracy(NUM_CLASSES)
+        MCA = torchmetrics.classification.MulticlassAccuracy
+        self.top1 = MCA(NUM_CLASSES, top_k=1, average="micro")
+        self.top5 = MCA(NUM_CLASSES, top_k=5, average="micro")
 
     def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
-        self.acc.to(pl_module.device)
-        self.acc.update(batch["logits"], batch["label"].long())
+        self.top1.to(pl_module.device)
+        self.top5.to(pl_module.device)
+        self.top1.update(batch["logits"], batch["label"].long())
+        self.top5.update(batch["logits"], batch["label"].long())
 
     def on_validation_epoch_end(self, trainer, pl_module):
-        pl_module.log("val/top1", self.acc.compute(), prog_bar=True, sync_dist=True)
-        self.acc.reset()
+        pl_module.log("val/top1", self.top1.compute(), prog_bar=True, sync_dist=True)
+        pl_module.log("val/top5", self.top5.compute(), prog_bar=True, sync_dist=True)
+        self.top1.reset()
+        self.top5.reset()
 
 
-def build_loaders(batch_size, num_workers):
+def build_loaders(batch_size, num_workers, use_randaug=True, cpu_norm=False):
     # CPU does the bare minimum (decode + square resize) so the GPU aug pipeline
     # is the bottleneck-free fast path. Train images stay un-normalized [0,1]
     # floats for the kornia GPU ops; val is fully prepared on CPU.
-    train_cpu = transforms.Compose(
-        transforms.RGB(),
-        transforms.Resize((256, 256)),
-        transforms.ToImage(),  # -> float tensor in [0, 1], no normalization
-    )
+    if cpu_norm:
+        # debug path: full CPU pipeline producing normalized 224 crops, no gpu_transform
+        train_cpu = transforms.Compose(
+            transforms.RGB(),
+            transforms.RandomResizedCrop((224, 224), scale=(0.08, 1.0)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.ToImage(mean=_MEAN, std=_STD),
+        )
+    else:
+        train_cpu = transforms.Compose(
+            transforms.RGB(),
+            transforms.Resize((256, 256)),
+            transforms.ToImage(),  # [0,1], GPU pipeline normalizes
+        )
     val_cpu = transforms.Compose(
         transforms.RGB(),
         transforms.Resize((256, 256)),
@@ -84,14 +107,18 @@ def build_loaders(batch_size, num_workers):
         path="ILSVRC/imagenet-1k", split="train", transform=train_cpu
     )
     # Heavy augmentation, batched on GPU (DeiT/AugReg policy).
-    train_ds.gpu_transform = GPUCompose(
+    aug = [
         ToDevice(),
         GPURandomResizedCrop(224, scale=(0.08, 1.0)),
         GPURandomHorizontalFlip(p=0.5),
-        GPURandAugment(n=2, m=9),
-        GPUNormalize(mean=_MEAN, std=_STD),
-        GPURandomErasing(p=0.25),
-    )
+    ]
+    if use_randaug:
+        aug.append(GPURandAugment(n=2, m=9))
+    aug.append(GPUNormalize(mean=_MEAN, std=_STD))
+    if use_randaug:
+        aug.append(GPURandomErasing(p=0.25))
+    if not cpu_norm:
+        train_ds.gpu_transform = GPUCompose(aug)
     val_ds = spt.data.HFDataset(
         path="ILSVRC/imagenet-1k", split="validation", transform=val_cpu
     )
@@ -113,20 +140,43 @@ def build_loaders(batch_size, num_workers):
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--backbone", default="vit_large_patch16_224")
+    ap.add_argument("--strategy", default="fsdp2", help="fsdp2 | ddp | auto (debug)")
+    ap.add_argument("--devices", default="auto", help="'auto' or an int (debug)")
     ap.add_argument("--epochs", type=int, default=300)
+    ap.add_argument(
+        "--max-steps", type=int, default=0, help=">0 caps steps (smoke test)"
+    )
     ap.add_argument("--batch-size", type=int, default=64, help="per-GPU batch")
     ap.add_argument("--lr", type=float, default=1e-3)
     ap.add_argument("--weight-decay", type=float, default=0.05)
     ap.add_argument("--drop-path", type=float, default=0.4)  # DeiT-III ViT-L
     ap.add_argument("--label-smoothing", type=float, default=0.1)
     ap.add_argument("--num-workers", type=int, default=12)
+    ap.add_argument("--no-randaug", action="store_true", help="debug: drop RandAugment")
+    ap.add_argument("--no-mixup", action="store_true", help="debug: drop Mixup/CutMix")
+    ap.add_argument(
+        "--cpu-norm",
+        action="store_true",
+        help="debug: normalize on CPU, no gpu_transform",
+    )
+    ap.add_argument(
+        "--overfit-batches",
+        type=int,
+        default=0,
+        help="debug: Lightning overfit_batches",
+    )
     args = ap.parse_args()
 
     # Speed knobs (TF32 + autotuned cuDNN); bf16 comes from the Trainer.
     torch.set_float32_matmul_precision("high")
     torch.backends.cudnn.benchmark = True
 
-    train_loader, val_loader = build_loaders(args.batch_size, args.num_workers)
+    train_loader, val_loader = build_loaders(
+        args.batch_size,
+        args.num_workers,
+        use_randaug=not args.no_randaug,
+        cpu_norm=args.cpu_norm,
+    )
     data = spt.data.DataModule(train=train_loader, val=val_loader)
 
     backbone = getattr(spt.backbone, args.backbone)(
@@ -141,9 +191,18 @@ def main():
 
     def forward(self, batch, stage):
         if self.training:
-            images, soft = mixup(batch["image"], batch["label"])
+            if args.no_mixup:
+                images, target = batch["image"], batch["label"]
+            else:
+                images, target = mixup(batch["image"], batch["label"])
             logits = self.backbone(images)
-            batch["loss"] = F.cross_entropy(logits, soft)
+            loss = F.cross_entropy(logits, target)
+            batch["loss"] = loss
+            # Log train loss so we can actually see whether the backbone learns
+            # (CE starts near ln(1000)=6.9; flat there = frozen, dropping = learning).
+            self.log(
+                "train/loss", loss.detach(), prog_bar=True, on_step=True, sync_dist=True
+            )
         else:
             batch["logits"] = self.backbone(batch["image"])
         return batch
@@ -162,17 +221,22 @@ def forward(self, batch, stage):
         },
     )
 
-    trainer = dict(
-        strategy="fsdp2",
+    # Pass a real pl.Trainer (not a config dict): the Manager wraps a dict
+    # trainer in OmegaConf, which can't hold callback *instances*
+    # (UnsupportedValueType). Callback instances are fine on a built Trainer.
+    devices = args.devices if args.devices == "auto" else int(args.devices)
+    trainer = pl.Trainer(
+        strategy=args.strategy,
         precision="bf16-mixed",
         accelerator="gpu",
-        devices="auto",
+        devices=devices,
         max_epochs=args.epochs,
+        max_steps=args.max_steps if args.max_steps > 0 else -1,
+        overfit_batches=args.overfit_batches if args.overfit_batches > 0 else 0.0,
         callbacks=[
             pl.pytorch.callbacks.LearningRateMonitor(logging_interval="step"),
             TopkAccuracy(),
         ],
-        gradient_clip_val=1.0,
         num_sanity_val_steps=0,
         enable_checkpointing=True,
     )

stable_pretraining/manager.py

@@ -528,9 +528,13 @@ def __init__(
                     "regardless of where the process is launched from."
                 )
             p = p.with_suffix(".ckpt")
-            if not p.is_file():
+            # ``exists()`` not ``is_file()``: a distributed (FSDP2) checkpoint
+            # is a directory of ``*.distcp`` shards, not a single file. Accept
+            # both so an FSDP2 ``last.ckpt`` can be passed as a fresh-run
+            # ``ckpt_path`` too.
+            if not p.exists():
                 raise FileNotFoundError(
-                    f"`ckpt_path` was set to {p} but no such file exists. "
+                    f"`ckpt_path` was set to {p} but no such file/dir exists. "
                     "Refusing to silently start training from scratch."
                 )
             ckpt_path = p
@@ -1243,7 +1247,13 @@ def _resolve_load_path(self, run_dir: Path) -> tuple[Optional[str], Optional[boo
                     "picks up exactly where it left off. The user ckpt_path "
                     "is only consumed on the FIRST (fresh) invocation."
                 )
-            if not last_ckpt.is_file():
+            # ``exists()`` not ``is_file()``: an FSDP2 *distributed* checkpoint
+            # is saved as a DIRECTORY of ``*.distcp`` shards (one per rank),
+            # while a plain checkpoint is a single file. Both are valid
+            # ``last.ckpt`` forms — gating on ``is_file()`` made every sharded
+            # (FSDP2 / save_distributed_checkpoint=True) run refuse to resume
+            # on requeue even though the checkpoint was written correctly.
+            if not last_ckpt.exists():
                 raise RuntimeError(
                     f"REQUEUE but no last.ckpt to resume from at {last_ckpt}. "
                     "The original run was preempted before saving its first "

stable_pretraining/tests/unit/test_manager.py

@@ -502,3 +502,67 @@ def test_no_label(self):
     def test_with_label(self):
         # Should not raise; label is purely cosmetic.
         print_signal_info("post-fit")
+
+
+@pytest.mark.unit
+class TestDistributedCheckpointResume:
+    """FSDP2 distributed checkpoints are directories, not single files.
+
+    They are saved as a directory of ``*.distcp`` shards (one per rank).
+    Regression for a bug where requeue-resume and fresh-run ``ckpt_path``
+    validation both used ``Path.is_file()``, which is ``False`` for the
+    directory form — so every sharded (FSDP2 / ``save_distributed_checkpoint``)
+    run refused to resume on SLURM requeue even though ``last.ckpt`` was on disk.
+    """
+
+    @staticmethod
+    def _make_distributed_ckpt(path: Path) -> Path:
+        """Create a fake distributed checkpoint (a directory of shards)."""
+        path.mkdir(parents=True)
+        (path / "__0_0.distcp").touch()
+        (path / ".metadata").touch()
+        return path
+
+    def test_requeue_accepts_directory_last_ckpt(
+        self, manager_factory, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        run_dir = tmp_path / "run"
+        last_ckpt = self._make_distributed_ckpt(run_dir / "checkpoints" / "last.ckpt")
+        manager = manager_factory(
+            callbacks=[], ckpt_path=None, trainer_enable_checkpointing=True
+        )
+        monkeypatch.setattr(
+            "stable_pretraining.manager._is_slurm_requeue", lambda: True
+        )
+        path, weights_only = manager._resolve_load_path(run_dir)
+        assert path == str(last_ckpt)
+        # Requeue always full-restores (optimizer/scheduler/RNG).
+        assert weights_only is False
+
+    def test_requeue_missing_last_ckpt_still_raises(
+        self, manager_factory, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ):
+        run_dir = tmp_path / "run"
+        (run_dir / "checkpoints").mkdir(parents=True)  # no last.ckpt at all
+        manager = manager_factory(
+            callbacks=[], ckpt_path=None, trainer_enable_checkpointing=True
+        )
+        monkeypatch.setattr(
+            "stable_pretraining.manager._is_slurm_requeue", lambda: True
+        )
+        with pytest.raises(RuntimeError, match="no last.ckpt"):
+            manager._resolve_load_path(run_dir)
+
+    def test_init_accepts_directory_ckpt_path(self, tmp_path: Path):
+        """A distributed-checkpoint directory is a valid fresh-run ``ckpt_path``."""
+        ckpt_dir = self._make_distributed_ckpt(tmp_path / "last.ckpt")
+        trainer = BoringTrainer(
+            default_root_dir=str(tmp_path), enable_checkpointing=True
+        )
+        manager = Manager(
+            trainer=trainer,
+            module=BoringModule(),
+            data=BoringDataModule(),
+            ckpt_path=str(ckpt_dir),
+        )
+        assert manager.ckpt_path == ckpt_dir