Commit 7fc877b
Adjust torch.compile() best practices (#3336)
* Adjust torch.compile() best practices
1. Add best practice to prefer `mod.compile` over `torch.compile(mod)`, which avoids `_orig_` naming problems.
Repro steps:
- opt_mod = torch.compile(mod)
- train opt_mod
- save checkpoint
In another script, potentially on a machine that does NOT support `torch.compile`: load checkpoint.
This fails with an error, because the checkpoint on `opt_mod` got its params renamed by `torch.compile`:
```
RuntimeError: Error(s) in loading state_dict for VQVAE:
Missing key(s) in state_dict: "embedding.weight", "encoder.encoder.net.0.weight", "encoder.encoder.net.0.bias", ...
Unexpected key(s) in state_dict: "_orig_mod.embedding.weight", "_orig_mod.encoder.encoder.net.0.weight", "_orig_mod.encoder.encoder.net.0.bias", ...
```
- Add best practice to use, or at least try, `fullgraph=True`. This doesn't always work, but we should encourage it.
---------
Co-authored-by: Svetlana Karslioglu <svekars@meta.com>1 parent a5632da commit 7fc877b
1 file changed
Lines changed: 13 additions & 4 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
101 | 101 | | |
102 | 102 | | |
103 | 103 | | |
104 | | - | |
105 | | - | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
106 | 109 | | |
107 | 110 | | |
108 | 111 | | |
| |||
135 | 138 | | |
136 | 139 | | |
137 | 140 | | |
138 | | - | |
139 | | - | |
| 141 | + | |
| 142 | + | |
140 | 143 | | |
141 | 144 | | |
142 | 145 | | |
| |||
197 | 200 | | |
198 | 201 | | |
199 | 202 | | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
200 | 209 | | |
201 | 210 | | |
202 | 211 | | |
| |||
0 commit comments