Merge branch 'master' into feat/fsdp

pytorch_lightning/accelerators/accelerator.py

@@ -322,8 +322,7 @@ def clip_gradients(
         gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
     ) -> None:
         """clips all the optimizer parameters to the given value"""
-
-        self.precision_plugin.clip_gradients(self.model, optimizer, clip_val)
+        self.precision_plugin.clip_gradients(self.model, optimizer, clip_val, gradient_clip_algorithm)
 
     def on_train_epoch_end(self, outputs: Sequence[_STEP_OUTPUT_TYPE]) -> None:
         """Hook to do something on the end of an training epoch

pytorch_lightning/plugins/precision/deepspeed_precision.py

@@ -77,7 +77,11 @@ def backward(
         return closure_loss
 
     def clip_gradients(
-        self, model: Any, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0
+        self,
+        model: 'LightningModule',
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
     ) -> None:
         """
         DeepSpeed handles clipping gradients via the training type plugin.

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -107,9 +107,13 @@ def post_optimizer_step(self, optimizer: 'Optimizer', optimizer_idx: int) -> Non
         """Hook to do something after each optimizer step."""
 
     def clip_gradients(
-        self, model: Any, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0
+        self,
+        model: 'LightningModule',
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
     ) -> None:
-        """Clips the gradients to a specific value"""
+        """Clips the gradients"""
         if clip_val is None:
             return
 

pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -32,10 +32,11 @@ def __init__(self) -> None:
         super().__init__()
         self.scaler = ShardedGradScaler()
 
-    def clip_gradients(
-        self, model: Any, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0
+    def clip_grad_by_norm(
+        self,
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        norm_type: float = 2.0
     ) -> None:
-        if clip_val <= 0:
-            return
         optimizer = cast(OSS, optimizer)
         optimizer.clip_grad_norm(clip_val, norm_type=norm_type)

pytorch_lightning/plugins/training_type/ddp.py

@@ -270,9 +270,6 @@ def init_ddp_connection(self, global_rank: int, world_size: int) -> None:
             torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)
 
     def pre_dispatch(self):
-        if self.sync_batchnorm:
-            self.model = self.configure_sync_batchnorm(self.model)
-
         if self.move_to_device_in_prefetch:
             # move the model to the correct device
             self.model_to_device()

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -259,7 +259,8 @@ def use_dp(self) -> bool:
     def use_ddp(self) -> bool:
         return self._distrib_type in (
             DistributedType.DDP, DistributedType.DDP_SPAWN, DistributedType.DDP_SHARDED,
-            DistributedType.DDP_SHARDED_SPAWN, DistributedType.FULLY_SHARDED, DistributedType.DEEPSPEED
+            DistributedType.DDP_SHARDED_SPAWN, DistributedType.FULLY_SHARDED, DistributedType.DEEPSPEED,
+            DistributedType.TPU_SPAWN
         )
 
     @property

requirements/extra.txt

@@ -8,3 +8,5 @@ torchtext>=0.5
 onnxruntime>=1.3.0
 hydra-core>=1.0
 fairscale>=0.3.2
+jsonargparse[signatures]>=3.3.1
+deepspeed>=0.3.13