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 40ce41b)

Here are some key observations to aid the review process:

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

🧪 PR contains tests

⚡ Recommended focus areas for review Potential Logic Change The FusionDefinitionWrapper class introduces a new way of defining fusion without multidevice_schedule. This change may have implications on the existing logic and should be thoroughly reviewed. 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 _create_fusion_definition( self, in_dtensors: Iterable[DTensor] ) -> FusionDefinition: define_fn = self._define_fusion class Model(FusionDefinition): def definition(self) -> None: 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) -> None: 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.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 ) return Model() def __call__(self, in_dtensors: Iterable[DTensor]) -> list[DTensor]: fusion_def = self._create_fusion_definition(in_dtensors) in_tensors = [in_dtensor.to_local() for in_dtensor in in_dtensors] out_tensors = fusion_def.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 Function Signature Change The FusionDefinitionWrapper class has a new method _create_fusion_definition which changes the way fusion definitions are created. This change may affect the existing function signatures and should be reviewed. def _create_fusion_definition( self, in_dtensors: Iterable[DTensor] ) -> FusionDefinition: define_fn = self._define_fusion class Model(FusionDefinition): def definition(self) -> None: 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) -> None: 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.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 ) return Model() Potential Performance Impact The FusionDefinitionWrapper class introduces a new way of parallelizing tensors which may have performance implications. This change should be reviewed to ensure it does not introduce any performance regressions. 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.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 )

[wujingyue]

!test

[wujingyue]

cc @syed-ahmed FYI, I made nvFuser return output shardings as well as local tensors

[wujingyue]

!test

[wujingyue]

!test

[jjsjann123]

LGTM.

[jjsjann123]

QQ: are we still required in nvfuser to have the shard dimension with size == num_of_shard?

[wujingyue]

We are gradually adding DID loop split so eventually this won't be a requirement. See test_communication.py and several test_*_loop_splits in test_multidevice.py recently added by @Priya2698

[jjsjann123]

I guess it doesn't matter here since we only have 1d device_mesh for now.

But the logic doesn't feel right in this for loop with for parallel_type in [nvfuser.ParallelType.mesh_x]:. i.e. I think we need to check the rank of out_tensor.mesh. Which is only 1-d. So we can't really write a future proof thing here.

^^^ I realized those aren't really constructive comments, I'm just trying to point out that we might need to put more thought on how we want to expose out_tensor.mesh. 😜

[wujingyue]

Agreed. When nvFuser supports >1D mesh, we may want DeviceMesh to hold a torch.Tensor, similar to https://github.com/pytorch/pytorch/blob/3917053f63b75f14e3cb2f53805fad4ade5363df/torch/distributed/device_mesh.py#L420.

[rdspring1]

LGTM. Just some comments on FusionDefinitionWrapper.

[wujingyue]

!test

[wujingyue]

!test