{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/a-distributed-data-parallel-pytorch","title":"A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale","arxiv_id":"2309.06497","date":"2023-09-12","proceeding":null,"authors":["Hao-Jun Michael Shi","Tsung-Hsien Lee","Shintaro Iwasaki","Jose Gallego-Posada","Zhijing Li","Kaushik Rangadurai","Dheevatsa Mudigere","Michael Rabbat"],"abstract":"Shampoo is an online and stochastic optimization algorithm belonging to the AdaGrad family of methods for training neural networks. It constructs a block-diagonal preconditioner where each block consists of a coarse Kronecker product approximation to full-matrix AdaGrad for each parameter of the neural network. In this work, we provide a complete description of the algorithm as well as the performance optimizations that our implementation leverages to train deep networks at-scale in PyTorch. Our implementation enables fast multi-GPU distributed data-parallel training by distributing the memory and computation associated with blocks of each parameter via PyTorch's DTensor data structure and performing an AllGather primitive on the computed search directions at each iteration. This major performance enhancement enables us to achieve at most a 10% performance reduction in per-step wall-clock time compared against standard diagonal-scaling-based adaptive gradient methods. We validate our implementation by performing an ablation study on training ImageNet ResNet50, demonstrating Shampoo's superiority over standard training recipes with minimal hyperparameter tuning.","url_abs":"https://arxiv.org/abs/2309.06497v1","url_pdf":"https://arxiv.org/pdf/2309.06497v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"a-distributed-data-parallel-pytorch","repo_url":"https://github.com/facebookresearch/optimizers/tree/main/distributed_shampoo","is_official":1,"mentioned_in_paper":0,"mentioned_in_github":0,"framework":"pytorch","reach":{"status":"ok","spdx":"NOASSERTION"}},{"paper_slug":"a-distributed-data-parallel-pytorch","repo_url":"https://github.com/mlcommons/submissions_algorithms","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"a-distributed-data-parallel-pytorch","repo_url":"https://github.com/yorkerlin/remove-the-square-root","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok"}}],"tasks":[{"task_slug":null,"task_name":"GPU"},{"task_slug":"stochastic-optimization","task_name":"Stochastic Optimization"}],"methods":[{"method_slug":"adagrad","method_name":"AdaGrad"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2309.06497","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2309.06497"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/facebookresearch/optimizers/tree/main/distributed_shampoo","reach":{"status":"ok","spdx":"NOASSERTION"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/mlcommons/submissions_algorithms","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/yorkerlin/remove-the-square-root","reach":{"status":"ok"}}],"summary":{"unverified":2},"by_repo_kind":{"official":{"samples":2,"ran":0,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":2,"samples":[{"code_sha256_prefix":"509c55e7e1ef2138","entry":"create_hybrid_shard_process_groups","repo":"facebookresearch/optimizers","repo_kind":"official","path":"distributed_shampoo/distributor/shampoo_dist_utils.py","file_url":"https://github.com/facebookresearch/optimizers/blob/HEAD/distributed_shampoo/distributor/shampoo_dist_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NOASSERTION","inline_ok":false,"mcp_get_code":{"code_sha256":"509c55e7e1ef2138"}},{"code_sha256_prefix":"3e3c14dedb9c5a3b","entry":"get_device_mesh","repo":"facebookresearch/optimizers","repo_kind":"official","path":"distributed_shampoo/distributor/shampoo_dist_utils.py","file_url":"https://github.com/facebookresearch/optimizers/blob/HEAD/distributed_shampoo/distributor/shampoo_dist_utils.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NOASSERTION","inline_ok":false,"mcp_get_code":{"code_sha256":"3e3c14dedb9c5a3b"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}