A FusionDefinition wrapper that takes/produces DTensors.

[wujingyue]

This is a proof of concept for integrating nvFuser's model parallelism to the framework.

[wujingyue]

!test

[wujingyue]

Cc @jjsjann123

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit f9082a9)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Logic Change The new local property added to the FusionDefinition class may change the behavior of the axis_sharded_on method. Reviewers should verify that the logic is correct and consistent with the rest of the codebase. @property def local(self) -> torch.Tensor: """Returns the underlying local tensor.""" return self._dtensor.local Potential Bug The FusionDefinitionWrapper class assumes that the define_fusion function will always return a FusionDefinition object. However, if the function returns None or an object of a different type, the wrapper will fail. Reviewers should add error handling to ensure that the wrapper can handle such cases. class FusionDefinitionWrapper: def __init__(self, define_fusion: Callable[[FusionDefinition], None]): """Wraps a function that defines a fusion without `multidevice_schedule`.""" self._define_fusion = define_fusion def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]: define_fn = self._define_fusion class Model(FusionDefinition): def definition(self): define_fn(self) def _find_tensor_by_index(self, index: int) -> nvfuser.Tensor: for t in self.sched.tensors(): if t.index == index: return t return None def multidevice_schedule(self): for in_tensor_index, in_dtensor in zip(self.inputs(), in_dtensors): in_tensor = self._find_tensor_by_index(in_tensor_index) # Set the device mesh. assert ( in_dtensor.device_mesh.ndim == 1 ), "nvFuser's Python API only supports 1D meshes." mesh = nvfuser.DeviceMesh( in_dtensor.device_mesh.mesh.view(-1).tolist() ) self.sched._set_device_mesh(in_tensor, mesh) # Parallelize. assert len(in_dtensor.placements) == 1, "Expect a 1D mesh" placement: Placement = in_dtensor.placements[0] if placement.is_shard(): dim = cast(Shard, placement).dim self.sched.parallelize( in_tensor, dim, nvfuser.ParallelType.mesh_x ) in_tensors = [in_dtensor.to_local() for in_dtensor in in_dtensors] model = Model() out_tensors = model.execute(in_tensors) for i, out_tensor in enumerate(out_tensors): if isinstance(out_tensor, nvfuser.DistributedTensor): mesh = dist.device_mesh.init_device_mesh("cuda", [out_tensor.mesh.size]) placements: list[Placement] = [] for parallel_type in [nvfuser.ParallelType.mesh_x]: axis: int = out_tensor.axis_sharded_on(parallel_type) placements.append(Replicate() if axis == -1 else Shard(axis)) out_tensors[i] = DTensor.from_local(out_tensor.local, mesh, placements) return out_tensors

[wujingyue]

!test