github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer] Pipeline/whisper (#4456)

Browse Source 
* add some base tests and policies

* finish whisper base model

* add conditional generation

* finish basic tests

* whisper

* finish whisper

* finish whisper

* del useless  whisper test

* fix

* add argmin to replace

* finish revision
This commit is contained in:
Jianghai
2023-08-18 21:29:25 +08:00
committed by GitHub
parent a27e0bb494
commit 8739aa7fa0
8 changed files with 1156 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,6 +6,7 @@ from torch import distributed as dist
import colossalai
from colossalai.logging import disable_existing_loggers
from colossalai.shardformer.layer.utils import Randomizer
from colossalai.tensor.d_tensor.api import clear_layout_converter
from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
from tests.kit.model_zoo import model_zoo
@@ -143,6 +144,7 @@ def run_llama_test(test_config):
check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)
clear_layout_converter()
Randomizer.reset_index()
torch.cuda.empty_cache()