github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-09 04:50:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer] fix emerged bugs after updating transformers (#4526)

Browse Source
This commit is contained in:
Baizhou Zhang
2023-08-29 11:25:05 +08:00
committed by GitHub
parent c554b7f559
commit 0387a47e63
2 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/pipeline/schedule/_utils.py
View File

@@ -123,7 +123,10 @@ def merge_batch(data: List[Any]) -> Any:
merged_data = []
for elem_batch in zip(*flattened_data):
if isinstance(elem_batch[0], torch.Tensor):
merged_data.append(torch.cat(elem_batch, dim=0))
if len(elem_batch[0].shape) == 0: # set loss to None in pipeline outputs
merged_data.append(None)
else:
merged_data.append(torch.cat(elem_batch, dim=0))
else:
merged_data.append(list(elem_batch))
return tree_unflatten(merged_data, tree_spec)