search

tunib-ai / parallelformers

Parallelformers: An Efficient Model Parallelization Toolkit for Deployment
https://tunib-ai.github.io/parallelformers
Apache License 2.0
776 stars 61 forks source link

How can I parallelize the MegatronBertModel? #14

Closed kajyuuen closed 2 years ago

kajyuuen commented 2 years ago

Thanks for sharing the great code.

Let me ask you one question. In the list of Fully Supported Models, it says Megatron BERT, but the following code does not work.

from transformers import MegatronBertModel
from parallelformers import parallelize

model = MegatronBertModel.from_pretrained('nvidia/megatron-bert-cased-345m')
parallelize(model, num_gpus=2, verbose='detail')

Also, I could not find MegatronBertModel in the policies. How can I parallelize the MegatronBertModel?

hyunwoongko commented 2 years ago

We added MegatronBertModel