Add transition-aware R1 and R2 objectives

2026-08-26 17:18:50 +08:00
parent 0590df1e88
commit 5923b89d60
27 changed files with 979 additions and 22 deletions
--- a/train_omni_turn_lora.py
+++ b/train_omni_turn_lora.py
@@ -982,6 +982,7 @@ def _save_weight_checkpoint(
                "num_input_layers": raw_head.num_input_layers,
                "layer_fusion": raw_head.layer_fusion,
                "learnable_layer_mixture": raw_head.learnable_layer_mixture,
+                "transition_auxiliary": raw_head.transition_auxiliary,
            },
            "initial_head_checkpoint": (
                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None