lomo训练65b llama实测 Lomo is incompatible with pipeline parallelism

[zlh1992]

配置： CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --rdzv_backend=c10d --rdzv_endpoint=localhost:29402 --nnodes=1 --nproc_per_node=8 train.py config.tp_size = 1 config.dp_size = 1 # 8 无所谓 config.pp_size = 1 config.train_epochs = 1 config.eval_per_n_steps = 0 config.eval_per_n_epochs = 1 config.train_micro_batch_size = 1 config.eval_batch_size = 1 config.ds_config = { "fp16": { "enabled": True }, "zero_allow_untested_optimizer": True, "zero_force_ds_cpu_optimizer": False, "zero_optimization": { "stage": 3, "offload_optimizer": { "device": "cpu", "pin_memory": False } } } 8张A100 每张消耗在30gb左右 内存消耗130GB 加载模型offload峰值约400GB内存

想做pp试试，修改配置如下： config.tp_size = 4 config.dp_size = 1 config.pp_size = 2

collie/module.py中需要修改： self.parts = [int(i) for i in self.parts] os.environ["COLLIE_PP_PARTS"] = json.dumps(self.parts)

目前发现现在还不支持：Lomo is incompatible with pipeline parallelism

[KaiLv69]

你好，因为fused_backward过程，lomo不支持1f1b调度的流水线并行。 8张A100 每张消耗在30gb左右确实有点高，可能deepspeed的通信buffer占掉了许多显存。 lomo和deepspeed的offload兼容性还没测试，不知道其结果怎么样，也不确定有没有真正offload到cpu上。