Anuj040
commented
1 year ago I was able to solve the memory related issues. However I am still stuck with following issues: With the following strategy
elif args.strategy == 'colossalai_gemini':
strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5, shard_init=True)leads to
AssertionError: You should use `zero_ddp_wrapper` first
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /rlhf/applications/Chat/train_prompts.py:286 in <module> │
│ │
│ 283 │ parser.add_argument('--kl_coef', type=float, default=0.1) │
│ 284 │ parser.add_argument('--ptx_coef', type=float, default=0.9) │
│ 285 │ args = parser.parse_args() │
│ ❱ 286 │ main(args) │
│ 287 │
│ │
│ /rlhf/applications/Chat/train_prompts.py:213 in main │
│ │
│ 210 │ │ return {k: v.to(torch.cuda.current_device()) for k, v in batch │
│ 211 │ if is_rank_0(): │
│ 212 │ │ print("12") │
│ ❱ 213 │ (actor, actor_optim), (critic, critic_optim) = strategy.prepare((a │
│ 214 │ if is_rank_0(): │
│ 215 │ │ print("15") │
│ 216 │
│ │
│ /rlhf/applications/Chat/coati/trainer/strategies/base.py:84 in prepare │
│ │
│ 81 │ │ │ │ assert len(arg) == 2, f'Expect (model, optimizer) pair │
│ 82 │ │ │ │ model, optimizer = arg │
│ 83 │ │ │ │ model = prepare_model(model) │
│ ❱ 84 │ │ │ │ optimizer = self.setup_optimizer(optimizer, self._unwr │
│ 85 │ │ │ │ rets.append((model, optimizer)) │
│ 86 │ │ │ elif isinstance(arg, nn.Module): │
│ 87 │ │ │ │ rets.append(prepare_model(arg)) │
│ │
│ /rlhf/applications/Chat/coati/trainer/strategies/colossalai.py:147 in │
│ setup_optimizer │
│ │
│ 144 │ │
│ 145 │ def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Mo │
│ 146 │ │ assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsuppo │
│ ❱ 147 │ │ return zero_optim_wrapper(model, optimizer, optim_config=self. │
│ 148 │ │
│ 149 │ def backward(self, loss: torch.Tensor, model: nn.Module, optimizer │
│ 150 │ │ optimizer.backward(loss) │
│ │
│ /opt/conda/lib/python3.9/site-packages/colossalai/zero/wrapper.py:83 in │
│ zero_optim_wrapper │
│ │
│ 80 │ │ │ │ >>> zero2_config = dict(reduce_bucket_size=12 * 1024 * │
│ 81 │ │ │ │ >>> optim = zero_optim_wrapper(model, optim, optim_con │
│ 82 │ """ │
│ ❱ 83 │ assert hasattr(model, "_colo_zero_stage"), "You should use `zero_d │
│ 84 │ zero_stage = getattr(model, "_colo_zero_stage") │
│ 85 │ │
│ 86 │ assert norm_type == 2.0, "Current ZeRO optimizers only support 'no │
╰──────────────────────────────────────────────────────────────────────────────╯
🐛 Describe the bug
I am following this blog https://medium.com/pytorch/colossalchat-an-open-source-solution-for-cloning-chatgpt-with-a-complete-rlhf-pipeline-5edf08fb538b to train a 6.7B parameter model. The blog outlines training scheme for a 7B model. So I expect it to be more or less same.
When I use
colossalai_zero2_cpu, the execution stops with the following tracebackBelieving it to be the case of insufficient GPU memory, I thought perhaps the model needs to be sharded. With
colossalai_geministartegy, the job crashed with the same errors as before.Next I though of trying a mix and match of last two strategies
For the sake of testing, I reduced the size of reward model to 2.7B while keeping the actor model size to 6.7B. It led to following errors/warnings
Following the advice here, I implemented the following
Something to notice, even with
shard_init = TrueI would have expected the pipeline parallel size or tensor parallel size to be different from 1.
After implemented the above hacky sharding solution, the models were initialized alright but the job crashed with the following traceback
when both the models were 6.7B parameters.
All help will be appreciated.
Environment
8 Tesla V100 32GB RAM shm: 512 CUDA Version 11.6