【在 RX 6750 GRE 10GB 上预训练 GPT：一场 ROCm 生态的踩坑实录

在 RX 6750 GRE 10GB 上预训练 GPT：一场 ROCm 生态的踩坑实录

• 导言
• • 一、为什么做这件事
  • 二、硬件与环境
  • [三、模型设计：50M 参数的 Mini-GPT](#三、模型设计：50M 参数的 Mini-GPT)
  • 四、技术栈：把能用的优化全用上
  • 五、数据集：wikitext-2，小而精五、数据集：wikitext-2，小而精
  • [六、 训练日志与效果](#六、 训练日志与效果)
  • [7、 踩坑记录：那些只有跑过才知道的事](#7、 踩坑记录：那些只有跑过才知道的事)
  • 附录：核心代码

导言

用一张消费级 A 卡（RX 6750 GRE 10GB）在 Linux + ROCm 环境下成功预训练了一个 50M 参数的 GPT 模型。本文记录了从环境配置、ROCm 兼容性伪装、数据集选型到显存优化的完整过程，以及那些只有亲自踩过才知道的坑。

一、为什么做这件事

大模型训练似乎已经成了 NVIDIA 的专属游戏------A100/H100 一卡难求，消费级 RTX 4090 也要被 CUDA 生态绑定，而且现在NVIDIA的显卡涨价太厉害了。但手边恰好有一张 RX 6750 GRE 10GB（Navi 22, gfx1031），女朋友又送了我一张3T的机械硬盘，能不能用它跑起来？

AMD 的 ROCm 生态近年进步明显，PyTorch 官方也提供了 ROCm wheel。但消费级 RDNA2 显卡并不在 ROCm 的官方支持列表里。这篇文章记录的就是：如何让一张"不被支持"的 A 卡，稳定跑通 GPT 预训练。

二、硬件与环境

组件	配置
GPU	AMD Radeon RX 6750 GRE 10GB (Navi 22, gfx1031)
驱动	ROCm 6.x
系统	Ubuntu 22.04
PyTorch	2.x (ROCm 6.1 wheel)
显存	10 GB GDDR6

RX 6750 GRE 的架构是 gfx1031，ROCm 官方仅支持 gfx1030（RX 6700 XT）。直接运行会报 invalid device function。解决方案：环境变量伪装。

export HSA_OVERRIDE_GFX_VERSION=10.3.0

这行命令必须在 import torch 之前 执行，让 HIP runtime 把 gfx1031 识别为 gfx1030。绝大多数 kernel 可以兼容运行，但会触发一个无害的警告：

UserWarning: Attempting to use hipBLASLt on an unsupported architecture! 
Overriding blas backend to hipblas

三、模型设计：50M 参数的 Mini-GPT

10GB 显存无法支撑标准的 GPT-2 Small（124M），因此设计了一个轻量化架构：

参数	数值
层数	8
注意力头	8
嵌入维度	512
上下文长度	1024
词表大小	50304（GPT-2 BPE）
总参数量	~51M

在 BF16 + 梯度检查点 + SDPA 的组合下，实际显存占用非常健康：

[VRAM] 分配: 0.68 GB | 预留: 1.66 GB | 峰值: 1.44 GB

这意味着：

显存余量巨大：10GB 只用了不到 15%
理论上可以上更大的模型：355M（GPT-2 Medium）在 10GB 下也能跑，但 batch_size 必须降到 1
上下文可以更长：当前 1024，实测可以推到 2048 甚至 4096

四、技术栈：把能用的优化全用上

4.1 混合精度（BF16）

ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)

BF16 对 RDNA2 的兼容性比 FP16 更好，且不需要 loss scaling，训练更稳定。

4.2 PyTorch SDPA（替代 FlashAttention）

ROCm 上无法使用 CUDA 版的 FlashAttention-2，但 PyTorch 2.0+ 的 scaled_dot_product_attention 在 ROCm 下会自动调用 MIOpen memory efficient attention，实现类似的 O(n) 显存复杂度。

y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

4.3 梯度检查点（Gradient Checkpointing）

以时间换空间，只保存每层的输入，反向传播时重算中间激活。对于 8 层模型，显存节省约 60-70%。

x = x + checkpoint(lambda t: self.attn(self.ln_1(t)), x, use_reentrant=False)

五、数据集：wikitext-2，小而精五、数据集：wikitext-2，小而精

属性	数值
数据集	Salesforce/wikitext / wikitext-2-raw-v1
大小	~2 MB（极小，秒下）
语言	英文（维基百科）
训练样本	2,332 个（block_size=1024）
验证样本	241 个

六、 训练日志与效果

============================================================
  设备: AMD Radeon RX 6750 GRE 10GB
  显存: 10.0 GB | 架构: gfx103
============================================================
[!] 显存仅 10.0GB，强制 batch_size=1

[1/5] 加载 tiktoken...
    使用 tiktoken (gpt2)
[2/5] 加载数据集: Salesforce/wikitext/wikitext-2-raw-v1...
    [datasets] 加载成功, 样本数: 36718
    生成 2332 个训练样本
[3/5] 构建模型...
    参数量 : 50.98M
    上下文 : 1024
    层数   : 8 | 头数: 8 | 维度: 512
[4/5] 训练配置:
    微批次     : 1
    梯度累积   : 4
    有效批次   : 4
    混合精度   : bfloat16
    梯度检查点 : True
============================================================
iter      0 | loss 10.9044 | lr 0.00e+00 | 1759ms | mfu 3.7%
iter     10 | loss 10.6042 | lr 6.00e-06 | 11090ms | mfu 0.6%
iter     50 | loss 9.3630 | lr 3.00e-05 | 11116ms | mfu 0.6%
iter    100 | loss 8.3577 | lr 6.00e-05 | 11103ms | mfu 0.6%
iter    150 | loss 6.8392 | lr 9.00e-05 | 11125ms | mfu 0.6%
iter    200 | loss 6.9590 | lr 1.20e-04 | 11036ms | mfu 0.6%
iter    250 | loss 6.4357 | lr 1.50e-04 | 11017ms | mfu 0.6%
iter    300 | loss 6.3136 | lr 1.80e-04 | 11053ms | mfu 0.6%
iter    330 | loss 6.5860 | lr 1.92e-04 | 11024ms | mfu 0.6%

7、 踩坑记录：那些只有跑过才知道的事

7.1 invalid device function

原因：gfx1031 不被 ROCm 原生支持。

修复：export HSA_OVERRIDE_GFX_VERSION=10.3.0，且必须在 import torch 前设置。

7.2 ValueError: not enough values to unpack (expected 2, got 1)

原因：batch_size=1 时，next(dataset) 返回 (1024,) 一维张量，模型 forward 期望 (B, T) 二维。

修复：get_batch 函数使用 torch.stack 强制堆叠，确保 batch 维度始终存在。

7.3 KeyError: 'lm_head.weight'

原因：wte 和 lm_head 共享权重，named_parameters() 中只保留一个 key，但 configure_optimizers 通过 named_modules() 收集到了两个名字。

修复：与 param_dict 取交集，过滤掉因共享权重而不存在的 key。

7.4 Invalid HF URI 'hf://datasets/wikitext'

原因：huggingface_hub >= 1.16 强制要求 namespace/name 格式，拒绝单段 ID。

修复：wikitext → Salesforce/wikitext。

附录：核心代码

#!/usr/bin/env python3
"""
ROCm GPT-2 预训练脚本 (RX 6750 GRE 10GB, 最终稳定版)
=======================================================
数据集: Salesforce/wikitext (wikitext-2-raw-v1, 英文维基百科, ~2MB)
修复: huggingface_hub >= 1.16 要求 namespace/name 格式

安装:
    pip install torch --index-url https://download.pytorch.org/whl/rocm6.1
    pip install datasets tiktoken numpy tqdm

运行:
    export HSA_OVERRIDE_GFX_VERSION=10.3.0
    python train_gpt2_rocm_stable.py
"""

import os
import sys
import time
import math
import json
from dataclasses import dataclass, asdict
from contextlib import nullcontext

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint

# =====================================================================
# 0. 强制设置 HuggingFace 国内镜像
# =====================================================================
if not os.environ.get("HF_ENDPOINT"):
    os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
    print(f"[HF] 自动设置镜像: {os.environ['HF_ENDPOINT']}")

# =====================================================================
# 1. ROCm 环境自检
# =====================================================================
def rocm_sanity_check():
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA/ROCm 不可用")
    props = torch.cuda.get_device_properties(0)
    total_mem = props.total_memory / 1024**3
    print("=" * 60)
    print(f"  设备: {torch.cuda.get_device_name(0)}")
    print(f"  显存: {total_mem:.1f} GB | 架构: gfx{props.major}{props.minor}")
    print("=" * 60)
    return total_mem

# =====================================================================
# 2. 配置 (~50M 参数小模型)
# =====================================================================
@dataclass
class TrainConfig:
    n_layer: int = 8
    n_head: int = 8
    n_embd: int = 512
    block_size: int = 1024
    vocab_size: int = 50304
    dropout: float = 0.0
    bias: bool = True

    batch_size: int = 2
    gradient_accumulation_steps: int = 4
    max_iters: int = 50000
    learning_rate: float = 6e-4
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    warmup_iters: int = 1000
    lr_decay_iters: int = 50000
    min_lr: float = 6e-5

    use_mixed_precision: bool = True
    use_gradient_checkpointing: bool = True
    use_sdpa: bool = True

    # =====================================================================
    # 关键修复: 数据集名称必须用 namespace/name 格式
    # huggingface_hub >= 1.16 拒绝单段 ID (如 "wikitext")
    # =====================================================================
    dataset: str = "Salesforce/wikitext"      # 修复: 加命名空间
    dataset_config: str = "wikitext-2-raw-v1"  # 子集配置名
    text_field: str = "text"

    eval_interval: int = 500
    eval_iters: int = 100
    log_interval: int = 10
    checkpoint_interval: int = 2500
    sample_interval: int = 1000
    out_dir: str = "out_rocm_gpt2"
    device: str = "cuda"
    dtype: str = "bfloat16"
    seed: int = 1337

# =====================================================================
# 3. 模型定义
# =====================================================================
class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.register_buffer("mask", torch.triu(torch.ones(config.block_size, config.block_size), diagonal=1).bool(), persistent=False)
        self.use_sdpa = config.use_sdpa and hasattr(F, "scaled_dot_product_attention")

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        if self.use_sdpa:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0, is_causal=True)
        else:
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
            att = att.masked_fill(self.mask[:, :, :T, :T], float("-inf"))
            att = F.softmax(att, dim=-1)
            y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
    def forward(self, x):
        return self.c_proj(F.gelu(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)
        self.use_ckpt = config.use_gradient_checkpointing

    def forward(self, x):
        if self.use_ckpt and self.training:
            x = x + checkpoint(lambda t: self.attn(self.ln_1(t)), x, use_reentrant=False)
            x = x + checkpoint(lambda t: self.mlp(self.ln_2(t)), x, use_reentrant=False)
        else:
            x = x + self.attn(self.ln_1(x))
            x = x + self.mlp(self.ln_2(x))
        return x

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight
        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            if pn.endswith("c_proj.weight"):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device).unsqueeze(0)
        x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos))
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        return logits, loss

    def configure_optimizers(self, config):
        decay, no_decay = set(), set()
        whitelist, blacklist = (nn.Linear,), (nn.LayerNorm, nn.Embedding)
        for mn, m in self.named_modules():
            for pn, p in m.named_parameters(recurse=False):
                fpn = f"{mn}.{pn}" if mn else pn
                if pn.endswith('bias'):
                    no_decay.add(fpn)
                elif pn.endswith('weight') and isinstance(m, whitelist):
                    decay.add(fpn)
                elif pn.endswith('weight') and isinstance(m, blacklist):
                    no_decay.add(fpn)
        param_dict = {pn: p for pn, p in self.named_parameters()}
        decay = {pn for pn in decay if pn in param_dict}
        no_decay = {pn for pn in no_decay if pn in param_dict}
        inter = decay & no_decay
        union = decay | no_decay
        assert len(inter) == 0
        assert len(param_dict.keys() - union) == 0
        optim_groups = [
            {"params": [param_dict[pn] for pn in sorted(decay)], "weight_decay": config.weight_decay},
            {"params": [param_dict[pn] for pn in sorted(no_decay)], "weight_decay": 0.0},
        ]
        return torch.optim.AdamW(optim_groups, lr=config.learning_rate, betas=(config.beta1, config.beta2))

    def get_num_params(self, non_embedding=True):
        n = sum(p.numel() for p in self.parameters())
        if non_embedding:
            n -= self.transformer.wpe.weight.numel()
        return n

    def estimate_mfu(self, fwdbwd_per_iter, dt):
        N = self.get_num_params()
        cfg = self.config
        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size
        flops_per_token = 6 * N + 12 * L * H * Q * T
        flops_per_fwdbwd = flops_per_token * T
        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
        return flops_per_iter * (1.0 / dt) / 22.6e12

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

# =====================================================================
# 4. 数据加载 (datasets + hf-mirror, namespace/name 格式)
# =====================================================================
class WikitextDataset:
    """加载 wikitext 数据集, 使用 namespace/name 格式"""
    def __init__(self, dataset_name, config_name, text_field, tokenizer, block_size, split="train"):
        from datasets import load_dataset
        print(f"    [datasets] 加载 {dataset_name}/{config_name} ({split})...")
        print(f"    [datasets] 镜像: {os.environ.get('HF_ENDPOINT', '默认')}")

        try:
            ds = load_dataset(dataset_name, config_name, split=split, trust_remote_code=False)
            print(f"    [datasets] 加载成功, 样本数: {len(ds)}")
        except Exception as e:
            print(f"\n[!] 加载失败: {e}")
            print("\n    可能原因:")
            print("    1. 未安装 datasets: pip install datasets")
            print("    2. 网络问题, 检查是否能访问 hf-mirror.com")
            print("    3. datasets 版本过旧/过新，尝试: pip install -U datasets huggingface_hub")
            raise

        self.samples = []
        self.tokenizer = tokenizer
        self.block_size = block_size

        buffer = []
        for ex in ds:
            text = ex.get(text_field, "") if isinstance(ex, dict) else getattr(ex, text_field, "")
            if not text or len(text.strip()) < 10:
                continue

            ids = tokenizer.encode(text, allowed_special={"<|endoftext|>"})
            buffer.extend(ids)

            while len(buffer) >= block_size + 1:
                chunk = buffer[:block_size + 1]
                buffer = buffer[block_size + 1:]
                self.samples.append((
                    torch.tensor(chunk[:-1], dtype=torch.long),
                    torch.tensor(chunk[1:], dtype=torch.long)
                ))

        print(f"    生成 {len(self.samples)} 个训练样本")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx]

def get_batch(ds, batch_size, device):
    """确保 batch 维度始终存在"""
    ix = torch.randint(0, len(ds), (batch_size,))
    x = torch.stack([ds[i][0] for i in ix])
    y = torch.stack([ds[i][1] for i in ix])
    return x.to(device), y.to(device)

# =====================================================================
# 5. 工具函数
# =====================================================================
def get_lr(it, config):
    if it < config.warmup_iters:
        return config.learning_rate * it / config.warmup_iters
    if it > config.lr_decay_iters:
        return config.min_lr
    decay_ratio = (it - config.warmup_iters) / (config.lr_decay_iters - config.warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return config.min_lr + coeff * (config.learning_rate - config.min_lr)

def print_memory(tag=""):
    alloc = torch.cuda.memory_allocated() / 1024**3
    res = torch.cuda.memory_reserved() / 1024**3
    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"  [VRAM {tag}] 分配: {alloc:.2f} GB | 预留: {res:.2f} GB | 峰值: {peak:.2f} GB")

# =====================================================================
# 6. 主训练循环
# =====================================================================
def main():
    config = TrainConfig()
    os.makedirs(config.out_dir, exist_ok=True)

    with open(os.path.join(config.out_dir, "config.json"), "w") as f:
        json.dump(asdict(config), f, indent=2)

    torch.manual_seed(config.seed)
    torch.cuda.manual_seed(config.seed)

    total_mem = rocm_sanity_check()
    device = torch.device(config.device)

    if total_mem < 12 and config.batch_size > 1:
        print(f"[!] 显存仅 {total_mem:.1f}GB，强制 batch_size=1")
        config.batch_size = 1

    # 加载 tokenizer
    print("\n[1/5] 加载 tiktoken...")
    try:
        import tiktoken
        enc = tiktoken.get_encoding("gpt2")
        print("    使用 tiktoken (gpt2)")
    except ImportError:
        print("[!] 未安装 tiktoken: pip install tiktoken")
        sys.exit(1)

    # 加载数据集
    print(f"[2/5] 加载数据集: {config.dataset}/{config.dataset_config}...")
    try:
        train_ds = WikitextDataset(config.dataset, config.dataset_config, config.text_field, enc, config.block_size, "train")
        val_ds = WikitextDataset(config.dataset, config.dataset_config, config.text_field, enc, config.block_size, "validation")
    except Exception as e:
        print(f"\n[!] 加载数据集失败: {e}")
        sys.exit(1)

    # 构建模型
    print("[3/5] 构建模型...")
    model = GPT(config)
    model = model.to(device)
    n_params = model.get_num_params()
    print(f"    参数量 : {n_params / 1e6:.2f}M")
    print(f"    上下文 : {config.block_size}")
    print(f"    层数   : {config.n_layer} | 头数: {config.n_head} | 维度: {config.n_embd}")

    optimizer = model.configure_optimizers(config)

    ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[config.dtype]
    ctx = nullcontext() if not config.use_mixed_precision else torch.amp.autocast(device_type="cuda", dtype=ptdtype)

    print(f"[4/5] 训练配置:")
    print(f"    微批次     : {config.batch_size}")
    print(f"    梯度累积   : {config.gradient_accumulation_steps}")
    print(f"    有效批次   : {config.batch_size * config.gradient_accumulation_steps}")
    print(f"    混合精度   : {config.dtype}")
    print(f"    梯度检查点 : {config.use_gradient_checkpointing}")
    print("=" * 60)

    iter_num = 0
    best_val_loss = 1e9
    t0 = time.time()
    model.train()

    X, Y = get_batch(train_ds, config.batch_size, device)

    while iter_num < config.max_iters:
        lr = get_lr(iter_num, config)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        for micro_step in range(config.gradient_accumulation_steps):
            with ctx:
                logits, loss = model(X, Y)
                loss = loss / config.gradient_accumulation_steps

            loss.backward()
            X, Y = get_batch(train_ds, config.batch_size, device)

        torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

        if iter_num % config.log_interval == 0:
            dt = time.time() - t0
            t0 = time.time()
            lossf = loss.item() * config.gradient_accumulation_steps
            mfu = model.estimate_mfu(config.batch_size * config.gradient_accumulation_steps, dt)
            print(f"iter {iter_num:6d} | loss {lossf:.4f} | lr {lr:.2e} | {dt*1000:.0f}ms | mfu {mfu*100:.1f}%")

        if iter_num % (config.log_interval * 10) == 0:
            print_memory(f"iter {iter_num}")

        if iter_num > 0 and iter_num % config.eval_interval == 0:
            model.eval()
            losses = torch.zeros(config.eval_iters)
            with torch.no_grad():
                for k in range(config.eval_iters):
                    xv, yv = get_batch(val_ds, config.batch_size, device)
                    with ctx:
                        _, l = model(xv, yv)
                    losses[k] = l.item()
            model.train()
            val_loss = losses.mean()
            print(f"\n>>> 验证损失: {val_loss:.4f} <<<")

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                ckpt = {"model": model.state_dict(), "optimizer": optimizer.state_dict(),
                        "iter_num": iter_num, "config": asdict(config)}
                torch.save(ckpt, os.path.join(config.out_dir, "ckpt_best.pt"))
                print(f"    [保存] 最佳模型")

        if iter_num > 0 and iter_num % config.sample_interval == 0:
            model.eval()
            prompt = torch.zeros((1, 1), dtype=torch.long, device=device)
            with torch.no_grad():
                with ctx:
                    sample = model.generate(prompt, max_new_tokens=64, top_k=40)[0].tolist()
            text = enc.decode(sample)
            print(f"\n--- 生成样本 (iter {iter_num}) ---")
            print(text[:300])
            print("---\n")
            model.train()

        if iter_num > 0 and iter_num % config.checkpoint_interval == 0:
            ckpt = {"model": model.state_dict(), "optimizer": optimizer.state_dict(),
                    "iter_num": iter_num, "config": asdict(config)}
            torch.save(ckpt, os.path.join(config.out_dir, f"ckpt_{iter_num}.pt"))

        iter_num += 1

    print("\n[5/5] 训练完成!")
    print_memory("最终")
    print(f"最佳验证损失: {best_val_loss:.4f}")

if __name__ == "__main__":
    main()