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

Add low-drift LoRA training controls

Browse Source
This commit is contained in:
Xin Wang
2026-08-24 15:24:53 +08:00
parent b7856ece4d
commit 7acab44ea2
2 changed files with 209 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
README.md
View File

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

143
train_omni_turn_lora.py
View File

@@ -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 = (
(raw_head,) if args.freeze_thinker else (raw_language_model, raw_head)
synchronization_modules = tuple(
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