Kumar Ashutosh 405a0040cf Adds tool to visualize sharding (#114307) ... This pull request adds a tool to visualize sharding. It uses the device_mesh and placement details to construct a visualization of the split of a torch dtensor. Things to fix: - [x] This implementation only uses the first element of the placement tuple, when can there be more than one elements? - [x] The calculation of the split is happening here but maybe it is already done somewhere internally in Shard class and can we directly call that here? Fixes #108746 Pull Request resolved: https://github.com/pytorch/pytorch/pull/114307 Approved by: https://github.com/wanchaol		2023-12-12 06:18:03 +00:00
..
checkpoint_example.py	[TP] fully rewrite Tensor Parallel APIs (#114732)	2023-12-02 08:18:12 +00:00
convnext_example.py	[dtensor] support convolution ops (#113123)	2023-11-20 21:01:28 +00:00
visualize_sharding_example.py	Adds tool to visualize sharding (#114307)	2023-12-12 06:18:03 +00:00