veScale DModule (Distributed Module) provides a single-device abstraction for multiple-device nn.Module and empowers user to write distributed training/inference code as if on a single device (i.e., SPMD).
DModule unifies Module-level Tensor Parallelism and Sequence Parallelism by transparently handling distributed logic under the hood:
convert Tensor to DTensor within a nn.Module
manage DTensor sharding and resharding during forward and backward
configure (re)sharding of DTensor via Module-level API parallelize_module() with given sharding_plan
allow sharding_plan to be either:
imported from a pre-defined "plan zoo"
given by "manually written" json
[experimental] given by "automatical plan generation" of veScale
handles gradient synchronization automatically in backward
support deferred initialization with deferred_init() (i.e., initialize with Fake Tensor without allocating memory, then shard Fake Tensor with TP, and then materialize only a shard of Tensor on device)
support third-party plug-in Module (e.g. APEX)
provide patch interface for customized Module-level hacking
extend to optional DP, optional FSDP, and optional EP (in the future)
provide debuggability for easy dumping, printing, listing, etc.
Difference of veScale DModule from PyTorch parallelize_module?
veScale DModule is inspired by PyTorch's parallelize_module, but is developed with explicit Module-level abstraction with complete features for our production usage.
veScale DModule extends PyTorch parallelize_module with extra features as below (i.e., major differences from PyTorch-v2.2.0):
nD Tensor Parallelism
Sequence Parallelism
auto gradient synchronization
deferred initialization
third-party plug-in Module
module-level patch interface
[experimental] automatical plan generation
Credit to veScale DModule Team
This endeavor would not have been possible without the contribution of our DModule team which includes but not limited to: @SerailHydra @Vremold @JsBlueCat @jc-bytedance @MackZackA @leonardo0lyj.
Also thanks to the great guidance and leadership of: @liwenchangbdbz @pengyanghua @eric-haibin-lin @Meteorix
Credit to PyTorch DTensor Team
Once again, 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
DModule, Yo~What is veScale DModule?
veScale
DModule(Distributed Module) provides a single-device abstraction for multiple-devicenn.Moduleand empowers user to write distributed training/inference code as if on a single device (i.e., SPMD).DModuleunifies Module-level Tensor Parallelism and Sequence Parallelism by transparently handling distributed logic under the hood:TensortoDTensorwithin ann.ModuleDTensorsharding and resharding during forward and backwardDTensorvia Module-level APIparallelize_module()with givensharding_plansharding_planto be either:deferred_init()(i.e., initialize withFakeTensor without allocating memory, then shard Fake Tensor with TP, and then materialize only a shard ofTensoron device)Difference of veScale DModule from PyTorch
parallelize_module?veScale
DModuleis inspired by PyTorch'sparallelize_module, but is developed with explicit Module-level abstraction with complete features for our production usage.veScale
DModuleextends PyTorchparallelize_modulewith extra features as below (i.e., major differences from PyTorch-v2.2.0):Credit to veScale DModule Team
This endeavor would not have been possible without the contribution of our DModule team which includes but not limited to: @SerailHydra @Vremold @JsBlueCat @jc-bytedance @MackZackA @leonardo0lyj.
Also thanks to the great guidance and leadership of: @liwenchangbdbz @pengyanghua @eric-haibin-lin @Meteorix
Credit to PyTorch DTensor Team
Once again, 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.