Merge pull request #43 from krasserm/wip-pytorch-2

Upgrade to PyTorch 2.0 and PyTorch Lightning 2.0

environment.yml

@@ -5,6 +5,6 @@ channels:
 dependencies:
   - python=3.9
   - pytorch-cuda=11.7
-  - pytorch=1.13
-  - torchvision=0.14
+  - pytorch=2.0
+  - torchvision=0.15
   - pip>=22

examples/training/clm/train_fsdp.sh

@@ -30,7 +30,6 @@ python -m perceiver.scripts.text.clm_fsdp fit \
   --trainer.precision=bf16 \
   --trainer.max_steps=50000 \
   --trainer.accumulate_grad_batches=1 \
-  --trainer.track_grad_norm=2 \
   --trainer.check_val_every_n_epoch=null \
   --trainer.val_check_interval=500 \
   --trainer.limit_val_batches=20 \

perceiver/scripts/text/clm_fsdp.py

@@ -3,7 +3,8 @@
 
 import torch
 from pytorch_lightning.cli import LightningArgumentParser, LRSchedulerCallable, OptimizerCallable
-from pytorch_lightning.strategies import DDPFullyShardedNativeStrategy, StrategyRegistry
+from pytorch_lightning.strategies import FSDPStrategy, StrategyRegistry
+from pytorch_lightning.utilities import grad_norm
 from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
 
 from perceiver.model.core import CrossAttentionLayer, SelfAttentionLayer
@@ -27,7 +28,7 @@
 
 StrategyRegistry.register(
     name="fsdp_perceiver_ar",
-    strategy=DDPFullyShardedNativeStrategy,
+    strategy=FSDPStrategy,
     description="FSDP strategy optimized for Perceiver AR models",
     activation_checkpointing=[CrossAttentionLayer, SelfAttentionLayer],
     auto_wrap_policy=policy,
@@ -60,9 +61,10 @@ def configure_optimizers(self):
             "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
         }
 
-    def on_before_optimizer_step(self, optimizer, optimizer_idx):
+    def on_before_optimizer_step(self, optimizer):
         if self.hparams.max_grad_norm is not None:
             self.trainer.model.clip_grad_norm_(self.hparams.max_grad_norm)
+            self.log_dict(grad_norm(self, norm_type=2))
 
 
 class CausalLanguageModelCLI(CLI):