match_prefix for test_fsdp against compiled state dict (#190)

Summary:
Fix test/test_fsdp.py to match correct state dict entries

Pull Request resolved: #190

Reviewed By: bdhirsh

Differential Revision: D52846095

Pulled By: drisspg

fbshipit-source-id: e094575525820be91026bb445406234fb5ca33f1

Author: drisspg

test/test_fsdp.py

@@ -216,6 +216,9 @@ def forward_backward():
         sd_out_single_gpu = torch.load(sd_out_single_gpu_fname)
         sd_out_fsdp = torch.load(sd_out_fsdp_fname)
         for k, v1 in sd_out_single_gpu.items():
+            if compile_fsdp:
+                # The state-dict for compiled fsdp has a `_orig_mod` prefix
+                k = f"_orig_mod.{k}"
             v2 = sd_out_fsdp[k]
             v1, v2 = v1.cpu(), v2.cpu()
             if is_fp8 and "noop" in k: