Skip to content

Commit 6839657

Browse files
minhua-chenfacebook-github-bot
authored andcommitted
Make iter persistent for AdagradW
Summary: Make iter persistent for AdagradW optimizer state saving. This is to avoid potential loss of the iter information when training is restarted. Differential Revision: D74717848
1 parent cbc7dbb commit 6839657

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1359,6 +1359,13 @@ def __init__( # noqa C901
13591359
OptimType.EMAINPLACE_ROWWISE_ADAGRAD,
13601360
)
13611361
or self._used_rowwise_adagrad_with_global_weight_decay
1362+
or (
1363+
self._used_rowwise_adagrad_with_counter
1364+
and self.optimizer_args.regularization_mode
1365+
== WeightDecayMode.COUNTER.value
1366+
and self.optimizer_args.weight_decay_mode
1367+
== CounterWeightDecayMode.ADAGRADW.value
1368+
)
13621369
):
13631370
self.register_buffer(
13641371
"iter",

0 commit comments

Comments
 (0)