From 244f8e1d1f640b6e39a29c41086b8187a55c3bef Mon Sep 17 00:00:00 2001 From: Daniel King Date: Wed, 24 Jul 2024 19:00:46 -0700 Subject: [PATCH] oops --- llmfoundry/callbacks/hf_checkpointer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py index 2f858dd186..35508cc0c7 100644 --- a/llmfoundry/callbacks/hf_checkpointer.py +++ b/llmfoundry/callbacks/hf_checkpointer.py @@ -465,8 +465,7 @@ def tensor_hook( hooks = [] for _, module in state_dict_model.named_modules(): - if isinstance(module, FSDP): - hooks.append(module._register_state_dict_hook(tensor_hook),) + hooks.append(module._register_state_dict_hook(tensor_hook),) state_dict = get_model_state_dict( state_dict_model,