Skip to content

[FEATURE]: Deepspeed不开offload 与Colossalai Gemini对llama2_13b的测试对比 #4719

@wangbluo

Description

@wangbluo

Describe the feature

尊敬的开发者们你们好,前几天测试了你们关于llama2 gemini-auto和Deepspeed zero3+offload的性能,今天测试了一下gemini和Deepspeed两边都不开offload的一个性能情况,想跟你们分享一下。

测试脚本分别是pretrain.py和Deepspeed的Trainer训练框架,Deepspeed版本为0.10.0

测试模型是llama2 13b,sequence length为2048,pre device batch size都是8,梯度累积Deepspeed设为1,采用zero3

我测试是从一个step的开始打点一个time,到backward结束打点,计算一个step time的时间差。

colossalai gemini pretrain.py的step_time:13.5s,显存占用:52834MB,
而benchmark.py的steptime:9.8s左右,显存占用:39998MB
(ps:还在看这二者之间的区别,不懂为什么step time和显存占用有一定差距)

Deepspeed的step time:24-25s,显存占用:51438MB
如果Deepspeed开了offload,那么step_time为27s左右,显存占用为26445MB

也就是说,colossalai gemini 开offload 和不开offload的step_time和显存占用,都要优于Deepspeed的开offload 和不开offload的step time和显存占用情况。

因为看到你们给的https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/llama2 readme里面,给的是Deepspeed zero3+offload和colossalai的对比情况,并未给出二者都不开offload 的对比,所以想和你们确认一下这个情况。

我理解Deepspeed zero3+offload比不过colossalai是很正常的,毕竟他们的offload不是auto的,但是我不确定二者都不开offload的一个对比,你们有没有测过这种情况呢,可以找一个不会oom的配置。

如果以上测得没有问题的话我觉得是非常棒的。

其中ds_config配置为:

{
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 0,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 0,
        "stage3_max_reuse_distance": 0,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": "auto",
            "betas": [
                0.9,
                0.999
            ],
            "eps": 1e-8,
            "weight_decay": 0
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "last_batch_iteration": -1,
            "total_num_steps": "auto",
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "initial_scale_power": 32,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions