github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 04:24:47 +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

2
colossalai/device/device_mesh.py
View File

@@ -38,7 +38,7 @@ class DeviceMesh:
device (str): the device for the process groups used by the DeviceMesh instance. (default: 'cuda')
"""
_DIST_BACKEND = {"cuda": "nccl", "cpu": "gloo"}
_DIST_BACKEND = {"cuda": "nccl", "cpu": "gloo", "npu": "hccl"}
def __init__(
self,