Skip to content

Commit b024677

Browse files
aporialiaofacebook-github-bot
authored andcommitted
Fix optimizer check edge case (#3058)
Summary: There's an edge case for detecting an optimizer - where no optimizer is applied during backwards, but the compute kernel is fused, so techncially an EBC.optimizer would be created. In this case for dynamic sharding, we still don't need to store any optimizers - I've added the condition to properly check for optimizers to pass the unit test. Differential Revision: D76203188
1 parent 558f476 commit b024677

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

torchrec/distributed/embeddingbag.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1533,7 +1533,9 @@ def update_shards(
15331533
return
15341534

15351535
current_state = self.state_dict()
1536-
has_optimizer = len(self._optim._optims) > 0
1536+
has_optimizer = len(self._optim._optims) > 0 and all(
1537+
len(i) > 0 for i in self._optim.state_dict()["state"].values()
1538+
)
15371539

15381540
# TODO: Saving lookups tensors to CPU to eventually avoid recreating them completely again
15391541
# TODO: Ensure lookup tensors are actually being deleted

0 commit comments

Comments
 (0)