awslabs / awslabs/slapo

[Feature] Type Inference of primitives and module selection API change

Open
#91 7 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
153
Forks
17
PR merge metrics
No merged PRs in 30d

Description

Currently, we have the following schedule operations logic for gpt2.

```python3
# code snippet from slapo/model_schedule/gpt2.py
...
attn_op = []
for idx in range(model_config.num_hidden_layers):
sub_sch = sch[attn_path.replace("N", str(idx))]
with init_empty_weights(enable=delay_init):
new_mod = Attention(**init_config)
attn_op.append(new_mod.module.attn_op_name)
sub_sch.replace(new_mod)
cnt += 1
...
```

Issues of this code snippet from my view:
1. the primitive function `replace` cannot provide type inference features: difficult to know the options for the primitive and not possible to get the doc string for the primitive.
2. selection for sub graph is not intuitive due to the concept of the sub schedule. Treating the schedule as a dictionary/hash table is not that intuitive to me. For a single model, it is natural to me that we have a single schedule for this model. The schedule can only affect part of the model, and consider them as a list of tuples, e.g., `(module_part_id, "replace_with", new_module_obj)`. This can also facility debugging the schedule, e.g., removing entries of in the schedule to disable schedules.

Recommend to changes the APIs to the following
```python3
...
for idx in range(model_config.num_hidden_layers):
sub_module = slapo.select(model, "transformer.h."+str(idx))
with init_empty_weights(enable=delay_init):
new_mod = Attention(**init_config)
cur_schedule = slapo.replace(cur_schedule, sub_module, new_mod)
cnt += 1
...
```

And the select method can be further improved to consider fuzzy match
```python3
...
sub_modules = slapo.select("transformer.h.*")
with init_empty_weights(enable=delay_init):
new_mod = Attention(**init_config)
cur_schedule = slapo.replace(cur_schedule, sub_modules, new_mod)
cnt = len(sub_modules)
...
```

Contributor guide

Open the contributing guide

Research direction

Start with the scheduling example in slapo/model_schedule/gpt2.py and review the proposed slapo.select and slapo.replace APIs. Clarify how primitive type inference, documentation access, single-schedule selection, and fuzzy module matching should work. Done means the API design is agreed and the gpt2 scheduling usage reflects the resulting behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design, machine-learning, tooling
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.