github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-28 13:05:26 +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

2
examples/inference/benchmark_ops/benchmark_rmsnorm.py
View File

@@ -35,7 +35,7 @@ configs = [
styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
ylabel="ms",
plot_name=f"RMSNorm benchmarking results",
args={"HIDDEN_SIZE": 1024},
args={"HIDDEN_SIZE": 5120},
)
]

6
extensions/csrc/cuda/rms_layernorm_kernel.cu
View File

@@ -277,6 +277,9 @@ void rms_layernorm(
case 2:
RMSNORM_LAUNCHER(2, block);
break;
case 3:
RMSNORM_LAUNCHER(3, block);
break;
case 4:
RMSNORM_LAUNCHER(4, block);
break;
@@ -321,6 +324,9 @@ void fused_add_rms_layernorm(
case 2:
FUSED_ADD_RMSNORM_LAUNCHER(2, block);
break;
case 3:
FUSED_ADD_RMSNORM_LAUNCHER(3, block);
break;
case 4:
FUSED_ADD_RMSNORM_LAUNCHER(4, block);
break;

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)