Add low-drift LoRA training controls

README.md

@@ -1077,6 +1077,83 @@ process uses its `LOCAL_RANK`.  DDP intentionally rejects
 `--online-current-tail`; the online producer-consumer path remains a separate
 benchmark.
 
+### Dense low-drift LoRA probe
+
+The first follow-up to the sparse-LoRA regression should remove sparse
+supervision as a variable and constrain adaptation near the Thinker output.
+Use dense 5-Hz caches, reserve `Session_12` as development data, and train on
+`Session_4` through `Session_11`.  `--lora-last-n-layers 8` discovers the
+language-layer indices and stores PEFT `layers_to_transform=[28,...,35]` for
+the 36-layer Qwen2.5-Omni-3B Thinker.  Combining it with
+`--lora-target-modules q_proj,v_proj,o_proj` reduces the adapter from all 252
+attention/MLP projections to 24 output-side attention projections.
+
+Start with the successful frozen 8-second head fixed.  `--freeze-head`
+requires `--head-checkpoint`, excludes the head from DDP and AdamW, but still
+backpropagates through it into the Thinker LoRA.  This isolates representation
+adaptation from head decision-boundary drift.  Do not continue from a regressed
+adapter: create this adapter from the original Thinker.
+
+For four GPUs with global batch 16, compute exactly one epoch from the dense
+manifest anchor counts and run:
+
+```bash
+cd /data1/yfzx/wangx/Code/icassp2027-experiment
+
+MAX_STEPS=$(python - <<'PY'
+import json
+import math
+from pathlib import Path
+
+root = Path("/data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s")
+sessions = (4, 5, 6, 7, 8, 9, 10, 11)
+examples = sum(len(json.loads((root / f"Session_{s}" / "manifest.json").read_text())["anchors"])
+               for s in sessions)
+print(math.ceil(examples / 16))
+PY
+)
+WARMUP_STEPS=$(( (MAX_STEPS + 19) / 20 ))
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+OMP_NUM_THREADS=2 \
+MKL_NUM_THREADS=2 \
+OPENBLAS_NUM_THREADS=2 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+torchrun --standalone --nproc_per_node=4 train_omni_turn_lora.py \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_4 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_5 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_6 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_7 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_8 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_9 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_10 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_11 \
+  --head-checkpoint checkpoints/easycom_open_stream_2fps_8s_head_seed0.pth \
+  --freeze-head \
+  --lora-target-modules q_proj,v_proj,o_proj \
+  --lora-last-n-layers 8 \
+  --lora-rank 8 \
+  --lora-alpha 16 \
+  --learning-rate 5e-6 \
+  --output-dir checkpoints/omni_turn_lora_dense_8s_last8_qvo_frozen_head_epoch1 \
+  --max-steps "${MAX_STEPS}" \
+  --batch-size 2 \
+  --gradient-accumulation-steps 2 \
+  --num-workers 0 \
+  --block-cache-size 32 \
+  --warmup-steps "${WARMUP_STEPS}" \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --log-every 20
+```
+
+Select this recipe on dense `Session_12` using 0.2-second mAP as the primary
+metric and mean mAP as the tie-breaker.  Only after the recipe is fixed should
+it be evaluated on official test `Sessions_1--3`.  A second controlled run may
+omit `--freeze-head` and use `--head-learning-rate 1e-5`; all other settings
+must remain identical.
+
 ### Frozen-Thinker sparse head-only control
 
 To determine whether the regression comes from transition-sparse sampling or

train_omni_turn_lora.py

@@ -22,6 +22,7 @@ import math
 import os
 import queue
 import random
+import re
 import threading
 import time
 from contextlib import ExitStack
@@ -70,6 +71,8 @@ DEFAULT_LORA_TARGETS = (
     "down_proj",
 )
 
+_TRANSFORMER_LAYER_PATTERN = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
+
 
 @dataclass(frozen=True)
 class DistributedRun:
@@ -227,6 +230,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "Thinker sequences under no_grad and train only the turn head"
         ),
     )
+    parser.add_argument(
+        "--freeze-head",
+        action="store_true",
+        help=(
+            "Keep the initialized turn head fixed and update only Thinker "
+            "LoRA parameters; requires --head-checkpoint"
+        ),
+    )
     parser.add_argument("--head-checkpoint", type=Path)
     parser.add_argument("--max-examples", type=int, default=0)
     parser.add_argument("--max-steps", type=int, default=200)
@@ -300,6 +311,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         "--lora-target-modules",
         default=",".join(DEFAULT_LORA_TARGETS),
     )
+    parser.add_argument(
+        "--lora-last-n-layers",
+        type=int,
+        default=0,
+        help=(
+            "Apply a newly created LoRA adapter only to the final N Thinker "
+            "language layers; 0 keeps all matched layers"
+        ),
+    )
     parser.add_argument("--head-dropout", type=float, default=0.1)
     parser.add_argument("--gradient-checkpointing", action="store_true")
     parser.add_argument("--seed", type=int, default=1337)
@@ -717,6 +737,7 @@ def _prepare_language_lora(
     alpha: int,
     dropout: float,
     target_modules: Sequence[str],
+    last_n_layers: int,
     gradient_checkpointing: bool,
     verbose: bool = True,
 ) -> torch.nn.Module:
@@ -762,22 +783,62 @@ def _prepare_language_lora(
             )
         return language_model
 
-    matched = [
+    matched_all = [
         name
         for name, module in language_model.named_modules()
         if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules
     ]
-    if not matched:
+    if not matched_all:
         raise RuntimeError(
             "No language-model linear layers matched LoRA targets "
             f"{list(target_modules)}"
         )
+    layer_indices = sorted(
+        {
+            int(match.group(1))
+            for name in matched_all
+            if (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
+        }
+    )
+    layers_to_transform: list[int] | None = None
+    matched = matched_all
+    if last_n_layers:
+        if not layer_indices:
+            raise RuntimeError(
+                "Could not identify Thinker layer indices in matched LoRA modules"
+            )
+        if last_n_layers > len(layer_indices):
+            raise ValueError(
+                f"--lora-last-n-layers={last_n_layers} exceeds the "
+                f"{len(layer_indices)} matched Thinker layers"
+            )
+        layers_to_transform = layer_indices[-last_n_layers:]
+        selected = set(layers_to_transform)
+        matched = [
+            name
+            for name in matched_all
+            if (
+                (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
+                and int(match.group(1)) in selected
+            )
+        ]
+        if not matched:
+            raise RuntimeError(
+                "Final-layer selection removed every matched LoRA module"
+            )
+    config_kwargs: dict[str, Any] = {}
+    if layers_to_transform is not None:
+        config_kwargs.update(
+            layers_to_transform=layers_to_transform,
+            layers_pattern="layers",
+        )
     config = LoraConfig(
         r=rank,
         lora_alpha=alpha,
         lora_dropout=dropout,
         target_modules=list(target_modules),
         bias="none",
+        **config_kwargs,
     )
     language_model = get_peft_model(language_model, config)
     language_model.config.use_cache = False
@@ -794,6 +855,12 @@ def _prepare_language_lora(
             f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
             f"last={matched[-1]}"
         )
+        if layers_to_transform is not None:
+            print(
+                "LoRA layer restriction: "
+                f"{layers_to_transform[0]}--{layers_to_transform[-1]} "
+                f"({len(layers_to_transform)} of {len(layer_indices)} layers)"
+            )
     return language_model
 
 
@@ -912,6 +979,23 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         raise ValueError(
             "--freeze-thinker cannot be combined with --adapter-checkpoint"
         )
+    if args.freeze_thinker and args.freeze_head:
+        raise ValueError(
+            "--freeze-thinker and --freeze-head would leave no trainable parameters"
+        )
+    if args.freeze_head and args.head_checkpoint is None:
+        raise ValueError("--freeze-head requires --head-checkpoint")
+    if args.lora_last_n_layers < 0:
+        raise ValueError("--lora-last-n-layers must be non-negative")
+    if args.adapter_checkpoint is not None and args.lora_last_n_layers:
+        raise ValueError(
+            "Layer selection is stored in an existing adapter config; omit "
+            "--lora-last-n-layers when using --adapter-checkpoint"
+        )
+    if args.freeze_thinker and args.lora_last_n_layers:
+        raise ValueError(
+            "--lora-last-n-layers has no effect with --freeze-thinker"
+        )
     if args.freeze_thinker and args.gradient_checkpointing:
         raise ValueError(
             "--gradient-checkpointing is unnecessary with --freeze-thinker"
@@ -1035,6 +1119,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             alpha=args.lora_alpha,
             dropout=args.lora_dropout,
             target_modules=target_modules,
+            last_n_layers=args.lora_last_n_layers,
             gradient_checkpointing=args.gradient_checkpointing,
             verbose=run.is_main,
         )
@@ -1104,16 +1189,27 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         dropout=args.head_dropout,
         verbose=run.is_main,
     ).to(device=device, dtype=torch.float32)
+    if args.freeze_head:
+        head.requires_grad_(False)
+        head.eval()
+        if run.is_main:
+            print("Frozen-head control: only Thinker LoRA parameters are trainable")
 
     raw_language_model = language_model
     raw_head = head
-    synchronization_modules = (
+    synchronization_modules = tuple(
-        (raw_head,) if args.freeze_thinker else (raw_language_model, raw_head)
+        module
+        for module, trainable in (
+            (raw_language_model, not args.freeze_thinker),
+            (raw_head, not args.freeze_head),
+        )
+        if trainable
     )
     _synchronize_trainable_parameters(synchronization_modules, run)
     if not args.freeze_thinker:
         language_model = _wrap_ddp(raw_language_model, run)
-    head = _wrap_ddp(raw_head, run)
+    if not args.freeze_head:
+        head = _wrap_ddp(raw_head, run)
     # Model initialization is synchronized above; rank-specific RNG streams
     # avoid identical dropout masks for different distributed samples.
     _seed_everything(args.seed + run.rank)
@@ -1131,6 +1227,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             weight_decay=args.weight_decay,
         )
         head_group_index = 0
+    elif args.freeze_head:
+        optimizer = torch.optim.AdamW(
+            [{"params": lora_parameters, "lr": args.learning_rate}],
+            weight_decay=args.weight_decay,
+        )
+        head_group_index = None
     else:
         optimizer = torch.optim.AdamW(
             [
@@ -1165,7 +1267,10 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         language_model.eval()
     else:
         language_model.train()
-    head.train()
+    if args.freeze_head:
+        head.eval()
+    else:
+        head.train()
     optimizer.zero_grad(set_to_none=True)
     _distributed_barrier(run)
     torch.cuda.synchronize(device)
@@ -1219,11 +1324,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             )
             with ExitStack() as stack:
                 if run.enabled and not synchronize_gradients:
-                    assert isinstance(head, DDP)
                     if not args.freeze_thinker:
                         assert isinstance(language_model, DDP)
                         stack.enter_context(language_model.no_sync())
-                    stack.enter_context(head.no_sync())
+                    if not args.freeze_head:
+                        assert isinstance(head, DDP)
+                        stack.enter_context(head.no_sync())
                 if args.freeze_thinker:
                     with torch.no_grad():
                         states = forward_last_state(
@@ -1249,7 +1355,11 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         trainable_parameters = (
             list(raw_head.parameters())
             if args.freeze_thinker
-            else [*lora_parameters, *raw_head.parameters()]
+            else (
+                lora_parameters
+                if args.freeze_head
+                else [*lora_parameters, *raw_head.parameters()]
+            )
         )
         grad_norm = torch.nn.utils.clip_grad_norm_(
             trainable_parameters, args.max_grad_norm
@@ -1261,7 +1371,11 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         lora_lr = (
             0.0 if args.freeze_thinker else float(optimizer.param_groups[0]["lr"])
         )
-        head_lr = float(optimizer.param_groups[head_group_index]["lr"])
+        head_lr = (
+            0.0
+            if head_group_index is None
+            else float(optimizer.param_groups[head_group_index]["lr"])
+        )
         record = {
             "step": update_step,
             "loss": reported_loss,
@@ -1343,6 +1457,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                 "initial_head_checkpoint": (
                     str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
                 ),
+                "head_frozen": bool(args.freeze_head),
                 "training_step": args.max_steps,
             },
             output_dir / "turn_head.pt",
@@ -1357,9 +1472,15 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         "model_variant": (
             "frozen_base_thinker_head_only"
             if args.freeze_thinker
-            else "thinker_lora"
+            else (
+                "thinker_lora_frozen_head"
+                if args.freeze_head
+                else "thinker_lora"
+            )
         ),
         "thinker_frozen": bool(args.freeze_thinker),
+        "head_frozen": bool(args.freeze_head),
+        "lora_last_n_layers": int(args.lora_last_n_layers),
         "initial_adapter_checkpoint": (
             str(args.adapter_checkpoint.expanduser().resolve())
             if args.adapter_checkpoint is not None