Colossal-AI 低成本复现 ChatGPT，此次开源有什么技术亮点？

{"text": "\u7b2c1\u7ae0\n\u9884\u5907\u77e5\u8bc6\n\u77e5\u8bc6\u76ee\u6807\n1.\u638c\u63e1\u5b89\u5168\u6807\u5fd7\u3002\n2.\u4e86\u89e3\u5b89\u5168\u6587\u660e\u751f\u4ea7\u548c\u201c6S\u201d\u7ba1\u7406\u8981\u6c42\u3002\n3\uff0e\u4e86\u89e3\u94b3\u5de5\uff08\u673a\u68b0\u62c6\u88c5\uff09\u5b9e\u8bad\u573a\u5730\u7684\u8bbe\u5907\u548c\u64cd\u4f5c\u4e2d\u5e38\u7528\u7684\u5de5\u5177\u3001\u91cf\u5177\u548c\u5203\u5177\u3002\n4\uff0e\u4e86\u89e3\u8f66\u5e8a\u52a0\u5de5\u8303\u56f4\u3001\u578b\u53f7\u3001\u89c4\u683c\u548c\u8f66\u5e8a\u4e3b\u8981\u90e8\u4ef6\u7684\u7ec4\u6210\u53ca\u5176\u4f5c\u7528\u3002\n5\uff0e\u638c\u63e1\u5e8a\u978d\uff08\u5927\u62d6\u677f\uff09\u4e2d\u6ed1\u677f\uff08\u4e2d\u62d6\u677f\uff09\u3001\u5c0f\u6ed1\u677f\uff08\u5c0f\u62d6\u677f\uff09\u7684\u8fdb\u9000\u5200\u65b9\u5411\u3002\n6.\u4e86\u89e3\u8f66\u5e8a\u4f20\u52a8\u7cfb\u7edf\u3002\n1.1\u5b89\u5168\u6587\u660e\u64cd\u4f5c\u89c4\u7a0b(6S\u7ba1\u7406\uff09\n\u5b89\u5168\u4fdd\u8bc1\u751f\u4ea7\uff0c\u751f\u4ea7\u5fc5\u987b\u5b89\u5168\u3002\n1.1.1\n\u8ba4\u8bc6\u5e38\u7528\u7684\u5b89\u5168\u56fe\u6807\n\u5e38\u7528\u7684\u5b89\u5168\u56fe\u6807\u5982\u56fe1-1\u6240\u793a\u3002\n0\n\u6d88\u9632\u624b\u52a8\u542f\u52a8\u5668\n\u706b\u8b66\u7535\u8bdd\n\u6d88-\u9632-\u6c34-"}

maxlen=1024        
filenames = glob.glob('/root/bert4keras/data/machine/*')
for filename in filenames:
            with open(filename) as f:
                    l = f.read()
                    out=""
                    for l_one in l.split("。"):
                        if len(out+l_one+"。") > maxlen:
                            yield out
                            out=""
                        else:
                            out+=l_one+"。"

import contextlib
import os
import torch
import torch.nn as nn
from dataset.webtext import WebtextDataset
from titans.model.gpt import GPTLMLoss
import colossalai
import colossalai.utils as utils
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.logging import disable_existing_loggers, get_dist_logger
from colossalai.nn import LinearWarmupLR
from colossalai.trainer import Trainer, hooks
from colossalai.utils import colo_set_process_memory_fraction, is_using_pp
from colossalai.utils.timer import MultiTimer
from colossalai.zero.init_ctx import ZeroInitContext
def calc_local_model_size(model: torch.nn.Module):
    numel_per_device = 0
    for p in model.parameters():
        numel_per_device += p.numel()
    return numel_per_device
VOCAB_SIZE = 50257
def main():
    parser = colossalai.get_default_parser()
    parser.add_argument('--from_torch', default=False, action='store_true')
    parser.add_argument('--use_dummy_dataset', default=False, action='store_true')
    args = parser.parse_args()
    disable_existing_loggers()
    if args.from_torch:
        colossalai.launch_from_torch(config=args.config)
    else:
        colossalai.launch_from_slurm(config=args.config, host=args.host, port=29500, seed=42)
    logger = get_dist_logger()
    data_path = None if args.use_dummy_dataset else os.environ['DATA']
    logger.info(f'Build data loader from path {data_path}', ranks=[0])
    train_ds = WebtextDataset(path=data_path, seq_len=gpc.config.SEQ_LEN)
    train_dataloader = utils.get_dataloader(train_ds,
                                            seed=42,
                                            batch_size=gpc.config.BATCH_SIZE,
                                            pin_memory=True,
                                            shuffle=True,
                                            drop_last=True)
    logger.info('Build model', ranks=[0])
    use_pipeline = is_using_pp()
    use_interleaved = hasattr(gpc.config.model, 'num_chunks')
    use_zero3 = hasattr(gpc.config, 'zero')
    ctx = contextlib.nullcontext()
    if use_zero3:
        ctx = ZeroInitContext(target_device=torch.cuda.current_device(),
                              shard_strategy=gpc.config.zero.model_config.shard_strategy,
                              shard_param=True)
    with ctx:
        model = gpc.config.model.pop('type')(**gpc.config.model)
    if use_pipeline and use_interleaved and not isinstance(model, nn.ModuleList):
        model = nn.ModuleList([model])
    if use_zero3:
        numel = ctx.model_numel_tensor.item()
    else:
        numel = calc_local_model_size(model)
    tflop = numel * gpc.config.BATCH_SIZE * gpc.config.SEQ_LEN \
        * gpc.get_world_size(ParallelMode.MODEL) * gpc.get_world_size(ParallelMode.DATA) * 8 / (1024 ** 4)
    criterion = getattr(gpc.config, 'loss_fn', None)
    if criterion is not None:
        criterion = criterion.type()
    else:
        criterion = GPTLMLoss()
    logger.info('Build optimizer', ranks=[0])
    optimizer = gpc.config.optimizer.pop('type')(model.parameters(), **gpc.config.optimizer)
    lr_scheduler = LinearWarmupLR(optimizer, total_steps=gpc.config.NUM_EPOCHS, warmup_steps=5)
    engine, train_dataloader, _, lr_scheduler = colossalai.initialize(model,
                                                                      optimizer,
                                                                      criterion,
                                                                      train_dataloader=train_dataloader,
                                                                      lr_scheduler=lr_scheduler)
    global_batch_size = gpc.config.BATCH_SIZE * \
        gpc.get_world_size(ParallelMode.DATA) * getattr(gpc.config, "gradient_accumulation", 1)
    logger.info(f'Init done, global batch size = {global_batch_size}', ranks=[0])
    timier = MultiTimer()
    trainer = Trainer(engine=engine, logger=logger, timer=timier)
    hook_list = [
        hooks.LossHook(),
        hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=True),
        hooks.LogMetricByEpochHook(logger),
        hooks.ThroughputHook(ignored_steps=10, tflop_per_step=tflop),
        hooks.LogMetricByStepHook(),
        hooks.LogMemoryByEpochHook(logger),
    trainer.fit(train_dataloader=train_dataloader,
                epochs=gpc.config.NUM_EPOCHS,
                test_interval=1,
                hooks=hook_list,
                display_progress=True,
                return_output_label=False)
if __name__ == '__main__':
    main()

# 引入colossalai.utils包
from colossalai.utils import colo_set_process_memory_fraction, is_using_pp, save_checkpoint
# 保存我们进行训练过的参数   
save_checkpoint('模型名称.pt', 1, trainer.engine.model)

from model import GPT2_small_pipeline_hybrid
from colossalai.nn.optimizer import HybridAdam
from colossalai.zero.shard_utils import TensorShardStrategy
BATCH_SIZE = 12
NUM_EPOCHS = 2
SEQ_LEN = 1024
NUM_MICRO_BATCHES = 4
HIDDEN_SIZE = 768
TENSOR_SHAPE = (BATCH_SIZE // NUM_MICRO_BATCHES, SEQ_LEN, HIDDEN_SIZE)
# if you do no want zero, just comment out this dictionary
zero = dict(model_config=dict(tensor_placement_policy='cuda', shard_strategy=TensorShardStrategy()),
            optimizer_config=dict(initial_scale=2**5))
optimizer = dict(
    type=HybridAdam,
    lr=0.000015,
    weight_decay=1e-2,
model = dict(type=GPT2_small_pipeline_hybrid, checkpoint=True, num_chunks=1)
# pipeline parallel: modify integer value for the number of pipeline stages
# tensor parallel: modify size to set the tensor parallel size, usually the number of GPUs per node
# for the current model implementation, mode can only be 1D or None
parallel = dict(