Skip to content

support as_python_constant on PlacementClassVariable #124398

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

support as_python_constant on PlacementClassVariable #124398

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ac47c77
support as_python_constant on PlacementClassVariable
bdhirsh Apr 18, 2024
508b64c
Update on "support as_python_constant on PlacementClassVariable"
bdhirsh Apr 18, 2024
ea7f31d
Update on "support as_python_constant on PlacementClassVariable"
bdhirsh May 1, 2024
945073e
Update on "support as_python_constant on PlacementClassVariable"
bdhirsh May 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions torch/_dynamo/variables/distributed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,9 @@ def is_placement_type(value):

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@wanchaol Can you remind me why we needed this variable?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This boils down to the "custom fx type" that we've been discussing. So fundamentally we need to inline Placement and DeviceMesh as a constant in the closure function (i.e. from_local) and put that closure function to the fx graph as a call_function node.

But if we don't have this PlacementClassVariable, dynamo would trace the DTensor's metadata construction as a UserDefinedClass/Object (i.e. Shard(1)). UDTs are not constant variable (as it's hard to tell whether a UDT is a ConstantVariable unless user explicitly tell dynamo the objects are ConstantVariable). So this PlacementClassVariable allow us to turn the Shard(1) as a PlacementVariable (which is a constant variable), and then the sharding metadata can be inlined as a closure

return type(value) is type and issubclass(value, Placement)

def as_python_constant(self):
return self.value

def call_function(
self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
) -> "VariableTracker":
Expand Down