diff --git a/README.md b/README.md index 6fe665c..aea49eb 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index b467729..6b359f1 100644 --- a/train_omni_turn_lora.py +++ b/train_omni_turn_lora.py @@ -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