[DTensor] Open Source

[leonardo0lyj]

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.

[yiakwy-xpu-ml-framework-team]

Integrated pytorch distributed tensor and device mesh into a large scale parallel framework like megatron is awesome. Good job!