Initialization/warmstarting for RoBERTa pretraining

I am currently trying to further pretrain a RoBERTa model on a custom dataset, initializing the model with roberta-base weights. I am using this script which works fine for pretraining from scratch. The function load_training_checkpoint allows to resume training from an existing deepspeed checkpoint but I haven't found an easy way to initialize the pretraining with weights not obtained through deepspeed training (e.g. roberta-base weights). I have looked at the docs and I tried to set the load_module_only argument to True but I get a TypeError:
TypeError: load_checkpoint() got an unexpected keyword argument 'load_module_only'

Could you please help me on this? Is there any easy way to do warmstarting with DeepSpeed?

I am using deepspeed==0.4.3 and transformers==4.4.0.

Start with run_pretraining.py, especially load_training_checkpoint and the attempted load_module_only call, then compare that usage with the DeepSpeed model checkpointing documentation for version 0.4.3. Determine whether initialization from roberta-base weights is supported; done means a documented warmstarting path or a clear explanation of the limitation and required versions.