Dtensor oom (#1395)

mosaicml · Jul 25, 2024 · e882658 · e882658
1 parent cfab70e
commit e882658
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -455,9 +455,9 @@ def tensor_hook(
                         state_dict[fqn] = tensor
                     else:
                         state_dict[fqn] = None
-                # Convert the state dict to the requested precision
-                if isinstance(tensor, torch.Tensor):
-                    state_dict[fqn] = tensor.to(dtype=self.dtype)
+
+                if isinstance(state_dict[fqn], torch.Tensor):
+                    state_dict[fqn] = state_dict[fqn].to(dtype=self.dtype)
                 del tensor
             if dist.get_global_rank() != 0:
                 state_dict = {}