# Copyright (c) Meta Platforms, Inc. and affiliatesimporttorchimporttorch.distributed.tensor._ops# force import all built-in dtensor opsfromtorch.distributed.device_meshimportDeviceMesh,init_device_mesh# noqa: F401fromtorch.distributed.tensor._apiimport(distribute_module,distribute_tensor,DTensor,empty,full,ones,rand,randn,zeros,)fromtorch.distributed.tensor.placement_typesimport(Partial,Placement,Replicate,Shard,)fromtorch.optim.optimizerimport(_foreach_supported_typesas_optim_foreach_supported_types,)fromtorch.utils._foreach_utilsimport(_foreach_supported_typesas_util_foreach_supported_types,)# All public APIs from dtensor package__all__=["DTensor","distribute_tensor","distribute_module","Shard","Replicate","Partial","Placement","ones","empty","full","rand","randn","zeros",]# For weights_only torch.loadfrom._dtensor_specimportDTensorSpecas_DTensorSpec,TensorMetaas_TensorMetatorch.serialization.add_safe_globals([DeviceMesh,_DTensorSpec,_TensorMeta,DTensor,Partial,Replicate,Shard,])# Append DTensor to the list of supported types for foreach implementation for optimizer# and clip_grad_norm_ so that we will try to use foreach over the for-loop implementation on CUDA.ifDTensornotin_optim_foreach_supported_types:_optim_foreach_supported_types.append(DTensor)ifDTensornotin_util_foreach_supported_types:_util_foreach_supported_types.append(DTensor)# Set namespace for exposed private namesDTensor.__module__="torch.distributed.tensor"distribute_tensor.__module__="torch.distributed.tensor"distribute_module.__module__="torch.distributed.tensor"ones.__module__="torch.distributed.tensor"empty.__module__="torch.distributed.tensor"full.__module__="torch.distributed.tensor"rand.__module__="torch.distributed.tensor"randn.__module__="torch.distributed.tensor"zeros.__module__="torch.distributed.tensor"
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.