[Pipeline] Tie weight analysis

[comaniac]

Description

This PR analyzes whether parameters are tie together by looking at their nn.Paramter object IDs. Accordingly, this analysis has to be done before the consolidation (if needed); otherwise all nn.Parameter objects will be replaced so that their object IDs won't be the same anymore.

Note that the above analysis works only when parameters are pointing to the same Python object. In other words, for the distributed training that requires weight consolidation, users have to explicitly call tie_weights one more time when initializing the model with empty weights. For example:

with slapo.init_empty_weights():
    model = get_model() # tie_weights() is called implicitly when constructing this model.

sch = slapo.create_schedule(model, ...)
# The .build does the following in order:
# 1. Partition pipeline stages.
# 2. Tie weight analysis.
# 3. Consolidation (this wipes out the previous tie weights in the model)
# 4. Liveness analysis.
# 5. Pipeline module construction (this makes use the tie weight and liveness analysis results)
sch.build(...)

In addition, this PR also refactors the pipeline logic in schedu.build. The ultimate goal is to ensure that schedule.build doesn't have any framework specific logic (e.g., if target == "deepspeed":).

cc @szhengac @zarzen

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

[szhengac]

LGTM.

[szhengac]

Let me know if it is ready for merge

[comaniac]

Thanks @szhengac

[zarzen]

Is stage_modules a list of layers? Or is it partitioned by the Slapo schedule?