speed up fsdp

tests/trainer/test_fsdp_checkpoint.py

@@ -83,8 +83,8 @@ def get_trainer(
     save_folder: Optional[str] = None,
     save_filename: str = 'ba{batch}-rank{rank}.pt',
     save_overwrite: bool = False,
-    num_features: int = 4,
-    num_classes: int = 2,
+    num_features: int = 4,  # Reduced from default
+    num_classes: int = 2,   # Reduced from default
     load_path: Optional[str] = None,
     autoresume: bool = False,
     run_name: Optional[str] = None,
@@ -111,11 +111,11 @@ def get_trainer(
         val_metrics=val_metrics,
     )
     model.module.to(model_init_device)
-    dataset = RandomClassificationDataset(shape=(num_features,), num_classes=num_classes, size=128)
+    dataset = RandomClassificationDataset(shape=(num_features,), num_classes=num_classes, size=32)
     dataloader = DataLoader(
         dataset,
         sampler=dist.get_sampler(dataset),
-        batch_size=8,
+        batch_size=2,
     )
     if optimizer == 'adam':
         optim = torch.optim.Adam(params=model.parameters())