Update on "[FSDP][optim_state_dict][2/N] Add _get_fqn_to_fsdp_param_info to map from original FQN to flat_param"

fegin · fegin · commit 97ff6505d9d2 · 2022-12-07T17:10:24.000Z
**Motivation:**
Add a helper to map from the FQN to the corresponding flat_param. The helper will directly get flat_param from fsdp_state and flat_handler as flat_param is not registered to the module if `use_orig_params` is True.


[ghstack-poisoned]
diff --git a/torch/distributed/fsdp/_optim_utils.py b/torch/distributed/fsdp/_optim_utils.py
@@ -1293,11 +1293,12 @@ def _optim_state_dict(
 
 def _get_fqn_to_fsdp_param_info(model: nn.Module) -> Dict[str, FSDPParamInfo]:
     """
-    Construct the maaping from a param's fqn to its corresponding FSDPParamInfo
-    if the param is managed by FSDP. FlatParam only stores the first FQN of a
-    shared parameter. So the keys in the mapping are guranteed to map to unique
-    parameters.
+    Construct the mapping from a param's fqn to its corresponding ``FSDPParamInfo``
+    if the param is managed by FSDP. ``FlatParameter._fqns`` only stores the first
+    FQN of a shared parameter. So the keys in the mapping are guaranteed to map
+    to unique parameters.
     """
+
     def module_fn(module, prefix, fqn_to_param_info):
         # TODO: make it work with composable API.
         if not isinstance(module, fsdp_file.FullyShardedDataParallel):
@@ -1319,6 +1320,9 @@ def return_fn(fqn_to_param_info):
         return fqn_to_param_info
 
     fqn_to_param_info: Dict[str, FSDPParamInfo] = {}
+    # FlatParameter._fqns stores the local fqn, starting from the root of the
+    # FSDP. Using _apply_to_modules() with model (may not be the FSDP root
+    # module) allows us to construct the global fqn.
     return _apply_to_modules(
         model,
         module_fn,
