In this PR, we open source our DTensor and Dockerfile, Yo~
veScale is a PyTorch native framework rooted in PyTorch DTensor. veScale DTensor shares the majority of code of PyTorch DTensor, but extends it with extra features as below (i.e., major differences from PyTorch DTensor v2.2.0) for our production usage:
enabled DTensor support for third-party plug-in ops (e.g., APEX) by unleashing DTensor.data_ptr and handling asynchronous collective tensors (e.g., in from_local, to_local, redistribute)
make implicit _Partial to explicit Partial placement for optimized initialization, output, and checkpoint (with an extra dispatch mode)
enabled DTensor ops that were not implemented in PyTorch:
argmax and argmin
topk
_unique2
scatter_ and scatter
select
alias
index_put_ and index_put
_scaled_dot_product_flash_attention
_scaled_dot_product_efficient_attention
support uneven sharding in conversion between DTensor and torch.Tensor
decoupled special op handling that bypasses DTensor dispatching (_bypass_for_dispatch)
enabled patching before (_pre_patch_for_dispatch) and after (_post_patch_for_dispatch) DTensor dispatch, for adding user's custom dispatching logic without coupling original dispatch logic
enabled short-cut for ops to bypass sharding propagation entirely (_bypass_for_sharding_prop)
bypassed tensor_meta propagation for ops:
with output DTensor as pure Replicate, by using local output Tensor's tensor_meta
with registered tensor_meta propagation under dtensor/ops (e.g., conv, slice, copy, clone, bucketize, t)
excluding ops in recompute_tensor_meta_list (e.g., clone, native_dropout, nll_loss_forward)
enabled DeviceMesh on meta device type
enabled DeviceMesh initialization from an existing processs group
enabled DeviceMesh being split into a list of sub meshes
disabled redistributed input:
torch DTensor allows each op select its best sharding strategy for input-output sharding based on a cost model capturing input redistribution communication and then redistributes input DTensor to selected input-sharding.
But we currently disable this feature (via environment var VESCALE_DISABLE_REDISTRIBUTE), as we don't expect uncontrollable resharding and implicit communication in DTensor dispatch for production. (Ideally, all resharding and communication should be controlled by the end users.)
support deferred initiailization and materialization for DTensor with extended torchdistx
[experimental] developed InterleavedShard placement to support merged QKV in MHA
Credit to veScale DTensor Team
This endeavor would not have been possible without the contribution of our DTensor team which includes but not limited to: @Vremold @JsBlueCat @SerailHydra @wenlei-bao @Hao-Gong @jc-bytedance @yaochengji @Connor-XY @cheimu @MackZackA @leonardo0lyj.
Also thanks to the great guidance and leadership of: @liwenchangbdbz @pengyanghua @eric-haibin-lin @Meteorix
Credit to PyTorch DTensor Team
We would like to sincerely acknowledge the assistance of and collaboration with the PyTorch DTensor team which includes but not limited to: @wanchaol @XilunWu @wz337 @tianyu-l @fduwjj @awgu @yifuwang @wconstab @ezyang @mrshenli.
In this PR, we open source our
DTensor
andDockerfile
, Yo~veScale is a PyTorch native framework rooted in PyTorch
DTensor
. veScaleDTensor
shares the majority of code of PyTorch DTensor, but extends it with extra features as below (i.e., major differences from PyTorch DTensor v2.2.0) for our production usage:enabled DTensor support for third-party plug-in ops (e.g.,
APEX
) by unleashingDTensor.data_ptr
and handling asynchronous collective tensors (e.g., infrom_local
,to_local
,redistribute
)make implicit
_Partial
to explicitPartial
placement for optimized initialization, output, and checkpoint (with an extra dispatch mode)enabled DTensor ops that were not implemented in PyTorch:
argmax
andargmin
topk
_unique2
scatter_
andscatter
select
alias
index_put_
andindex_put
_scaled_dot_product_flash_attention
_scaled_dot_product_efficient_attention
support uneven sharding in conversion between
DTensor
andtorch.Tensor
decoupled special op handling that bypasses DTensor dispatching (
_bypass_for_dispatch
)enabled patching before (
_pre_patch_for_dispatch
) and after (_post_patch_for_dispatch
) DTensor dispatch, for adding user's custom dispatching logic without coupling original dispatch logicenabled short-cut for ops to bypass sharding propagation entirely (
_bypass_for_sharding_prop
)bypassed
tensor_meta
propagation for ops:Replicate
, by using local output Tensor'stensor_meta
tensor_meta
propagation underdtensor/ops
(e.g.,conv
,slice
,copy
,clone
,bucketize
,t
)recompute_tensor_meta_list
(e.g.,clone
,native_dropout
,nll_loss_forward
)enabled
DeviceMesh
onmeta
device typeenabled
DeviceMesh
initialization from an existing processs groupenabled
DeviceMesh
being split into a list of sub meshesdisabled redistributed input:
VESCALE_DISABLE_REDISTRIBUTE
), as we don't expect uncontrollable resharding and implicit communication in DTensor dispatch for production. (Ideally, all resharding and communication should be controlled by the end users.)support deferred initiailization and materialization for DTensor with extended
torchdistx
[experimental] developed
InterleavedShard
placement to support merged QKV in MHACredit to veScale DTensor Team
This endeavor would not have been possible without the contribution of our DTensor team which includes but not limited to: @Vremold @JsBlueCat @SerailHydra @wenlei-bao @Hao-Gong @jc-bytedance @yaochengji @Connor-XY @cheimu @MackZackA @leonardo0lyj.
Also thanks to the great guidance and leadership of: @liwenchangbdbz @pengyanghua @eric-haibin-lin @Meteorix
Credit to PyTorch DTensor Team
We would like to sincerely acknowledge the assistance of and collaboration with the PyTorch DTensor team which includes but not limited to: @wanchaol @XilunWu @wz337 @tianyu-l @fduwjj @awgu @yifuwang @wconstab @ezyang @mrshenli.