多模态——从 0 训练简化版 Omni 模型

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

seeker-omni/
└─ data/
   ├─ raw/
   │  ├─ minimind/
   │  │  ├─ pretrain_hq.jsonl
   │  │  └─ sft_mini_512.jsonl
   │  └─ flickr8k/
   │     ├─ flickr8kzhc.caption.txt
   │     ├─ Flickr8k_Dataset/
   │     │  ├─ *.jpg
   │     │  └─ ...
   │     └─ text/
   │        ├─ Flickr_8k.trainImages.txt
   │        └─ ...
   └─ interim/
      ├─ tokenizer_corpus/
      │  └─ minimind_pretrain_text.txt
      ├─ sft_converted/
      │  └─ minimind_sft_chatml.jsonl
      └─ packs/
         └─ mm/
            └─ train_imgonly.jsonl

1

我知道底边长度和高，分别是5cm和8cm，请告诉我计算出来的面积是多少。三角形的面积为20平方厘米。<|im_end|> 为什么理智告诉我们不要跟陌生人接触？因为陌生人可能会对我们的个人安全造成威胁、危害或诱骗等恶劣行为。<|im_end|> 上一次科比在NBA总决赛中获得总冠军是哪一年？科比在NBA总决赛中最后一次获得总冠军是在2010年。<|im_end|> 写一首五言诗，以“秋天”为主题。秋叶飘飘降，秋风肆虐长街，落叶纷纷掉下，百鸟南飞不归，唯有寒蝉孤嘶。<|im_end|> 好的。那你可以跟我聊一下最近的天气吗？当然可以。最近的天气因地区而异，你想知道哪个城市的天气情况呢？<|im_end|> 好的，那么请你接着把这个句子进一步改写，加入一个介绍该书的短语。这本畅销书是由名著名文学家所著写的。<|im_end|> 从以下描述中归纳出这是哪种动物：该动物四肢行走，有毛发，是哺乳动物。这是哺乳动物，具体种类无法确定。<|im_end|> 请问这周末的天气怎么样？我需要您告诉我您所在的城市地区，我可以通过查询天气API来提供相关天气信息。<|im_end|>

1

{"id": "line-676605", "conversations": [{"role": "system", "content": "你是一个知识丰富的AI，尽力为用户提供准确的信息。"}, {"role": "user", "content": "请向我解释在 GitHub 的公共代码库中使用的 MIT 许可证。"}, {"role": "assistant", "content": "MIT 许可证是一种广泛使用的开源软件许可证，允许用户自由地使用、修改和分发软件。它通常用于 GitHub 上的公共存储库，因为它提供了一个简单和宽松的许可证，允许其他人在不需要担心严格要求或复杂的法律条款的情况下使用和建立在该软件之上。该许可证还要求用户保留软件中的版权和许可证通知，这有助于确保该软件的原始创建者获得信誉，并且其他人知道可以使用该软件的条款。总的来说，MIT 许可证为他人提供开源软件的简单有效方式，同时保护原始创建者的权利。"}, {"role": "user", "content": "你觉得 MIT 许可证好吗？"}, {"role": "assistant", "content": "MIT许可证的好处最终取决于使用它的个人或组织的具体情况和目标。一般而言，MIT许可证被认为是一种非常宽松和简单的开源许可证，允许用户自由使用、修改和分发软件，而不需要其他开源许可证中存在的严格要求和法律复杂性。这使它成为许多开源项目的好选择，特别是那些专注于协作和社区参与的项目。"}]}

1

{"id": "flickr8k-2208631481_3e4a5675e1-zhc-0", "system": "你是一个只用中文回答的助手。", "prompt": "请描述这张图片。", "answer": "在草地上赛跑的两只黑狗。", "image": "data/raw/flickr8k/Flickr8k_Dataset/2208631481_3e4a5675e1.jpg"}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

# dataprep/prepare/tokenizer.py
import json
import random
import shutil
import time
from pathlib import Path

from tokenizers import Tokenizer

from .text_bpe import train_text_bpe
from ..data_paths import DATA_INTERIM, MINIMIND_TEXT_CORPUS, TOKENIZER_DIR, TOKENIZER_VOCAB_SIZE

MINIMIND2_CHATML_TOKENS = [
    "<|endoftext|>",
    "<|im_start|>",
    "<|im_end|>",
    "<img_bos>",
    "<img>",
    "<img_eos>",
]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

# dataprep/data_paths.py
from pathlib import Path

DATA_RAW = Path("data/raw")
DATA_INTERIM = Path("data/interim")
DATA_PROCESSED = Path("data/processed")
ARTIFACTS = Path("artifacts")

# 默认设置常量
SEED = 42
OVERWRITE = False

# (此处省略中间各种有关 Flickr8k 图文和 minimind 基础语料的大量源地址定义和图文的交接文件配置)
...

# 分词器语料与默认导出目录
MINIMIND_TEXT_CORPUS = DATA_INTERIM / "tokenizer_corpus" / "minimind_pretrain_text.txt"
TOKENIZER_VOCAB_SIZE = 6400
TOKENIZER_DIR = ARTIFACTS / "tokenizers" / "bpe_m2chatml_6400"

def default_dataprep_cfg():
    """返回 dataprep 默认配置 dict"""
    return {
        "seed": int(SEED),
        "overwrite": bool(OVERWRITE),
        "tokenizer": {
            "sample_ratio": 0.6,
            "sample_seed": int(SEED),
        },
        # 其他限流参数与占位配置省略...
    }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

# dataprep/prepare/text_bpe.py
import json
from pathlib import Path

from tokenizers import Tokenizer
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.trainers import BpeTrainer

# 此处省略 ChatML Jinja 模板字符串，完整内容参考配套代码
_MINIMIND_CHAT_TEMPLATE = """{%- if tools %}..."""

def _infer_core_tokens(special_tokens):
    """从 special_tokens 推断 scheme 名称与 pad/bos/eos/unk 的字符串形式。
    
    当前 dataprep 只支持 minimind2_chatml。
    """
    s = set(map(str, special_tokens))
    required = {"<|endoftext|>", "<|im_start|>", "<|im_end|>"}
    if not required.issubset(s):
        raise ValueError(
            "special_tokens must include minimind2_chatml core tokens: "
            "<|endoftext|>, <|im_start|>, <|im_end|>"
        )
    return "minimind2_chatml", "<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|endoftext|>"

def _write_tokenizer_config(
    *, out_dir, tok, special_tokens,
    scheme_name, pad_token, bos_token, eos_token, unk_token,
):
    added_tokens_decoder = {
        str(i): {
            "content": t, "lstrip": False, "normalized": False, 
            "rstrip": False, "single_word": False, "special": True
        } for i, t in enumerate(special_tokens)
    }

    cfg = {
        "legacy": True, "model_max_length": 32768,
        "tokenizer_class": "PreTrainedTokenizerFast",
        "added_tokens_decoder": added_tokens_decoder,
        "bos_token": str(bos_token), "eos_token": str(eos_token),
        "pad_token": str(pad_token), "unk_token": str(unk_token),
        "vocab_size": int(tok.get_vocab_size()),
    }
    if str(scheme_name) == "minimind2_chatml":
        cfg["chat_template"] = _MINIMIND_CHAT_TEMPLATE

    (out_dir / "tokenizer_config.json").write_text(json.dumps(cfg, ensure_ascii=False, indent=2), encoding="utf-8")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

# dataprep/prepare/text_bpe.py
def train_text_bpe(
    *, input_path, out_dir, vocab_size=6400, special_tokens,
):
    inp = Path(input_path)
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    special_tokens = list(special_tokens)
    scheme_name, pad_token, bos_token, eos_token, unk_token = _infer_core_tokens(special_tokens)

    tok = Tokenizer(BPE(unk_token=str(unk_token), byte_fallback=True))
    tok.pre_tokenizer = ByteLevel(add_prefix_space=False)

    trainer = BpeTrainer(
        vocab_size=int(vocab_size), special_tokens=special_tokens,
        show_progress=True, initial_alphabet=ByteLevel.alphabet(),
    )
    tok.train([str(inp)], trainer=trainer)

    # 保证特殊 token 的 ID 稳定
    for i, t in enumerate(special_tokens):
        tid = tok.token_to_id(t)
        if tid is None or int(tid) != int(i):
            raise ValueError(f"special token id mismatch: {t} -> {tid} (expected {i})")
    
    tok.decoder = ByteLevelDecoder()
    tok.save(str(out / "tokenizer.json"))
    _write_tokenizer_config(
        out_dir=out, tok=tok, special_tokens=special_tokens,
        scheme_name=str(scheme_name), pad_token=str(pad_token),
        bos_token=str(bos_token), eos_token=str(eos_token), unk_token=str(unk_token),
    )

    meta = {
        "special_tokens_scheme": str(scheme_name), "vocab_size": tok.get_vocab_size(),
        "special_tokens": special_tokens, "special_ids": {t: tok.token_to_id(t) for t in special_tokens},
    }
    (out / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

# dataprep/prepare/tokenizer.py
def _count_nonempty_lines(path):
    n = 0
    with path.open("r", encoding="utf-8-sig") as f:
        for line in f:
            if line.strip():
                n += 1
    return n

def _sample_text_corpus(*, src, dst, ratio, seed, overwrite):
    # 随机采样逻辑，详见配套源码
    pass

def _ensure_tokenizer(*, text_corpus, out_dir, vocab_size, overwrite):
    meta_path = out_dir / "meta.json"
    if overwrite and out_dir.exists():
        shutil.rmtree(out_dir)

    if meta_path.exists():
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
        if meta.get("special_tokens_scheme") == "minimind2_chatml" and int(meta.get("vocab_size", -1)) == int(vocab_size):
            return

    out_dir.parent.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    train_text_bpe(
        input_path=text_corpus,
        out_dir=out_dir,
        vocab_size=int(vocab_size),
        special_tokens=list(MINIMIND2_CHATML_TOKENS),
    )
    dt = int(time.time() - t0)
    print(f"ok: tokenizer -> {out_dir} (sec={dt})")

def ensure_tokenizer_and_load(cfg):
    seed = int(cfg.get("seed", 42))
    overwrite = bool(cfg.get("overwrite", False))
    tok_cfg = cfg.get("tokenizer", {})

    vocab_size = int(TOKENIZER_VOCAB_SIZE)
    tokenizer_out = Path(TOKENIZER_DIR)
    text_corpus = Path(MINIMIND_TEXT_CORPUS)
    if not text_corpus.exists():
        raise FileNotFoundError(text_corpus)

    tmp_root = Path(DATA_INTERIM) / "tmp"
    sample_ratio = float(tok_cfg.get("sample_ratio", 1.0))
    sample_seed = int(tok_cfg.get("sample_seed", seed))

    tokenizer_corpus_for_train = text_corpus
    if sample_ratio < 1.0:
        safe_name = text_corpus.name.replace(".", "_")
        ratio_str = f"{sample_ratio:.6f}".rstrip("0").rstrip(".")
        sampled = tmp_root / f"{safe_name}.sample_p{ratio_str}_seed{sample_seed}.txt"
        tokenizer_corpus_for_train = _sample_text_corpus(
            src=text_corpus, dst=sampled, ratio=sample_ratio, seed=sample_seed, overwrite=overwrite
        )

    print(f"using tokenizer corpus: {tokenizer_corpus_for_train} (lines={_count_nonempty_lines(tokenizer_corpus_for_train)})")
    print("== prepare: tokenizer ==")
    
    _ensure_tokenizer(
        text_corpus=tokenizer_corpus_for_train,
        out_dir=tokenizer_out,
        vocab_size=vocab_size,
        overwrite=overwrite,
    )

    if tokenizer_corpus_for_train != text_corpus and tokenizer_corpus_for_train.exists():
        tokenizer_corpus_for_train.unlink(missing_ok=True)

    tok = Tokenizer.from_file(str(tokenizer_out / "tokenizer.json"))
    if int(tok.get_vocab_size()) != int(vocab_size):
        raise RuntimeError(f"tokenizer vocab mismatch: tok={tok.get_vocab_size()} expected={vocab_size}")
    return tok

1
2
3
4
5

# dataprep/prepare/__init__.py
from .tokenizer import ensure_tokenizer_and_load

def run(cfg):
    tok = ensure_tokenizer_and_load(cfg)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

# dataprep/prepare/memmap.py
import shutil
import time
from pathlib import Path
from tokenizers import Tokenizer
from .memmap_io import MemmapMeta

def _ensure_packed_text(*, text_path, out_dir, tokenizer, vocab_size, max_seq_len, max_samples, seed, overwrite):
    # 1. 校验是否已存在符合要求的历史特征，避免重复生成
    meta_path = out_dir / "meta.json"
    if meta_path.exists() and not overwrite:
        meta = MemmapMeta.load(meta_path)
        if int(meta.seq_len) == int(max_seq_len) and int(meta.vocab_size) == int(vocab_size):
            return

    # 2. 调用底层的构造器正式生成二进制特征阵列
    # (具体封装逻辑 `build_packed_text_memmap_dataset` 省略，功能是读取文本逐行Token化并写入硬盘)
    ...

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

# dataprep/prepare/memmap_io.py
import json
from dataclasses import dataclass
from pathlib import Path
import numpy as np

@dataclass(frozen=True)
class MemmapMeta:
    """memmap 数据集的元信息（用于数据准备阶段写入/校验）。"""
    # ...省略字段定义（num_samples, seq_len 等）...

    @staticmethod
    def load(path):
        # ...从 JSON 加载元数据逻辑...
        pass
        
    def save(self, path):
        # ...保存至 meta.json 逻辑...
        pass

class MemmapWriter:
    """将 token/labels/attention 写入到 memmap 目录。"""
    def __init__(self, out_dir, *, num_samples, seq_len, vocab_size, pad_id, image_tokens=None, image_feat_dim=None, mode="w+"):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        
        # 初始化与校验元数据
        self.meta = MemmapMeta(num_samples, seq_len, vocab_size, pad_id, image_tokens=image_tokens, image_feat_dim=image_feat_dim)
        if mode == "w+":
            self.meta.save(self.out_dir / "meta.json")
            
        n = int(self.meta.num_samples)
        s = int(self.meta.seq_len)

        # 基于元信息开辟 Numpy 内存映射空间
        self.input_ids = np.memmap(self.out_dir / "input_ids.bin", dtype=np.dtype(self.meta.input_dtype), mode=mode, shape=(n, s))
        self.labels = np.memmap(self.out_dir / "labels.bin", dtype=np.dtype(self.meta.label_dtype), mode=mode, shape=(n, s))
        self.attention_mask = np.memmap(self.out_dir / "attention_mask.bin", dtype=np.dtype(self.meta.attention_dtype), mode=mode, shape=(n, s))

        self.image_feats = None
        if self.meta.image_tokens is not None and self.meta.image_feat_dim is not None:
            self.image_feats = np.memmap(
                self.out_dir / "image_feats.bin", dtype=np.float16, mode=mode,
                shape=(n, int(self.meta.image_tokens), int(self.meta.image_feat_dim)),
            )

    def write(self, idx, *, input_ids, labels, attention_mask, image_feats=None):
        self.input_ids[idx] = input_ids.astype(self.input_ids.dtype, copy=False)
        self.labels[idx] = labels.astype(self.labels.dtype, copy=False)
        self.attention_mask[idx] = attention_mask.astype(self.attention_mask.dtype, copy=False)
        
        if self.image_feats is not None and image_feats is not None:
            self.image_feats[idx] = image_feats.astype(np.float16, copy=False)

    def flush(self):
        self.input_ids.flush()
        self.labels.flush()
        self.attention_mask.flush()
        if self.image_feats is not None:
            self.image_feats.flush()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

# dataprep/prepare/memmap.py
from .sft_builder import build_sft_text_memmap_dataset

def _ensure_sft_text(*, in_jsonl, out_dir, tokenizer, vocab_size, max_seq_len, overwrite):
    # 第一步同理，先通过 meta.json 校验是否可复用缓存跳过...

    # 主要差异：由专属的 SFT 构造器来处理带角色限定的对话语料
    build_sft_text_memmap_dataset(
        jsonl_path=str(in_jsonl),
        out_dir=str(out_dir),
        text_tokenizer=tokenizer,
        vocab_size=int(vocab_size),
        max_seq_len=int(max_seq_len),
        resume=False,
        flush_every=200,
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# dataprep/prepare/sft_builder.py
import json
import numpy as np
from .memmap_io import MemmapWriter

def build_sft_text_memmap_dataset(*, jsonl_path, out_dir, text_tokenizer, max_seq_len, vocab_size, ...):
    # 1. 挂载词表中的特殊控制 Token（如 pad 占位，以及 ChatML 体系的体系化标记）
    pad_id = text_tokenizer.token_to_id("<|endoftext|>")
    bos_id = text_tokenizer.token_to_id("<|im_start|>")
    eos_id = text_tokenizer.token_to_id("<|im_end|>")
    # ...省略繁琐的 Token ID 校验与断言...
    
    # 2. 初始化底层流式存储类
    writer = MemmapWriter(out_dir, num_samples=n, seq_len=max_seq_len, vocab_size=vocab_size, pad_id=pad_id, mode='w+')
    
    # 3. 逐行读取 jsonl 格式的微调语料，通过 ChatML 角色控制符硬编码包裹拼接
    for idx, obj in enumerate(_iter_json_objects(jsonl_path)):
        conversations = obj.get("conversations")
        
        tokens = []       # 喂给大模型的所有字符索引
        labels_list = []  # 训练掩码对齐数组
        
        for m in conversations:
            role = m["role"]
            content = m["content"]
            
            # 使用 ChatML 特有的边界符包裹具体的发言内容，帮助基座快速建立角色认知
            header = [bos_id] + text_tokenizer.encode(f"{role}\n").ids
            body = text_tokenizer.encode(content).ids
            footer = [eos_id] + text_tokenizer.encode("\n").ids
            
            msg_tokens = header + body + footer
            
            # 只让大模型对 Assistant（它自己的答案）去计算 Loss
            # 对于 System、User 的上下文前提，统统打上 -100 (PyTorch CrossEntropyLoss 默认的 ignores_index)
            msg_labels = [-100] * len(msg_tokens)
            if role == "assistant":
                for j in range(len(header), len(msg_tokens)):
                    msg_labels[j] = msg_tokens[j]
                    
            tokens.extend(msg_tokens)
            labels_list.extend(msg_labels)

        # 4. 根据设定的最大序列长度进行截断与 Padding（以 pad_id 填补余位），最后交由 MemmapWriter 写盘
        # ...省略截断补齐的矩阵赋值...
        writer.write(idx, input_ids=input_ids, labels=labels, attention_mask=attention_mask)
        
    writer.flush()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

# dataprep/prepare/memmap.py
from pathlib import Path
from ..data_paths import MINIMIND_TEXT_CORPUS, MINIMIND_SFT_SEEKER, TEXT_PRETRAIN_340, TEXT_SFT_340

def ensure_memmaps(cfg: dict, tokenizer):
    overwrite = bool(cfg.get("overwrite", False))
    vocab_size = int(tokenizer.get_vocab_size())
    
    print("== prepare: packed text (340) ==")
    _ensure_packed_text(
        text_path=Path(MINIMIND_TEXT_CORPUS),
        out_dir=Path(TEXT_PRETRAIN_340),
        tokenizer=tokenizer,
        vocab_size=vocab_size,
        max_seq_len=340,
        overwrite=overwrite,
    )

    print("== prepare: sft text (340) ==")
    _ensure_sft_text(
        in_jsonl=Path(MINIMIND_SFT_SEEKER),
        out_dir=Path(TEXT_SFT_340),
        tokenizer=tokenizer,
        vocab_size=vocab_size,
        max_seq_len=340,
        overwrite=overwrite,
    )

1
2
3
4
5
6
7

# dataprep/prepare/__init__.py
from .tokenizer import ensure_tokenizer_and_load
from .memmap import ensure_memmaps  # 引入新增的内存映射逻辑

def run(cfg):
    tok = ensure_tokenizer_and_load(cfg)
    ensure_memmaps(cfg, tok)  # 词表构建完毕，正式执行所有数据的二进制分配

1
2
3
4
5
6
7

# dataprep/prepare/__main__.py
from . import run
from ..data_paths import default_dataprep_cfg

if __name__ == "__main__":
    cfg = default_dataprep_cfg()
    run(cfg)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

graph LR
    subgraph 输入["输入"]
        TXT["Token IDs"]
        IMG["Image Features"]
    end

    subgraph 嵌入["嵌入层"]
        EMB["词嵌入<br/><i>special + base</i>"]
        PROJ["Linear Projection<br/>× tanh(门控)"]
        ADD_E(("Add"))
    end

    TXT --> EMB --> ADD_E
    IMG -.-> PROJ -.-> ADD_E
    ADD_E --> DROP["Dropout"]

    subgraph BLOCK["Transformer Block × N"]

        subgraph 自注意力["Grouped-Query Attention"]
            QKV["Q / K / V 投影"] --> ROPE["RoPE"]
            ROPE --> SDPA["Scaled Dot-Product<br/>Attention<br/><i>+ Causal Mask</i>"]
            SDPA --> O_PROJ["Output Projection"]
        end

        subgraph 前馈网络["SwiGLU Feed-Forward"]
            GU["W_gate / W_up"] --> SW["SiLU(gate) × up"]
            SW --> WD["W_down"]
        end

        DROP --> NORM1["RMSNorm"]
        NORM1 --> QKV
        O_PROJ --> ADD1(("Add"))
        DROP -. "残差" .-> ADD1

        ADD1 --> NORM2["RMSNorm"]
        NORM2 --> GU
        WD --> ADD2(("Add"))
        ADD1 -. "残差" .-> ADD2
    end

    subgraph 输出["输出层"]
        FNORM["RMSNorm"]
        HEAD["LM Head<br/><i>与嵌入权重共享</i>"]
        LOGITS["Output Logits"]
    end

    ADD2 --> FNORM --> HEAD --> LOGITS

    style BLOCK fill:#fef9e7,stroke:#f39c12,stroke-width:2px
    style 自注意力 fill:#eaf2f8,stroke:#2980b9
    style 前馈网络 fill:#f5eef8,stroke:#8e44ad
    style 嵌入 fill:#fdf2e9,stroke:#e67e22
    style 输出 fill:#e8f8f5,stroke:#1abc9c
    style ADD1 fill:#d5f5e3,stroke:#27ae60
    style ADD2 fill:#d5f5e3,stroke:#27ae60
    style ADD_E fill:#d5f5e3,stroke:#27ae60

1
2
3
4
5

# seeker_omni/paths.py
from pathlib import Path
DATA_PROCESSED = Path("data/processed")
TEXT_PRETRAIN_340 = DATA_PROCESSED / "text_pretrain_packed_340_u16_offline"
# ... 其他训练相关的数据集与分词器默认常量路径定义

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# seeker_omni/config.py
from dataclasses import dataclass
from .special_tokens import DEFAULT_SPECIAL_TOKENS_SCHEME

@dataclass(frozen=True)
class ModelConfig:
    name: str
    vocab_size: int
    max_seq_len: int
    hidden_size: int
    num_layers: int
    num_heads: int
    # ... 其他架构参数
    special_tokens_scheme: str = DEFAULT_SPECIAL_TOKENS_SCHEME

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

# seeker_omni/special_tokens.py
from dataclasses import dataclass

@dataclass(frozen=True)
class TokenSchemeSpec:
    name: str
    special_tokens: list[str]
    pad_token: str
    bos_token: str
    eos_token: str
    unk_token: str


MINIMIND2_CHATML = TokenSchemeSpec(
    name="minimind2_chatml",
    special_tokens=[
        "<|endoftext|>",
        "<|im_start|>",
        "<|im_end|>",
        "<img_bos>",
        "<img>",
        "<img_eos>",
    ],
    pad_token="<|endoftext|>",
    bos_token="<|im_start|>",
    eos_token="<|im_end|>",
    unk_token="<|endoftext|>",
)


_SPECS: dict[str, TokenSchemeSpec] = {
    MINIMIND2_CHATML.name: MINIMIND2_CHATML,
}


def get_token_scheme_spec(scheme: str | None) -> TokenSchemeSpec:
    scheme = (scheme or MINIMIND2_CHATML.name).strip()
    spec = _SPECS.get(scheme)
    if spec is None:
        opts = ", ".join(sorted(_SPECS.keys()))
        raise ValueError(f"unknown special_tokens_scheme={scheme!r} (expected one of: {opts})")
    return spec


@dataclass(frozen=True)
class SpecialTokenIds:
    pad: int
    unk: int
    bos: int
    eos: int
    img_bos: int
    img: int
    img_eos: int


def build_special_token_ids(spec: TokenSchemeSpec) -> SpecialTokenIds:
    tok2id = {t: i for i, t in enumerate(spec.special_tokens)}

    def _id(t: str) -> int:
        if t not in tok2id:
            raise ValueError(f"required special token missing from scheme={spec.name!r}: {t}")
        return int(tok2id[t])

    return SpecialTokenIds(
        pad=_id(spec.pad_token),
        unk=_id(spec.unk_token),
        bos=_id(spec.bos_token),
        eos=_id(spec.eos_token),
        img_bos=_id("<img_bos>"),
        img=_id("<img>"),
        img_eos=_id("<img_eos>"),
    )


DEFAULT_SPECIAL_TOKENS_SCHEME = MINIMIND2_CHATML.name

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

# seeker_omni/model/lm.py
import torch
import torch.nn as nn
from ..config import ModelConfig
from ..special_tokens import get_token_scheme_spec, build_special_token_ids

class SeekerOmniLM(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.special_spec = get_token_scheme_spec(cfg.special_tokens_scheme)
        self.special = build_special_token_ids(self.special_spec)
        self.n_special = int(len(self.special_spec.special_tokens))

        # 将特殊 token（如 bos, eos, img 占位符等）与普通文本 token 分离
        self.special_embed = nn.Embedding(self.n_special, cfg.hidden_size)
        self.base_embed = nn.Embedding(cfg.vocab_size - self.n_special, cfg.hidden_size)

    def _embed_tokens(self, input_ids):
        # 分别根据 ID 范围查特殊词表和普通词表，最后拼合成统一的隐藏层张量
        n_special = int(self.n_special)
        is_special = input_ids < n_special

        x = torch.empty(
            (*input_ids.shape, self.cfg.hidden_size),
            device=input_ids.device,
            dtype=self.special_embed.weight.dtype,
        )

        if is_special.any():
            x[is_special] = self.special_embed(input_ids[is_special])
        if (~is_special).any():
            base_ids = (input_ids[~is_special] - n_special).clamp(min=0)
            x[~is_special] = self.base_embed(base_ids)
        return x

    def reset_parameters(self):
        # 初始化常规嵌入以及线性层
        nn.init.normal_(self.special_embed.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.base_embed.weight, mean=0.0, std=0.02)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

        # 关键：将模态占位符的词嵌入强行初始化为 0，确保模型初期不会因占位符引入额外噪声
        with torch.no_grad():
            for tid in (
                self.special.img_bos,
                self.special.img,
                self.special.img_eos,
            ):
                self.special_embed.weight[tid].zero_()

    def _lm_head_weight(self):
        # 语言模型头的权重与底层的文本嵌入权重共享（Weight Tying）以节约显存
        return torch.cat([self.special_embed.weight, self.base_embed.weight], dim=0)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

# seeker_omni/model/resampler.py
import torch
from torch import nn

class PerceiverResampler(nn.Module):
    def __init__(self, *, dim, num_latents, num_layers=2, num_heads=8, ff_mult=4):
        super().__init__()
        self.dim = int(dim)
        
        # 初始化固定数量的可学习潜在特征 (Latents)
        lat = torch.empty((num_latents, dim), dtype=torch.float32)
        nn.init.trunc_normal_(lat, std=0.02)
        self.latents = nn.Parameter(lat)
        
        # 构建交叉注意力层与前馈网络
        # ... (具体层的构建代码省略)
        self.layers = nn.ModuleList([...])

    def forward(self, x):
        # x 为 SigLIP 提取的原始长序列视觉特征: [B, T_img, D]
        b = int(x.shape[0])
        # 扩展潜在特征以匹配批次大小: [B, num_latents, D]
        latents = self.latents.unsqueeze(0).expand(b, -1, -1)

        for layer in self.layers:
            q = layer["ln_q"](latents)
            kv = layer["ln_kv"](x)
            # 通过交叉注意力，用少量的 Latents 去 Query 密集的视觉特征 KV
            attn_out, _ = layer["attn"](q, kv, kv, need_weights=False)
            latents = latents + attn_out
            latents = latents + layer["ff"](layer["ln_ff"](latents))

        # 返回定长的压缩视觉特征: [B, num_latents, D]
        return latents

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

# seeker_omni/model/projector.py
import torch

def inject_feature_tokens(
    x,
    *,
    input_ids,
    image_feats,
    img_token_id,
    img_proj,
    img_gate,
):
    # x: [B, S, hidden_size]
    if image_feats is not None:
        img_mask = input_ids == int(img_token_id)
        if img_mask.any():
            # 将视觉特征投影至 LLM 维度
            img_tokens = img_proj(image_feats.to(dtype=x.dtype))
            # 门控机制：初始阶段 tanh(0)=0，使得模型初期等同于纯文本模型，有助于更稳定的联合训练
            img_tokens = img_tokens * torch.tanh(img_gate)[None, None, :]
            
            # 使用 nonzero 获取展平索引，进行快路径累加注入
            pos = img_mask.nonzero(as_tuple=False)  # [N,2] row-major
            flat = img_tokens.reshape(-1, img_tokens.shape[-1])
            if pos.shape[0] == flat.shape[0]:
                x[pos[:, 0], pos[:, 1]] = x[pos[:, 0], pos[:, 1]] + flat
            else:
                for b in range(int(x.shape[0])):
                    idx = torch.where(img_mask[b])[0]
                    if idx.numel() == 0:
                        continue
                    take = min(idx.numel(), img_tokens.shape[1])
                    x[b, idx[:take]] = x[b, idx[:take]] + img_tokens[b, :take]
    return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# seeker_omni/model/lm.py

    def _inject_modality_tokens(self, x, *, input_ids, image_feats):
        # 将传入的投影参数绑定好，方便 forward 调用
        from .projector import inject_feature_tokens
        return inject_feature_tokens(
            x,
            input_ids=input_ids,
            image_feats=image_feats,
            img_token_id=int(self.special.img),
            img_proj=self.img_proj,
            img_gate=self.img_gate,
        )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

# seeker_omni/model/block.py
import torch
from torch import nn

from .attention import PastKeyValue, SeekerSelfAttention
from .mlp import SeekerMLP
from .norm import RMSNorm

class SeekerBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, num_kv_heads, **kwargs): # 用 **kwargs 代替其他传入的初始化超参数
        super().__init__()
        self.attn_norm = RMSNorm(hidden_size)
        self.attn = SeekerSelfAttention(hidden_size, num_heads, num_kv_heads, **kwargs)
        self.mlp_norm = RMSNorm(hidden_size)
        self.mlp = SeekerMLP(hidden_size, **kwargs)

    def forward(self, x, cos, sin, attention_mask=None, past_kv=None, use_cache=False):
        # 如果开启流式缓存，则调用注意力机制时抽取并封装 KV 缓存
        if use_cache:
            attn_out, present_kv = self.attn(
                self.attn_norm(x),
                cos=cos, sin=sin, attention_mask=attention_mask,
                past_kv=past_kv, use_cache=True,
            )
            x = x + attn_out
            x = x + self.mlp(self.mlp_norm(x))
            return x, present_kv

        # 否则普通的残差堆叠计算
        x = x + self.attn(self.attn_norm(x), cos=cos, sin=sin, attention_mask=attention_mask)
        x = x + self.mlp(self.mlp_norm(x))
        return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# seeker_omni/model/norm.py
import torch
from torch import nn

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        scale = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (x * scale) * self.weight

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

# seeker_omni/model/rope.py
import torch
from torch import nn

def _build_rope_cache(seq_len, dim, theta):
    # 构建复数域的旋转角度缓存矩阵
    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    positions = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.einsum("i,j->ij", positions, inv_freq)
    return torch.cos(freqs), torch.sin(freqs)

def apply_rope(x, cos, sin):
    # 将预计算好的 cos 和 sin 旋转矩阵作用到张量的奇偶维度上
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    
    out_even = x_even * cos - x_odd * sin
    out_odd = x_even * sin + x_odd * cos
    return torch.stack((out_even, out_odd), dim=-1).flatten(-2)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

# seeker_omni/model/attention.py
import torch.nn.functional as F
from .rope import apply_rope

class SeekerSelfAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, num_kv_heads, dropout=0.0):
        super().__init__()
        # ... 维度整除校验等省略 ...
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden_size // num_heads
        
        # 分组查询注意力（GQA）投影层
        q_out = num_heads * self.head_dim
        kv_out = num_kv_heads * self.head_dim
        self.q_proj = nn.Linear(hidden_size, q_out, bias=False)
        self.k_proj = nn.Linear(hidden_size, kv_out, bias=False)
        self.v_proj = nn.Linear(hidden_size, kv_out, bias=False)
        self.o_proj = nn.Linear(q_out, hidden_size, bias=False)

    def forward(self, x, *, cos, sin, attention_mask=None, past_kv=None, use_cache=False):
        bsz, seq_len, _ = x.shape

        # 1. 投影并重塑维度: [B, H, S, D]
        q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # 2. 对 Q 和 K 注入 RoPE 旋转位置编码
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        # 3. 如果使用 KV Cache，则与历史缓存进行维度拼接
        if past_kv is not None:
            past_k, past_v = past_kv
            k_all = torch.cat([past_k.to(k.dtype), k], dim=2)
            v_all = torch.cat([past_v.to(v.dtype), v], dim=2)
        else:
            k_all, v_all = k, v

        # 4. GQA：将 K, V 的头数广播扩展至与 Q 一致
        if self.num_kv_heads != self.num_heads:
            repeat = self.num_heads // self.num_kv_heads
            k_rep = k_all.repeat_interleave(repeat, dim=1)
            v_rep = v_all.repeat_interleave(repeat, dim=1)
        else:
            k_rep, v_rep = k_all, v_all

        # 5. 求解自回归掩码与 Padding 掩码偏差（attn_bias）的逻辑省略 ...
        attn_bias = None 

        # 6. 调用 PyTorch 原生的 Flash Attention / 缩放点积注意力加速计算
            out = F.scaled_dot_product_attention(
                q, k_rep, v_rep,
                attn_mask=attn_bias,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=False, # 手动传入掩码时需关闭原生 causal 标志
            )

        # 6. 维度还原与输出投影
        out = out.transpose(1, 2).contiguous().view(bsz, seq_len, -1)
        out = self.o_proj(out)

        if use_cache:
            return out, (k_all, v_all)
        return out

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# seeker_omni/model/mlp.py
import torch
import torch.nn.functional as F
from torch import nn

def _round_up(x, multiple):
    # 向上取整到 multiple 的倍数，用于硬件显存对齐
    return ((x + multiple - 1) // multiple) * multiple

class SeekerMLP(nn.Module):
    def __init__(self, hidden_size, intermediate_size=None, dropout=0.0):
        super().__init__()
        # 如果未指定中间层维度，则默认缩放为 8/3 并对齐 64
        if intermediate_size is None:
            intermediate_size = int(hidden_size * 8 / 3)
            intermediate_size = _round_up(intermediate_size, 64)

        # SwiGLU 需要平行的 Gate 和 Up 两个上采样矩阵
        self.w_gate = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.w_up = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.w_down = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.dropout = dropout

    def forward(self, x):
        # 将门控输出经过 SiLU 后与 Up 输出逐元素相乘
        x = F.silu(self.w_gate(x)) * self.w_up(x)
        x = self.w_down(x)
        return F.dropout(x, p=self.dropout, training=self.training)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

# seeker_omni/model/lm.py

class SeekerOmniLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        # ... 前文已初始化的特殊文本嵌入层等代码 (此处略)
        
        # 初始化模态投影层与门控
        self.img_proj = nn.Linear(cfg.image_feat_dim, cfg.hidden_size, bias=False)
        self.img_gate = nn.Parameter(torch.zeros(cfg.hidden_size))
        
        # 初始化旋转位置编码与自注意力解码块序列
        head_dim = cfg.hidden_size // cfg.num_heads
        self.rope = RotaryEmbedding(dim=head_dim, max_seq_len=cfg.max_seq_len, theta=cfg.rope_theta)
        self.blocks = nn.ModuleList([
            SeekerBlock(
                cfg.hidden_size,
                cfg.num_heads,
                cfg.num_kv_heads,
                dropout=cfg.dropout,
                intermediate_size=cfg.mlp_intermediate_size,
            ) for _ in range(cfg.num_layers)
        ])
        
        self.norm = RMSNorm(cfg.hidden_size)

    def forward(self, input_ids, image_feats=None, **kwargs): 
        # 1. 文本 Token 化为词汇嵌入矩阵
        x = self._embed_tokens(input_ids)
        
        # 2. 模态特征注入: 使用 Projector 将视觉特征注入到指定的图像占位符位置
        x = self._inject_modality_tokens(x, input_ids=input_ids, image_feats=image_feats)
        x = self.drop(x)
        
        # 3. 准备 RoPE 旋转位置编码
        seq_len = int(input_ids.shape[1])
        cos, sin = self.rope.get_cos_sin(seq_len) # 省略对具体最大序列长度和精度的适配

        # 4. 逐层通过标准的 Transformer Decoder Blocks 进行自注意力推理
        present_kv = [] if use_cache else None

        for i, block in enumerate(self.blocks):
            if use_cache:
                pkv = past_kv[i] if past_kv is not None else None
                x, kv = block(x, cos=cos, sin=sin, attention_mask=attention_mask, past_kv=pkv, use_cache=True)
                present_kv.append(kv)
            else:
                x = block(x, cos=cos, sin=sin, attention_mask=attention_mask)

        # 5. 计算最后的输出与 LM Loss
        x = self.norm(x)
        logits = F.linear(x, self._lm_head_weight())
        
        # 返回带分布概率的 Logits 和更新后的 KV 缓存
        return SeekerOmniOutput(logits=logits, loss=loss, kv_cache=present_kv) # 封装为 dataclass，省略损失计算部分