diff --git a/README.md b/README.md index 193c512..8c935ab 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 7fa60f3..b467729 100644 --- a/train_omni_turn_lora.py +++ b/train_omni_turn_lora.py @@ -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: "