Can I train a pre-train bert model using pytorchLightning?

[wa008]

I can't find code or docs abount train pre-train model..
I don't know how to add hugging face transformer's Trainer code to pytorch-lightning class, please help
train a pre-train transformer on gpu is below

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir=bert_file,
    overwrite_output_dir=True,
    num_train_epochs=20,
    per_device_train_batch_size=64,
    save_steps=100000,
    save_total_limit=2,
    prediction_loss_only=False,
)

trainer = Trainer(
    model = model,
    args = training_args,
    data_collator = data_collator,
    train_dataset = train_dataset
)
trainer.train()

[tchaton]

Dear @wa008,

To do so, you need to wrap your model into LightningModule and use Pytorch Lightning Trainer. Works just fine :)

Best,
T.C

[wa008]

The Trainer of bert is a complete class wraped by huggingface, It can't be breaked down to forward function, train_step function and so on, So I don't know how to warp The Trainer of bert into LightningModule... Is there a demo or ducument about it ? Thanks very much!

[LoicGrobol]

@wa008 Apologies for tooting my own horn, but I have a tool for doing exactly that, it might give some inspiration: https://github.com/LoicGrobol/zeldarose :-)

[wa008]

Thanks very much