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

Fix : get_balanced_memory when using multi gpus with small models or quantized models with a large vocabulary #3244

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+7 −1
Open

Fix : get_balanced_memory when using multi gpus with small models or quantized models with a large vocabulary #3244

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7b85592
first fix
MekkCyber Nov 18, 2024
c5eb33c
fix style
MekkCyber Nov 18, 2024
a7ea07c
max leave size instead of embed layer
MekkCyber Nov 19, 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
8 changes: 7 additions & 1 deletion src/accelerate/utils/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1065,8 +1065,14 @@ def get_balanced_memory(
)
)
# The last device is left with max_memory just in case the buffer is not enough.
max_leave_size = max([module_sizes[leave] for leave in leaves])
for idx in gpus_idx_list[:-1]:
max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_gpu, max_memory[idx])
if idx == 0 and not low_zero and max_leave_size > per_gpu * 0.9:
max_memory[idx] = min(max_leave_size * 1.3, max_memory[idx])
elif idx == 1 and low_zero and max_leave_size > per_gpu * 0.9:
max_memory[idx] = min(max_leave_size * 1.3, max_memory[idx])
Comment on lines +1070 to +1073
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are taking the minimum, is this expected ?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes we take the minimum between the max_memory of the gpu and the gpu space needed on the device. So if the space needed exceeds the space available on the gpu, we only allocate the space available.

else:
max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_gpu, max_memory[idx])

if low_zero:
min_zero = max(0, module_sizes[""] - sum([max_memory[i] for i in range(1, num_devices)]))
Expand Down
Loading