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

臭小子2222026-08-01 11:58

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

导言

用一张消费级 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。解决方案：环境变量伪装。

bash 复制代码 
export HSA_OVERRIDE_GFX_VERSION=10.3.0

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

bash 复制代码 
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 的组合下，实际显存占用非常健康：

bash 复制代码 
[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）

python 复制代码 
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) 显存复杂度。

python 复制代码 
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

4.3 梯度检查点（Gradient Checkpointing）

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

python 复制代码 
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 个

六、 训练日志与效果

python 复制代码 
============================================================
  设备: 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。

附录：核心代码

python 复制代码 
#!/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()
相关推荐
AAASilverwolf11 小时前
OpenAI 新增 12 小时提示缓存：跨小时任务能省多少？
人工智能·gpt·缓存·api
Caroline51615 小时前
AI Logo 设计实战：如何用 GPT Image 2.5 生成透明背景 Logo？附 Prompt 示例
人工智能·gpt·prompt·图像生成
Caroline51616 小时前
GPT Image 2.5 + DeepSeek V4.1 Flash-1017：从游戏 UI 设计到 3D 项目生成
gpt·游戏·ui·ai 编程·gpt image 2.5
变量棱镜17 小时前
ChatGPT 遭到起诉，被控伤害儿童和其他脆弱用户
gpt·ai·openai
cc.ChenLy1 天前
Codex 无法安装的终极解决方案:应用安装失败，错误消息: 错误 0x8007177E
gpt·codex·安装codex
TMT星球1 天前
六项评测超越GPT-6 Astra！大晓ACE-Brain-0.6以数据Scaling推动具身空间智能跃迁
gpt
梦想的初衷~2 天前
GPT × ArcGIS 强强联合：AI助力数据处理、空间分析、可视化及多案例综合应用实战教程
人工智能·gpt·arcgis
BUTFL2 天前
一次 429 打挂全站，我连夜重写了整个熔断器
gpt·claude
AI砖家2 天前
DeepSeek V4.1 Flash vs. GPT-5.6 Sol：前沿闭源与高效开源，路线已经分叉
gpt·开源·deepseek
小粉粉hhh3 天前
大模型RAG和Agent智能体（一）——初了解
gpt·langchain·aigc
热门推荐
01ai安全工具:CyberStrikeAI安装部署与使用02国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）03如何新建文件夹？ 电脑新建文件夹的4种方法04GitHub 镜像站点05历年考研数学一、数学二、数学三真题试卷及答案PDF06几个好用的ip纯净度检测网站07乙二醇水溶液物性参数08昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战09「全语录+省流版」邵艾伦 × 孙宇晨：聊风口、聊个人IP、聊年轻人如何抓住 AI 时代的机会10苹果iPA游戏软件资源下载网站