train-llm-from-scratch

BASE_URL = "https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main"
VAL_URL = f"{BASE_URL}/val.jsonl.zst"                              # 验证集（固定一个文件）
TRAIN_URLS = [f"{BASE_URL}/train/{i:02d}.jsonl.zst" for i in range(65)]  # 训练集：00~64 共 65 个分片

def download_file(url: str, file_name: str) -> None:
    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))
    block_size = 1024
    with open(file_name, 'wb') as f:
        for chunk in tqdm(response.iter_content(block_size), total=total_size // block_size, desc="Downloading", leave=True):
            f.write(chunk)

def process_files(input_dir, output_file, tokenizer_name, max_data=None):
    enc = tiktoken.get_encoding(tokenizer_name)   # 默认 'r50k_base'（GPT-3 用的分词器）

    with h5py.File(output_file, 'w') as out_f:
        # 创建一个一维、可动态扩容的 dataset，存所有 token
        dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')
        start_index = 0

        for filename in sorted(os.listdir(input_dir)):
            if filename.endswith(".jsonl.zst"):
                with zstd.open(in_file, 'rt', encoding='utf-8') as in_f:  # 流式解压 + 文本模式读
                    for line in tqdm(in_f, desc=f"Processing {filename}", total=max_data):
                        data = json.loads(line)
                        text = data.get('text')
                        if text:
                            # 每条文本末尾追加 <|endoftext|>，再编码
                            encoded = enc.encode(text + "<|endoftext|>", allowed_special={'<|endoftext|>'})
                            encoded_len = len(encoded)
                            end_index = start_index + encoded_len
                            dataset.resize(dataset.shape[0] + encoded_len, axis=0)  # 扩容
                            dataset[start_index:end_index] = encoded                # 追加写入
                            start_index = end_index
                        processed_lines += 1
                        if max_data is not None and processed_lines >= max_data:
                            break  # 每个文件最多处理 max_data 行

def get_batch_iterator(data_path, batch_size, context_length, device="cpu"):
    with h5py.File(data_path, 'r') as hdf5_file:
        dataset = hdf5_file['tokens']
        dataset_size = dataset.shape[0]
        # 能切出多少个不重叠的样本（-1 是因为目标要往后错一位）
        n_examples = (dataset_size - 1) // context_length
        example_idxs = np.arange(n_examples)
        np.random.shuffle(example_idxs)   # 打乱样本顺序

        epochs = 0
        counter = 0
        while True:                       # 无限循环，训练循环自己控制何时停
            if counter + batch_size > n_examples:
                np.random.shuffle(example_idxs)   # 一个 epoch 用完，重新打乱
                counter = 0
                print(f"Finished epoch {epochs}")
                epochs += 1

            # 把"第几个样本"换算成 token 流里的起始下标
            random_indices = example_idxs[counter:counter+batch_size] * context_length
            # 每个样本取 context_length+1 个 token（多取一个用于错位）
            random_samples = torch.tensor(np.array(
                [dataset[idx:idx+context_length+1] for idx in random_indices]
            ))
            xb = random_samples[:, :context_length].to(device)    # 输入：前 context_length 个
            yb = random_samples[:, 1:context_length+1].to(device) # 目标：错后一位
            counter += batch_size
            yield xb, yb

token 流片段：  [The] [cat] [sat] [on] [the] [mat]
xb (输入):      [The] [cat] [sat] [on] [the]            # 位置 0..T-1
yb (目标):            [cat] [sat] [on] [the] [mat]      # 位置 1..T

class MLP(nn.Module):
    def __init__(self, n_embed):
        super().__init__()
        self.hidden = nn.Linear(n_embed, 4 * n_embed)  # 升维到 4 倍
        self.relu = nn.ReLU()
        self.proj = nn.Linear(4 * n_embed, n_embed)    # 投影回原维度

    def forward(self, x):
        x = self.forward_embedding(x)   # hidden + relu
        x = self.project_embedding(x)   # proj
        return x

    def forward_embedding(self, x):
        return self.relu(self.hidden(x))

    def project_embedding(self, x):
        return self.proj(x)

class Head(nn.Module):
    def __init__(self, head_size, n_embed, context_length):
        super().__init__()
        self.key   = nn.Linear(n_embed, head_size, bias=False)   # K 投影
        self.query = nn.Linear(n_embed, head_size, bias=False)   # Q 投影
        self.value = nn.Linear(n_embed, head_size, bias=False)   # V 投影
        # 下三角矩阵，注册成 buffer（不是参数，不训练，但随模型存取/搬设备）
        self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))

    def forward(self, x):
        B, T, C = x.shape
        head_size = self.key.out_features
        k = self.key(x)      # (B, T, head_size)
        q = self.query(x)    # (B, T, head_size)
        scale_factor = 1 / math.sqrt(head_size)
        # 注意力分数：q 和 k 做点积
        attn_weights = q @ k.transpose(-2, -1) * scale_factor       # (B, T, T)
        # 因果掩码：把"未来"位置置成 -inf
        attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        attn_weights = F.softmax(attn_weights, dim=-1)              # 每行归一化成概率
        v = self.value(x)    # (B, T, head_size)
        out = attn_weights @ v                                      # (B, T, head_size)
        return out

class MultiHeadAttention(nn.Module):
    def __init__(self, n_head, n_embed, context_length):
        super().__init__()
        # n_head 个头，每个头维度 n_embed // n_head
        self.heads = nn.ModuleList(
            [Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)]
        )

    def forward(self, x):
        # 每个头独立处理，再沿最后一维拼接
        x = torch.cat([h(x) for h in self.heads], dim=-1)
        return x

class Block(nn.Module):
    def __init__(self, n_head, n_embed, context_length):
        super().__init__()
        self.ln1  = nn.LayerNorm(n_embed)
        self.attn = MultiHeadAttention(n_head, n_embed, context_length)
        self.ln2  = nn.LayerNorm(n_embed)
        self.mlp  = MLP(n_embed)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # 注意力子层 + 残差
        x = x + self.mlp(self.ln2(x))    # 前馈子层 + 残差
        return x

class Transformer(nn.Module):
    def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
        super().__init__()
        self.context_length = context_length
        self.N_BLOCKS = N_BLOCKS
        self.token_embed    = nn.Embedding(vocab_size, n_embed)        # token → 向量
        self.position_embed = nn.Embedding(context_length, n_embed)    # 位置 → 向量
        self.attn_blocks    = nn.ModuleList(
            [Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)]
        )
        self.layer_norm = nn.LayerNorm(n_embed)              # 最后一层归一化
        self.lm_head    = nn.Linear(n_embed, vocab_size)     # 投影到词表，得到 logits
        self.register_buffer('pos_idxs', torch.arange(context_length))

    def _pre_attn_pass(self, idx):
        B, T = idx.shape
        tok_embedding = self.token_embed(idx)                 # (B, T, C)
        pos_embedding = self.position_embed(self.pos_idxs[:T]) # (T, C)，广播相加
        return tok_embedding + pos_embedding

    def forward(self, idx, targets=None):
        x = self._pre_attn_pass(idx)         # 词嵌入 + 位置嵌入
        for block in self.attn_blocks:       # 逐层 Transformer Block
            x = block(x)
        x = self.layer_norm(x)               # 最终归一化
        logits = self.lm_head(x)             # (B, T, vocab_size)
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            flat_logits = logits.view(B * T, C)
            targets = targets.view(B * T).long()
            loss = F.cross_entropy(flat_logits, targets)   # 交叉熵
        return logits, loss

idx (B,T) ──token_embed──▶ (B,T,C) ┐
                                    ├─加──▶ (B,T,C) ──Block×N──▶ (B,T,C) ──LN──▶ ──lm_head──▶ logits (B,T,V)
pos_idxs[:T] ─position_embed─▶ (T,C)┘

def generate(self, idx, max_new_tokens):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -self.context_length:]   # 只保留最近 context_length 个 token
        logits, _ = self(idx_cond)                 # 前向
        logits = logits[:, -1, :]                  # 只取最后一个位置的预测 (B, V)
        probs = F.softmax(logits, dim=-1)          # 转概率
        idx_next = torch.multinomial(probs, num_samples=1)  # 按概率采样一个 token
        idx = torch.cat((idx, idx_next), dim=1)    # 拼回序列，继续下一轮
    return idx

class DemoCorpusDataset(Dataset):
    def __getitem__(self, idx):
        text = self.samples[idx]
        full_text_with_bos = self.tokenizer.bos_token + text     # 句首加 BOS
        encoding = self.tokenizer(full_text_with_bos, max_length=self.max_length,
                                  padding="max_length", truncation=True, return_tensors='pt')
        input_ids = encoding.input_ids.squeeze(0)                # (max_length,)

        effective_loss_mask = (input_ids != self.tokenizer.pad_token_id).long()  # 非 pad 处为 1
        X = input_ids[:-1]                       # 输入
        Y = input_ids[1:]                        # 目标（错位一格，同主代码库）
        mask_for_loss_calculation = effective_loss_mask[1:]      # 对齐 Y
        return X, Y, mask_for_loss_calculation

class DemoChatDataset(Dataset):
    def __getitem__(self, idx):
        conversations = self.samples[idx]
        # 用 chat template 把整段对话拼接 + 分词
        input_ids = self.tokenizer.apply_chat_template(
            conversations, tokenize=True, add_generation_prompt=False,
            return_tensors="pt", max_length=self.max_length,
            truncation=True, padding="max_length"
        ).squeeze(0)

        loss_mask = torch.zeros_like(input_ids, dtype=torch.long)   # 默认全 0（都不算 loss）
        bos_assistant_ids = self.tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
        eos_ids = self.tokenizer.encode("<|im_end|>", add_special_tokens=False)

        # 扫描 token 序列，把每段 "assistant\n ... <|im_end|>" 之间的 token 标成 1
        i = 0
        input_ids_list = input_ids.tolist()
        while i < len(input_ids_list):
            # 命中一段 assistant 回复的起始标记
            if input_ids_list[i : i+len(bos_assistant_ids)] == bos_assistant_ids:
                start_of_response = i + len(bos_assistant_ids)
                # 往后找该回复的结束标记 <|im_end|>
                end_marker = -1
                j = start_of_response
                while j < len(input_ids_list):
                    if input_ids_list[j : j+len(eos_ids)] == eos_ids:
                        end_marker = j
                        break
                    j += 1
                if end_marker != -1:                                # 找到结束标记
                    loss_mask[start_of_response : end_marker + len(eos_ids)] = 1
                    i = end_marker + len(eos_ids)                   # 跳到这段之后继续扫
                    continue
                else:                                               # 没找到（被截断），标到末尾
                    loss_mask[start_of_response:] = 1
                    break
            i += 1

        loss_mask[input_ids == self.tokenizer.pad_token_id] = 0   # padding 不算
        X = input_ids[:-1]; Y = input_ids[1:]
        mask_for_loss_calculation = loss_mask[1:]
        return X, Y, mask_for_loss_calculation

self.head_dim = self.hidden_size // self.num_attention_heads   # 1024 // 16 = 64
self.num_key_value_heads = num_key_value_heads or num_attention_heads
# 约束：Q 头数必须能被 KV 头数整除（GQA 分组的前提）
if num_attention_heads % num_key_value_heads != 0:
    raise ValueError("num_attention_heads 必须能被 num_key_value_heads 整除")
hidden_act = "silu"          # SwiGLU 用的激活
rms_norm_eps = 1e-5          # RMSNorm 数值稳定项
rope_theta = 10000.0         # RoPE 的频率基数

DEMO_INTERMEDIATE_SIZE = int(DEMO_HIDDEN_SIZE * 8 / 3)              # 1024 × 8/3 ≈ 2730
DEMO_INTERMEDIATE_SIZE = 32 * ((DEMO_INTERMEDIATE_SIZE + 31) // 32) # 向上对齐到 32 的倍数 = 2752

class DemoRMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))   # 可学习的缩放，无 bias

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

    def forward(self, x):
        output_dtype = x.dtype
        x = x.to(torch.float32)          # 在 float32 下算，保证数值稳定
        output = self._norm(x)
        return (output * self.weight).to(output_dtype)

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len, theta=10000.0, device=None):
        super().__init__()
        # θ_i = 1 / theta^(2i/d)，i = 0,1,...,d/2-1
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))   # (d/2,)
        t = torch.arange(max_seq_len).float()                             # (max_seq_len,)
        freqs = torch.outer(t, freqs)                                     # (max_seq_len, d/2)，元素 = m·θ_i
        # 用模长 1、角度 m·θ_i 构造复数 cos(m·θ_i) + i·sin(m·θ_i)
        self.register_buffer("freqs_cis", torch.polar(torch.ones_like(freqs), freqs), persistent=False)

    def forward(self, xq, xk, seq_len):
        # xq/xk: (bsz, num_heads, seq_len, head_dim)
        # 把相邻两维拼成复数：(..., head_dim) → (..., head_dim/2) 复数
        xq_c = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
        xk_c = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
        freqs_cis_pos = self.freqs_cis[:seq_len].unsqueeze(0).unsqueeze(0)  # (1,1,seq_len,d/2)
        # 复数乘法 = 旋转
        xq_out = torch.view_as_real(xq_c * freqs_cis_pos).flatten(3)
        xk_out = torch.view_as_real(xk_c * freqs_cis_pos).flatten(3)
        return xq_out.type_as(xq), xk_out.type_as(xk)

class DemoAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_q_heads  = config.num_attention_heads     # Query 头数 16
        self.num_kv_heads = config.num_key_value_heads      # KV 头数（可 < Q 头数）
        self.num_kv_groups = self.num_q_heads // self.num_kv_heads  # 每组共享的 Q 头数
        self.head_dim = config.head_dim                     # 64
        # 注意 K/V 投影的输出维度是 num_kv_heads × head_dim（可能小于 Q）
        self.q_proj = nn.Linear(hidden, num_q_heads  * head_dim, bias=False)
        self.k_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)
        self.v_proj = nn.Linear(hidden, num_kv_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(num_q_heads * head_dim, hidden, bias=False)  # ← 输出投影（主代码库没有）
        self.rotary_emb = RotaryEmbedding(head_dim, config.max_position_embeddings, theta=config.rope_theta)
        self.flash_available = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn

# 辅助方法：把 KV 头复制 n_rep 次，(bsz, num_kv_heads, seq, d) → (bsz, num_kv_heads*n_rep, seq, d)
def _repeat_kv(self, x, n_rep):
    bs, num_kv_heads, slen, head_dim = x.shape
    if n_rep == 1:
        return x
    return (x[:, :, None, :, :]
            .expand(bs, num_kv_heads, n_rep, slen, head_dim)
            .reshape(bs, num_kv_heads * n_rep, slen, head_dim))

def forward(self, hidden_states, attention_mask=None, position_ids=None,
            past_key_value=None, use_cache=False):
    bsz, q_len, _ = hidden_states.shape
    # 1) 投影 + 拆头：(bsz, heads, q_len, head_dim)
    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_q_heads,  self.head_dim).transpose(1, 2)
    key_states   = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

    # 总的 K/V 序列长度（含 KV cache 里的历史 token）
    kv_seq_len = q_len
    if past_key_value is not None:
        kv_seq_len += past_key_value[0].shape[2]

    # 2) 给 Q、K 施加 RoPE（按当前新 token 的长度）
    query_states, key_states = self.rotary_emb(query_states, key_states, seq_len=q_len)

    # 3) KV cache：把历史的 K/V 拼到前面（推理加速）
    if past_key_value is not None:
        key_states   = torch.cat([past_key_value[0], key_states], dim=2)
        value_states = torch.cat([past_key_value[1], value_states], dim=2)
    current_key_value = (key_states, value_states) if use_cache else None

    # 4) GQA：把 KV 头复制 num_kv_groups 次，匹配 Q 头数
    key_states   = self._repeat_kv(key_states,   self.num_kv_groups)
    value_states = self._repeat_kv(value_states, self.num_kv_groups)

    # 5) 算注意力：优先 Flash，否则手动
    if self.flash_available and attention_mask is None:
        attn_output = F.scaled_dot_product_attention(
            query_states, key_states, value_states,
            dropout_p=self.config.dropout if self.training else 0.0,
            is_causal=(q_len == kv_seq_len))        # 无 KV cache 时即标准因果注意力
    else:
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
        # 因果掩码；causal_shift 兼容 KV cache（q_len < kv_seq_len）的情形
        causal_shift = kv_seq_len - q_len
        mask = torch.triu(torch.full((q_len, kv_seq_len), float("-inf"),
                                     device=query_states.device), diagonal=1 + causal_shift)
        attn_weights = attn_weights + mask[None, None, :, :]
        if attention_mask is not None:              # 额外的 padding 掩码（加性）
            attn_weights = attn_weights + attention_mask
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(query_states)
        attn_output = torch.matmul(attn_weights, value_states)

    # 6) 合并头 + 输出投影
    attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
    return self.o_proj(attn_output), current_key_value

class DemoFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]   # SiLU
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))

class DemoTransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.self_attn = DemoAttention(config)
        self.mlp = DemoFeedForward(config)
        self.input_layernorm = DemoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = DemoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, position_ids=None,
                past_key_value=None, use_cache=False):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        attn_out, present_kv = self.self_attn(hidden_states, attention_mask=attention_mask,
                                              position_ids=position_ids,
                                              past_key_value=past_key_value, use_cache=use_cache)
        hidden_states = residual + attn_out                  # 残差 1

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = residual + self.mlp(hidden_states)   # 残差 2
        return hidden_states, present_kv

class DemoLLMModel(PreTrainedModel):
    config_class = DemoLLMConfig

    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
        self.layers = nn.ModuleList([DemoTransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.norm = DemoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                past_key_values=None, use_cache=None, **kwargs):
        batch_size, seq_length = input_ids.shape
        past_len = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        # 没传 position_ids 就自动生成（有 KV cache 时从 past_len 往后排）
        if position_ids is None:
            position_ids = torch.arange(past_len, seq_length + past_len,
                                        dtype=torch.long, device=input_ids.device).unsqueeze(0)

        inputs_embeds = self.embed_tokens(input_ids)
        hidden_states = self.dropout(inputs_embeds)

        # 把 (bsz, seq) 的 padding mask 扩成 (bsz,1,q,kv) 加性掩码；因果掩码在 attention 内部处理
        _expanded_mask = None
        if attention_mask is not None:
            expanded = attention_mask[:, None, None, :].expand(batch_size, 1, seq_length, seq_length + past_len)
            _expanded_mask = torch.zeros_like(expanded, dtype=hidden_states.dtype)
            _expanded_mask.masked_fill_(expanded == 0, float("-inf"))

        next_cache = [] if use_cache else None
        for i, layer in enumerate(self.layers):
            past_kv = past_key_values[i] if past_key_values is not None else None
            hidden_states, kv = layer(hidden_states, attention_mask=_expanded_mask,
                                      position_ids=position_ids, past_key_value=past_kv, use_cache=use_cache)
            if use_cache:
                next_cache.append(kv)

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=next_cache)

class DemoLLMForCausalLM(PreTrainedModel):
    config_class = DemoLLMConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = DemoLLMModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # self.model.embed_tokens.weight = self.lm_head.weight   # 权重共享（注释掉了，可选）
        self.post_init()                                          # HF 的统一权重初始化

    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                past_key_values=None, labels=None, use_cache=None, **kwargs):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask,
                             position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache)
        logits = self.lm_head(outputs.last_hidden_state)
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()       # 错位：第 t 个位置预测第 t+1 个 token
            shift_labels = labels[..., 1:].contiguous()
            loss = nn.CrossEntropyLoss()(shift_logits.view(-1, self.config.vocab_size),
                                         shift_labels.view(-1))
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values)

    # 供 HF generate() 调用：有 KV cache 时只需喂最后一个 token
    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
        if past_key_values:
            input_ids = input_ids[:, -1:]
        return {"input_ids": input_ids, "past_key_values": past_key_values,
                "attention_mask": attention_mask, "use_cache": kwargs.get("use_cache")}

预训练 (Pretrain)  ──保存 demo_llm_pretrained.pth──▶  学会基本语言规律
       │ load
       ▼
监督微调 (SFT)      ──保存 demo_llm_sft.pth──────────▶  学会跟随指令、对话
       │ load
       ▼
推理训练 (Reasoning)──保存 demo_llm_reasoning.pth─────▶  学会输出 <think>/<answer> 结构