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

[example] simplify opt example (#2344)

Browse Source
This commit is contained in:
Jiarui Fang
2023-01-06 10:08:41 +08:00
committed by GitHub
parent 7080a8edb0
commit 35e22be2f6
10 changed files with 234 additions and 684 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/language/gpt/train_gpt_demo.py
View File

@@ -5,7 +5,6 @@ from time import time
import psutil
import torch
import torch.nn as nn
from model_zoo import model_builder
from packaging import version
from torch.nn.parallel import DistributedDataParallel as DDP
from utils import get_data, get_tflops
@@ -16,6 +15,7 @@ from colossalai.nn.parallel import ZeroDDP
from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
from colossalai.utils import get_current_device
from colossalai.utils.model.colo_init_context import ColoInitContext
from model_zoo import model_builder
CAI_VERSION = colossalai.__version__