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

feat baichuan2 rmsnorm whose hidden size equals to 5120 (#5611)

Browse Source
This commit is contained in:
Steve Luo
2024-04-19 15:34:53 +08:00
committed by GitHub
parent e37ee2fb65
commit ccf72797e3
3 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_infer/test_ops/cuda/test_rms_layernorm.py
View File

@@ -9,7 +9,7 @@ inference_ops = InferenceOpsLoader().load()
@pytest.mark.parametrize("M", [2, 4, 8, 16])
@pytest.mark.parametrize("N", [64, 128, 512])
@pytest.mark.parametrize("N", [64, 128, 512, 5120])
def test_rms_layernorm(M: int, N: int):
torch.manual_seed(123)
torch.cuda.empty_cache()
@@ -48,4 +48,4 @@ def test_rms_layernorm(M: int, N: int):
if __name__ == "__main__":
test_rms_layernorm(16, 512)
test_rms_layernorm(16, 5120)