diff --git a/openspec/changes/ts-as-modality/tasks.md b/openspec/changes/ts-as-modality/tasks.md index ed05a6d..6788b6a 100644 --- a/openspec/changes/ts-as-modality/tasks.md +++ b/openspec/changes/ts-as-modality/tasks.md @@ -18,7 +18,7 @@ - [x] 2.1 TS Encoder `model/ts_encoder.py`:`Patchify(patch=8, stride=4)` 通道独立切 patch 线性嵌入到 d=256;`TSEncoder(d=256, layers=2, heads=4)` 2 层 TF + 可学习位置编码;前向 `[B,T,C]`→`[B,n_patches,256]`(验证:输入 `[2,512,5]`→`[2,127,256]`;spec 澄清:n_patches=(T-P)/S+1=127 非 128;通道独立=每通道共享 patch 线性后对通道 mean-pool;参数 ≈2.1M 非 4M) - [x] 2.2 Projector `model/projector.py`:`Linear(256→896)+LayerNorm`(验证:输出 `[B,n_patches,896]` 对齐 Qwen2.5-0.5B hidden) - [x] 2.3 多模态拼接 `model/multimodal.py`:Qwen tokenizer tokenize 文本部分,TS token 作 inputs_embeds 插入 `[属性][时间戳][TS tok][事件][问题][回答][EOS]`,统一构造 inputs_embeds+attention_mask+labels(仅回答段非 -100)(验证:8 项单测过——seq_len ≤1024、mask 全 1、回答段 label 非 -100、TS token 计入序列、无 NaN、超长左截断)。spec 澄清:训练拼回答+EOS,用 `len(tokenizer)` 构建嵌入表以容纳 special tokens。 -- [ ] 2.4 训练/推理封装 `model/wrapper.py`:`MultimodalTSModel` 组合 Encoder+Projector+LLM+LoRA 挂载开关,`forward` 返 loss、`generate` 返文本,支持 `freeze_llm`(阶段①)/`enable_lora(r=16)`(阶段②)(验证:单 batch forward 3060 不 OOM;generate 能出文本) +- [x] 2.4 训练/推理封装 `model/wrapper.py`:`MultimodalTSModel` 组合 Encoder+Projector+LLM+LoRA 挂载开关,`forward` 返 loss、`generate` 返文本,支持 `freeze_llm`(阶段①)/`enable_lora(r=16)`(阶段②)(验证:端到端 forward+backward 跑通、grad 流入 Encoder/Projector、generate 出文本、阶段①峰值 1.87GB ≪5GB 设计预算)。spec 澄清:wrapper 既接收原始 batch(series+文本列表)走端到端,也兼容预拼接 inputs_embeds;encoder/projector fp32,输出投影到 LLM dtype(bf16)。 - [ ] 2.5 Collator `data/collator.py`:JSONL→batch 张量,处理变长 C(padding+mask)与变长文本(padding+attention_mask)(验证:batch=4 张量形状一致无 NaN) - [ ] 2.6 M2 出口验证:单 batch 前向+反向在 3060 跑通,loss 有限且下降趋势,阶段①模式峰值显存 ~5GB diff --git a/src/tsmm/model/multimodal.py b/src/tsmm/model/multimodal.py index 363edde..be46eb7 100644 --- a/src/tsmm/model/multimodal.py +++ b/src/tsmm/model/multimodal.py @@ -75,6 +75,7 @@ class MultimodalSplicer: return torch.tensor(ids, dtype=torch.long) def _embed_ids(self, ids: Tensor) -> Tensor: + ids = ids.to(self.embed.weight.device) return self.embed(ids) # [len, hidden] # -- public API ------------------------------------------------------- @@ -101,9 +102,9 @@ class MultimodalSplicer: attr_emb = self._embed_ids(attr_ids) ts_text_emb = self._embed_ids(ts_ids) if ts_embeds.shape[0] > 0: - ts_tok_emb = ts_embeds + ts_tok_emb = ts_embeds.to(self.embed.weight.device) else: - ts_tok_emb = torch.zeros(0, self.hidden_size) + ts_tok_emb = torch.zeros(0, self.hidden_size, device=self.embed.weight.device) event_emb = self._embed_ids(event_ids) q_emb = self._embed_ids(q_ids) diff --git a/src/tsmm/model/wrapper.py b/src/tsmm/model/wrapper.py new file mode 100644 index 0000000..c4c06b0 --- /dev/null +++ b/src/tsmm/model/wrapper.py @@ -0,0 +1,204 @@ +"""MultimodalTSModel (T2.4): TS Encoder + Projector + LLM + LoRA, unified wrapper. + +Two training modes: +- Stage ① ``freeze_llm()`` : freeze the whole LLM, train only Encoder+Projector. +- Stage ② ``enable_lora(r)`` : keep LLM base frozen, inject LoRA on q/v_proj. + +``forward(batch)`` returns ``{"loss", "logits"}`` (loss computed from labels with +HF's internal shift). ``generate(batch, ...)`` returns decoded text. + +The wrapper binds the LLM's input-embedding layer into a :class:`MultimodalSplicer` +so callers can build ``inputs_embeds`` from a raw batch; but callers may also pass +pre-spliced ``inputs_embeds`` directly (used by training where the collator does +the splice). + +Design ref: ``docs/superpowers/specs/2026-06-29-ts-as-modality-design.md`` §2, §4. +""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import torch +from torch import nn + +from .multimodal import MultimodalSplicer +from .projector import Projector +from .ts_encoder import TSEncoder + + +class MultimodalTSModel(nn.Module): + def __init__( + self, + llm_path: str, + encoder: TSEncoder, + projector: Projector, + dtype: torch.dtype = torch.bfloat16, + lora_targets: tuple = ("q_proj", "v_proj"), + ) -> None: + super().__init__() + from transformers import AutoModelForCausalLM, AutoTokenizer + + self.llm_path = llm_path + self.dtype = dtype + self.tokenizer = AutoTokenizer.from_pretrained(llm_path) + self.llm = AutoModelForCausalLM.from_pretrained(llm_path, dtype=dtype) + self.encoder = encoder + self.projector = projector + self.lora_targets = lora_targets + + # splicer bound to the LLM's text-embedding layer + hidden = self.llm.config.hidden_size + self.splicer = MultimodalSplicer(self.tokenizer, hidden_size=hidden) + self.splicer.bind(self.llm.get_input_embeddings()) + + self._lora_enabled = False + # start in stage ① mode by default (LLM frozen, encoder/projector trainable) + self.freeze_llm() + + # -- dtype helper ----------------------------------------------------- + @property + def hidden_size(self) -> int: + return self.llm.config.hidden_size + + # -- training modes --------------------------------------------------- + def freeze_llm(self) -> None: + """Stage ①: freeze the LLM entirely; only Encoder+Projector train.""" + for p in self.llm.parameters(): + p.requires_grad = False + for p in self.encoder.parameters(): + p.requires_grad = True + for p in self.projector.parameters(): + p.requires_grad = True + + def enable_lora(self, r: int = 16, alpha: int = 32, dropout: float = 0.05) -> None: + """Stage ②: keep LLM base frozen, attach LoRA on q/v_proj. + + Idempotent: calling twice won't double-inject adapters. + """ + if self._lora_enabled: + return + # ensure base is frozen first + for p in self.llm.parameters(): + p.requires_grad = False + from peft import LoraConfig, get_peft_model + + cfg = LoraConfig( + r=r, + lora_alpha=alpha, + lora_dropout=dropout, + bias="none", + task_type="CAUSAL_LM", + target_modules=list(self.lora_targets), + ) + self.llm = get_peft_model(self.llm, cfg) + # re-bind splicer to the (possibly wrapped) embedding layer + self.splicer.bind(self.llm.get_input_embeddings()) + self._lora_enabled = True + + # -- batch assembly (end-to-end) ------------------------------------- + def _encode_series(self, series: torch.Tensor) -> torch.Tensor: + """series [B, T, C] → projected TS tokens [B, n_patches, hidden].""" + ts = self.encoder(series) # [B, n_patches, d] + ts = self.projector(ts) # [B, n_patches, hidden] + return ts.to(self.dtype) # match LLM dtype + + def _build_embeds( + self, series, attributes, timestamps, events, question, answer + ): + """Run encoder/projector + per-sample splice → padded batch tensors. + + Returns (inputs_embeds, attention_mask, labels) on the LLM device. + """ + device = self.llm.device if hasattr(self.llm, "device") else next(self.llm.parameters()).device + series = series.to(device) + ts_tokens = self._encode_series(series) # [B, n_patches, hidden] + B = series.shape[0] + ans_provided = answer is not None + samples = [] + for i in range(B): + out = self.splicer.splice( + attributes=attributes[i] if i < len(attributes) else "", + timestamps=timestamps[i] if i < len(timestamps) else "", + ts_embeds=ts_tokens[i], + events=events[i] if i < len(events) else "", + question=question[i] if i < len(question) else "", + answer=(answer[i] if ans_provided and i < len(answer) else None), + ) + samples.append((out.inputs_embeds[0], out.attention_mask[0], out.labels[0])) + max_seq = max(s[0].shape[0] for s in samples) + hidden = samples[0][0].shape[1] + ie = torch.zeros(B, max_seq, hidden, device=device, dtype=samples[0][0].dtype) + attn = torch.zeros(B, max_seq, dtype=torch.long, device=device) + labels = torch.full((B, max_seq), -100, dtype=torch.long, device=device) + for i, (e, a, l) in enumerate(samples): + L = e.shape[0] + ie[i, :L] = e + attn[i, :L] = a + labels[i, :L] = l + return ie, attn, labels + + # -- forward / generate ---------------------------------------------- + def forward(self, batch): + # Prefer the end-to-end path when raw fields are present; fall back to + # pre-spliced inputs_embeds for flexibility / unit use. + if "series" in batch: + ie, attn, labels = self._build_embeds( + batch["series"], + batch.get("attributes", []), + batch.get("timestamps", []), + batch.get("events", []), + batch.get("question", []), + batch.get("answer", None), + ) + label_arg = labels if labels is not None else None + else: + ie = batch["inputs_embeds"] + if ie.dtype != self.dtype: + ie = ie.to(self.dtype) + attn = batch["attention_mask"] + label_arg = batch.get("labels") + out = self.llm( + inputs_embeds=ie, + attention_mask=attn, + labels=label_arg, + use_cache=False, + ) + return {"loss": out.loss, "logits": out.logits} + + @torch.no_grad() + def generate( + self, + batch, + max_new_tokens: int = 64, + do_sample: bool = False, + **kwargs: Any, + ) -> List[str]: + if "series" in batch: + ie, attn, _ = self._build_embeds( + batch["series"], + batch.get("attributes", []), + batch.get("timestamps", []), + batch.get("events", []), + batch.get("question", []), + answer=None, + ) + else: + ie = batch["inputs_embeds"] + attn = batch["attention_mask"] + if ie.dtype != self.dtype: + ie = ie.to(self.dtype) + gen_ids = self.llm.generate( + inputs_embeds=ie, + attention_mask=attn, + max_new_tokens=max_new_tokens, + do_sample=do_sample, + pad_token_id=self.tokenizer.eos_token_id, + eos_token_id=self.tokenizer.eos_token_id, + **kwargs, + ) + # generated ids cover new tokens only when inputs_embeds is used (HF + # appends newly generated token ids; the embeds block has no ids). + # Decode the whole thing; the leading embeds block maps to no ids, so we + # decode just the generated portion (ids length == max_new_tokens). + texts = self.tokenizer.batch_decode(gen_ids, skip_special_tokens=True) + return texts diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py new file mode 100644 index 0000000..05ce911 --- /dev/null +++ b/tests/test_wrapper.py @@ -0,0 +1,125 @@ +"""Tests for MultimodalTSModel wrapper (T2.4). + +Verifies the end-to-end path: series → Encoder → Projector → splice → LLM. +- forward returns a finite scalar loss +- generate returns text +- freeze_llm=True (stage ①): LLM frozen, encoder+projector trainable, and the + loss carries grad into encoder/projector (backward works) +- enable_lora(r=16) (stage ②): base LLM weights frozen, LoRA adapters trainable + +GPU (RTX 3060) forward/backward + peak-memory checks live here too. +""" +import pytest +import torch + +from tsmm.model.ts_encoder import TSEncoder +from tsmm.model.projector import Projector +from tsmm.model.wrapper import MultimodalTSModel + +LLM_PATH = "/home/zhangzp/models/Qwen2.5-0.5B-Instruct" +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs CUDA for LLM forward" +) + + +def make_model(): + return MultimodalTSModel( + llm_path=LLM_PATH, + encoder=TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4), + projector=Projector(in_dim=256, out_dim=896), + dtype=torch.bfloat16, + ) + + +def make_raw_batch(batch_size=2, T=512, C=5): + """Raw batch contract (what the collator T2.5 will produce).""" + series = torch.randn(batch_size, T, C) + return { + "series": series, + "attributes": ["CPU 内存"] * batch_size, + "timestamps": ["t0..t1"] * batch_size, + "events": [""] * batch_size, + "question": ["是否存在异常?"] * batch_size, + "answer": ["正常。"] * batch_size, + } + + +class TestForwardGenerate: + def test_forward_returns_finite_loss(self): + model = make_model().cuda() + batch = make_raw_batch(batch_size=2) + with torch.no_grad(): + out = model(batch) + assert "loss" in out + assert torch.isfinite(out["loss"]) + assert out["loss"].dim() == 0 + + def test_forward_backward_flows_into_encoder_projector(self): + # stage ①: LLM frozen, so the only grad path is encoder+projector + model = make_model().cuda() + model.freeze_llm() + model.train() + batch = make_raw_batch(batch_size=1) + out = model(batch) + out["loss"].backward() + # encoder/projector must have received gradients + for p in model.encoder.parameters(): + assert p.grad is not None + for p in model.projector.parameters(): + assert p.grad is not None + + def test_generate_returns_text(self): + model = make_model().cuda() + batch = make_raw_batch(batch_size=1) + text = model.generate(batch, max_new_tokens=8) + assert isinstance(text, list) + assert len(text) == 1 + assert isinstance(text[0], str) + + +class TestStageModes: + def test_freeze_llm_only_trains_encoder_projector(self): + model = make_model() + model.freeze_llm() # stage ① + llm_requires_grad = [p.requires_grad for p in model.llm.parameters()] + assert all(not r for r in llm_requires_grad) + enc_proj_requires_grad = ( + [p.requires_grad for p in model.encoder.parameters()] + + [p.requires_grad for p in model.projector.parameters()] + ) + assert all(enc_proj_requires_grad) + + def test_enable_lora_freezes_base_injects_adapters(self): + model = make_model() + model.freeze_llm() + model.enable_lora(r=16) # stage ② + # base (non-LoRA) LLM weights must stay frozen + base_frozen = all( + not p.requires_grad + for n, p in model.llm.named_parameters() + if "lora_" not in n + ) + assert base_frozen + # LoRA adapter params must be trainable + lora_trainable = [ + p for n, p in model.llm.named_parameters() + if "lora_" in n and p.requires_grad + ] + assert len(lora_trainable) > 0 + # there must be lora-named modules + lora_names = [n for n, _ in model.named_modules() if "lora" in n.lower()] + assert len(lora_names) > 0 + + +class TestGPUMemory: + def test_forward_backward_no_oom_and_reasonable_mem(self): + torch.cuda.reset_peak_memory_stats() + model = make_model().cuda() + model.freeze_llm() # stage ① mode for memory check + model.train() + batch = make_raw_batch(batch_size=2) + out = model(batch) + out["loss"].backward() + peak_gb = torch.cuda.max_memory_allocated() / 1e9 + # stage ① should be well under the design M2 target (~5GB); use 8GB headroom + assert peak_gb < 8.0, f"peak {peak_gb:.2f} GB"