[pipelining] Add back support for multi-use parameters/buffers #126626
Labels
module: pipelining
Pipeline Parallelism
oncall: distributed
Add this issue/PR to distributed oncall triage queue
triaged
This issue has been looked at a team member, and triaged and prioritized into an appropriate module
🚀 The feature, motivation and pitch
When running tracer mode with torchtitan, the follow
NotImplementedError
was raised:The source code that causes the multi-use is in
Transformer
's forward function:The support was temporarily dropped when we refactor the tracer (_IR.py) to use unflattener. We should add it back.
Alternatives
No response
Additional context
No response
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k
The text was updated successfully, but these errors were encountered: