github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-04 11:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fix] fix detach_output_obj clone;

Browse Source
This commit is contained in:
duanjunwen 2024-09-24 08:08:32 +00:00
parent 6c1e1550ae
commit 7e6f793c51
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/pipeline/schedule/zero_bubble_pp.py
View File

@ -622,10 +622,10 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
else: else:
# detach output # detach output
detached_output_obj = tree_map(detach, output_obj) detached_output_obj = tree_map(detach, output_obj)
# 3-2 clone output # 3-2 clone detached_output_obj
output_obj = tree_map(clone, output_obj) detached_output_obj = tree_map(clone, detached_output_obj)
# 3-3 release cloned output.data; release_tensor_data output for bwd b & w; (do not detach output) # 3-3 release cloned output.data; release_tensor_data output for bwd b & w; (do not detach output)
# output_obj = tree_map(clone, output_obj)
if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True):
# We should not release_tensor_data bwd LOSS # We should not release_tensor_data bwd LOSS
pass pass