diff --git a/README.md b/README.md index 9c37ef5..2faff93 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ def split_model(model_name): device_map['language_model.model.embed_tokens'] = 0 device_map['language_model.output'] = 0 device_map['language_model.model.norm'] = 0 + device_map['language_model.model.rotary_emb'] = 0 device_map['language_model.lm_head'] = 0 device_map[f'language_model.model.layers.{num_layers - 1}'] = 0