Skip to content

Latest commit

 

History

History
32 lines (24 loc) · 1.18 KB

README.md

File metadata and controls

32 lines (24 loc) · 1.18 KB
Raw

Examples

This folder includes some examples of using model schedule. Before running examples, please follow the guidance in the benchmark folder to install the required Python packages.

These examples can be run on Megatron-LM and DeepSpeed. If you attempt to run the scheduled model on other frameworks, you can invoke the scheduled model as follows:

# my_script.py
from transformers import AutoConfig
from transformers.models.modeling_bert import BertLMHeadModel

model_name = "bert-large-uncased"
config = AutoConfig.from_pretrained(model_name)
model = BertLMHeadModel(config)

def apply_and_build_schedule(model, config):
    from slapo.model_schedule import apply_schedule

    sch = apply_schedule(
        model, "bert", model_config=config, prefix="bert", fp16=True, ckpt_ratio=0
    )
    opt_model, _ = slapo.build(sch, init_weights=model._init_weights)
    return opt_model

opt_model = apply_and_build_schedule(model, config)
# ... training logic