Future 1/n: package in src/ folder

src/pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -22,6 +22,8 @@
 if _FAIRSCALE_AVAILABLE:
     from fairscale.optim import OSS
     from fairscale.optim.grad_scaler import ShardedGradScaler
+else:
+    OSS = ShardedGradScaler = object
 
 
 class ShardedNativeMixedPrecisionPlugin(NativeMixedPrecisionPlugin):

src/pytorch_lightning/strategies/collaborative.py

@@ -26,6 +26,8 @@
 
 if _HIVEMIND_AVAILABLE:
     import hivemind
+else:
+    hivemind = None
 
 log = logging.getLogger(__name__)
 

src/pytorch_lightning/strategies/ddp.py

@@ -64,6 +64,8 @@
 
 if _FAIRSCALE_AVAILABLE:
     from fairscale.optim import OSS
+else:
+    OSS = object
 if _TORCH_GREATER_EQUAL_1_10 and torch.distributed.is_available():
     from torch.distributed.algorithms.model_averaging.averagers import ModelAverager
 

src/pytorch_lightning/strategies/sharded.py

@@ -33,6 +33,8 @@
     from fairscale.optim import OSS
 
     from pytorch_lightning.overrides.fairscale import LightningShardedDataParallel, unwrap_lightning_module_sharded
+else:
+    OSS = ShardedDataParallel = object
 
 
 class DDPShardedStrategy(DDPStrategy):

src/pytorch_lightning/strategies/sharded_spawn.py

@@ -31,6 +31,8 @@
     from fairscale.optim import OSS
 
     from pytorch_lightning.overrides.fairscale import LightningShardedDataParallel, unwrap_lightning_module_sharded
+else:
+    OSS = ShardedDataParallel = object
 
 
 class DDPSpawnShardedStrategy(DDPSpawnStrategy):