Add frozen Thinker sparse head-only control

README.md

@@ -1021,6 +1021,79 @@ process uses its `LOCAL_RANK`.  DDP intentionally rejects
 `--online-current-tail`; the online producer-consumer path remains a separate
 benchmark.
 
+### Frozen-Thinker sparse head-only control
+
+To determine whether the regression comes from transition-sparse sampling or
+from updating the Thinker representation, rerun the same sparse-anchor
+training with `--freeze-thinker`.  This still forwards every reconstructed
+8-second sequence through the original Thinker, but runs that forward under
+`no_grad` and updates only the turn head initialized from the frozen 8-second
+checkpoint.  It does not create an adapter directory.
+
+For a schedule-matched comparison with a one-epoch LoRA run at global batch
+16, use the same cache directories, head learning rate, warmup, seed, and
+approximately `ceil(num_examples / 16)` updates:
+
+```bash
+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_sparse_8s/Session_7 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_11 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_10 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_6 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_9 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_4 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_8 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_5 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_12 \
+  --freeze-thinker \
+  --head-checkpoint checkpoints/easycom_open_stream_2fps_8s_head_seed0.pth \
+  --output-dir checkpoints/easycom_sparse_head_only_from_frozen_8s_ddp4_epoch1 \
+  --max-steps 3390 \
+  --batch-size 2 \
+  --gradient-accumulation-steps 2 \
+  --num-workers 0 \
+  --block-cache-size 32 \
+  --head-learning-rate 1e-4 \
+  --warmup-steps 500 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --log-every 20
+```
+
+Evaluate the resulting head on the same dense test caches without an adapter:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2 \
+OMP_NUM_THREADS=1 \
+MKL_NUM_THREADS=1 \
+OPENBLAS_NUM_THREADS=1 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \
+torchrun --standalone --nproc_per_node=3 eval_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_3 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_2 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_1 \
+  --head-checkpoint checkpoints/easycom_sparse_head_only_from_frozen_8s_ddp4_epoch1/turn_head.pt \
+  --output-dir results/easycom_sparse_head_only_dense_eval_epoch1 \
+  --batch-size 2 \
+  --block-cache-size 32 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --ddp-timeout-minutes 60
+```
+
+If this head-only control also falls well below the frozen `58.16` dense-cache
+control, sparse sampling/head optimization is sufficient to explain the
+regression.  If it stays near `58.16` while LoRA falls, the main failure is
+Thinker representation drift.
+
 To continue from a completed run, pass both its `adapter/` directory through
 `--adapter-checkpoint` and its `turn_head.pt` through `--head-checkpoint` while
 writing to a new output directory.  This is weights-only continuation: the

train_omni_turn_lora.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches.
+"""Train a frozen or LoRA Qwen2.5-Omni Thinker + 10x3 head from AV caches.
 
 This is a custom task runner intended to execute inside the ms-swift image.
 It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI
@@ -219,6 +219,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
             "scheduler are reinitialized"
         ),
     )
+    parser.add_argument(
+        "--freeze-thinker",
+        action="store_true",
+        help=(
+            "Do not create or update LoRA parameters; run the same cached "
+            "Thinker sequences under no_grad and train only the turn head"
+        ),
+    )
     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)
@@ -900,6 +908,14 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                 "Warning: --adapter-checkpoint was provided without "
                 "--head-checkpoint; the turn head will be reinitialized."
             )
+    if args.freeze_thinker and args.adapter_checkpoint is not None:
+        raise ValueError(
+            "--freeze-thinker cannot be combined with --adapter-checkpoint"
+        )
+    if args.freeze_thinker and args.gradient_checkpointing:
+        raise ValueError(
+            "--gradient-checkpointing is unnecessary with --freeze-thinker"
+        )
     if (
         args.audio_right_padding_samples is not None
         and args.audio_right_padding_samples < 0
@@ -1001,16 +1017,27 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     )
     if not target_modules:
         raise ValueError("--lora-target-modules is empty")
-    language_model = _prepare_language_lora(
+    if args.freeze_thinker:
-        base_model,
+        base_model.requires_grad_(False)
-        adapter_checkpoint=args.adapter_checkpoint,
+        language_model = base_model.model
-        rank=args.lora_rank,
+        language_model.config.use_cache = False
-        alpha=args.lora_alpha,
+        language_model.eval()
-        dropout=args.lora_dropout,
+        if run.is_main:
-        target_modules=target_modules,
+            print(
-        gradient_checkpointing=args.gradient_checkpointing,
+                "Frozen-Thinker control: no LoRA adapter; only the turn head "
-        verbose=run.is_main,
+                "is trainable"
-    )
+            )
+    else:
+        language_model = _prepare_language_lora(
+            base_model,
+            adapter_checkpoint=args.adapter_checkpoint,
+            rank=args.lora_rank,
+            alpha=args.lora_alpha,
+            dropout=args.lora_dropout,
+            target_modules=target_modules,
+            gradient_checkpointing=args.gradient_checkpointing,
+            verbose=run.is_main,
+        )
     materialize_batch: OnlineCurrentTailMaterializer | None = None
     producer_model: torch.nn.Module | None = None
     if args.online_current_tail:
@@ -1080,8 +1107,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
 
     raw_language_model = language_model
     raw_head = head
-    _synchronize_trainable_parameters((raw_language_model, raw_head), run)
+    synchronization_modules = (
-    language_model = _wrap_ddp(raw_language_model, run)
+        (raw_head,) if args.freeze_thinker else (raw_language_model, raw_head)
+    )
+    _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)
     # Model initialization is synchronized above; rank-specific RNG streams
     # avoid identical dropout masks for different distributed samples.
@@ -1092,15 +1123,23 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         for parameter in raw_language_model.parameters()
         if parameter.requires_grad
     ]
-    if not lora_parameters:
+    if not args.freeze_thinker and not lora_parameters:
         raise RuntimeError("LoRA preparation produced no trainable parameters")
-    optimizer = torch.optim.AdamW(
+    if args.freeze_thinker:
-        [
+        optimizer = torch.optim.AdamW(
-            {"params": lora_parameters, "lr": args.learning_rate},
+            [{"params": raw_head.parameters(), "lr": args.head_learning_rate}],
-            {"params": raw_head.parameters(), "lr": args.head_learning_rate},
+            weight_decay=args.weight_decay,
-        ],
+        )
-        weight_decay=args.weight_decay,
+        head_group_index = 0
-    )
+    else:
+        optimizer = torch.optim.AdamW(
+            [
+                {"params": lora_parameters, "lr": args.learning_rate},
+                {"params": raw_head.parameters(), "lr": args.head_learning_rate},
+            ],
+            weight_decay=args.weight_decay,
+        )
+        head_group_index = 1
     scheduler = torch.optim.lr_scheduler.LambdaLR(
         optimizer,
         lr_lambda=lambda step: _lr_lambda(
@@ -1122,7 +1161,10 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         print(f"Initial cached loss: {initial_loss:.6f}")
     if materialize_batch is not None:
         materialize_batch.reset_timing()
-    language_model.train()
+    if args.freeze_thinker:
+        language_model.eval()
+    else:
+        language_model.train()
     head.train()
     optimizer.zero_grad(set_to_none=True)
     _distributed_barrier(run)
@@ -1144,7 +1186,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     history: list[dict[str, float | int]] = []
     progress = tqdm(
         total=args.max_steps,
-        desc="Thinker LoRA",
+        desc=("Frozen Thinker head" if args.freeze_thinker else "Thinker LoRA"),
         unit="update",
         disable=args.no_progress or not run.is_main,
     )
@@ -1177,23 +1219,38 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
             )
             with ExitStack() as stack:
                 if run.enabled and not synchronize_gradients:
-                    assert isinstance(language_model, DDP)
                     assert isinstance(head, DDP)
-                    stack.enter_context(language_model.no_sync())
+                    if not args.freeze_thinker:
+                        assert isinstance(language_model, DDP)
+                        stack.enter_context(language_model.no_sync())
                     stack.enter_context(head.no_sync())
-                states = forward_last_state(
+                if args.freeze_thinker:
-                    language_model,
+                    with torch.no_grad():
-                    inputs_embeds=inputs_embeds,
+                        states = forward_last_state(
-                    position_ids=position_ids,
+                            language_model,
-                    attention_mask=attention_mask,
+                            inputs_embeds=inputs_embeds,
-                    readout_indices=readout_indices,
+                            position_ids=position_ids,
-                )
+                            attention_mask=attention_mask,
+                            readout_indices=readout_indices,
+                        )
+                else:
+                    states = forward_last_state(
+                        language_model,
+                        inputs_embeds=inputs_embeds,
+                        position_ids=position_ids,
+                        attention_mask=attention_mask,
+                        readout_indices=readout_indices,
+                    )
                 logits = head(states.float())
                 loss = egospeak_multihorizon_loss(logits, targets)
                 (loss / args.gradient_accumulation_steps).backward()
             accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps
 
-        trainable_parameters = [*lora_parameters, *raw_head.parameters()]
+        trainable_parameters = (
+            list(raw_head.parameters())
+            if args.freeze_thinker
+            else [*lora_parameters, *raw_head.parameters()]
+        )
         grad_norm = torch.nn.utils.clip_grad_norm_(
             trainable_parameters, args.max_grad_norm
         )
@@ -1201,25 +1258,33 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
         reported_loss = _distributed_average(accumulated_loss, run)
+        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"])
         record = {
             "step": update_step,
             "loss": reported_loss,
             "grad_norm": float(grad_norm),
-            "lora_lr": float(optimizer.param_groups[0]["lr"]),
+            "lora_lr": lora_lr,
-            "head_lr": float(optimizer.param_groups[1]["lr"]),
+            "head_lr": head_lr,
         }
         if run.is_main:
             history.append(record)
             progress.update(1)
             progress.set_postfix(
                 loss=f"{reported_loss:.4f}",
-                lr=f"{optimizer.param_groups[0]['lr']:.2g}",
+                lr=f"{head_lr:.2g}" if args.freeze_thinker else f"{lora_lr:.2g}",
             )
         if run.is_main and (update_step == 1 or update_step % args.log_every == 0):
             print(
                 f"step={update_step} loss={reported_loss:.6f} "
                 f"grad_norm={float(grad_norm):.4f} "
-                f"lora_lr={optimizer.param_groups[0]['lr']:.3g}",
+                + (
+                    f"head_lr={head_lr:.3g}"
+                    if args.freeze_thinker
+                    else f"lora_lr={lora_lr:.3g} head_lr={head_lr:.3g}"
+                ),
                 flush=True,
             )
     progress.close()
@@ -1264,7 +1329,8 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
 
     adapter_dir = output_dir / "adapter"
     if run.is_main:
-        raw_language_model.save_pretrained(adapter_dir)
+        if not args.freeze_thinker:
+            raw_language_model.save_pretrained(adapter_dir)
         torch.save(
             {
                 "head": raw_head.state_dict(),
@@ -1288,6 +1354,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         "schema_version": 1,
         "status": "pilot" if args.max_examples else "training",
         "model_id": args.model_id,
+        "model_variant": (
+            "frozen_base_thinker_head_only"
+            if args.freeze_thinker
+            else "thinker_lora"
+        ),
+        "thinker_frozen": bool(args.freeze_thinker),
         "initial_adapter_checkpoint": (
             str(args.adapter_checkpoint.expanduser().resolve())
             if args.adapter_checkpoint is not None
@@ -1351,7 +1423,10 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                     f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
                     f"producer batches={online_timing['producer_batches']}"
                 )
-        print(f"Saved LoRA adapter to {adapter_dir}")
+        if args.freeze_thinker:
+            print("Frozen Thinker: no LoRA adapter was created")
+        else:
+            print(f"Saved LoRA adapter to {adapter_dir}")
         print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
         print(
             "Peak CUDA allocated across ranks: "