From 0ff8406b009d15298e5b1d6ca884742630f6d91c Mon Sep 17 00:00:00 2001 From: ver217 Date: Fri, 3 Mar 2023 16:27:59 +0800 Subject: [PATCH] [chatgpt] allow shard init and display warning (#2986) --- .../ChatGPT/chatgpt/trainer/strategies/colossalai.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py index bf4ecdfdf..b6ed1d451 100644 --- a/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py +++ b/applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py @@ -79,8 +79,9 @@ class ColossalAIStrategy(DDPStrategy): self.stage = stage # TODO(ver217): support shard_init when using from_pretrained() if shard_init: - warnings.warn(f'Shard init is not supported yet. Ignore.') - shard_init = False + warnings.warn( + f'Shard init is not supported model.from_pretrained() yet. Please load weights after strategy.prepare()' + ) self.shard_init = shard_init self.gemini_config = dict(device=get_current_device(), placement_policy=placement_policy,