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

[example] add TP to GPT example (#1828)

Browse Source
This commit is contained in:
Jiarui Fang
2022-11-08 17:17:19 +08:00
committed by GitHub
parent 49216d7ab1
commit a25f755331
4 changed files with 113 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/language/opt/run_clm.py
View File

@@ -36,7 +36,6 @@ from datasets import load_dataset
from packaging import version
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from utils import colo_memory_cap
import colossalai
import transformers
@@ -47,7 +46,6 @@ from colossalai.nn.optimizer import HybridAdam
from colossalai.nn.parallel import ZeroDDP
from colossalai.tensor import ProcessGroup
from colossalai.utils import get_current_device, get_dataloader
from colossalai.utils.checkpoint import load_checkpoint, save_checkpoint
from colossalai.utils.model.colo_init_context import ColoInitContext
from colossalai.zero import ZeroOptimizer
from transformers import (
@@ -249,12 +247,20 @@ def parse_args():
return args
def colo_memory_cap(size_in_GB):
from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device
cuda_capacity = colo_device_memory_capacity(get_current_device())
if size_in_GB * (1024**3) < cuda_capacity:
colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity)
print("Using {} GB of GPU memory".format(size_in_GB))
def main():
args = parse_args()
disable_existing_loggers()
colossalai.launch_from_torch(config=dict())
logger = get_dist_logger()
is_main_process = gpc.get_local_rank(ParallelMode.DATA) == 0
is_main_process = dist.get_rank() == 0
if is_main_process:
datasets.utils.logging.set_verbosity_warning()