[BUG]: 使用llama2增量预训练失败

[zryowen123]

🐛 Describe the bug

我使用examples/language/llama2中的代码预训练llama2-70b。使用gemini.sh直接跑benchmark.py是成功的，但是我想基于训好的模型进行增量预训练，训练参数和gemini.sh中给出的参数一致，只是修改了如下代码读取已有的模型：
with init_ctx:
# model = LlamaForCausalLM(config)
model = LlamaForCausalLM.from_pretrained(args.model_path)
然后跑gemini.sh报了OOM的错误：
outputs = model(**batch)

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl

return forward_call(*input, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/gemini_ddp.py", line 247, in forward

outputs = self.module(*args, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl

return forward_call(*input, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 806, in forward

outputs = self.model(

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl

return forward_call(*input, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 693, in forward

layer_outputs = decoder_layer(

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl

return forward_call(*input, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 405, in forward

hidden_states = self.input_layernorm(hidden_states)

File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl

return forward_call(*input, **kwargs)

File "/opt/conda/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 89, in forward

return self.weight * hidden_states.to(input_dtype)

File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/colo_parameter.py", line 63, in torch_function

new_args = ColoParamOpHookManager.pre_op(params, *args, *kwargs.values())

File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/param_op_hook.py", line 82, in pre_op

ColoParamOpHookManager._trigger_pre_forward(params)

File "/opt/conda/lib/python3.8/site-packages/colossalai/tensor/param_op_hook.py", line 63, in _trigger_pre_forward

hook.pre_forward(params)

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/gemini_hook.py", line 47, in pre_forward

self.pre_op(params)

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/gemini_hook.py", line 35, in pre_op

self._chunk_manager.access_chunk(chunk)

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/chunk/manager.py", line 110, in access_chunk

self.__add_accessed_chunk(chunk)

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/chunk/manager.py", line 246, in __add_accessed_chunk

chunk.access_chunk()

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/chunk/chunk.py", line 359, in access_chunk

self.__gather()

File "/opt/conda/lib/python3.8/site-packages/colossalai/zero/gemini/chunk/chunk.py", line 480, in __gather

dist.all_gather(gather_list, self.cuda_shard, self.torch_pg)

File "/opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2277, in all_gather

work = group.allgather([tensor_list], [tensor])

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 500.00 MiB (GPU 3; 79.33 GiB total capacity; 73.44 GiB already allocated; 469.81 MiB free; 77.47 GiB reserved in total by PyTorch) If reserved memory is > > allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

辛苦开发者看下这个是怎么回事

Environment

No response

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Title: [BUG]: Failed to use llama2 incremental pre-training

[Fridge003]

您好，目前GeminiPlugin还不支持用from_pretrained加载模型。请保持model = LlamaForCausalLM(config) 这一行不变，在219行的booster.boost封装完模型以后，使用booster.load_model(model, model_checkpoint_path)来从存档点加载模型。我们会在不久后支持from_pretrained，请您关注我们的更新。

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Hello, currently GeminiPlugin does not support loading models with from_pretrained. Please keep the line model = LlamaForCausalLM(config) unchanged. After booster.boost encapsulates the model in line 219, use booster.load_model(model, save_dir) to load the model from the save point. We will support from_pretrained soon, please pay attention to our updates.

[zryowen123]

好的 ， 谢谢开发者

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

ok thanks developer

[huxiaosheng123]

可以参考：https://github.com/huxiaosheng123/open-llama2

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

You can refer to: https://github.com/huxiaosheng123/open-llama2

[lulia0228]

您好，目前GeminiPlugin还不支持用from_pretrained加载模型。请保持model = LlamaForCausalLM(config) 这一行不变，在219行的booster.boost封装完模型以后，使用booster.load_model(model, model_checkpoint_path)来从存档点加载模型。我们会在不久后支持from_pretrained，请您关注我们的更新。

请问这里是加载hf格式的，还是其他格式？咱们这边checkpoint存储也是hf吗？

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Hello, currently GeminiPlugin does not support loading models with from_pretrained. Please keep the line model = LlamaForCausalLM(config) unchanged. After booster.boost encapsulates the model in line 219, use booster.load_model(model, model_checkpoint_path) to load the model from the archive point. We will support from_pretrained soon, please pay attention to our updates.

Is the hf format loaded here, or other formats? Is our checkpoint storage also hf?

[Fridge003]

您好，目前GeminiPlugin还不支持用from_pretrained加载模型。请保持model = LlamaForCausalLM(config) 这一行不变，在219行的booster.boost封装完模型以后，使用booster.load_model(model, model_checkpoint_path)来从存档点加载模型。我们会在不久后支持from_pretrained，请您关注我们的更新。

请问这里是加载hf格式的，还是其他格式？咱们这边checkpoint存储也是hf吗？

您好，我们checkpoint的文件夹的格式和huggingface是一样的，只是目前有bug，需要用booster.load_model 才能加载模型。这个bug在这两天会被修复，之后就可以用from_pretrained 直接从我们的存档文件夹里加载模型了。

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Hello, currently GeminiPlugin does not support loading models with from_pretrained. Please keep the line model = LlamaForCausalLM(config) unchanged. After booster.boost encapsulates the model in line 219, use booster.load_model(model, model_checkpoint_path) to load the model from the archive point. We will support from_pretrained soon, please pay attention to our updates.

May I ask if the hf format is loaded here, or other formats? Is our checkpoint storage also hf?

Hello, the format of our checkpoint folder is the same as that of huggingface, but there is currently a bug that requires booster.load_model to load the model. This bug will be fixed in these two days, and then we can use from_pretrained to load the model directly from our archive folder.

[lulia0228]

您好，目前GeminiPlugin还不支持用from_pretrained加载模型。请保持model = LlamaForCausalLM(config) 这一行不变，在219行的booster.boost封装完模型以后，使用booster.load_model(model, model_checkpoint_path)来从存档点加载模型。我们会在不久后支持from_pretrained，请您关注我们的更新。

请问这里是加载hf格式的，还是其他格式？咱们这边checkpoint存储也是hf吗？

您好，我们checkpoint的文件夹的格式和huggingface是一样的，只是目前有bug，需要用booster.load_model 才能加载模型。这个bug在这两天会被修复，之后就可以用from_pretrained 直接从我们的存档文件夹里加载模型了。

Great！

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Hello, currently GeminiPlugin does not support loading models with from_pretrained. Please keep the line model = LlamaForCausalLM(config) unchanged. After booster.boost encapsulates the model in line 219, use booster.load_model(model, model_checkpoint_path) to load the model from the archive point. We will support from_pretrained soon, please pay attention to our updates.

Is the hf format loaded here, or other formats? Is our checkpoint storage also hf?

Hello, the format of our checkpoint folder is the same as that of huggingface, but there is currently a bug that requires booster.load_model to load the model. This bug will be fixed in these two days, and then we can use from_pretrained to load the model directly from our archive folder.

Great!

[wangbluo]

您好，目前GeminiPlugin还不支持用from_pretrained加载模型。请保持model = LlamaForCausalLM(config) 这一行不变，在219行的booster.boost封装完模型以后，使用booster.load_model(model, model_checkpoint_path)来从存档点加载模型。我们会在不久后支持from_pretrained，请您关注我们的更新。

请问这里是加载hf格式的，还是其他格式？咱们这边checkpoint存储也是hf吗？

您好，我们checkpoint的文件夹的格式和huggingface是一样的，只是目前有bug，需要用booster.load_model 才能加载模型。这个bug在这两天会被修复，之后就可以用from_pretrained 直接从我们的存档文件夹里加载模型了。

您好，这几天碰到了同样的问题，请问这个bug修复了吗