[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"doc-detail-83013-en":3,"doc-seo-83013-105":30,"detail-sidebar-cat-0-en-105":91},{"code":4,"msg":5,"data":6},0,"success",{"doc_id":7,"user_id":8,"nickname":9,"user_avatar":10,"doc_module":4,"category_id":11,"category_name":12,"doc_title":13,"doc_description":14,"doc_content":15,"file_id":16,"file_url":17,"file_type":18,"file_size":19,"view_count":20,"is_deleted":4,"is_public":21,"is_downloadable":21,"audit_status":21,"page_count":22,"language":23,"language_code":24,"site_id":25,"html_lang":24,"table_of_contents":26,"faqs":27,"seo_title":13,"seo_description":14,"update_tm":28,"read_time":29},83013,7971461740909,"Levi","https://ap-avatar.wpscdn.com/davatar_155a257f0dc6eb9ab79c44ca47cae57d",8,"Research & Report","MatrixFSDP: Communication-Free Matrix Optimizers Under Zero-3 Parameter Sharding","MatrixFSDP targets the systems mismatch between matrix-structured optimizers and FSDP/ZeRO-3 sharding. Muon improves convergence and token efficiency using Newton–Schulz spectrum-balanced orthogonalization, but it requires full 2D matrix inputs, conflicting with shard-only optimizer views in ZeRO-3. MatrixFSDP reshapes ZeRO-3 shard ownership so one data-parallel rank holds whole matrices while others keep empty shards, enabling local Newton–Schulz execution without optimizer-step matrix collectives. It adds MatrixShard metadata, deterministic owner-segment P2P collectives, and owner-buffer pinning. On 64 A100s it reduces optimizer-step latency 4.2× versus FSDP2-Muon on one node and 54.6× on eight nodes, reaching up to 2.15× end-to-end speedup and enabling model sizes beyond ZeRO-1’s 80 GB limit.","arXiv :2607 .05895v 1 [ cs .DC] 7 Jul 2026  \nMATRIXFSDP: COMMUNICATION-FREE MATRIX OPTIMIZERS UNDER  \nZERO-3 PARAMETER SHARDING  \nMing Gao 1 Yanwu Xu 2 Hao Zhang 3  \nABSTRACT  \nMatrix optimizers such as Muon are attractive for large-scale training because they can improve convergence and token efficiency over coordinate-wise optimizers. Muon does this by orthogonalizing momentum-smoothed matrix updates with Newton–Schulz, producing spectrum-balanced updates that require the complete 2D matrix as input. This exposes a systems mismatch: FSDP/ZeRO-3 saves memory by making the optimizer see shards, not whole matrices. Existing systems therefore either reconstruct matrices at every optimizer step, paying weightsized communication after backward, or make the update local by using ZeRO-1 owner placement with full parameters resident. MatrixFSDP takes a third path: it changes where ZeRO-3 shards live, not the optimizer being computed. For each 2D weight, one data-parallel rank owns the whole matrix and the other ranks hold empty shards; non-matrix tensors are packed into tail owners and stay on AdamW. The ordinary backward reduction then lands the full Muon input on the owner, so Newton–Schulz runs locally with no optimizer-step matrix collective. Forward and backward still materialize and reshard parameters; the runtime challenge is to make that uneven layout efficient and correct. MatrixFSDP does so with MatrixShard metadata, a balance-aware owner planner, deterministic owner-segment P2P collectives, owner-buffer pinning, and owner-shard checkpoint resharding. The resulting update matches full-matrix Muon while preserving ZeRO-3-scale memory: on 64 A100s, MatrixFSDP reduces optimizer-step latency over stock FSDP2-Muon by 4.2 × on one node and 54.6 × on eight nodes, reaches up to 2.15 × end-to-end speedup, and runs model sizes where ZeRO-1 owner placement exceeds an 80 GB GPU.  \n1 INTRODUCTION  \nAs language models grow, training systems must do more than fit larger parameter counts: they must also reduce the time and tokens needed to reach a target quality. Matrix-structured optimizers are one promising direction. Muon (Jordan et al., 2024), for example, improves token efficiency over AdamW and has been adopted in recent large-scale training runs (Liu et al., 2025 ; Essential AI, 2025); Shampoo (Gupta et al., 2018 ; Anil et al., 2020) and SOAP (Vyas et al., 2024) pursue the same goal with matrix preconditioners. The catch is that these optimizers ask the distributed runtime for a different object. FSDP2 / ZeRO-3 (Rajbhandari et al., 2020 ; Zhao et al., 2023) shard parameters, gradients, and optimizer state into per-rank slices, which is exactly the right contract for coordinatewise AdamW. Muon instead applies Newton–Schulz to a whole 2D gradient: a row-band is not the same optimizer input as the full matrix. Step efficiency thus becomes a ZeRO-3 systems problem.  \n1University of Pittsburgh 2 Google 3Tsinghua University. Correspondence to: Ming Gao \u003C[dujinshidai30@gmail.com](dujinshidai30@gmail.com)>, Yanwu Xu \u003C[yvansxu@gmail.com](yvansxu@gmail.com)>, Hao Zhang \u003Czhang[hao_990127@163.com](hao_990127@163.com)> .  \nPreprint.  \nExisting systems expose two incomplete choices. FSDP2-Muon keeps ZeRO-3 memory but reconstructs each matrix at optimizer time, placing weight-sized communication after backward. ZeRO-1 owner placement makes the matrix operation local, but keeps full parameters resident on every rank. The missing combination is whole-matrix optimizer inputs without abandoning ZeRO-3 parameter memory.  \nThe hard part is the runtime contract around that placement. At the optimizer boundary, Muon should see the same complete 2D gradient as a gathered reference; between steps, no rank should retain a full model copy; and forward/backward must still use FSDP2’s transient materialize/reshard buffers.  \nWe therefore present MatrixFSDP, an FSDP2 runtime for owner-shaped ZeRO-3 shards:  \n▷ Owner-shaped placement. Each data-parallel-sharded 2","cbCaiolYt1AOsGYo","https://ap.wps.com/l/cbCaiolYt1AOsGYo","pdf",289525,4,1,10,"English","en",105,"# Introduction\n# Motivation","[{\"question\":\"What core problem does MatrixFSDP address in distributed training?\",\"answer\":\"MatrixFSDP addresses the mismatch where matrix-structured optimizers like Muon require full 2D matrix inputs, while FSDP/ZeRO-3 exposes only sharded views of parameters and optimizer state to each rank.\"},{\"question\":\"How does MatrixFSDP enable Muon’s Newton–Schulz steps without matrix all-gathers?\",\"answer\":\"MatrixFSDP reshapes ZeRO-3 shard ownership so that for each 2D weight, one data-parallel owner rank holds the whole matrix while other ranks keep empty shards; the backward reduction deposits the full Muon input on that owner, so Newton–Schulz runs locally.\"},{\"question\":\"What runtime components make MatrixFSDP practical and efficient?\",\"answer\":\"MatrixFSDP uses MatrixShard metadata to carry an owner plan through optimizer routing and resharding, employs deterministic owner-segment P2P collectives to avoid invalid equal-size assumptions, and applies owner-buffer pinning plus owner-shard checkpoint resharding.\"}]",1784184666,25,{"code":4,"msg":31,"data":32},"ok",{"site_id":25,"language":24,"slug":33,"title":13,"keywords":34,"description":14,"schema_data":35,"social_meta":86,"head_meta":88,"extra_data":90,"updated_unix":28},"matrixfsdp-communication-free-matrix-optimizers-under-zero-3-parameter-sharding","",{"@graph":36,"@context":85},[37,53,68],{"@type":38,"itemListElement":39},"BreadcrumbList",[40,44,48,51],{"item":41,"name":42,"@type":43,"position":21},"https://docshare.wps.com","Home","ListItem",{"item":45,"name":46,"@type":43,"position":47},"https://docshare.wps.com/document/","Document",2,{"item":49,"name":12,"@type":43,"position":50},"https://docshare.wps.com/document/research-report/",3,{"item":52,"name":13,"@type":43,"position":20},"https://docshare.wps.com/document/matrixfsdp-communication-free-matrix-optimizers-under-zero-3-parameter-sharding/83013/",{"url":52,"name":13,"@type":54,"author":55,"headline":13,"publisher":57,"fileFormat":60,"inLanguage":24,"description":14,"dateModified":61,"datePublished":62,"encodingFormat":60,"isAccessibleForFree":63,"interactionStatistic":64},"DigitalDocument",{"name":9,"@type":56},"Person",{"url":41,"name":58,"@type":59},"DocShare","Organization","application/pdf","2026-07-23","2026-07-16",true,{"@type":65,"interactionType":66,"userInteractionCount":20},"InteractionCounter",{"@type":67},"ViewAction",{"@type":69,"mainEntity":70},"FAQPage",[71,77,81],{"name":72,"@type":73,"acceptedAnswer":74},"What core problem does MatrixFSDP address in distributed training?","Question",{"text":75,"@type":76},"MatrixFSDP addresses the mismatch where matrix-structured optimizers like Muon require full 2D matrix inputs, while FSDP/ZeRO-3 exposes only sharded views of parameters and optimizer state to each rank.","Answer",{"name":78,"@type":73,"acceptedAnswer":79},"How does MatrixFSDP enable Muon’s Newton–Schulz steps without matrix all-gathers?",{"text":80,"@type":76},"MatrixFSDP reshapes ZeRO-3 shard ownership so that for each 2D weight, one data-parallel owner rank holds the whole matrix while other ranks keep empty shards; the backward reduction deposits the full Muon input on that owner, so Newton–Schulz runs locally.",{"name":82,"@type":73,"acceptedAnswer":83},"What runtime components make MatrixFSDP practical and efficient?",{"text":84,"@type":76},"MatrixFSDP uses MatrixShard metadata to carry an owner plan through optimizer routing and resharding, employs deterministic owner-segment P2P collectives to avoid invalid equal-size assumptions, and applies owner-buffer pinning plus owner-shard checkpoint resharding.","https://schema.org",{"og:url":52,"og:type":87,"og:title":13,"og:site_name":58,"og:description":14},"article",{"robots":89,"canonical":52},"index,follow",{"doc_id":7,"site_id":25},{"code":4,"msg":5,"data":92},[93,97,101,105,110,115,120,123,128,131,134],{"id":21,"doc_module":4,"doc_module_name":46,"category_name":94,"show_sort_weight":95,"slug":96},"Story & Novel",90,"story-novel",{"id":47,"doc_module":4,"doc_module_name":46,"category_name":98,"show_sort_weight":99,"slug":100},"Literature",80,"literature",{"id":20,"doc_module":4,"doc_module_name":46,"category_name":102,"show_sort_weight":103,"slug":104},"Exam",70,"exam",{"id":106,"doc_module":4,"doc_module_name":46,"category_name":107,"show_sort_weight":108,"slug":109},5,"Comic",60,"comic",{"id":111,"doc_module":4,"doc_module_name":46,"category_name":112,"show_sort_weight":113,"slug":114},6,"Technology",50,"technology",{"id":116,"doc_module":4,"doc_module_name":46,"category_name":117,"show_sort_weight":118,"slug":119},7,"Healthcare",40,"healthcare",{"id":11,"doc_module":4,"doc_module_name":46,"category_name":12,"show_sort_weight":121,"slug":122},30,"research-report",{"id":124,"doc_module":4,"doc_module_name":46,"category_name":125,"show_sort_weight":126,"slug":127},9,"Religion & Spirituality",20,"religion-spirituality",{"id":126,"doc_module":4,"doc_module_name":46,"category_name":129,"show_sort_weight":126,"slug":130},"World Cup","world-cup",{"id":22,"doc_module":4,"doc_module_name":46,"category_name":132,"show_sort_weight":22,"slug":133},"Lifestyle","lifestyle",{"id":135,"doc_module":4,"doc_module_name":46,"category_name":136,"show_sort_weight":106,"slug":137},19,"General","general"]