cofe-ai / msg

Masked Structural Growth for 2x Faster Language Model Pre-training

Masked Structural Growth

We grow up language models in pre-training with efficient schedules and function-preserving operators that yields 2x speedup.

MSG paper: https://arxiv.org/abs/2305.02869

Quick Start

The following example shows how to run MSG on public Bert Pre-training data.

1. Pre-processing

    preprocess_bert_data.py

This generates static masks for raw data.

2. Run MSG

For Bert-base:

    sh grow_bert_base.sh

For Bert-large:

    sh grow_bert_large.sh

3. Evaluation

    cd glue_eval
    sh run_glue_together_with_stat.sh

Notes

You can modify configs/*.json and set "attention_probs_dropout_prob" and "hidden_dropout_prob" to 0.0 in order to check function preservation. However, according to different pytorch versions, there can still be negligible differences of loss before and after growth.

References

If this project helps you, please cite us, thanks!

@inproceedings{
yao2024masked,
title={Masked Structural Growth for 2x Faster Language Model Pre-training},
author={Yiqun Yao and Zheng Zhang and Jing Li and Yequan Wang},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=rL7xsg1aRn}
}

msg's Issues

Can't work when gradient_accumulation_steps > 1

Note that the update of completed_steps and grow_step_count are not consistent

Reproducing experiments from the paper

Hi, I was wondering how to reproduce the base experiments from the paper. In particular, the current dataset and maximal number of training steps are too small to be right. How should I set those values?

can this method used for fine tuned model?

能用在微调过的模型上吗？增加参数后更推荐进行指令微调还是知识注入？