# Training a Search-R1 agent with Flame: warm GPU services, gang scheduling, and incremental weight sync **URL:** \ **Category:** reinforcement-learning\ **Created:** [October 20, 2026, 5:28am UTC](https://discuss.pytorch.org/t/training-a-search-r1-agent-with-flame-warm-gpu-services-gang-scheduling-and-incremental-weight-sync/225581 "2026-10-20T05:28:58Z")\ **Posts on this page:** 3\ **Page:** 1
**Author:** ![k82cn](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/k82cn/32/79644_2.png) [@k82cn](https://discuss.pytorch.org/u/k82cn)\ **Post date:** [October 20, 2026, 5:28am UTC](https://discuss.pytorch.org/t/training-a-search-r1-agent-with-flame-warm-gpu-services-gang-scheduling-and-incremental-weight-sync/225581/1 "2026-10-20T05:28:58Z")
I work on [Flame](https://github.com/xflops/flame), an open-source distributed engine for elastic workloads. We added a **Search-R1** -style training example that teaches **Qwen2.5-3B-Instruct** to search a BM25 corpus while answering HotpotQA questions. The training loop runs as Python services: vLLM executors generate trajectories, PyTorch trainer replicas update the policy using group-relative advantages, and Flame’s object cache carries batches and weight updates. The interesting infrastructure pieces: - **Warm executors** : models, retrieval indexes, and optimizer state stay loaded between calls. - **Elastic rollouts** : generation capacity can grow with pending work. - **Gang-scheduled trainers** : the trainer group receives its minimum resources together, avoiding partially scheduled collective participants. - **Incremental weight sync** : after the first full export, existing rollout executors fetch only changed bf16 values. In a 100-step run using 14 rollout GPUs and two trainer GPUs, held-out exact match improved **from 14.5% to 47%** , averaging 24.3 seconds per step after the first two steps. Steady-state weight publication was approximately 165 MB per step, versus a 6.17 GB full model. Caveats: one seed, 200 evaluation questions, and a corpus built from the selected questions’ context paragraphs. These scores aren’t directly comparable with published Search-R1 benchmarks. The synchronous loop also leaves substantial GPU idle time. - [Blog and results](https://xflops.io/blog/training-search-r1-with-flame/) - [Runnable example](https://github.com/xflops/flame/tree/main/examples/rl/search-r1) For others building agent-training systems: how are you handling weight synchronization and the tradeoff between synchronous updates and policy lag? ---
**Author:** ![aicost\_tools](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/b5e925/32.png) [@aicost\_tools](https://discuss.pytorch.org/u/aicost_tools)\ **Post date:** [October 26, 2026, 10:13am UTC](https://discuss.pytorch.org/t/training-a-search-r1-agent-with-flame-warm-gpu-services-gang-scheduling-and-incremental-weight-sync/225581/2 "2026-10-26T10:13:08Z")
The 165 MB vs 6.17 GB weight sync number is the part I’d love to see more of. With 14 rollout GPUs, the idle time in the synchronous loop is probably the biggest line on the bill, more than the sync itself. I work on aicost.tools, so I tend to read setups like this as GPU-hours per step. Have you tried letting rollouts run one or two policy versions behind (with importance-ratio clipping) and measured how much exact match you lose versus the idle time you get back? ---
**Author:** ![drezaptmp](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/drezaptmp/32/80182_2.png) [@drezaptmp](https://discuss.pytorch.org/u/drezaptmp)\ **Post date:** [November 2, 2026, 3:26am UTC](https://discuss.pytorch.org/t/training-a-search-r1-agent-with-flame-warm-gpu-services-gang-scheduling-and-incremental-weight-sync/225581/3 "2026-11-02T03:26:16Z")
and then quick question, what questions are we training them on? And, I’m parsing the blog post and repo, what questions are we training them on? Sound like linguistic questions? If it’s a commonly used dataset, aren’t we just overfitting these models to that dataset? Also, in a linguistic setting, couldn’t “non-exact” matches of text still be correct answers? Like a written response on the SAT/ACT, or something. I’m looking for the set of training and test questions here: [GitHub - xflops/flame: A distributed engine for elastic workload, e.g. AI Agent, RL, Quant · GitHub](https://github.com/xflops/flame), can you point me to the directory?