From 505fa02d1920ca2fb8362e856567c76ed8830b01 Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Sat, 22 Aug 2026 10:20:37 +0100
Subject: [PATCH] Preserve requires_grad in SimpleFSDP parameters

---
 .../experiments/graph_trainer/simple_fsdp.py  |  3 ++-
 .../graph_trainer/tests/test_simple_fsdp.py   | 26 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py
index c108da4f25..059a612cf6 100644
--- a/torchtitan/experiments/graph_trainer/simple_fsdp.py
+++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py
@@ -288,7 +288,8 @@ def data_parallel(
                 mod.register_parameter(
                     p_name,
                     nn.Parameter(
-                        distribute_tensor_func(p, device_mesh, param_sharding)
+                        distribute_tensor_func(p, device_mesh, param_sharding),
+                        requires_grad=p.requires_grad,
                     ),
                 )
 
diff --git a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
index dd9379580c..4464f7cc68 100644
--- a/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
+++ b/torchtitan/experiments/graph_trainer/tests/test_simple_fsdp.py
@@ -75,6 +75,32 @@ def test_param_cast_to_bf16_at_ngpu_1(self):
         y = model(x)
         self.assertEqual(y.dtype, torch.bfloat16)
 
+    @patch("torchtitan.distributed.parallel_dims.device_type", "cpu")
+    def test_preserves_parameter_requires_grad(self):
+        parallel_dims = ParallelDims(
+            dp_replicate=1,
+            dp_shard=1,
+            cp=1,
+            tp=1,
+            pp=1,
+            ep=1,
+            world_size=1,
+            spmd_backend="partial_dtensor",
+        )
+        model = nn.Linear(8, 8)
+        model.weight.requires_grad_(False)
+
+        model = apply_simple_fsdp(
+            model,
+            parallel_dims=parallel_dims,
+            training=TrainingConfig(),
+        )
+
+        self.assertFalse(model._parameters["weight"].requires_grad)
+        self.assertTrue(model._parameters["bias"].requires_grad)
+        self.assertFalse(model.weight.requires_grad)
+        self.assertTrue(model.bias.requires_grad)
+
 
 if __name__ == "__main__":
     unittest.main()