Add .boxed() to c10d::ProcessGroup and c10d::Work's pybind

[yifuwang]

Stack from ghstack (oldest at bottom):

• -> Add .boxed() to c10d::ProcessGroup and c10d::Work's pybind #111997

Summary:
When passed from C++ to Python, c10d::ProcessGroup and c10d::Work are automatically converted to their pybind class which can't be used for dispatcher ops. .boxed() exposes c10d::ProcessGroup and c10d::Work as boxed custom class object to Python.

import tempfile

import torch
import torch.distributed as dist


if __name__ == "__main__":
    with tempfile.NamedTemporaryFile(delete=False) as tmpf:
        dist.init_process_group(
            backend="nccl", init_method=f"file://{tmpf.name}", rank=0, world_size=1
        )
        group = dist.group.WORLD
        print(group)
        print(group.boxed())

<torch.distributed.distributed_c10d.ProcessGroup object at 0x7fe42fb78d30>
ScriptObject <__torch__.torch.classes.c10d.ProcessGroup>

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/111997

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit c889d69 with merge base 31c0ef9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lw]

Thanks for getting this out so promptly!

I have a hard time building PyTorch from source, but I copied that code into my own C++ extensions and it works!

I have however noticed that, inside my operator, I'm now receiving an instance of torch.classes.c10d.ProcessGroup, and I need to convert it back to a regular ProcessGroup, thus I also need an "unboxing" function.

Do you think it would be possible to add this as a static method to the pybind of c10d::ProcessGroup?

Based on my own trial and error, this would be the code:

c10::intrusive_ptr<c10d::ProcessGroup> unbox_process_group(
    const py::object& obj) {
  return torch::jit::toIValue(
             obj,
             c10::getCustomClassType<c10::intrusive_ptr<c10d::ProcessGroup>>())
      .toCustomClass<c10d::ProcessGroup>();
}

[yifuwang]

I copied that code into my own C++ extensions and it works!

That's great to know! Thanks for testing it out!

Do you think it would be possible to add this as a static method to the pybind of c10d::ProcessGroup?

This is possible. I understand you want a convenient unboxing function on cpp side, or unboxing c10d::Work on python side (since you might want to wait from python). Curious if/why you want to unbox a c10d::ProcessGroup on python side since the unboxed version is already available.

[lw]

My end-to-end use case is something like this:

def my_op(tensor, boxed_process_group):
    process_group = torch.distributed.ProcessGroup.unbox(boxed_process_group)
    torch.distributed.all_reduce(tensor, group=process_group)

lib = torch.library.Library("my_lib", "DEF")
lib.define("my_op(Tensor tensor, __torch__.torch.classes.c10d.ProcessGroup process_group) -> Tensor")
lib.impl("my_op", my_op, "CUDA")

def main():
    torch.distributed.init_process_group("nccl")
    group = torch.distributed.new_group()
    input_ = torch.ones((1,), device="cuda")
    output = torch.ops.my_lib.my_op(input_, group.boxed())

In short: I have a Python function, which I need to invoke from some other Python code, however I need this call to go though the dispatcher, thus I need to register my function as an operator. Therefore I need to box the PG when invoking the operator, and I need to unbox it within the operator.

The reason I need the function to be an operator is because I want it to be intercepted by xFormers's selective activation checkpointing mechanism.

[yifuwang]

@lw Got it! I overlooked the part where the custom kernels could be in Python. Will add unbox before landing.

[yifuwang]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here