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

[Inference]Fused the gate and up proj in mlp，and optimized the autograd process. (#5365)

Browse Source 
* fused the gate and up proj in mlp

* fix code styles

* opt auto_grad

* rollback test_inference_engine.py

* modifications based on the review feedback.

* fix bugs in flash attn

* Change reshape to view

* fix test_rmsnorm_triton.py
This commit is contained in:
yuehuayingxueluo
2024-02-06 19:38:25 +08:00
committed by GitHub
parent 1dedb57747
commit 35382a7fbf
10 changed files with 484 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/sampler.py
View File

@@ -10,7 +10,7 @@ def greedy_sample(
"""
Sample tokens greedyly.
"""
results = torch.argmax(logprobs, dim=-1).cpu()
results = torch.argmax(logprobs, dim=-1)
return results