Facilities for PyTorch DTensor
Installation
Install the project using the following command:
pip install -e .
Visualize Sharding
Since transitioning from the PyTorch Distributed team at Meta to working on Apple Foundation Models at Apple, I've been utilizing AXLearn, which is built on top of JAX.
To understand JAX's distributed arrays and how sharding across a device mesh could implement various parallelism strategies, I've used jax.debug.visualize_array_sharding. In this LinkedIn post, co-authored with Wanchao Liang, we demonstrate the benefits of visualizing sharding.
It's encouraging to see my former team releasing PyTorch's DTensor and showcasing its value through the development of FSDP2 and TorchTitan. DTensor also includes a sharding visualization feature that utilizes tabulate to produce a basic visual representation.
This project introduces an enhanced visualization for DTensor's sharding, inspired by JAX's approach. Run the following command to see how it works.
OMP_NUM_THREADS=1 torchrun --nproc_per_node=4 dtensor/visualize_example.py
This will display a list of sharding/mesh combinations. For each combination, visualizations generated by both the original and the new methods will be presented.
1 maintainer
Metadata
1 maintainer
Release files for dtensor 0.0.1
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | |
|---|---|---|---|---|
| dtensor-0.0.1.tar.gz | 6.1 kB | Without provenance | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| dtensor-0.0.1-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 12.9 kB
Release files / dtensor-0.0.1.tar.gz
| Download URL | dtensor-0.0.1.tar.gz |
|---|---|
| Size | 6.1 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
cef095b6e6d3bac56ddab72fdc31d6f0bdc9343f78aa94a1f11ec335f09fe22b
|
|
BLAKE2b-256 checksum How to use checksums |
6f26b27cfb2100d05ded6c53c7ec815d0e89cb5e62c8345795e59b1fed4bf9b5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.12.9
|
Release files / dtensor-0.0.1-py3-none-any.whl
| Download URL | dtensor-0.0.1-py3-none-any.whl |
|---|---|
| Size | 6.8 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
e038fdf8a4cb1870aef09fcd532e8cfb5b053c091cd816304b9cc5ee3aebc960
|
|
BLAKE2b-256 checksum How to use checksums |
f46b55d7a24136232f71136144838859fe6ecd515937f85d00a6435a19019f2b
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.12.9
|
Security for dtensor 0.0.1
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
