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

[fix] remove unnecessary dp_size assert (#5351)

Browse Source 
* fix: remove unnecessary assert

* test: add more 3d plugin tests

* fix: add warning
This commit is contained in:
Wenhao Chen
2024-02-02 14:40:20 +08:00
committed by GitHub
parent ffffc32dc7
commit 1c790c0877
2 changed files with 21 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_booster/test_plugin/test_3d_plugin.py
View File

@@ -118,6 +118,20 @@ def check_3d_plugin(init_method: str = "none", early_stop: bool = True):
@parameterize(
"test_args",
[
{
"batch_size": 8,
"num_steps": 4,
"tp": 2,
"pp": 2,
"pp_style": "1f1b",
"num_model_chunks": 1,
"num_microbatches": 4,
"zero": 1,
"precision": "fp16",
"initial_scale": 1,
"max_length": 512,
"gradient_accumulation_step": 2,
},
{
"batch_size": 8,
"num_steps": 4,