首先，创建一个名为llm_tarin的容器

docker run --gpus all -dit --name llm_train -p 2668:22 -v /home/kemove/workspace:/root/workdir base_image:ubuntu_22.04_cuda12.8_pytorch_2.8.0 /bin/bash

然后，使用如下命令安装LLM训练需要用到依赖库

pip install torch torchvision transformers -i https://pypi.tuna.tsinghua.edu.cn/simple/

pip install deepspeed accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple/

pip install datasets tokenizers -i https://pypi.tuna.tsinghua.edu.cn/simple/

然后，使用如下代码构建简单的LLM模型并进行训练，代码中包含数据集的加载（r拖）

# 导入必要的库
# 导入必要的库
import torch
from torch.optim import AdamW
from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    AutoTokenizer,
    get_scheduler,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm
import math

# 确保使用 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 1. 加载数据集
# 使用 Hugging Face 的 WikiText-2 数据集作为示例
print("Loading dataset...")
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")  # WikiText-2 是一个小型语言建模数据集

# 2. 加载分词器
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# 为分词器添加一个 pad_token
tokenizer.pad_token = tokenizer.eos_token

# 定义分词函数
def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",  # 填充到固定长度
        max_length=1024,  # 最大长度为 1024
    )

# 对数据集进行分词
print("Tokenizing dataset...")
tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 使用 DataCollator 自动处理填充和张量转换
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=False  # 自回归模型，不使用掩码语言模型
)

# 创建 DataLoader
train_dataloader = DataLoader(
    tokenized_datasets["train"], batch_size=8, shuffle=True, collate_fn=data_collator
)
valid_dataloader = DataLoader(
    tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator
)

# 3. 构建模型
print("Initializing model...")
config = GPT2Config(
    vocab_size=50257,  # GPT-2 默认词汇表大小
    n_positions=1024,  # 最大上下文长度
    n_embd=768,  # 嵌入维度
    n_layer=12,  # Transformer 层数
    n_head=12,  # 注意力头数
)
model = GPT2LMHeadModel(config)
model.to(device)

# 4. 配置优化器和学习率调度器
print("Setting up optimizer and scheduler...")
optimizer = AdamW(model.parameters(), lr=5e-4)  # AdamW 优化器
num_training_steps = len(train_dataloader) * 3  # 假设训练 3 个 epoch
lr_scheduler = get_scheduler(
    "linear", optimizer=optimizer, num_warmup_steps=500, num_training_steps=num_training_steps
)

# 5. 训练模型
print("Starting training...")
num_epochs = 3
model.train()

for epoch in range(num_epochs):
    print(f"Epoch {epoch + 1}/{num_epochs}")
    progress_bar = tqdm(train_dataloader)

    for batch in progress_bar:
        # 数据已经是张量，直接移动到 GPU
        inputs = batch["input_ids"].to(device)
        labels = batch["input_ids"].to(device)  # 自回归训练，输入即标签

        # 前向传播
        outputs = model(inputs, labels=labels)
        loss = outputs.loss

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()

        # 更新进度条
        progress_bar.set_description(f"Loss: {loss.item():.4f}")

# 6. 保存模型
print("Saving model...")
model.save_pretrained("./my_small_gpt_model")
tokenizer.save_pretrained("./my_small_gpt_model")

# 7. 生成文本
print("Generating text...")
from transformers import pipeline

generator = pipeline("text-generation", model="./my_small_gpt_model", tokenizer="./my_small_gpt_model")
generated_text = generator("今天是一个好日子", max_length=50)
print("Generated Text:")
print(generated_text)

# 8. 计算困惑度（Perplexity）
print("Calculating perplexity...")
model.eval()  # 切换到评估模式
total_loss = 0
num_batches = 0

for batch in valid_dataloader:
    inputs = batch["input_ids"].to(device)
    labels = batch["input_ids"].to(device)

    with torch.no_grad():
        outputs = model(inputs, labels=labels)
        loss = outputs.loss
        total_loss += loss.item()  # 累加损失
        num_batches += 1

# 计算平均损失和困惑度
avg_loss = total_loss / num_batches
perplexity = math.exp(avg_loss)
print(f"验证集困惑度: {perplexity:.2f}")

使用如下代码进行上下文推理测试，输入为What is the capital of France?：

# 导入必要的库
from transformers import GPT2LMHeadModel, AutoTokenizer, pipeline
import torch

# 设置设备 (GPU 或 CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 1. 加载训练好的模型和分词器
model_path = "./my_small_gpt_model"  # 训练好的模型保存路径

print("Loading model and tokenizer...")
model = GPT2LMHeadModel.from_pretrained(model_path)  # 加载模型
tokenizer = AutoTokenizer.from_pretrained(model_path)  # 加载分词器

# 确保分词器包含 pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token  # 设置特殊 token 为 eos_token

# 将模型移动到设备 (GPU/CPU)
model.to(device)

# 2. 使用模型的 generate() 方法进行推理
print("Running inference using generate() method...")

# 输入提示文本
prompt = "What is the capital of France?"

# 将提示文本转化为 token IDs，并显式设置 attention_mask
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
inputs = {key: value.to(device) for key, value in inputs.items()}

# 使用模型生成文本
outputs = model.generate(
    inputs["input_ids"],                        # 输入 token IDs
    attention_mask=inputs["attention_mask"],    # 显式传递 attention_mask
    max_new_tokens=50,                          # 生成最多 50 个新 token
    temperature=0.7,                            # 降低随机性，生成更确定的答案
    top_k=50,                                   # 限制前 k 个概率最高的 token
    top_p=0.9,                                  # 核采样，保留累计概率达到 90% 的 token
    repetition_penalty=1.2,                     # 惩罚重复生成的 token
    do_sample=True,                             # 使用采样而非贪心搜索
    num_return_sequences=1,                     # 返回的生成序列数量
    pad_token_id=tokenizer.pad_token_id,        # 显式设置 pad_token_id
)

# 将生成的 token IDs 转化为文本
generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print("Generated Text (generate):")
print(generated_text)

# **单独提取回答部分**
# 假设答案是生成文本中第一个句子（以句号或问号为分隔符）
answer_part = generated_text[0].split(".")[0]  # 取第一个句子
print("Answer Part:")
print(answer_part)

# 3. 使用 Hugging Face 的 pipeline 进行快速问答推理（备用解决方案）
print("Running question-answering pipeline...")
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")

# 问答的上下文和问题
context = "France is a country in Europe. The capital of France is Paris. It is known for the Eiffel Tower."
question = "What is the capital of France?"

# 使用 pipeline 进行问答
qa_result = qa_pipeline(question=question, context=context)
print("Generated Answer (pipeline):")
print(qa_result)

# **单独提取问答结果的回答部分**
qa_answer_part = qa_result["answer"]
print("Answer Part (pipeline):")
print(qa_answer_part)

# 4. 可选：保存生成的文本到文件
output_file = "generated_text.txt"
with open(output_file, "w", encoding="utf-8") as f:
    f.write("Generated Text (generate):\n")
    f.write(str(generated_text) + "\n\n")
    f.write("Answer Part (generate):\n")
    f.write(answer_part + "\n\n")
    f.write("Generated Answer (pipeline):\n")
    f.write(str(qa_result) + "\n\n")
    f.write("Answer Part (pipeline):\n")
    f.write(qa_answer_part + "\n")
print(f"Generated text saved to {output_file}")

推理结果如下，输出为Paris