Merge branch 'master' into fix/no-grad-amp-bug

Author: baskrahmer

requirements/pytorch/extra.txt

@@ -5,7 +5,7 @@
 matplotlib>3.1, <3.10.0
 omegaconf >=2.2.3, <2.4.0
 hydra-core >=1.2.0, <1.4.0
-jsonargparse[signatures] >=4.39.0, <4.41.0
+jsonargparse[signatures,jsonnet] >=4.39.0, <4.41.0
 rich >=12.3.0, <14.1.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.45.2,<0.47.0; platform_system != "Darwin"

src/lightning/pytorch/callbacks/pruning.py

@@ -458,7 +458,10 @@ def sanitize_parameters_to_prune(
 
         if not parameters_to_prune:
             parameters_to_prune = [
-                (m, p) for p in parameters for m in current_modules if getattr(m, p, None) is not None
+                (m, p)
+                for p in parameters
+                for m in current_modules
+                if getattr(m, p, None) is not None and isinstance(getattr(m, p, None), nn.Parameter)
             ]
         elif (
             isinstance(parameters_to_prune, (list, tuple))

src/lightning/pytorch/cli.py

@@ -554,7 +554,9 @@ def parse_arguments(self, parser: LightningArgumentParser, args: ArgsType) -> No
     def _dump_config(self) -> None:
         if hasattr(self, "config_dump"):
             return
-        self.config_dump = yaml.safe_load(self.parser.dump(self.config, skip_link_targets=False, skip_none=False))
+        self.config_dump = yaml.safe_load(
+            self.parser.dump(self.config, skip_link_targets=False, skip_none=False, format="yaml")
+        )
         if "subcommand" in self.config:
             self.config_dump = self.config_dump[self.config.subcommand]
 

tests/tests_pytorch/callbacks/test_pruning.py

@@ -338,3 +338,70 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint):
     assert not hasattr(model.layer.mlp_3, "weight_orig")
     model = TestModel.load_from_checkpoint(trainer.checkpoint_callback.last_model_path)
     assert not hasattr(model.layer.mlp_3, "weight_orig")
+
+
+def test_sanitize_parameters_explicit_check():
+    """Test the sanitize_parameters_to_prune method with various attribute types."""
+
+    class TestModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.randn(5, 5))
+            self.bias = nn.Parameter(torch.randn(5))
+            self.some_bool = True
+            self.some_tensor = torch.randn(3, 3)  # Regular tensor, not parameter
+            self.some_string = "test"
+            self.some_none = None
+
+    class TestModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.test_module = TestModule()
+
+    model = TestModel()
+
+    parameters_to_prune = ModelPruning.sanitize_parameters_to_prune(
+        model,
+        parameters_to_prune=(),
+        parameter_names=["weight", "bias", "some_bool", "some_tensor", "some_string", "some_none"],
+    )
+
+    param_names_found = set()
+    for module, param_name in parameters_to_prune:
+        param = getattr(module, param_name)
+        assert isinstance(param, nn.Parameter), f"Expected Parameter, got {type(param)}"
+        param_names_found.add(param_name)
+
+    assert "weight" in param_names_found
+    assert "bias" in param_names_found
+    assert "some_bool" not in param_names_found
+    assert "some_tensor" not in param_names_found
+    assert "some_string" not in param_names_found
+    assert "some_none" not in param_names_found
+
+
+def test_original_issue_reproduction():
+    """Issue: https://github.com/Lightning-AI/pytorch-lightning/issues/10835."""
+
+    class ProblematicModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.layer = Sequential(
+                OrderedDict([
+                    ("mlp_1", nn.Linear(32, 32)),
+                    ("mlp_2", nn.Linear(32, 2)),
+                ])
+            )
+            # Add boolean attributes that would cause the original error
+            self.layer.mlp_1.training = True
+            self.layer.mlp_2.requires_grad = True
+
+    model = ProblematicModel()
+
+    parameters_to_prune = ModelPruning.sanitize_parameters_to_prune(
+        model, parameters_to_prune=(), parameter_names=["weight", "bias", "training", "requires_grad"]
+    )
+
+    for module, param_name in parameters_to_prune:
+        param = getattr(module, param_name)
+        assert isinstance(param, nn.Parameter), f"Non-parameter found: {type(param)}"

tests/tests_pytorch/test_cli.py

@@ -1858,3 +1858,23 @@ def test_lightning_cli_with_args_given(args):
 def test_lightning_cli_args_and_sys_argv_warning():
     with mock.patch("sys.argv", ["", "--model.foo=456"]), pytest.warns(Warning, match="LightningCLI's args parameter "):
         LightningCLI(TestModel, run=False, args=["--model.foo=789"])
+
+
+def test_lightning_cli_jsonnet(cleandir):
+    class MainModule(BoringModel):
+        def __init__(self, main_param: int = 1):
+            super().__init__()
+
+    config = """{
+        "model":{
+            "main_param": 2
+        }
+    }"""
+    config_path = Path("config.jsonnet")
+    config_path.write_text(config)
+
+    cli_args = [f"--config={config_path}"]
+    with mock.patch("sys.argv", ["any.py"] + cli_args):
+        cli = LightningCLI(MainModule, run=False, parser_kwargs={"parser_mode": "jsonnet"})
+
+    assert cli.config["model"]["main_param"] == 2