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

[npu] add npu support for hybrid plugin and llama (#5090)

Browse Source 
* llama 3d

* update

* fix autocast
This commit is contained in:
Xuanlei Zhao
2023-11-22 19:23:21 +08:00
committed by GitHub
parent aae496631c
commit 3acbf6d496
9 changed files with 61 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
examples/language/llama2/benchmark.py
View File

@@ -131,7 +131,7 @@ def main():
tp_size=args.tp,
pp_size=args.pp,
zero_stage=args.zero,
enable_fused_normalization=True,
enable_fused_normalization=torch.cuda.is_available(),
num_microbatches=args.mbs,
precision="bf16",
)
@@ -141,7 +141,7 @@ def main():
pp_size=args.pp,
zero_stage=args.zero,
cpu_offload=True,
enable_fused_normalization=True,
enable_fused_normalization=torch.cuda.is_available(),
num_microbatches=args.mbs,
initial_scale=2**8,
precision="bf16",