Update

fegin · fegin · commit a4d47b18fb32 · 2024-05-03T14:11:34.000-07:00
[ghstack-poisoned]
diff --git a/torch/distributed/checkpoint/state_dict.py b/torch/distributed/checkpoint/state_dict.py
@@ -76,7 +76,6 @@ def gc_context():
         yield
     finally:
         # TODO: add logging for the gc details/time
-        gc.collect()
         if is_enabled:
             gc.enable()
 
@@ -133,6 +132,7 @@ class _StateDictInfo(StateDictOptions):
     fsdp_modules: List[nn.Module] = field(default_factory=list)
 
 
+@functools.lru_cache(maxsize=None)
 def _get_fqns(
     model: nn.Module,
     name: str,
