fix set_model_state_dict _IncompatibleKeys when full_state_dict

YassineYousfi · web-flow · commit b557e69a8b77 · 2025-05-11T14:20:36.000-07:00
fixes #153350
diff --git a/torch/distributed/checkpoint/state_dict.py b/torch/distributed/checkpoint/state_dict.py
@@ -595,7 +595,9 @@ def _load_model_state_dict(
             )
         elif info.full_state_dict:
             _distribute_state_dict(state_dict, local_state_dict, device=devices.pop())
-        state_dict.update(local_state_dict)
+        for fqn, local_state in local_state_dict.items():
+            if fqn in state_dict:
+                state_dict[fqn] = local_state
 
     with info.fsdp_context():
         return cast(
