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
train_omni_turn_lora.py
+114 -39
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,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(
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,
)
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,
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)
language_model = _wrap_ddp(raw_language_model, 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
# 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(
[
{"params": lora_parameters, "lr": args.learning_rate},
{"params": raw_head.parameters(), "lr": args.head_learning_rate},
],
weight_decay=args.weight_decay,
)
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},
{"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(
language_model,
inputs_embeds=inputs_embeds,
position_ids=position_ids,
attention_mask=attention_mask,
readout_indices=readout_indices,
)
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,
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"]),
"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,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: "