07b weight sync deep dive

import%20marimo%0A%0A__generated_with%20%3D%20%220.19.9%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20Shared%20imports%20for%20the%20notebook%0A%20%20%20%20import%20time%0A%20%20%20%20import%20torch%0A%20%20%20%20from%20monarch.actor%20import%20Actor%2C%20endpoint%2C%20this_host%2C%20current_rank%0A%20%20%20%20from%20monarch.rdma%20import%20RDMABuffer%2C%20RDMAAction%2C%20is_rdma_available%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20Actor%2C%0A%20%20%20%20%20%20%20%20RDMAAction%2C%0A%20%20%20%20%20%20%20%20RDMABuffer%2C%0A%20%20%20%20%20%20%20%20current_rank%2C%0A%20%20%20%20%20%20%20%20endpoint%2C%0A%20%20%20%20%20%20%20%20is_rdma_available%2C%0A%20%20%20%20%20%20%20%20this_host%2C%0A%20%20%20%20%20%20%20%20time%2C%0A%20%20%20%20%20%20%20%20torch%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20RDMA%20Deep%20Dive%0A%0A%20%20%20%20Notebook%2007%20covered%20weight%20synchronization%20end-to-end%3A%20why%20RDMA%20matters%20for%20async%20RL%2C%0A%20%20%20%20the%20magic%20pointer%20pattern%2C%20circular%20buffers%2C%20and%20re-sharding.%20We%20used%20%60RDMABuffer%60%20and%0A%20%20%20%20%60read_into()%60%20as%20black%20boxes%20%E2%80%94%20call%20the%20API%2C%20weights%20appear.%0A%0A%20%20%20%20This%20notebook%20goes%20deeper%20into%20how%20RDMA%20actually%20works.%20We'll%20walk%20through%20the%20ibverbs%0A%20%20%20%20primitives%20that%20power%20every%20transfer%2C%20understand%20the%20costs%20that%20separate%20a%20fast%0A%20%20%20%20implementation%20from%20a%20slow%20one%2C%20and%20see%20the%20concrete%20patterns%20for%20managing%20RDMA%20buffers%0A%20%20%20%20in%20production.%0A%0A%20%20%20%20**The%20central%20question%3A%20where%20do%20the%20milliseconds%20go%3F**%0A%0A%20%20%20%201.%20**ibverbs%20Internals**%20%E2%80%94%20Queue%20Pairs%2C%20Memory%20Registration%2C%20Completion%20Queues%0A%20%20%20%202.%20**RDMA%20Buffer%20Patterns**%20%E2%80%94%20Three%20approaches%20to%20managing%20registration%20and%20transfers%0A%0A%20%20%20%20*We%20focus%20on%20**ibverbs**%20because%20that's%20what%20Monarch's%20RDMA%20subsystem%20supports%20today%0A%20%20%20%20(InfiniBand%20and%20RoCE%20via%20Mellanox%2FNVIDIA%20ConnectX%20NICs).%20EFA%20(AWS%20Elastic%20Fabric%20Adapter)%0A%20%20%20%20is%20a%20relevant%20transport%20but%20not%20yet%20supported%20%E2%80%94%20it's%20actively%20being%20worked%20on.*%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%201.%20ibverbs%20Internals%0A%0A%20%20%20%20RDMA%20(Remote%20Direct%20Memory%20Access)%20lets%20one%20machine%20read%2Fwrite%20another%20machine's%20memory%0A%20%20%20%20directly%2C%20bypassing%20the%20kernel%20and%20CPU%20on%20both%20sides.%0A%0A%20%20%20%20%23%23%23%20The%20ibverbs%20Stack%0A%0A%20%20%20%20%60%60%60%0A%20%20%20%20%E2%94%8C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%90%0A%20%20%20%20%E2%94%82%20%20Application%20(PyTorch%2C%20Monarch%2C%20etc.)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%9C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%A4%0A%20%20%20%20%E2%94%82%20%20libibverbs%20%20(userspace%20RDMA%20API)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%9C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%A4%0A%20%20%20%20%E2%94%82%20%20Provider%20driver%20(mlx5%2C%20efa%2C%20rxe%2C%20etc.)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%9C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%A4%0A%20%20%20%20%E2%94%82%20%20Hardware%20(InfiniBand%20NIC%2C%20RoCE%20NIC%2C%20etc.)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%94%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%98%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20We%20focus%20on%20**InfiniBand**%20and%20**RoCE**%20(RDMA%20over%20Converged%20Ethernet)%20%E2%80%94%20the%20two%0A%20%20%20%20transports%20Monarch%20supports%20today%20via%20the%20ibverbs%20API.%0A%0A%20%20%20%20%23%23%23%20Key%20RDMA%20Operations%0A%0A%20%20%20%20%7C%20Operation%20%7C%20Description%20%7C%0A%20%20%20%20%7C-----------%7C-------------%7C%0A%20%20%20%20%7C%20%60RDMA_WRITE%60%20%7C%20Write%20to%20remote%20memory%20(one-sided)%20%7C%0A%20%20%20%20%7C%20%60RDMA_READ%60%20%7C%20Read%20from%20remote%20memory%20(one-sided)%20%7C%0A%20%20%20%20%7C%20%60SEND%2FRECV%60%20%7C%20Two-sided%20messaging%20(like%20TCP)%20%7C%0A%0A%20%20%20%20The%20magic%20is%20in%20%60RDMA_WRITE%60%20and%20%60RDMA_READ%60%20-%20they're%20**one-sided**%3A%0A%20%20%20%20-%20Remote%20CPU%20is%20not%20involved%0A%20%20%20%20-%20Remote%20application%20doesn't%20need%20to%20call%20anything%0A%20%20%20%20-%20NIC%20handles%20everything%20in%20hardware%0A%0A%20%20%20%20%23%23%23%20Memory%20Registration%0A%0A%20%20%20%20Before%20RDMA%2C%20memory%20must%20be%20**registered**%20with%20the%20NIC%3A%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20%23%20Conceptually%20(actual%20ibverbs%20API%20is%20in%20C)%0A%20%20%20%20mr%20%3D%20rdma_register_memory(buffer%2C%20size)%0A%20%20%20%20%23%20Returns%3A%0A%20%20%20%20%23%20%20%20-%20lkey%3A%20local%20access%20key%20(for%20local%20operations)%0A%20%20%20%20%23%20%20%20-%20rkey%3A%20remote%20access%20key%20(share%20with%20remote%20peer)%0A%20%20%20%20%23%20%20%20-%20addr%3A%20physical%2Fvirtual%20address%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20%60(addr%2C%20rkey)%60%20pair%20is%20a%20**remote-accessible%20pointer**.%20Share%20it%20with%20a%20peer%2C%0A%20%20%20%20and%20they%20can%20read%2Fwrite%20your%20memory%20directly.%0A%0A%20%20%20%20%23%23%23%20Queue%20Pair%20Setup%0A%0A%20%20%20%20Before%20any%20RDMA%20operations%2C%20you%20need%20to%20establish%20a%20**Queue%20Pair%20(QP)**%20between%0A%20%20%20%20sender%20and%20receiver.%20This%20is%20a%20one-time%20connection%20setup%3A%0A%0A%20%20%20%20%60%60%60%0A%20%20%20%20%E2%94%8C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%90%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%8C%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%90%0A%20%20%20%20%E2%94%82%20%20%20Sender%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20Receiver%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20Create%20QP%20%20%E2%94%82%20%E2%94%80%E2%94%80%E2%94%80%20exchange%20QP%20info%20%E2%94%80%E2%94%80%E2%94%80%E2%96%BA%20%E2%94%82%20%20Create%20QP%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20(qp_num%2C%20%20%20%E2%94%82%20%E2%97%84%E2%94%80%E2%94%80%20(qp_num%2C%20lid%2C%20gid)%20%E2%94%80%E2%94%80%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%20lid%2C%20gid)%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20Move%20QP%20to%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20Move%20QP%20to%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20RTR%20%E2%86%92%20RTS%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20RTR%20%E2%86%92%20RTS%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20Now%20ready%20%20%E2%94%82%20%E2%95%90%E2%95%90%E2%95%90%20RDMA%20operations%20%E2%95%90%E2%95%90%E2%95%90%E2%95%90%E2%96%BA%20%E2%94%82%20%20Now%20ready%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20for%20RDMA!%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20%20for%20RDMA!%20%20%E2%94%82%0A%20%20%20%20%E2%94%94%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%98%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%94%94%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%98%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20This%20is%20where%20**Monarch%20actors**%20shine.%20Because%20you%20can%20spawn%20arbitrary%20actors%2C%0A%20%20%20%20you%20can%20create%20**RDMA%20Manager%20actors**%20that%3A%0A%20%20%20%20-%20Initialize%20QPs%20on%20their%20respective%20hosts%0A%20%20%20%20-%20Exchange%20QP%20info%20via%20actor%20messages%0A%20%20%20%20-%20Manage%20the%20connection%20lifecycle%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20%23%20Monarch%20pattern%3A%20RDMA%20managers%20as%20actors%0A%20%20%20%20class%20RDMAManager(Actor)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.qp%20%3D%20create_queue_pair()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.qp_info%20%3D%20get_qp_info(self.qp)%20%20%23%20(qp_num%2C%20lid%2C%20gid)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_qp_info(self)%20-%3E%20QpInfo%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.qp_info%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20connect(self%2C%20remote_qp_info%3A%20QpInfo)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Transition%20QP%3A%20INIT%20%E2%86%92%20RTR%20%E2%86%92%20RTS%0A%20%20%20%20%20%20%20%20%20%20%20%20connect_qp(self.qp%2C%20remote_qp_info)%0A%0A%20%20%20%20%23%20Setup%3A%20exchange%20QP%20info%20via%20actor%20messages%2C%20then%20RDMA%20is%20ready%0A%20%20%20%20trainer_info%20%3D%20trainer_rdma.get_qp_info.call_one().get()%0A%20%20%20%20generator_rdma.connect.call_one(trainer_info).get()%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20actor%20abstraction%20makes%20RDMA%20connection%20management%20natural%20and%20composable.%0A%0A%20%20%20%20%23%23%23%20Completion%20Queues%20(CQs)%0A%0A%20%20%20%20How%20does%20the%20initiator%20know%20an%20operation%20finished%3F%20Every%20QP%20is%20associated%20with%20a%0A%20%20%20%20**Completion%20Queue%20(CQ)**.%20When%20the%20NIC%20finishes%20an%0A%20%20%20%20RDMA%20operation%2C%20it%20posts%20a%20**completion%20event**%20to%20the%20CQ%3A%0A%0A%20%20%20%20%60%60%60%0A%20%20%20%20App%3A%20post%20RDMA_READ%20to%20Send%20Queue%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%86%93%0A%20%20%20%20NIC%3A%20executes%20the%20read%20(bypasses%20remote%20CPU)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%86%93%0A%20%20%20%20NIC%3A%20posts%20completion%20to%20CQ%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%E2%86%93%0A%20%20%20%20App%3A%20poll%20CQ%20%E2%86%92%20%22done%2C%204096%20bytes%20transferred%22%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Monarch's%20Rust%20layer%20(%60RdmaManagerActor%60)%20handles%20CQ%20polling%20internally.%20When%20you%0A%20%20%20%20call%20%60.get()%60%20on%20an%20RDMA%20future%2C%20you're%20ultimately%20waiting%20for%20a%20CQ%20completion%20event.%0A%20%20%20%20This%20is%20why%20RDMA%20is%20%22one-sided%22%20but%20not%20%22zero-sided%22%20-%20the%20*initiator*%20still%20needs%0A%20%20%20%20to%20know%20when%20the%20transfer%20is%20done.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Monarch%20Using%20Monarch%3A%20RdmaController%0A%0A%20%20%20%20Here's%20the%20cool%20part%3A%20**Monarch%20uses%20itself**%20to%20manage%20RDMA%20infrastructure.%20Looking%20at%0A%20%20%20%20the%20actual%20Python%20code%20in%20%60monarch%2F_src%2Frdma%2Frdma.py%60%3A%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20%23%20From%20Monarch's%20RDMA%20implementation%0A%20%20%20%20from%20monarch._src.actor.proc_mesh%20import%20get_or_spawn_controller%0A%0A%20%20%20%20class%20RdmaController(Actor)%3A%0A%20%20%20%20%20%20%20%20'''Singleton%20controller%20that%20coordinates%20RDMA%20initialization.'''%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Track%20which%20proc%20meshes%20have%20RDMA%20initialized%0A%20%20%20%20%20%20%20%20%20%20%20%20self._manager_futures%3A%20dict%5BProcMesh%2C%20Future%5BRdmaManager%5D%5D%20%3D%20%7B%7D%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20async%20def%20init_rdma_on_mesh(self%2C%20proc_mesh%3A%20ProcMesh)%20-%3E%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20'''Lazily%20initialize%20RDMA%20on%20a%20proc%20mesh.'''%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20proc_mesh%20not%20in%20self._manager_futures%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._manager_futures%5Bproc_mesh%5D%20%3D%20Future(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20coro%3DRdmaManager.create(proc_mesh)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20await%20self._manager_futures%5Bproc_mesh%5D%0A%0A%20%20%20%20%23%20Cached%20initialization%20-%20only%20runs%20once%20per%20process%0A%20%20%20%20%40functools.cache%0A%20%20%20%20def%20_ensure_init_rdma_manager()%3A%0A%20%20%20%20%20%20%20%20async%20def%20task()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20controller%20%3D%20await%20get_or_spawn_controller(%22rdma_controller%22%2C%20RdmaController)%0A%20%20%20%20%20%20%20%20%20%20%20%20await%20controller.init_rdma_on_mesh.call_one(current_proc_mesh())%0A%20%20%20%20%20%20%20%20return%20spawn_task(task())%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20This%20is%20**Monarch%20building%20Monarch**%20-%20the%20RDMA%20subsystem%20uses%20the%20same%20patterns%3A%0A%0A%20%20%20%20-%20%60get_or_spawn_controller(%22rdma_controller%22%2C%20RdmaController)%60%20ensures%20one%20global%20controller%0A%20%20%20%20-%20The%20controller%20lazily%20initializes%20RDMA%20managers%20per%20proc%20mesh%0A%20%20%20%20-%20%60%40functools.cache%60%20ensures%20we%20only%20bootstrap%20once%20per%20process%0A%20%20%20%20-%20Under%20the%20hood%2C%20the%20actual%20RDMA%20operations%20are%20in%20Rust%20(%60RdmaManagerActor%60)%0A%0A%20%20%20%20It's%20actors%20all%20the%20way%20down.%0A%0A%20%20%20%20Now%20that%20we%20understand%20the%20primitives%20-%20QPs%2C%20MRs%2C%20CQs%2C%20and%20how%20Monarch%20wraps%20them%20-%0A%20%20%20%20the%20next%20question%20is%3A%20**what's%20the%20cost%20of%20getting%20them%20wrong%3F**%20Memory%20registration%0A%20%20%20%20turns%20out%20to%20be%20the%20biggest%20hidden%20tax.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%202.%20RDMA%20Buffer%20Patterns%0A%0A%20%20%20%20Memory%20registration%20is%20the%20hidden%20cost%20of%20RDMA.%20Before%20the%20NIC%20can%20read%20or%20write%20a%0A%20%20%20%20buffer%2C%20that%20buffer%20must%20be%20**registered**%20%E2%80%94%20the%20OS%20pins%20its%20physical%20pages%20and%20creates%0A%20%20%20%20DMA%20mappings%20so%20the%20NIC%20can%20access%20them%20directly.%20This%20can%20take%20milliseconds%20for%20large%0A%20%20%20%20buffers.%0A%0A%20%20%20%20The%20question%20isn't%20*whether*%20to%20pay%20this%20cost%20%E2%80%94%20it's%20*when%20and%20how%20often*.%20We'll%0A%20%20%20%20look%20at%20three%20patterns%3A%0A%0A%20%20%20%20%7C%20Pattern%20%7C%20Registration%20%7C%20Transfer%20%7C%20Trade-off%20%7C%0A%20%20%20%20%7C---------%7C-------------%7C----------%7C-----------%7C%0A%20%20%20%20%7C%20**Naive**%20%7C%20Every%20call%20%7C%20Per-buffer%20%7C%20Baseline%20%E2%80%94%20pays%20MR%20cost%20every%20step%20%7C%0A%20%20%20%20%7C%20**Contiguous**%20%7C%20Once%20at%20init%20%7C%20One%20bulk%20read%20%7C%20Fastest%2C%20but%20requires%20copying%20to%20a%20contiguous%20region%20%7C%0A%20%20%20%20%7C%20**Scattered%20%2B%20RDMAAction**%20%7C%20Once%20at%20init%20%7C%20Batched%20plan%20%7C%20Practical%20%E2%80%94%20works%20with%20non-contiguous%20layouts%20%7C%0A%0A%20%20%20%20The%20benchmark%20at%20the%20end%20shows%20the%20performance%20difference.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Pattern%201%3A%20Naive%0A%0A%20%20%20%20Create%20a%20new%20%60RDMABuffer%60%20on%20every%20transfer%20%E2%80%94%20pays%20memory%20registration%20cost%20each%20time%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Actor%2C%20RDMABuffer%2C%20current_rank%2C%20endpoint%2C%20time%2C%20torch)%3A%0A%20%20%20%20class%20NaiveSender(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Creates%20new%20RDMABuffer%20handles%20every%20transfer.%20Expensive!%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20layer_sizes%3A%20list)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layer_sizes%20%3D%20layer_sizes%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layers%20%3D%20%5Btorch.zeros(size%2C%20dtype%3Dtorch.float32)%20for%20size%20in%20layer_sizes%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20layer%20in%20enumerate(self.layers)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layer.fill_(float(i%20%2B%201))%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_fresh_handles(self)%20-%3E%20list%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20handles%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20size%2C%20layer%20in%20zip(self.layer_sizes%2C%20self.layers)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20layer.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20handles.append((size%2C%20RDMABuffer(byte_view)))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20handles%0A%0A%0A%20%20%20%20class%20NaiveReceiver(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Receives%20from%20naive%20sender%20-%20pays%20MR%20cost%20every%20step.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20layer_sizes%3A%20list)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layer_sizes%20%3D%20layer_sizes%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layers%20%3D%20%5Btorch.zeros(size%2C%20dtype%3Dtorch.float32)%20for%20size%20in%20layer_sizes%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rank%20%3D%20current_rank().rank%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20warmup(self%2C%20sender%3A%20NaiveSender)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Force%20RDMA%20subsystem%20init%20%2B%20QP%20establishment%20(not%20timed).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20handles%20%3D%20sender.get_fresh_handles.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20(size%2C%20handle)%20in%20enumerate(handles)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.layers%5Bi%5D.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20handle.read_into(byte_view).get()%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20receive_step(self%2C%20sender%3A%20NaiveSender)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20handles%20%3D%20sender.get_fresh_handles.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20(size%2C%20handle)%20in%20enumerate(handles)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.layers%5Bi%5D.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20handle.read_into(byte_view).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20elapsed_ms%20%3D%20(time.perf_counter()%20-%20start)%20*%201000%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22elapsed_ms%22%3A%20elapsed_ms%7D%0A%0A%20%20%20%20print(%22NaiveSender%3A%20Re-registers%20all%20parameters%20on%20every%20call%22)%0A%20%20%20%20return%20NaiveReceiver%2C%20NaiveSender%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Pattern%202%3A%20Contiguous%20Buffer%0A%0A%20%20%20%20Pack%20all%20layers%20into%20one%20buffer%2C%20register%20once%20%E2%80%94%20one%20MR%2C%20one%20transfer%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Actor%2C%20RDMABuffer%2C%20current_rank%2C%20endpoint%2C%20time%2C%20torch)%3A%0A%20%20%20%20class%20ContiguousSender(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22One%20buffer%2C%20one%20MR%2C%20registered%20at%20startup.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20layer_sizes%3A%20list)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layer_sizes%20%3D%20layer_sizes%0A%20%20%20%20%20%20%20%20%20%20%20%20total_size%20%3D%20sum(layer_sizes)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20One%20contiguous%20buffer%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer%20%3D%20torch.zeros(total_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20offset%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20size%20in%20enumerate(layer_sizes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.buffer%5Boffset%20%3A%20offset%20%2B%20size%5D.fill_(float(i%20%2B%201))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20offset%20%2B%3D%20size%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Register%20ONCE%20at%20startup%0A%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.buffer.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.handle%20%3D%20RDMABuffer(byte_view)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_handle(self)%20-%3E%20tuple%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20(len(self.buffer)%2C%20self.handle)%20%20%23%20Same%20handle%20every%20time!%0A%0A%0A%20%20%20%20class%20ContiguousReceiver(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Receives%20from%20contiguous%20sender%20-%20one%20big%20read.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20total_size%3A%20int)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer%20%3D%20torch.zeros(total_size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rank%20%3D%20current_rank().rank%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20warmup(self%2C%20sender%3A%20ContiguousSender)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Force%20QP%20establishment%20(not%20timed).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20size%2C%20handle%20%3D%20sender.get_handle.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.buffer.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20handle.read_into(byte_view).get()%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20receive_step(self%2C%20sender%3A%20ContiguousSender)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20size%2C%20handle%20%3D%20sender.get_handle.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.buffer.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20handle.read_into(byte_view).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20elapsed_ms%20%3D%20(time.perf_counter()%20-%20start)%20*%201000%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22elapsed_ms%22%3A%20elapsed_ms%7D%0A%0A%20%20%20%20print(%22ContiguousSender%3A%20Registers%20once%2C%20reuses%20same%20handle%22)%0A%20%20%20%20return%20ContiguousReceiver%2C%20ContiguousSender%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Pattern%203%3A%20Scattered%20%2B%20RDMAAction%0A%0A%20%20%20%20Register%20each%20buffer%20at%20init%2C%20build%20a%20transfer%20plan%20once%20via%20handshake%3A%0A%0A%20%20%20%20**What%20is%20RDMAAction%3F**%0A%0A%20%20%20%20Think%20of%20%60RDMAAction%60%20as%20a%20**transfer%20plan**.%20You%20describe%20all%20the%20reads%2Fwrites%20you%20want%0A%20%20%20%20to%20do%2C%20then%20%60submit()%60%20executes%20the%20whole%20plan%20at%20once%3A%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20%23%20Build%20the%20plan%20once%0A%20%20%20%20action%20%3D%20RDMAAction()%0A%20%20%20%20action.read_into(handle1%2C%20local_buffer1)%0A%20%20%20%20action.read_into(handle2%2C%20local_buffer2)%0A%20%20%20%20action.read_into(handle3%2C%20local_buffer3)%0A%0A%20%20%20%20%23%20Execute%20whenever%20you%20want%20-%20just%20one%20call%0A%20%20%20%20action.submit().get()%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20This%20is%20useful%20when%20you%20have%20many%20scattered%20buffers%20(like%20model%20parameters)%20and%20want%0A%20%20%20%20to%20batch%20them%20into%20a%20single%20logical%20operation.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Actor%2C%20RDMAAction%2C%20RDMABuffer%2C%20current_rank%2C%20endpoint%2C%20time%2C%20torch)%3A%0A%20%20%20%20class%20ScatteredSender(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Multiple%20buffers%2C%20each%20registered%20once%20at%20startup.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20layer_sizes%3A%20list)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layer_sizes%20%3D%20layer_sizes%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layers%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20self.handles%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20size%20in%20enumerate(layer_sizes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layer%20%3D%20torch.zeros(size%2C%20dtype%3Dtorch.float32)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layer.fill_(float(i%20%2B%201))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.layers.append(layer)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Register%20ONCE%20at%20startup%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20layer.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.handles.append(RDMABuffer(byte_view))%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_handles(self)%20-%3E%20list%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%5B(size%2C%20handle)%20for%20size%2C%20handle%20in%20zip(self.layer_sizes%2C%20self.handles)%5D%0A%0A%20%20%20%20class%20ScatteredReceiver(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Receives%20from%20scattered%20sender%20with%20RDMAAction%20batching.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20layer_sizes%3A%20list)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layer_sizes%20%3D%20layer_sizes%0A%20%20%20%20%20%20%20%20%20%20%20%20self.layers%20%3D%20%5Btorch.zeros(size%2C%20dtype%3Dtorch.float32)%20for%20size%20in%20layer_sizes%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20self.rank%20%3D%20current_rank().rank%0A%20%20%20%20%20%20%20%20%20%20%20%20self.action%20%3D%20None%20%20%23%20Built%20on%20handshake%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20handshake(self%2C%20sender%3A%20ScatteredSender)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Call%20once%20to%20build%20the%20RDMAAction%20transfer%20plan.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20handles%20%3D%20sender.get_handles.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.action%20%3D%20RDMAAction()%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20(size%2C%20handle)%20in%20enumerate(handles)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self.layers%5Bi%5D.view(torch.uint8).flatten()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.action.read_into(handle%2C%20byte_view)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%22Transfer%20plan%20ready%22%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20receive_step(self)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Execute%20the%20cached%20transfer%20plan.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.action.submit().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20elapsed_ms%20%3D%20(time.perf_counter()%20-%20start)%20*%201000%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22elapsed_ms%22%3A%20elapsed_ms%7D%0A%0A%20%20%20%20print(%22ScatteredSender%3A%20Registers%20each%20layer%20once%22)%0A%20%20%20%20print(%22ScatteredReceiver%3A%20handshake()%20builds%20plan%2C%20receive_step()%20executes%20it%22)%0A%20%20%20%20return%20ScatteredReceiver%2C%20ScatteredSender%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Running%20the%20Benchmark%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20ContiguousReceiver%2C%0A%20%20%20%20ContiguousSender%2C%0A%20%20%20%20NaiveReceiver%2C%0A%20%20%20%20NaiveSender%2C%0A%20%20%20%20ScatteredReceiver%2C%0A%20%20%20%20ScatteredSender%2C%0A%20%20%20%20is_rdma_available%2C%0A%20%20%20%20this_host%2C%0A)%3A%0A%20%20%20%20def%20run_benchmark()%3A%0A%20%20%20%20%20%20%20%20%22%22%22Compare%20the%20three%20approaches%20over%20multiple%20steps.%22%22%22%0A%20%20%20%20%20%20%20%20if%20not%20is_rdma_available()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(%22%E2%9A%A0%20RDMA%20not%20available%20on%20this%20machine.%20Skipping%20benchmark.%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20print(%22%20%20Run%20on%20an%20RDMA-capable%20host%20to%20see%20real%20results.%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20None%0A%0A%20%20%20%20%20%20%20%20layer_sizes%20%3D%20%5B10000%2C%2050000%2C%2020000%5D%20%20%23%2080000%20floats%20total%0A%20%20%20%20%20%20%20%20total_size%20%3D%20sum(layer_sizes)%0A%20%20%20%20%20%20%20%20num_steps%20%3D%205%0A%0A%20%20%20%20%20%20%20%20host%20%3D%20this_host()%0A%20%20%20%20%20%20%20%20sender_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%201%7D)%0A%20%20%20%20%20%20%20%20receiver_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%201%7D)%0A%0A%20%20%20%20%20%20%20%20%23%20Spawn%20all%20actors%0A%20%20%20%20%20%20%20%20naive_sender%20%3D%20sender_procs.spawn(%22naive_s%22%2C%20NaiveSender%2C%20layer_sizes)%0A%20%20%20%20%20%20%20%20naive_receiver%20%3D%20receiver_procs.spawn(%22naive_r%22%2C%20NaiveReceiver%2C%20layer_sizes)%0A%0A%20%20%20%20%20%20%20%20cont_sender%20%3D%20sender_procs.spawn(%22cont_s%22%2C%20ContiguousSender%2C%20layer_sizes)%0A%20%20%20%20%20%20%20%20cont_receiver%20%3D%20receiver_procs.spawn(%22cont_r%22%2C%20ContiguousReceiver%2C%20total_size)%0A%0A%20%20%20%20%20%20%20%20scat_sender%20%3D%20sender_procs.spawn(%22scat_s%22%2C%20ScatteredSender%2C%20layer_sizes)%0A%20%20%20%20%20%20%20%20scat_receiver%20%3D%20receiver_procs.spawn(%22scat_r%22%2C%20ScatteredReceiver%2C%20layer_sizes)%0A%0A%20%20%20%20%20%20%20%20%23%20%E2%94%80%E2%94%80%20WARMUP%20%E2%94%80%E2%94%80%0A%20%20%20%20%20%20%20%20%23%20Force%20the%20full%20RDMA%20initialization%20chain%20(RdmaController%20spawn%2C%0A%20%20%20%20%20%20%20%20%23%20RdmaManagerActor%20spawn%2C%20QP%20establishment)%20so%20benchmarks%20measure%0A%20%20%20%20%20%20%20%20%23%20only%20what%20we%20intend%20to%20measure.%0A%20%20%20%20%20%20%20%20print(%22Warming%20up%20RDMA%20subsystem...%22)%0A%20%20%20%20%20%20%20%20naive_receiver.warmup.call_one(naive_sender).get()%0A%20%20%20%20%20%20%20%20cont_receiver.warmup.call_one(cont_sender).get()%0A%20%20%20%20%20%20%20%20scat_receiver.handshake.call_one(scat_sender).get()%0A%20%20%20%20%20%20%20%20scat_receiver.receive_step.call_one().get()%20%20%23%20warm%20the%20transfer%20path%0A%20%20%20%20%20%20%20%20print(%22Warmup%20complete.%5Cn%22)%0A%0A%20%20%20%20%20%20%20%20print(%22%3D%3D%3D%20RDMA%20Buffer%20Pattern%20Benchmark%20%3D%3D%3D%22)%0A%20%20%20%20%20%20%20%20print(f%22Transferring%20%7Btotal_size%7D%20floats%20(%7Btotal_size%20*%204%20%2F%201024%3A.1f%7D%20KB)%20x%20%7Bnum_steps%7D%20steps%5Cn%22)%0A%0A%20%20%20%20%20%20%20%20results%20%3D%20%7B%7D%0A%0A%20%20%20%20%20%20%20%20%23%20Pattern%201%3A%20Naive%20(re-register%20each%20step)%0A%20%20%20%20%20%20%20%20times%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20step%20in%20range(num_steps)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20r%20%3D%20naive_receiver.receive_step.call_one(naive_sender).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20times.append(r%5B%22elapsed_ms%22%5D)%0A%20%20%20%20%20%20%20%20results%5B%22Naive%22%5D%20%3D%20times%0A%20%20%20%20%20%20%20%20print(f%22Pattern%201%20%E2%80%94%20Naive%20(re-register%20each%20step)%3A%22)%0A%20%20%20%20%20%20%20%20for%20i%2C%20t%20in%20enumerate(times)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%20%20Step%20%7Bi%2B1%7D%3A%20%7Bt%3A.2f%7Dms%22)%0A%20%20%20%20%20%20%20%20print(f%22%20%20Average%3A%20%7Bsum(times)%2Flen(times)%3A.2f%7Dms%5Cn%22)%0A%0A%20%20%20%20%20%20%20%20%23%20Pattern%202%3A%20Contiguous%20(one%20buffer%2C%20one%20read)%0A%20%20%20%20%20%20%20%20times%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20step%20in%20range(num_steps)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20r%20%3D%20cont_receiver.receive_step.call_one(cont_sender).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20times.append(r%5B%22elapsed_ms%22%5D)%0A%20%20%20%20%20%20%20%20results%5B%22Contiguous%22%5D%20%3D%20times%0A%20%20%20%20%20%20%20%20print(f%22Pattern%202%20%E2%80%94%20Contiguous%20(one%20buffer%2C%20one%20read)%3A%22)%0A%20%20%20%20%20%20%20%20for%20i%2C%20t%20in%20enumerate(times)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%20%20Step%20%7Bi%2B1%7D%3A%20%7Bt%3A.2f%7Dms%22)%0A%20%20%20%20%20%20%20%20print(f%22%20%20Average%3A%20%7Bsum(times)%2Flen(times)%3A.2f%7Dms%5Cn%22)%0A%0A%20%20%20%20%20%20%20%20%23%20Pattern%203%3A%20Scattered%20%2B%20RDMAAction%20(batched%20plan)%0A%20%20%20%20%20%20%20%20scat_receiver.handshake.call_one(scat_sender).get()%20%20%23%20rebuild%20plan%0A%20%20%20%20%20%20%20%20times%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20step%20in%20range(num_steps)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20r%20%3D%20scat_receiver.receive_step.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20times.append(r%5B%22elapsed_ms%22%5D)%0A%20%20%20%20%20%20%20%20results%5B%22Scattered%22%5D%20%3D%20times%0A%20%20%20%20%20%20%20%20print(f%22Pattern%203%20%E2%80%94%20Scattered%20%2B%20RDMAAction%20(batched%20plan)%3A%22)%0A%20%20%20%20%20%20%20%20for%20i%2C%20t%20in%20enumerate(times)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%20%20Step%20%7Bi%2B1%7D%3A%20%7Bt%3A.2f%7Dms%22)%0A%20%20%20%20%20%20%20%20print(f%22%20%20Average%3A%20%7Bsum(times)%2Flen(times)%3A.2f%7Dms%5Cn%22)%0A%0A%20%20%20%20%20%20%20%20print(%22%3D%3D%3D%20What's%20Happening%20%3D%3D%3D%22)%0A%20%20%20%20%20%20%20%20print(%22Naive%3A%20Re-registers%20MRs%20%2B%20sequential%20per-layer%20reads%20(~9ms)%22)%0A%20%20%20%20%20%20%20%20print(%22Contiguous%3A%20One%20MR%2C%20one%20read%20%E2%80%94%20fewest%20round%20trips%20(~3ms)%22)%0A%20%20%20%20%20%20%20%20print(%22Scattered%20%2B%20RDMAAction%3A%20Pre-built%20transfer%20plan%20(~8ms%2C%20Python%20overhead)%22)%0A%0A%20%20%20%20%20%20%20%20return%20results%0A%0A%20%20%20%20benchmark_results%20%3D%20run_benchmark()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20**What%20the%20benchmark%20shows%3A**%0A%0A%20%20%20%20-%20**Naive**%20(~9ms)%3A%20Re-registers%20MRs%20each%20step%2C%20plus%20sequential%20per-layer%20reads.%0A%20%20%20%20%20%20At%20larger%20buffer%20sizes%20the%20MR%20cost%20grows%2C%20but%20here%20it's%20masked%20by%20Python%20overhead.%0A%20%20%20%20-%20**Contiguous**%20(~3ms)%3A%20One%20MR%2C%20one%20%60read_into()%60%20%E2%80%94%20fewest%20round%20trips%20wins.%0A%20%20%20%20%20%20The%20trade-off%3A%20in%20practice%2C%20model%20parameters%20aren't%20contiguous%2C%20so%20you'd%20need%0A%20%20%20%20%20%20to%20copy%20into%20a%20contiguous%20staging%20buffer%20first.%0A%20%20%20%20-%20**Scattered%20%2B%20RDMAAction**%20(~8ms)%3A%20Pre-built%20transfer%20plan%2C%20works%20with%20non-contiguous%0A%20%20%20%20%20%20layouts.%20Currently%20has%20Python%20overhead%20from%20the%20batching%20logic%3B%20lowering%20%60RDMAAction%60%0A%20%20%20%20%20%20into%20Rust%20will%20close%20the%20gap%20with%20Contiguous.%0A%0A%20%20%20%20The%20key%20lesson%3A%20**minimize%20round%20trips.**%20Whether%20that%20means%20packing%20into%20one%20buffer%0A%20%20%20%20or%20batching%20via%20%60RDMAAction%60%2C%20the%20goal%20is%20the%20same%20%E2%80%94%20fewer%20Python-to-Rust-to-NIC%0A%20%20%20%20transitions%20per%20sync.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Summary%0A%0A%20%20%20%20We%20started%20with%20a%20question%20%E2%80%94%20*where%20do%20the%20milliseconds%20go%3F*%20%E2%80%94%20and%20traced%20the%20answer%0A%20%20%20%20through%20two%20layers%3A%0A%0A%20%20%20%201.%20**ibverbs%20internals**%20%E2%80%94%20QPs%2C%20MRs%2C%20CQs%3A%20the%20primitives%20that%20make%20RDMA%20work%2C%20and%20how%0A%20%20%20%20%20%20%20Monarch's%20actor%20model%20wraps%20them%20into%20composable%20connections%0A%20%20%20%202.%20**RDMA%20buffer%20patterns**%20%E2%80%94%20From%20naive%20re-registration%20to%20contiguous%20buffers%0A%20%20%20%20%20%20%20and%20batched%20transfer%20plans%20via%20%60RDMAAction%60%20%E2%80%94%20minimizing%20Python%20round%20trips%0A%0A%20%20%20%20We%20focused%20on%20**ibverbs**%20specifically%20because%20that's%20what%20Monarch's%20RDMA%20subsystem%0A%20%20%20%20supports%20today%20(InfiniBand%20and%20RoCE%20via%20Mellanox%2FNVIDIA%20ConnectX%20NICs).%20EFA%20(Elastic%0A%20%20%20%20Fabric%20Adapter%2C%20used%20on%20AWS)%20is%20a%20relevant%20transport%20but%20not%20yet%20supported%20%E2%80%94%20it's%0A%20%20%20%20actively%20being%20worked%20on.%0A%0A%20%20%20%20The%20main%20notebook%20(06)%20covers%20the%20high-level%20patterns%20for%20async%20RL%20weight%20sync%20%E2%80%94%0A%20%20%20%20magic%20pointers%2C%20CPU%20staging%2C%20circular%20buffers%2C%20and%20DTensor%20re-sharding.%20This%20notebook%0A%20%20%20%20showed%20the%20RDMA%20specifics%20underneath%3A%20what%20the%20NIC%20is%20actually%20doing%20and%20how%20to%20avoid%0A%20%20%20%20paying%20unnecessary%20costs%20on%20the%20hot%20path.%0A%0A%20%20%20%20---%0A%0A%20%20%20%20**Previous%3A**%20%5BNB07%20%E2%80%94%20RDMA%20Weight%20Sync%5D(.%2F07_rdma_weight_sync.html)%20%C2%B7%20**Next%3A**%20%5BNB08%20%E2%80%94%20Async%20RL%20E2E%5D(.%2F08_rl_e2e.html)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A