Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dtensor oom #1395

Merged
merged 10 commits into from
Jul 25, 2024
Merged

Dtensor oom #1395

Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
56e93f9
try if statement
dakinggg Jul 25, 2024
aea32d6
test
dakinggg Jul 25, 2024
e4c9691
put back
dakinggg Jul 25, 2024
3de4088
go back
dakinggg Jul 25, 2024
bff6de3
try
dakinggg Jul 25, 2024
04bbfb9
try
dakinggg Jul 25, 2024
96e5194
try
dakinggg Jul 25, 2024
710a3ed
cleaner
dakinggg Jul 25, 2024
b7c44ef
clean up
dakinggg Jul 25, 2024
244f8e1
oops
dakinggg Jul 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions llmfoundry/callbacks/hf_checkpointer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down
Loading