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

Added MoE parallel (#127)

Browse Source
This commit is contained in:
HELSON
2022-01-07 15:08:36 +08:00
committed by GitHub
parent 42741dd4a3
commit dceae85195
26 changed files with 858 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/engine/schedule/_base_schedule.py
View File

@@ -38,8 +38,9 @@ class BaseSchedule(ABC):
return data
@staticmethod
def _check_sanity(data, tag):
assert isinstance(data, (torch.Tensor, dict)), f'{tag} must be torch.Tensor or dict'
def _check_sanity(data, tag: str):
assert isinstance(data, (torch.Tensor, dict)), \
f'{tag} must be torch.Tensor or dict'
def load_batch(self, data_iter, to_gpu=True):
"""Loads a batch from data iterator. It returns the data and labels which are