【CISC7021】Build a Toy Llama-2 Language Model

一、背景

在本实验中，我们将准备一个采用 Llama-2架构的玩具语言模型，并评估数据集的困惑度。我们将学习如何使用 PyTorch 和 Hugging Face 库对基座语言模型进行继续预训练。其核心目标是用 42M 参数的玩具 Llama-2 完成三大任务------解码探索、困惑度评估和继续预训练。模型权重由 Karpathy 的 llama2.c 转换而来，数据来自 TinyStories 儿童故事数据集。

二、环境准备

1、依赖安装说明

!pip install torch torchvision torchaudio
!pip install transformers
!pip install datasets accelerate -q
!pip install ipywidgets

针对 Blackwell 架构显卡的 Nightly 版本：

# 先彻底卸载
!pip uninstall torch torchvision torchaudio -y
# 清理可能残留的缓存
!pip cache purge
# 重新安装支持 Blackwell 的 Nightly 版本
!pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

第二组命令针对 RTX 50 系列（Blackwell 架构）显卡，因为旧版 CUDA 12.1 编译的 PyTorch 不兼容。

2、包导入

围绕 Hugging Face 生态 + PyTorch 的常规依赖。chain 用于 group_texts 中展平 token 列表，CrossEntropyLoss 用于手动计算 PPL。

import torch  # 核心张量库
import pandas as pd
import matplotlib.pyplot as plt
import os
import glob

from transformers import LlamaForCausalLM, AutoTokenizer
from transformers import Trainer, TrainingArguments
from itertools import chain  # 用于拼接多个 token 序列
from datasets import load_dataset  # 加载 JSONL 数据

from tqdm import tqdm  # Notebook 进度条
from torch.nn import CrossEntropyLoss  # PPL 计算

抑制 transformers 的生成警告，避免 max_new_tokens 与 max_length 冲突时刷屏。

import warnings
import logging

logging.getLogger("transformers.generation.utils").setLevel(logging.ERROR)
warnings.filterwarnings("ignore", message=".*max_new_tokens.*max_length.*")

3、路径配置

# TRAIN_FILE = '/content/drive/MyDrive/assignment1/data/zh_train.jsonl'
# VALIDATION_FILE = '/content/drive/MyDrive/assignment1/data/zh_dev.jsonl'
# TEST_FILE = '/content/drive/MyDrive/assignment1/data/zh_test.jsonl'
# EN_TEST_FILE = '/content/drive/MyDrive/assignment1/data/en_test.jsonl'
# MODEL_FOLDER = "/content/drive/MyDrive/assignment1/llama-42m"
TRAIN_FILE = './data/zh_train.jsonl'
VALIDATION_FILE = './data/zh_dev.jsonl'
TEST_FILE = './data/zh_test.jsonl'
EN_TEST_FILE = './data/en_test.jsonl'
MODEL_FOLDER = "./llama-42m"  # 含 config.json / pytorch_model.bin
# 评估葡语的额外定义
PT_TEST_FILE = './data/pt_test.jsonl'
PT_TRAIN_FILE = './data/pt_train.jsonl'
PT_VALID_FILE = './data/pt_dev.jsonl'

三、模型加载

1、加载模型

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device type: {device}")  # 预期：cuda

model_path = MODEL_FOLDER
# Load model from local files
model = LlamaForCausalLM.from_pretrained(model_path).to(device)
# Load tokenizer from local files
tokenizer = AutoTokenizer.from_pretrained(model_path, device=device)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

2、模型参数统计

# Exclude embedding, only count non-embedding parameters
total_para = sum(v.numel() for k, v in model.state_dict().items() if k != 'model.embed_tokens.weight') / 1e6
print(model)
print(f"#Parameters: {total_para:.2f}M")

输出如下：

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 512)
    (layers): ModuleList(
      (0-7): 8 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=512, out_features=512, bias=False)
          (k_proj): Linear(in_features=512, out_features=512, bias=False)
          (v_proj): Linear(in_features=512, out_features=512, bias=False)
          (o_proj): Linear(in_features=512, out_features=512, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=512, out_features=1376, bias=False)
          (up_proj): Linear(in_features=512, out_features=1376, bias=False)
          (down_proj): Linear(in_features=1376, out_features=512, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((512,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((512,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((512,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=512, out_features=32000, bias=False)
)
#Parameters: 41.69M

模型与 Llama-2 完全同构，但规模缩小。8 层、512 隐维度、8 注意力头、32000 词表、512 上下文长度。非嵌入参数 41.69M（含 embedding 总计约 58M）。

四、任务1：解码

1、任务说明

最小示例：

prompt = "Once upon a time, " # Input, prefix of generation

步骤 1：使用分词器模型编码原始文本

tokenized_input = tokenizer.encode(prompt, return_tensors='pt').to(device)

步骤 2：设置解码超参数，获取模型输出

output_ids = model.generate(tokenized_input, do_sample=True, max_new_tokens=300, temperature=0.6)

重要参数：

（1）max_new_tokens：要生成的最大 token 数量，忽略 prompt 中的 token 数。

（2）temperature：用于调节下一个 token 概率的温度值。更高的温度 → 生成更多样的文本。更低的温度 → 生成更确定的文本。

（3）do_sample：do_sample=False 使用贪婪解码策略。要启用贪婪解码，我们还需要将其他采样参数 top_p、temperature 设置为 None。

步骤 3：将模型输出转换为原始文本

output_text = tokenizer.decode(output_ids[0])

或（当输入实例 >= 1 时）

output_text = tokenizer.batch_decode(output_ids)

重要参数：

设置 skip_special_tokens=True 将防止特殊 token（如 <s>）出现在结果中。

2、单条生成演示

prompt = "Once upon a time, Stella Lou had a dream."  # Feel free to use other generation prefix

# Step 1: Encode raw text using tokenizer model. Run tokenization and covert strings into token ids in vocabulary.
# tokenizer.encode() automatically adds BOS token (id=1) at the beginning
# return_tensors='pt' returns a PyTorch tensor with shape [1, seq_len]
tokenized_input = tokenizer.encode(prompt, return_tensors='pt').to(device)
# See the tokenized results.
print(tokenized_input)

输出如下：

tensor([[    1,  9038,  2501,   263,   931, 29892,   624,  3547,  4562,   750,
           263, 12561, 29889]], device='cuda:0')

开头的 1 是自动加的 BOS token。

# Step 2: Set decoding hyperparameters.
# Greedy decoding: pick the highest probability token each step, deterministic result
# For greedy decoding
max_new_tokens = 300
do_sample = False  # `do_sample=False` means using greedy decoding strategy`.
temperature = None  # To enable greedy decoding, we also need to set `top_p`, `temperature` as `None

# call generation function model.generate()
output_ids = model.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,  # Stop generation when id=1 is encountered
    do_sample=do_sample,
    temperature=temperature,
    top_p=None,  # Must also be set to None in greedy mode
)

# The decoded results are token ids.
print("=" * 20 + "Token IDs" + "=" * 20)
print(output_ids)

输出如下：

====================Token IDs====================
tensor([[    1,  9038,  2501,   263,   931, 29892,   624,  3547,  4562,   750,
           263, 12561, 29889,  2296,  5131,   304,   367,   263, 12456,   985,
         29889,  2296,  5131,   304, 19531,   263,  9560, 10714,   322,   263,
           528,  4901, 20844, 29889,  1205,  1183,   471,  2086,  2319,   322,
           278, 10714,   471,  2086,  4802, 29889,    13,  6716,  2462, 29892,
           624,  3547,  4446,   263,  4802, 29892,   528,  4901, 10714,   297,
           263,  3787, 29889,  2296,  4433,   902, 16823,   565,  1183,  1033,
           505,   372, 29889,  2439, 16823,  1497,  4874,   322, 18093,   372,
           363,   902, 29889,    13,   855,  3547,   471,   577,  9796, 29889,
          2296,  1925,   373,   278, 10714,   322,  3252,   381,   839,  2820,
         29889,  2296,  7091,   763,   263,  1855, 12456,   985, 29889,    13,
          6246,   769, 29892,  1554,  8515,  9559, 29889,   624,  3547,  4687,
           304,  4459,   270,   466,  1537, 29889,  2296,  8496, 29915, 29873,
          2317,   701,  7812, 29889,  2296,  7091,   763,  1183,   471, 10917,
          1076,  2820,   322,  2820, 29889,    13,   855,  3547, 29915, 29879,
         16823,  4446,   902,   322,  1497, 29892,   376,   855,  3547, 29892,
           366,   817,   304,  2125,   263,  2867, 29889,   887,  1106,   270,
           466,  1537,  1213,    13,   855,  3547,  3614,  1283,   278, 10714,
           322,  6568,  1623,   373,   278, 11904, 29889,  2296,  5764,   902,
          5076,   322,  3614,   263,  6483, 16172, 29889,  2860,   263,  2846,
          6233, 29892,  1183,  7091,  2253, 29889,    13,   855,  3547, 25156,
           322,  1497, 29892,   376, 29924,   290, 29892,   306, 29915, 29885,
          7960,   304,   367,   263, 12456,   985,  1449,  3850,     1]],
       device='cuda:0')

# Step 3: Convert model outputs into raw text.
# decode token ids into tokens
print("=" * 20 + "Decoded Results" + "=" * 20)
# We only have one input instance. So we directly decode the first item of model output, i.e., `output_ids[0]`.
# skip_special_tokens=True filters out special tokens like <s>, </s>
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

生成结果如下：

====================Decoded Results====================
Once upon a time, Stella Lou had a dream. She wanted to be a princess. She wanted to wear a beautiful dress and a shiny crown. But she was too small and the dress was too big.
One day, Stella saw a big, shiny dress in a store. She asked her mom if she could have it. Her mom said yes and bought it for her.
Stella was so happy. She put on the dress and twirled around. She felt like a real princess.
But then, something strange happened. Stella started to feel dizzy. She couldn't stand up straight. She felt like she was spinning around and around.
Stella's mom saw her and said, "Stella, you need to take a break. You look dizzy."
Stella took off the dress and lay down on the floor. She closed her eyes and took a deep breath. After a few minutes, she felt better.
Stella smiled and said, "Mom, I'm ready to be a princess again!"

观察可得贪婪解码输出确定、连贯，但句式比较单一。

3、采样解码演示

# Difference from greedy: sample from the probability distribution at each step, result is random
prompt = "Once upon a time, Stella Lou had a dream."

# Decoding hyperparameters
max_new_tokens = 300
do_sample = True  # True = 采样 / True = sampling
# The value of temperature used to modulate the next token probabilities.
# Higher temperature -> generate more diverse text. Lower temperature -> generate more deterministic text.
temperature = 0.3  # Low temperature → sharper distribution, still favors high-probability words

tokenized_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
output_ids = model.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,
    do_sample=do_sample,
    temperature=temperature,
)
output_text = tokenizer.decode(output_ids[0])
print(output_text)

生成结果如下：

<s> Once upon a time, Stella Lou had a dream. In her dream, she was playing in a big field with lots of flowers. She was so happy and excited!
Suddenly, Stella heard a voice. It was her mom. She said, "Stella Lou, it's time to go home now."
Stella was sad. She wanted to stay and play in the field. But she knew she had to go home.
When Stella got home, her mom said, "It's time to eat dinner." Stella was hungry, so she ate her dinner.
After dinner, Stella's mom said, "Let's go back to the field." Stella was so happy! She ran outside and played in the field.
Stella was so happy. She had a dream, and she was able to play in the field. She was so happy!<s>

do_sample=True + temperature=0.3 得到既有多样性又保持连贯的输出。注意这里没加 skip_special_tokens=True，所以输出中保留了 <s> 和 </s>。

4、单条输入

空 prompt 时模型仍会从 BOS 开始生成一个完整故事，体现了它作为语言模型的能力。

# Skeleton Code: Single input (same as previous code blocks)
prompt = ""  # ⬅️ try to construct different prompts.

# ⬇️ Try to tune different decoding hyperparameters.
# You can also add more hyperparameters like `top_p`, `top_k`.
max_new_tokens = 300
do_sample = True
temperature = 0.3

tokenized_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
output_ids = model.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,
    do_sample=do_sample,
    temperature=temperature,
)
output_text = tokenizer.decode(output_ids[0])
print(output_text)

生成结果如下：

<s> Once upon a time, there was a little girl named Lily. She loved to play with her toys and her friends. One day, Lily's mom asked her to help with the dishes. Lily didn't want to help because she wanted to play. But her mom said, "If you help me, we can play together later."
Lily helped her mom with the dishes and then they played together. After they finished, Lily's mom said, "Let's go outside and play!" They went outside and played in the sun. Lily's mom said, "Let's go to the park!" Lily was so excited and said, "Yes, let's go!"
When they got to the park, they saw a big slide. Lily's mom said, "Let's go down the slide!" Lily was scared but her mom said, "Don't worry, it's fun!" They went down the slide together and Lily laughed and laughed. She said, "That was so much fun! I'm glad we came to the park."<s>

5、批量输入

# Skeleton Code: Batched input-output

prompts = ["Once upon a time,", "Tom is a cute kitty."]  # ⬅️ try to construct different prompts.

# If you have multiple data inputs, please control the batch size to prevent out-of-memory issues.
batch_size = 2  # batch_size=2 means generating two prompts at once

# ⬇️ Try to tune different decoding hyperparameters.
# You can also add more hyperparameters like `top_p`, `top_k`.
max_new_tokens = 300
do_sample = True
temperature = 0.6

for i in range(0, len(prompts), batch_size):
    batch_input = prompts[i:i + batch_size]
    # padding=True: pad prompts of different lengths
    tokenized_input = tokenizer(batch_input, return_tensors="pt", padding=True).to(device)

    # For decoder-only models, batched inputs of model.generate() should be in the format of input_ids.
    output_ids = model.generate(
        tokenized_input["input_ids"],
        max_new_tokens=max_new_tokens,
        eos_token_id=1,
        do_sample=do_sample,
        temperature=temperature,
    )
    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)

    for idx, result in enumerate(output_text):
        print(f"{result}\n")

生成结果如下：

Once upon a time, and her mommy went to the mall. It was a very special mall, because it had a really big door with a big lock. 
When they got inside, mommy said, "Let's go to the toy store, okay?" 
But when they got to the store, mommy said, "Oh no, the door is locked!" 
Lily was very sad. She said, "Mommy, why can't we go in?" 
Mommy said, "I'm sorry, sweetheart. The door is locked, and we can't get in." 
Lily was so sad. She wanted to go in the store, but she couldn't. 
Then, mommy said, "Let's go home, and I'll make you a special treat." 
Lily was so happy. She said, "Yay!" 
And they went home, and mommy made her a delicious treat.

Tom is a cute kitty. He likes to play with his toys and his owner, Sam. Sam is a nice boy who gives Tom food and water and hugs him.
One day, Sam and Tom go to the park. There are many other kittens there. Some are big, some are small, some are fluffy, some are noisy. Tom wants to play with them, but Sam says no. He says they are not like the cats he knows.
Tom is sad. He does not understand why Sam does not let him play with the other kittens. He thinks they are his friends. He sneaks away when Sam is not looking. He finds a ball and a mouse. He chases them and has fun.
But the other kittens do not like Tom. They are scared of him. They think he is mean and noisy. They run away and hide. Tom does not see them. He is too busy chasing them.
Sam sees Tom and gets angry. He shouts at him and tells him to come back. He says he is not his friend anymore. He takes him home and puts him in a cage. He says he is a bad boy. He says he does not love him.
Tom is scared. He does not understand why Sam is mad. He does not know he is a bad boy. He wishes he had a friend. He wishes he had a real kitten. He c

解析：

（1）循环切 batch 防止 OOM

（2）padding=True 补齐不同长度的 prompt

（3）传入 input_ids 而非完整 dict

（4）batch_decode 批量转文本

6、中文 prompt 探针

# Key phenomenon: the model replies with a few Chinese characters, then abruptly switches back to English
prompt = "从前有一只小兔子乖乖"

# Decoding hyperparameters
max_new_tokens = 300
do_sample = True
temperature = 0.3

tokenized_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
output_ids = model.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,
    do_sample=do_sample,
    temperature=temperature,
)
output_text = tokenizer.decode(output_ids[0])
print(output_text)

生成结果如下：

<s> 从前有一只小兔子乖乖ons were playing in the park. They were having so much fun running around and playing tag. Suddenly, a big dog came running towards them. The dog was very angry and growled loudly.
Little Timmy was scared and he started to cry. His mom came running over and said, "Don't worry Timmy, I'm here to protect you." She held Timmy close and the dog ran away.
Timmy was so happy and he hugged his mom tightly. She said, "It's okay, I'm here to protect you."
Timmy looked up at his mom and said, "Mom, I'm sorry for being scared. I'm so happy you're here."
His mom smiled and said, "It's okay Timmy, I'm here to protect you." She hugged him again and they continued to play in the park.<s>

这是最关键的诊断性实验。模型只能勉强输出几个汉字，然后立刻切换回英文------直接证明基座模型不具备中文能力，为 Task 2 的 PPL 分析埋下伏笔。

7、客观指标与实验矩阵

# Explore different prompts and decoding parameters

# ---------- Define objective metrics ----------
def compute_ttr(text):
    """Type-Token Ratio: unique words / total words"""
    words = text.split()
    return len(set(words)) / len(words) if words else 0


def compute_repetition(text):
    """Repetition rate: proportion of repeated 3-grams"""
    words = text.split()
    trigrams = [tuple(words[i:i + 3]) for i in range(len(words) - 2)]
    if not trigrams:
        return 0
    return 1 - len(set(trigrams)) / len(trigrams)


# ---------- Experiment matrix ----------
prompts = [
    "Once upon a time,",
    "Tom is a cute kitty.",
    "The little girl was",
    "In a dark forest,",
]

configs = [
    {"name": "greedy", "do_sample": False, "temperature": None},
    {"name": "T=0.1", "do_sample": True, "temperature": 0.1},
    {"name": "T=0.6", "do_sample": True, "temperature": 0.6},
    {"name": "T=1.0", "do_sample": True, "temperature": 1.0},
    {"name": "T=1.5", "do_sample": True, "temperature": 1.5},
    {"name": "T=0.6+k50", "do_sample": True, "temperature": 0.6, "top_k": 50},
    {"name": "T=0.6+p0.9", "do_sample": True, "temperature": 0.6, "top_p": 0.9},
]

# ---------- Loop generation ----------
results = []
for p in prompts:
    for c in configs:
        tokenized_input = tokenizer.encode(p, return_tensors="pt").to(device)
        kwargs = dict(
            max_new_tokens=300, eos_token_id=1,
            do_sample=c["do_sample"], temperature=c["temperature"],
        )
        if "top_k" in c: kwargs["top_k"] = c["top_k"]
        if "top_p" in c: kwargs["top_p"] = c["top_p"]

        output_ids = model.generate(tokenized_input, **kwargs)
        text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

        results.append({
            "prompt": p,
            "config": c["name"],
            "length": len(text.split()),
            "ttr": round(compute_ttr(text), 3),
            "repetition": round(compute_repetition(text), 3),
        })

df_task1 = pd.DataFrame(results)
df_task1

输出如下：

prompt	config	length	ttr	repetition
Once upon a time,	greedy	141	0.645	0.007
Once upon a time,	T=0.1	187	0.610	0.000
Once upon a time,	T=0.6	127	0.622	0.016
Once upon a time,	T=1.0	156	0.590	0.019
Once upon a time,	T=1.5	198	0.631	0.010
Once upon a time,	T=0.6+k50	181	0.558	0.017
Once upon a time,	T=0.6+p0.9	158	0.620	0.000
Tom is a cute kitty.	greedy	177	0.480	0.051
Tom is a cute kitty.	T=0.1	226	0.451	0.085
Tom is a cute kitty.	T=0.6	223	0.498	0.027
Tom is a cute kitty.	T=1.0	230	0.448	0.066
Tom is a cute kitty.	T=1.5	245	0.506	0.045
Tom is a cute kitty.	T=0.6+k50	216	0.519	0.033
Tom is a cute kitty.	T=0.6+p0.9	180	0.556	0.022
The little girl was	greedy	168	0.500	0.036
The little girl was	T=0.1	124	0.556	0.033
The little girl was	T=0.6	146	0.527	0.069
The little girl was	T=1.0	149	0.611	0.034
The little girl was	T=1.5	177	0.605	0.017
The little girl was	T=0.6+k50	146	0.493	0.104
The little girl was	T=0.6+p0.9	210	0.486	0.087
In a dark forest,	greedy	118	0.585	0.026
In a dark forest,	T=0.1	100	0.620	0.020
In a dark forest,	T=0.6	145	0.676	0.035
In a dark forest,	T=1.0	197	0.543	0.031
In a dark forest,	T=1.5	182	0.632	0.006
In a dark forest,	T=0.6+k50	114	0.561	0.045
In a dark forest,	T=0.6+p0.9	137	0.620	0.000

解析如下：

（1）TTR：类型-标记比，越高越多样

（2）Repetition：3-gram 重复率，越高越冗余

（3）实验矩阵：4 prompt × 7 config = 28 组

（4）通过 **kwargs 灵活注入 top_k / top_p

关键结果如下：

（1）贪婪解码长度最短但 TTR 也低

（2）T=1.5 TTR 最高但语法错误增多

（3）T=0.6+k50 在多样性和重复率间取得最佳平衡

（4）The little girl was 在 T=0.6 时重复率高达 0.104

8、可视化

# Task 1 TTR and repetition visualization
# Output: one horizontally arranged twin subplot
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# TTR comparison: one curve per prompt, x-axis is config
for p in prompts:
    sub = df_task1[df_task1["prompt"] == p]
    axes[0].plot(sub["config"], sub["ttr"], marker="o", label=p[:15])
axes[0].set_title("TTR across decoding configs")
axes[0].set_ylabel("Type-Token Ratio")
axes[0].tick_params(axis="x", rotation=45)
axes[0].legend(fontsize=8)
axes[0].grid(True)

# Repetition comparison: one curve per prompt
for p in prompts:
    sub = df_task1[df_task1["prompt"] == p]
    axes[1].plot(sub["config"], sub["repetition"], marker="s", label=p[:15])
axes[1].set_title("Repetition across decoding configs")
axes[1].set_ylabel("Repetition rate (3-gram)")
axes[1].tick_params(axis="x", rotation=45)
axes[1].legend(fontsize=8)
axes[1].grid(True)

plt.tight_layout()
plt.show()

图表对比如下：

五、任务2：困惑度评估

1、背景说明

困惑度是评估语言模型的关键指标。它量化模型预测样本的好坏，困惑度越低表示性能越好。对于分词序列 ，困惑度的数学定义为：

这里表示给定前序 token 时 token 的概率，该公式包含序列上的平均对数概率。

2、PPL 计算函数

# The following code was adapted from the `evaluate` library. Licensed under the Apache License, Version 2.0 (the "License").
# We modify them to avoid causing serious memory issues in the Colab environment.

def compute_ppl(
        model, tokenizer, inputs, device, batch_size: int = 16, add_start_token: bool = True, max_length=None
):
    # Parameter validation
    if device is not None:
        assert device in ["gpu", "cpu", "cuda"], "device should be either gpu or cpu."
        # device 应为 gpu 或 cpu
        if device == "gpu":
            device = "cuda"
    else:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    # if batch_size > 1 (which generally leads to padding being required), and
    # if there is not an already assigned pad_token, assign an existing special token to also be the padding token
    if tokenizer.pad_token is None and batch_size > 1:
        existing_special_tokens = list(tokenizer.special_tokens_map_extended.values())
        # check that the model already has at least one special token defined
        assert (
                len(existing_special_tokens) > 0
        ), "If batch_size > 1, model must have at least one special token to use for padding. Please use a different model or set batch_size=1."
        # assign one of the special tokens to also be the pad token
        tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})

    if add_start_token and max_length:
        # leave room for <BOS> token to be added:
        assert (
                tokenizer.bos_token is not None
        ), "Input model must already have a BOS token if using add_start_token=True. Please use a different model, or set add_start_token=False"
        max_tokenized_len = max_length - 1
    else:
        max_tokenized_len = max_length

    # Encode input text (without special tokens, since BOS is added manually later)
    encodings = tokenizer(
        inputs,
        add_special_tokens=False,
        padding=True,
        truncation=True if max_tokenized_len else False,
        max_length=max_tokenized_len,
        return_tensors="pt",
        return_attention_mask=True,
    )

    encoded_texts = encodings["input_ids"]
    attn_masks = encodings["attention_mask"]

    # check that each input is long enough:
    # Each input has at least 1 token
    if add_start_token:
        assert torch.all(torch.ge(attn_masks.sum(1), 1)), "Each input text must be at least one token long."
    else:
        assert torch.all(
            torch.ge(attn_masks.sum(1), 2)
        ), "When add_start_token=False, each input text must be at least two tokens long. Run with add_start_token=True if inputting strings of only one token, and remove all empty input strings."

    ppls = []
    loss_fct = CrossEntropyLoss(reduction="none")  # Retain per-token loss

    for start_index in tqdm(range(0, len(encoded_texts), batch_size)):
        end_index = min(start_index + batch_size, len(encoded_texts))
        encoded_batch = encoded_texts[start_index:end_index].to(device)
        attn_mask = attn_masks[start_index:end_index].to(device)

        # Manually add BOS at the beginning
        if add_start_token:
            bos_tokens_tensor = torch.tensor([[tokenizer.bos_token_id]] * encoded_batch.size(dim=0)).to(device)
            encoded_batch = torch.cat([bos_tokens_tensor, encoded_batch], dim=1)
            attn_mask = torch.cat(
                [torch.ones(bos_tokens_tensor.size(), dtype=torch.int64).to(device), attn_mask], dim=1
            )

        labels = encoded_batch

        with torch.no_grad():
            out_logits = model(encoded_batch, attention_mask=attn_mask).logits

            # Causal LM shift: position t predicts position t+1
            shift_logits = out_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            shift_attention_mask_batch = attn_mask[..., 1:].contiguous()

            # PPL = exp(average negative log-likelihood)
            perplexity_batch = torch.exp(
                (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).sum(1)
                / shift_attention_mask_batch.sum(1)
            )

            ppls += perplexity_batch.tolist()

    del encoded_batch, attn_mask
    if device == "cuda":
        torch.cuda.empty_cache()

    return {"perplexities": ppls, "mean_perplexity": sum(ppls) / float(len(ppls))}

（1）最小示例

test_dataset = ["Once upon a time,"]

compute_ppl(
    model=model,
    tokenizer=tokenizer,
    device=device,
    inputs=test_dataset,
    batch_size = 16
)

（2）重要参数

① inputs: list of input text, each separate text snippet is one list entry.

② inputs：输入文本列表，每个独立的文本片段是一个列表项。

③ batch_size: the batch size to run evaluations.

④ batch_size：运行评估的批量大小。

（3）返回值

perplexity：{"perplexities": x.x, x.x, ..., "mean_perplexity": x.x} 字典，包含输入列表中文本的困惑度分数以及平均困惑度。

3、PPL 核心代码

# Skeleton Code: Evaluate the perplexity (PPL) on a list of raw text.

test_dataset = ["Once upon a time,",
                "Tom is a cute kitty."]  # ⬅️ you can use your examples / or read from raw text file

results = compute_ppl(model=model, tokenizer=tokenizer, device=device, inputs=test_dataset, batch_size=16)
dataset_ppl = results['mean_perplexity']
print(f"Perplexity: {dataset_ppl:.2f}")

# Skeleton Code: Evaluate the perplexity (PPL) on an external test set file (JSONL).

# English test set.
data_file = EN_TEST_FILE  # ⬅️ you can change your file path
test_dataset = list(load_dataset('json', data_files={'test': data_file})["test"]["text"])

results = compute_ppl(model=model, tokenizer=tokenizer, device=device, inputs=test_dataset, batch_size=16)
dataset_ppl = results['mean_perplexity']
print(f"(English Text) Test Perplexity: {dataset_ppl:.2f}")

# Chinese test set.（token 长，batch_size 调小）
data_file = TEST_FILE  # ⬅️ you can change your file path
test_dataset = list(load_dataset('json', data_files={'test': data_file})["test"]["text"])

results = compute_ppl(model=model, tokenizer=tokenizer, device=device, inputs=test_dataset, batch_size=4)
dataset_ppl = results['mean_perplexity']
print(f"(Chinese Text) Test Perplexity: {dataset_ppl:.2f}")

# Portuguese test set.（葡语 token 数比英文多，batch_size 用 8（OOM 时降到 4））
data_file = PT_TEST_FILE
test_dataset = list(load_dataset('json', data_files={'test': data_file})["test"]["text"])

results = compute_ppl(model=model, tokenizer=tokenizer, device=device, inputs=test_dataset, batch_size=8)
dataset_ppl = results['mean_perplexity']
print(f"(Portuguese Text) Test Perplexity: {results['mean_perplexity']:.2f}")

# Try your own data file!

输出如下：

Perplexity: 10.68
(English Text) Test Perplexity: 4.14
(Chinese Text) Test Perplexity: 70030.35
(Portuguese Text) Test Perplexity: 25050.81

解析如下：

（1）批大小分层：英文 16 / 葡萄牙语 8 / 中文 4，避免 OOM（中文 token 序列最长）

（2）结果验证了作者预言：中文 PPL 比英文高 4 个数量级，葡语介于中间

4、多领域 PPL 扩展

# Task 2 multi-domain, multilingual PPL evaluation, collecting more diverse text

extra_data = {"English - News": [
    "The president announced a new policy yesterday.",
    "Stock markets rallied on positive earnings reports.",
    "Scientists discovered a new species in the Amazon.",
], "English - Dialogue": [
    "Hello, how are you today?",
    "I'm fine, thank you. And you?",
    "Could you please pass me the salt?",
],
    "English - Poetry": [
        "The road not taken winds through yellow wood.",
        "Shall I compare thee to a summer's day?",
        "Two roads diverged in a yellow wood.",
    ],
    "Chinese - News": [
        "总统昨天宣布了一项新政策。",
        "股市因利好的财报而上涨。",
        "科学家在亚马逊发现了一个新物种。",
    ],
    "Chinese - Dialogue": [
        "你好，你今天怎么样？",
        "我很好，谢谢。你呢？",
        "请把盐递给我好吗？",
    ],
    "Portuguese - News": [
        "O presidente anunciou uma nova política ontem.",
        "Os mercados subiram com relatórios positivos.",
        "Cientistas descobriram uma nova espécie na Amazônia.",
    ],
}

# Add the original English, Chinese and Portuguese TinyStories test sets (first 200 each)
extra_data["English - TinyStories"] = load_dataset(
    'json', data_files={'test': EN_TEST_FILE})["test"]["text"][:200]
extra_data["Chinese - TinyStories"] = load_dataset(
    'json', data_files={'test': TEST_FILE})["test"]["text"][:200]
extra_data["Portuguese - TinyStories"] = list(load_dataset(
    'json', data_files={'test': PT_TEST_FILE})["test"]["text"])[:200]

# Iterate and compute
ppl_results = []
for domain, texts in extra_data.items():
    # Select batch_size based on language
    if "English" in domain:
        bs = 16
    elif "Portuguese" in domain:
        bs = 8
    else:
        bs = 4
    r = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                    inputs=texts, batch_size=bs)
    ppl_results.append({
        "domain": domain,
        "mean_ppl": round(r["mean_perplexity"], 2),
        "median_ppl": round(sorted(r["perplexities"])[len(r["perplexities"]) // 2], 2),
        "max_ppl": round(max(r["perplexities"]), 2),
        "min_ppl": round(min(r["perplexities"]), 2),
        "n": len(texts),
    })

df_task2 = pd.DataFrame(ppl_results)
df_task2

输出如下：

domain	mean_ppl	median_ppl	max_ppl	min_ppl	n
English - News	14592.04	3815.55	38110.80	1849.76	3
English - Dialogue	409.97	209.92	864.89	155.09	3
English - Poetry	7167.12	4109.78	14093.12	3298.46	3
Chinese - News	96239.73	98418.82	98468.02	91832.36	3
Chinese - Dialogue	126390.04	114709.27	157660.42	106800.43	3
Portuguese - News	253455.61	182352.88	441655.44	136358.52	3
English - TinyStories	2.50	2.43	6.05	1.68	200
Chinese - TinyStories	70815.81	71543.40	108639.77	23556.29	200
Portuguese - TinyStories	25073.46	23455.91	72198.93	5396.22	200

解析如下：

（1）自收集新闻、对话、诗歌三个领域数据

（2）每个领域都报 mean/median/max/min，能看出长尾分布

（3）结果显示模型在训练分布（TinyStories）上 PPL 最低，分布外（新闻、诗歌）急剧升高

5、三语 PPL 对比表

# Trilingual PPL comparison (before training)
# English
data_en = list(load_dataset('json', data_files={'test': EN_TEST_FILE})["test"]["text"])
en_ppl_before = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                            inputs=data_en, batch_size=16)["mean_perplexity"]

# Chinese
data_zh = list(load_dataset('json', data_files={'test': TEST_FILE})["test"]["text"])
zh_ppl_before = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                            inputs=data_zh, batch_size=4)["mean_perplexity"]

# Portuguese
data_pt = list(load_dataset('json', data_files={'test': PT_TEST_FILE})["test"]["text"])
pt_ppl_before = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                            inputs=data_pt, batch_size=8)["mean_perplexity"]

df_multiling = pd.DataFrame({
    "language": ["English", "Chinese", "Portuguese"],
    "before_training": [
        round(en_ppl_before, 2),
        round(zh_ppl_before, 2),
        round(pt_ppl_before, 2),
    ],
})
df_multiling

输出如下：

language	before_training
English	4.14
Chinese	70030.35
Portuguese	25050.81

表格保存训练前的三语 PPL 作为基线，用于和 Task 3 训练后对比，量化灾难性遗忘程度。

6、多领域 PPL 可视化

# Multi-domain PPL comparison (log scale)
plt.figure(figsize=(12, 5))
colors = []
for d in df_task2["domain"]:
    if "English" in d:
        colors.append("#4C72B0")
    elif "Chinese" in d:
        colors.append("#C44E52")
    else:
        colors.append("#55A868")
plt.bar(df_task2["domain"], df_task2["mean_ppl"], color=colors)
plt.yscale("log")
plt.ylabel("Mean Perplexity (log scale)")
plt.title("PPL across languages and domains (before continual pre-training)")
plt.xticks(rotation=45, ha="right")
plt.tight_layout()
plt.show()

图表对比如下：

7、释放显存

# 🚨 Release gpu cache before training the model
if device == "cuda":
    torch.cuda.empty_cache()

六、任务 3：继续预训练

1、任务说明

目前，英文基准语言模型擅长英文，但缺乏生成或理解其他语言（例如中文）的能力。本实验任务的目标是通过用另一种语言的文本继续预训练来增强英文基准语言模型。这个过程旨在使模型能够理解并生成另一种语言的迷你故事。

2、数据预处理说明

这里为后续的评估和预训练阶段预处理（分词和分组）文本：

chinese_dataset = load_dataset('json',
                               data_files={'train': TRAIN_FILE, 'validation': VALIDATION_FILE, 'test': TEST_FILE})
print(chinese_dataset)
print(chinese_dataset["test"][2]["text"])

输出如下：

DatasetDict({
    train: Dataset({
        features: ['text'],
        num_rows: 10000
    })
    validation: Dataset({
        features: ['text'],
        num_rows: 500
    })
    test: Dataset({
        features: ['text'],
        num_rows: 1000
    })
})
从前，有一个小女孩名叫莉莉。她喜欢和家人一起去度假。有一天，她的家人决定去海边旅行。莉莉非常兴奋，她跳起来又跳下去，像发了疯一样。

当他们到达海滩时，他们搭起了遮阳伞和毯子。莉莉想立刻去游泳，但她的父母告诉她要等吃完午餐再说。莉莉感到很不耐烦，她说："我现在就想去游泳！"她妈妈回答："莉莉，我们需要先吃东西。游泳需要能量。"

莉莉意识到妈妈说得对，于是耐心地等待午餐结束。她学会了有时候要克制自己的激动情绪，并听从父母的意见。从那天起，莉莉变得更善于倾听，也更加享受她的假期时光。

3、分组函数

我们使用 Llama-2 的分词器对原始文本进行分词，并将分词后的文本分组作为输入：

block_size = 512


def group_texts(examples):
    # Concatenate all tokens from all samples in the batch into one long sequence
    concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # Truncate to an integer multiple of block_size
    if total_length >= block_size:
        total_length = (total_length // block_size) * block_size
    # Split into multiple blocks of size block_size
    result = {
        k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    # For causal LM, labels are input_ids themselves (shifted internally)
    result["labels"] = result["input_ids"].copy()
    return result

解析如下：

（1）拼接切块策略：把 batch 内所有 token 拼成长序列，再切成长度 512 的密集块，避免 padding 浪费

（2）chain(*examplesk)：展平 list of lists

（3）labels = input_ids.copy()：LlamaForCausalLM 内部自动 shift

# tokenize + group
# On Windows, num_proc must be changed to 1, otherwise multi-process errors will occur
tokenized_zh_datasets = chinese_dataset.map(
    lambda examples: tokenizer(examples["text"]),
    batched=True,
    num_proc=None,  # Changed to None on Windows (original was 4), disable multiprocessing
    remove_columns=["text"])
lm_datasets = tokenized_zh_datasets.map(
    group_texts,
    batched=True,
    batch_size=512,
    num_proc=None,  # Changed to None on Windows (original was 4), disable multiprocessing
)

Windows 环境下 num_proc=4 会触发多进程 pickle 错误，改为 None（单进程）即可。

4、教程：TrainingArguments()

重要训练超参数如下：

（1）learning_rate：优化器的初始学习率。

（2）num_train_epochs：要执行的训练轮数总数（如果不是整数，将在停止训练前执行最后一个 epoch 的小数部分百分比）。

（3）*_strategy：训练期间采用的评估/保存策略。可能的值有：

① "no"：训练期间不进行任何评估/保存。

② "steps"：每 eval_steps 步进行（并记录）评估/保存。

③ "epoch"：在每个 epoch 结束时进行评估/保存。

（4）per_device_train_batch_size：训练时每个 GPU/XPU/TPU/MPS/NPU 核心/CPU 的批量大小。

（5）per_device_eval_batch_size：评估时每个 GPU/XPU/TPU/MPS/NPU 核心/CPU 的批量大小。

（6）save_total_limit：如果传入一个值，将限制检查点的总数。

优化器超参数如下：

（1）weight_decay：应用于 AdamW 优化器中除所有 bias 和 LayerNorm 权重之外的所有层的权重衰减（如果非零）。

（2）adam_beta1：AdamW 优化器的 beta1 超参数。

（3）adam_beta2：AdamW 优化器的 beta2 超参数。

学习调度如下：

（1）lr_scheduler：要使用的调度器类型。

（2）warmup_ratio：从 0 到 learning_rate 线性预热的训练总步数比例。

5、训练超参数

# =========Pre-training hyperparameters, please feel free to tune them~=========
# =Important=
lr = 1e-4  # Learning rate
epochs = 4  # Training epochs (reduced from 8 to 4)
save_steps = 100  # Save/eval every 100 steps (reduced from 200 to 100)
strategy = "steps"  # Evaluate by steps
train_bsz = 16  # Training batch size (reduce to 16 or 8 if out of memory) (reduced from 32 to 16)
eval_bsz = 8  # Eval batch size (reduced from 16 to 8)

# If you do not understand AdamW optimizer and learning scheduler, you may use default settings.
# =Optimizer=
optimizer = "adamw_torch"  # PyTorch native AdamW
weight_decay = 0.01
adam_beta1 = 0.9  # Llama paper recommended value
adam_beta2 = 0.98  # Llama paper recommended value (default is 0.999)
# =Learning scheduler=
lr_scheduler = "linear"  # Linear decay
warmup_steps = 25  # 1% warmup ≈ 25 steps
# =========End of pre-training hyperparameters=========


training_args = TrainingArguments(
    "llama-42m-zh-fairytales",  # Output directory
    eval_strategy=strategy,
    eval_steps=save_steps,
    save_strategy=strategy,
    save_steps=save_steps,
    logging_strategy="steps",
    logging_steps=10,
    learning_rate=lr,
    weight_decay=weight_decay,
    seed=42,
    per_device_train_batch_size=train_bsz,
    per_device_eval_batch_size=eval_bsz,
    save_total_limit=1,  # Keep only the latest 1 checkpoint
    optim=optimizer,
    lr_scheduler_type=lr_scheduler,
    adam_beta1=adam_beta1,
    adam_beta2=adam_beta2,
    warmup_steps=warmup_steps,
    num_train_epochs=epochs,
    report_to=[],
    dataloader_pin_memory=True,  # Enable in GPU environments, accelerates data transfer
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_datasets["train"],
    eval_dataset=lm_datasets["validation"],
)

trainer.train()

解析如下：

（1）作者已将原默认值调小（epochs 8→4，save_steps 200→100，train_bsz 32→16），保证不会 OOM

（2）adam_beta2=0.98 采用 Llama 论文推荐值

（3）report_to=\[\]（而非 None）是新版 transformers 兼容性要求

（4）dataloader_pin_memory=True 加速 GPU 数据传输

输出训练 1996 步，训练损失 4.65→1.19，验证损失 4.61→1.30。

6、超参数对比实验

# Compare multiple hyperparameter groups, each group takes about 6-10 minutes in GPU environment, run 2-3 groups depending on time

def run_experiment(exp_name, lr, epochs, train_bsz, eval_bsz=8):
    """Train and return log_history, each group loads the base model independently"""
    _model = LlamaForCausalLM.from_pretrained(MODEL_FOLDER).to(device)
    _tok = AutoTokenizer.from_pretrained(MODEL_FOLDER)
    if _tok.pad_token is None:
        _tok.pad_token = _tok.eos_token

    args = TrainingArguments(
        f"exp_{exp_name}",
        eval_strategy="steps", eval_steps=100,
        save_strategy="steps", save_steps=100,
        logging_strategy="steps", logging_steps=10,
        learning_rate=lr, weight_decay=0.01, seed=42,
        per_device_train_batch_size=train_bsz,
        per_device_eval_batch_size=eval_bsz,
        save_total_limit=1, optim="adamw_torch",
        lr_scheduler_type="linear",
        warmup_steps=25,
        num_train_epochs=epochs,
        report_to=[],
        dataloader_pin_memory=True,
    )
    trainer = Trainer(
        model=_model, args=args,
        train_dataset=lm_datasets["train"],
        eval_dataset=lm_datasets["validation"],
    )
    trainer.train()
    return trainer.state.log_history


# ---- Three comparisons ----
log_A = run_experiment("A_lr1e-4_ep4_bs16", lr=1e-4, epochs=4, train_bsz=16)
log_B = run_experiment("B_lr5e-5_ep4_bs16", lr=5e-5, epochs=4, train_bsz=16)
log_C = run_experiment("C_lr1e-4_ep4_bs32", lr=1e-4, epochs=4, train_bsz=32)  # GPU can use bs32


def summarize(log_history, name):
    train_loss = [e["loss"] for e in log_history if "loss" in e and "eval_loss" not in e]
    val_entries = [(e["step"], e["eval_loss"]) for e in log_history if "eval_loss" in e]
    final_train = round(train_loss[-1], 4) if train_loss else None
    final_val = round(val_entries[-1][1], 4) if val_entries else None
    min_val = round(min([v for _, v in val_entries]), 4) if val_entries else None
    return {"experiment": name, "final_train_loss": final_train,
            "final_val_loss": final_val, "min_val_loss": min_val}


df_task3 = pd.DataFrame([
    summarize(log_A, "A: lr=1e-4, bs=16"),
    summarize(log_B, "B: lr=5e-5, bs=16"),
    summarize(log_C, "C: lr=1e-4, bs=32"),
])
df_task3

输出如下：

实验 / Exp	最终训练损失 / Final Train Loss	最终验证损失 / Final Val Loss
A: lr=1e-4, bs=16	1.1521	1.2761
B: lr=5e-5, bs=16	1.5148	1.5779
C: lr=1e-4, bs=32	1.4486	1.4946

每次 run_experiment 都重新加载基准模型，防止实验间交叉污染。

7、训练曲线可视化

# Compare training curves of three experiments
def extract_curve(log_history):
    steps, train_loss, val_loss = [], [], []
    for entry in log_history:
        if "loss" in entry and "eval_loss" not in entry:
            steps.append(entry["step"])
            train_loss.append(entry["loss"])
        if "eval_loss" in entry:
            val_loss.append((entry["step"], entry["eval_loss"]))
    return steps, train_loss, val_loss


plt.figure(figsize=(10, 6))
for name, log in [("A: lr=1e-4, bs=16", log_A),
                  ("B: lr=5e-5, bs=16", log_B),
                  ("C: lr=1e-4, bs=32", log_C)]:
    steps, train_loss, val_loss = extract_curve(log)
    plt.plot(steps, train_loss, label=f"{name} (train)", linestyle="--")
    if val_loss:
        vs, vl = zip(*val_loss)
        plt.plot(vs, vl, label=f"{name} (val)", marker="o")

plt.xlabel("Step")
plt.ylabel("Loss")
plt.legend(fontsize=10, loc="upper right")
plt.title("Training curves across hyperparameter settings")
plt.grid(True)
plt.show()

图表对比如下：

8、加载最优 checkpoint 并评估

加载预训练模型，尝试用另一种语言生成迷你故事：

# If there are multiple experiments, select the best one
new_model_path = "llama-42m-zh-fairytales/checkpoint-1996"  # saved checkpoint path
model = LlamaForCausalLM.from_pretrained(new_model_path).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_path, device=device)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

再次评估中文文本（或其他语言）的 PPL，可以注意到，在继续预训练之后，实际上获得了低得多的 PPL。

data_file = TEST_FILE
test_dataset = list(load_dataset('json', data_files={'test': data_file})["test"]["text"])

results = compute_ppl(model=model, tokenizer=tokenizer, device=device, inputs=test_dataset, batch_size=16)
dataset_ppl = results['mean_perplexity']
print(f"Test Perplexity: {dataset_ppl:.2f}")

输出如下：

Test Perplexity: 3.58

9、灾难性遗忘量化

# PPL evaluation after Chinese training
data_zh = list(load_dataset('json', data_files={'test': TEST_FILE})["test"]["text"])
zh_ppl_after = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                           inputs=data_zh, batch_size=4)["mean_perplexity"]

data_en = list(load_dataset('json', data_files={'test': EN_TEST_FILE})["test"]["text"])
en_ppl_after = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                           inputs=data_en, batch_size=16)["mean_perplexity"]

data_pt = list(load_dataset('json', data_files={'test': PT_TEST_FILE})["test"]["text"])
pt_ppl_after = compute_ppl(model=model, tokenizer=tokenizer, device=device,
                           inputs=data_pt, batch_size=8)["mean_perplexity"]

df_zh_ppl = pd.DataFrame({
    "language": ["English", "Chinese", "Portuguese"],
    "before_training": [
        round(en_ppl_before, 2),
        round(zh_ppl_before, 2),
        round(pt_ppl_before, 2),
    ],
    "after_zh_training": [
        round(en_ppl_after, 2),
        round(zh_ppl_after, 2),
        round(pt_ppl_after, 2),
    ],
})
df_zh_ppl

输出如下：

语言 / Language	训练前 / Before	中文训练后 / After ZH Training
英文 / English	4.14	304.68
中文 / Chinese	70030.35	3.58
葡萄牙语 / Portuguese	25050.81	729813.11

中文 PPL 骤降 19562 倍，但英文和葡萄牙语 PPL 均大幅上升，说明 42M 模型容量不足以同时容纳多种语言。

10、PPL 前后对比可视化

# PPL before/after comparison bar chart (log scale)
x = range(len(df_zh_ppl))
width = 0.35

plt.figure(figsize=(8, 5))
plt.bar([i - width / 2 for i in x], df_zh_ppl["before_training"],
        width, label="Before", color="#4C72B0")
plt.bar([i + width / 2 for i in x], df_zh_ppl["after_zh_training"],
        width, label="After (Chinese training)", color="#C44E52")
plt.xticks(x, df_zh_ppl["language"])
plt.yscale("log")
plt.ylabel("Perplexity (log scale)")
plt.title("PPL before vs after Chinese continual pre-training")
plt.legend()
plt.tight_layout()
plt.show()

图表对比如下：

原始英文基准模型在 200 万条数据样本上预训练。考虑到只使用了 10,000 条训练样本（原始预训练数据的 0.5%），模型可以生成一些流畅的句子，但在长文本生成或其他语言的常识方面可能仍然吃力，可以根据计算资源尝试使用更多数据或训练步数。

11、中文生成测试

prompt = "从前，有一只叫做汤姆的猫。汤姆和他的朋友们一起玩。"

# Decoding hyperparameters
max_new_tokens = 300
do_sample = True
temperature = 0.3

tokenized_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
output_ids = model.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,
    do_sample=do_sample,
    temperature=temperature,
)
output_text = tokenizer.decode(output_ids[0])
print(output_text)

输出如下：

<s> 从前，有一只叫做汤姆的猫。汤姆和他的朋友们一起玩。他们喜欢一起玩。有一天，汤姆和他的朋友们在公园里玩。
有一天，汤姆和他的朋友们决定去散步。他们在那里看到了一个大大的、可怕的玩具。汤姆说："嗨，汤姆！我们可以玩吗？" 汤姆说："好的，我们可以玩！"
汤姆和他的朋友们玩了一整天。他们玩了一整天。他们在一起玩得很开心。他们玩得很开心。他们玩得很开心。
...

中文流利但存在明显重复（"他们玩得很开心"多次出现），体现 42M 模型的容量瓶颈。

12、葡语训练完整流程

# Portuguese data loading and preprocessing
portuguese_dataset = load_dataset('json',
                                  data_files={'train': PT_TRAIN_FILE,
                                              'validation': PT_VALID_FILE,
                                              'test': PT_TEST_FILE})
print(portuguese_dataset)

tokenized_pt_datasets = portuguese_dataset.map(
    lambda examples: tokenizer(examples["text"]),
    batched=True, num_proc=None, remove_columns=["text"])
lm_pt_datasets = tokenized_pt_datasets.map(
    group_texts, batched=True, batch_size=512, num_proc=None)

复用 Task 3 中文部分的 group_texts 和分词逻辑，只换数据文件。

# Portuguese model training (load base model independently to avoid Chinese training contamination)
model_pt = LlamaForCausalLM.from_pretrained(MODEL_FOLDER).to(device)

# Training hyperparameters
lr_pt = 1e-4
epochs_pt = 4
train_bsz_pt = 16
eval_bsz_pt = 8
save_steps_pt = 100
warmup_steps_pt = 25

training_args_pt = TrainingArguments(
    "llama-42m-pt-fairytales",  # Portuguese training output directory
    eval_strategy="steps",
    eval_steps=save_steps_pt,
    save_strategy="steps",
    save_steps=save_steps_pt,
    logging_strategy="steps",
    logging_steps=10,
    learning_rate=lr_pt,
    weight_decay=0.01,
    seed=42,
    per_device_train_batch_size=train_bsz_pt,
    per_device_eval_batch_size=eval_bsz_pt,
    save_total_limit=1,
    optim="adamw_torch",
    lr_scheduler_type="linear",
    adam_beta1=0.9,
    adam_beta2=0.98,
    warmup_steps=warmup_steps_pt,
    num_train_epochs=epochs_pt,
    report_to=[],
    dataloader_pin_memory=True,
)

trainer_pt = Trainer(
    model=model_pt,
    args=training_args_pt,
    train_dataset=lm_pt_datasets["train"],
    eval_dataset=lm_pt_datasets["validation"],
)

trainer_pt.train()

重新加载基座模型（from_pretrained(MODEL_FOLDER)），保证葡语实验不会受中文训练污染。超参数直接沿用实验 A 的配置。

训练 1256 步，训练损失 3.14→1.45，验证损失 3.07→1.69。

# Load Portuguese post-training checkpoint
print("葡语训练目录内容:", os.listdir("llama-42m-pt-fairytales"))
# Portuguese training directory contents:

ckpts_pt = sorted(glob.glob("llama-42m-pt-fairytales/checkpoint-*"))
pt_model_path = ckpts_pt[-1] if ckpts_pt else "llama-42m-pt-fairytales"
print(f"加载: {pt_model_path}")  # Loading:

model_pt = LlamaForCausalLM.from_pretrained(pt_model_path).to(device)
tokenizer = AutoTokenizer.from_pretrained(MODEL_FOLDER, device=device)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

用 glob.glob 动态查找最新 checkpoint，比硬编码更稳健。

# PPL evaluation after Portuguese training
data_pt = list(load_dataset('json', data_files={'test': PT_TEST_FILE})["test"]["text"])
pt_ppl_after = compute_ppl(model=model_pt, tokenizer=tokenizer, device=device,
                           inputs=data_pt, batch_size=8)["mean_perplexity"]

data_en = list(load_dataset('json', data_files={'test': EN_TEST_FILE})["test"]["text"])
en_ppl_after_pt = compute_ppl(model=model_pt, tokenizer=tokenizer, device=device,
                              inputs=data_en, batch_size=16)["mean_perplexity"]

df_pt_ppl = pd.DataFrame({
    "language": ["English", "Portuguese"],
    "before_training": [round(en_ppl_before, 2), round(pt_ppl_before, 2)],
    "after_pt_training": [round(en_ppl_after_pt, 2), round(pt_ppl_after, 2)],
})
display(df_pt_ppl)

输出如下：

语言 / Language	训练前 / Before	葡语训练后 / After PT Training
英文 / English	4.14	182.35
葡萄牙语 / Portuguese	25050.81	5.34

# Portuguese PPL before/after comparison bar chart
x = range(len(df_pt_ppl))
width = 0.35

plt.figure(figsize=(7, 5))
plt.bar([i - width / 2 for i in x], df_pt_ppl["before_training"],
        width, label="Before", color="#4C72B0")
plt.bar([i + width / 2 for i in x], df_pt_ppl["after_pt_training"],
        width, label="After (Portuguese training)", color="#55A868")
plt.xticks(x, df_pt_ppl["language"])
plt.yscale("log")
plt.ylabel("Perplexity (log scale)")
plt.title("PPL before vs after Portuguese continual pre-training")
plt.legend()
plt.tight_layout()
plt.show()

图表对比如下：

13、葡语生成测试

# Portuguese prompt generation
prompt = "Era uma vez, um gato chamado Tom."  # Once upon a time, a cat named Tom
max_new_tokens = 300
do_sample = True
temperature = 0.3

tokenized_input = tokenizer.encode(prompt, return_tensors="pt").to(device)
output_ids = model_pt.generate(
    tokenized_input,
    max_new_tokens=max_new_tokens,
    eos_token_id=1,
    do_sample=do_sample,
    temperature=temperature,
)
print(tokenizer.decode(output_ids[0]))

输出如下：

<s> Era uma vez, um gato chamado Tom. Tom adorava brincar com seus brinquedos. Um dia, Tom viu uma grande caixa. Ele queria brincar com ela, mas não conseguiu encontrar a caixa.
Tom começou a encontrar uma maneira de brincar com a caixa. Ele estava muito feliz. Ele brincava com a maneira de brincar com a caixa. Tom brincava com a maneira de brincar o dia todo.
Mas então, algo inesperado aconteceu. A caixa caiu do chão e partiu em pedaços. Tom estava triste. Ele não queria brincar com a maneira de brincar. Tom aprendeu a lição. Ele aprendeu que é importante ser gentil com os outros.<s>

七、实验总结

本实验围绕一个42M参数的玩具级Llama-2模型，完成了解码、困惑度评估与继续预训练三项任务。

在解码任务中，通过4种提示词与7种解码配置的28组实验发现，温度0.6配合top_k=50能在多样性与连贯性之间取得最佳平衡；同时中文提示词会让模型输出几个汉字后立即切回英文，暴露出基座模型的英文单一性。

在困惑度评估中，模型在英文测试集上PPL仅为4.14，但中文和葡萄牙语分别高达70030.35和25050.81，多领域结果也表明模型在训练分布TinyStories上表现最好，分布外数据则显著退化。

在继续预训练阶段，使用学习率1e-4、批大小16、训练4个epoch的最优配置后，中文PPL降至3.58，葡萄牙语降至5.34，但英文PPL分别升至304.68和182.35，出现严重的灾难性遗忘，说明42M的模型容量不足以同时容纳多种语言。

整体而言，本实验完整展现了小模型在多语言场景下"学新忘旧"的根本局限。