wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add frozen Thinker sparse head-only control

Browse Source
This commit is contained in:
Xin Wang
2026-08-24 11:32:03 +08:00
parent f6dbce30e4
commit 43cf1f386e
2 changed files with 187 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
README.md
View File

@@ -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

95
train_omni_turn_lora.py
View File

@@ -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,6 +1017,17 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
)
if not target_modules:
raise ValueError("--lora-target-modules is empty")
if args.freeze_thinker:
base_model.requires_grad_(False)
language_model = base_model.model
language_model.config.use_cache = False
language_model.eval()
if run.is_main:
print(
"Frozen-Thinker control: no LoRA adapter; only the turn head "
"is trainable"
)
else:
language_model = _prepare_language_lora(
base_model,
adapter_checkpoint=args.adapter_checkpoint,
@@ -1080,7 +1107,11 @@ 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 = (
(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
@@ -1092,8 +1123,15 @@ 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")
if args.freeze_thinker:
optimizer = torch.optim.AdamW(
[{"params": raw_head.parameters(), "lr": args.head_learning_rate}],
weight_decay=args.weight_decay,
)
head_group_index = 0
else:
optimizer = torch.optim.AdamW(
[
{"params": lora_parameters, "lr": args.learning_rate},
@@ -1101,6 +1139,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
],
weight_decay=args.weight_decay,
)
head_group_index = 1
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda step: _lr_lambda(
@@ -1122,6 +1161,9 @@ 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()
if args.freeze_thinker:
language_model.eval()
else:
language_model.train()
head.train()
optimizer.zero_grad(set_to_none=True)
@@ -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,10 +1219,21 @@ 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)
if not args.freeze_thinker:
assert isinstance(language_model, DDP)
stack.enter_context(language_model.no_sync())
stack.enter_context(head.no_sync())
if args.freeze_thinker:
with torch.no_grad():
states = forward_last_state(
language_model,
inputs_embeds=inputs_embeds,
position_ids=position_ids,
attention_mask=attention_mask,
readout_indices=readout_indices,
)
else:
states = forward_last_state(
language_model,
inputs_embeds=inputs_embeds,
@@ -1193,7 +1246,11 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
(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"]),
"head_lr": float(optimizer.param_groups[1]["lr"]),
"lora_lr": lora_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,6 +1329,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
adapter_dir = output_dir / "adapter"
if run.is_main:
if not args.freeze_thinker:
raw_language_model.save_pretrained(adapter_dir)
torch.save(
{
@@ -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,6 +1423,9 @@ 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']}"
)
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(