Modular Post-Training

BAR (Branch-Adapt-Route) is a modular post-training recipe that trains domain-specific experts each through an independent pipeline all the way to completion, then merges them into an MoE to be used as a single model. Each expert can be independently upgraded or replaced, without needing to rerun one giant pipeline.

With the existing FlexOlmo approach of freezing all shared layers, things work well for pretraining, but it can't keep up with the behavioral changes needed during post-training. In fact, when training with shared parameters frozen under RLVR, the reward curve was nearly flat, which is why a post-training-specific design was needed.

Stage 1: independent expert training

• Each domain expert starts as a 2-expert MoE consisting of a frozen anchor expert and a trainable expert.
• During mid-training, all shared layers are frozen.
• During SFT, the embedding layer and LM head are unfrozen. For tool use, which requires new special tokens, the BFCL score rose from 20.3 → 46.4.
• During RLVR, all shared parameters, including attention, are unfrozen.
• Domain-only SFT boosts in-domain performance but significantly damages general instruction following and knowledge, so mixing in general SFT data is essential.

Stage 2: expert merging Once training is complete, the experts are merged into a single MoE, and shared parameters that diverged during SFT/RLVR are simply averaged. This averaging produced almost no loss in domain-specific evaluations.

Stage 3: router training The router is trained while all other experts and shared weights are frozen, and a stratified 5% SFT sample was sufficient. This makes this stage fast and cheap.

After training math, code, tool use, and safety experts on top of a 7B-class Olmo 2 base, BAR outperformed all baselines that don't rerun mid-training. The overall average was higher than post-training-only retraining, at 49.1 vs 47.8, with particularly notable improvements of +7.8 in math and +4.7 in code. On the other hand, dense model averaging after mid-training was nearly broken, reaching only 6.5, and BTX also scored 46.7, lower than BAR's 49.1. Full retraining still holds the top score of 50.5, but it requires the original pretraining checkpoint and full-scale reprocessing, making it less practical.

Practicality was also confirmed. Replacing the code expert with better data and RL raised the combined model's code performance by +16.5, and adding RL to the math expert improved it by +13. Since only the changed expert and router need to be retrained, BAR can iterate on domain updates repeatedly at a much lower cost than monolithic retraining.

The key point is that post-training needs to be far more flexible than pretraining. Progressive unfreezing, which unfreezes shared layers differently at each stage, data mixing that avoids domain-only SFT, and the combination of averaging and routing are what make modular post-training actually work.