GeezerVisonThe World Wide Web, restored to 800 pixels of glory. |
| Home|About|Gaming cabinet|Archives|How it works|Starter sites |
GeezerVison is showing github.com — Fuser/python/python_direct/multidevice.cpp at main · NVIDIA/Fuser Expand file tree/ Copy pathmultidevice.cppMore file actions 303 lines (281 loc) · 9.13 KB / Copy pathmultidevice.cppFile metadata and controls303 lines (281 loc) · 9.13 KB 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 // clang-format off /* * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES. * All rights reserved. * SPDX-License-Identifier: BSD-3-Clause */ // clang-format on #include <bindings.h> #include <direct_utils.h> #include <python_utils.h> #include <multidevice/communicator.h> #include <multidevice/device_mesh.h> #include <multidevice/execution_utils.h> #include <multidevice/executor.h> #include <multidevice/multidevice.h> #include <runtime/fusion_kernel_runtime.h> #include <python_common/distributed_tensor.h> namespace nvfuser::python { namespace { void bindCommunicator(py::module& nvfuser) { // py::nodelete is necessary because Communicator does not have a destructor. // https://pybind11.readthedocs.io/en/stable/advanced/classes.html#non-public-destructors py::class_<Communicator, std::unique_ptr<Communicator, py::nodelete>> communicator(nvfuser, "Communicator", py::module_local()); communicator.def( "instance", &Communicator::getInstance, R"( Returns the singleton communicator instance. )", py::return_value_policy::reference); communicator.def( "size", &Communicator::size, R"( Returns the number of processes in the communicator. )"); communicator.def( "rank", &Communicator::deviceId, R"( Returns the device ID associated with the current process. )"); communicator.def( "local_size", &Communicator::local_size, R"( Returns the number of processes within the node. )"); communicator.def( "local_rank", &Communicator::local_rank, R"( Returns the in-node rank associated with the current process. )"); communicator.def( "barrier", [](Communicator& self) { // Communicator::barrier takes an optional backend argument, which we // don't use yet. self.barrier(); }, R"( Performs a blocking barrier across all ranks. )"); } void bindDeviceMesh(py::module& nvfuser) { py::class_<DeviceMesh> device_mesh(nvfuser, "DeviceMesh", py::module_local()); device_mesh.def( py::init([](at::Tensor devices) { return std::make_unique<DeviceMesh>(std::move(devices)); }), py::arg("devices"), R"( Create a new DeviceMesh from torch.Tensor. )"); device_mesh.def( py::init([](const std::vector<int64_t>& devices) { return std::make_unique<DeviceMesh>(at::tensor(devices)); }), py::arg("devices"), R"( Create a new DeviceMesh from an integer list, for backward compatibility. )"); device_mesh.def("__repr__", [](const DeviceMesh& self) { std::stringstream ss; ss << self; return ss.str(); }); device_mesh .def_property_readonly( "size", [](const DeviceMesh& self) -> int64_t { return self.size(); }, R"( Returns the number of devices in the mesh. )") .def_property_readonly( "shape", [](const DeviceMesh& self) -> at::IntArrayRef { return self.shape(); }, R"( Returns the shape of the mesh. )"); device_mesh.def( "shard_tensor", [](const DeviceMesh& self, at::Tensor tensor, const int64_t axis) -> at::Tensor { return shardTensor1D(tensor, axis, self); }, py::arg("tensor"), py::arg("axis"), R"( Shards the input tensor along `axis`. Returns the sharded tensor.)"); } void bindSharding(py::module& nvfuser) { py::class_<Sharding>(nvfuser, "Sharding", py::module_local()) .def_property_readonly( "mesh", &Sharding::mesh, R"( Returns the device mesh of the sharding. )", py::return_value_policy::reference) .def( "axis_sharded_on", &Sharding::axisShardedOn, py::arg("parallel_type"), R"( Returns the axis sharded on the given parallel type. If the distributed tensor is replicated on that parallel type, returns -1. )"); } void bindMultiDeviceExecutor(py::module& nvfuser) { // Bind params type under the multidevice submodule. We'll alias it to the // top-level module in bindMultiDevice to allow direct imports. py::class_<MultiDeviceExecutorParams>(nvfuser, "MultiDeviceExecutorParams") .def(py::init<>()) .def_property( "use_allocation_cache", [](const MultiDeviceExecutorParams& self) { return self.executor.use_allocation_cache; }, [](MultiDeviceExecutorParams& self, bool value) { self.executor.use_allocation_cache = value; }) .def_property( "number_of_streams", [](const MultiDeviceExecutorParams& self) { return self.executor.number_of_streams; }, [](MultiDeviceExecutorParams& self, int64_t value) { self.executor.number_of_streams = value; }) .def_property( "backend_type", [](const MultiDeviceExecutorParams& self) { return self.lower.communicator_backend; }, [](MultiDeviceExecutorParams& self, CommunicatorBackend value) { self.lower.communicator_backend = value; }) .def_property( "offset_stream_indexing_by_rank", [](const MultiDeviceExecutorParams& self) { return self.lower.offset_stream_indexing_by_rank; }, [](MultiDeviceExecutorParams& self, bool value) { self.lower.offset_stream_indexing_by_rank = value; }) .def_property( "inter_stream_synchronization", [](const MultiDeviceExecutorParams& self) { return self.lower.inter_stream_synchronization; }, [](MultiDeviceExecutorParams& self, bool value) { self.lower.inter_stream_synchronization = value; }); py::class_<MultiDeviceExecutor> multi_device_executor( nvfuser, "MultiDeviceExecutor"); multi_device_executor.def( py::init( [](const Fusion& fusion, const MultiDeviceExecutorParams& params) { return std::make_unique<MultiDeviceExecutor>( std::make_unique<Fusion>(fusion), Communicator::getInstance(), params); }), R"( Create a new MultiDeviceExecutor. Parameters ---------- fusion : Fusion The fusion to be executed. params : MultiDeviceExecutorParams Parameters configuring the executor and communicator backend. Examples -------- >>> params = MultiDeviceExecutorParams() >>> params.backend_type = CommunicatorBackend.nccl >>> multi_device_executor = MultiDeviceExecutor(fusion, params) >>> outputs = multi_device_executor.run(inputs) )", py::arg("fusion"), py::arg("params")); multi_device_executor.def( "__str__", [](MultiDeviceExecutor& self) { std::stringstream ss; self.print(ss); return ss.str(); }, R"( Return a string representing the MultiDeviceExecutor. )"); multi_device_executor.def( "run", [](MultiDeviceExecutor& self, const py::iterable& args) { KernelArgumentHolder outputs = self.runWithInput(from_pyiterable(args)); return to_tensor_vector(outputs); }, R"( Run the fusion with the given arguments. Parameters ---------- args : KernelArgumentHolder The input arguments for the fusion. Returns ------- list of Tensor The output tensors containing the results. )", py::arg("args")); } } // namespace void bindMultiDevice(py::module& nvfuser) { py::module_ nvf_multidevice = nvfuser.def_submodule( "multidevice", "This submodule contains all multi-device features for NvFuser."); bindCommunicator(nvf_multidevice); bindDeviceMesh(nvf_multidevice); bindSharding(nvf_multidevice); bindMultiDeviceExecutor(nvf_multidevice); // Standalone shard_tensor function that takes TensorView nvf_multidevice.def( "shard_tensor", [](at::Tensor tensor, TensorView* tv) -> at::Tensor { return shardTensor(tensor, tv); }, py::arg("tensor"), py::arg("tv"), R"( Shards the input tensor according to the TensorView's parallelization and device mesh. This function automatically: - Determines the sharding axis from the TensorView's DIDx parallelization - Extracts the device mesh from the TensorView - Uses the current rank from the Communicator singleton - Returns the original tensor if the TensorView is not sharded Parameters ---------- tensor : torch.Tensor The unsharded tensor to shard tv : TensorView The TensorView that defines the sharding pattern (must have device mesh and DIDx parallelization if sharded) Returns ------- torch.Tensor The sharded tensor for the current device Examples -------- >>> mesh = nvfuser.multidevice.DeviceMesh(torch.arange(num_devices)) >>> with nvfuser.FusionDefinition() as fd: ... inp_tv = fd.define_tensor([-1, -1]) ... inp_tv.set_device_mesh(mesh) ... inp_tv.axis(0).parallelize(nvfuser.ParallelType.mesh_x) ... # ... rest of fusion definition >>> >>> unsharded_tensor = torch.randn(num_devices, 4) >>> sharded_tensor = nvfuser.multidevice.shard_tensor(unsharded_tensor, inp_tv) )"); } } // namespace nvfuser::python |
| You are visitor 01402503 — GeezerVison does not store the pages you look up. Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required. |