Address feedback for Fabric.init_module() (2/4) (#17722)

Author: awaelchli

src/lightning/fabric/fabric.py

@@ -898,7 +898,7 @@ def _validate_setup_dataloaders(self, dataloaders: Sequence[DataLoader]) -> None
 @contextmanager
 def _old_sharded_model_context(strategy: Strategy) -> Generator:
     if isinstance(strategy, _Sharded):
-        with strategy.init_sharded_context():
+        with strategy.module_sharded_context():
             yield
     else:
         yield

src/lightning/fabric/strategies/deepspeed.py

@@ -340,7 +340,7 @@ def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
         raise NotImplementedError(self._err_msg_joint_setup_required())
 
     @contextmanager
-    def init_sharded_context(self) -> Generator[None, None, None]:
+    def module_sharded_context(self) -> Generator[None, None, None]:
         # Current limitation in Fabric: The config needs to be fully determined at the time of calling the context
         # manager, which happens at the start of `Fabric.run()`. Later modifications through e.g. `Fabric.setup()`
         # won't have an effect here.

src/lightning/fabric/strategies/fsdp.py

@@ -260,7 +260,7 @@ def init_context(self) -> Generator[None, None, None]:
             yield
 
     @contextmanager
-    def init_sharded_context(self) -> Generator:
+    def module_sharded_context(self) -> Generator:
         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
         from torch.distributed.fsdp.wrap import enable_wrap
 

src/lightning/fabric/strategies/strategy.py

@@ -396,7 +396,7 @@ class _Sharded(ABC):
 
     @abstractmethod
     @contextmanager
-    def init_sharded_context(self) -> Generator:
+    def module_sharded_context(self) -> Generator:
         """A context manager that goes over the instantiation of an :class:`torch.nn.Module` and handles sharding
         of parameters on creation.
 

tests/tests_fabric/test_fabric.py

@@ -718,15 +718,15 @@ def test_module_sharding_context():
     """Test that the sharding context manager gets applied when the strategy supports it and is a no-op
     otherwise."""
     fabric = Fabric()
-    fabric._strategy = MagicMock(spec=DDPStrategy, init_sharded_context=Mock())
+    fabric._strategy = MagicMock(spec=DDPStrategy, module_sharded_context=Mock())
     with pytest.warns(DeprecationWarning, match="sharded_model"), fabric.sharded_model():
         pass
-    fabric._strategy.init_sharded_context.assert_not_called()
+    fabric._strategy.module_sharded_context.assert_not_called()
 
     fabric._strategy = MagicMock(spec=_Sharded)
     with pytest.warns(DeprecationWarning, match="sharded_model"), fabric.sharded_model():
         pass
-    fabric._strategy.init_sharded_context.assert_called_once()
+    fabric._strategy.module_sharded_context.assert_called_once()
 
 
 def test_init_module_context(monkeypatch):