基于昇思 MindSpore 开展大模型无监督预训练,完整链路包含数据集处理、Transformer 模型搭建、分布式并行配置、优化器与学习率调度、预训练循环、断点续训。大规模无标注预训练依靠海量无标注文本,通过自回归语言建模(AR‑LM)让模型学习语言知识,本文基于 MindSpore 2.3,支持自动并行,适配昇腾 910B 硬件,完整展示预训练核心流程。

一、整体流程概述

无标注原始文本预处理、分词、构造自回归训练样本;

搭建 Decoder‑Only Transformer 大模型;

配置 MindSpore 分布式自动并行策略;

定义损失函数、优化器、余弦学习率预热衰减;

训练循环执行预训练,支持 Checkpoint 断点保存加载;

训练监控、loss 打印,完成大规模无标注预训练。

说明:示例为核心极简可运行代码,真实工业级大模型还需结合 MindFormers 套件、MindRecord 大格式数据集、FlashAttention、MoE、重计算、梯度累积。

二、完整核心代码

python

运行

import mindspore as ms

import mindspore.nn as nn

from mindspore import Tensor, Model

from mindspore.dataset import GeneratorDataset

from mindspore.communication import init, get_rank, get_group_size

from mindspore.nn.optim import AdamW

from mindspore.nn.loss import CrossEntropyLoss

import numpy as np

# ===================== 1. 初始化分布式环境(昇腾多卡预训练) =====================

ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")

init()

rank_id = get_rank()

world_size = get_group_size()

# 开启MindSpore自动并行,自动切分模型参数、数据并行

ms.set_auto_parallel_context(

parallel_mode=ms.ParallelMode.AUTO_PARALLEL,

gradients_mean=True,

full_batch=True

)

# ===================== 2. 大规模无标注数据集模拟:自回归LM样本 =====================

VOCAB_SIZE = 32000

SEQ_LEN = 1024

# 模拟海量无标注文本数据集,实际业务替换为MindRecord读取海量语料

class TextGenDataset:

def __init__(self, dataset_size=1000000):

self.dataset_size = dataset_size

def __getitem__(self, idx):

# input_ids:原始文本token;labels向右移位,自回归预测下一个token

input_ids = np.random.randint(0, VOCAB_SIZE, size=(SEQ_LEN,), dtype=np.int32)

labels = np.copy(input_ids)

return input_ids, labels

def __len__(self):

return self.dataset_size

def create_pretrain_ds(batch_size=8):

dataset = GeneratorDataset(TextGenDataset(), column_names=["input_ids", "labels"])

dataset = dataset.batch(batch_size, drop_remainder=True)

return dataset

# ===================== 3. Decoder‑Only Transformer模块搭建 =====================

class MLP(nn.Cell):

def __init__(self, hidden_dim, ffn_dim):

super().__init__()

self.up_proj = nn.Dense(hidden_dim, ffn_dim)

self.down_proj = nn.Dense(ffn_dim, hidden_dim)

self.act = nn.GELU()

def construct(self, x):

x = self.up_proj(x)

x = self.act(x)

x = self.down_proj(x)

return x

class DecoderLayer(nn.Cell):

def __init__(self, hidden_dim, head_num, ffn_dim):

super().__init__()

self.norm1 = nn.LayerNorm((hidden_dim,))

self.attn = nn.MultiheadAttention(hidden_dim, head_num)

self.norm2 = nn.LayerNorm((hidden_dim,))

self.mlp = MLP(hidden_dim, ffn_dim)

def construct(self, x, mask):

residual = x

x = self.norm1(x)

attn_out, _ = self.attn(x, x, x, mask)

x = residual + attn_out

residual = x

x = self.norm2(x)

mlp_out = self.mlp(x)

x = residual + mlp_out

return x

class DecoderOnlyModel(nn.Cell):

def __init__(self, vocab_size, hidden_dim=1024, head_num=16, layer_num=24, ffn_dim=4096):

super().__init__()

self.hidden_dim = hidden_dim

self.embedding = nn.Embedding(vocab_size, hidden_dim)

self.position_embedding = nn.Embedding(SEQ_LEN, hidden_dim)

self.decoder_layers = nn.CellList([DecoderLayer(hidden_dim, head_num, ffn_dim) for _ in range(layer_num)])

self.norm_final = nn.LayerNorm((hidden_dim,))

self.lm_head = nn.Dense(hidden_dim, vocab_size)

def construct(self, input_ids):

b, s = input_ids.shape

pos = Tensor(np.arange(0, s), ms.int32)

hidden = self.embedding(input_ids) + self.position_embedding(pos)

# 构造下三角因果mask,防止看到未来token

mask = Tensor(np.triu(np.ones((s, s)), k=1), ms.bool_)

for layer in self.decoder_layers:

hidden = layer(hidden, mask)

hidden = self.norm_final(hidden)

logits = self.lm_head(hidden)

return logits

# ===================== 4. 预训练网络:损失封装 =====================

class PretrainNetWithLoss(nn.Cell):

def __init__(self, backbone):

super().__init__()

self.backbone = backbone

self.loss_fn = CrossEntropyLoss()

def construct(self, input_ids, labels):

logits = self.backbone(input_ids)

logits = logits.reshape((-1, VOCAB_SIZE))

labels = labels.reshape((-1,))

loss = self.loss_fn(logits, labels)

return loss

# ===================== 5. 优化器、学习率调度、训练执行 =====================

def run_pretrain():

batch_size = 8

lr_max = 1e-4

warmup_step = 2000

total_step = 100000

model = DecoderOnlyModel(vocab_size=VOCAB_SIZE)

net_loss = PretrainNetWithLoss(model)

# 学习率:warmup + 余弦衰减

lr_schedule = nn.CosineDecayLR(warmup_steps=warmup_step, decay_steps=total_step, max_lr=lr_max)

optimizer = AdamW(net_loss.trainable_params(), learning_rate=lr_schedule, weight_decay=0.01)

train_net = nn.TrainOneStepCell(net_loss, optimizer)

train_net.set_train(True)

dataset = create_pretrain_ds(batch_size)

ds_iter = dataset.create_tuple_iterator()

for step in range(1, total_step+1):

input_ids, labels = next(ds_iter)

loss = train_net(input_ids, labels)

if rank_id == 0 and step % 100 == 0:

print(f"step: {step}, loss: {loss.asnumpy():.4f}")

# 每5000步保存checkpoint断点

if step % 5000 == 0 and rank_id == 0:

ms.save_checkpoint(model, f"./ckpt/pretrain_step_{step}.ckpt")

if __name__ == "__main__":

run_pretrain()

三、代码要点解析

运行模式与分布式

GRAPH_MODE静态图模式是大规模预训练必备,性能远高于 PYNATIVE。AUTO_PARALLEL自动并行自动完成数据并行、模型并行混合切分,不需要手动编写复杂并行逻辑;多卡环境必须执行init()初始化通信域。

无标注数据处理

大规模预训练不使用人工标签,采用自回归建模,输入文本右移一位作为预测标签。工业场景不能用 GeneratorDataset,应使用 MindRecord 格式存储 TB 级无标注语料,配合分词器做文本切分、过滤、去重。

因果掩码

Decoder‑Only 模型必须构造下三角 mask,保证每个 token 只能看到历史位置,看不到未来 token,是自回归预训练的核心。

训练封装

TrainOneStepCell封装单步训练,自动完成前向、损失计算、反向传播、参数更新。AdamW + 余弦学习率 + warmup 是大模型预训练标准配置,weight_decay 抑制过拟合。

断点续训

通过save_checkpoint保存,load_checkpoint加载,中断预训练后可继续训练;多卡场景仅 rank0 保存 checkpoint,避免重复 IO。

四、大规模预训练优化方向

上述代码是最小原型。真实大模型预训练,需要基于 MindFormers 套件:开启 Flash‑Attention 降低算力开销;启用梯度重计算节省显存;梯度累积模拟更大 batch;混合精度ms.set_amp_level;开启 MoE 稀疏架构;数据侧做数据去重、质量过滤;配合昇腾 910B 集群执行分布式预训练。

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐