From 46956cffc89dc2a321e0fc39d23cc83722b90530 Mon Sep 17 00:00:00 2001 From: NouamaneTazi Date: Mon, 25 Nov 2024 09:40:58 +0000 Subject: [PATCH] quick fix check_optim_state_in_sync --- src/nanotron/serialize/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nanotron/serialize/main.py b/src/nanotron/serialize/main.py index e87d8dbb..e9ad04d8 100644 --- a/src/nanotron/serialize/main.py +++ b/src/nanotron/serialize/main.py @@ -149,7 +149,7 @@ def save( tensor=tied_param, pg=group, msg=lambda err: f"Tied {tied_info.name} are not synced {err}" ) if not optimizer.inherit_from(optim.ZeroDistributedOptimizer): - check_optim_state_in_sync(optimizer, parallel_context.dp_pg) + check_optim_state_in_sync(optimizer.state_dict(), parallel_context.dp_pg) # SANITY CHECK: tied parameters have their optimizer states synchronized # Compute a mapping from id_ to index in the optimizer sense