github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix(optimizer/hybrid_zero_optim.py): fix bucket size full judge condition when reduce scatter overlap

pull/436/head
huangting4201 2023-11-02 10:30:16 +08:00
parent 10b5056e1e
commit 4851291356
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internlm/solver/optimizer/hybrid_zero_optim.py
View File

@ -404,7 +404,7 @@ class HybridZeroOptimizer(BaseOptimizer):
# check if the bucket is full
# if full, will reduce the grads already in the bucket
# after reduction, the bucket will be empty
if current_bucket.num_elements_in_bucket(reduce_rank) >= self._reduce_bucket_size:
if current_bucket.num_elements_in_bucket(reduce_rank) + param_size > self._reduce_bucket_size:
self._accum_grads_store_in_bucket(current_bucket, reduce_rank)
# otherwise, add the parameter into bucket.

2
internlm/train/training_internlm.py
View File

@ -600,12 +600,12 @@ def record_current_batch_training_metrics(
step_count=batch_count,
cur_step_loss=loss.item(),
)
loss_list.append(loss.item())
if batch_count >= 5:
tgs_list.append(tgs_origin)
tflops_list.append(tflops)
tflops_list_2.append(tflops_2)
if batch_count == gpc.config.data.total_steps - 1:
print(tgs_list, flush=True)
avg_tgs = sum(tgs_list) / len(tgs_list)