%23!%2Fusr%2Fbin%2Fenv%20python3%0A%22%22%22%0ANotebook%2008%3A%20Closing%20the%20Loop%20-%20Async%20RL%20Training%0A%0AIn%20Notebook%2005%2C%20we%20saw%20that%20Qwen%200.5B%20struggles%20with%20compositional%20Zorplex%20tasks.%0ANow%20we%20close%20the%20loop%3A%20train%20the%20model%20using%20async%20RL%20with%20Monarch.%0A%0AThis%20notebook%20demonstrates%3A%0A-%20Building%20RL%20actors%20(Trainer%2C%20Generator%2C%20ReplayBuffer)%0A-%20Running%20concurrent%20generation%20and%20training%0A-%20Weight%20synchronization%20between%20actors%20(circular%20buffer%20%2B%20CPU%20staging)%0A-%20Real%20training%20metrics%20and%20before%2Fafter%20evaluation%0A%0ARun%20with%3A%20uv%20run%20marimo%20edit%20notebooks%2F08_rl_e2e.py%0A%22%22%22%0A%0A%23%20CRITICAL%3A%20Set%20allocator%20config%20BEFORE%20any%20PyTorch%20imports%20(including%20in%20subprocesses)%0A%23%20Set%20both%20names%20for%20compatibility%20(old%20and%20new%20PyTorch%20versions)%0A%0Aimport%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%20Environment%20setup%20for%20Monarch%20subprocesses%0A%20%20%20%20import%20os%0A%20%20%20%20os.environ%5B%22TOKENIZERS_PARALLELISM%22%5D%20%3D%20%22false%22%0A%20%20%20%20os.environ%5B%22HF_HUB_OFFLINE%22%5D%20%3D%20%221%22%0A%20%20%20%20os.environ%5B%22TRANSFORMERS_OFFLINE%22%5D%20%3D%20%221%22%0A%20%20%20%20%23%20Note%3A%20CUDA_VISIBLE_DEVICES%20is%20set%20per-actor%20in%20setup()%0A%20%20%20%20%23%20Note%3A%20PYTORCH_ALLOC_CONF%20is%20set%20at%20module%20level%20for%20RDMA%0A%0A%20%20%20%20import%20sys%0A%20%20%20%20_src_dir%20%3D%20os.path.abspath(os.path.join(os.path.dirname(__file__)%20if%20%22__file__%22%20in%20dir()%20else%20os.getcwd()%2C%20%22..%22%2C%20%22src%22))%0A%20%20%20%20if%20_src_dir%20not%20in%20sys.path%3A%0A%20%20%20%20%20%20%20%20sys.path.insert(0%2C%20_src_dir)%0A%0A%20%20%20%20%23%20Set%20PYTHONPATH%20for%20Monarch%20subprocesses%0A%20%20%20%20_existing%20%3D%20os.environ.get(%22PYTHONPATH%22%2C%20%22%22)%0A%20%20%20%20os.environ%5B%22PYTHONPATH%22%5D%20%3D%20f%22%7B_src_dir%7D%3A%7B_existing%7D%22%20if%20_existing%20else%20_src_dir%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%20Closing%20the%20Loop%3A%20Async%20RL%20Training%0A%0A%20%20%20%20In%20**Notebook%2005**%2C%20we%20introduced%20the%20Zorplex%20benchmark%20and%20identified%20three%0A%20%20%20%20failure%20modes%3A%20wrong%20format%20(no%20%60%5BANSWER%5D%60%20tag)%2C%20tool%20spam%2C%20and%20wrong%20answers.%0A%20%20%20%20The%20model%20often%20gets%20the%20right%20value%20but%20fails%20to%20emit%20it%20correctly.%0A%0A%20%20%20%20**Now%20we%20close%20the%20loop**%3A%20train%20the%20model%20to%20get%20better%20at%20these%20tasks%2C%20and%0A%20%20%20%20track%20which%20failure%20modes%20improve%20during%20training.%0A%0A%20%20%20%20We'll%20build%20on%20patterns%20from%20across%20the%20series.%20The%20architecture%20we're%20building%3A%20multiple%20generators%0A%20%20%20%20feed%20trajectories%20into%20a%20replay%20buffer%20while%20a%20trainer%20continuously%20samples%20and%20updates%20the%20policy.%0A%0A%20%20%20%20We'll%20measure%20*before*%20and%20*after*%20accuracy%20--%20and%20failure%20mode%20breakdown%20--%20to%0A%20%20%20%20see%20if%20training%20actually%20helps.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20collections%20import%20deque%0A%20%20%20%20import%20random%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn.functional%20as%20F%0A%0A%20%20%20%20%23%20Model%20imports%0A%20%20%20%20from%20transformers%20import%20AutoModelForCausalLM%2C%20AutoTokenizer%0A%0A%20%20%20%20%23%20Zorplex%20imports%0A%20%20%20%20from%20zorplex_rl%20import%20get_spec%2C%20Task%0A%20%20%20%20from%20zorplex_rl.evaluate%20import%20generate_with_tools%0A%0A%20%20%20%20%23%20RL%20primitives%20(shared%20dataclasses)%0A%20%20%20%20from%20rl_primitives%20import%20Trajectory%2C%20TrainMetrics%0A%0A%20%20%20%20%23%20RDMA%20imports%20(with%20fallback)%0A%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20from%20monarch.rdma%20import%20RDMABuffer%2C%20is_rdma_available%0A%20%20%20%20%20%20%20%20_rdma_available%20%3D%20is_rdma_available()%0A%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20RDMABuffer%20%3D%20None%0A%20%20%20%20%20%20%20%20_rdma_available%20%3D%20False%0A%0A%20%20%20%20def%20rdma_available()%3A%0A%20%20%20%20%20%20%20%20return%20_rdma_available%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20AutoModelForCausalLM%2C%0A%20%20%20%20%20%20%20%20AutoTokenizer%2C%0A%20%20%20%20%20%20%20%20F%2C%0A%20%20%20%20%20%20%20%20RDMABuffer%2C%0A%20%20%20%20%20%20%20%20Task%2C%0A%20%20%20%20%20%20%20%20TrainMetrics%2C%0A%20%20%20%20%20%20%20%20Trajectory%2C%0A%20%20%20%20%20%20%20%20deque%2C%0A%20%20%20%20%20%20%20%20generate_with_tools%2C%0A%20%20%20%20%20%20%20%20get_spec%2C%0A%20%20%20%20%20%20%20%20random%2C%0A%20%20%20%20%20%20%20%20rdma_available%2C%0A%20%20%20%20%20%20%20%20torch%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20monarch.actor%20import%20Actor%2C%20endpoint%2C%20current_rank%0A%0A%20%20%20%20return%20Actor%2C%20current_rank%2C%20endpoint%0A%0A%0A%40app.cell%0Adef%20_(TrainMetrics%2C%20Trajectory%2C%20mo)%3A%0A%20%20%20%20import%20dataclasses%20as%20_dc%0A%20%20%20%20_traj_fields%20%3D%20%5B(f.name%2C%20f.type.__name__%20if%20hasattr(f.type%2C%20'__name__')%20else%20str(f.type))%20for%20f%20in%20_dc.fields(Trajectory)%5D%0A%20%20%20%20_metrics_fields%20%3D%20%5B(f.name%2C%20f.type.__name__%20if%20hasattr(f.type%2C%20'__name__')%20else%20str(f.type))%20for%20f%20in%20_dc.fields(TrainMetrics)%5D%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%20Shared%20Data%20Structures%0A%0A%20%20%20%20For%20clarity%2C%20we%20are%20using%20the%20following%20data%20structures%20in%20this%20notebook%3A%0A%0A%20%20%20%20**Trajectory**%20--%20one%20rollout%20from%20a%20generator%3A%0A%0A%20%20%20%20%7C%20Field%20%7C%20Type%20%7C%0A%20%20%20%20%7C-------%7C------%7C%0A%20%20%20%20%7B%22%22.join(f%22%7C%20%60%7Bn%7D%60%20%7C%20%60%7Bt%7D%60%20%7C%7Bchr(10)%7D%22%20for%20n%2C%20t%20in%20_traj_fields)%7D%0A%0A%20%20%20%20**TrainMetrics**%20--%20returned%20after%20each%20training%20step%3A%0A%0A%20%20%20%20%7C%20Field%20%7C%20Type%20%7C%0A%20%20%20%20%7C-------%7C------%7C%0A%20%20%20%20%7B%22%22.join(f%22%7C%20%60%7Bn%7D%60%20%7C%20%60%7Bt%7D%60%20%7C%7Bchr(10)%7D%22%20for%20n%2C%20t%20in%20_metrics_fields)%7D%0A%0A%20%20%20%20Key%20fields%3A%0A%20%20%20%20-%20%60model_only_text%60%20stores%20the%20model's%20generated%20tokens%20without%20injected%20tool%0A%20%20%20%20%20%20results%2C%20so%20the%20trainer%20can%20compute%20log-probabilities%20on%20exactly%20what%20the%20model%20produced.%0A%20%20%20%20-%20%60has_answer_tag%60%20tracks%20whether%20the%20model%20emitted%20%60%5BANSWER%5D%60%20--%20this%20is%20the%0A%20%20%20%20%20%20format%20compliance%20signal%20from%20%5BNB05%5D(.%2F05_rl_intro.html)'s%20failure%20mode%20analysis.%0A%20%20%20%20-%20%60failure_mode%60%20classifies%20each%20trajectory%20as%20%60%22success%22%60%2C%20%60%22wrong_format%22%60%2C%0A%20%20%20%20%20%20%60%22tool_spam%22%60%2C%20or%20%60%22wrong_answer%22%60%20so%20we%20can%20track%20which%20failure%20modes%20improve%0A%20%20%20%20%20%20during%20training.%0A%20%20%20%20-%20%60correct_rate%60%20and%20%60format_rate%60%20on%20%60TrainMetrics%60%20let%20us%20track%20these%20signals%0A%20%20%20%20%20%20per%20training%20step.%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%20Service%20Infrastructure%0A%0A%20%20%20%20We%20import%20a%20**custom%20Service%20abstraction**%20from%20%60monarch_utils%60%20that%20manages%20worker%0A%20%20%20%20replicas%20with%20health%20tracking%20and%20round-robin%20routing.%20This%20is%20a%20utility%20we%20built%0A%20%20%20%20for%20this%20notebook%20series%20--%20the%20canonical%20Monarch%20pattern%20uses%20direct%20actor%0A%20%20%20%20references%20and%20slicing%2C%20which%20is%20what%20the%20Service%20wraps%20internally.%0A%0A%20%20%20%20(See%20notebook%2005%20for%20the%20full%20implementation.)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20monarch_utils.services%20import%20Service%2C%20register_service%0A%0A%20%20%20%20return%20Service%2C%20register_service%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%20The%20%60setup()%60%20Pattern%0A%0A%20%20%20%20Actors%20in%20this%20notebook%20use%20a%20two-phase%20initialization%3A%0A%0A%20%20%20%201.%20**%60__init__%60**%20runs%20during%20%60spawn()%60%20--%20keep%20it%20lightweight%20(store%20config%2C%20set%20rank)%0A%20%20%20%202.%20**%60setup()%60**%20is%20an%20endpoint%20called%20explicitly%20after%20spawn%20--%20do%20heavy%20work%20here%0A%20%20%20%20%20%20%20(load%20models%2C%20allocate%20GPU%20memory%2C%20register%20RDMA%20buffers)%0A%0A%20%20%20%20Why%20not%20do%20everything%20in%20%60__init__%60%3F%20Two%20reasons%3A%0A%0A%20%20%20%20-%20**%60spawn()%60%20is%20asynchronous**%20--%20it%20returns%20immediately%2C%20and%20%60__init__%60%20runs%20in%0A%20%20%20%20%20%20the%20remote%20process%20before%20the%20first%20endpoint%20call.%20But%20you%20don't%20control%20*when*%2C%0A%20%20%20%20%20%20and%20you%20can't%20confirm%20it%20completed.%20An%20explicit%20%60setup()%60%20call%20lets%20you%20sequence%0A%20%20%20%20%20%20initialization%20(e.g.%2C%20set%20%60CUDA_VISIBLE_DEVICES%60%20and%20confirm%20it%20took%20effect%20before%0A%20%20%20%20%20%20loading%20a%20model).%0A%20%20%20%20-%20**Coordination**%20--%20you%20often%20need%20to%20initialize%20actors%20in%20a%20specific%20order%20(set%20up%0A%20%20%20%20%20%20the%20trainer%20before%20generators%20try%20to%20sync%20weights).%20Endpoint%20calls%20give%20you%20that%0A%20%20%20%20%20%20sequencing%3B%20%60__init__%60%20doesn't.%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%20Actor%201%3A%20ZorplexWorker%0A%0A%20%20%20%20Tool%20execution%20environments%20(docker%20containers%2C%20sandboxes%2C%20API%20endpoints)%20naturally%0A%20%20%20%20form%20a%20fleet%20--%20you%20want%20many%20instances%20running%20in%20parallel%20to%20keep%20up%20with%0A%20%20%20%20generation%20throughput.%20That%20makes%20them%20a%20good%20fit%20for%20a%20**Service**%20(from%20%5BNB06%5D(.%2F06_services.html))%0A%20%20%20%20with%20health%20tracking%20and%20round-robin%20routing.%0A%0A%20%20%20%20Our%20ZorplexWorker%20actors%20handle%20Zorplex%20tasks%3A%0A%20%20%20%20-%20%60generate_task()%60%20--%20creates%20a%20new%20problem%0A%20%20%20%20-%20%60execute_tool()%60%20--%20handles%20LOOKUP%20calls%0A%20%20%20%20-%20%60check_answer()%60%20--%20verifies%20correctness%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Actor%2C%20current_rank%2C%20endpoint%2C%20get_spec)%3A%0A%20%20%20%20class%20ZorplexWorker(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Worker%20actor%20that%20handles%20Zorplex%20tool%20execution.%0A%0A%20%20%20%20%20%20%20%20Managed%20by%20a%20Service%20for%20load%20balancing%20across%20replicas.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20difficulty%3A%20str%20%3D%20%22easy%22%2C%20seed%3A%20int%20%3D%2042)%3A%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.spec%20%3D%20get_spec(%22compositional%22%2C%20difficulty%3Ddifficulty%2C%20seed%3Dseed%20%2B%20self.rank)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.calls_served%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BZorplexWorker%3A%7Bself.rank%7D%5D%20Initialized%20with%20difficulty%3D%7Bdifficulty%7D%22)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20ping(self)%20-%3E%20bool%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20True%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20generate_task(self)%20-%3E%20tuple%5Bstr%2C%20int%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Generate%20a%20new%20task.%20Returns%20(question%2C%20correct_answer).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20task%20%3D%20self.spec.generate_task()%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20task.question%2C%20task.correct_answer%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20execute_tool(self%2C%20tool_name%3A%20str%2C%20argument%3A%20str)%20-%3E%20str%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Execute%20a%20tool%20call.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20from%20zorplex_rl.task_specs%20import%20ToolCall%0A%20%20%20%20%20%20%20%20%20%20%20%20tc%20%3D%20ToolCall(tool_name%2C%20argument)%0A%20%20%20%20%20%20%20%20%20%20%20%20result%20%3D%20self.spec.execute_tool(tc)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.calls_served%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20str(result)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_system_prompt(self)%20-%3E%20str%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Get%20the%20system%20prompt%20with%20tool%20hints.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.spec.get_system_prompt(with_hint%3DTrue)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20check_answer(self%2C%20model_output%3A%20str%2C%20correct_answer%3A%20int)%20-%3E%20tuple%5Bbool%2C%20int%20%7C%20None%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Check%20if%20model%20output%20contains%20the%20correct%20answer.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20extracted%20%3D%20self.spec.extract_answer(model_output%2C%20%5B%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20is_correct%20%3D%20extracted%20%3D%3D%20correct_answer%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20is_correct%2C%20extracted%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20stats(self)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22rank%22%3A%20self.rank%2C%20%22calls_served%22%3A%20self.calls_served%7D%0A%0A%20%20%20%20return%20(ZorplexWorker%2C)%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%20Actor%202%3A%20ReplayBuffer%0A%0A%20%20%20%20A%20simple%20actor%20that%20stores%20trajectories.%20Generators%20push%20trajectories%20in%2C%0A%20%20%20%20the%20trainer%20samples%20batches%20out.%0A%0A%20%20%20%20Recall%20our%20intro%20to%20async%20RL%20in%20notebook%204%20--%20the%20replay%20buffer%20is%20the%20decoupling%0A%20%20%20%20point%20that%20enables%20asynchronous%20execution.%20Generators%20push%2C%20trainer%20pulls%2C%20neither%0A%20%20%20%20waits%20for%20the%20other.%20A%20secondary%20benefit%20is%20decorrelation%3A%20random%20sampling%20breaks%0A%20%20%20%20the%20correlation%20between%20consecutive%20trajectories%20from%20the%20same%20generator%2C%20giving%0A%20%20%20%20better%20gradient%20estimates%20(especially%20when%20mixing%20tasks%20of%20different%20difficulties).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Actor%2C%20Trajectory%2C%20deque%2C%20endpoint%2C%20random)%3A%0A%20%20%20%20class%20ReplayBuffer(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Stores%20trajectories%20for%20async%20RL%20training.%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20max_size%3A%20int%20%3D%201000)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer%3A%20deque%5BTrajectory%5D%20%3D%20deque(maxlen%3Dmax_size)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.total_added%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BReplayBuffer%5D%20Initialized%20with%20max_size%3D%7Bmax_size%7D%22)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20add(self%2C%20trajectory%3A%20Trajectory)%20-%3E%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Add%20a%20trajectory%20to%20the%20buffer.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer.append(trajectory)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.total_added%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20sample(self%2C%20batch_size%3A%20int)%20-%3E%20list%5BTrajectory%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Sample%20a%20batch%20of%20trajectories.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20len(self.buffer)%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20n%20%3D%20min(batch_size%2C%20len(self.buffer))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20random.sample(list(self.buffer)%2C%20n)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20size(self)%20-%3E%20int%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20len(self.buffer)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20clear(self)%20-%3E%20int%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Clear%20the%20buffer.%20Returns%20number%20of%20items%20removed.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20count%20%3D%20len(self.buffer)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.buffer.clear()%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20count%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20stats(self)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20len(self.buffer)%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22size%22%3A%200%2C%20%22total_added%22%3A%20self.total_added%2C%20%22avg_reward%22%3A%200.0%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20rewards%20%3D%20%5Bt.reward%20for%20t%20in%20self.buffer%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%20%3D%20%7B%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20t%20in%20self.buffer%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20fm%20%3D%20t.failure_mode%20or%20%22unknown%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%5Bfm%5D%20%3D%20failure_modes.get(fm%2C%200)%20%2B%201%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22size%22%3A%20len(self.buffer)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22total_added%22%3A%20self.total_added%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22avg_reward%22%3A%20sum(rewards)%20%2F%20len(rewards)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22correct_rate%22%3A%20sum(1%20for%20t%20in%20self.buffer%20if%20t.is_correct)%20%2F%20len(self.buffer)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22format_rate%22%3A%20sum(1%20for%20t%20in%20self.buffer%20if%20t.has_answer_tag)%20%2F%20len(self.buffer)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22failure_modes%22%3A%20failure_modes%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20return%20(ReplayBuffer%2C)%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%20Actor%203%3A%20TrainerActor%0A%0A%20%20%20%20The%20trainer%20loads%20the%20model%2C%20receives%20batches%20of%20trajectories%2C%20and%20computes%0A%20%20%20%20policy%20gradient%20updates.%20We%20use%20**REINFORCE**%20--%20the%20simplest%20policy%20gradient%0A%20%20%20%20method.%20Production%20systems%20typically%20use%20PPO%20or%20GRPO%2C%20which%20are%20variations%20that%20improve%0A%20%20%20%20stability%2C%20but%20the%20approach%20looks%20similar%20from%20a%20systems%20perspective.%20In%20other%20words%2C%0A%20%20%20%20REINFORCE%20lets%20us%20focus%20on%20the%20*system*%20(actors%2C%20weight%20sync%2C%20async%20coordination)%20rather%0A%20%20%20%20than%20the%20algorithm.%0A%0A%20%20%20%20The%20loss%20for%20each%20trajectory%20is%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20loss%20%3D%20-sum(log_prob(response_token_i))%20*%20(reward%20-%20baseline)%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20trainer%20is%20the%20most%20complex%20actor%2C%20with%20several%20responsibilities%3A%0A%0A%20%20%20%20-%20**%60setup()%60**%20%E2%80%94%20loads%20the%20model%20onto%20GPU%200%2C%20creates%20the%20optimizer%2C%20and%0A%20%20%20%20%20%20registers%20RDMA%20circular%20buffer%20slots%0A%20%20%20%20-%20**%60train_step()%60**%20%E2%80%94%20REINFORCE%20policy%20gradient%20on%20a%20batch%20of%20trajectories%0A%20%20%20%20-%20**%60get_weight_handle()%60**%20%E2%80%94%20returns%20an%20RDMA%20handle%20to%20the%20current%20circular%0A%20%20%20%20%20%20buffer%20slot%20for%20generators%20to%20pull%20from%0A%20%20%20%20-%20**%60evaluate_zorplex()%60**%20%E2%80%94%20runs%20deterministic%20evaluation%20for%20before%2Fafter%20comparison%0A%0A%20%20%20%20**GPU%20assignment%20note%3A**%20Monarch%20doesn't%20assign%20GPUs%20automatically%20%E2%80%94%0A%20%20%20%20%60spawn_procs%60%20creates%20processes%2C%20but%20it's%20up%20to%20you%20to%20set%0A%20%20%20%20%60CUDA_VISIBLE_DEVICES%60%20in%20%60setup()%60.%20Here%2C%20the%20trainer%20hardcodes%20GPU%200%0A%20%20%20%20and%20generators%20use%20GPU%201%2B.%0A%0A%20%20%20%20**Circular%20buffer%20with%20CPU%20staging**%20(from%20%5BNB07%5D(.%2F07_rdma_weight_sync.html))%3A%20After%20each%20training%20step%2C%0A%20%20%20%20weights%20are%20copied%20GPU%20-%3E%20CPU%20into%20a%20circular%20buffer%20slot.%20Generators%20read%0A%20%20%20%20from%20CPU%20via%20RDMA%2C%20then%20copy%20to%20their%20own%20GPU.%20This%20decouples%20training%20from%0A%20%20%20%20weight%20distribution.%0A%0A%20%20%20%20%60%60%60%0A%20%20%20%20Trainer%20GPU%20--D2H--%3E%20CPU%20slot%5Bv%20%25%203%5D%20--RDMA--%3E%20Generator%20CPU%20staging%20--H2D--%3E%20Generator%20GPU%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Each%20slot%20is%20a%20single%20**contiguous**%20CPU%20buffer%20%E2%80%94%20all%20parameters%20packed%0A%20%20%20%20end-to-end.%20This%20means%20one%20RDMA%20read%20transfers%20the%20entire%20model.%20An%0A%20%20%20%20alternative%20is%20keeping%20parameters%20scattered%20and%20batching%20reads%20with%0A%20%20%20%20%60RDMAAction%60.%20We%20go%20into%20the%20different%20patterns%20and%20trade-offs%20in%20%5BNB07b%5D(.%2F07b_weight_sync_deep_dive.html).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Actor%2C%0A%20%20%20%20AutoModelForCausalLM%2C%0A%20%20%20%20AutoTokenizer%2C%0A%20%20%20%20F%2C%0A%20%20%20%20RDMABuffer%2C%0A%20%20%20%20TrainMetrics%2C%0A%20%20%20%20Trajectory%2C%0A%20%20%20%20current_rank%2C%0A%20%20%20%20endpoint%2C%0A%20%20%20%20generate_with_tools%2C%0A%20%20%20%20get_spec%2C%0A%20%20%20%20rdma_available%2C%0A%20%20%20%20torch%2C%0A)%3A%0A%20%20%20%20class%20TrainerActor(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Trains%20the%20model%20on%20trajectories.%0A%0A%20%20%20%20%20%20%20%20Uses%20setup()%20for%20heavy%20initialization%20(model%20loading%2C%20RDMA%20registration).%0A%20%20%20%20%20%20%20%20Implements%20circular%20buffer%20with%20CPU%20staging%20for%20weight%20distribution.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(%0A%20%20%20%20%20%20%20%20%20%20%20%20self%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20model_name%3A%20str%20%3D%20%22Qwen%2FQwen2.5-0.5B-Instruct%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20lr%3A%20float%20%3D%201e-5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20device%3A%20str%20%3D%20%22cuda%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20n_buffer_slots%3A%20int%20%3D%203%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Lightweight%20init%20-%20just%20store%20config%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model_name%20%3D%20model_name%0A%20%20%20%20%20%20%20%20%20%20%20%20self.lr%20%3D%20lr%0A%20%20%20%20%20%20%20%20%20%20%20%20self.device_config%20%3D%20device%0A%20%20%20%20%20%20%20%20%20%20%20%20self.n_buffer_slots%20%3D%20n_buffer_slots%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._ready%20%3D%20False%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTrainer%3A%7Bself.rank%7D%5D%20Spawned%2C%20waiting%20for%20setup()...%22)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20setup(self)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Heavy%20initialization%3A%20load%20model%2C%20create%20optimizer%2C%20set%20up%20circular%20buffer.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20os%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self._ready%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22status%22%3A%20%22already_ready%22%7D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Trainer%20always%20uses%20GPU%200%0A%20%20%20%20%20%20%20%20%20%20%20%20os.environ%5B%22CUDA_VISIBLE_DEVICES%22%5D%20%3D%20%220%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.device%20%3D%20%22cuda%22%20if%20torch.cuda.is_available()%20else%20%22cpu%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_version%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20self.train_steps%20%3D%200%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTrainer%3A%7Bself.rank%7D%5D%20Loading%20model%20%7Bself.model_name%7D%20on%20GPU%200...%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.tokenizer%20%3D%20AutoTokenizer.from_pretrained(self.model_name)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model%20%3D%20AutoModelForCausalLM.from_pretrained(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.model_name%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20torch_dtype%3Dtorch.bfloat16%20if%20self.device%20%3D%3D%20%22cuda%22%20else%20torch.float32%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20).to(self.device)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self.tokenizer.pad_token%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.tokenizer.pad_token%20%3D%20self.tokenizer.eos_token%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.optimizer%20%3D%20torch.optim.AdamW(self.model.parameters()%2C%20lr%3Dself.lr)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20---%20Circular%20buffer%20with%20CPU%20staging%20---%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Each%20slot%20is%20a%20single%20contiguous%20CPU%20buffer%20that%20holds%20ALL%20model%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20parameters%20packed%20end-to-end.%20We%20copy%20the%20full%20state_dict%20into%20one%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20flat%20region%2C%20which%20means%20one%20RDMA%20read%20per%20sync%20(fewest%20round%20trips).%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Alternative%3A%20keep%20parameters%20scattered%20(one%20buffer%20per%20param)%20and%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20use%20RDMAAction%20to%20batch%20multiple%20reads%20into%20one%20operation.%20See%20NB07b%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20for%20a%20comparison%20of%20the%20different%20patterns%20and%20trade-offs.%0A%20%20%20%20%20%20%20%20%20%20%20%20total_bytes%20%3D%20sum(p.numel()%20*%20p.element_size()%20for%20p%20in%20self.model.parameters())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._slots%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20torch.empty(total_bytes%2C%20dtype%3Dtorch.uint8)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20_%20in%20range(self.n_buffer_slots)%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._slot_handles%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20rdma_available()%20and%20RDMABuffer%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20slot%20in%20self._slots%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._slot_handles.append(RDMABuffer(slot))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTrainer%3A%7Bself.rank%7D%5D%20RDMA%20handles%20registered%20for%20%7Bself.n_buffer_slots%7D%20circular%20buffer%20slots%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%20as%20e%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTrainer%3A%7Bself.rank%7D%5D%20RDMA%20registration%20failed%3A%20%7Be%7D%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._slot_handles%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._param_meta%20%3D%20%7B%7D%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%20name%2C%20p%20in%20self.model.named_parameters()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._param_meta%5Bname%5D%20%3D%20(offset%2C%20tuple(p.shape)%2C%20p.dtype)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20offset%20%2B%3D%20p.numel()%20*%20p.element_size()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._publish_weights()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._ready%20%3D%20True%0A%20%20%20%20%20%20%20%20%20%20%20%20param_count%20%3D%20sum(p.numel()%20for%20p%20in%20self.model.parameters())%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTrainer%3A%7Bself.rank%7D%5D%20Ready!%20%7Bparam_count%3A%2C%7D%20params%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22RDMA%3D%7Blen(self._slot_handles)%20%3E%200%7D%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22buffer_slots%3D%7Bself.n_buffer_slots%7D%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22status%22%3A%20%22ready%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22params%22%3A%20param_count%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22rdma%22%3A%20len(self._slot_handles)%20%3E%200%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22buffer_slots%22%3A%20self.n_buffer_slots%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20%20%20%20%20def%20_publish_weights(self)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Copy%20GPU%20params%20to%20the%20current%20circular%20buffer%20slot%20(D2H).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20slot_idx%20%3D%20self.policy_version%20%25%20self.n_buffer_slots%0A%20%20%20%20%20%20%20%20%20%20%20%20slot%20%3D%20self._slots%5Bslot_idx%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20name%2C%20p%20in%20self.model.named_parameters()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20off%2C%20shape%2C%20dtype%20%3D%20self._param_meta%5Bname%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nbytes%20%3D%20p.numel()%20*%20p.element_size()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20slot%5Boff%3Aoff%20%2B%20nbytes%5D.copy_(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p.data.view(-1).view(torch.uint8).cpu()%2C%20non_blocking%3DTrue%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%20torch.cuda.synchronize()%20%20%23%20Ensure%20D2H%20complete%20before%20RDMA%20reads%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_weight_handle(self)%20-%3E%20tuple%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Get%20RDMA%20handle%20for%20the%20latest%20weight%20slot.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Returns%20(handle_or_None%2C%20param_meta%2C%20version%2C%20total_bytes).%0A%20%20%20%20%20%20%20%20%20%20%20%20If%20RDMA%20unavailable%2C%20handle%20is%20None%20and%20caller%20should%20use%20get_state_dict().%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20total_bytes%20%3D%20sum(p.numel()%20*%20p.element_size()%20for%20p%20in%20self.model.parameters())%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self._slot_handles%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20slot_idx%20%3D%20self.policy_version%20%25%20self.n_buffer_slots%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20self._slot_handles%5Bslot_idx%5D%2C%20self._param_meta%2C%20self.policy_version%2C%20total_bytes%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20None%2C%20self._param_meta%2C%20self.policy_version%2C%20total_bytes%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_state_dict(self)%20-%3E%20tuple%5Bdict%2C%20int%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Fallback%3A%20get%20state%20dict%20directly%20(when%20RDMA%20not%20available).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.model.state_dict()%2C%20self.policy_version%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_version(self)%20-%3E%20int%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.policy_version%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20train_step(self%2C%20trajectories%3A%20list%5BTrajectory%5D%2C%20baseline%3A%20float)%20-%3E%20TrainMetrics%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Train%20on%20a%20batch%20of%20trajectories%20using%20REINFORCE.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Each%20trajectory%20carries%20pre-tokenized%20input_ids%20and%20a%20prompt_length%0A%20%20%20%20%20%20%20%20%20%20%20%20boundary%20from%20the%20generator%2C%20so%20we%20just%20slice%20and%20compute%20log-probs.%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20len(trajectories)%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20TrainMetrics(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20step%3Dself.train_steps%2C%20loss%3D0.0%2C%20batch_size%3D0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_reward%3D0.0%2C%20policy_version%3Dself.policy_version%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model.train()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.optimizer.zero_grad()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20losses%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20valid_count%20%3D%200%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20traj%20in%20trajectories%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20traj.input_ids%20or%20traj.prompt_length%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Step%201%3A%20Load%20pre-tokenized%20sequence%20from%20the%20generator%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20full_ids%20%3D%20torch.tensor(traj.input_ids%2C%20device%3Dself.device).unsqueeze(0)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prompt_len%20%3D%20traj.prompt_length%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20full_ids.shape%5B1%5D%20%3C%3D%20prompt_len%20%2B%201%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Step%202-3%3A%20Forward%20pass%2C%20then%20slice%20at%20prompt_length%20for%20response-only%20log-probs%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20with%20torch.amp.autocast('cuda'%2C%20enabled%3Dself.device%20%3D%3D%20%22cuda%22)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self.model(full_ids).logits%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20logits%5Bi%5D%20predicts%20token%5Bi%2B1%5D%2C%20so%20start%20at%20prompt_len%20-%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20shift_logits%20%3D%20logits%5B%3A%2C%20prompt_len%20-%201%3A-1%2C%20%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20shift_labels%20%3D%20full_ids%5B%3A%2C%20prompt_len%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20log_probs%20%3D%20F.log_softmax(shift_logits%2C%20dim%3D-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20token_log_probs%20%3D%20log_probs.gather(2%2C%20shift_labels.unsqueeze(-1)).squeeze(-1)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Step%204%3A%20REINFORCE%20loss%20%3D%20-log_prob%20*%20advantage%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20advantage%20%3D%20traj.reward%20-%20baseline%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20losses.append(-token_log_probs.sum()%20*%20advantage)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20valid_count%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Step%205%3A%20Optimizer%20step%2C%20then%20publish%20weights%20to%20circular%20buffer%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20valid_count%20%3E%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_loss%20%3D%20torch.stack(losses).sum()%20%2F%20valid_count%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_loss.backward()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20torch.nn.utils.clip_grad_norm_(self.model.parameters()%2C%20max_norm%3D1.0)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.optimizer.step()%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_loss%20%3D%20torch.tensor(0.0)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Bump%20version%2C%20then%20publish%20weights%20to%20the%20new%20slot.%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Safe%20because%20Monarch%20actors%20process%20endpoints%20sequentially.%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_version%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20self._publish_weights()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.train_steps%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20avg_reward%20%3D%20sum(t.reward%20for%20t%20in%20trajectories)%20%2F%20len(trajectories)%0A%20%20%20%20%20%20%20%20%20%20%20%20correct_rate%20%3D%20sum(1%20for%20t%20in%20trajectories%20if%20t.is_correct)%20%2F%20len(trajectories)%0A%20%20%20%20%20%20%20%20%20%20%20%20format_rate%20%3D%20sum(1%20for%20t%20in%20trajectories%20if%20t.has_answer_tag)%20%2F%20len(trajectories)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20TrainMetrics(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20step%3Dself.train_steps%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20loss%3Davg_loss.item()%20if%20torch.is_tensor(avg_loss)%20else%20avg_loss%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_size%3Dlen(trajectories)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20avg_reward%3Davg_reward%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_version%3Dself.policy_version%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20correct_rate%3Dcorrect_rate%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20format_rate%3Dformat_rate%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20evaluate_zorplex(self%2C%20num_samples%3A%20int%20%3D%2010%2C%20seed%3A%20int%20%3D%2042)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Evaluate%20current%20model%20on%20compositional%20Zorplex%20tasks.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20re%20as%20_re%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model.eval()%0A%20%20%20%20%20%20%20%20%20%20%20%20torch.manual_seed(seed)%20%20%23%20Deterministic%20evaluation%0A%20%20%20%20%20%20%20%20%20%20%20%20spec%20%3D%20get_spec(%22compositional%22%2C%20seed%3Dseed)%0A%20%20%20%20%20%20%20%20%20%20%20%20correct%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20total_turns%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20total_tools%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20format_ok%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%20%3D%20%7B%22success%22%3A%200%2C%20%22wrong_format%22%3A%200%2C%20%22tool_spam%22%3A%200%2C%20%22wrong_answer%22%3A%200%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20_%20in%20range(num_samples)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20task%20%3D%20spec.generate_task()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20result%20%3D%20generate_with_tools(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.model%2C%20self.tokenizer%2C%20spec%2C%20task%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.device%2C%20max_turns%3D5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20temperature%3D0.0%2C%20do_sample%3DFalse%2C%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%20%20%20%20%20correct%20%2B%3D%20int(result.is_correct)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20total_turns%20%2B%3D%20len(result.turns)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20total_tools%20%2B%3D%20result.total_tool_calls%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20has_tag%20%3D%20bool(_re.search(r'%5C%5BANSWER%5C%5D'%2C%20result.final_text))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20format_ok%20%2B%3D%20int(has_tag)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20result.is_correct%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%5B%22success%22%5D%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20elif%20not%20has_tag%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%5B%22wrong_format%22%5D%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20elif%20result.total_tool_calls%20%3E%203%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%5B%22tool_spam%22%5D%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_modes%5B%22wrong_answer%22%5D%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22accuracy%22%3A%20correct%20%2F%20num_samples%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22correct%22%3A%20correct%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22total%22%3A%20num_samples%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22avg_turns%22%3A%20total_turns%20%2F%20num_samples%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22avg_tools%22%3A%20total_tools%20%2F%20num_samples%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22format_rate%22%3A%20format_ok%20%2F%20num_samples%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22failure_modes%22%3A%20failure_modes%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20return%20(TrainerActor%2C)%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%20How%20%60train_step%60%20Works%0A%0A%20%20%20%20Each%20trajectory%20arrives%20with%20pre-tokenized%20%60input_ids%60%20and%20a%20%60prompt_length%60%0A%20%20%20%20boundary%20(computed%20by%20the%20generator%20at%20generation%20time).%20The%20trainer%3A%0A%0A%20%20%20%201.%20**Loads%20the%20token%20sequence**%20directly%20from%20%60traj.input_ids%60%20--%20no%20re-tokenization.%0A%20%20%20%202.%20**Slices%20at%20%60prompt_length%60**%20to%20separate%20prompt%20from%20response%20tokens.%0A%20%20%20%203.%20**Computes%20log-probs**%20on%20response%20tokens%20only%20(%60logits%5Bi%5D%60%20predicts%20%60token%5Bi%2B1%5D%60%2C%0A%20%20%20%20%20%20%20so%20we%20start%20at%20%60prompt_length%20-%201%60).%0A%20%20%20%204.%20**Computes%20loss**%3A%20%60loss%20%3D%20-sum(log_probs)%20*%20advantage%60%20where%0A%20%20%20%20%20%20%20%60advantage%20%3D%20reward%20-%20baseline%60.%20Positive%20advantage%20reinforces%20the%20response.%0A%20%20%20%205.%20**Steps%20the%20optimizer**%20once%20for%20the%20whole%20batch%2C%20then%20publishes%20new%20weights%0A%20%20%20%20%20%20%20to%20the%20circular%20buffer.%0A%0A%20%20%20%20Look%20for%20the%20%60%23%20Step%20N%3A%60%20comments%20in%20%60train_step%60%20above%20--%20they%20correspond%20to%0A%20%20%20%20these%20steps.%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%20Actor%204%3A%20GeneratorWorker%0A%0A%20%20%20%20Each%20generator%20loads%20its%20own%20copy%20of%20the%20model%2C%20generates%20its%20own%20tasks%20from%0A%20%20%20%20its%20seeded%20spec%2C%20and%20runs%20inference%20independently.%20The%20key%20endpoint%20is%0A%20%20%20%20%60generate_trajectory()%60%20--%20it%20generates%20a%20task%2C%20runs%20multi-turn%20inference%0A%20%20%20%20with%20tool%20execution%2C%20and%20returns%20a%20complete%20%60Trajectory%60%20with%20pre-tokenized%0A%20%20%20%20%60input_ids%60%20and%20%60prompt_length%60%20for%20the%20trainer.%0A%0A%20%20%20%20**Reward%20shaping.**%20Instead%20of%20a%20binary%200%2F1%20reward%2C%20we%20decompose%20rewards%20from%0A%20%20%20%20the%20failure%20modes%20identified%20in%20%5BNB05%5D(.%2F05_rl_intro.html)%3A%0A%0A%20%20%20%20%7C%20Component%20%7C%20Value%20%7C%20Why%20%7C%0A%20%20%20%20%7C-----------%7C-------%7C-----%7C%0A%20%20%20%20%7C%20Correct%20answer%20%7C%20%2B1.0%20%7C%20The%20main%20signal%20%7C%0A%20%20%20%20%7C%20Format%20compliance%20(%60%5BANSWER%5D%60%20tag)%20%7C%20%2B0.2%20%7C%20Learnable%20even%20when%20wrong%20%7C%0A%20%20%20%20%7C%20Tool%20spam%20penalty%20%7C%20-0.1%20per%20call%20beyond%202%20%7C%20Discourages%20degenerate%20loops%20%7C%0A%0A%20%20%20%20This%20means%20a%20correct%2C%20well-formatted%20response%20earns%20up%20to%201.2%2C%20while%20a%0A%20%20%20%20format-only%20success%20(wrong%20answer%20but%20used%20%60%5BANSWER%5D%60)%20earns%200.2.%20The%0A%20%20%20%20gradient%20signal%20is%20richer%20than%20binary%3A%20the%20model%20gets%20*partial%20credit*%20for%0A%20%20%20%20good%20formatting%20even%20before%20it%20learns%20the%20right%20answers.%0A%0A%20%20%20%20Weight%20sync%20uses%20the%20pattern%20from%20%5BNB07%5D(.%2F07_rdma_weight_sync.html)%3A%20the%20trainer%20publishes%20weights%20to%20CPU%0A%20%20%20%20slots%20(circular%20buffer)%2C%20and%20generators%20pull%20via%20RDMA%20into%20a%20CPU%20staging%0A%20%20%20%20buffer%2C%20then%20scatter%20into%20GPU%20parameters%20(H2D%20copy).%20Ideally%20we'd%20load%0A%20%20%20%20directly%20from%20the%20trainer's%20CPU%20buffer%20into%20the%20model's%20%60state_dict%60%20to%0A%20%20%20%20avoid%20the%20extra%20copy%2C%20but%20we%20hit%20%60RDMABuffer%60%20bugs%20doing%20that%20%E2%80%94%20will%20fix.%0A%20%20%20%20Fallback%20path%20(%60sync_weights%60%20using%20%60state_dict%60)%20stays%20for%20when%20RDMA%20is%20unavailable.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Actor%2C%0A%20%20%20%20AutoModelForCausalLM%2C%0A%20%20%20%20AutoTokenizer%2C%0A%20%20%20%20Task%2C%0A%20%20%20%20Trajectory%2C%0A%20%20%20%20current_rank%2C%0A%20%20%20%20endpoint%2C%0A%20%20%20%20generate_with_tools%2C%0A%20%20%20%20get_spec%2C%0A%20%20%20%20torch%2C%0A)%3A%0A%20%20%20%20class%20GeneratorWorker(Actor)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Individual%20generator%20worker.%0A%0A%20%20%20%20%20%20%20%20Uses%20setup()%20for%20heavy%20initialization%20(model%20loading).%0A%20%20%20%20%20%20%20%20Weight%20sync%20uses%20CPU%20staging%20buffer%20for%20explicit%20RDMA%20-%3E%20H2D%20flow.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%0A%20%20%20%20%20%20%20%20def%20__init__(%0A%20%20%20%20%20%20%20%20%20%20%20%20self%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20model_name%3A%20str%20%3D%20%22Qwen%2FQwen2.5-0.5B-Instruct%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20difficulty%3A%20str%20%3D%20%22easy%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20device%3A%20str%20%3D%20%22cuda%22%2C%0A%20%20%20%20%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Lightweight%20init%20-%20just%20store%20config%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model_name%20%3D%20model_name%0A%20%20%20%20%20%20%20%20%20%20%20%20self.difficulty%20%3D%20difficulty%0A%20%20%20%20%20%20%20%20%20%20%20%20self.device_config%20%3D%20device%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._ready%20%3D%20False%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BGeneratorWorker%3A%7Bself.rank%7D%5D%20Spawned%2C%20waiting%20for%20setup()...%22)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20setup(self)%20-%3E%20dict%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Heavy%20initialization%3A%20load%20model%2C%20create%20weight%20buffer.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20os%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self._ready%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22status%22%3A%20%22already_ready%22%7D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Generators%20use%20GPU%201%20%2B%20rank%20(trainer%20uses%20GPU%200)%0A%20%20%20%20%20%20%20%20%20%20%20%20gpu_id%20%3D%201%20%2B%20self.rank%0A%20%20%20%20%20%20%20%20%20%20%20%20os.environ%5B%22CUDA_VISIBLE_DEVICES%22%5D%20%3D%20str(gpu_id)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.device%20%3D%20%22cuda%22%20if%20torch.cuda.is_available()%20else%20%22cpu%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_version%20%3D%200%0A%20%20%20%20%20%20%20%20%20%20%20%20self.generations%20%3D%200%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BGeneratorWorker%3A%7Bself.rank%7D%5D%20Loading%20model%20%7Bself.model_name%7D%20on%20GPU%20%7Bgpu_id%7D...%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.tokenizer%20%3D%20AutoTokenizer.from_pretrained(self.model_name)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model%20%3D%20AutoModelForCausalLM.from_pretrained(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.model_name%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20torch_dtype%3Dtorch.bfloat16%20if%20self.device%20%3D%3D%20%22cuda%22%20else%20torch.float32%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20).to(self.device)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self.tokenizer.pad_token%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.tokenizer.pad_token%20%3D%20self.tokenizer.eos_token%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.spec%20%3D%20get_spec(%22compositional%22%2C%20difficulty%3Dself.difficulty%2C%20seed%3D42%20%2B%20self.rank)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._sync_buf%20%3D%20None%20%20%23%20CPU%20staging%20buffer%20for%20RDMA%20weight%20sync%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self._ready%20%3D%20True%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BGeneratorWorker%3A%7Bself.rank%7D%5D%20Ready%20on%20GPU%20%7Bgpu_id%7D!%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%7B%22status%22%3A%20%22ready%22%2C%20%22rank%22%3A%20self.rank%2C%20%22gpu%22%3A%20gpu_id%7D%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20get_version(self)%20-%3E%20int%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.policy_version%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20sync_weights_from_buffer(self%2C%20handle%2C%20param_meta%3A%20dict%2C%20version%3A%20int%2C%20total_bytes%3A%20int)%20-%3E%20bool%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Sync%20weights%20via%20RDMA%20from%20trainer's%20circular%20buffer.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Flow%3A%20Trainer%20CPU%20slot%20--RDMA--%3E%20Generator%20CPU%20staging%20--H2D--%3E%20Generator%20GPU%20params%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20NOTE%3A%20Ideally%20we'd%20load%20weights%20from%20the%20trainer's%20CPU%20buffer%20directly%0A%20%20%20%20%20%20%20%20%20%20%20%20into%20the%20model's%20state_dict%20parameters%2C%20avoiding%20the%20intermediate%20copy.%0A%20%20%20%20%20%20%20%20%20%20%20%20We%20hit%20bugs%20with%20RDMABuffer%20targeting%20model%20tensors%20directly%2C%20so%20for%0A%20%20%20%20%20%20%20%20%20%20%20%20now%20we%20read%20into%20a%20separate%20CPU%20buffer%20and%20then%20do%20a%20H2D%20copy.%20Will%20fix.%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20version%20%3C%3D%20self.policy_version%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20False%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Allocate%20CPU%20staging%20buffer%20on%20first%20sync%20(reuse%20thereafter).%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Ideally%20we'd%20load%20from%20the%20trainer's%20CPU%20buffer%20straight%20into%20the%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20model's%20state_dict%20to%20skip%20this%20copy%2C%20but%20we%20hit%20RDMABuffer%20bugs%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20doing%20that%20--%20so%20for%20now%2C%20separate%20CPU%20buffer%20%2B%20H2D%20scatter.%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20self._sync_buf%20is%20None%20or%20self._sync_buf.numel()%20%3C%20total_bytes%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self._sync_buf%20%3D%20torch.empty(total_bytes%2C%20dtype%3Dtorch.uint8)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20RDMA%20read%3A%20trainer%20CPU%20slot%20-%3E%20generator%20CPU%20staging%20buffer%0A%20%20%20%20%20%20%20%20%20%20%20%20byte_view%20%3D%20self._sync_buf%5B%3Atotal_bytes%5D.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%20%20%20%20%23%20Scatter%20from%20CPU%20staging%20into%20GPU%20model%20params%20(H2D%20copy%20per%20parameter)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20name%2C%20p%20in%20self.model.named_parameters()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20off%2C%20shape%2C%20dtype%20%3D%20param_meta%5Bname%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nbytes%20%3D%20p.numel()%20*%20p.element_size()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20src%20%3D%20self._sync_buf%5Boff%3Aoff%20%2B%20nbytes%5D.view(dtype).view(shape)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20p.data.copy_(src)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_version%20%3D%20version%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20True%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20sync_weights(self%2C%20state_dict%3A%20dict%2C%20version%3A%20int)%20-%3E%20bool%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Sync%20weights%20directly%20(fallback%20when%20RDMA%20unavailable).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20version%20%3C%3D%20self.policy_version%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20False%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model.load_state_dict(state_dict)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.policy_version%20%3D%20version%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20True%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20generate(self%2C%20question%3A%20str%2C%20answer%3A%20int%2C%20max_turns%3A%20int%20%3D%205)%20-%3E%20Trajectory%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Generate%20a%20trajectory%20for%20a%20given%20task%20(used%20for%20examples%2Fdebugging).%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model.eval()%0A%20%20%20%20%20%20%20%20%20%20%20%20task%20%3D%20Task(question%3Dquestion%2C%20correct_answer%3Danswer%2C%20metadata%3D%7B%7D)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self._run_generation(task%2C%20max_turns)%0A%0A%20%20%20%20%20%20%20%20%40endpoint%0A%20%20%20%20%20%20%20%20def%20generate_trajectory(self%2C%20max_turns%3A%20int%20%3D%205)%20-%3E%20Trajectory%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Generate%20a%20trajectory%20using%20a%20self-generated%20task.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Each%20generator%20has%20its%20own%20seeded%20spec%2C%20so%20broadcasting%20this%20endpoint%0A%20%20%20%20%20%20%20%20%20%20%20%20to%20all%20generators%20produces%20diverse%20trajectories%20from%20different%20tasks.%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20self.model.eval()%0A%20%20%20%20%20%20%20%20%20%20%20%20task%20%3D%20self.spec.generate_task()%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self._run_generation(task%2C%20max_turns)%0A%0A%20%20%20%20%20%20%20%20def%20_run_generation(self%2C%20task%3A%20Task%2C%20max_turns%3A%20int)%20-%3E%20Trajectory%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Shared%20generation%20logic%3A%20run%20inference%2C%20compute%20tokens%2C%20return%20Trajectory.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20import%20re%20as%20_re%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20result%20%3D%20generate_with_tools(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.model%2C%20self.tokenizer%2C%20self.spec%2C%20task%2C%20self.device%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20max_turns%3Dmax_turns%2C%20max_tokens_per_turn%3D150%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.generations%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Build%20model-only%20text%20(generated%20tokens%20without%20injected%20tool%20results)%0A%20%20%20%20%20%20%20%20%20%20%20%20model_only_text%20%3D%20%22%22.join(t.generated_text%20for%20t%20in%20result.turns)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Detect%20%5BANSWER%5D%20tag%20and%20classify%20failure%20mode%0A%20%20%20%20%20%20%20%20%20%20%20%20has_answer_tag%20%3D%20bool(_re.search(r'%5C%5BANSWER%5C%5D'%2C%20result.final_text))%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20result.is_correct%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_mode%20%3D%20%22success%22%0A%20%20%20%20%20%20%20%20%20%20%20%20elif%20not%20has_answer_tag%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_mode%20%3D%20%22wrong_format%22%0A%20%20%20%20%20%20%20%20%20%20%20%20elif%20result.total_tool_calls%20%3E%203%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_mode%20%3D%20%22tool_spam%22%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_mode%20%3D%20%22wrong_answer%22%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Pre-tokenize%20for%20the%20trainer%3A%20prompt%20%2B%20model_only_text%0A%20%20%20%20%20%20%20%20%20%20%20%20messages%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%22role%22%3A%20%22system%22%2C%20%22content%22%3A%20self.spec.get_system_prompt(with_hint%3DTrue)%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%22role%22%3A%20%22user%22%2C%20%22content%22%3A%20task.question%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20prompt_text%20%3D%20self.tokenizer.apply_chat_template(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20messages%2C%20tokenize%3DFalse%2C%20add_generation_prompt%3DTrue%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20prompt_ids%20%3D%20self.tokenizer(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prompt_text%2C%20return_tensors%3D%22pt%22%2C%20add_special_tokens%3DFalse%0A%20%20%20%20%20%20%20%20%20%20%20%20)%5B%22input_ids%22%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20prompt_length%20%3D%20prompt_ids.shape%5B1%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20full_ids%20%3D%20self.tokenizer(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prompt_text%20%2B%20model_only_text%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return_tensors%3D%22pt%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20add_special_tokens%3DFalse%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20truncation%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20max_length%3D1024%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%5B%22input_ids%22%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Reward%20shaping%20(see%20NB05%20%22From%20Failure%20Modes%20to%20RL%20Rewards%22)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20%20%20%2B1.0%20for%20correct%20answer%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20%20%20%2B0.2%20for%20format%20compliance%20(%5BANSWER%5D%20tag)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20%20%20-0.1%20per%20tool%20call%20beyond%202%20(discourages%20tool%20spam)%0A%20%20%20%20%20%20%20%20%20%20%20%20reward%20%3D%200.0%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20result.is_correct%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20reward%20%2B%3D%201.0%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20has_answer_tag%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20reward%20%2B%3D%200.2%0A%20%20%20%20%20%20%20%20%20%20%20%20excess_tools%20%3D%20max(0%2C%20result.total_tool_calls%20-%202)%0A%20%20%20%20%20%20%20%20%20%20%20%20reward%20-%3D%200.1%20*%20excess_tools%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20Trajectory(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20task_question%3Dtask.question%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20task_answer%3Dtask.correct_answer%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20response_text%3Dresult.final_text%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20reward%3Dreward%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20is_correct%3Dresult.is_correct%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20num_turns%3Dlen(result.turns)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20num_tool_calls%3Dresult.total_tool_calls%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20generator_id%3Dself.rank%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20policy_version%3Dself.policy_version%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20model_only_text%3Dmodel_only_text%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20has_answer_tag%3Dhas_answer_tag%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20failure_mode%3Dfailure_mode%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20input_ids%3Dfull_ids%5B0%5D.tolist()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20prompt_length%3Dprompt_length%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20return%20(GeneratorWorker%2C)%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%20Architecture%20Overview%0A%0A%20%20%20%20Now%20we%20have%20all%20our%20actors%20defined.%20Here's%20how%20they%20connect%20--%20this%20is%20the%0A%20%20%20%20**single-controller%20paradigm**%20from%20%5BNB01%5D(.%2F01_history_and_vision.html)%3A%20the%20notebook%20process%20orchestrates%0A%20%20%20%20everything%2C%20but%20actors%20do%20the%20heavy%20lifting%20on%20their%20own%20GPUs.%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%90%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%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%20GeneratorMesh%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%E2%94%82%20%20ZorplexService%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20(ActorMesh)%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%E2%94%82%20%20(Service)%20%20%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%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%90%20%20%E2%94%82%20%20tool%20%20%20%E2%94%82%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%90%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%E2%94%82%20Generator%200%20%20%20%E2%94%82%E2%94%80%E2%94%80%E2%94%BC%E2%94%80calls%E2%94%80%E2%94%80%E2%96%BA%E2%94%82%20%20%E2%94%82%20ZorplexWorker%20%E2%94%82%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%E2%94%82%20Generator%201%20%20%20%E2%94%82%E2%94%80%E2%94%80%E2%94%BC%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%94%80%E2%96%BA%E2%94%82%20%20%E2%94%82%20ZorplexWorker%20%E2%94%82%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%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%98%20%20%E2%94%82%E2%97%84%E2%94%80results%E2%94%82%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%98%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%E2%94%82%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%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%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%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%BC%E2%94%80%E2%94%80%E2%94%80%E2%94%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%20%20%20%20%20%20%20%20%20%20%E2%94%82%20trajectories%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20v%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%90%0A%20%20%20%20%E2%94%82%20%20%20%20ReplayBuffer%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%AC%E2%94%80%E2%94%80%E2%94%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%20%20%20%20%20%20%20%20%20%20%20%E2%94%82%20sample%20batch%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20v%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%90%0A%20%20%20%20%E2%94%82%20%20%20%20%20%20Trainer%20%20%20%20%20%20%20%20%E2%94%82%0A%20%20%20%20%E2%94%82%20%20(circular%20buffer)%20%20%E2%94%82%E2%94%80%E2%94%80%3E%20RDMA%20weight%20sync%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%98%20%20%20%20%20%20to%20GeneratorMesh%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Each%20generator%20calls%20zorplex%20tool%20endpoints%20during%20multi-turn%20inference%0A%20%20%20%20(e.g.%2C%20%60lookup_value%60%2C%20%60compute%60).%20The%20Service%20routes%20these%20calls%20round-robin%0A%20%20%20%20across%20ZorplexWorkers.%0A%0A%20%20%20%20**ActorMesh%20vs%20Service.**%20Generators%20are%20a%20plain%20**ActorMesh**%20--%20we%20address%0A%20%20%20%20them%20directly%20via%20%60.call()%60%20(broadcast%20to%20all)%20or%20%60.slice()%60%20(individual%0A%20%20%20%20access).%20This%20is%20natural%20for%20sync%20RL%20(broadcast%20generate%2C%20then%20train)%20and%0A%20%20%20%20for%20async%20RL%20(each%20thread%20slices%20its%20own%20generator).%20ZorplexWorkers%20are%0A%20%20%20%20wrapped%20in%20a%20**Service**%20(%5BNB06%5D(.%2F06_services.html)%20pattern)%20because%20they're%20stateless%3A%20any%0A%20%20%20%20worker%20can%20handle%20any%20request%2C%20so%20round-robin%20routing%20and%20health%20tracking%0A%20%20%20%20are%20useful.%20In%20production%20async%20RL%2C%20you%20might%20wrap%20generators%20in%20a%20Service%0A%20%20%20%20too%20--%20that%20gives%20you%20auto-scaling%20and%20health%20tracking%20--%20but%20here%20the%0A%20%20%20%20ActorMesh%20is%20simpler%20and%20lets%20us%20demonstrate%20both%20addressing%20patterns.%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%20Sync%20vs%20Async%20RL%0A%0A%20%20%20%20**Sync%20RL**%20(traditional)%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20%7C--generate--%7C--train--%7C--generate--%7C--train--%7C--generate--%7C--train--%7C%0A%20%20%20%20%60%60%60%0A%20%20%20%20Only%20ONE%20thing%20happens%20at%20a%20time.%20GPU%20sits%20idle%20during%20generation%2C%0A%20%20%20%20generator%20sits%20idle%20during%20training.%0A%0A%20%20%20%20**Async%20RL**%20(what%20we're%20building)%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20Gen0%3A%20%20%7C%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%7C%0A%20%20%20%20Gen1%3A%20%20%7C%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%E2%96%88%7C%0A%20%20%20%20Train%3A%20%20%20%20%20%20%7C%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%E2%96%93%7C%0A%20%20%20%20%60%60%60%0A%20%20%20%20Everything%20runs%20concurrently.%20More%20data%20collected%2C%20better%20GPU%20utilization.%0A%0A%20%20%20%20We'll%20run%20BOTH%20modes%20with%20the%20**same%20actors**%20and%20compare%20wall%20time%2C%20throughput%2C%0A%20%20%20%20and%20utilization.%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%20num_steps_slider%20%3D%20mo.ui.slider(10%2C%20100%2C%20value%3D20%2C%20label%3D%22Training%20steps%22)%0A%20%20%20%20num_generators_slider%20%3D%20mo.ui.slider(1%2C%204%2C%20value%3D2%2C%20label%3D%22Generators%22)%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%20Configuration%0A%0A%20%20%20%20Adjust%20parameters%20for%20the%20training%20run.%20**Marimo%20is%20reactive**%3A%20changing%20a%20slider%0A%20%20%20%20re-runs%20all%20downstream%20cells%20that%20depend%20on%20it.%20This%20means%20actors%20will%20be%0A%20%20%20%20re-spawned%20and%20both%20training%20loops%20will%20re-execute%20with%20the%20new%20values.%0A%0A%20%20%20%20%7Bnum_steps_slider%7D%0A%0A%20%20%20%20%7Bnum_generators_slider%7D%0A%0A%20%20%20%20**Batch%20size**%20is%20set%20to%20match%20the%20number%20of%20generators%20--%20each%20training%20step%0A%20%20%20%20trains%20on%20exactly%20one%20round%20of%20generation.%20This%20keeps%20the%20comparison%20fair%3A%0A%20%20%20%20sync%20and%20async%20train%20on%20the%20same%20amount%20of%20data%20per%20step.%0A%0A%20%20%20%20**Suggestions%3A**%20Start%20with%20defaults%20(20%20steps%2C%202%20generators)%20to%20see%20the%0A%20%20%20%20full%20pipeline.%20Then%20try%20increasing%20generators%20to%203-4%20to%20see%20the%20async%20throughput%0A%20%20%20%20advantage%20grow.%20Increasing%20training%20steps%20gives%20the%20model%20more%20updates%20but%20adds%0A%20%20%20%20wall%20time.%20Note%20that%20re-spawning%20actors%20(loading%20models%20onto%20GPUs)%20is%20the%20most%0A%20%20%20%20expensive%20part%20of%20the%20setup%20--%20the%20training%20loops%20themselves%20are%20relatively%20fast.%0A%0A%20%20%20%20**Try%20this%3A**%20Set%20generators%20to%201%20and%20watch%20the%20async%20timeline%20--%20with%20only%20one%0A%20%20%20%20generator%2C%20async%20degrades%20to%20near-sync%20performance%20because%20there's%20no%20parallel%0A%20%20%20%20generation%20to%20overlap%20with%20training.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%20num_generators_slider%2C%20num_steps_slider%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20threading%0A%20%20%20%20import%20time%0A%20%20%20%20from%20dataclasses%20import%20dataclass%2C%20field%0A%0A%20%20%20%20%40dataclass%0A%20%20%20%20class%20TimingEvent%3A%0A%20%20%20%20%20%20%20%20%22%22%22A%20single%20timed%20event%20for%20timeline%20visualization.%22%22%22%0A%20%20%20%20%20%20%20%20actor_id%3A%20str%0A%20%20%20%20%20%20%20%20event_type%3A%20str%20%20%23%20%22generate%22%2C%20%22train%22%2C%20%22sync%22%0A%20%20%20%20%20%20%20%20start_time%3A%20float%0A%20%20%20%20%20%20%20%20duration%3A%20float%0A%0A%20%20%20%20%40dataclass%0A%20%20%20%20class%20TimingStats%3A%0A%20%20%20%20%20%20%20%20%22%22%22Timing%20statistics%20for%20a%20training%20run.%22%22%22%0A%20%20%20%20%20%20%20%20mode%3A%20str%0A%20%20%20%20%20%20%20%20num_generators%3A%20int%0A%20%20%20%20%20%20%20%20num_steps%3A%20int%0A%20%20%20%20%20%20%20%20total_generations%3A%20int%0A%20%20%20%20%20%20%20%20wall_time%3A%20float%0A%20%20%20%20%20%20%20%20gen_times%3A%20list%20%3D%20field(default_factory%3Dlist)%0A%20%20%20%20%20%20%20%20train_times%3A%20list%20%3D%20field(default_factory%3Dlist)%0A%20%20%20%20%20%20%20%20events%3A%20list%20%3D%20field(default_factory%3Dlist)%20%20%23%20List%20of%20TimingEvent%0A%20%20%20%20%20%20%20%20rdma_syncs%3A%20int%20%3D%200%0A%20%20%20%20%20%20%20%20direct_syncs%3A%20int%20%3D%200%0A%20%20%20%20%20%20%20%20staleness%3A%20list%20%3D%20field(default_factory%3Dlist)%20%20%23%20policy_version%20gaps%20per%20train%20batch%0A%0A%20%20%20%20%20%20%20%20%40property%0A%20%20%20%20%20%20%20%20def%20gens_per_second(self)%20-%3E%20float%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.total_generations%20%2F%20self.wall_time%20if%20self.wall_time%20%3E%200%20else%200%0A%0A%20%20%20%20%20%20%20%20%40property%0A%20%20%20%20%20%20%20%20def%20steps_per_second(self)%20-%3E%20float%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.num_steps%20%2F%20self.wall_time%20if%20self.wall_time%20%3E%200%20else%200%0A%0A%20%20%20%20return%20TimingEvent%2C%20TimingStats%2C%20threading%2C%20time%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%20Spawning%20and%20Initializing%20Actors%0A%0A%20%20%20%20This%20is%20the%20**single-controller%20paradigm**%20in%20action.%20The%20notebook%20process%0A%20%20%20%20orchestrates%20a%20careful%20initialization%20sequence%3A%0A%0A%20%20%20%201.%20Spawn%20ZorplexWorkers%20via%20a%20**Service**%20(%5BNB06%5D(.%2F06_services.html)%20pattern%20--%20health%20tracking%2C%20round-robin)%0A%20%20%20%202.%20Spawn%20GeneratorWorkers%20as%20a%20plain%20**ActorMesh**%20and%20call%20%60setup()%60%20on%20all%20via%0A%20%20%20%20%20%20%20%60.call()%60%20broadcast%20(loads%20model%20onto%20each%20GPU)%0A%20%20%20%203.%20Spawn%20ReplayBuffer%20(CPU-only%2C%20ready%20immediately)%0A%20%20%20%204.%20Spawn%20Trainer%2C%20then%20call%20%60setup()%60%20(loads%20model%20onto%20GPU%200%2C%20registers%20RDMA%20buffers)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20GeneratorWorker%2C%0A%20%20%20%20ReplayBuffer%2C%0A%20%20%20%20Service%2C%0A%20%20%20%20TrainerActor%2C%0A%20%20%20%20ZorplexWorker%2C%0A%20%20%20%20num_generators_slider%2C%0A%20%20%20%20num_steps_slider%2C%0A%20%20%20%20register_service%2C%0A)%3A%0A%20%20%20%20from%20monarch.actor%20import%20this_host%0A%0A%20%20%20%20NUM_STEPS%20%3D%20num_steps_slider.value%0A%20%20%20%20NUM_GENERATORS%20%3D%20num_generators_slider.value%0A%20%20%20%20NUM_ZORPLEX%20%3D%202%0A%20%20%20%20BATCH_SIZE%20%3D%20NUM_GENERATORS%20%20%23%20Train%20on%20exactly%20one%20round%20of%20generation%20per%20step%0A%0A%20%20%20%20def%20setup_actors()%3A%0A%20%20%20%20%20%20%20%20%22%22%22Spawn%20and%20initialize%20all%20actors.%20Returns%20them%20for%20reuse.%22%22%22%0A%20%20%20%20%20%20%20%20host%20%3D%20this_host()%0A%0A%20%20%20%20%20%20%20%20%23%201.%20ZorplexWorkers%20--%20wrapped%20in%20a%20Service%20(NB06%20pattern)%20for%0A%20%20%20%20%20%20%20%20%23%20%20%20%20health%20tracking%20and%20round-robin%20routing%0A%20%20%20%20%20%20%20%20zorplex_worker_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%20NUM_ZORPLEX%7D)%0A%20%20%20%20%20%20%20%20zorplex_svc_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%201%7D)%0A%20%20%20%20%20%20%20%20zorplex_svc%20%3D%20zorplex_svc_procs.spawn(%22zorplex_svc%22%2C%20Service%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20service_name%3D%22zorplex%22%2C%20worker_class%3DZorplexWorker%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20procs%3Dzorplex_worker_procs%2C%20procs_per_replica%3D1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20difficulty%3D%22easy%22)%0A%0A%20%20%20%20%20%20%20%20%23%202.%20Generators%20--%20plain%20ActorMesh%20(no%20Service%20wrapper).%0A%20%20%20%20%20%20%20%20%23%20%20%20%20Each%20generator%20has%20its%20own%20GPU%20and%20model%20copy%3B%20we%20address%20them%0A%20%20%20%20%20%20%20%20%23%20%20%20%20via%20.call()%20broadcast%20or%20.slice()%20for%20individual%20access.%0A%20%20%20%20%20%20%20%20gen_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%20NUM_GENERATORS%7D)%0A%20%20%20%20%20%20%20%20generators%20%3D%20gen_procs.spawn(%22generators%22%2C%20GeneratorWorker)%0A%0A%20%20%20%20%20%20%20%20%23%203.%20ReplayBuffer%0A%20%20%20%20%20%20%20%20buffer_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%201%7D)%0A%20%20%20%20%20%20%20%20buffer%20%3D%20buffer_procs.spawn(%22buffer%22%2C%20ReplayBuffer%2C%20max_size%3D500)%0A%0A%20%20%20%20%20%20%20%20%23%204.%20Trainer%0A%20%20%20%20%20%20%20%20trainer_procs%20%3D%20host.spawn_procs(per_host%3D%7B%22procs%22%3A%201%7D)%0A%20%20%20%20%20%20%20%20trainer%20%3D%20trainer_procs.spawn(%22trainer%22%2C%20TrainerActor)%0A%0A%20%20%20%20%20%20%20%20%23%20Initialize%20actors%20that%20need%20setup%0A%20%20%20%20%20%20%20%20zorplex_svc.ping.call_one().get()%0A%0A%20%20%20%20%20%20%20%20print(%22%5BSETUP%5D%20Setting%20up%20generator%20workers...%22)%0A%20%20%20%20%20%20%20%20generators.setup.call().get()%20%20%23%20broadcast%20setup%20to%20all%20generators%0A%0A%20%20%20%20%20%20%20%20buffer.stats.call_one().get()%0A%0A%20%20%20%20%20%20%20%20print(%22%5BSETUP%5D%20Setting%20up%20trainer...%22)%0A%20%20%20%20%20%20%20%20trainer.setup.call_one().get()%0A%0A%20%20%20%20%20%20%20%20register_service(%22zorplex%22%2C%20zorplex_svc)%0A%0A%20%20%20%20%20%20%20%20print(f%22%5BSETUP%5D%20All%20actors%20ready!%20%7BNUM_GENERATORS%7D%20generators%2C%20%7BNUM_ZORPLEX%7D%20zorplex%20workers%22)%0A%0A%20%20%20%20%20%20%20%20%23%20Track%20ProcMeshes%20for%20cleanup%0A%20%20%20%20%20%20%20%20proc_meshes%20%3D%20%5Bzorplex_worker_procs%2C%20zorplex_svc_procs%2C%20gen_procs%2C%20buffer_procs%2C%20trainer_procs%5D%0A%0A%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22trainer%22%3A%20trainer%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22buffer%22%3A%20buffer%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22generators%22%3A%20generators%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22zorplex_svc%22%3A%20zorplex_svc%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22_proc_meshes%22%3A%20proc_meshes%2C%0A%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20def%20teardown_actors(actors)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Stop%20all%20ProcMeshes%2C%20releasing%20processes%20and%20GPU%20memory.%22%22%22%0A%20%20%20%20%20%20%20%20for%20pm%20in%20actors.get(%22_proc_meshes%22%2C%20%5B%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pm.stop(%22teardown%20for%20re-init%22).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pass%20%20%23%20Best-effort%20cleanup%0A%20%20%20%20%20%20%20%20print(%22%5BTEARDOWN%5D%20All%20actors%20stopped.%22)%0A%0A%20%20%20%20actors%20%3D%20setup_actors()%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20BATCH_SIZE%2C%0A%20%20%20%20%20%20%20%20NUM_GENERATORS%2C%0A%20%20%20%20%20%20%20%20NUM_STEPS%2C%0A%20%20%20%20%20%20%20%20actors%2C%0A%20%20%20%20%20%20%20%20setup_actors%2C%0A%20%20%20%20%20%20%20%20teardown_actors%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%23%20Before%20Training%3A%20Zorplex%20Baseline%0A%0A%20%20%20%20Let's%20evaluate%20the%20model%20*before*%20any%20training%20to%20establish%20a%20baseline.%0A%20%20%20%20We%20run%2010%20compositional%20Zorplex%20tasks%20and%20record%20accuracy%2C%20average%20turns%2C%0A%20%20%20%20and%20tool%20usage.%20This%20gives%20us%20a%20concrete%20%22before%22%20snapshot%20to%20compare%20against.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(actors%2C%20mo)%3A%0A%20%20%20%20print(%22Evaluating%20pre-training%20baseline...%22)%0A%20%20%20%20pre_eval%20%3D%20actors%5B%22trainer%22%5D.evaluate_zorplex.call_one(num_samples%3D10%2C%20seed%3D42).get()%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Pre-Training%20Results%0A%0A%20%20%20%20**Metrics%20refresher**%20(from%20%5BNB05%5D(.%2F05_rl_intro.html))%3A%20*Accuracy*%20is%20how%20often%20the%20model%20gets%20the%0A%20%20%20%20correct%20answer.%20*Format%20compliance*%20tracks%20whether%20it%20emits%20the%20%60%5BANSWER%5D%60%20tag%0A%20%20%20%20we%20trained%20it%20to%20use.%20*Avg%20turns%2Ftool%20calls*%20measure%20how%20many%20interaction%0A%20%20%20%20steps%20the%20model%20takes%20%E2%80%94%20lower%20is%20more%20efficient.%0A%0A%20%20%20%20%7C%20Metric%20%7C%20Value%20%7C%0A%20%20%20%20%7C--------%7C-------%7C%0A%20%20%20%20%7C%20Accuracy%20%7C%20%7Bpre_eval%5B'accuracy'%5D%3A.0%25%7D%20(%7Bpre_eval%5B'correct'%5D%7D%2F%7Bpre_eval%5B'total'%5D%7D)%20%7C%0A%20%20%20%20%7C%20Format%20compliance%20%7C%20%7Bpre_eval%5B'format_rate'%5D%3A.0%25%7D%20%7C%0A%20%20%20%20%7C%20Avg%20turns%20%7C%20%7Bpre_eval%5B'avg_turns'%5D%3A.1f%7D%20%7C%0A%20%20%20%20%7C%20Avg%20tool%20calls%20%7C%20%7Bpre_eval%5B'avg_tools'%5D%3A.1f%7D%20%7C%0A%0A%20%20%20%20**Failure%20mode%20breakdown%3A**%0A%0A%20%20%20%20%7C%20Mode%20%7C%20Count%20%7C%0A%20%20%20%20%7C------%7C-------%7C%0A%20%20%20%20%7C%20Success%20%7C%20%7Bpre_eval%5B'failure_modes'%5D%5B'success'%5D%7D%20%7C%0A%20%20%20%20%7C%20Wrong%20format%20%7C%20%7Bpre_eval%5B'failure_modes'%5D%5B'wrong_format'%5D%7D%20%7C%0A%20%20%20%20%7C%20Tool%20spam%20%7C%20%7Bpre_eval%5B'failure_modes'%5D%5B'tool_spam'%5D%7D%20%7C%0A%20%20%20%20%7C%20Wrong%20answer%20%7C%20%7Bpre_eval%5B'failure_modes'%5D%5B'wrong_answer'%5D%7D%20%7C%0A%0A%20%20%20%20This%20is%20our%20starting%20point.%20Let's%20see%20if%20training%20improves%20it%20at%20all...%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%20(pre_eval%2C)%0A%0A%0A%40app.cell%0Adef%20_(actors%2C%20mo)%3A%0A%20%20%20%20%23%20Generate%20one%20example%20trajectory%20so%20the%20reader%20can%20see%20what%20the%20pipeline%20produces%0A%20%20%20%20_gen%20%3D%20actors%5B%22generators%22%5D.slice(procs%3D0)%0A%20%20%20%20_example_traj%20%3D%20_gen.generate_trajectory.call_one().get()%0A%0A%20%20%20%20_status%20%3D%20%22Correct%22%20if%20_example_traj.is_correct%20else%20%22Wrong%22%0A%20%20%20%20_reward%20%3D%20_example_traj.reward%0A%0A%20%20%20%20%23%20Truncate%20long%20responses%20for%20display%0A%20%20%20%20_resp_display%20%3D%20_example_traj.response_text%5B%3A500%5D%0A%20%20%20%20if%20len(_example_traj.response_text)%20%3E%20500%3A%0A%20%20%20%20%20%20%20%20_resp_display%20%2B%3D%20%22...%22%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Example%20Trajectory%0A%0A%20%20%20%20Here's%20what%20a%20single%20generation%20looks%20like%20--%20this%20is%20the%20data%20unit%20flowing%0A%20%20%20%20through%20the%20pipeline%3A%0A%0A%20%20%20%20%7C%20Field%20%7C%20Value%20%7C%0A%20%20%20%20%7C-------%7C-------%7C%0A%20%20%20%20%7C%20Question%20%7C%20%7B_example_traj.task_question%5B%3A100%5D%7D...%20%7C%0A%20%20%20%20%7C%20Correct%20answer%20%7C%20%60%7B_example_traj.task_answer%7D%60%20%7C%0A%20%20%20%20%7C%20Result%20%7C%20**%7B_status%7D**%20(reward%3D%7B_reward%3A.2f%7D)%20%7C%0A%20%20%20%20%7C%20Failure%20mode%20%7C%20%60%7B_example_traj.failure_mode%7D%60%20%7C%0A%20%20%20%20%7C%20Format%20(%60%5BANSWER%5D%60%20tag)%20%7C%20%7B%22Yes%22%20if%20_example_traj.has_answer_tag%20else%20%22No%22%7D%20%7C%0A%20%20%20%20%7C%20Turns%20%7C%20%7B_example_traj.num_turns%7D%20%7C%0A%20%20%20%20%7C%20Tool%20calls%20%7C%20%7B_example_traj.num_tool_calls%7D%20%7C%0A%0A%20%20%20%20**Model%20response**%20(first%20500%20chars)%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20%7B_resp_display%7D%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Each%20generator%20produces%20trajectories%20like%20this%2C%20which%20flow%20into%20the%20replay%20buffer%0A%20%20%20%20for%20the%20trainer%20to%20sample%20from.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(NUM_GENERATORS%2C%20NUM_STEPS%2C%20TimingEvent%2C%20TimingStats%2C%20time)%3A%0A%20%20%20%20def%20run_sync_loop(actors)%20-%3E%20TimingStats%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20SYNC%20MODE%3A%20Broadcast%20generate%20to%20all%20generators%2C%20then%20train.%0A%20%20%20%20%20%20%20%20Pattern%3A%20generate%20batch%20-%3E%20train%20-%3E%20generate%20batch%20-%3E%20train%20...%0A%0A%20%20%20%20%20%20%20%20Uses%20.call()%20to%20broadcast%20generate_trajectory%20to%20all%20generators%0A%20%20%20%20%20%20%20%20simultaneously.%20Each%20generator%20produces%20a%20different%20trajectory%0A%20%20%20%20%20%20%20%20(different%20seed%2C%20stochastic%20sampling)%2C%20but%20the%20call%20is%20synchronous%20--%0A%20%20%20%20%20%20%20%20we%20wait%20for%20ALL%20generators%20to%20finish%20before%20training.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20print(%22%5Cn%22%20%2B%20%22%3D%22%20*%2060)%0A%20%20%20%20%20%20%20%20print(%22SYNC%20MODE%3A%20Broadcast%20Generate%20-%3E%20Train%22)%0A%20%20%20%20%20%20%20%20print(%22%3D%22%20*%2060)%0A%0A%20%20%20%20%20%20%20%20trainer%20%3D%20actors%5B%22trainer%22%5D%0A%20%20%20%20%20%20%20%20generators%20%3D%20actors%5B%22generators%22%5D%0A%0A%20%20%20%20%20%20%20%20stats%20%3D%20TimingStats(%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22SYNC%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20num_generators%3DNUM_GENERATORS%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20num_steps%3DNUM_STEPS%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20total_generations%3D0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20wall_time%3D0%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20baseline%20%3D%200.5%0A%20%20%20%20%20%20%20%20t0%20%3D%20time.perf_counter()%0A%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%20%23%20Generate%20trajectories%20--%20broadcast%20to%20ALL%20generators%0A%20%20%20%20%20%20%20%20%20%20%20%20gen_start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20traj_mesh%20%3D%20generators.generate_trajectory.call().get()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Collect%20into%20a%20plain%20list%20--%20no%20buffer%20in%20sync%20mode.%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20We%20train%20on%20exactly%20what%20we%20just%20generated%2C%20so%20staleness%20is%200.%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20(Async%20mode%20uses%20the%20buffer%20because%20generators%20and%20trainer%20are%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20decoupled%20in%20time%20--%20that's%20where%20the%20buffer%20matters.)%0A%20%20%20%20%20%20%20%20%20%20%20%20batch%20%3D%20list(traj_mesh.values())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20gen_time%20%3D%20time.perf_counter()%20-%20gen_start%0A%20%20%20%20%20%20%20%20%20%20%20%20stats.gen_times.append(gen_time)%0A%20%20%20%20%20%20%20%20%20%20%20%20stats.total_generations%20%2B%3D%20NUM_GENERATORS%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Record%20one%20event%20per%20generator%20(they%20ran%20in%20parallel)%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20gi%20in%20range(NUM_GENERATORS)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.events.append(TimingEvent(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20actor_id%3Df%22Gen%7Bgi%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20event_type%3D%22generate%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20start_time%3Dgen_start%20-%20t0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20duration%3Dgen_time%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Train%20directly%20on%20the%20batch%20we%20just%20generated%20(no%20buffer)%0A%20%20%20%20%20%20%20%20%20%20%20%20train_start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20batch%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20metrics%20%3D%20trainer.train_step.call_one(batch%2C%20baseline).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20baseline%20%3D%200.9%20*%20baseline%20%2B%200.1%20*%20metrics.avg_reward%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Staleness%20should%20be%200%3A%20we%20generated%20with%20current%20policy%20and%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20train%20immediately.%20This%20contrasts%20with%20async%20mode.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_staleness%20%3D%20%5Bmetrics.policy_version%20-%20t.policy_version%20for%20t%20in%20batch%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.staleness.extend(batch_staleness)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Sync%20weights%20to%20all%20generators%20(broadcast)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20handle%2C%20param_meta%2C%20version%2C%20total_bytes%20%3D%20trainer.get_weight_handle.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20handle%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20generators.sync_weights_from_buffer.call(handle%2C%20param_meta%2C%20version%2C%20total_bytes).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state_dict%2C%20ver%20%3D%20trainer.get_state_dict.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20generators.sync_weights.call(state_dict%2C%20ver).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pass%20%20%23%20Non-fatal%3A%20generators%20will%20use%20slightly%20stale%20weights%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20train_time%20%3D%20time.perf_counter()%20-%20train_start%0A%20%20%20%20%20%20%20%20%20%20%20%20stats.train_times.append(train_time)%0A%20%20%20%20%20%20%20%20%20%20%20%20stats.events.append(TimingEvent(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20actor_id%3D%22Train%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20event_type%3D%22train%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20start_time%3Dtrain_start%20-%20t0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20duration%3Dtrain_time%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20correct_count%20%3D%20sum(1%20for%20t%20in%20traj_mesh.values()%20if%20t.is_correct)%0A%20%20%20%20%20%20%20%20%20%20%20%20format_count%20%3D%20sum(1%20for%20t%20in%20traj_mesh.values()%20if%20t.has_answer_tag)%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BSYNC%20%7Bstep%20%2B%201%3A2d%7D%5D%20%7Bcorrect_count%7D%2F%7BNUM_GENERATORS%7D%20correct%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22%7Bformat_count%7D%2F%7BNUM_GENERATORS%7D%20formatted%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22gen%3D%7Bgen_time%20*%201000%3A.0f%7Dms%20train%3D%7Btrain_time%20*%201000%3A.0f%7Dms%22)%0A%0A%20%20%20%20%20%20%20%20stats.wall_time%20%3D%20time.perf_counter()%20-%20t0%0A%20%20%20%20%20%20%20%20return%20stats%0A%0A%20%20%20%20return%20(run_sync_loop%2C)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20BATCH_SIZE%2C%0A%20%20%20%20NUM_GENERATORS%2C%0A%20%20%20%20NUM_STEPS%2C%0A%20%20%20%20TimingEvent%2C%0A%20%20%20%20TimingStats%2C%0A%20%20%20%20threading%2C%0A%20%20%20%20time%2C%0A)%3A%0A%20%20%20%20def%20run_async_loop(actors)%20-%3E%20TimingStats%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20ASYNC%20MODE%3A%20All%20generators%20running%20concurrently%20with%20trainer.%0A%20%20%20%20%20%20%20%20-%201%20thread%20per%20generator%20(each%20uses%20.slice()%20to%20address%20its%20generator)%0A%20%20%20%20%20%20%20%20-%20Training%20in%20main%20thread%0A%20%20%20%20%20%20%20%20-%20Each%20generator%20pulls%20latest%20weights%20before%20each%20trajectory%0A%0A%20%20%20%20%20%20%20%20Uses%20try%2Fexcept%20pattern%20from%20NB03%20for%20fault%20tolerance%20in%20generation%20loops.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20print(%22%5Cn%22%20%2B%20%22%3D%22%20*%2060)%0A%20%20%20%20%20%20%20%20print(f%22ASYNC%20MODE%3A%20%7BNUM_GENERATORS%7D%20Generators%20%2B%201%20Trainer%20(Concurrent)%22)%0A%20%20%20%20%20%20%20%20print(%22%3D%22%20*%2060)%0A%0A%20%20%20%20%20%20%20%20trainer%20%3D%20actors%5B%22trainer%22%5D%0A%20%20%20%20%20%20%20%20buffer%20%3D%20actors%5B%22buffer%22%5D%0A%20%20%20%20%20%20%20%20generators%20%3D%20actors%5B%22generators%22%5D%0A%0A%20%20%20%20%20%20%20%20stats%20%3D%20TimingStats(%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22ASYNC%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20num_generators%3DNUM_GENERATORS%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20num_steps%3DNUM_STEPS%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20total_generations%3D0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20wall_time%3D0%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20lock%20%3D%20threading.Lock()%0A%20%20%20%20%20%20%20%20stop_flag%20%3D%20threading.Event()%0A%20%20%20%20%20%20%20%20t0%20%3D%20time.perf_counter()%0A%0A%20%20%20%20%20%20%20%20def%20generation_loop(gen_idx)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22Each%20generator%20gets%20its%20own%20thread%2C%20using%20.slice()%20for%20individual%20access.%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20gen%20%3D%20generators.slice(procs%3Dgen_idx)%0A%20%20%20%20%20%20%20%20%20%20%20%20while%20not%20stop_flag.is_set()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20gen_start%20%3D%20time.perf_counter()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Pull%20latest%20weights%20before%20generating.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20sync_weights_from_buffer%20short-circuits%20if%20version%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20hasn't%20changed%2C%20so%20this%20is%20cheap%20when%20there's%20nothing%20new.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20handle%2C%20param_meta%2C%20version%2C%20total_bytes%20%3D%20trainer.get_weight_handle.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20handle%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20synced%20%3D%20gen.sync_weights_from_buffer.call_one(handle%2C%20param_meta%2C%20version%2C%20total_bytes).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20synced%3A%0A%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%20with%20lock%3A%0A%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%20%20%20%20%20stats.rdma_syncs%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20state_dict%2C%20ver%20%3D%20trainer.get_state_dict.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20synced%20%3D%20gen.sync_weights.call_one(state_dict%2C%20ver).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20synced%3A%0A%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%20with%20lock%3A%0A%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%20%20%20%20%20stats.direct_syncs%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Generate%20trajectory%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20traj%20%3D%20gen.generate_trajectory.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20buffer.add.call_one(traj).get()%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20gen_time%20%3D%20time.perf_counter()%20-%20gen_start%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20with%20lock%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.gen_times.append(gen_time)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.total_generations%20%2B%3D%201%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20count%20%3D%20stats.total_generations%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.events.append(TimingEvent(%0A%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%20actor_id%3Df%22Gen%7Bgen_idx%7D%22%2C%0A%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%20event_type%3D%22generate%22%2C%0A%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%20start_time%3Dgen_start%20-%20t0%2C%0A%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%20duration%3Dgen_time%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20status%20%3D%20%22correct%22%20if%20traj.is_correct%20else%20traj.failure_mode%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BGEN%7Bgen_idx%7D%20%23%7Bcount%3A2d%7D%5D%20%7Bstatus%7D%20gen%3D%7Bgen_time%20*%201000%3A.0f%7Dms%22)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20except%20Exception%20as%20e%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20try%2Fexcept%20pattern%20from%20NB03%20--%20log%20and%20continue%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BGEN%7Bgen_idx%7D%5D%20Error%3A%20%7Be%7D%2C%20retrying...%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%0A%20%20%20%20%20%20%20%20%23%20Start%201%20thread%20per%20generator%2C%20each%20using%20.slice()%20for%20its%20worker%0A%20%20%20%20%20%20%20%20gen_threads%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20idx%20in%20range(NUM_GENERATORS)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20t%20%3D%20threading.Thread(target%3Dgeneration_loop%2C%20args%3D(idx%2C)%2C%20daemon%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20t.start()%0A%20%20%20%20%20%20%20%20%20%20%20%20gen_threads.append(t)%0A%0A%20%20%20%20%20%20%20%20%23%20Training%20in%20main%20thread%0A%20%20%20%20%20%20%20%20train_steps_done%20%3D%200%0A%20%20%20%20%20%20%20%20baseline%20%3D%200.5%0A%0A%20%20%20%20%20%20%20%20while%20train_steps_done%20%3C%20NUM_STEPS%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Wait%20for%20enough%20samples%0A%20%20%20%20%20%20%20%20%20%20%20%20while%20True%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20size%20%3D%20buffer.size.call_one().get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20size%20%3E%3D%20BATCH_SIZE%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20break%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20time.sleep(0.02)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20train_start%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20%20%20%20%20batch%20%3D%20buffer.sample.call_one(BATCH_SIZE).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20batch%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20metrics%20%3D%20trainer.train_step.call_one(batch%2C%20baseline).get()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20baseline%20%3D%200.9%20*%20baseline%20%2B%200.1%20*%20metrics.avg_reward%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Measure%20staleness%3A%20in%20async%20mode%2C%20trajectories%20may%20have%20been%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20generated%20with%20an%20older%20policy%20version.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20batch_staleness%20%3D%20%5Bmetrics.policy_version%20-%20t.policy_version%20for%20t%20in%20batch%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20with%20lock%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.staleness.extend(batch_staleness)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20train_time%20%3D%20time.perf_counter()%20-%20train_start%0A%20%20%20%20%20%20%20%20%20%20%20%20with%20lock%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.train_times.append(train_time)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stats.events.append(TimingEvent(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20actor_id%3D%22Train%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20event_type%3D%22train%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20start_time%3Dtrain_start%20-%20t0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20duration%3Dtrain_time%2C%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%20train_steps_done%20%2B%3D%201%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20print(f%22%5BTRAIN%20%7Btrain_steps_done%3A2d%7D%5D%20time%3D%7Btrain_time%20*%201000%3A.0f%7Dms%20buffer%3D%7Bsize%7D%22)%0A%0A%20%20%20%20%20%20%20%20stop_flag.set()%0A%0A%20%20%20%20%20%20%20%20for%20t%20in%20gen_threads%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20t.join(timeout%3D2.0)%0A%0A%20%20%20%20%20%20%20%20stats.wall_time%20%3D%20time.perf_counter()%20-%20t0%0A%20%20%20%20%20%20%20%20return%20stats%0A%0A%20%20%20%20return%20(run_async_loop%2C)%0A%0A%0A%40app.cell%0Adef%20_(actors%2C%20run_sync_loop)%3A%0A%20%20%20%20sync_stats%20%3D%20run_sync_loop(actors)%0A%20%20%20%20print(f%22%5CnSync%20complete%3A%20%7Bsync_stats.wall_time%3A.2f%7Ds%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22%7Bsync_stats.total_generations%7D%20generations%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22%7Bsync_stats.gens_per_second%3A.2f%7D%20gens%2Fs%22)%0A%0A%20%20%20%20%23%20Evaluate%20immediately%20after%20sync%20training%0A%20%20%20%20print(%22Evaluating%20post-sync%20performance...%22)%0A%20%20%20%20sync_post_eval%20%3D%20actors%5B%22trainer%22%5D.evaluate_zorplex.call_one(num_samples%3D10%2C%20seed%3D42).get()%0A%20%20%20%20print(f%22Post-sync%20accuracy%3A%20%7Bsync_post_eval%5B'accuracy'%5D%3A.0%25%7D%22)%0A%20%20%20%20return%20sync_post_eval%2C%20sync_stats%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%20Re-initializing%20for%20Async%0A%0A%20%20%20%20To%20compare%20fairly%2C%20we%20tear%20down%20all%20actors%20and%20re-spawn%20from%20scratch%20so%20async%0A%20%20%20%20training%20starts%20from%20the%20same%20untrained%20baseline.%20%60ProcMesh.stop()%60%20releases%0A%20%20%20%20the%20processes%20and%20frees%20GPU%20memory%20before%20we%20spawn%20fresh%20ones.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(actors%2C%20run_async_loop%2C%20setup_actors%2C%20teardown_actors)%3A%0A%20%20%20%20%23%20Tear%20down%20sync%20actors%20to%20free%20GPU%20memory%0A%20%20%20%20teardown_actors(actors)%0A%0A%20%20%20%20%23%20Re-spawn%20everything%20so%20async%20starts%20from%20the%20same%20untrained%20baseline%0A%20%20%20%20print(%22Re-spawning%20actors%20for%20async%20run...%22)%0A%20%20%20%20async_actors%20%3D%20setup_actors()%0A%20%20%20%20print(%22Actors%20re-initialized.%20Starting%20async%20loop...%22)%0A%0A%20%20%20%20async_stats%20%3D%20run_async_loop(async_actors)%0A%20%20%20%20print(f%22%5CnAsync%20complete%3A%20%7Basync_stats.wall_time%3A.2f%7Ds%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22%7Basync_stats.total_generations%7D%20generations%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22%7Basync_stats.gens_per_second%3A.2f%7D%20gens%2Fs%22)%0A%0A%20%20%20%20%23%20Evaluate%20immediately%20after%20async%20training%0A%20%20%20%20print(%22Evaluating%20post-async%20performance...%22)%0A%20%20%20%20async_post_eval%20%3D%20async_actors%5B%22trainer%22%5D.evaluate_zorplex.call_one(num_samples%3D10%2C%20seed%3D42).get()%0A%20%20%20%20print(f%22Post-async%20accuracy%3A%20%7Basync_post_eval%5B'accuracy'%5D%3A.0%25%7D%22)%0A%20%20%20%20return%20async_post_eval%2C%20async_stats%0A%0A%0A%40app.cell%0Adef%20_(async_stats%2C%20mo%2C%20sync_stats)%3A%0A%20%20%20%20def%20_build_comparison(sync_s%2C%20async_s)%20-%3E%20str%3A%0A%20%20%20%20%20%20%20%20speedup%20%3D%20sync_s.wall_time%20%2F%20async_s.wall_time%20if%20async_s.wall_time%20%3E%200%20else%200%0A%20%20%20%20%20%20%20%20gen_ratio%20%3D%20async_s.gens_per_second%20%2F%20sync_s.gens_per_second%20if%20sync_s.gens_per_second%20%3E%200%20else%200%0A%0A%20%20%20%20%20%20%20%20avg_sync_gen%20%3D%20sum(sync_s.gen_times)%20%2F%20len(sync_s.gen_times)%20*%201000%20if%20sync_s.gen_times%20else%200%0A%20%20%20%20%20%20%20%20avg_async_gen%20%3D%20sum(async_s.gen_times)%20%2F%20len(async_s.gen_times)%20*%201000%20if%20async_s.gen_times%20else%200%0A%20%20%20%20%20%20%20%20avg_sync_train%20%3D%20sum(sync_s.train_times)%20%2F%20len(sync_s.train_times)%20*%201000%20if%20sync_s.train_times%20else%200%0A%20%20%20%20%20%20%20%20avg_async_train%20%3D%20sum(async_s.train_times)%20%2F%20len(async_s.train_times)%20*%201000%20if%20async_s.train_times%20else%200%0A%0A%20%20%20%20%20%20%20%20async_syncs%20%3D%20async_s.rdma_syncs%20%2B%20async_s.direct_syncs%0A%0A%20%20%20%20%20%20%20%20return%20f%22%22%22%0A%20%20%20%20%23%23%20Sync%20vs%20Async%20Comparison%0A%0A%20%20%20%20%7C%20Metric%20%7C%20SYNC%20%7C%20ASYNC%20%7C%20Ratio%20%7C%0A%20%20%20%20%7C--------%7C------%7C-------%7C-------%7C%0A%20%20%20%20%7C%20Wall%20time%20%7C%20%7Bsync_s.wall_time%3A.2f%7Ds%20%7C%20%7Basync_s.wall_time%3A.2f%7Ds%20%7C%20**%7Bspeedup%3A.2f%7Dx**%20speedup%20%7C%0A%20%20%20%20%7C%20Generations%20%7C%20%7Bsync_s.total_generations%7D%20%7C%20%7Basync_s.total_generations%7D%20%7C%20%7Basync_s.total_generations%20%2F%20max(sync_s.total_generations%2C%201)%3A.1f%7Dx%20%7C%0A%20%20%20%20%7C%20Gens%2Fsecond%20%7C%20%7Bsync_s.gens_per_second%3A.2f%7D%20%7C%20%7Basync_s.gens_per_second%3A.2f%7D%20%7C%20**%7Bgen_ratio%3A.1f%7Dx**%20throughput%20%7C%0A%20%20%20%20%7C%20Avg%20gen%20time%20%7C%20%7Bavg_sync_gen%3A.0f%7Dms%20%7C%20%7Bavg_async_gen%3A.0f%7Dms%20%7C%20%7C%0A%20%20%20%20%7C%20Avg%20train%20time%20%7C%20%7Bavg_sync_train%3A.0f%7Dms%20%7C%20%7Bavg_async_train%3A.0f%7Dms%20%7C%20%7C%0A%20%20%20%20%7C%20Weight%20syncs%20%7C%20%7Bsync_s.total_generations%7D%20(every%20step)%20%7C%20%7Basync_syncs%7D%20(per-generator)%20%7C%20%7C%0A%0A%20%20%20%20%23%23%23%20Key%20Observations%0A%0A%20%20%20%20-%20**Data%20throughput**%3A%20Async%20collected%20**%7Bgen_ratio%3A.1f%7Dx**%20more%20trajectories%20per%20second.%0A%20%20%20%20%20%20More%20data%20means%20better%20gradient%20estimates.%0A%20%20%20%20-%20**GPU%20utilization**%3A%20In%20sync%20mode%2C%20the%20trainer%20GPU%20sits%20idle%20during%20generation%20and%0A%20%20%20%20%20%20vice%20versa.%20Async%20keeps%20both%20busy.%0A%20%20%20%20-%20**Generators%20ran%20in%20parallel**%3A%20%7Basync_s.num_generators%7D%20generators%20each%20had%20their%20own%0A%20%20%20%20%20%20thread%2C%20producing%20data%20independently.%0A%20%20%20%20-%20The%20trainer%20consumed%20from%20the%20replay%20buffer%20continuously%2C%20never%20waiting%20for%20a%20specific%0A%20%20%20%20%20%20generator%20to%20finish.%0A%0A%20%20%20%20In%20production%20with%20more%20generators%2C%20the%20throughput%20advantage%20grows%20further.%0A%20%20%20%20%22%22%22%0A%0A%20%20%20%20comparison_md%20%3D%20_build_comparison(sync_stats%2C%20async_stats)%0A%20%20%20%20mo.md(comparison_md)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(async_stats%2C%20mo%2C%20sync_stats)%3A%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20matplotlib.patches%20as%20mpatches%0A%0A%20%20%20%20def%20_plot_timeline(stats%2C%20ax%2C%20title)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Plot%20a%20Gantt%20chart%20of%20timing%20events.%22%22%22%0A%20%20%20%20%20%20%20%20color_map%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22generate%22%3A%20%22%234CAF50%22%2C%20%20%23%20green%0A%20%20%20%20%20%20%20%20%20%20%20%20%22train%22%3A%20%22%23E91E63%22%2C%20%20%20%20%20%23%20pink%0A%20%20%20%20%20%20%20%20%20%20%20%20%22sync%22%3A%20%22%239C27B0%22%2C%20%20%20%20%20%20%23%20purple%0A%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20%20%20%20%20%23%20Collect%20unique%20actor%20IDs%20and%20assign%20y%20positions%0A%20%20%20%20%20%20%20%20actor_ids%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20ev%20in%20stats.events%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20ev.actor_id%20not%20in%20actor_ids%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20actor_ids.append(ev.actor_id)%0A%0A%20%20%20%20%20%20%20%20%23%20Sort%3A%20Gen0%2C%20Gen1%2C%20...%2C%20Train%2C%20Sync%0A%20%20%20%20%20%20%20%20gen_ids%20%3D%20sorted(%5Ba%20for%20a%20in%20actor_ids%20if%20a.startswith(%22Gen%22)%5D)%0A%20%20%20%20%20%20%20%20other_ids%20%3D%20%5Ba%20for%20a%20in%20%5B%22Train%22%2C%20%22Sync%22%5D%20if%20a%20in%20actor_ids%5D%0A%20%20%20%20%20%20%20%20actor_ids%20%3D%20gen_ids%20%2B%20other_ids%0A%0A%20%20%20%20%20%20%20%20y_map%20%3D%20%7Baid%3A%20i%20for%20i%2C%20aid%20in%20enumerate(actor_ids)%7D%0A%0A%20%20%20%20%20%20%20%20for%20ev%20in%20stats.events%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20ev.actor_id%20in%20y_map%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20y%20%3D%20y_map%5Bev.actor_id%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%20%3D%20color_map.get(ev.event_type%2C%20%22%23999999%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ax.barh(y%2C%20ev.duration%2C%20left%3Dev.start_time%2C%20height%3D0.6%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%3Dcolor%2C%20alpha%3D0.8%2C%20edgecolor%3D%22white%22%2C%20linewidth%3D0.5)%0A%0A%20%20%20%20%20%20%20%20ax.set_yticks(range(len(actor_ids)))%0A%20%20%20%20%20%20%20%20ax.set_yticklabels(actor_ids)%0A%20%20%20%20%20%20%20%20ax.set_xlabel(%22Wall%20time%20(seconds)%22)%0A%20%20%20%20%20%20%20%20ax.set_title(title)%0A%20%20%20%20%20%20%20%20ax.invert_yaxis()%0A%0A%20%20%20%20fig%2C%20(ax1%2C%20ax2)%20%3D%20plt.subplots(2%2C%201%2C%20figsize%3D(12%2C%206)%2C%20sharex%3DFalse)%0A%0A%20%20%20%20_plot_timeline(sync_stats%2C%20ax1%2C%20f%22SYNC%20(%7Bsync_stats.wall_time%3A.1f%7Ds)%22)%0A%20%20%20%20_plot_timeline(async_stats%2C%20ax2%2C%20f%22ASYNC%20(%7Basync_stats.wall_time%3A.1f%7Ds)%22)%0A%0A%20%20%20%20%23%20Legend%0A%20%20%20%20legend_patches%20%3D%20%5B%0A%20%20%20%20%20%20%20%20mpatches.Patch(color%3D%22%234CAF50%22%2C%20label%3D%22Generate%22)%2C%0A%20%20%20%20%20%20%20%20mpatches.Patch(color%3D%22%23E91E63%22%2C%20label%3D%22Train%22)%2C%0A%20%20%20%20%5D%0A%20%20%20%20fig.legend(handles%3Dlegend_patches%2C%20loc%3D%22upper%20right%22%2C%20framealpha%3D0.9)%0A%0A%20%20%20%20plt.tight_layout()%0A%20%20%20%20_timeline_desc%20%3D%20mo.md(%22%22%22%23%23%23%20Timeline%20Visualization%0A%0A%20%20%20%20The%20Gantt%20charts%20below%20show%20what%20each%20actor%20was%20doing%20over%20time.%20In%20sync%20mode%2C%0A%20%20%20%20bars%20are%20strictly%20sequential%20--%20notice%20the%20gaps%20between%20generation%20and%20training%20bars.%0A%20%20%20%20In%20async%20mode%2C%20generators%20and%20trainer%20overlap%20--%20that%20overlap%20is%20where%20the%20throughput%0A%20%20%20%20gain%20comes%20from.%0A%0A%20%20%20%20**Try%20this%3A**%20Look%20at%20the%20sync%20chart%20and%20count%20the%20idle%20gaps.%20Each%20gap%20is%20wasted%20GPU%0A%20%20%20%20time.%20Then%20look%20at%20the%20async%20chart%20--%20the%20trainer%20bar%20starts%20almost%20immediately%20because%0A%20%20%20%20generators%20are%20pre-filling%20the%20buffer%20concurrently.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20mo.vstack(%5B_timeline_desc%2C%20fig%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(async_stats%2C%20mo%2C%20sync_stats)%3A%0A%20%20%20%20def%20_avg(lst)%3A%0A%20%20%20%20%20%20%20%20return%20sum(lst)%20%2F%20len(lst)%20if%20lst%20else%200.0%0A%0A%20%20%20%20_sync_avg%20%3D%20_avg(sync_stats.staleness)%0A%20%20%20%20_async_avg%20%3D%20_avg(async_stats.staleness)%0A%20%20%20%20_async_max%20%3D%20max(async_stats.staleness)%20if%20async_stats.staleness%20else%200%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Policy%20Staleness%3A%20The%20Cost%20of%20Async%0A%0A%20%20%20%20Async%20mode%20gives%20us%20better%20hardware%20utilization%2C%20but%20there's%20a%20trade-off%3A%0A%20%20%20%20**policy%20staleness**.%20Generators%20produce%20trajectories%20using%20an%20older%20version%0A%20%20%20%20of%20the%20policy%20while%20the%20trainer%20has%20already%20moved%20on.%20This%20is%20*off-policy*%0A%20%20%20%20data%20--%20the%20log-probabilities%20computed%20during%20training%20don't%20match%20the%20policy%0A%20%20%20%20that%20generated%20the%20trajectory.%0A%0A%20%20%20%20We%20measure%20staleness%20as%20%60trainer_version%20-%20trajectory_version%60%20at%20each%0A%20%20%20%20training%20step%3A%0A%0A%20%20%20%20%7C%20Metric%20%7C%20SYNC%20%7C%20ASYNC%20%7C%0A%20%20%20%20%7C--------%7C------%7C-------%7C%0A%20%20%20%20%7C%20Avg%20staleness%20%7C%20%7B_sync_avg%3A.1f%7D%20%7C%20%7B_async_avg%3A.1f%7D%20%7C%0A%20%20%20%20%7C%20Max%20staleness%20%7C%20%7Bmax(sync_stats.staleness)%20if%20sync_stats.staleness%20else%200%7D%20%7C%20%7B_async_max%7D%20%7C%0A%0A%20%20%20%20Sync%20mode%20shows%20~0%20staleness%20because%20we%20sync%20weights%20to%20generators%20after%0A%20%20%20%20every%20training%20step.%20Async%20mode%20shows%20%3E0%20because%20generators%20keep%20producing%0A%20%20%20%20with%20older%20weights%20while%20the%20trainer%20advances.%0A%0A%20%20%20%20With%20REINFORCE%2C%20this%20introduces%20some%20bias.%20More%20sophisticated%20algorithms%0A%20%20%20%20(PPO%2C%20GRPO)%20address%20this%20with%20importance%20sampling%20ratios%20(%60pi_new%20%2F%20pi_old%60)%0A%20%20%20%20and%20clipping%2C%20but%20that's%20beyond%20our%20scope%20here.%20For%20a%20small%20model%20with%20few%0A%20%20%20%20steps%2C%20the%20staleness%20is%20mild%20--%20and%20the%20throughput%20gain%20from%20async%20more%20than%0A%20%20%20%20compensates.%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%20After%20Training%3A%20Did%20It%20Improve%3F%0A%0A%20%20%20%20We%20ran%20sync%20and%20async%20training%20**independently**%20--%20each%20started%20from%20the%20same%0A%20%20%20%20untrained%20model%20(we%20re-spawned%20actors%20between%20runs).%20This%20lets%20us%20compare%0A%20%20%20%20both%20the%20throughput%20characteristics%20(above)%20and%20the%20training%20outcomes.%0A%0A%20%20%20%20Note%3A%20We're%20using%20a%20small%20model%20(0.5B)%20with%20few%20training%20steps%2C%20so%20dramatic%0A%20%20%20%20improvement%20isn't%20guaranteed.%20The%20point%20is%20the%20*infrastructure*%20--%20showing%20that%0A%20%20%20%20the%20full%20loop%20works%20end%20to%20end.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(async_post_eval%2C%20mo%2C%20pre_eval%2C%20sync_post_eval)%3A%0A%20%20%20%20def%20_delta(post%2C%20pre%2C%20key)%3A%0A%20%20%20%20%20%20%20%20return%20post%5Bkey%5D%20-%20pre%5Bkey%5D%0A%0A%20%20%20%20def%20_dir(delta)%3A%0A%20%20%20%20%20%20%20%20if%20delta%20%3E%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%22improved%22%0A%20%20%20%20%20%20%20%20elif%20delta%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20%22unchanged%22%0A%20%20%20%20%20%20%20%20return%20%22decreased%22%0A%0A%20%20%20%20_sync_acc_d%20%3D%20_delta(sync_post_eval%2C%20pre_eval%2C%20%22accuracy%22)%0A%20%20%20%20_async_acc_d%20%3D%20_delta(async_post_eval%2C%20pre_eval%2C%20%22accuracy%22)%0A%20%20%20%20_sync_fmt_d%20%3D%20_delta(sync_post_eval%2C%20pre_eval%2C%20%22format_rate%22)%0A%20%20%20%20_async_fmt_d%20%3D%20_delta(async_post_eval%2C%20pre_eval%2C%20%22format_rate%22)%0A%0A%20%20%20%20_pre_fm%20%3D%20pre_eval%5B%22failure_modes%22%5D%0A%20%20%20%20_sync_fm%20%3D%20sync_post_eval%5B%22failure_modes%22%5D%0A%20%20%20%20_async_fm%20%3D%20async_post_eval%5B%22failure_modes%22%5D%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Training%20Results%3A%20Baseline%20vs%20Sync%20vs%20Async%0A%0A%20%20%20%20Both%20runs%20started%20from%20the%20same%20untrained%20model%20and%20ran%20for%20the%20same%20number%0A%20%20%20%20of%20training%20steps.%0A%0A%20%20%20%20%7C%20Metric%20%7C%20Baseline%20%7C%20After%20Sync%20%7C%20After%20Async%20%7C%0A%20%20%20%20%7C--------%7C----------%7C------------%7C-------------%7C%0A%20%20%20%20%7C%20Accuracy%20%7C%20%7Bpre_eval%5B'accuracy'%5D%3A.0%25%7D%20%7C%20%7Bsync_post_eval%5B'accuracy'%5D%3A.0%25%7D%20(%7B_sync_acc_d%3A%2B.0%25%7D)%20%7C%20%7Basync_post_eval%5B'accuracy'%5D%3A.0%25%7D%20(%7B_async_acc_d%3A%2B.0%25%7D)%20%7C%0A%20%20%20%20%7C%20Format%20compliance%20%7C%20%7Bpre_eval%5B'format_rate'%5D%3A.0%25%7D%20%7C%20%7Bsync_post_eval%5B'format_rate'%5D%3A.0%25%7D%20(%7B_sync_fmt_d%3A%2B.0%25%7D)%20%7C%20%7Basync_post_eval%5B'format_rate'%5D%3A.0%25%7D%20(%7B_async_fmt_d%3A%2B.0%25%7D)%20%7C%0A%20%20%20%20%7C%20Avg%20turns%20%7C%20%7Bpre_eval%5B'avg_turns'%5D%3A.1f%7D%20%7C%20%7Bsync_post_eval%5B'avg_turns'%5D%3A.1f%7D%20%7C%20%7Basync_post_eval%5B'avg_turns'%5D%3A.1f%7D%20%7C%0A%20%20%20%20%7C%20Avg%20tool%20calls%20%7C%20%7Bpre_eval%5B'avg_tools'%5D%3A.1f%7D%20%7C%20%7Bsync_post_eval%5B'avg_tools'%5D%3A.1f%7D%20%7C%20%7Basync_post_eval%5B'avg_tools'%5D%3A.1f%7D%20%7C%0A%0A%20%20%20%20**Failure%20mode%20breakdown%3A**%0A%0A%20%20%20%20%7C%20Mode%20%7C%20Baseline%20%7C%20After%20Sync%20%7C%20After%20Async%20%7C%0A%20%20%20%20%7C------%7C----------%7C------------%7C-------------%7C%0A%20%20%20%20%7C%20Success%20%7C%20%7B_pre_fm%5B'success'%5D%7D%20%7C%20%7B_sync_fm%5B'success'%5D%7D%20%7C%20%7B_async_fm%5B'success'%5D%7D%20%7C%0A%20%20%20%20%7C%20Wrong%20format%20%7C%20%7B_pre_fm%5B'wrong_format'%5D%7D%20%7C%20%7B_sync_fm%5B'wrong_format'%5D%7D%20%7C%20%7B_async_fm%5B'wrong_format'%5D%7D%20%7C%0A%20%20%20%20%7C%20Tool%20spam%20%7C%20%7B_pre_fm%5B'tool_spam'%5D%7D%20%7C%20%7B_sync_fm%5B'tool_spam'%5D%7D%20%7C%20%7B_async_fm%5B'tool_spam'%5D%7D%20%7C%0A%20%20%20%20%7C%20Wrong%20answer%20%7C%20%7B_pre_fm%5B'wrong_answer'%5D%7D%20%7C%20%7B_sync_fm%5B'wrong_answer'%5D%7D%20%7C%20%7B_async_fm%5B'wrong_answer'%5D%7D%20%7C%0A%0A%20%20%20%20Sync%20accuracy%20%7B_dir(_sync_acc_d)%7D%20by%20%7Babs(_sync_acc_d)%3A.0%25%7D.%0A%20%20%20%20Async%20accuracy%20%7B_dir(_async_acc_d)%7D%20by%20%7Babs(_async_acc_d)%3A.0%25%7D.%0A%0A%20%20%20%20With%20a%200.5B%20model%20and%20only%20a%20few%20training%20steps%2C%20large%20gains%20are%20unlikely.%0A%20%20%20%20The%20key%20result%20is%20that%20the%20full%20pipeline%20works%3A%20generation%2C%20training%2C%0A%20%20%20%20weight%20sync%2C%20and%20evaluation%20all%20compose%20correctly%20through%20Monarch%20actors.%20The%0A%20%20%20%20failure%20mode%20breakdown%20shows%20*where*%20the%20model%20is%20improving%20(or%20not)%20--%20watch%0A%20%20%20%20for%20format%20compliance%20changes%20in%20particular%2C%20since%20that's%20the%20easiest%20RL%20win.%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%20What's%20Happening%20Under%20the%20Hood%0A%0A%20%20%20%20When%20you%20run%20the%20training%20loop%2C%20here's%20what%20each%20layer%20does%3A%0A%0A%20%20%20%20**Actor%20isolation**%3A%20Each%20actor%20(trainer%2C%20generators%2C%20buffer%2C%20zorplex%20workers)%0A%20%20%20%20runs%20in%20its%20own%20process%20with%20its%20own%20GPU%20assignment.%20%60CUDA_VISIBLE_DEVICES%60%20is%0A%20%20%20%20set%20in%20%60setup()%60%2C%20not%20at%20spawn%20time%20--%20the%20%60procs%60%20dimension%20in%20%60spawn_procs%60%0A%20%20%20%20is%20just%20a%20dimension%20name%2C%20not%20a%20GPU%20assignment.%0A%0A%20%20%20%20**Weight%20sync%20data%20flow**%20(circular%20buffer%20%2B%20CPU%20staging%20from%20%5BNB07%5D(.%2F07_rdma_weight_sync.html))%3A%0A%20%20%20%20%60%60%60%0A%20%20%20%20Trainer%20GPU%20%20--D2H--%3E%20%20CPU%20slot%5Bv%20%25%203%5D%20%20--RDMA--%3E%20%20Generator%20CPU%20staging%20%20--H2D--%3E%20%20Generator%20GPU%0A%20%20%20%20%60%60%60%0A%20%20%20%20-%20Trainer%20publishes%20weights%20to%20a%20circular%20buffer%20after%20each%20train%20step%0A%20%20%20%20-%20Generators%20pull%20from%20the%20buffer%20via%20RDMA%20into%20a%20CPU%20staging%20buffer%0A%20%20%20%20-%20Explicit%20H2D%20copy%20scatters%20into%20GPU%20model%20parameters%0A%20%20%20%20-%20The%20circular%20buffer%20has%203%20slots%2C%20so%20training%20never%20blocks%20on%20reads%0A%20%20%20%20-%20**Future%20improvement**%3A%20ideally%20we'd%20load%20from%20the%20trainer's%20CPU%20buffer%0A%20%20%20%20%20%20directly%20into%20the%20model's%20%60state_dict%60%2C%20skipping%20the%20staging%20copy.%0A%20%20%20%20%20%20We%20hit%20%60RDMABuffer%60%20bugs%20doing%20that%2C%20so%20for%20now%20we%20use%20the%20extra%20buffer.%0A%0A%20%20%20%20**Async%20concurrency**%20(via%20threads)%3A%0A%20%20%20%20-%201%20thread%20per%20generator%2C%20each%20using%20%60.slice(procs%3Di)%60%20to%20address%20its%20generator%0A%20%20%20%20-%20Each%20generator%20pulls%20latest%20weights%20from%20the%20trainer%20before%20each%20trajectory%0A%20%20%20%20%20%20(%60sync_weights_from_buffer%60%20short-circuits%20if%20version%20hasn't%20changed)%0A%20%20%20%20-%20Training%20in%20the%20main%20thread%0A%20%20%20%20-%20%60threading.Event%60%20coordinates%20shutdown%20when%20training%20completes%0A%20%20%20%20-%20GIL%20is%20released%20during%20I%2FO%20(actor%20calls)%20and%20CUDA%20(GPU%20compute)%2C%20so%20threads%0A%20%20%20%20%20%20achieve%20real%20concurrency%0A%0A%20%20%20%20**Sync%20vs%20Async%20generation**%3A%0A%20%20%20%20-%20Sync%20mode%20uses%20%60.call()%60%20broadcast%20to%20trigger%20all%20generators%20at%20once%2C%20then%20waits%0A%20%20%20%20%20%20for%20all%20to%20finish%20before%20training%0A%20%20%20%20-%20Async%20mode%20uses%20%60.slice()%60%20per%20thread%20so%20each%20generator%20runs%20independently%20--%0A%20%20%20%20%20%20no%20generator%20waits%20for%20another%0A%0A%20%20%20%20**Fault%20tolerance**%20(from%20%5BNB03%5D(.%2F03_fault_tolerance.html))%3A%0A%20%20%20%20-%20Generation%20loops%20wrap%20%60generate_trajectory.call_one().get()%60%20in%20%60try%2Fexcept%60%0A%20%20%20%20-%20On%20failure%2C%20the%20generator%20logs%20and%20retries%20instead%20of%20crashing%20the%20loop%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%20Scaling%20Up%0A%0A%20%20%20%20What%20we%20built%20here%20scales%20naturally%20with%20Monarch%3A%0A%0A%20%20%20%20%7C%20Scale%20%7C%20What%20Changes%20%7C%0A%20%20%20%20%7C-------%7C--------------%7C%0A%20%20%20%20%7C%20More%20generators%20%7C%20Increase%20%60num_generators%60%20slider%20--%20spawns%20larger%20ActorMesh%2C%20%60.call()%60%20broadcast%20scales%20automatically%20%7C%0A%20%20%20%20%7C%20More%20zorplex%20workers%20%7C%20Increase%20%60NUM_ZORPLEX%60%20--%20parallel%20task%20generation%20via%20Service%20%7C%0A%20%20%20%20%7C%20Multi-node%20%7C%20Use%20%60SlurmJob%60%20instead%20of%20%60this_host()%60%20%7C%0A%20%20%20%20%7C%20Better%20algorithms%20%7C%20Swap%20REINFORCE%20for%20PPO%2FGRPO%20(add%20importance%20sampling)%20%7C%0A%20%20%20%20%7C%20Production%20generators%20%7C%20Wrap%20generators%20in%20a%20Service%20too%20(health%20tracking%2C%20auto-scaling)%20%7C%0A%20%20%20%20%7C%20More%20services%20%7C%20Add%20reward%20models%2C%20search%20APIs%20as%20actors%20%7C%0A%0A%20%20%20%20**The%20patterns%20stay%20the%20same%3A**%0A%20%20%20%20-%20Actors%20for%20isolation%20and%20GPU%20assignment%0A%20%20%20%20-%20Endpoints%20for%20communication%20(%60.call_one().get()%60)%0A%20%20%20%20-%20RDMA%20%2B%20circular%20buffer%20for%20efficient%20weight%20transfer%0A%20%20%20%20-%20Version%20tracking%20for%20consistency%20across%20actors%0A%0A%20%20%20%20This%20is%20the%20foundation%20for%20which%20you%20could%20build%20production%20systems%2C%20using%20monarch%20at%20scale.%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%20Recap%3A%20The%20Full%20Journey%0A%0A%20%20%20%20We've%20come%20a%20long%20way%20in%20this%20notebook%20series%3A%0A%0A%20%20%20%20%7C%20Notebook%20%7C%20What%20We%20Learned%20%7C%0A%20%20%20%20%7C----------%7C-----------------%7C%0A%20%20%20%20%7C%2001%20%7C%20Monarch's%20history%20and%20the%20single-controller%20paradigm%20%7C%0A%20%20%20%20%7C%2002%20%7C%20Interactive%20development%20with%20%60this_host()%60%20%7C%0A%20%20%20%20%7C%2003%20%7C%20Fault%20tolerance%20with%20%60try%2Fexcept%60%20on%20actor%20calls%20%7C%0A%20%20%20%20%7C%2004%20%7C%20Distributed%20tensors%20--%20Monarch's%20tensor%20engine%20%7C%0A%20%20%20%20%7C%2005%20%7C%20Zorplex%20benchmark%20--%20where%20Qwen%200.5B%20struggles%20%7C%0A%20%20%20%20%7C%2006%20%7C%20Services%20for%20managing%20worker%20pools%20with%20health%20tracking%20%7C%0A%20%20%20%20%7C%2007%20%7C%20RDMA%20weight%20sync%2C%20circular%20buffers%2C%20CPU%20staging%20%7C%0A%20%20%20%20%7C%20**08**%20%7C%20**Closing%20the%20loop%3A%20async%20RL%20training%20end%20to%20end**%20%7C%0A%0A%20%20%20%20**Key%20takeaways%20from%20this%20notebook%3A**%0A%0A%20%20%20%20-%20Monarch%20makes%20distributed%20RL%20feel%20like%20local%20Python%20--%20actors%2C%20endpoints%2C%0A%20%20%20%20%20%20and%20slicing%20compose%20naturally%20into%20a%20full%20training%20system%0A%20%20%20%20-%20Async%20RL%20collects%20more%20data%20per%20unit%20wall%20time%20by%20running%20generators%0A%20%20%20%20%20%20and%20trainer%20concurrently%0A%20%20%20%20-%20The%20circular%20buffer%20%2B%20CPU%20staging%20pattern%20from%20%5BNB07%5D(.%2F07_rdma_weight_sync.html)%20decouples%20training%0A%20%20%20%20%20%20from%20weight%20distribution%0A%20%20%20%20-%20Before%2Fafter%20evaluation%20closes%20the%20loop%3A%20we%20can%20measure%20whether%20training%0A%20%20%20%20%20%20actually%20improves%20the%20model%0A%0A%20%20%20%20**Where%20to%20go%20next%3A**%20Forge%20GRPO%20implements%20these%20same%20patterns%20at%20production%0A%20%20%20%20scale%20--%20multiple%20nodes%2C%20larger%20models%2C%20PPO%2FGRPO%20instead%20of%20REINFORCE%2C%20and%0A%20%20%20%20proper%20reward%20modeling.%20The%20Monarch%20primitives%20you've%20learned%20here%20are%20the%0A%20%20%20%20building%20blocks%20for%20all%20of%20it.%0A%0A%20%20%20%20---%0A%0A%20%20%20%20**Previous%3A**%20%5BNB07b%20%E2%80%94%20RDMA%20Deep%20Dive%5D(.%2F07b_weight_sync_deep_dive.html)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%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
aa949da010076d0a08a87fc91e9a14b9