[Feature] Type Inference of primitives and module selection API change
- 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
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