Add candidate improvement for explicit working and long memory representation in Omni LoRA

README.md

@@ -1020,6 +1020,16 @@ process uses its `LOCAL_RANK`.  DDP intentionally rejects
 `--online-current-tail`; the online producer-consumer path remains a separate
 benchmark.
 
+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
+adapter and head weights resume, but AdamW moments and the cosine scheduler are
+reinitialized.  Preserve the previous global batch when changing GPU count;
+for example, change 8 GPUs with batch 2 and accumulation 1 to 4 GPUs with batch
+2 and accumulation 2.  Use zero warmup and a reduced learning rate for a short
+post-hoc continuation because a one-epoch cosine schedule has already reached
+zero.
+
 For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU
 producer-consumer benchmark.  The producer GPU runs the frozen AV towers and
 places current-tail embeddings in a bounded CPU queue while the consumer GPU
@@ -1112,6 +1122,43 @@ semantics. Run the following controlled experiments:
 The frozen-Omni head is the cleanest representation test. LoRA can be reported
 as a second-stage ceiling, not as the only result.
 
+## Candidate improvement: explicit working and long memory
+
+Motivated by EgoSpeak's separation of long-term and working memory, test an
+explicit two-timescale representation instead of compressing the entire causal
+window into one undifferentiated final Thinker state.  At prediction time `t`,
+align long memory to fixed global 2-second boundaries and let working memory
+slide with the current prediction time:
+
+```text
+b = floor(t / 2 s) * 2 s
+long memory:     complete 2-second chunks ending at b, ..., [b-4, b-2], [b-2, b]
+working memory:  the sliding recent interval [t-2, t], updated every 0.2 s
+```
+
+For example, at `t = 20.4 s`, long memory contains the fixed chunks
+`[16,18]`, `[18,20]`, and any earlier retained chunks, while the working
+partial is `[18.4,20.4]`.  The overlap between `[18,20]` and `[18.4,20.4]` is
+intentional: long memory provides a stable completed-chunk representation,
+whereas working memory preserves the exact high-resolution context ending at
+the current prediction time.  Cache one representation per completed
+long-memory chunk, recompute the working-memory representation at each 5 Hz
+prediction step, and fuse the two branches with a small gated/attention
+pooling module before the same `10 x 3` prediction head.  Keep the Omni
+backbone frozen in the first test so that any gain can be attributed to memory
+organization rather than additional backbone capacity.
+
+This is a candidate score-improvement idea, not yet an experimental result.
+The key hypothesis is that a dedicated recent branch preserves local prosody,
+speaker-change, and utterance-completion cues while long memory supplies
+conversation context without diluting them in one final state.  Evaluate it
+against matched-compute baselines: (1) the current single-state 8-second model,
+(2) working-memory-only `[t-2,t]`, and (3) long+working memory with the same
+total causal context.  Report parameter count, throughput, and per-horizon and
+per-class AP over identical seeds.  A useful result requires the fused model to
+beat both controls consistently; otherwise a gain could merely come from the
+extra fusion parameters or the stronger recent-window shortcut.
+
 ## Single-window feature inspection
 
 `qwen_omni_features.py` is useful for inspecting a short window or layer, but

train_omni_turn_lora.py

@@ -211,6 +211,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         help="Override target path for a single --cache-dir",
     )
     parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--adapter-checkpoint",
+        type=Path,
+        help=(
+            "Existing PEFT adapter directory to continue from; optimizer and "
+            "scheduler are reinitialized"
+        ),
+    )
     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)
@@ -696,6 +704,7 @@ def _load_initial_head(
 def _prepare_language_lora(
     model: torch.nn.Module,
     *,
+    adapter_checkpoint: Path | None,
     rank: int,
     alpha: int,
     dropout: float,
@@ -704,7 +713,7 @@ def _prepare_language_lora(
     verbose: bool = True,
 ) -> torch.nn.Module:
     try:
-        from peft import LoraConfig, get_peft_model
+        from peft import LoraConfig, PeftModel, get_peft_model
     except ImportError as exc:
         raise RuntimeError(
             "PEFT is required. The ms-swift 3.12.5 image should include it."
@@ -712,6 +721,39 @@ def _prepare_language_lora(
 
     model.requires_grad_(False)
     language_model = model.model
+    if adapter_checkpoint is not None:
+        adapter_checkpoint = adapter_checkpoint.expanduser().resolve()
+        config_path = adapter_checkpoint / "adapter_config.json"
+        if not config_path.is_file():
+            raise FileNotFoundError(
+                f"Missing PEFT adapter config: {config_path}"
+            )
+        language_model = PeftModel.from_pretrained(
+            language_model,
+            str(adapter_checkpoint),
+            is_trainable=True,
+        )
+        language_model.config.use_cache = False
+        if gradient_checkpointing:
+            try:
+                language_model.gradient_checkpointing_enable(
+                    gradient_checkpointing_kwargs={"use_reentrant": False}
+                )
+            except TypeError:
+                language_model.gradient_checkpointing_enable()
+        if verbose:
+            language_model.print_trainable_parameters()
+            trainable_lora = sum(
+                1
+                for name, parameter in language_model.named_parameters()
+                if parameter.requires_grad and "lora_" in name
+            )
+            print(
+                f"Continued trainable LoRA from {adapter_checkpoint}; "
+                f"trainable LoRA tensors={trainable_lora}"
+            )
+        return language_model
+
     matched = [
         name
         for name, module in language_model.named_modules()
@@ -852,6 +894,12 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         raise ValueError("--prefetch-depth must be positive")
     if args.ddp_timeout_minutes <= 0:
         raise ValueError("--ddp-timeout-minutes must be positive")
+    if args.adapter_checkpoint is not None and args.head_checkpoint is None:
+        if run.is_main:
+            print(
+                "Warning: --adapter-checkpoint was provided without "
+                "--head-checkpoint; the turn head will be reinitialized."
+            )
     if (
         args.audio_right_padding_samples is not None
         and args.audio_right_padding_samples < 0
@@ -955,6 +1003,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         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,
@@ -1239,6 +1288,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,
+        "initial_adapter_checkpoint": (
+            str(args.adapter_checkpoint.expanduser().resolve())
+            if args.adapter_checkpoint is not None
+            else None
+        ),
+        "optimizer_state_resumed": False,
         "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir],
         "num_examples": len(dataset),
         "initial_loss": initial_loss,