github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-27 15:57:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fix] fix test_shard_llama ci;

Browse Source
This commit is contained in:
duanjunwen 2024-10-28 02:42:33 +00:00
parent 03fa79a55c
commit 6377aa0fff
2 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/shardformer/modeling/llama.py
View File

@ -82,7 +82,7 @@ class LlamaPipelineForwards:
elif input_ids is not None: elif input_ids is not None:
batch_size, seq_length = input_ids.shape[:2] batch_size, seq_length = input_ids.shape[:2]
elif inputs_embeds is not None: elif inputs_embeds is not None:
batch_size, seq_length = inputs_embeds.shape[:2] batch_size, seq_length, _ = inputs_embeds.shape[:2]
else: else:
raise ValueError("You have to specify either input_ids or inputs_embeds") raise ValueError("You have to specify either input_ids or inputs_embeds")
if inputs_embeds is None: if inputs_embeds is None:

1
tests/test_shardformer/test_model/test_shard_llama.py
View File

@ -325,7 +325,6 @@ def run_llama_test(test_config):
).get_v_schedule() ).get_v_schedule()
test_config["scheduler_nodes"] = scheduler_nodes test_config["scheduler_nodes"] = scheduler_nodes
for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
print(f"name {name}")
if test_config.get("sequence_parallelism_mode", None) == "ring_attn" and "causal" not in name: if test_config.get("sequence_parallelism_mode", None) == "ring_attn" and "causal" not in name:
continue continue
try: try: