From 10b908d04fdf860bd6df6b0e7abe65f3f3c2df72 Mon Sep 17 00:00:00 2001
From: tianyitang <2201534206@qq.com>
Date: Wed, 8 Jul 2026 15:48:53 +0800
Subject: [PATCH] Add the required grads parameter

---
 torchtitan/experiments/graph_trainer/simple_fsdp.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py
index 19516edd64..99aa36c8f3 100644
--- a/torchtitan/experiments/graph_trainer/simple_fsdp.py
+++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py
@@ -297,7 +297,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,
                     ),
                 )