Skip to content

Commit d5ac1fe

Browse files
authored
fix per_gpu args (#1809)
1 parent a5cfaaf commit d5ac1fe

File tree

2 files changed

+7
-7
lines changed

2 files changed

+7
-7
lines changed

mindnlp/accelerate/utils/modeling.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ def get_balanced_memory(
213213
break # only one device
214214

215215
module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes)
216-
per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
216+
per_device = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
217217

218218
# We can't just set the memory to model_size // num_devices as it will end being too small: each GPU will get
219219
# slightly less layers and some layers will end up offload at the end. So this function computes a buffer size to
@@ -251,7 +251,7 @@ def get_balanced_memory(
251251
leaves = get_module_leaves(module_sizes)
252252
mean_leaves = int(sum(module_sizes[n] for n in leaves) / max(len(leaves), 1))
253253
buffer = int(1.25 * max(buffer, mean_leaves))
254-
per_gpu += buffer
254+
per_device += buffer
255255

256256
# Sorted list of GPUs id (we may have some gpu ids not included in the our max_memory list - let's ignore them)
257257
gpus_idx_list = list(
@@ -261,7 +261,7 @@ def get_balanced_memory(
261261
)
262262
# The last device is left with max_memory just in case the buffer is not enough.
263263
for idx in gpus_idx_list[:-1]:
264-
max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_gpu, max_memory[idx])
264+
max_memory[idx] = min(max_memory[0] if low_zero and idx == 0 else per_device, max_memory[idx])
265265

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

mindnlp/engine/train_args/base.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1164,8 +1164,8 @@ def __str__(self):
11641164

11651165
# Remove deprecated arguments. That code should be removed once
11661166
# those deprecated arguments are removed from TrainingArguments. (TODO: v5)
1167-
del self_as_dict["per_gpu_train_batch_size"]
1168-
del self_as_dict["per_gpu_eval_batch_size"]
1167+
del self_as_dict["per_device_train_batch_size"]
1168+
del self_as_dict["per_device_eval_batch_size"]
11691169

11701170
self_as_dict = {k: f"<{k.upper()}>" if k.endswith("_token") else v for k, v in self_as_dict.items()}
11711171

@@ -1193,7 +1193,7 @@ def n_device(self):
11931193
@property
11941194
def train_batch_size(self) -> int:
11951195
"""
1196-
The actual batch size for training (may differ from `per_gpu_train_batch_size` in distributed training).
1196+
The actual batch size for training (may differ from `per_device_train_batch_size` in distributed training).
11971197
"""
11981198
per_device_batch_size = self.per_device_train_batch_size
11991199
train_batch_size = per_device_batch_size * max(1, self.n_device)
@@ -1202,7 +1202,7 @@ def train_batch_size(self) -> int:
12021202
@property
12031203
def eval_batch_size(self) -> int:
12041204
"""
1205-
The actual batch size for evaluation (may differ from `per_gpu_eval_batch_size` in distributed training).
1205+
The actual batch size for evaluation (may differ from `per_device_eval_batch_size` in distributed training).
12061206
"""
12071207
per_device_batch_size = self.per_device_eval_batch_size
12081208
eval_batch_size = per_device_batch_size * max(1, self.n_device)

0 commit comments

Comments
 (0)