Merge pull request #25 from MotifTechnologies/fix/invalidate_cache_adamw

torch-ext/optimizer/muon.py

@@ -8,7 +8,7 @@ import torch.distributed as dist
 from torch.distributed.tensor import DTensor, Replicate, Shard
 from torch.profiler import record_function
 
-from .adamw import step_adamw
+from .adamw import _placement_cache, _tensor_cache, step_adamw
 from .async_utils import run_pipeline
 from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho,
                    get_default_muon_param_groups, is_expert_param, update_p)
@@ -1058,6 +1058,11 @@ class Muon(torch.optim.Optimizer):
 
         super().load_state_dict(state_dict)
 
+        # Invalidate adamw.py's module-level tensor caches so that
+        # the next step rebuilds them with the newly loaded state tensors.
+        _placement_cache.clear()
+        _tensor_cache.clear()
+
         if self.cpu_offload:
             # Re-create the offload pool since state tensors may be new
             # objects after load_state_dict.