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

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