如何同时训练两个模型？

[wangqiang9]

Is there an existing issue for this bug?

• I have searched the existing issues

🐛 Describe the bug

在官方文档中给出了训练一个model的例子：

colossalai.launch(...)
plugin = GeminiPlugin(...)
booster = Booster(precision='fp16', plugin=plugin)

model = GPT2()
optimizer = HybridAdam(model.parameters())
dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
lr_scheduler = LinearWarmupScheduler()
criterion = GPTLMLoss()

model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model, optimizer, criterion, dataloader, lr_scheduler)

for epoch in range(max_epochs):
    for input_ids, attention_mask in dataloader:
        outputs = model(input_ids.cuda(), attention_mask.cuda())
        loss = criterion(outputs.logits, input_ids)
        booster.backward(loss, optimizer)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

如果在我的训练中，有两个模型model1和model2都需要被训练，应该如何使用训练呢？

Environment

No response

[Issues-translate-bot]

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

---

Title: How to train two models at the same time?

[wangqiang9]

请教一下，能否写成这样：

model1, _, _, _, _ = booster.boost(model=model1)
model2, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model2, optimizer, criterion, dataloader, lr_scheduler)

[Issues-translate-bot]

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

---

Please tell me if you can write it like this:

model1, _ = booster.boost(model1)
model2, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model2, optimizer, criterion, dataloader, lr_scheduler)

[Edenzzzz]

Looks like the right way if you intend to freeze the first model. Is this throwing any error?